From ac2317906ccb7cf04a04efbddc8a7f27bc29a249 Mon Sep 17 00:00:00 2001 From: Angus Hollands Date: Thu, 14 Sep 2023 19:04:35 +0200 Subject: [PATCH] refactor: use public typetracer API (#894) * refactor: use public typetracer API * refactor: use newer API * chore: bump dask-awkward version * refactor: use `ak.almost_equal` instead of dask-awkward's testutils --------- Co-authored-by: Jim Pivarski --- pyproject.toml | 2 +- src/uproot/_dask.py | 16 ++++++---------- src/uproot/extras.py | 8 +++++++- tests/test_0700-dask-empty-arrays.py | 9 +++------ .../test_0755-dask-awkward-column-projection.py | 9 ++++----- tests/test_0876-uproot-dask-blind-steps.py | 9 +++++---- 6 files changed, 26 insertions(+), 27 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 831c26577..78eba667e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,7 +50,7 @@ dynamic = [ [project.optional-dependencies] dev = [ "boost_histogram>=0.13", - "dask-awkward>=2022.12a3;python_version >= \"3.8\"", + "dask-awkward>=2023.9.0;python_version >= \"3.8\"", "dask[array];python_version >= \"3.8\"", "hist>=1.2", "pandas", diff --git a/src/uproot/_dask.py b/src/uproot/_dask.py index 315ae9fae..4503b9ff2 100644 --- a/src/uproot/_dask.py +++ b/src/uproot/_dask.py @@ -806,11 +806,9 @@ def project_columns(self, common_keys=None, original_form=None): if self.form_mapping is not None: awkward = uproot.extras.awkward() - ( - new_meta_labelled, - report, - ) = awkward._nplikes.typetracer.typetracer_with_report(self.rendered_form) - tt = awkward.Array(new_meta_labelled) + tt, report = awkward.typetracer.typetracer_with_report( + self.rendered_form, highlevel=True + ) if common_keys is not None: for key in common_keys: @@ -950,11 +948,9 @@ def project_columns(self, columns=None, original_form=None): if self.form_mapping is not None: awkward = uproot.extras.awkward() - ( - new_meta_labelled, - report, - ) = awkward._nplikes.typetracer.typetracer_with_report(self.rendered_form) - tt = awkward.Array(new_meta_labelled) + tt, report = awkward.typetracer.typetracer_with_report( + self.rendered_form, highlevel=True + ) if columns is not None: for key in columns: diff --git a/src/uproot/extras.py b/src/uproot/extras.py index 84aae0e0b..10b6c0970 100644 --- a/src/uproot/extras.py +++ b/src/uproot/extras.py @@ -325,8 +325,14 @@ def dask_awkward(): or conda install -c conda-forge dask dask-awkward""" ) from err - else: + if parse_version("2023.9.0") <= parse_version(dask_awkward.__version__): return dask_awkward + else: + raise ModuleNotFoundError( + "Uproot 5.x can only be used with dask-awkward 2023.9.0 or newer; you have dask-awkward {}".format( + dask_awkward.__version__ + ) + ) def awkward_pandas(): diff --git a/tests/test_0700-dask-empty-arrays.py b/tests/test_0700-dask-empty-arrays.py index becb0f2cf..3b07b2d7e 100644 --- a/tests/test_0700-dask-empty-arrays.py +++ b/tests/test_0700-dask-empty-arrays.py @@ -3,15 +3,12 @@ import numpy import pytest import skhep_testdata - +import awkward as ak import uproot dask = pytest.importorskip("dask") dask_awkward = pytest.importorskip("dask_awkward") -pytest.importorskip("pyarrow") # dask_awkward.lib.testutils needs pyarrow -from dask_awkward.lib.testutils import assert_eq - def test_dask_numpy_empty_arrays(): test_path = skhep_testdata.data_path("uproot-issue-697.root") + ":tree" @@ -52,7 +49,7 @@ def test_dask_awkward_empty_arrays(): ak_array = ttree.arrays() dak_array = uproot.dask(test_path, library="ak") - assert_eq(dak_array, ak_array) + assert ak.almost_equal(dak_array.compute(scheduler="synchronous"), ak_array) def test_dask_delayed_open_awkward(): @@ -62,7 +59,7 @@ def test_dask_delayed_open_awkward(): ak_array = ttree.arrays() dak_array = uproot.dask(test_path, library="ak", open_files=False) - assert_eq(dak_array, ak_array) + ak.almost_equal(dak_array.compute(scheduler="synchronous"), ak_array) def test_no_common_tree_branches(): diff --git a/tests/test_0755-dask-awkward-column-projection.py b/tests/test_0755-dask-awkward-column-projection.py index 37ca88933..070d2ad88 100644 --- a/tests/test_0755-dask-awkward-column-projection.py +++ b/tests/test_0755-dask-awkward-column-projection.py @@ -1,6 +1,7 @@ # BSD 3-Clause License; see https://github.com/scikit-hep/uproot5/blob/main/LICENSE import numpy +import awkward as ak import pytest import skhep_testdata @@ -9,9 +10,6 @@ dask = pytest.importorskip("dask") dask_awkward = pytest.importorskip("dask_awkward") -pytest.importorskip("pyarrow") # dask_awkward.lib.testutils needs pyarrow -from dask_awkward.lib.testutils import assert_eq - def test_column_projection_sanity_check(): test_path = skhep_testdata.data_path("uproot-Zmumu.root") + ":events" @@ -20,6 +18,7 @@ def test_column_projection_sanity_check(): ak_array = ttree.arrays() dak_array = uproot.dask(test_path, library="ak") - assert_eq( - dak_array[["px1", "px2", "py1", "py2"]], ak_array[["px1", "px2", "py1", "py2"]] + assert ak.almost_equal( + dak_array[["px1", "px2", "py1", "py2"]].compute(scheduler="synchronous"), + ak_array[["px1", "px2", "py1", "py2"]], ) diff --git a/tests/test_0876-uproot-dask-blind-steps.py b/tests/test_0876-uproot-dask-blind-steps.py index 678ef1c15..b742e5cbe 100644 --- a/tests/test_0876-uproot-dask-blind-steps.py +++ b/tests/test_0876-uproot-dask-blind-steps.py @@ -3,12 +3,11 @@ import numpy import pytest import skhep_testdata +import awkward as ak import uproot dask_awkward = pytest.importorskip("dask_awkward") -pytest.importorskip("pyarrow") # dask_awkward.lib.testutils needs pyarrow -from dask_awkward.lib.testutils import assert_eq @pytest.mark.parametrize("library", ["np", "ak"]) @@ -60,7 +59,9 @@ def test_uproot_dask_steps(library, step_size, steps_per_file, open_files): assert all(comp), f"Incorrect array at key {key}" else: - assert_eq( - dask_arrays[["px1", "px2", "py1", "py2"]], + assert ak.almost_equal( + dask_arrays[["px1", "px2", "py1", "py2"]].compute( + scheduler="synchronous" + ), arrays[["px1", "px2", "py1", "py2"]], )