refactor: use form_with_unique_keys from awkward

pyproject.toml

@@ -36,7 +36,7 @@ classifiers = [
   "Topic :: Software Development",
 ]
 dependencies = [
-  "awkward >=2.5.1",
+  "awkward >=2.8.2",
   "dask >=2023.04.0,<2025.4.0",
   "cachetools",
   "typing_extensions >=4.8.0",

src/dask_awkward/lib/io/columnar.py

@@ -17,7 +17,6 @@
     METADATA_ATTRIBUTES,
     FormStructure,
     buffer_keys_required_to_compute_shapes,
-    form_with_unique_keys,
     parse_buffer_key,
     render_buffer_key,
     trace_form_structure,
@@ -80,7 +79,7 @@ def mock_empty(self: S, backend: BackendT = "cpu") -> AwkwardArray:
     def prepare_for_projection(
         self: S,
     ) -> tuple[AwkwardArray, TypeTracerReport, FormStructure]:
-        form = form_with_unique_keys(self.form, "@")
+        form = ak.forms.form_with_unique_keys(self.form, ("@",))
 
         # Build typetracer and associated report object
         (meta, report) = typetracer_with_report(

src/dask_awkward/lib/utils.py

@@ -6,8 +6,6 @@
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, TypedDict, TypeVar
 
-import awkward as ak
-
 if TYPE_CHECKING:
     from awkward.forms import Form
 
@@ -137,36 +135,6 @@ def parse_buffer_key(buffer_key: str) -> tuple[str, str]:
     return head, tail
 
 
-def form_with_unique_keys(form: Form, key: str) -> Form:
-    def impl(form: Form, key: str) -> None:
-        # Set form key
-        form.form_key = key
-
-        # If the form is a record we need to loop over all fields in the
-        # record and set form that include the field name; this will keep
-        # recursing as well.
-        if form.is_record:
-            for field in form.fields:
-                impl(form.content(field), f"{key}.{field}")
-
-        elif form.is_union:
-            for i, entry in enumerate(form.contents):
-                impl(entry, f"{key}#{i}")
-
-        # NumPy like array is easy
-        elif form.is_numpy or form.is_unknown:
-            pass
-
-        # Anything else grab the content and keep recursing
-        else:
-            impl(form.content, f"{key}.content")
-
-    # Perform a "deep" copy without preserving references
-    form = ak.forms.from_dict(form.to_dict())
-    impl(form, key)
-    return form
-
-
 @contextmanager
 def typetracer_nochecks():
     from awkward._nplikes.typetracer import TypeTracerArray

tests/test_inspect.py

@@ -18,9 +18,9 @@ def test_report_necessary_buffers(
         assert v == (
             frozenset(
                 {
-                    "@.points-offsets",
-                    "@.points.content.y-data",
-                    "@.points.content.x-data",
+                    repr(("@", "points")) + "-offsets",
+                    repr(("@", "points", None, "y")) + "-data",
+                    repr(("@", "points", None, "x")) + "-data",
                 }
             ),
             frozenset(),
@@ -31,7 +31,13 @@ def test_report_necessary_buffers(
     )
     for k, v in dak.report_necessary_buffers(w).items():
         assert v == (
-            frozenset({"@.points-offsets", "@.points.content.x-data"}),
+            # frozenset({"@.points-offsets", "@.points.content.x-data"}),
+            frozenset(
+                {
+                    repr(("@", "points")) + "-offsets",
+                    repr(("@", "points", None, "x")) + "-data",
+                }
+            ),
             frozenset(),
         )
 
@@ -40,9 +46,9 @@ def test_report_necessary_buffers(
         assert v == (
             frozenset(
                 {
-                    "@.points-offsets",
-                    "@.points.content.x-data",
-                    "@.points.content.y-data",
+                    repr(("@", "points")) + "-offsets",
+                    repr(("@", "points", None, "y")) + "-data",
+                    repr(("@", "points", None, "x")) + "-data",
                 }
             ),
             frozenset(),
@@ -51,8 +57,8 @@ def test_report_necessary_buffers(
     z = dak.zeros_like(daa.points.x)
     for k, v in dak.report_necessary_buffers(z).items():
         assert v == (
-            frozenset({"@.points-offsets"}),
-            frozenset({"@.points.content.x-data"}),
+            frozenset({repr(("@", "points")) + "-offsets"}),
+            frozenset({repr(("@", "points", None, "x")) + "-data"}),
         )
 
 

tests/test_operations.py

@@ -24,7 +24,8 @@ def test_concatenate_axis_0_logical_same(daa):
     buffers = next(iter(buffers_report.values()))
 
     assert buffers.data_and_shape == frozenset(
-        ["@.points.content.x-data", "@.points-offsets"]
+        # ["@.points.content.x-data", "@.points-offsets"]
+        [repr(("@", "points", None, "x")) + "-data", repr(("@", "points")) + "-offsets"]
     )
     assert buffers.shape_only == frozenset()
 
@@ -61,7 +62,12 @@ def test_concatenate_axis_0_logical_different(daa):
 
         buffers = next(iter(buffers_report.values()))
         assert buffers.data_and_shape == frozenset(
-            ["@.points.content.x-data", "@.points.content.y-data", "@.points-offsets"]
+            # ["@.points.content.x-data", "@.points.content.y-data", "@.points-offsets"]
+            [
+                repr(("@", "points", None, "x")) + "-data",
+                repr(("@", "points", None, "y")) + "-data",
+                repr(("@", "points")) + "-offsets",
+            ]
         )
         assert buffers.shape_only == frozenset()