scaffolding

dask-contrib · Sep 18, 2023 · 9e7b806 · 9e7b806
1 parent fd86f66
commit 9e7b806
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 2 deletions.
diff --git a/src/dask_awkward/layers/layers.py b/src/dask_awkward/layers/layers.py
@@ -158,7 +158,12 @@ def mock(self) -> tuple[AwkwardInputLayer, TypeTracerReport]:
         )
         return new_input_layer, report
 
-    def project_columns(self, columns: list[str]) -> AwkwardInputLayer:
+    def project_columns(
+        self,
+        columns: list[str],
+        *,
+        necessary_shape_columns: list[str] | None = None,
+    ) -> AwkwardInputLayer:
         if hasattr(self.io_func, "project_columns"):
             # TODO: make project_columns call sites never pass in an
             # empty list.
@@ -182,6 +187,7 @@ def project_columns(self, columns: list[str]) -> AwkwardInputLayer:
                 io_func = self.io_func.project_columns(
                     columns,
                     original_form=original_form,
+                    necessary_shape_columns=necessary_shape_columns,
                 )
             except TypeError:
                 io_func = self.io_func.project_columns(columns)

diff --git a/src/dask_awkward/lib/io/json.py b/src/dask_awkward/lib/io/json.py
@@ -68,6 +68,7 @@ def _default_project_columns(
         self,
         columns: Sequence[str],
         original_form: Form | None = None,
+        necessary_shape_columns: Sequence[str] | None = None,
     ) -> Self:
         if self.schema is not None:
             return self
@@ -126,13 +127,15 @@ def project_columns(
         self,
         columns: Sequence[str],
         original_form: Form | None = None,
+        necessary_shape_columns: Sequence[str] | None = None,
     ) -> Self:
         if not _use_optimization():
             return self
 
         return self._default_project_columns(
             columns=columns,
             original_form=original_form,
+            necessary_shape_columns=necessary_shape_columns,
         )
 
 
@@ -177,13 +180,15 @@ def project_columns(
         self,
         columns: Sequence[str],
         original_form: Form | None = None,
+        necessary_shape_columns: Sequence[str] | None = None,
     ) -> Self:
         if not _use_optimization():
             return self
 
         return self._default_project_columns(
             columns=columns,
             original_form=original_form,
+            necessary_shape_columns=necessary_shape_columns,
         )
 
 
@@ -236,13 +241,15 @@ def project_columns(
         self,
         columns: Sequence[str],
         original_form: Form | None = None,
+        necessary_shape_columns: Sequence[str] | None = None,
     ) -> Self:
         if not _use_optimization():
             return self
 
         return self._default_project_columns(
             columns=columns,
             original_form=original_form,
+            necessary_shape_columns=necessary_shape_columns,
         )
 
 

diff --git a/src/dask_awkward/lib/io/parquet.py b/src/dask_awkward/lib/io/parquet.py
@@ -70,6 +70,7 @@ def project_columns(
         self,
         columns: Sequence[str] | None,
         orignal_form: Form | None = None,
+        necessary_shape_columns: Sequence[str] | None = None,
     ) -> _FromParquetFn:
         ...
 
@@ -130,6 +131,7 @@ def project_columns(
         self,
         columns: Sequence[str] | None,
         original_form: Form | None = None,
+        necessary_shape_columns: Sequence[str] | None = None,
     ) -> _FromParquetFileWiseFn:
         if not _use_optimization():
             return self
@@ -192,6 +194,7 @@ def project_columns(
         self,
         columns: Sequence[str] | None,
         original_form: Form | None = None,
+        necessary_shape_columns: Sequence[str] | None = None,
     ) -> _FromParquetFragmentWiseFn:
         if not _use_optimization():
             return self

diff --git a/src/dask_awkward/lib/optimize.py b/src/dask_awkward/lib/optimize.py
@@ -123,7 +123,10 @@ def optimize_columns(dsk: HighLevelGraph) -> HighLevelGraph:
     for name, neccols in layer_to_necessary_columns.items():
         meta = layers[name]._meta
         neccols = _prune_wildcards(neccols, meta)
-        layers[name] = layers[name].project_columns(neccols)
+        layers[name] = layers[name].project_columns(
+            neccols,
+            necessary_shape_columns=None,
+        )
 
     return HighLevelGraph(layers, deps)