Dask awkward support for uproot.dask

.github/workflows/build-test.yml

@@ -57,6 +57,7 @@ jobs:
         run: |
           conda env list
           conda install root
+          pip install dask-awkward
           conda list
 
       - name: "Install XRootD"

setup.py

@@ -28,6 +28,7 @@ def get_version():
         "xxhash",
         "requests",
         "dask[array]",
+        "dask-awkward",
     ],
 }
 extras["all"] = sum(extras.values(), [])

src/uproot/behaviors/TBranch.py

@@ -598,8 +598,8 @@ def dask(
     """
     files = uproot._util.regularize_files(files)
     library = uproot.interpretation.library._regularize_library(library)
-    if library.name != "np":
-        raise NotImplementedError()
+    # if library.name != "np":
+    #     raise NotImplementedError()
 
     real_options = dict(options)
     if "num_workers" not in real_options:
@@ -609,31 +609,60 @@ def dask(
 
     filter_branch = uproot._util.regularize_filter(filter_branch)
 
-    if open_files:
-        return _get_dask_array(
-            files,
-            filter_name,
-            filter_typename,
-            filter_branch,
-            recursive,
-            full_paths,
-            step_size,
-            custom_classes,
-            allow_missing,
-            real_options,
-        )
+    if library.name == "np":
+        if open_files:
+            return _get_dask_array(
+                files,
+                filter_name,
+                filter_typename,
+                filter_branch,
+                recursive,
+                full_paths,
+                step_size,
+                custom_classes,
+                allow_missing,
+                real_options,
+            )
+        else:
+            return _get_dask_array_delay_open(
+                files,
+                filter_name,
+                filter_typename,
+                filter_branch,
+                recursive,
+                full_paths,
+                custom_classes,
+                allow_missing,
+                real_options,
+            )
+    elif library.name == "ak":
+        if open_files:
+            return _get_dak_array(
+                files,
+                filter_name,
+                filter_typename,
+                filter_branch,
+                recursive,
+                full_paths,
+                step_size,
+                custom_classes,
+                allow_missing,
+                real_options,
+            )
+        else:
+            return _get_dak_array_delay_open(
+                files,
+                filter_name,
+                filter_typename,
+                filter_branch,
+                recursive,
+                full_paths,
+                custom_classes,
+                allow_missing,
+                real_options,
+            )
     else:
-        return _get_dask_array_delay_open(
-            files,
-            filter_name,
-            filter_typename,
-            filter_branch,
-            recursive,
-            full_paths,
-            custom_classes,
-            allow_missing,
-            real_options,
-        )
+        raise NotImplementedError()
 
 
 class Report:
@@ -3549,6 +3578,187 @@ def delayed_get_array(ttree, key):
     return dask_dict
 
 
+def _get_dak_array(
+    files,
+    filter_name=no_filter,
+    filter_typename=no_filter,
+    filter_branch=no_filter,
+    recursive=True,
+    full_paths=False,
+    step_size="100 MB",
+    custom_classes=None,
+    allow_missing=False,
+    real_options=None,
+):
+    import dask_awkward
+
+    hasbranches = []
+    common_keys = None
+    is_self = []
+
+    count = 0
+    for file_path, object_path in files:
+        obj = uproot._util.regularize_object_path(
+            file_path, object_path, custom_classes, allow_missing, real_options
+        )
+
+        if obj is not None:
+            count += 1
+
+            if isinstance(obj, TBranch) and len(obj.keys(recursive=True)) == 0:
+                original = obj
+                obj = obj.parent
+                is_self.append(True)
+
+                def real_filter_branch(branch):
+                    return branch is original and filter_branch(branch)
+
+            else:
+                is_self.append(False)
+                real_filter_branch = filter_branch
+
+            hasbranches.append(obj)
+
+            new_keys = obj.keys(
+                recursive=recursive,
+                filter_name=filter_name,
+                filter_typename=filter_typename,
+                filter_branch=real_filter_branch,
+                full_paths=full_paths,
+            )
+
+            if common_keys is None:
+                common_keys = new_keys
+            else:
+                new_keys = set(new_keys)
+                common_keys = [key for key in common_keys if key in new_keys]
+
+    if count == 0:
+        raise ValueError(
+            "allow_missing=True and no TTrees found in\n\n    {}".format(
+                "\n    ".join(
+                    "{"
+                    + "{}: {}".format(
+                        repr(f.file_path if isinstance(f, HasBranches) else f),
+                        repr(f.object_path if isinstance(f, HasBranches) else o),
+                    )
+                    + "}"
+                    for f, o in files
+                )
+            )
+        )
+
+    if len(common_keys) == 0 or not (all(is_self) or not any(is_self)):
+        raise ValueError(
+            "TTrees in\n\n    {}\n\nhave no TBranches in common".format(
+                "\n    ".join(
+                    "{"
+                    + "{}: {}".format(
+                        repr(f.file_path if isinstance(f, HasBranches) else f),
+                        repr(f.object_path if isinstance(f, HasBranches) else o),
+                    )
+                    + "}"
+                    for f, o in files
+                )
+            )
+        )
+
+    class _UprootRead:
+        def __init__(self, hasbranches, branches) -> None:
+            self.hasbranches = hasbranches
+            self.branches = branches
+
+        def __call__(self, i_start_stop):
+            i, start, stop = i_start_stop
+            return self.hasbranches[i].arrays(
+                self.branches, entry_start=start, entry_stop=stop
+            )
+
+    partition_args = []
+    for i, ttree in enumerate(hasbranches):
+        entry_start, entry_stop = _regularize_entries_start_stop(
+            ttree.num_entries, None, None
+        )
+        entry_step = 0
+        if uproot._util.isint(step_size):
+            entry_step = step_size
+        else:
+            entry_step = ttree.num_entries_for(step_size, expressions=common_keys)
+
+        def foreach(start):
+            stop = min(start + entry_step, entry_stop)
+            partition_args.append((i, start, stop))
+
+        for start in range(entry_start, entry_stop, entry_step):
+            foreach(start)
+
+    first5 = hasbranches[0].arrays(common_keys, entry_start=0, entry_stop=5)
+    meta = dask_awkward.core.typetracer_array(first5)
+
+    return dask_awkward.from_map(
+        _UprootRead(hasbranches, common_keys),
+        partition_args,
+        label="from-uproot",
+        meta=meta,
+    )
+
+
+def _get_dak_array_delay_open(
+    files,
+    filter_name=no_filter,
+    filter_typename=no_filter,
+    filter_branch=no_filter,
+    recursive=True,
+    full_paths=False,
+    custom_classes=None,
+    allow_missing=False,
+    real_options=None,  # NOTE: a comma after **options breaks Python 2
+):
+    import dask_awkward
+
+    ffile_path, fobject_path = files[0]
+    obj = uproot._util.regularize_object_path(
+        ffile_path, fobject_path, custom_classes, allow_missing, real_options
+    )
+    common_keys = obj.keys(
+        recursive=recursive,
+        filter_name=filter_name,
+        filter_typename=filter_typename,
+        filter_branch=filter_branch,
+        full_paths=full_paths,
+    )
+
+    class _UprootOpenAndRead:
+        def __init__(
+            self, custom_classes, allow_missing, real_options, common_keys
+        ) -> None:
+            self.custom_classes = custom_classes
+            self.allow_missing = allow_missing
+            self.real_options = real_options
+            self.common_keys = common_keys
+
+        def __call__(self, file_path_object_path):
+            file_path, object_path = file_path_object_path
+            ttree = uproot._util.regularize_object_path(
+                file_path,
+                object_path,
+                self.custom_classes,
+                self.allow_missing,
+                self.real_options,
+            )
+            return ttree.arrays(self.common_keys)
+
+    first5 = obj.arrays(common_keys, entry_start=0, entry_stop=5)
+    meta = dask_awkward.core.typetracer_array(first5)
+
+    return dask_awkward.from_map(
+        _UprootOpenAndRead(custom_classes, allow_missing, real_options, common_keys),
+        files,
+        label="from-uproot",
+        meta=meta,
+    )
+
+
 class _WrapDict(MutableMapping):
     def __init__(self, dict):
         self.dict = dict

tests/test_0652_dask-for-awkward.py

@@ -0,0 +1,96 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/uproot4/blob/main/LICENSE
+
+import numpy
+import pytest
+import skhep_testdata
+
+import uproot
+
+dask = pytest.importorskip("dask")
+dask_awkward = pytest.importorskip("dask_awkward")
+
+
+def test_single_dask_awkward_array():
+    test_path = skhep_testdata.data_path("uproot-Zmumu.root") + ":events"
+    ttree = uproot.open(test_path)
+
+    ak_array = ttree.arrays()
+    dak_array = uproot.dask(test_path, library="ak")
+
+    assert len(ak_array) == len(dak_array)
+
+    assert dak_array.compute().to_list() == ak_array.to_list()
+
+
+def test_dask_concatenation():
+    test_path1 = skhep_testdata.data_path("uproot-Zmumu.root") + ":events"
+    test_path2 = skhep_testdata.data_path("uproot-Zmumu-uncompressed.root") + ":events"
+    test_path3 = skhep_testdata.data_path("uproot-Zmumu-zlib.root") + ":events"
+    test_path4 = skhep_testdata.data_path("uproot-Zmumu-lzma.root") + ":events"
+
+    ak_array = uproot.concatenate([test_path1, test_path2, test_path3, test_path4])
+    dak_array = uproot.dask(
+        [test_path1, test_path2, test_path3, test_path4], library="ak"
+    )
+
+    assert len(ak_array) == len(dak_array)
+
+    assert dak_array.compute().to_list() == ak_array.to_list()
+
+
+def test_multidim_array():
+    test_path = (
+        skhep_testdata.data_path("uproot-sample-6.08.04-uncompressed.root") + ":sample"
+    )
+    ttree = uproot.open(test_path)
+
+    ak_array = ttree.arrays()
+    dak_array = uproot.dask(test_path, library="ak")
+
+    assert len(ak_array) == len(dak_array)
+
+    assert dak_array.compute().to_list() == ak_array.to_list()
+
+
+def test_chunking_single_num():
+    test_path = skhep_testdata.data_path("uproot-Zmumu.root") + ":events"
+    assert uproot.dask(test_path, step_size=10, library="ak").npartitions == 231
+
+
+def test_chunking_single_string():
+    test_path = skhep_testdata.data_path("uproot-Zmumu.root") + ":events"
+    assert uproot.dask(test_path, step_size="500B", library="ak").npartitions == 330
+
+
+def test_chunking_multiple_num():
+    filename1 = skhep_testdata.data_path("uproot-Zmumu.root") + ":events"
+    filename2 = skhep_testdata.data_path("uproot-Zmumu-uncompressed.root") + ":events"
+    assert (
+        uproot.dask([filename1, filename2], step_size=10, library="ak").npartitions
+        == 462
+    )
+
+
+def test_chunking_multiple_string():
+    filename1 = skhep_testdata.data_path("uproot-Zmumu.root") + ":events"
+    filename2 = skhep_testdata.data_path("uproot-Zmumu-uncompressed.root") + ":events"
+    assert (
+        uproot.dask([filename1, filename2], step_size="500B", library="ak").npartitions
+        == 1098
+    )
+
+
+def test_delay_open():
+    test_path1 = skhep_testdata.data_path("uproot-Zmumu.root") + ":events"
+    test_path2 = skhep_testdata.data_path("uproot-Zmumu-uncompressed.root") + ":events"
+    test_path3 = skhep_testdata.data_path("uproot-Zmumu-zlib.root") + ":events"
+    test_path4 = skhep_testdata.data_path("uproot-Zmumu-lzma.root") + ":events"
+
+    ak_array = uproot.concatenate([test_path1, test_path2, test_path3, test_path4])
+    dak_array = uproot.dask(
+        [test_path1, test_path2, test_path3, test_path4], open_files=False, library="ak"
+    )
+
+    assert len(ak_array) == len(dak_array)
+
+    assert dak_array.compute().to_list() == ak_array.to_list()