feat: One-pass column optimization

src/dask_awkward/__init__.py

@@ -20,9 +20,6 @@
 )
 from dask_awkward.lib.describe import fields
 from dask_awkward.lib.inspect import sample
-
-necessary_columns = None  # TODO
-
 from dask_awkward.lib.io.io import (
     from_awkward,
     from_dask_array,
@@ -38,6 +35,7 @@
 from dask_awkward.lib.io.parquet import from_parquet, to_parquet
 from dask_awkward.lib.io.text import from_text
 from dask_awkward.lib.operations import concatenate
+from dask_awkward.lib.optimize import necessary_columns
 from dask_awkward.lib.reducers import (
     all,
     any,

src/dask_awkward/lib/__init__.py

@@ -24,6 +24,7 @@
 from dask_awkward.lib.io.parquet import from_parquet, to_parquet
 from dask_awkward.lib.io.text import from_text
 from dask_awkward.lib.operations import concatenate
+from dask_awkward.lib.optimize import necessary_columns
 from dask_awkward.lib.reducers import (
     all,
     any,

src/dask_awkward/lib/core.py

@@ -719,8 +719,8 @@ def __getitem__(self, where):
         token = tokenize(self, where)
         new_name = f"{where}-{token}"
         report = self.report
-        [_.commit(new_name) for _ in report]
         new_meta = self._meta[where]
+        [_.commit(new_name) for _ in report]
 
         # first check for array type return
         if isinstance(new_meta, ak.Array):

src/dask_awkward/lib/optimize.py

@@ -111,10 +111,30 @@ def optimize_columns(dsk: HighLevelGraph, keys: Sequence[Key]) -> HighLevelGraph
             all_reps.update(getattr(dsk.layers[ln].meta, "_report", ()))
     name = tokenize("output", lays)
     [_.commit(name) for _ in all_reps]
+    all_layers = tuple(dsk.layers) + (name,)
+
+    for k, lay, cols in _optimize_columns(dsk.layers, all_layers):
+        new_lay = lay.project(cols)
+        dsk2[k] = new_lay
+
+    return HighLevelGraph(dsk2, dsk.dependencies)
+
+
+def _buf_to_col(s):
+    return (
+        s[2:]
+        .replace(".content", "")
+        .replace("-offsets", "")
+        .replace("-data", "")
+        .replace("-index", "")
+        .replace("-mask", "")
+    )
+
+
+def _optimize_columns(dsk, all_layers):
 
     # this loop is necessary_columns
-    all_layers = tuple(dsk.layers) + (name,)
-    for k, lay in dsk.layers.items():
+    for k, lay in dsk.copy().items():
         if not isinstance(lay, AwkwardInputLayer) or not hasattr(
             lay.io_func, "_column_report"
         ):
@@ -125,29 +145,38 @@ def optimize_columns(dsk: HighLevelGraph, keys: Sequence[Key]) -> HighLevelGraph
         for ln in all_layers:
             # this loop not required after next ak release
             try:
-                cols |= set(rep.data_touched_in((ln,)))
-                for col in rep.shape_touched_in((ln,)):
-                    if col in cols or any(_.startswith(col) for _ in cols):
-                        # loopy loop?
+                cols |= {_buf_to_col(s) for s in rep.data_touched_in((ln,))}
+                for col in (_buf_to_col(s) for s in rep.shape_touched_in((ln,))):
+                    if col in cols:
+                        continue
+                    if any(_.startswith(col) for _ in cols):
                         continue
-                    col2 = (
-                        col[2:]
-                        .replace(".content", "")
-                        .replace("-offsets", "")
-                        .replace("-data", "")
-                        .replace("-index", "")
-                        .replace("-mask", "")
-                    )
-                    ll = list(_ for _ in all_cols if _.startswith(col2))
+                    ll = list(_ for _ in all_cols if _.startswith(col))
                     if ll:
-                        cols.add("@." + ll[0])
+                        cols.add(ll[0])
 
             except KeyError:
                 pass
-        new_lay = lay.project([c.replace("@.", "") for c in cols])
-        dsk2[k] = new_lay
+        yield k, lay, cols
 
-    return HighLevelGraph(dsk2, dsk.dependencies)
+
+def necessary_columns(*args):
+    dsk = {}
+    all_reps = set()
+    all_layers = set()
+    for arg in args:
+        dsk.update(arg.dask.layers)
+        all_layers.add(arg.name)
+        touch_data(arg._meta)
+        all_reps.update(getattr(arg.dask.layers[arg.name].meta, "_report", ()))
+    name = tokenize("output", args)
+    [_.commit(name) for _ in all_reps]
+    all_layers = tuple(all_layers) + (name,)
+
+    out = {}
+    for k, _, cols in _optimize_columns(dsk, all_layers):
+        out[k] = cols
+    return out
 
 
 def rewrite_layer_chains(dsk: HighLevelGraph, keys: Sequence[Key]) -> HighLevelGraph: