Function test_blockwise_cull

dask/tests/test_highgraph.py:235–265

(flat)

@pytest.mark.parametrize("flat", [True, False])
def test_blockwise_cull(flat):
    np = pytest.importorskip("numpy")
    da = pytest.importorskip("dask.array")
    if flat:
        # Simple "flat" mapping between input and
        # output indices
        x = da.from_array(np.arange(40).reshape((4, 10)), (2, 4)) + 100
    else:
        # Complex mapping between input and output
        # indices (outer product and transpose)
        x = da.from_array(np.arange(10).reshape((10,)), (4,))
        y = da.from_array(np.arange(10).reshape((10,)), (4,))
        x = da.outer(x, y).transpose()

    # Check that blockwise culling results in correct
    # output keys and that full graph is not materialized
    dsk = x.__dask_graph__()
    select = (1, 1)  # Select a single chunk
    keys = {(x._name, *select)}
    dsk_cull = dsk.cull(keys)
    for name, layer in dsk_cull.layers.items():
        old_name = name.rsplit("-", 1)[0]
        if not isinstance(layer, dask.blockwise.Blockwise):
            # The original layer shouldn't be Blockwise if the new one isn't
            assert not isinstance(dsk.layers[old_name], dask.blockwise.Blockwise)
            continue
        assert isinstance(dsk.layers[old_name], dask.blockwise.Blockwise)
        assert not layer.is_materialized()
        out_keys = layer.get_output_keys()
        assert out_keys == {(layer.output, *select)}
        assert not layer.is_materialized()


def test_len_does_not_materialize():