def test_replace_name_in_keys():
assert replace_name_in_key("foo", "bar") == "bar"
assert replace_name_in_key("foo-123", "bar-456") == "bar-456"
h1 = object() # Arbitrary hashables
h2 = object()
assert replace_name_in_key(("foo-123", h1, h2), "bar") == ("bar", h1, h2)
with pytest.raises(TypeError):
replace_name_in_key(1, "foo")
def _build_map_layer(
func: Callable,
prev_name: str,
new_name: str,
collection,
dependencies: tuple[Delayed, ...] = (),
) -> Layer:
"""Apply func to all keys of collection. Create a Blockwise layer whenever possible;
fall back to MaterializedLayer otherwise.
Parameters
----------
func
Callable to be invoked on the graph node
prev_name : str
name of the layer to map from; in case of dask base collections, this is the
collection name. Note how third-party collections, e.g. xarray.Dataset, can
have multiple names.
new_name : str
name of the layer to map to
collection
Arbitrary dask collection
dependencies
Zero or more Delayed objects, which will be passed as arbitrary variadic args to
func after the collection's chunk
"""
if _can_apply_blockwise(collection):
# Use a Blockwise layer
try:
numblocks = collection.numblocks
except AttributeError:
numblocks = (collection.npartitions, )
indices = tuple(i for i, _ in enumerate(numblocks))
kwargs = {
"_deps": [d.key for d in dependencies]
} if dependencies else {}
return blockwise(
func,
new_name,
indices,
prev_name,
indices,
numblocks={prev_name: numblocks},
dependencies=dependencies,
**kwargs,
)
else:
# Delayed, bag.Item, dataframe.core.Scalar, or third-party collection;
# fall back to MaterializedLayer
dep_keys = tuple(d.key for d in dependencies)
return MaterializedLayer({
replace_name_in_key(k, {prev_name: new_name}): (func, k) + dep_keys
for k in flatten(collection.__dask_keys__())
if get_name_from_key(k) == prev_name
})
def _rebuild(self, dsk, *, rename=None):
key = replace_name_in_key(self.key, rename) if rename else self.key
if isinstance(dsk, HighLevelGraph) and len(dsk.layers) == 1:
# FIXME Delayed is currently the only collection type that supports both high- and low-level graphs.
# The HLG output of `optimize` will have a layer name that doesn't match `key`.
# Remove this when Delayed is HLG-only (because `optimize` will only be passed HLGs, so it won't have
# to generate random layer names).
layer = next(iter(dsk.layers))
else:
layer = None
return Delayed(key, dsk, self._length, layer=layer)
def _rebuild(dsk, keys, *, rename=None):
if rename:
keys = [replace_name_in_key(key, rename) for key in keys]
return Tuple(dsk, keys)
def test_replace_name_in_keys():
assert replace_name_in_key("foo", {}) == "foo"
assert replace_name_in_key("foo", {"bar": "baz"}) == "foo"
assert replace_name_in_key("foo", {"foo": "bar", "baz": "asd"}) == "bar"
assert replace_name_in_key("foo-123", {"foo-123": "bar-456"}) == "bar-456"
h1 = object() # Arbitrary hashables
h2 = object()
assert replace_name_in_key(("foo-123", h1, h2), {"foo-123": "bar"}) == (
"bar",
h1,
h2,
)
with pytest.raises(TypeError):
replace_name_in_key(1, {})
with pytest.raises(TypeError):
replace_name_in_key((), {})
with pytest.raises(TypeError):
replace_name_in_key((1, ), {})
def _rebuild(dsk, keys, name=None):
if name is not None:
keys = [replace_name_in_key(key, name) for key in keys]
return Tuple(dsk, keys)
