Merge pull request #351 from dask-contrib/agoose77/refactor-highlevel…

…-typetracer

refactor: use `ak.typetracer.XX` high-level functions

src/dask_awkward/__init__.py

@@ -13,7 +13,6 @@
     compatible_partitions,
     map_partitions,
     partition_compatibility,
-    typetracer_from_form,
 )
 from dask_awkward.lib.describe import fields
 from dask_awkward.lib.inspect import necessary_columns, sample

src/dask_awkward/layers/layers.py

@@ -138,17 +138,12 @@ def mock(self) -> tuple[AwkwardInputLayer, TypeTracerReport]:
         from dask_awkward.lib._utils import set_form_keys
 
         starting_form = copy.deepcopy(self._meta.layout.form)
-        starting_layout = starting_form.length_zero_array(highlevel=False)
-        new_meta = ak.Array(
-            starting_layout.to_typetracer(forget_length=True),
-            behavior=self._behavior,
-        )
-        form = new_meta.layout.form
+        set_form_keys(starting_form, key=self.name)
 
-        set_form_keys(form, key=self.name)
+        new_meta_array, report = ak.typetracer.typetracer_with_report(
+            starting_form, highlevel=True, behavior=self._behavior
+        )
 
-        new_meta_labelled, report = ak.typetracer.typetracer_with_report(form)
-        new_meta_array = ak.Array(new_meta_labelled, behavior=self._behavior)
         new_input_layer = AwkwardInputLayer(
             name=self.name,
             columns=self.columns,

src/dask_awkward/lib/__init__.py

@@ -4,7 +4,6 @@
     compatible_partitions,
     map_partitions,
     partition_compatibility,
-    typetracer_from_form,
 )
 from dask_awkward.lib.describe import fields
 from dask_awkward.lib.inspect import necessary_columns, sample

src/dask_awkward/lib/core.py

@@ -2051,10 +2051,7 @@ def to_meta(objects):
 
 def length_zero_array_or_identity(obj: Any) -> Any:
     if is_awkward_collection(obj):
-        return ak.Array(
-            obj._meta.layout.form.length_zero_array(highlevel=False),
-            behavior=obj.behavior,
-        )
+        return ak.typetracer.length_zero_if_typetracer(obj._meta, behavior=obj.behavior)
     return obj
 
 
@@ -2090,7 +2087,7 @@ def map_meta(fn: ArgsKwargsPackedFunction, *deps: Any) -> ak.Array | None:
         pass
     try:
         arg_lzas = to_length_zero_arrays(deps)
-        meta = typetracer_from_form(fn(*arg_lzas).layout.form)
+        meta = ak.typetracer.typetracer_from_form(fn(*arg_lzas).layout.form)
         return meta
     except Exception:
         # if compute-unknown-meta is True and we've gotten to this
@@ -2204,24 +2201,6 @@ def normalize_single_outer_inner_index(
     return (int(partition_index), int(new_index))
 
 
-def typetracer_from_form(form: Form) -> ak.Array:
-    """Create a typetracer Array from an awkward form.
-
-    Parameters
-    ----------
-    form : awkward.form.Form
-        Form that the resulting Array will have.
-
-    Returns
-    -------
-    awkward.Array
-        Resulting highlevel typetracer Array
-
-    """
-    layout = form.length_zero_array(highlevel=False)
-    return ak.Array(layout.to_typetracer(forget_length=True))
-
-
 def make_unknown_length(array: ak.Array) -> ak.Array:
     """Make any highlevel Array a highlevel typetracer Array with unknown length.
 

src/dask_awkward/lib/io/io.py

@@ -418,7 +418,7 @@ def to_dataframe(
         **kwargs,
     )
     meta = ak.to_dataframe(
-        array._meta.layout.form.length_zero_array(highlevel=False), **kwargs
+        ak.typetracer.length_zero_if_typetracer(array._meta), **kwargs
     )
     return new_dd_object(
         intermediate.dask,

src/dask_awkward/lib/io/parquet.py

@@ -302,10 +302,7 @@ def from_parquet(
     if split_row_groups is None:
         split_row_groups = row_counts is not None and len(row_counts) > 1
 
-    meta = ak.Array(
-        subform.length_zero_array(highlevel=False).to_typetracer(forget_length=True),
-        behavior=behavior,
-    )
+    meta = ak.typetracer.typetracer_from_form(subform, behavior=behavior)
 
     if split_row_groups is False or subrg is None:
         # file-wise
@@ -324,7 +321,7 @@ def from_parquet(
             actual_paths,
             label=label,
             token=token,
-            meta=typetracer_array(meta),
+            meta=meta,
         )
     else:
         # row-group wise

tests/test_io.py

@@ -407,7 +407,7 @@ def test_from_awkward_empty_array(daa) -> None:
     assert len(a1) == 0
 
     # with a form
-    c2 = ak.Array(daa.layout.form.length_zero_array(highlevel=False))
+    c2 = ak.typetracer.length_zero_if_typetracer(daa.layout)
     assert len(c2) == 0
     a2 = dak.from_awkward(c2, npartitions=1)
     assert len(a2) == 0