diff --git a/pyproject.toml b/pyproject.toml index d5514159..f37d7535 100644 --- a/pyproject.toml +++ b/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,7 +126,7 @@ warn_unused_ignores = true warn_unreachable = true [[tool.mypy.overrides]] - module = [ +module = [ "awkward.*", "IPython.*", "fsspec.*", @@ -138,9 +134,9 @@ warn_unreachable = true "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"] } diff --git a/src/dask_awkward/__init__.py b/src/dask_awkward/__init__.py index 34b5c4f5..5d1c00f8 100644 --- a/src/dask_awkward/__init__.py +++ b/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, diff --git a/src/dask_awkward/layers/__init__.py b/src/dask_awkward/layers/__init__.py index d4ba4c5e..6ea8c1c3 100644 --- a/src/dask_awkward/layers/__init__.py +++ b/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", ) diff --git a/src/dask_awkward/layers/layers.py b/src/dask_awkward/layers/layers.py index 92441443..f9320da9 100644 --- a/src/dask_awkward/layers/layers.py +++ b/src/dask_awkward/layers/layers.py @@ -1,23 +1,14 @@ from __future__ import annotations -import copy from collections.abc import Callable, Mapping -from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar, Union, cast +from typing import Any, Protocol, TypeVar from dask.blockwise import Blockwise, BlockwiseDepDict, blockwise_token from dask.highlevelgraph import MaterializedLayer from dask.layers import DataFrameTreeReduction -from typing_extensions import TypeAlias from dask_awkward.utils import LazyInputsDict -if TYPE_CHECKING: - from awkward import Array as AwkwardArray - from awkward._nplikes.typetracer import TypeTracerReport - - -BackendT: TypeAlias = Union[Literal["cpu"], Literal["jax"], Literal["cuda"]] - class AwkwardBlockwiseLayer(Blockwise): """Just like upstream Blockwise, except we override pickling""" @@ -30,17 +21,11 @@ def from_blockwise(cls, layer: Blockwise) -> AwkwardBlockwiseLayer: ob.__dict__.update(layer.__dict__) return ob - def mock(self) -> AwkwardBlockwiseLayer: - layer = copy.copy(self) - nb = layer.numblocks - layer.numblocks = {k: tuple(1 for _ in v) for k, v in nb.items()} - layer.__dict__.pop("_dims", None) - return layer - def __getstate__(self) -> dict: - # Indicator that this layer has been serialised state = self.__dict__.copy() + # Indicator that this layer has been serialised state["has_been_unpickled"] = True + state.pop("meta", None) # this is a typetracer return state def __repr__(self) -> str: @@ -54,59 +39,13 @@ def __call__(self, *args, **kwargs): ... T = TypeVar("T") -class ImplementsMocking(ImplementsIOFunction, Protocol): - def mock(self) -> AwkwardArray: ... - - -class ImplementsMockEmpty(ImplementsIOFunction, Protocol): - def mock_empty(self, backend: BackendT) -> AwkwardArray: ... - - class ImplementsReport(ImplementsIOFunction, Protocol): @property def return_report(self) -> bool: ... -class ImplementsProjection(ImplementsMocking, Protocol[T]): - def prepare_for_projection(self) -> tuple[AwkwardArray, TypeTracerReport, T]: ... - - def project(self, report: TypeTracerReport, state: T) -> ImplementsIOFunction: ... - - -class ImplementsNecessaryColumns(ImplementsProjection[T], Protocol): - def necessary_columns( - self, report: TypeTracerReport, state: T - ) -> frozenset[str]: ... - - -class IOFunctionWithMocking(ImplementsMocking, ImplementsIOFunction): - def __init__(self, meta: AwkwardArray, io_func: ImplementsIOFunction): - self._meta = meta - self._io_func = io_func - - def __getstate__(self) -> dict: - state = self.__dict__.copy() - state["_meta"] = None - return state - - def __call__(self, *args, **kwargs): - return self._io_func(*args, **kwargs) - - def mock(self) -> AwkwardArray: - assert self._meta is not None - return self._meta - - def io_func_implements_projection(func: ImplementsIOFunction) -> bool: - return hasattr(func, "prepare_for_projection") - - -def io_func_implements_mocking(func: ImplementsIOFunction) -> bool: - return hasattr(func, "mock") - - -def io_func_implements_mock_empty(func: ImplementsIOFunction) -> bool: - return hasattr(func, "mock_empty") + return hasattr(func, "project") def io_func_implements_columnar(func: ImplementsIOFunction) -> bool: @@ -172,94 +111,15 @@ def __init__( def __repr__(self) -> str: return f"AwkwardInputLayer<{self.output}>" - @property - def is_projectable(self) -> bool: - # isinstance(self.io_func, ImplementsProjection) - return ( - io_func_implements_projection(self.io_func) and not self.has_been_unpickled - ) - - @property - def is_mockable(self) -> bool: - # isinstance(self.io_func, ImplementsMocking) - return io_func_implements_mocking(self.io_func) - @property def is_columnar(self) -> bool: return io_func_implements_columnar(self.io_func) - def mock(self) -> AwkwardInputLayer: - assert self.is_mockable - return AwkwardInputLayer( - name=self.name, - inputs=[None][: int(list(self.numblocks.values())[0][0])], - io_func=lambda *_, **__: cast(ImplementsMocking, self.io_func).mock(), - label=self.label, - produces_tasks=self.produces_tasks, - creation_info=self.creation_info, - annotations=self.annotations, - ) - - def prepare_for_projection(self) -> tuple[AwkwardInputLayer, TypeTracerReport, T]: - """Mock the input layer as starting with a data-less typetracer. - This method is used to create new dask task graphs that - operate purely on typetracer Arrays (that is, array with - awkward structure but without real data buffers). This allows - us to test which parts of a real awkward array will be used in - a real computation. We do this by running a graph which starts - with mocked AwkwardInputLayers. - - We mock an AwkwardInputLayer in these steps: - 1. Ask the IO function to prepare a new meta array, and return - any transient state. - 2. Build a new AwkwardInputLayer whose IO function just returns - this meta (typetracer) array - 3. Return the new input layer and the transient state - - When this new layer is added to a dask task graph and that - graph is computed, the report object will be mutated. - Inspecting the report object after the compute tells us which - buffers from the original form would be required for a real - compute with the same graph. - Returns - ------- - AwkwardInputLayer - Copy of the input layer with data-less input. - TypeTracerReport - The report object used to track touched buffers. - Any - The black-box state object returned by the IO function. - """ - assert self.is_projectable - new_meta_array, report, state = cast( - ImplementsProjection, self.io_func - ).prepare_for_projection() - - new_return = new_meta_array - if io_func_implements_report(self.io_func): - if cast(ImplementsReport, self.io_func).return_report: - new_return = (new_meta_array, type(new_meta_array)([])) - - new_input_layer = AwkwardInputLayer( - name=self.name, - inputs=[None][: int(list(self.numblocks.values())[0][0])], - io_func=AwkwardTokenizable(new_return, self.name), - label=self.label, - produces_tasks=self.produces_tasks, - creation_info=self.creation_info, - annotations=self.annotations, - ) - return new_input_layer, report, state - - def project( - self, - report: TypeTracerReport, - state: T, - ) -> AwkwardInputLayer: - assert self.is_projectable - io_func = cast(ImplementsProjection, self.io_func).project( - report=report, state=state - ) + def project(self, columns: list[str]) -> AwkwardInputLayer: + if hasattr(self.io_func, "project"): + io_func = self.io_func.project(columns) + else: + return self return AwkwardInputLayer( name=self.name, inputs=self.inputs, @@ -270,12 +130,6 @@ def project( annotations=self.annotations, ) - def necessary_columns(self, report: TypeTracerReport, state: T) -> frozenset[str]: - assert self.is_columnar - return cast(ImplementsNecessaryColumns, self.io_func).necessary_columns( - report=report, state=state - ) - class AwkwardMaterializedLayer(MaterializedLayer): def __init__( @@ -290,68 +144,5 @@ def __init__( self.fn = fn super().__init__(mapping, **kwargs) - def mock(self) -> MaterializedLayer: - mapping = copy.copy(self.mapping) - if not mapping: - # no partitions at all - return self - name = next(iter(mapping))[0] - - npln = len(self.previous_layer_names) - # one previous layer name - # - # this case is used for mocking repartition or slicing where - # we maybe have multiple partitions that need to be included - # in a task. - if npln == 1: - prev_name: str = self.previous_layer_names[0] - if (name, 0) in mapping: - task = mapping[(name, 0)] - task = tuple( - ( - (prev_name, 0) - if isinstance(v, tuple) and len(v) == 2 and v[0] == prev_name - else v - ) - for v in task - ) - - # when using Array.partitions we need to mock that we - # just want the first partition. - if len(task) == 2 and isinstance(task[1], int) and task[1] > 0: - task = (task[0], 0) - return MaterializedLayer({(name, 0): task}) - return self - - # zero previous layers; this is likely a known scalar. - # - # we just use the existing mapping - elif npln == 0: - return MaterializedLayer({(name, 0): mapping[(name, 0)]}) - # more than one previous_layer_names - # - # this case is needed for dak.concatenate on axis=0; we need - # the first partition of _each_ of the previous layer names! - else: - if self.fn is None: - raise ValueError( - "For multiple previous layers the fn argument cannot be None." - ) - name0s = tuple((name, 0) for name in self.previous_layer_names) - task = (self.fn, *name0s) - return MaterializedLayer({(name, 0): task}) - - -class AwkwardTreeReductionLayer(DataFrameTreeReduction): - def mock(self) -> AwkwardTreeReductionLayer: - return AwkwardTreeReductionLayer( - name=self.name, - name_input=self.name_input, - npartitions_input=1, - concat_func=self.concat_func, - tree_node_func=self.tree_node_func, - finalize_func=self.finalize_func, - split_every=self.split_every, - tree_node_name=self.tree_node_name, - ) +class AwkwardTreeReductionLayer(DataFrameTreeReduction): ... diff --git a/src/dask_awkward/lib/__init__.py b/src/dask_awkward/lib/__init__.py index 74d16d6c..879568cc 100644 --- a/src/dask_awkward/lib/__init__.py +++ b/src/dask_awkward/lib/__init__.py @@ -8,11 +8,7 @@ partition_compatibility, ) from dask_awkward.lib.describe import backend, fields -from dask_awkward.lib.inspect import ( - report_necessary_buffers, - report_necessary_columns, - sample, -) +from dask_awkward.lib.inspect import sample from dask_awkward.lib.io.io import ( from_awkward, from_dask_array, @@ -28,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, diff --git a/src/dask_awkward/lib/core.py b/src/dask_awkward/lib/core.py index d7c1a4e0..95c07364 100644 --- a/src/dask_awkward/lib/core.py +++ b/src/dask_awkward/lib/core.py @@ -27,6 +27,7 @@ TypeTracerArray, create_unknown_scalar, is_unknown_scalar, + touch_data, ) from dask.base import ( DaskMethodsMixin, @@ -47,6 +48,7 @@ from dask_awkward.layers import AwkwardBlockwiseLayer, AwkwardMaterializedLayer from dask_awkward.lib.optimize import all_optimizations +from dask_awkward.lib.utils import commit_to_reports from dask_awkward.utils import ( ConcretizationTypeError, DaskAwkwardNotImplemented, @@ -399,6 +401,10 @@ def name(self) -> str: def key(self) -> Key: return (self._name, 0) + @property + def report(self): + return getattr(self._meta, "_report", set()) + def _check_meta(self, m): if isinstance(m, MaybeNone): return ak.Array(m.content) @@ -533,6 +539,7 @@ def f(self, other): meta = op(other, self._meta) else: meta = op(self._meta, other) + commit_to_reports(name, self.report) return new_scalar_object(graph, name, meta=meta) return f @@ -675,7 +682,7 @@ def new_known_scalar( Examples -------- - >>> from dask_awkward.core import new_known_scalar + >>> from dask_awkward.lib.core import new_known_scalar >>> a = new_known_scalar(5, label="five") >>> a dask.awkward @@ -728,7 +735,9 @@ def _check_meta(self, m: Any | None) -> Any | None: def __getitem__(self, where): token = tokenize(self, where) new_name = f"{where}-{token}" + report = self.report new_meta = self._meta[where] + commit_to_reports(new_name, report) # first check for array type return if isinstance(new_meta, ak.Array): @@ -738,6 +747,8 @@ def __getitem__(self, where): graphlayer, dependencies=[self], ) + new_meta._report = report + hlg.layers[new_name].meta = new_meta return new_array_object(hlg, new_name, meta=new_meta, npartitions=1) # then check for scalar (or record) type @@ -748,6 +759,8 @@ def __getitem__(self, where): dependencies=[self], ) if isinstance(new_meta, ak.Record): + new_meta._report = report + hlg.layers[new_name].meta = new_meta return new_record_object(hlg, new_name, meta=new_meta) else: return new_scalar_object(hlg, new_name, meta=new_meta) @@ -821,7 +834,7 @@ def new_record_object(dsk: HighLevelGraph, name: str, *, meta: Any) -> Record: raise TypeError( f"meta Record must have a typetracer backend, not {ak.backend(meta)}" ) - return Record(dsk, name, meta) + return out def _is_numpy_or_cupy_like(arr: Any) -> bool: @@ -969,6 +982,10 @@ def reset_meta(self) -> None: """Assign an empty typetracer array as the collection metadata.""" self._meta = empty_typetracer() + @property + def report(self): + return getattr(self._meta, "_report", set()) + def repartition( self, npartitions: int | None = None, @@ -1050,6 +1067,7 @@ def repartition( new_graph = HighLevelGraph.from_collections( key, new_layer, dependencies=(self,) ) + commit_to_reports(key, self.report) return new_array_object( new_graph, key, @@ -1235,11 +1253,13 @@ def _partitions(self, index: Any) -> Array: name = f"partitions-{token}" new_keys = self.keys_array[index].tolist() dsk = {(name, i): tuple(key) for i, key in enumerate(new_keys)} + layer = AwkwardMaterializedLayer(dsk, previous_layer_names=[self.name]) graph = HighLevelGraph.from_collections( name, - AwkwardMaterializedLayer(dsk, previous_layer_names=[self.name]), + layer, dependencies=(self,), ) + layer.meta = self._meta # if a single partition was requested we trivially know the new divisions. if len(raw) == 1 and isinstance(raw[0], int) and self.known_divisions: @@ -1251,7 +1271,7 @@ def _partitions(self, index: Any) -> Array: # otherwise nullify the known divisions else: new_divisions = (None,) * (len(new_keys) + 1) # type: ignore - + commit_to_reports(name, self.report) return new_array_object( graph, name, meta=self._meta, divisions=tuple(new_divisions) ) @@ -1473,6 +1493,7 @@ def _getitem_slice_on_zero(self, where): AwkwardMaterializedLayer(dask, previous_layer_names=[self.name]), dependencies=[self], ) + commit_to_reports(name, self.report) return new_array_object( hlg, name, @@ -1583,9 +1604,14 @@ def __getitem__(self, where): raise RuntimeError("Lists containing integers are not supported.") if isinstance(where, tuple): - return self._getitem_tuple(where) - - return self._getitem_single(where) + out = self._getitem_tuple(where) + else: + out = self._getitem_single(where) + if self.report: + commit_to_reports(out.name, self.report) + out._meta._report = self._meta._report + out.dask.layers[out.name].meta = out._meta + return out def _is_method_heuristic(self, resolved: Any) -> bool: return callable(resolved) @@ -1912,10 +1938,12 @@ def partitionwise_layer( """ pairs: list[Any] = [] numblocks: dict[str, tuple[int, ...]] = {} + reps = set() for arg in args: if isinstance(arg, Array): pairs.extend([arg.name, "i"]) numblocks[arg.name] = (arg.npartitions,) + reps.update(arg.report) elif isinstance(arg, BlockwiseDep): if len(arg.numblocks) == 1: pairs.extend([arg, "i"]) @@ -1935,6 +1963,8 @@ def partitionwise_layer( ) else: pairs.extend([arg, None]) + commit_to_reports(name, reps) + layer = dask_blockwise( func, name, @@ -2018,8 +2048,23 @@ def _map_partitions( **kwargs, ) - if meta is None: - meta = map_meta(fn, *args, **kwargs) + reps = set() + try: + if meta is None: + meta = map_meta(fn, *args, **kwargs) + else: + # To do any touching?? + map_meta(fn, *args, **kwargs) + meta._report = reps + lay.meta = meta + except (AssertionError, TypeError, NotImplementedError): + [touch_data(_._meta) for _ in dak_arrays] + + for dep in dak_arrays: + for rep in dep.report: + if rep not in reps: + rep.commit(name) + reps.add(rep) hlg = HighLevelGraph.from_collections( name, @@ -2033,6 +2078,8 @@ def _map_partitions( "should be a dask_awkward.Array collection." ) dak_cache[name] = hlg, meta + if name in dak_cache: + hlg0, meta0 = dak_cache[name] in_npartitions = dak_arrays[0].npartitions in_divisions = dak_arrays[0].divisions if output_divisions is not None: @@ -2042,7 +2089,6 @@ def _map_partitions( new_divisions = tuple(map(lambda x: x * output_divisions, in_divisions)) else: new_divisions = in_divisions - if output_divisions is not None: return new_array_object( hlg, @@ -2273,10 +2319,6 @@ def non_trivial_reduction( if combiner is None: combiner = reducer - # is_positional == True is not implemented - # if is_positional: - # assert combiner is reducer - # For `axis=None`, we prepare each array to have the following structure: # [[[ ... [x1 x2 x3 ... xN] ... ]]] (length-1 outer lists) # This makes the subsequent reductions an `axis=-1` reduction @@ -2343,14 +2385,16 @@ def non_trivial_reduction( ) graph = HighLevelGraph.from_collections(name_finalize, trl, dependencies=(chunked,)) - meta = reducer( array._meta, axis=axis, keepdims=keepdims, mask_identity=mask_identity, ) + trl.meta = meta + commit_to_reports(name_finalize, array.report) if isinstance(meta, ak.highlevel.Array): + meta._report = array.report return new_array_object(graph, name_finalize, meta=meta, npartitions=1) else: return new_scalar_object(graph, name_finalize, meta=meta) @@ -2471,7 +2515,7 @@ def meta_or_identity(obj: Any) -> Any: -------- >>> import awkward as ak >>> import dask_awkward as dak - >>> from dask_awkward.core import meta_or_identity + >>> from dask_awkward.lib.core import meta_or_identity >>> x = ak.from_iter([[1, 2, 3], [4]]) >>> x = dak.from_awkward(x, npartitions=2) >>> x @@ -2667,7 +2711,7 @@ def normalize_single_outer_inner_index( Examples -------- - >>> from dask_awkward.utils import normalize_single_outer_inner_index + >>> from dask_awkward.lib.core import normalize_single_outer_inner_index >>> divisions = (0, 3, 6, 9) >>> normalize_single_outer_inner_index(divisions, 0) (0, 0) diff --git a/src/dask_awkward/lib/inspect.py b/src/dask_awkward/lib/inspect.py index 3f63fc22..b0d180bf 100644 --- a/src/dask_awkward/lib/inspect.py +++ b/src/dask_awkward/lib/inspect.py @@ -13,210 +13,6 @@ from dask_awkward.lib.core import Array -class NecessaryBuffers(NamedTuple): - data_and_shape: frozenset[str] - shape_only: frozenset[str] - - -def report_necessary_buffers( - *args: Any, traverse: bool = True -) -> dict[str, NecessaryBuffers | None]: - r"""Determine the buffer keys necessary to compute a collection. - - Parameters - ---------- - *args : Dask collections or HighLevelGraphs - The collection (or collection graph) of interest. These can be - individual objects, lists, sets, or dictionaries. - traverse : bool, optional - If True (default), builtin Python collections are traversed - looking for any Dask collections they might contain. - - Returns - ------- - dict[str, NecessaryBuffers | None] - Mapping that pairs the input layers in the graph to objects - describing the data and shape buffers that have been tagged - as required by column optimisation of the given layer. - - Examples - -------- - If we have a hypothetical parquet dataset (``ds``) with the fields - - - "foo" - - "bar" - - "baz" - - And the "baz" field has fields - - - "x" - - "y" - - The calculation of ``ds.bar + ds.baz.x`` will only require the - ``bar`` and ``baz.x`` columns from the parquet file. - - >>> import dask_awkward as dak - >>> ds = dak.from_parquet("some-dataset") - >>> ds.fields - ["foo", "bar", "baz"] - >>> ds.baz.fields - ["x", "y"] - >>> x = ds.bar + ds.baz.x - >>> dak.report_necessary_buffers(x) - { - "from-parquet-abc123": NecessaryBuffers( - data_and_shape=frozenset(...), shape_only=frozenset(...) - ) - } - - """ - import dask_awkward.lib.optimize as o - - collections, _ = unpack_collections(*args, traverse=traverse) - if not collections: - return {} - - seen_names = set() - - name_to_necessary_buffers: dict[str, NecessaryBuffers | None] = {} - for obj in collections: - dsk = obj.__dask_graph__() - keys = obj.__dask_keys__() - projection_data = o._prepare_buffer_projection(dsk, keys) - - # If the projection failed, or there are no input layers - if projection_data is None: - # Ensure that we have a record of the seen layers, if they're inputs - for name, layer in dsk.items(): - if isinstance(layer, AwkwardInputLayer): - seen_names.add(name) - continue - - # Unpack projection information - layer_to_reports, _ = projection_data - for name, report in layer_to_reports.items(): - existing_buffers = name_to_necessary_buffers.setdefault( - name, NecessaryBuffers(frozenset(), frozenset()) - ) - # Compute the shape-only keys in addition to the data and shape - data_and_shape = frozenset(report.data_touched) - shape_only = frozenset(report.shape_touched) - data_and_shape - - # Update set of touched keys - assert existing_buffers is not None - name_to_necessary_buffers[name] = NecessaryBuffers( - data_and_shape=existing_buffers.data_and_shape | data_and_shape, - shape_only=existing_buffers.shape_only | shape_only, - ) - - # Populate result with names of seen layers - for k in seen_names: - name_to_necessary_buffers.setdefault(k, None) - return name_to_necessary_buffers - - -def report_necessary_columns( - *args: Any, traverse: bool = True -) -> dict[str, frozenset[str] | None]: - r"""Get columns necessary to compute a collection - - This function is specific to sources that are columnar (e.g. Parquet). - - Parameters - ---------- - *args : Dask collections or HighLevelGraphs - The collection (or collection graph) of interest. These can be - individual objects, lists, sets, or dictionaries. - traverse : bool, optional - If True (default), builtin Python collections are traversed - looking for any Dask collections they might contain. - - Returns - ------- - dict[str, frozenset[str] | None] - Mapping that pairs the input layers in the graph to the - set of necessary IO columns that have been identified by column - optimisation of the given layer. If the layer is not backed by a - columnar source, then None is returned instead of a set. - - Examples - -------- - If we have a hypothetical parquet dataset (``ds``) with the fields - - - "foo" - - "bar" - - "baz" - - And the "baz" field has fields - - - "x" - - "y" - - The calculation of ``ds.bar + ds.baz.x`` will only require the - ``bar`` and ``baz.x`` columns from the parquet file. - - >>> import dask_awkward as dak - >>> ds = dak.from_parquet("some-dataset") - >>> ds.fields - ["foo", "bar", "baz"] - >>> ds.baz.fields - ["x", "y"] - >>> x = ds.bar + ds.baz.x - >>> dak.report_necessary_columns(x) - { - "from-parquet-abc123": frozenset({"bar", "baz.x"}) - } - - """ - import dask_awkward.lib.optimize as o - - collections, _ = unpack_collections(*args, traverse=traverse) - if not collections: - return {} - - seen_names = set() - - name_to_necessary_columns: dict[str, frozenset | None] = {} - with o.typetracer_nochecks(): - for obj in collections: - dsk = obj.__dask_graph__() - keys = obj.__dask_keys__() - projection_data = o._prepare_buffer_projection(dsk, keys) - - # If the projection failed, or there are no input layers - if projection_data is None: - # Ensure that we have a record of the seen layers, if they're inputs - for name, layer in dsk.items(): - if isinstance(layer, AwkwardInputLayer): - seen_names.add(name) - continue - - # Unpack projection information - layer_to_reports, layer_to_projection_state = projection_data - for name, report in layer_to_reports.items(): - layer = dsk.layers[name] - if not (isinstance(layer, AwkwardInputLayer) and layer.is_columnar): - continue - - existing_columns = name_to_necessary_columns.setdefault( - name, frozenset() - ) - - assert existing_columns is not None - # Update set of touched keys - name_to_necessary_columns[name] = ( - existing_columns - | layer.necessary_columns( - report=report, state=layer_to_projection_state[name] - ) - ) - - # Populate result with names of seen layers - for k in seen_names: - name_to_necessary_columns.setdefault(k, None) - return name_to_necessary_columns - - def _random_boolean_like(array_like: AwkArray, probability: float) -> AwkArray: import awkward as ak diff --git a/src/dask_awkward/lib/io/columnar.py b/src/dask_awkward/lib/io/columnar.py index 32ffb3ae..752e2369 100644 --- a/src/dask_awkward/lib/io/columnar.py +++ b/src/dask_awkward/lib/io/columnar.py @@ -1,56 +1,13 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Protocol, TypeVar, cast import awkward as ak -from awkward import Array as AwkwardArray -from awkward.forms import Form -from awkward.typetracer import typetracer_from_form, typetracer_with_report - -from dask_awkward.layers.layers import ( - BackendT, - ImplementsIOFunction, - ImplementsNecessaryColumns, -) -from dask_awkward.lib.utils import ( - METADATA_ATTRIBUTES, - FormStructure, - buffer_keys_required_to_compute_shapes, - form_with_unique_keys, - parse_buffer_key, - render_buffer_key, - trace_form_structure, - walk_graph_depth_first, -) - -if TYPE_CHECKING: - from awkward._nplikes.typetracer import TypeTracerReport log = logging.getLogger(__name__) -T = TypeVar("T") - - -class ImplementsColumnProjectionMixin(ImplementsNecessaryColumns, Protocol): - @property - def form(self) -> Form: ... - - @property - def behavior(self) -> dict | None: ... - - @property - def attrs(self) -> dict | None: ... - - def project_columns(self: T, columns: frozenset[str]) -> T: ... - def __call__(self, *args, **kwargs): ... - - -S = TypeVar("S", bound=ImplementsColumnProjectionMixin) - - -class ColumnProjectionMixin(ImplementsNecessaryColumns[FormStructure]): +class ColumnProjectionMixin: """A mixin to add column-centric buffer projection to an IO function. Classes that inherit from this mixin are assumed to be able to read at the @@ -59,120 +16,16 @@ class ColumnProjectionMixin(ImplementsNecessaryColumns[FormStructure]): when only metadata buffers are required. """ - def mock(self: S) -> AwkwardArray: - return cast( - AwkwardArray, - typetracer_from_form(self.form, behavior=self.behavior, attrs=self.attrs), - ) + def project(self, *args, **kwargs): + # default implementation does nothing + return self - def mock_empty(self: S, backend: BackendT = "cpu") -> AwkwardArray: - return cast( - AwkwardArray, + def mock_empty(self, backend: str = "cpu"): + # used by failure report generation + return ( ak.to_backend( self.form.length_zero_array(highlevel=False, behavior=self.behavior), backend, highlevel=True, ), ) - - def prepare_for_projection( - self: S, - ) -> tuple[AwkwardArray, TypeTracerReport, FormStructure]: - form = form_with_unique_keys(self.form, "@") - - # Build typetracer and associated report object - (meta, report) = typetracer_with_report( - form, - highlevel=True, - behavior=self.behavior, - buffer_key=render_buffer_key, - ) - - return ( - cast(AwkwardArray, meta), - report, - trace_form_structure(form, buffer_key=render_buffer_key), - ) - - def necessary_columns( - self: S, - report: TypeTracerReport, - state: FormStructure, - ) -> frozenset[str]: - ## Read from stash - # Form hierarchy information - form_key_to_parent_form_key = state["form_key_to_parent_form_key"] - form_key_to_child_form_keys: dict[str, list[str]] = {} - for child_key, parent_key in form_key_to_parent_form_key.items(): - form_key_to_child_form_keys.setdefault(parent_key, []).append(child_key) # type: ignore - form_key_to_form = state["form_key_to_form"] - # Buffer hierarchy information - form_key_to_buffer_keys = state["form_key_to_buffer_keys"] - # Column hierarchy information - form_key_to_path = state["form_key_to_path"] - - # Require the data of metadata buffers above shape-only requests - data_buffers = { - *report.data_touched, - *buffer_keys_required_to_compute_shapes( - parse_buffer_key, - report.shape_touched, - form_key_to_parent_form_key, - form_key_to_buffer_keys, - ), - } - - # We can't read buffers directly, but if we encounter a metadata - # buffer, then we should be able to pick any child. - paths = set() - wildcard_form_key = set() - for buffer_key in data_buffers: - form_key, attribute = parse_buffer_key(buffer_key) - if attribute in METADATA_ATTRIBUTES: - wildcard_form_key.add(form_key) - else: - paths.add(form_key_to_path[form_key]) - - # Select the most appropriate column for each wildcard - for form_key in wildcard_form_key: - # Find (DFS) any non-empty record form in any child - recursive_child_forms = ( - form_key_to_form[k] - for k in walk_graph_depth_first(form_key, form_key_to_child_form_keys) - ) - record_form_keys_with_contents = ( - f.form_key - for f in recursive_child_forms - if isinstance(f, ak.forms.RecordForm) and f.contents - ) - # Now find the deepest of such records - try: - last_record_form_key = next(record_form_keys_with_contents) - except StopIteration: - # This is a leaf! Therefore, we read this column - paths.add(form_key_to_path[form_key]) - continue - else: - # Ensure we get the "actual" last form key - for last_record_form_key in record_form_keys_with_contents: - ... - - # First see if any child is already included - for any_child_form_key in form_key_to_child_form_keys[last_record_form_key]: - any_child_path = form_key_to_path[any_child_form_key] - if any_child_path in paths: - break - # Otherwise, add the last child - else: - paths.add(any_child_path) - return frozenset({".".join(p) for p in paths if p}) - - def project( - self: S, - report: TypeTracerReport, - state: FormStructure, - ) -> ImplementsIOFunction: - if not self.use_optimization: # type: ignore[attr-defined] - return self - - return self.project_columns(self.necessary_columns(report, state)) diff --git a/src/dask_awkward/lib/io/io.py b/src/dask_awkward/lib/io/io.py index d8fd60b1..117cdb06 100644 --- a/src/dask_awkward/lib/io/io.py +++ b/src/dask_awkward/lib/io/io.py @@ -4,17 +4,14 @@ import math from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass -from functools import partial from typing import TYPE_CHECKING, Any, cast import awkward as ak -import dask.config import numpy as np from awkward.types.numpytype import primitive_to_dtype -from awkward.typetracer import length_zero_if_typetracer +from awkward.typetracer import length_zero_if_typetracer, typetracer_with_report from dask.base import flatten, tokenize from dask.highlevelgraph import HighLevelGraph -from dask.local import identity from dask.utils import funcname, is_integer, parse_bytes from fsspec.utils import infer_compression @@ -22,12 +19,7 @@ AwkwardBlockwiseLayer, AwkwardInputLayer, AwkwardMaterializedLayer, - AwkwardTreeReductionLayer, - ImplementsMocking, - ImplementsReport, - IOFunctionWithMocking, - io_func_implements_mocking, - io_func_implements_report, + io_func_implements_projection, ) from dask_awkward.lib.core import ( Array, @@ -37,7 +29,7 @@ typetracer_array, ) from dask_awkward.lib.io.columnar import ColumnProjectionMixin -from dask_awkward.utils import first, second +from dask_awkward.lib.utils import form_with_unique_keys, render_buffer_key if TYPE_CHECKING: from dask.array.core import Array as DaskArray @@ -133,10 +125,13 @@ def from_awkward( ) -class _FromListsFn: - def __init__(self, behavior: Mapping | None, attrs: Mapping[str, Any] | None): +class _FromListsFn(ColumnProjectionMixin): + def __init__( + self, behavior: Mapping | None, attrs: Mapping[str, Any] | None, form=None + ): self.behavior = behavior self.attrs = attrs + self.form = form def __call__(self, x: list) -> ak.Array: return ak.Array(x, behavior=self.behavior, attrs=self.attrs) @@ -178,12 +173,13 @@ def from_lists( """ lists = list(source) divs = (0, *np.cumsum(list(map(len, lists)))) + meta = typetracer_array(ak.Array(lists[0], attrs=attrs, behavior=behavior)) return cast( Array, from_map( - _FromListsFn(behavior=behavior, attrs=attrs), + _FromListsFn(behavior=behavior, attrs=attrs, form=meta.layout.form), lists, - meta=typetracer_array(ak.Array(lists[0], attrs=attrs, behavior=behavior)), + meta=meta, divisions=divs, label="from-lists", ), @@ -425,6 +421,8 @@ def from_dask_array( concatenate=True, ) layer = AwkwardBlockwiseLayer.from_blockwise(layer) + layer.meta = meta + meta._report = set() # just because we can't project, we shouldn't track? hlg = HighLevelGraph.from_collections(name, layer, dependencies=[array]) if np.any(np.isnan(array.chunks)): return new_array_object( @@ -620,16 +618,35 @@ def from_map( packed=packed, ) - # Special `io_func` implementations can implement mocking and optionally - # support buffer projection. - if io_func_implements_mocking(func): + kw = {} + if io_func_implements_projection(func): + # Special `io_func` implementations can do buffer projection - choosing columns + # so here we start with a blank report io_func = func - array_meta = cast(ImplementsMocking, func).mock() - # If we know the meta, we can spoof mocking - elif meta is not None: - io_func = IOFunctionWithMocking(meta, func) - array_meta = meta + array_meta, report = typetracer_with_report( + form_with_unique_keys(io_func.form, "@"), + highlevel=True, + behavior=io_func.behavior, + buffer_key=render_buffer_key, + attrs=meta._attrs if meta is not None else None, + ) + report.commit(name) + # column tracking report, not failure report, below + array_meta._report = {report} # Without `meta`, the meta will be computed by executing the graph + elif meta is not None: + # we can still track necessary columns even if we can't project + io_func = func + array_meta, report = typetracer_with_report( + form_with_unique_keys(meta.layout.form, "@"), + highlevel=True, + behavior=meta._behavior, + attrs=meta._attrs, + buffer_key=render_buffer_key, + ) + report.commit(name) + # column tracking report, not failure report, below + array_meta._report = {report} else: io_func = func array_meta = None @@ -637,56 +654,23 @@ def from_map( dsk = AwkwardInputLayer(name=name, inputs=inputs, io_func=io_func) hlg = HighLevelGraph.from_collections(name, dsk) + making_report = getattr(io_func, "return_report", False) + if making_report: + array_meta = ak.Array( + {"ioreport": ak.Array([0]).layout.to_typetracer(True), "data": array_meta} + ) + array_meta._report = {report} + if divisions is not None: - result = new_array_object(hlg, name, meta=array_meta, divisions=divisions) + result = new_array_object(hlg, name, meta=array_meta, divisions=divisions, **kw) else: - result = new_array_object(hlg, name, meta=array_meta, npartitions=len(inputs)) - - if io_func_implements_report(io_func): - if cast(ImplementsReport, io_func).return_report: - res = result.map_partitions( - first, meta=array_meta, label=label, output_divisions=1 - ) - - concat_fn = partial( - ak.concatenate, - axis=0, - ) - - split_every = dask.config.get("awkward.aggregation.split-every", 8) - - rep_trl_label = f"{label}-report" - rep_trl_token = tokenize(result, second, concat_fn, split_every) - rep_trl_name = f"{rep_trl_label}-{rep_trl_token}" - rep_trl_tree_node_name = f"{rep_trl_label}-tree-node-{rep_trl_token}" - - rep_part = result.map_partitions( - second, meta=empty_typetracer(), label=f"{label}-partitioned-report" - ) - - rep_trl = AwkwardTreeReductionLayer( - name=rep_trl_name, - name_input=rep_part.name, - npartitions_input=rep_part.npartitions, - concat_func=concat_fn, - tree_node_func=identity, - finalize_func=identity, - split_every=split_every, - tree_node_name=rep_trl_tree_node_name, - ) - - rep_graph = HighLevelGraph.from_collections( - rep_trl_name, rep_trl, dependencies=[rep_part] - ) - - rep = new_array_object( - rep_graph, - rep_trl_name, - meta=empty_typetracer(), - npartitions=len(rep_trl.output_partitions), - ) + result = new_array_object( + hlg, name, meta=array_meta, npartitions=len(inputs), **kw + ) + dsk.meta = result._meta - return res, rep + if making_report: + return result.data, result.ioreport return result diff --git a/src/dask_awkward/lib/io/json.py b/src/dask_awkward/lib/io/json.py index 9b63a493..61288d0d 100644 --- a/src/dask_awkward/lib/io/json.py +++ b/src/dask_awkward/lib/io/json.py @@ -9,6 +9,7 @@ import awkward as ak import dask from awkward.forms.form import Form +from awkward.typetracer import touch_data from dask.base import tokenize from dask.blockwise import BlockIndex from dask.core import flatten @@ -31,6 +32,7 @@ _BytesReadingInstructions, from_map, ) +from dask_awkward.lib.utils import _buf_to_col, commit_to_reports if TYPE_CHECKING: from awkward.contents.content import Content @@ -74,8 +76,9 @@ def use_optimization(self) -> bool: and self.schema is None ) - def project_columns(self, columns): - form = self.form.select_columns(columns) + def project(self, columns: list[str]): + cols = [_buf_to_col(s) for s in columns] + form = self.form.select_columns(cols) assert form is not None schema = layout_to_jsonschema(form.length_zero_array(highlevel=False)) @@ -769,6 +772,8 @@ def to_json( map_res.dask.layers[map_res.name].annotations = {"ak_output": True} name = f"to-json-{tokenize(array, path)}" dsk = {(name, 0): (lambda *_: None, map_res.__dask_keys__())} + touch_data(array._meta) + commit_to_reports(name, array.report) graph = HighLevelGraph.from_collections( name, AwkwardMaterializedLayer(dsk, previous_layer_names=[map_res.name]), diff --git a/src/dask_awkward/lib/io/parquet.py b/src/dask_awkward/lib/io/parquet.py index 9ff2a58d..8e8c17b5 100644 --- a/src/dask_awkward/lib/io/parquet.py +++ b/src/dask_awkward/lib/io/parquet.py @@ -6,12 +6,12 @@ import math import operator from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast +from typing import Any, Literal, cast import awkward as ak import awkward.operations.ak_from_parquet as ak_from_parquet -import dask from awkward.forms.form import Form +from awkward.typetracer import touch_data from dask.base import tokenize from dask.blockwise import BlockIndex from dask.highlevelgraph import HighLevelGraph @@ -23,16 +23,11 @@ from dask_awkward.lib.io.columnar import ColumnProjectionMixin from dask_awkward.lib.io.io import from_map from dask_awkward.lib.unproject_layout import unproject_layout - -if TYPE_CHECKING: - pass +from dask_awkward.lib.utils import _buf_to_col, commit_to_reports log = logging.getLogger(__name__) -T = TypeVar("T") - - def report_failure(exception, *args, **kwargs): return ak.Array( [ @@ -94,19 +89,12 @@ def __init__( def __call__(self, *args, **kwargs): ... @abc.abstractmethod - def project_columns(self, columns): ... + def project(self, columns): ... @property def return_report(self) -> bool: return self.report - @property - def use_optimization(self) -> bool: - return "parquet" in dask.config.get( - "awkward.optimization.columns-opt-formats", - default=[], - ) - def __repr__(self) -> str: s = ( "\nFromParquetFn(\n" @@ -170,16 +158,20 @@ def __call__(self, *args, **kwargs): if self.return_report: try: result = self.read_fn(source) - return result, report_success(self.columns, source) + return { + "data": result, + "ioreport": report_success(self.columns, source), + } except self.allowed_exceptions as err: - return self.mock_empty(), report_failure(err, source) + return {"data": ak.Array([]), "ioreport": report_failure(err, source)} return self.read_fn(source) - def project_columns(self, columns): + def project(self, columns): + cols = [_buf_to_col(s) for s in columns] return FromParquetFileWiseFn( fs=self.fs, - form=self.form.select_columns(columns), + form=self.form.select_columns(cols), listsep=self.listsep, unnamed_root=self.unnamed_root, original_form=self.form, @@ -235,10 +227,11 @@ def __call__(self, pair: Any) -> ak.Array: attrs=self.attrs, ) - def project_columns(self, columns): + def project(self, columns): + cols = [_buf_to_col(s) for s in columns] return FromParquetFragmentWiseFn( fs=self.fs, - form=self.form.select_columns(columns), + form=self.form.select_columns(cols), unnamed_root=self.unnamed_root, original_form=self.form, report=self.report, @@ -697,23 +690,15 @@ def to_parquet( out = new_scalar_object(graph, final_name, dtype="f8") else: final_name = name + "-finalize" - from dask_awkward.layers import AwkwardTreeReductionLayer - - layer = AwkwardTreeReductionLayer( - name=final_name, - concat_func=none_to_none, - tree_node_func=none_to_none, - name_input=map_res.name, - npartitions_input=map_res.npartitions, - finalize_func=none_to_none, - ) - graph = HighLevelGraph.from_collections( - final_name, - layer, - dependencies=[map_res], - ) - out = new_scalar_object(graph, final_name, dtype="f8") - + dsk[(final_name, 0)] = (lambda *_: None, map_res.__dask_keys__()) + graph = HighLevelGraph.from_collections( + final_name, + AwkwardMaterializedLayer(dsk, previous_layer_names=[map_res.name]), + dependencies=[map_res], + ) + touch_data(array._meta) + commit_to_reports(name, array.report) + out = new_scalar_object(graph, final_name, dtype="f8") if compute: out.compute() return None diff --git a/src/dask_awkward/lib/operations.py b/src/dask_awkward/lib/operations.py index 6f1da8e5..0b4c7ad5 100644 --- a/src/dask_awkward/lib/operations.py +++ b/src/dask_awkward/lib/operations.py @@ -1,5 +1,6 @@ from __future__ import annotations +import copy from collections.abc import Mapping from typing import TYPE_CHECKING, Any @@ -7,7 +8,6 @@ from awkward.operations.ak_concatenate import ( enforce_concatenated_form as enforce_layout_to_concatenated_form, ) -from awkward.typetracer import typetracer_from_form from dask.base import tokenize from dask.highlevelgraph import HighLevelGraph @@ -36,17 +36,10 @@ def __call__(self, *args): def _enforce_concatenated_form(array: AwkwardArray, form: Form) -> AwkwardArray: layout = ak.to_layout(array) - # TODO: should this check whether the form agrees first, or assume that the - # operation is harmless if not required? result = enforce_layout_to_concatenated_form(layout, form) return ak.Array(result, behavior=array._behavior, attrs=array._attrs) -def _concatenate_axis_0_meta(*arrays: AwkwardArray) -> AwkwardArray: - # At this stage, the metas have all been enforced to the same type - return arrays[0] - - def concatenate( arrays: list[Array], axis: int = 0, @@ -60,33 +53,23 @@ def concatenate( name = f"{label}-{token}" metas = [c._meta for c in arrays] + report = set.union(*(getattr(m, "_report", set()) for m in metas)) if len(metas) == 0: raise ValueError("Need at least one array to concatenate") - # Are we performing a _logical_ concatenation? if axis == 0: # There are two possible cases here: # 1. all arrays have identical metas — just grow the Dask collection # 2. some arrays have different metas — coerce arrays to same form - - # Drop reports from metas to avoid later touching any buffers - metas_no_report = [ - typetracer_from_form(x.layout.form, behavior=x.behavior, attrs=x._attrs) - for x in metas - ] - # Concatenate metas to determine result form - meta_no_report = ak.concatenate( - metas_no_report, axis=0, behavior=behavior, attrs=attrs - ) - intended_form = meta_no_report.layout.form + intended_form = metas[0].layout.form # If any forms aren't equal to this form, we must enforce each form to the same type if any( not m.layout.form.is_equal_to( intended_form, all_parameters=True, form_key=False ) - for m in metas + for m in metas[1:] ): arrays = [ map_partitions( @@ -109,7 +92,6 @@ def concatenate( aml = AwkwardMaterializedLayer( g, previous_layer_names=[a.name for a in arrays], - fn=_concatenate_axis_0_meta, ) else: g = { @@ -119,13 +101,22 @@ def concatenate( ) } - aml = AwkwardMaterializedLayer(g, previous_layer_names=[arrays[0].name]) + aml = AwkwardMaterializedLayer( + g, previous_layer_names=[a.name for a in arrays] + ) + from awkward.typetracer import touch_data + # TODO: touching all metas we don't pass on should not be necessary + [touch_data(m) for m in metas[1:]] + [r.commit(name) for r in report] + new_meta = copy.copy(metas[0]) + new_meta._report = report hlg = HighLevelGraph.from_collections(name, aml, dependencies=arrays) + aml.meta = new_meta return new_array_object( hlg, name, - meta=meta_no_report, + meta=new_meta, npartitions=sum(a.npartitions for a in arrays), ) diff --git a/src/dask_awkward/lib/optimize.py b/src/dask_awkward/lib/optimize.py index 6ad2e132..51cf7435 100644 --- a/src/dask_awkward/lib/optimize.py +++ b/src/dask_awkward/lib/optimize.py @@ -1,24 +1,22 @@ from __future__ import annotations -import copy import logging -import warnings -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, cast, no_type_check +import awkward as ak import dask.config from awkward.typetracer import touch_data +from dask.base import tokenize from dask.blockwise import Blockwise, fuse_roots, optimize_blockwise from dask.core import flatten from dask.highlevelgraph import HighLevelGraph -from dask.local import get_sync from dask_awkward.layers import AwkwardBlockwiseLayer, AwkwardInputLayer -from dask_awkward.lib.utils import typetracer_nochecks +from dask_awkward.lib.utils import _buf_to_col, commit_to_reports, typetracer_nochecks from dask_awkward.utils import first if TYPE_CHECKING: - from awkward._nplikes.typetracer import TypeTracerReport from dask.typing import Key log = logging.getLogger(__name__) @@ -61,8 +59,8 @@ def all_optimizations(dsk: Mapping, keys: Sequence[Key], **_: Any) -> Mapping: def optimize(dsk: HighLevelGraph, keys: Sequence[Key], **_: Any) -> Mapping: """Run optimizations specific to dask-awkward. - This is currently limited to determining the necessary columns for - input layers. + - determine the necessary columns for input layers + - fuse linear chains of blockwise operations in linear time """ if dask.config.get("awkward.optimization.enabled"): @@ -75,81 +73,9 @@ def optimize(dsk: HighLevelGraph, keys: Sequence[Key], **_: Any) -> Mapping: return dsk -def _prepare_buffer_projection( - dsk: HighLevelGraph, keys: Sequence[Key] -) -> tuple[dict[str, TypeTracerReport], dict[str, Any]] | None: - """Pair layer names with lists of necessary columns.""" - import awkward as ak - - if not _has_projectable_awkward_io_layer(dsk): - return None - - layer_to_projection_state: dict[str, Any] = {} - layer_to_reports: dict[str, TypeTracerReport] = {} - projection_layers = dict(dsk.layers) - - for name, lay in dsk.layers.items(): - if isinstance(lay, AwkwardInputLayer): - if lay.is_projectable: - # Insert mocked array into layers, replacing generation func - # Keep track of mocked state - ( - projection_layers[name], - layer_to_reports[name], - layer_to_projection_state[name], - ) = lay.prepare_for_projection() - elif lay.is_mockable: - projection_layers[name] = lay.mock() - elif hasattr(lay, "mock"): - projection_layers[name] = lay.mock() - - for name in _ak_output_layer_names(dsk): - projection_layers[name] = _mock_output(projection_layers[name]) - - hlg = HighLevelGraph(projection_layers, dsk.dependencies) - - minimal_keys: set[Key] = set() - for k in keys: - if isinstance(k, tuple) and len(k) == 2: - minimal_keys.add((k[0], 0)) - else: - minimal_keys.add(k) - - # now we try to compute for each possible output layer key (leaf - # node on partition 0); this will cause the typetacer reports to - # get correct fields/columns touched. If the result is a record or - # an array we of course want to touch all of the data/fields. - try: - for layer in hlg.layers.values(): - layer.__dict__.pop("_cached_dict", None) - results = get_sync(hlg, list(minimal_keys)) - for out in results: - if isinstance(out, (ak.Array, ak.Record)): - touch_data(out) - except Exception as err: - on_fail = dask.config.get("awkward.optimization.on-fail") - # this is the default, throw a warning but skip the optimization. - if on_fail == "warn": - warnings.warn( - COLUMN_OPT_FAILED_WARNING_MSG.format(exception=type(err), message=err) - ) - # option "pass" means do not throw warning but skip the optimization. - elif on_fail == "pass": - log.debug("Column projection optimization failed; optimization skipped.") - # option "raise" to raise the exception here - elif on_fail == "raise": - raise - else: - raise ValueError( - f"Invalid awkward.optimization.on-fail option: {on_fail}.\n" - "Valid options are 'warn', 'pass', or 'raise'." - ) - return None - else: - return layer_to_reports, layer_to_projection_state - - -def optimize_columns(dsk: HighLevelGraph, keys: Sequence[Key]) -> HighLevelGraph: +def optimize_columns( + dsk: HighLevelGraph, keys: Sequence[Key], dryrun=False +) -> HighLevelGraph: """Run column projection optimization. This optimization determines which columns from an @@ -178,70 +104,87 @@ def optimize_columns(dsk: HighLevelGraph, keys: Sequence[Key]) -> HighLevelGraph New, optimized task graph with column-projected ``AwkwardInputLayer``. """ - projection_data = _prepare_buffer_projection(dsk, keys) - if projection_data is None: - return dsk - - # Unpack result - layer_to_reports, layer_to_projection_state = projection_data - - # Project layers using projection state - layers = dict(dsk.layers) - for name, state in layer_to_projection_state.items(): - layers[name] = cast(AwkwardInputLayer, layers[name]).project( - report=layer_to_reports[name], state=state - ) - - return HighLevelGraph(layers, dsk.dependencies) + dsk2 = dsk.layers.copy() + + lays = {_[0] for _ in keys if isinstance(_, tuple)} + all_reps = set() + for ln in lays: + if ln in dsk.layers and hasattr(dsk.layers[ln], "meta"): + m = dsk.layers[ln].meta + if not isinstance(m, ak._nplikes.typetracer.MaybeNone): + # maybenone cases should already have been all touched + # but we could extract the .content here + touch_data(m) + rep = getattr(dsk.layers[ln].meta, "_report", ()) + if rep: + all_reps.update(rep) + name = tokenize("output", lays) + commit_to_reports(name, all_reps) + all_layers = tuple(dsk.layers) + (name,) + + if dryrun: + out = {} + for k, lay, cols in _optimize_columns(dsk.layers, all_layers): + if dryrun: + out[k] = cols + else: + new_lay = lay.project(cols) + dsk2[k] = new_lay + if dryrun: + return out + return HighLevelGraph(dsk2, dsk.dependencies) -def _layers_with_annotation(dsk: HighLevelGraph, key: str) -> list[str]: - return [n for n, v in dsk.layers.items() if (v.annotations or {}).get(key)] +def _optimize_columns(dsk, all_layers): + for k, lay in dsk.copy().items(): + if not isinstance(lay, AwkwardInputLayer) or not hasattr(lay, "meta"): + continue + rep = getattr(lay.meta, "_report", None) + if not rep: + continue + rep = first(rep) # each meta of an IO layer should have just one report + cols = rep.data_touched_in(all_layers) + if cols: + yield k, lay, cols -def _ak_output_layer_names(dsk: HighLevelGraph) -> list[str]: - """Get a list output layer names. - Output layer names are annotated with 'ak_output'. +def necessary_columns(*args, normalize: bool = True, trim: bool = True): + """Find the columns in each input layer that are needed by given collections Parameters ---------- - dsk : HighLevelGraph - Graph of interest. + args: dask-awkward colections or other dask objects baseed on them + normalize: if True, will transform the internal buffer-oriented representation + to column names similar to the convention used for instance by parquet. The + raw representation is the one actually passed to the IO backends during + optimization, and includes information about which component of a field + is needed (data, offsets, index, etc.) + trim: if normalize is True, setting this True will remove parent columns Returns ------- - list[str] - Names of the output layers. - + dict: the keys are the dask names of IO layers contained in the combined graph, + and for each there is a set of required columns """ - return _layers_with_annotation(dsk, "ak_output") - - -def _has_projectable_awkward_io_layer(dsk: HighLevelGraph) -> bool: - """Check if a graph at least one AwkwardInputLayer that is project-able.""" - for _, v in dsk.layers.items(): - if isinstance(v, AwkwardInputLayer) and v.is_projectable: - return True - return False - - -def _touch_all_data(*args, **kwargs): - """Mock writing an ak.Array to disk by touching data buffers.""" - for arg in args + tuple(kwargs.values()): - touch_data(arg) - - -def _mock_output(layer): - """Update a layer to run the _touch_all_data.""" - assert len(layer.dsk) == 1 - - new_layer = copy.deepcopy(layer) - mp = new_layer.dsk.copy() - for k in iter(mp.keys()): - mp[k] = (_touch_all_data,) + mp[k][1:] - new_layer.dsk = mp - return new_layer + dsk = {} + keys = [] + for arg in args: + dsk.update(arg.dask.layers) + keys.append((arg.name, 0)) + hlg = HighLevelGraph(dsk, {}) + out = optimize_columns(hlg, keys, dryrun=True) + if normalize: + for k in list(out): + # `startswith` to clobber attributes of unnamed root field + col1 = {_buf_to_col(_) for _ in out[k] if _.startswith("@.")} + if trim: + parents = {_.rsplit(".", 1)[0] for _ in col1 if "." in _} + out[k] = {_ for _ in col1 if _ not in parents} + else: + out[k] = col1 + # TODO: remove columns included in children? + return out @no_type_check @@ -298,7 +241,8 @@ def rewrite_layer_chains(dsk: HighLevelGraph, keys: Sequence[Key]) -> HighLevelG len(children) == 1 and dsk.dependencies[first(children)] == {current_layer_key} and isinstance(dsk.layers[first(children)], AwkwardBlockwiseLayer) - and len(dsk.layers[current_layer_key]) == len(dsk.layers[first(children)]) + and len(dsk.layers[current_layer_key]) + == len(dsk.layers[first(children)]) # SLOW?! and current_layer_key not in required_layers ): # walk forwards @@ -332,7 +276,7 @@ def rewrite_layer_chains(dsk: HighLevelGraph, keys: Sequence[Key]) -> HighLevelG for chain in chains: # inputs are the inputs of chain[0] # outputs are the outputs of chain[-1] - # .dsk is composed from the .dsk of each layer + # .dsk is composed of the .dsk of each layer outkey = chain[-1] layer0 = cast(Blockwise, dsk.layers[chain[0]]) outlayer = layers[outkey] @@ -392,9 +336,3 @@ def _recursive_replace(args, layer, parent, indices): else: args2.append(arg) return args2 - - -def _buffer_keys_for_layer( - buffer_keys: Iterable[str], known_buffer_keys: frozenset[str] -) -> set[str]: - return {k for k in buffer_keys if k in known_buffer_keys} diff --git a/src/dask_awkward/lib/structure.py b/src/dask_awkward/lib/structure.py index 6b70c8b9..2e14cd48 100644 --- a/src/dask_awkward/lib/structure.py +++ b/src/dask_awkward/lib/structure.py @@ -24,6 +24,7 @@ new_scalar_object, partition_compatibility, ) +from dask_awkward.lib.utils import commit_to_reports from dask_awkward.utils import ( DaskAwkwardNotImplemented, IncompatiblePartitions, @@ -337,6 +338,7 @@ def copy(array: Array) -> Array: # dask-awkward's copy is metadata-only old_meta = array._meta new_meta = ak.Array(old_meta.layout, behavior=deepcopy(old_meta._behavior)) + new_meta._report = old_meta._report return Array( array._dask, @@ -931,8 +933,10 @@ def _array_with_rebuilt_meta( behavior = array._meta.behavior new_meta = ak.Array(array._meta, behavior=behavior, attrs=attrs) - - return Array(array.dask, array.name, new_meta, array.divisions) + new_meta._report = array.report + out = Array(array.dask, array.name, new_meta, array.divisions) + commit_to_reports(out.name, array.report) + return out @borrow_docstring(ak.unzip) diff --git a/src/dask_awkward/lib/testutils.py b/src/dask_awkward/lib/testutils.py index cc55d858..ed82e11a 100644 --- a/src/dask_awkward/lib/testutils.py +++ b/src/dask_awkward/lib/testutils.py @@ -7,7 +7,6 @@ import awkward as ak import numpy as np -from awkward.typetracer import typetracer_from_form from dask.base import is_dask_collection from packaging.version import Version @@ -288,15 +287,18 @@ def __call__(self, *args, **kwargs): if self.return_report: try: result, time = time_it(self.read_fn)(*args, **kwargs) - return result, self.make_success_report(time, *args, **kwargs) + return { + "data": result, + "ioreport": self.make_success_report(time, *args, **kwargs), + } except self.allowed_exceptions as err: - return self.mock_empty(), self.make_fail_report(err, *args, **kwargs) + return { + "data": self.mock_empty(), + "ioreport": self.make_fail_report(err, *args, **kwargs), + } return self.read_fn(*args, **kwargs) - def mock(self): - return typetracer_from_form(self.form) - def mock_empty(self, backend="cpu"): return ak.to_backend( self.form.length_zero_array(highlevel=False), diff --git a/src/dask_awkward/lib/utils.py b/src/dask_awkward/lib/utils.py index 7b067386..65b8841f 100644 --- a/src/dask_awkward/lib/utils.py +++ b/src/dask_awkward/lib/utils.py @@ -10,6 +10,7 @@ if TYPE_CHECKING: from awkward.forms import Form + from awkward.typetracer import TypeTracerReport KNOWN_LENGTH_ATTRIBUTES = frozenset(("mask",)) UNKNOWN_LENGTH_ATTRIBUTES = frozenset(("offsets", "starts", "stops", "index", "tags")) @@ -178,3 +179,20 @@ def typetracer_nochecks(): TypeTracerArray.runtime_typechecks = oldval else: del TypeTracerArray.runtime_typechecks + + +def commit_to_reports(name: str, reports: Iterable[TypeTracerReport]): + for report in reports: + report.commit(name) + + +def _buf_to_col(s): + return ( + s[2:] + .replace("content.", "") + .replace("-offsets", "") + .replace("-data", "") + .replace("-index", "") + .replace("-mask", "") + .replace("-tags", "") + ) diff --git a/tests/conftest.py b/tests/conftest.py index 142980de..a96f6200 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -11,7 +11,12 @@ import dask_awkward.lib.testutils as daktu -@pytest.fixture(scope="session") +@pytest.fixture(autouse=True) +def clear_cache(): + dak.lib.core.dak_cache.clear() + + +@pytest.fixture() def single_record_file(tmp_path_factory: pytest.TempPathFactory) -> str: fname = tmp_path_factory.mktemp("data") / "single_record.json" record = {"record": [1, 2, 3]} @@ -20,7 +25,7 @@ def single_record_file(tmp_path_factory: pytest.TempPathFactory) -> str: return str(fname) -@pytest.fixture(scope="session") +@pytest.fixture() def ndjson_points1(tmp_path_factory: pytest.TempPathFactory) -> str: array = daktu.awkward_xy_points() fname = tmp_path_factory.mktemp("data") / "points_ndjson1.json" @@ -30,7 +35,7 @@ def ndjson_points1(tmp_path_factory: pytest.TempPathFactory) -> str: return str(fname) -@pytest.fixture(scope="session") +@pytest.fixture() def ndjson_points1_str(tmp_path_factory: pytest.TempPathFactory) -> str: array = daktu.awkward_xy_points_str() fname = tmp_path_factory.mktemp("data") / "points_ndjson1.json" @@ -40,7 +45,7 @@ def ndjson_points1_str(tmp_path_factory: pytest.TempPathFactory) -> str: return str(fname) -@pytest.fixture(scope="session") +@pytest.fixture() def ndjson_points2(tmp_path_factory: pytest.TempPathFactory) -> str: array = daktu.awkward_xy_points() fname = tmp_path_factory.mktemp("data") / "points_ndjson2.json" @@ -50,77 +55,77 @@ def ndjson_points2(tmp_path_factory: pytest.TempPathFactory) -> str: return str(fname) -@pytest.fixture(scope="session") +@pytest.fixture() def ndjson_points_file(ndjson_points1: str) -> str: return ndjson_points1 -@pytest.fixture(scope="session") +@pytest.fixture() def ndjson_points_file_str(ndjson_points1_str: str) -> str: return ndjson_points1_str -@pytest.fixture(scope="session") +@pytest.fixture() def daa(ndjson_points1: str) -> dak.Array: return dak.from_json([ndjson_points1] * 3) -@pytest.fixture(scope="session") +@pytest.fixture() def pq_points_dir(daa: dak.Array, tmp_path_factory: pytest.TempPathFactory) -> str: pqdir = tmp_path_factory.mktemp("pqfiles") dak.to_parquet(daa, str(pqdir)) return str(pqdir) -@pytest.fixture(scope="session") +@pytest.fixture() def daa_parquet(pq_points_dir: str) -> dak.Array: return cast(dak.Array, dak.from_parquet(pq_points_dir)) -@pytest.fixture(scope="session") +@pytest.fixture() def daa_str(ndjson_points1_str: str) -> dak.Array: return dak.from_json([ndjson_points1_str] * 3) -@pytest.fixture(scope="session") +@pytest.fixture() def caa(ndjson_points1: str) -> ak.Array: with open(ndjson_points1, "rb") as f: a = ak.from_json(f, line_delimited=True) return ak.concatenate([a, a, a]) -@pytest.fixture(scope="session") +@pytest.fixture() def caa_str(ndjson_points1_str: str) -> ak.Array: with open(ndjson_points1_str, "rb") as f: a = ak.from_json(f, line_delimited=True) return ak.concatenate([a, a, a]) -@pytest.fixture(scope="session") +@pytest.fixture() def daa_p1(ndjson_points1: str) -> dak.Array: return dak.from_json([ndjson_points1] * 3) -@pytest.fixture(scope="session") +@pytest.fixture() def daa_p2(ndjson_points2: str) -> dak.Array: return dak.from_json([ndjson_points2] * 3) -@pytest.fixture(scope="session") +@pytest.fixture() def caa_p1(ndjson_points1: str) -> ak.Array: with open(ndjson_points1) as f: lines = [json.loads(line) for line in f] return ak.Array(lines * 3) -@pytest.fixture(scope="session") +@pytest.fixture() def caa_p2(ndjson_points2: str) -> ak.Array: with open(ndjson_points2) as f: lines = [json.loads(line) for line in f] return ak.Array(lines * 3) -@pytest.fixture(scope="session") +@pytest.fixture() def L1() -> list[list[dict[str, float]]]: return [ [{"x": 1.0, "y": 1.1}, {"x": 2.0, "y": 2.2}, {"x": 3, "y": 3.3}], @@ -131,7 +136,7 @@ def L1() -> list[list[dict[str, float]]]: ] -@pytest.fixture(scope="session") +@pytest.fixture() def L2() -> list[list[dict[str, float]]]: return [ [{"x": 0.9, "y": 1.0}, {"x": 2.0, "y": 2.2}, {"x": 2.9, "y": 3.0}], @@ -142,7 +147,7 @@ def L2() -> list[list[dict[str, float]]]: ] -@pytest.fixture(scope="session") +@pytest.fixture() def L3() -> list[list[dict[str, float]]]: return [ [{"x": 1.9, "y": 9.0}, {"x": 2.0, "y": 8.2}, {"x": 9.9, "y": 9.0}], @@ -153,7 +158,7 @@ def L3() -> list[list[dict[str, float]]]: ] -@pytest.fixture(scope="session") +@pytest.fixture() def L4() -> list[list[dict[str, float]] | None]: return [ [{"x": 1.9, "y": 9.0}, {"x": 2.0, "y": 8.2}, {"x": 9.9, "y": 9.0}], @@ -164,14 +169,14 @@ def L4() -> list[list[dict[str, float]] | None]: ] -@pytest.fixture(scope="session") +@pytest.fixture() def caa_parquet(caa: ak.Array, tmp_path_factory: pytest.TempPathFactory) -> str: fname = tmp_path_factory.mktemp("parquet_data") / "caa.parquet" ak.to_parquet(caa, str(fname), extensionarray=False) return str(fname) -@pytest.fixture(scope="session") +@pytest.fixture() def unnamed_root_parquet_file(tmp_path_factory: pytest.TempPathFactory) -> str: from dask_awkward.lib.testutils import unnamed_root_ds diff --git a/tests/test_core.py b/tests/test_core.py index 3ebf0777..50b21d02 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -57,7 +57,7 @@ def test_dunder_str(caa: ak.Array) -> None: def test_calculate_known_divisions(ndjson_points_file: str) -> None: daa = dak.from_json([ndjson_points_file] * 3) target = (0, 5, 10, 15) - assert calculate_known_divisions(daa) == target + # assert calculate_known_divisions(daa) == target assert calculate_known_divisions(daa.points) == target assert calculate_known_divisions(daa.points.x) == target assert calculate_known_divisions(daa["points"]["y"]) == target @@ -418,6 +418,7 @@ def test_record_dir(daa: Array) -> None: # import pickle # r = daa[0] + # assert type(r) is dak.Record # assert isinstance(r._meta, ak.Record) @@ -489,8 +490,8 @@ def test_scalar_pickle(daa: Array) -> None: # TODO: workaround since dask un/pack disappeared for lay2, lay1 in zip(s2.dask.layers.values(), s1.dask.layers.values()): - if hasattr(lay1, "_meta"): - lay2._meta = lay1._meta + if hasattr(lay1, "meta"): + lay2.meta = lay1.meta assert_eq(s1.compute(), s2.compute()) assert s1.known_value is None @@ -864,6 +865,7 @@ def test_map_partitions_args_and_kwargs_have_collection(): def test_dask_array_in_map_partitions(daa, caa): + daa.eager_compute_divisions() x1 = dak.zeros_like(daa.points.x) y1 = da.ones(len(x1), chunks=x1.divisions[1]) z1 = x1 + y1 diff --git a/tests/test_inspect.py b/tests/test_inspect.py index b4e2651b..baea4438 100644 --- a/tests/test_inspect.py +++ b/tests/test_inspect.py @@ -10,68 +10,6 @@ test_uproot_path = Path(__file__).parent / "test-uproot" -def test_report_necessary_buffers( - daa: dak.Array, tmpdir_factory: pytest.TempdirFactory -) -> None: - z = daa.points.x + daa.points.y - for k, v in dak.report_necessary_buffers(z).items(): - assert v == ( - frozenset( - { - "@.points-offsets", - "@.points.content.y-data", - "@.points.content.x-data", - } - ), - frozenset(), - ) - - w = dak.to_parquet( - daa.points.x, str(Path(tmpdir_factory.mktemp("pq")) / "out"), compute=False - ) - for k, v in dak.report_necessary_buffers(w).items(): - assert v == ( - frozenset({"@.points-offsets", "@.points.content.x-data"}), - frozenset(), - ) - - q = {"z": z, "w": w} - for k, v in dak.report_necessary_buffers(q).items(): - assert v == ( - frozenset( - { - "@.points-offsets", - "@.points.content.x-data", - "@.points.content.y-data", - } - ), - frozenset(), - ) - - z = dak.zeros_like(daa.points.x) - for k, v in dak.report_necessary_buffers(z).items(): - assert v == ( - frozenset({"@.points-offsets"}), - frozenset({"@.points.content.x-data"}), - ) - - -def test_report_necessary_columns(daa: dak.Array) -> None: - result = dak.min(daa.points.x, axis=1) - rep = dak.report_necessary_columns(result) - for k, v in rep.items(): - assert v is not None - assert sorted(["points.x"]) == sorted(v) - - result = dak.zeros_like(daa.points.y) - rep = dak.report_necessary_columns(result) - for k, v in rep.items(): - assert v is not None - points, coord = list(v)[0].split(".") - assert points == "points" - assert coord in ["x", "y"] - - def test_visualize_works(daa): query = daa.points.x @@ -97,9 +35,9 @@ def test_basic_root_works(): .fill(events.MET_pt) ) - columns = list(dak.necessary_columns(q1_hist).values())[0] - assert columns == frozenset({"MET_pt"}) dask.compute(q1_hist) + columns = list(dak.necessary_columns(q1_hist).values())[0] + assert columns == {"MET_pt"} def test_sample(daa): diff --git a/tests/test_io.py b/tests/test_io.py index fa1eac06..9c289ff7 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -360,14 +360,14 @@ def test_from_map_random_fail_from_lists(): many, meta=typetracer_array(ak.Array(many[0])), divisions=divs, - label="from-lists", + label="from-lists01", ) assert len(array.compute()) < (len(single) * len(many)) computed_report = report.compute() # we expect the 'args' field in the report to be empty if the - # from_map node succeded; so we use ak.num(..., axis=1) to filter + # from_map node succeeded; so we use ak.num(..., axis=1) to filter # those out. succ = ak.num(computed_report["args"], axis=1) == 0 fail = np.invert(succ) @@ -382,7 +382,7 @@ def test_from_map_random_fail_from_lists(): many, meta=typetracer_array(ak.Array(many[0])), divisions=divs, - label="from-lists", + label="from-lists02", ) array.compute() @@ -392,7 +392,7 @@ def test_from_map_random_fail_from_lists(): many, meta=typetracer_array(ak.Array(many[0])), divisions=divs, - label="from-lists", + label="from-lists03", ) cast(dak.Array, array).compute() @@ -434,7 +434,7 @@ def point_abs(self): caa = ak.with_name(caa.points, name="Point", behavior=behavior) daa = dak.from_awkward(caa, npartitions=2, behavior=behavior) assert_eq(caa.xsq, daa.xsq) - assert set(first(dak.necessary_columns(daa.xsq).items())[1]) == {"x"} - assert set(first(dak.necessary_columns(daa).items())[1]) == {"x", "y"} - assert set(first(dak.necessary_columns(np.abs(daa)).items())[1]) == {"x", "y"} + assert set(first(dak.necessary_columns(daa.xsq).values())) == {"x"} + assert set(first(dak.necessary_columns(daa).values())) == {"x", "y"} + assert set(first(dak.necessary_columns(np.abs(daa)).values())) == {"x", "y"} assert_eq(np.abs(caa), np.abs(daa)) diff --git a/tests/test_io_json.py b/tests/test_io_json.py index 688fb550..d82cba1e 100644 --- a/tests/test_io_json.py +++ b/tests/test_io_json.py @@ -91,26 +91,14 @@ def input_layer_array_partition0(collection: Array) -> ak.Array: """ with dask.config.set({"awkward.optimization.which": ["columns"]}): - optimized_hlg = dak_optimize(collection.dask, collection.keys) # type: ignore - layers = list(optimized_hlg.layers) # type: ignore + optimized_hlg = dak_optimize(collection.dask, collection.keys) + layers = list(optimized_hlg.layers) layer_name = [name for name in layers if name.startswith("from-json")][0] sgc, arg = optimized_hlg[(layer_name, 0)] array = sgc.dsk[layer_name][0](arg) return array -def test_json_column_projection_off(json_data_dir: Path) -> None: - source = os.path.join(str(json_data_dir), "*.json") - ds = dak.from_json(source) - fields_to_keep = ["name", "goals"] - - ds2 = ds[fields_to_keep] - with dask.config.set({"awkward.optimization.columns-opt-formats": []}): - array = input_layer_array_partition0(ds2) - - assert array.fields == ["name", "team", "goals"] - - def test_json_column_projection1(json_data_dir: Path) -> None: source = os.path.join(str(json_data_dir), "*.json") ds = dak.from_json(source) diff --git a/tests/test_operations.py b/tests/test_operations.py index 14a43fd9..352063e6 100644 --- a/tests/test_operations.py +++ b/tests/test_operations.py @@ -10,25 +10,13 @@ @pytest.mark.parametrize("axis", [0, 1]) def test_concatenate_simple(daa, caa, axis): + # daa.x and daa.y have the same types assert_eq( ak.concatenate([caa.points.x, caa.points.y], axis=axis), dak.concatenate([daa.points.x, daa.points.y], axis=axis), ) -def test_concatenate_axis_0_logical_same(daa): - result = dak.concatenate([daa, daa], axis=0) - buffers_report = dak.report_necessary_buffers(result.points.x) - assert len(buffers_report) == 1 - - buffers = next(iter(buffers_report.values())) - - assert buffers.data_and_shape == frozenset( - ["@.points.content.x-data", "@.points-offsets"] - ) - assert buffers.shape_only == frozenset() - - def test_concatenate_axis_0_logical_different(daa): import dask.config @@ -55,15 +43,10 @@ def test_concatenate_axis_0_logical_different(daa): empty_array = ak.Array(empty_form.length_zero_array(highlevel=False)) empty_dak_array = dak.from_awkward(empty_array, npartitions=1) result = dak.concatenate([daa, empty_dak_array], axis=0) - - buffers_report = dak.report_necessary_buffers(result.points.x) - assert len(buffers_report) == 2 - - buffers = next(iter(buffers_report.values())) - assert buffers.data_and_shape == frozenset( - ["@.points.content.x-data", "@.points.content.y-data", "@.points-offsets"] - ) - assert buffers.shape_only == frozenset() + expected = ak.concatenate( + [daa.compute(), empty_dak_array.compute()], axis=0 + ).tolist() + assert expected == result.compute().tolist() @pytest.mark.parametrize("axis", [0, 1, 2]) diff --git a/tests/test_optimize.py b/tests/test_optimize.py index 80fb975f..c737421b 100644 --- a/tests/test_optimize.py +++ b/tests/test_optimize.py @@ -144,10 +144,10 @@ def something(j, k): a_compute = something(a, b) nc1 = dak.necessary_columns(a_compute) - assert sorted(list(nc1.items())[0][1]) == ["points.x"] + assert list(nc1.values())[0] == {"points.x"} - nc2 = dak.necessary_columns(a_compute, (a, b)) - assert sorted(list(nc2.items())[0][1]) == ["points.x", "points.y", "points.z"] + nc2 = dak.necessary_columns(a_compute, a, b) + assert list(nc2.items())[0][1] == {"points.x", "points.y", "points.z"} x, (y, z) = dask.compute(a_compute, (a, b)) assert str(x) diff --git a/tests/test_utils.py b/tests/test_utils.py index 3d1b38c7..9a65e912 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -84,7 +84,8 @@ def test_field_access_to_front(pairs): def test_nocheck_context(): from awkward._nplikes.typetracer import TypeTracerArray - assert getattr(TypeTracerArray, "runtime_typechecks", True) + old = getattr(TypeTracerArray, "runtime_typechecks", True) with typetracer_nochecks(): assert not TypeTracerArray.runtime_typechecks - assert getattr(TypeTracerArray, "runtime_typechecks", True) + new = getattr(TypeTracerArray, "runtime_typechecks", True) + assert old == new