import pytest pytest.importorskip('numpy') import numpy as np import dask import dask.array as da from dask.optimize import fuse from dask.utils import SerializableLock from dask.array.core import getter, getter_nofancy from dask.array.optimization import (getitem, optimize, optimize_slices, fuse_slice) from dask.array.utils import assert_eq def test_fuse_getitem(): pairs = [((getter, (getter, 'x', slice(1000, 2000)), slice(15, 20)), (getter, 'x', slice(1015, 1020))), ((getitem, (getter, 'x', (slice(1000, 2000), slice(100, 200))), (slice(15, 20), slice(50, 60))), (getter, 'x', (slice(1015, 1020), slice(150, 160)))), ((getitem, (getter_nofancy, 'x', (slice(1000, 2000), slice(100, 200))), (slice(15, 20), slice(50, 60))), (getter_nofancy, 'x', (slice(1015, 1020), slice(150, 160)))), ((getter, (getter, 'x', slice(1000, 2000)), 10), (getter, 'x', 1010)), ((getitem, (getter, 'x', (slice(1000, 2000), 10)), (slice(15, 20),)), (getter, 'x', (slice(1015, 1020), 10))), ((getitem, (getter_nofancy, 'x', (slice(1000, 2000), 10)), (slice(15, 20),)), (getter_nofancy, 'x', (slice(1015, 1020), 10))), ((getter, (getter, 'x', (10, slice(1000, 2000))), (slice(15, 20), )), (getter, 'x', (10, slice(1015, 1020)))), ((getter, (getter, 'x', (slice(1000, 2000), slice(100, 200))), (slice(None, None), slice(50, 60))), (getter, 'x', (slice(1000, 2000), slice(150, 160)))), ((getter, (getter, 'x', (None, slice(None, None))), (slice(None, None), 5)), (getter, 'x', (None, 5))), ((getter, (getter, 'x', (slice(1000, 2000), slice(10, 20))), (slice(5, 10),)), (getter, 'x', (slice(1005, 1010), slice(10, 20)))), ((getitem, (getitem, 'x', (slice(1000, 2000),)), (slice(5, 10), slice(10, 20))), (getitem, 'x', (slice(1005, 1010), slice(10, 20)))), ((getter, (getter, 'x', slice(1000, 2000), False, False), slice(15, 20)), (getter, 'x', slice(1015, 1020))), ((getter, (getter, 'x', slice(1000, 2000)), slice(15, 20), False, False), (getter, 'x', slice(1015, 1020))), ((getter, (getter_nofancy, 'x', slice(1000, 2000), False, False), slice(15, 20), False, False), (getter_nofancy, 'x', slice(1015, 1020), False, False)), ] for inp, expected in pairs: result = optimize_slices({'y': inp}) assert result == {'y': expected} def test_fuse_getitem_lock(): lock1 = SerializableLock() lock2 = SerializableLock() pairs = [((getter, (getter, 'x', slice(1000, 2000), True, lock1), slice(15, 20)), (getter, 'x', slice(1015, 1020), True, lock1)), ((getitem, (getter, 'x', (slice(1000, 2000), slice(100, 200)), True, lock1), (slice(15, 20), slice(50, 60))), (getter, 'x', (slice(1015, 1020), slice(150, 160)), True, lock1)), ((getitem, (getter_nofancy, 'x', (slice(1000, 2000), slice(100, 200)), True, lock1), (slice(15, 20), slice(50, 60))), (getter_nofancy, 'x', (slice(1015, 1020), slice(150, 160)), True, lock1)), ((getter, (getter, 'x', slice(1000, 2000), True, lock1), slice(15, 20), True, lock2), (getter, (getter, 'x', slice(1000, 2000), True, lock1), slice(15, 20), True, lock2))] for inp, expected in pairs: result = optimize_slices({'y': inp}) assert result == {'y': expected} def test_optimize_with_getitem_fusion(): dsk = {'a': 'some-array', 'b': (getter, 'a', (slice(10, 20), slice(100, 200))), 'c': (getter, 'b', (5, slice(50, 60)))} result = optimize(dsk, ['c']) expected_task = (getter, 'some-array', (15, slice(150, 160))) assert any(v == expected_task for v in result.values()) assert len(result) < len(dsk) def test_optimize_slicing(): dsk = {'a': (range, 10), 'b': (getter, 'a', (slice(None, None, None),)), 'c': (getter, 'b', (slice(None, None, None),)), 'd': (getter, 'c', (slice(0, 5, None),)), 'e': (getter, 'd', (slice(None, None, None),))} expected = {'e': (getter, (range, 10), (slice(0, 5, None),))} result = optimize_slices(fuse(dsk, [], rename_keys=False)[0]) assert result == expected # protect output keys expected = {'c': (getter, (range, 10), (slice(0, None, None),)), 'd': (getter, 'c', (slice(0, 5, None),)), 'e': (getter, 'd', (slice(None, None, None),))} result = optimize_slices(fuse(dsk, ['c', 'd', 'e'], rename_keys=False)[0]) assert result == expected def test_fuse_slice(): assert fuse_slice(slice(10, 15), slice(0, 5, 2)) == slice(10, 15, 2) assert (fuse_slice((slice(100, 200),), (None, slice(10, 20))) == (None, slice(110, 120))) assert (fuse_slice((slice(100, 200),), (slice(10, 20), None)) == (slice(110, 120), None)) assert (fuse_slice((1,), (None,)) == (1, None)) assert (fuse_slice((1, slice(10, 20)), (None, None, 3, None)) == (1, None, None, 13, None)) with pytest.raises(NotImplementedError): fuse_slice(slice(10, 15, 2), -1) def test_fuse_slice_with_lists(): assert fuse_slice(slice(10, 20, 2), [1, 2, 3]) == [12, 14, 16] assert fuse_slice([10, 20, 30, 40, 50], [3, 1, 2]) == [40, 20, 30] assert fuse_slice([10, 20, 30, 40, 50], 3) == 40 assert fuse_slice([10, 20, 30, 40, 50], -1) == 50 assert fuse_slice([10, 20, 30, 40, 50], slice(1, None, 2)) == [20, 40] assert fuse_slice((slice(None), slice(0, 10), [1, 2, 3]), (slice(None), slice(1, 5), slice(None))) == (slice(0, None), slice(1, 5), [1, 2, 3]) assert fuse_slice((slice(None), slice(None), [1, 2, 3]), (slice(None), slice(1, 5), 1)) == (slice(0, None), slice(1, 5), 2) def test_nonfusible_fancy_indexing(): nil = slice(None) cases = [# x[:, list, :][int, :, :] ((nil, [1, 2, 3], nil), (0, nil, nil)), # x[int, :, :][:, list, :] ((0, nil, nil), (nil, [1, 2, 3], nil)), # x[:, list, :, :][:, :, :, int] ((nil, [1, 2], nil, nil), (nil, nil, nil, 0))] for a, b in cases: with pytest.raises(NotImplementedError): fuse_slice(a, b) def test_hard_fuse_slice_cases(): dsk = {'x': (getter, (getter, 'x', (None, slice(None, None))), (slice(None, None), 5))} assert optimize_slices(dsk) == {'x': (getter, 'x', (None, 5))} def test_dont_fuse_numpy_arrays(): x = np.ones(10) for chunks in [(5,), (10,)]: y = da.from_array(x, chunks=(10,)) dsk = y._optimize(y.dask, y._keys()) assert sum(isinstance(v, np.ndarray) for v in dsk.values()) == 1 def test_minimize_data_transfer(): x = np.ones(100) y = da.from_array(x, chunks=25) z = y + 1 dsk = z._optimize(z.dask, z._keys()) keys = list(dsk) results = dask.get(dsk, keys) big_key = [k for k, r in zip(keys, results) if r is x][0] dependencies, dependents = dask.core.get_deps(dsk) deps = dependents[big_key] assert len(deps) == 4 for dep in deps: assert dsk[dep][0] in (getitem, getter) assert dsk[dep][1] == big_key def test_fuse_slices_with_alias(): dsk = {'x': np.arange(16).reshape((4, 4)), ('dx', 0, 0): (getter, 'x', (slice(0, 4), slice(0, 4))), ('alias', 0, 0): ('dx', 0, 0), ('dx2', 0): (getitem, ('alias', 0, 0), (slice(None), 0))} keys = [('dx2', 0)] dsk2 = optimize(dsk, keys) assert len(dsk2) == 3 fused_key = set(dsk2).difference(['x', ('dx2', 0)]).pop() assert dsk2[fused_key] == (getter, 'x', (slice(0, 4), 0)) def test_dont_fuse_fancy_indexing_in_getter_nofancy(): dsk = {'a': (getitem, (getter_nofancy, 'x', (slice(10, 20, None), slice(100, 200, None))), ([1, 3], slice(50, 60, None)))} assert optimize_slices(dsk) == dsk dsk = {'a': (getitem, (getter_nofancy, 'x', [1, 2, 3]), 0)} assert optimize_slices(dsk) == dsk @pytest.mark.parametrize('chunks', [10, 5, 3]) def test_fuse_getter_with_asarray(chunks): x = np.ones(10) * 1234567890 y = da.ones(10, chunks=chunks) z = x + y dsk = z._optimize(z.dask, z._keys()) assert any(v is x for v in dsk.values()) for v in dsk.values(): s = str(v) assert s.count('getitem') + s.count('getter') <= 1 if v is not x: assert '1234567890' not in s n_getters = len([v for v in dsk.values() if v[0] in (getitem, getter)]) if y.npartitions > 1: assert n_getters == y.npartitions else: assert n_getters == 0 assert_eq(z, x + 1) def test_turn_off_fusion(): x = da.ones(10, chunks=(5,)) y = da.sum(x + 1 + 2 + 3) a = y._optimize(y.dask, y._keys()) with dask.set_options(fuse_ave_width=0): b = y._optimize(y.dask, y._keys()) assert dask.get(a, y._keys()) == dask.get(b, y._keys()) assert len(a) < len(b)