Skip to content

Commit

Permalink
Remove opt_touch_all from map_parititions
Browse files Browse the repository at this point in the history
  • Loading branch information
douglasdavis committed Oct 27, 2023
1 parent c9f10ab commit 260a996
Showing 3 changed files with 0 additions and 33 deletions.
8 changes: 0 additions & 8 deletions src/dask_awkward/lib/core.py
Original file line number Diff line number Diff line change
@@ -1539,7 +1539,6 @@ def partitionwise_layer(
func: Callable,
name: str,
*args: Any,
opt_touch_all: bool = False,
**kwargs: Any,
) -> AwkwardBlockwiseLayer:
"""Create a partitionwise graph layer.
@@ -1594,8 +1593,6 @@ def partitionwise_layer(
**kwargs,
)
layer = AwkwardBlockwiseLayer.from_blockwise(layer)
if opt_touch_all:
layer._opt_touch_all = True
return layer


@@ -1640,7 +1637,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:
@@ -1678,9 +1674,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
@@ -1762,7 +1755,6 @@ def map_partitions(
name,
*arg_flat_deps_expanded,
*kwarg_flat_deps,
opt_touch_all=opt_touch_all,
)

if meta is None:
24 changes: 0 additions & 24 deletions src/dask_awkward/lib/optimize.py
Original file line number Diff line number Diff line change
@@ -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.
1 change: 0 additions & 1 deletion src/dask_awkward/lib/structure.py
Original file line number Diff line number Diff line change
@@ -1102,7 +1102,6 @@ def zip(
*colls,
label="zip",
meta=meta,
opt_touch_all=True,
)

elif isinstance(arrays, (list, tuple)):

0 comments on commit 260a996

Please # to comment.