feat: support shape_touched from Dask

pyproject.toml

@@ -37,11 +37,11 @@ classifiers = [
   "Topic :: Utilities",
 ]
 dependencies = [
-  "awkward>=2.4.2",
-  "uproot>=5.0.10",
+  "awkward>=2.4.5",
+  "uproot>=5.1.0rc1",
   "dask[array]>=2023.4.0",
-  "dask-awkward>=2023.7.1,!=2023.8.0",
-  "dask-histogram>=2023.6.0",
+  "dask-awkward>=2023.10a1,!=2023.8.0",
+  "dask-histogram>=2023.7a0",
   "correctionlib>=2.0.0",
   "pyarrow>=6.0.0",
   "fsspec",

src/coffea/lookup_tools/rochester_lookup.py

@@ -1,4 +1,5 @@
 import awkward
+import dask_awkward as dak
 import numpy
 
 from coffea.lookup_tools.dense_lookup import dense_lookup
@@ -75,7 +76,7 @@ def _error(self, func, *args):
 
         newargs = args + (0, 0)
         default = func(*newargs)
-        result = numpy.zeros_like(default)
+        result = awkward.zeros_like(default)
         for s in range(self._nsets):
             oneOver = 1.0 / self._members[s]
             for m in range(self._members[s]):
@@ -226,12 +227,27 @@ def _kExtra(self, kpt, eta, nl, u, s=0, m=0):
         cbN_flat = awkward.flatten(cbN)
         cbS_flat = awkward.flatten(cbS)
 
-        invcdf = awkward.unflatten(
-            doublecrystalball.ppf(
-                u_flat, cbA_flat, cbA_flat, cbN_flat, cbN_flat, loc, cbS_flat
-            ),
-            counts,
-        )
+        args = (u_flat, cbA_flat, cbA_flat, cbN_flat, cbN_flat, loc, cbS_flat)
+
+        if any(isinstance(arg, dak.Array) for arg in args):
+
+            def apply(*args):
+                args_lz = [
+                    awkward.typetracer.length_zero_if_typetracer(arg) for arg in args
+                ]
+                out = awkward.Array(doublecrystalball.ppf(*args_lz))
+                if awkward.backend(args[0]) == "typetracer":
+                    out = awkward.Array(
+                        out.layout.to_typetracer(forget_length=True),
+                        behavior=out.behavior,
+                    )
+                return out
+
+            invcdf = dak.map_partitions(apply, *args)
+        else:
+            invcdf = doublecrystalball.ppf(*args)
+
+        invcdf = awkward.unflatten(invcdf, counts)
 
         x = awkward.where(
             mask,

src/coffea/nanoevents/factory.py

@@ -77,6 +77,41 @@ def __init__(
         self.metadata = metadata
         self.version = version
 
+    def keys_for_buffer_keys(self, buffer_keys):
+        base_columns = set()
+        for buffer_key in buffer_keys:
+            form_key, attribute = self.parse_buffer_key(buffer_key)
+            operands = urllib.parse.unquote(form_key).split(",")
+
+            it_operands = iter(operands)
+            next(it_operands)
+
+            base_columns.update(
+                [
+                    name
+                    for name, maybe_transform in zip(operands, it_operands)
+                    if maybe_transform == "!load"
+                ]
+            )
+        return base_columns
+
+    def parse_buffer_key(self, buffer_key):
+        prefix, attribute, form_key = buffer_key.rsplit("/", maxsplit=2)
+        if attribute == "offsets":
+            return (form_key[: -len("%2C%21offsets")], attribute)
+        else:
+            return (form_key, attribute)
+
+    @property
+    def buffer_key(self):
+        return partial(self._key_formatter, "")
+
+    def _key_formatter(self, prefix, form_key, form, attribute):
+        if attribute == "offsets":
+            form_key += "%2C%21offsets"
+        return prefix + f"/{attribute}/{form_key}"
+
+    # TODO: deprecate
     def extract_form_keys_base_columns(self, form_keys):
         base_columns = []
         for form_key in form_keys:
@@ -89,11 +124,6 @@ def extract_form_keys_base_columns(self, form_keys):
             )
         return list(set(base_columns))
 
-    def _key_formatter(self, prefix, form_key, form, attribute):
-        if attribute == "offsets":
-            form_key += "%2C%21offsets"
-        return prefix + f"/{attribute}/{form_key}"
-
 
 class _map_schema_uproot(_map_schema_base):
     def __init__(
@@ -125,7 +155,41 @@ def __call__(self, form):
             },
             "form_key": None,
         }
-        return awkward.forms.form.from_dict(self.schemaclass(lform, self.version).form)
+        return (
+            awkward.forms.form.from_dict(self.schemaclass(lform, self.version).form),
+            self,
+        )
+
+    def load_buffers(self, tree, keys, start, stop, interp_options):
+        from functools import partial
+
+        from coffea.nanoevents.util import tuple_to_key
+
+        partition_key = (
+            str(tree.file.uuid),
+            tree.object_path,
+            f"{start}-{stop}",
+        )
+        uuidpfn = {partition_key[0]: tree.file.file_path}
+        mapping = UprootSourceMapping(
+            TrivialUprootOpener(uuidpfn, interp_options),
+            start,
+            stop,
+            cache={},
+            access_log=None,
+            use_ak_forth=True,
+        )
+        mapping.preload_column_source(partition_key[0], partition_key[1], tree)
+        buffer_key = partial(self._key_formatter, tuple_to_key(partition_key))
+
+        class TranslateBufferKeys:
+            def __getitem__(this, key):
+                form_key, attribute = self.parse_buffer_key(key)
+                return mapping[
+                    buffer_key(form_key=form_key, attribute=attribute, form=None)
+                ]
+
+        return TranslateBufferKeys()
 
     def create_column_mapping_and_key(self, tree, start, stop, interp_options):
         from functools import partial

src/coffea/nanoevents/mapping/base.py

@@ -111,14 +111,19 @@ def __getitem__(self, key):
         if len(stack) != 1:
             raise RuntimeError(f"Syntax error in form key {nodes}")
         out = stack.pop()
-        try:
-            out = numpy.array(out)
-        except ValueError:
-            if self._debug:
-                print(out)
-            raise RuntimeError(
-                f"Left with non-bare array after evaluating form key {nodes}"
-            )
+        import awkward
+
+        if isinstance(out, awkward.contents.Content):
+            out = awkward.to_numpy(out)
+        else:
+            try:
+                out = numpy.array(out)
+            except ValueError:
+                if self._debug:
+                    print(out)
+                raise RuntimeError(
+                    f"Left with non-bare array after evaluating form key {nodes}"
+                )
         return out
 
     @abstractmethod

src/coffea/nanoevents/transforms.py

@@ -13,6 +13,15 @@ def to_layout(array):
     return array.layout
 
 
+def ensure_array(arraylike):
+    if isinstance(arraylike, (awkward.contents.Content, awkward.Array)):
+        return awkward.to_numpy(arraylike)
+    elif isinstance(arraylike, awkward.index.Index):
+        return arraylike.data
+    else:
+        return numpy.asarray(arraylike)
+
+
 def data(stack):
     """Extract content from array
     (currently a noop, can probably take place of !content)
@@ -96,7 +105,7 @@ def counts2offsets(stack):
     Signature: counts,!counts2offsets
     Outputs an array with length one larger than input
     """
-    counts = numpy.array(stack.pop())
+    counts = ensure_array(stack.pop())
     offsets = numpy.empty(len(counts) + 1, dtype=numpy.int64)
     offsets[0] = 0
     numpy.cumsum(counts, out=offsets[1:])
@@ -123,11 +132,11 @@ def local2global(stack):
     Signature: index,target_offsets,!local2global
     Outputs a content array with same shape as index content
     """
-    target_offsets = numpy.asarray(stack.pop())
+    target_offsets = ensure_array(stack.pop())
     index = stack.pop()
     index = index.mask[index >= 0] + target_offsets[:-1]
     index = index.mask[index < target_offsets[1:]]
-    out = numpy.array(awkward.flatten(awkward.fill_none(index, -1), axis=None))
+    out = ensure_array(awkward.flatten(awkward.fill_none(index, -1), axis=None))
     if out.dtype != numpy.int64:
         raise RuntimeError
     stack.append(out)

src/coffea/processor/executor.py

@@ -1718,7 +1718,7 @@ def _work_function(
                     import dask_awkward
 
                     to_compute = processor_instance.process(events)
-                    materialized = dask_awkward.necessary_columns(to_compute)
+                    # materialized = dask_awkward.report_necessary_buffers(to_compute)
                     out = dask.compute(to_compute, scheduler="single-threaded")[0]
             except Exception as e:
                 raise Exception(f"Failed processing file: {item!r}") from e
@@ -1734,11 +1734,11 @@ def _work_function(
                     metrics = {}
                     if isinstance(file, uproot.ReadOnlyDirectory):
                         metrics["bytesread"] = file.file.source.num_requested_bytes
+                    # metrics["data_and_shape_buffers"] = set(materialized)
+                    # metrics["shape_only_buffers"] = set(materialized)
                     if schema is not None and issubclass(schema, schemas.BaseSchema):
-                        metrics["columns"] = set(materialized)
                         metrics["entries"] = len(events)
                     else:
-                        metrics["columns"] = set(materialized)
                         metrics["entries"] = events.size
                     metrics["processtime"] = toc - tic
                     return {"out": out, "metrics": metrics, "processed": {item}}

tests/test_jetmet_tools.py

@@ -837,9 +837,9 @@ def test_corrected_jets_factory(optimization_enabled):
             **{name: evaluator[name] for name in jec_stack_names[5:6]}
         )
 
-        print(dak.necessary_columns(jets.eta))
+        print(dak.report_necessary_buffers(jets.eta))
         print(
-            dak.necessary_columns(
+            dak.report_necessary_buffers(
                 resosf.getScaleFactor(
                     JetEta=jets.eta,
                 )

tests/test_lookup_tools.py

@@ -372,8 +372,6 @@ def test_jec_txt_effareas():
 
 
 def test_rochester():
-    pytest.xfail("weird side effect from running other tests... passes by itself")
-
     rochester_data = lookup_tools.txt_converters.convert_rochester_file(
         "tests/samples/RoccoR2018.txt.gz", loaduncs=True
     )
@@ -390,27 +388,29 @@ def test_rochester():
 
     # test against nanoaod
     events = NanoEventsFactory.from_root(
-        os.path.abspath("tests/samples/nano_dimuon.root")
+        {os.path.abspath("tests/samples/nano_dimuon.root"): "Events"},
+        permit_dask=True,
     ).events()
 
     data_k = rochester.kScaleDT(
         events.Muon.charge, events.Muon.pt, events.Muon.eta, events.Muon.phi
     )
-    data_k = np.array(ak.flatten(data_k))
+    data_k = ak.flatten(data_k).compute().to_numpy()
     assert all(np.isclose(data_k, official_data_k))
     data_err = rochester.kScaleDTerror(
         events.Muon.charge, events.Muon.pt, events.Muon.eta, events.Muon.phi
     )
-    data_err = np.array(ak.flatten(data_err), dtype=float)
+    data_err = ak.flatten(data_err).compute().to_numpy()
     assert all(np.isclose(data_err, official_data_err, atol=1e-8))
 
     # test against mc
     events = NanoEventsFactory.from_root(
-        os.path.abspath("tests/samples/nano_dy.root")
+        {os.path.abspath("tests/samples/nano_dy.root"): "Events"},
+        permit_dask=True,
     ).events()
 
     hasgen = ~np.isnan(ak.fill_none(events.Muon.matched_gen.pt, np.nan))
-    mc_rand = ak.unflatten(mc_rand, ak.num(hasgen))
+    mc_rand = ak.unflatten(dak.from_awkward(ak.Array(mc_rand), 1), ak.num(hasgen))
     mc_kspread = rochester.kSpreadMC(
         events.Muon.charge[hasgen],
         events.Muon.pt[hasgen],
@@ -426,10 +426,10 @@ def test_rochester():
         events.Muon.nTrackerLayers[~hasgen],
         mc_rand[~hasgen],
     )
-    mc_k = np.array(ak.flatten(ak.ones_like(events.Muon.pt)))
-    hasgen_flat = np.array(ak.flatten(hasgen))
-    mc_k[hasgen_flat] = np.array(ak.flatten(mc_kspread))
-    mc_k[~hasgen_flat] = np.array(ak.flatten(mc_ksmear))
+    mc_k = ak.flatten(ak.ones_like(events.Muon.pt)).compute().to_numpy()
+    hasgen_flat = ak.flatten(hasgen).compute().to_numpy()
+    mc_k[hasgen_flat] = ak.flatten(mc_kspread).compute().to_numpy()
+    mc_k[~hasgen_flat] = ak.flatten(mc_ksmear).compute().to_numpy()
     assert all(np.isclose(mc_k, official_mc_k))
 
     mc_errspread = rochester.kSpreadMCerror(
@@ -447,9 +447,9 @@ def test_rochester():
         events.Muon.nTrackerLayers[~hasgen],
         mc_rand[~hasgen],
     )
-    mc_err = np.array(ak.flatten(ak.ones_like(events.Muon.pt)))
-    mc_err[hasgen_flat] = np.array(ak.flatten(mc_errspread))
-    mc_err[~hasgen_flat] = np.array(ak.flatten(mc_errsmear))
+    mc_err = ak.flatten(ak.ones_like(events.Muon.pt)).compute().to_numpy()
+    mc_err[hasgen_flat] = ak.flatten(mc_errspread).compute().to_numpy()
+    mc_err[~hasgen_flat] = ak.flatten(mc_errsmear).compute().to_numpy()
     assert all(np.isclose(mc_err, official_mc_err, atol=1e-8))