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 <[email protected]>

Author: agoose77

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",

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:

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():

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():

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"]],
     )

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"]],
             )