Repository: dask/dask
File: dask/array/tests/test_array_core.py (Python)

Code before:
    assert isinstance(v, Delayed)

    # You need a well-formed HighLevelgraph for e.g. dask.graph_manipulation.bind
    for layer in v.__dask_layers__():
        assert layer in v.dask.layers

    assert (at == 0).all() and (bt == 0).all()
    assert all([ev is None for ev in v.compute()])
    assert (at == 2).all() and (bt == 3).all()

    at = np.zeros(shape=(4, 4))
    bt = np.zeros(shape=(4, 4))

    dat, dbt = store([a, b], [at, bt], compute=False, return_stored=True)
    assert isinstance(dat, Array) and isinstance(dbt, Array)
    assert (at == 0).all() and (bt == 0).all()
    assert (dat.compute() == at).all() and (dbt.compute() == bt).all()
    assert (at == 2).all() and (bt == 3).all()


def test_store_nocompute_regions():
    x = da.ones(10, chunks=1)
    y = np.zeros((2, 10))
    d1 = da.store(x, y, regions=(0,), compute=False)
    d2 = da.store(x, y, regions=(1,), compute=False)
    assert d1.key != d2.key


class ThreadSafetyError(Exception):
    pass


class NonthreadSafeStore:
    def __init__(self):
        self.in_use = False

    def __setitem__(self, key, value):
        if self.in_use:
            raise ThreadSafetyError()
        self.in_use = True
        time.sleep(0.001)
        self.in_use = False


class ThreadSafeStore:
    def __init__(self):
        self.concurrent_uses = 0
        self.max_concurrent_uses = 0

    def __setitem__(self, key, value):
        self.concurrent_uses += 1

Review comment (refactor):
Note this will return an (8,8) array not (4,4) for `dat` and `bat` since each chunk returns a (4,4) array.

https://github.com/dask/dask/blob/fa8fecf10a94971f2f31df57d504d25bef4dd57e/dask/array/core.py#L4550-L4555

This is ugly, but totally backwards compatible, and I don't expect anyone will use it other than icechunk :).

I use this like-so in icechunk, where I extract the Zarr store from the Zarr array that is returned in each chunk.

https://github.com/earth-mover/icechunk/blob/dcdfa82e23427facebc921919c27dd615f11db72/icechunk-python/python/icechunk/dask.py#L31

I used this approach because I didn't want to copy over all the optimization logic to icechunk.



Code after:
    assert isinstance(v, Delayed)

    # You need a well-formed HighLevelgraph for e.g. dask.graph_manipulation.bind
    for layer in v.__dask_layers__():
        assert layer in v.dask.layers

    assert (at == 0).all() and (bt == 0).all()
    assert all([ev is None for ev in v.compute()])
    assert (at == 2).all() and (bt == 3).all()

    at = np.zeros(shape=(4, 4))
    bt = np.zeros(shape=(4, 4))
    dat, dbt = store([a, b], [at, bt], compute=False, return_stored=True)
    assert isinstance(dat, Array) and isinstance(dbt, Array)
    assert (at == 0).all() and (bt == 0).all()
    assert (dat.compute() == at).all() and (dbt.compute() == bt).all()
    assert (at == 2).all() and (bt == 3).all()

    at = np.zeros(shape=(4, 4))
    bt = np.zeros(shape=(4, 4))
    dat, dbt = store(
        [a, b], [at, bt], compute=False, return_stored=True, load_stored=False
    )
    assert isinstance(dat, Array) and isinstance(dbt, Array)
    assert (at == 0).all() and (bt == 0).all()
    dask.compute(dat, dbt)
    assert (at == 2).all() and (bt == 3).all()


def test_store_nocompute_regions():
    x = da.ones(10, chunks=1)
    y = np.zeros((2, 10))
    d1 = da.store(x, y, regions=(0,), compute=False)
    d2 = da.store(x, y, regions=(1,), compute=False)
    assert d1.key != d2.key


class ThreadSafetyError(Exception):
    pass


class NonthreadSafeStore:
    def __init__(self):
        self.in_use = False

    def __setitem__(self, key, value):
        if self.in_use:
            raise ThreadSafetyError()
        self.in_use = True
        time.sleep(0.001)
        self.in_use = False