diff --git a/src/dask_awkward/lib/core.py b/src/dask_awkward/lib/core.py index 017b705d..0babd032 100644 --- a/src/dask_awkward/lib/core.py +++ b/src/dask_awkward/lib/core.py @@ -1568,7 +1568,6 @@ def partitionwise_layer( func: Callable, name: str, *args: Any, - opt_touch_all: bool = False, **kwargs: Any, ) -> AwkwardBlockwiseLayer: """Create a partitionwise graph layer. @@ -1623,8 +1622,6 @@ def partitionwise_layer( **kwargs, ) layer = AwkwardBlockwiseLayer.from_blockwise(layer) - if opt_touch_all: - layer._opt_touch_all = True return layer @@ -1669,7 +1666,6 @@ def map_partitions( token: str | None = None, meta: Any | None = None, output_divisions: int | None = None, - opt_touch_all: bool = False, traverse: bool = True, **kwargs: Any, ) -> Array: @@ -1707,9 +1703,6 @@ def map_partitions( value greater than 1 means the divisions were expanded by some operation. This argument is mainly for internal library function implementations. - opt_touch_all : bool - Touch all layers in this graph during typetracer based - optimization. traverse : bool Unpack basic python containers to find dask collections. **kwargs : Any @@ -1747,6 +1740,14 @@ def map_partitions( This is effectively the same as `d = c * a` """ + opt_touch_all = kwargs.pop("opt_touch_all", None) + if opt_touch_all is not None: + warnings.warn( + "The opt_touch_all argument does nothing.\n" + "This warning will be removed in a future version of dask-awkward " + "and the function call will likely fail." + ) + token = token or tokenize(base_fn, *args, meta, **kwargs) label = hyphenize(label or funcname(base_fn)) name = f"{label}-{token}" @@ -1791,7 +1792,6 @@ def map_partitions( name, *arg_flat_deps_expanded, *kwarg_flat_deps, - opt_touch_all=opt_touch_all, ) if meta is None: diff --git a/src/dask_awkward/lib/optimize.py b/src/dask_awkward/lib/optimize.py index d4d3c619..2f971873 100644 --- a/src/dask_awkward/lib/optimize.py +++ b/src/dask_awkward/lib/optimize.py @@ -115,9 +115,6 @@ def _prepare_buffer_projection( for name in _ak_output_layer_names(dsk): projection_layers[name] = _mock_output(projection_layers[name]) - for name in _opt_touch_all_layer_names(dsk): - projection_layers[name] = _touch_and_call(projection_layers[name]) - hlg = HighLevelGraph(projection_layers, dsk.dependencies) # this loop builds up what are the possible final leaf nodes by @@ -235,11 +232,6 @@ def _ak_output_layer_names(dsk: HighLevelGraph) -> list[str]: return _layers_with_annotation(dsk, "ak_output") -def _opt_touch_all_layer_names(dsk: HighLevelGraph) -> list[str]: - return [n for n, v in dsk.layers.items() if hasattr(v, "_opt_touch_all")] - # return _layers_with_annotation(dsk, "ak_touch_all") - - 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(): @@ -266,22 +258,6 @@ def _mock_output(layer): return new_layer -def _touch_and_call_fn(fn, *args, **kwargs): - _touch_all_data(*args, **kwargs) - return fn(*args, **kwargs) - - -def _touch_and_call(layer): - assert len(layer.dsk) == 1 - - new_layer = copy.deepcopy(layer) - mp = new_layer.dsk.copy() - for k in iter(mp.keys()): - mp[k] = (_touch_and_call_fn,) + mp[k] - new_layer.dsk = mp - return new_layer - - def rewrite_layer_chains(dsk: HighLevelGraph, keys: Any) -> HighLevelGraph: """Smush chains of blockwise layers into a single layer. diff --git a/src/dask_awkward/lib/structure.py b/src/dask_awkward/lib/structure.py index 0ceebf32..025b9375 100644 --- a/src/dask_awkward/lib/structure.py +++ b/src/dask_awkward/lib/structure.py @@ -1261,7 +1261,6 @@ def zip( *colls, label="zip", meta=meta, - opt_touch_all=True, ) elif isinstance(arrays, Sequence):