feat: Implement make_ak_form for IO factories

pyproject.toml

@@ -11,7 +11,9 @@ dev = [
   "build",
   "twine",
   "pre-commit",
-  "ipykernel"
+  "ipykernel",
+  "dask-awkward",
+  "dask[complete]"
 ]
 docs = ["mkdocs", "mkdocs-material", "mkdocstrings", "mkdocstrings-python"]
 

src/pybes3/besio/root_io.py

@@ -222,6 +222,11 @@ def make_awkward_content(self, raw_data: dict):
             awkward.contents.RecordArray(record_contents, list(raw_data.keys())),
         )
 
+    def make_awkward_form(self):
+        raise NotImplementedError(
+            "make_awkward_form is not implemented for Bes3CgemClusterColFactory"
+        )
+
 
 class Bes3SymMatrixArrayFactory(Factory):
     target_items = {
@@ -299,6 +304,9 @@ def make_awkward_content(self, raw_data: np.ndarray):
             )
         )
 
+    def make_awkward_form(self):
+        return awkward.forms.NumpyForm("float64", inner_shape=[self.full_dim, self.full_dim])
+
 
 uproot_custom.registered_factories |= {
     Bes3TObjArrayFactory,

tests/test_besio.py

@@ -232,6 +232,20 @@ def test_symetric_matrix(arr):
         test_symetric_matrix(tmp_arr)
 
 
+def test_bes3_tobjarray_factory_dask(data_dir):
+    dask_arr = uproot.dask({data_dir / "test_full_mc_evt_1.rtraw": "Event/m_mdcDigiCol"})
+
+    dask_arr.compute()
+
+
+def test_symetric_matrix_expansion_dask(data_dir):
+    dask_arr = uproot.dask(
+        {data_dir / "test_full_mc_evt_1.dst": "Event/TDstEvent/m_mdcTrackCol"}
+    )
+
+    dask_arr.compute()
+
+
 def test_digi_expand_TRawData(data_dir):
     f_rec = uproot.open(data_dir / "test_full_mc_evt_1.rec")
     arr_digi = f_rec["Event/TDigiEvent"].arrays()