diff --git a/pyproject.toml b/pyproject.toml index c7ea29d4..dd717bb8 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/dask_awkward/lib/io/columnar.py b/src/dask_awkward/lib/io/columnar.py index 7fa5be6d..85824bb3 100644 --- a/src/dask_awkward/lib/io/columnar.py +++ b/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( diff --git a/src/dask_awkward/lib/utils.py b/src/dask_awkward/lib/utils.py index 7b067386..6f44f90c 100644 --- a/src/dask_awkward/lib/utils.py +++ b/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 diff --git a/tests/test_inspect.py b/tests/test_inspect.py index b4e2651b..e47dd4d0 100644 --- a/tests/test_inspect.py +++ b/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"}), ) diff --git a/tests/test_operations.py b/tests/test_operations.py index 14a43fd9..f68e9674 100644 --- a/tests/test_operations.py +++ b/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()