feat: add dask_to_root

pyproject.toml

@@ -39,9 +39,7 @@ dependencies = [
   "numpy",
   "fsspec",
   "packaging",
-  "typing_extensions>=4.1.0; python_version < \"3.11\"",
-  "dask >=2023.04.0",
-  "dask-awkward>=2023.12.1"
+  "typing_extensions>=4.1.0; python_version < \"3.11\""
 ]
 description = "ROOT I/O in pure Python and NumPy."
 dynamic = [

src/uproot/writing/_dask_write.py

@@ -3,11 +3,6 @@
 import math
 from typing import Any
 
-from dask.base import tokenize
-from dask.blockwise import BlockIndex
-from dask.highlevelgraph import HighLevelGraph
-from dask_awkward.layers.layers import AwkwardMaterializedLayer
-from dask_awkward.lib.core import map_partitions, new_scalar_object
 from fsspec import AbstractFileSystem
 from fsspec.core import url_to_fs
 
@@ -42,7 +37,7 @@ def __call__(self, data, block_index):
         if self.prefix is not None:
             filename = f"{self.prefix}-{filename}"
         filename = f"{self.protocol}://{self.path}/{filename}"
-        return to_root(
+        return ak_to_root(
             filename, data, **self.kwargs, storage_options=self.storage_options
         )
 
@@ -95,6 +90,12 @@ def dask_write(
         uproot.dask_write(d)
 
     """
+    from dask.base import tokenize
+    from dask.blockwise import BlockIndex
+    from dask.highlevelgraph import HighLevelGraph
+    from dask_awkward.layers.layers import AwkwardMaterializedLayer
+    from dask_awkward.lib.core import map_partitions, new_scalar_object
+
     fs, path = url_to_fs(destination, **(storage_options or {}))
     name = f"write-root-{tokenize(fs, array, destination)}"
 
@@ -114,9 +115,7 @@ def dask_write(
             initial_basket_capacity=initial_basket_capacity,
         ),
         array,
-        BlockIndex(
-            (array.npartitions,)
-        ),  # class to provide current block index at each block of the operation...
+        BlockIndex((array.npartitions,)),
         label="to-root",
         meta=array._meta,
     )
@@ -138,7 +137,7 @@ def dask_write(
         return out
 
 
-def to_root(  # user-defined groups for ak.zip?
+def ak_to_root(
     destination,
     array,
     tree_name,

tests/test_1085_dask_write.py

@@ -0,0 +1,87 @@
+import math
+import pytest
+import awkward as ak
+import skhep_testdata
+import os
+
+dask = pytest.importorskip("dask")
+da = pytest.importorskip("dask.array")
+dd = pytest.importorskip("dask.distributed")
+dask_awkward = pytest.importorskip("dask_awkward")
+
+from distributed import Client
+import uproot
+
+
+def simple_test(tmp_path):
+    partitions = 2
+    arr = ak.Array([{"a": [1, 2, 3]}, {"a": [4, 5]}])
+    dask_arr = dask_awkward.from_awkward(arr, npartitions=partitions)
+    uproot.dask_write(dask_arr, tmp_path, prefix="data")
+    file_1 = uproot.open(os.path.join(tmp_path, "data-part0.root"))
+    assert len(file_1["tree"]["a"].arrays()) == math.ceil(len(arr) / partitions)
+    assert ak.all(file_1["tree"]["a"].arrays()["a"][0] == arr[0]["a"])
+    file_2 = uproot.open(os.path.join(tmp_path, "data-part1.root"))
+    assert len(file_2["tree"]["a"].arrays()) == int(len(arr) / partitions)
+    assert ak.all(file_2["tree"]["a"].arrays()["a"][0] == arr[1]["a"])
+
+
+def HZZ_test(tmp_path):
+    """
+    Write data from HZZ with 3 partitions.
+    """
+    partitions = 2
+    arr = uproot.open(skhep_testdata.data_path("uproot-HZZ.root"))["events"].arrays()
+    dask_arr = dask_awkward.from_awkward(ak.from_iter(arr), partitions)
+    uproot.dask_write(dask_arr, tmp_path, prefix="data")
+    file_1 = uproot.open(os.path.join(tmp_path, "data-part0.root"))
+    assert len(file_1["tree"]["Jet_Px"].arrays()) == math.ceil(len(arr) / partitions)
+    assert ak.all(file_1["tree"]["Jet_Px"].arrays()["Jet_Px"][0] == arr["Jet_Px"])
+    file_2 = uproot.open(os.path.join(tmp_path, "data-part1.root"))
+    assert len(file_2["tree"]["Jet_Px"].arrays()) == int(len(arr) / partitions)
+    assert ak.all(file_2["tree"]["Jet_Px"].arrays()["Jet_Px"][0] == arr["Jet_Px"])
+
+
+def test_graph(tmp_path):
+    """
+    Test compute parameter (want to return highlevelgraph)
+    """
+    partitions = 2
+    arr = ak.Array([{"a": [1, 2, 3]}, {"a": [4, 5]}])
+    dask_arr = dask_awkward.from_awkward(arr, npartitions=partitions)
+    graph = uproot.dask_write(
+        dask_arr,
+        str(tmp_path),
+        prefix="compute",
+        compute=False,
+    )
+    graph.compute()
+    file_1 = uproot.open(os.path.join(tmp_path, "compute-part0.root"))
+    assert len(file_1["tree"]["a"].arrays()) == math.ceil(len(arr) / partitions)
+    assert ak.all(file_1["tree"]["a"].arrays()["a"][0] == arr[0]["a"])
+    file_2 = uproot.open(os.path.join(tmp_path, "compute-part1.root"))
+    assert len(file_2["tree"]["a"].arrays()) == int(len(arr) / partitions)
+    assert ak.all(file_2["tree"]["a"].arrays()["a"][0] == arr[1]["a"])
+
+
+def test_compute(tmp_path):
+    print("here")
+    partitions = 2
+    arr = uproot.open(skhep_testdata.data_path("uproot-HZZ.root"))["events"].arrays()
+    dask_arr = dask_awkward.from_awkward(ak.from_iter(arr), partitions)
+    with Client() as _:
+        graph = uproot.dask_write(
+            dask_arr, str(tmp_path), prefix="distribute", compute=False
+        )
+        dask.compute(graph)
+    dask_arr = dask_awkward.from_awkward(ak.from_iter(arr), partitions)
+    file_1 = uproot.open(os.path.join(tmp_path, "distribute-part0.root"))
+    assert len(file_1["tree"]["Jet_Px"].arrays()) == math.ceil(len(arr) / partitions)
+    assert ak.all(file_1["tree"]["Jet_Px"].arrays()["Jet_Px"][0] == arr["Jet_Px"])
+    file_2 = uproot.open(os.path.join(tmp_path, "distribute-part1.root"))
+    assert len(file_2["tree"]["Jet_Px"].arrays()) == int(len(arr) / partitions)
+    assert ak.all(file_2["tree"]["Jet_Px"].arrays()["Jet_Px"][0] == arr["Jet_Px"])
+
+
+# if __name__ == "__main__":
+#     test_compute("\my-output")