feat: One-pass column optimization

pyproject.toml

@@ -6,40 +6,40 @@ build-backend = "hatchling.build"
 name = "dask-awkward"
 description = "Awkward Array meets Dask"
 readme = "README.md"
-license = {text = "BSD-3-Clause"}
+license = { text = "BSD-3-Clause" }
 requires-python = ">=3.8"
 authors = [
-  { name = "Doug Davis", email = "ddavis@ddavis.io" },
-  { name = "Martin Durant", email = "mdurant@anaconda.com" },
+    { name = "Doug Davis", email = "ddavis@ddavis.io" },
+    { name = "Martin Durant", email = "mdurant@anaconda.com" },
 ]
 maintainers = [
-  { name = "Doug Davis", email = "ddavis@ddavis.io" },
-  { name = "Martin Durant", email = "mdurant@anaconda.com" },
+    { name = "Doug Davis", email = "ddavis@ddavis.io" },
+    { name = "Martin Durant", email = "mdurant@anaconda.com" },
 ]
 classifiers = [
-  "Development Status :: 5 - Production/Stable",
-  "Intended Audience :: Developers",
-  "Intended Audience :: Information Technology",
-  "Intended Audience :: Science/Research",
-  "License :: OSI Approved :: BSD License",
-  "Programming Language :: Python :: 3",
-  "Programming Language :: Python :: 3 :: Only",
-  "Programming Language :: Python :: 3.8",
-  "Programming Language :: Python :: 3.9",
-  "Programming Language :: Python :: 3.10",
-  "Programming Language :: Python :: 3.11",
-  "Programming Language :: Python :: 3.12",
-  "Topic :: Scientific/Engineering",
-  "Topic :: Scientific/Engineering :: Information Analysis",
-  "Topic :: Scientific/Engineering :: Mathematics",
-  "Topic :: Scientific/Engineering :: Physics",
-  "Topic :: Software Development",
+    "Development Status :: 5 - Production/Stable",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Information Technology",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: BSD License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering",
+    "Topic :: Scientific/Engineering :: Information Analysis",
+    "Topic :: Scientific/Engineering :: Mathematics",
+    "Topic :: Scientific/Engineering :: Physics",
+    "Topic :: Software Development",
 ]
 dependencies = [
-  "awkward >=2.5.1",
-  "dask >=2023.04.0",
-  "cachetools",
-  "typing_extensions >=4.8.0",
+    "awkward >=2.6.7",
+    "dask >=2023.04.0",
+    "cachetools",
+    "typing_extensions >=4.8.0",
 ]
 dynamic = ["version"]
 
@@ -48,33 +48,29 @@ Homepage = "https://github.com/dask-contrib/dask-awkward"
 "Bug Tracker" = "https://github.com/dask-contrib/dask-awkward/issues"
 
 [project.optional-dependencies]
-io = [
-  "pyarrow",
-]
-complete = [
-  "dask-awkward[io]",
-]
+io = ["pyarrow"]
+complete = ["dask-awkward[io]"]
 # `docs` and `test` are separate from user installs
 docs = [
-  "dask-awkward[complete]",
-  "sphinx-book-theme",
-  "sphinx-design",
-  "sphinx-codeautolink",
-  # broken see PR 451
-  # "dask-sphinx-theme",
+    "dask-awkward[complete]",
+    "sphinx-book-theme",
+    "sphinx-design",
+    "sphinx-codeautolink",
+    # broken see PR 451
+    # "dask-sphinx-theme",
 ]
 test = [
-  "aiohttp;python_version<\"3.12\"",
-  "dask[dataframe]",
-  "dask-awkward[complete]",
-  "dask-histogram",
-  "distributed",
-  "hist",
-  "pandas",
-  "pytest >=6.0,<8",
-  "pytest-cov >=3.0.0",
-  "requests",
-  "uproot >=5.1.0",
+    "aiohttp;python_version<\"3.12\"",
+    "dask[dataframe]",
+    "dask-awkward[complete]",
+    "dask-histogram",
+    "distributed",
+    "hist",
+    "pandas",
+    "pytest >=6.0,<8",
+    "pytest-cov >=3.0.0",
+    "requests",
+    "uproot >=5.1.0",
 ]
 
 [project.entry-points."dask.sizeof"]
@@ -130,17 +126,17 @@ warn_unused_ignores = true
 warn_unreachable = true
 
 [[tool.mypy.overrides]]
-  module = [
+module = [
     "awkward.*",
     "IPython.*",
     "fsspec.*",
     "pyarrow.*",
     "tlz.*",
     "uproot.*",
     "cloudpickle.*",
-    "cachetools.*"
-  ]
-  ignore_missing_imports = true
+    "cachetools.*",
+]
+ignore_missing_imports = true
 
 [tool.pyright]
 include = ["src"]
@@ -149,27 +145,27 @@ reportPrivateImportUsage = false
 
 [tool.coverage.report]
 exclude_lines = [
-  "pragma: no cover",
-  "if TYPE_CHECKING:",
-  "except ImportError:",
-  "NotImplementedError",
-  "DaskAwkwardNotImplemented",
-  "_ipython_key_completions_",
-  "Only highlevel=True is supported",
-  "\\.\\.\\.$",
+    "pragma: no cover",
+    "if TYPE_CHECKING:",
+    "except ImportError:",
+    "NotImplementedError",
+    "DaskAwkwardNotImplemented",
+    "_ipython_key_completions_",
+    "Only highlevel=True is supported",
+    "\\.\\.\\.$",
 ]
 fail_under = 90
 show_missing = true
 
 [tool.coverage.run]
 omit = [
-  "*/dask_awkward/lib/unproject_layout.py",
-  "*/tests/test_*.py",
-  "*/tests/__init__.py",
-  "*/version.py",
+    "*/dask_awkward/lib/unproject_layout.py",
+    "*/tests/test_*.py",
+    "*/tests/__init__.py",
+    "*/version.py",
 ]
 source = ["src/"]
 
 [tool.ruff]
-ignore = ["E501", "E402"]
-per-file-ignores = {"__init__.py" = ["E402", "F401"]}
+lint.ignore = ["E501", "E402"]
+lint.per-file-ignores = { "__init__.py" = ["E402", "F401"] }

src/dask_awkward/__init__.py

@@ -19,14 +19,7 @@
     partition_compatibility,
 )
 from dask_awkward.lib.describe import backend, fields
-from dask_awkward.lib.inspect import (
-    report_necessary_buffers,
-    report_necessary_columns,
-    sample,
-)
-
-necessary_columns = report_necessary_columns  # Export for backwards compatibility.
-
+from dask_awkward.lib.inspect import sample
 from dask_awkward.lib.io.io import (
     from_awkward,
     from_dask_array,
@@ -42,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/layers/__init__.py

@@ -4,8 +4,6 @@
     AwkwardMaterializedLayer,
     AwkwardTreeReductionLayer,
     ImplementsIOFunction,
-    ImplementsProjection,
-    IOFunctionWithMocking,
     io_func_implements_projection,
 )
 
@@ -14,8 +12,6 @@
     "AwkwardBlockwiseLayer",
     "AwkwardMaterializedLayer",
     "AwkwardTreeReductionLayer",
-    "ImplementsProjection",
     "ImplementsIOFunction",
-    "IOFunctionWithMocking",
     "io_func_implements_projection",
 )