feat: support shape_touched from Dask

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,23 +77,48 @@ def __init__(
         self.metadata = metadata
         self.version = version
 
-    def extract_form_keys_base_columns(self, form_keys):
-        base_columns = []
-        for form_key in form_keys:
-            base_columns.extend(
+    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)
+            base_columns.update(
                 [
                     acolumn
                     for acolumn in urllib.parse.unquote(form_key).split(",")
                     if not acolumn.startswith("!")
                 ]
             )
-        return list(set(base_columns))
+        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:
+            base_columns.extend(
+                [
+                    acolumn
+                    for acolumn in urllib.parse.unquote(form_key).split(",")
+                    if not acolumn.startswith("!")
+                ]
+            )
+        return list(set(base_columns))
+
 
 class _map_schema_uproot(_map_schema_base):
     def __init__(
@@ -125,7 +150,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/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))