Poor performance when slicing with steps #843

"""Demonstrate low performance of slicing zarr array with step size."""

import sys
import json
import time
from collections.abc import MutableMapping

import numpy
import zarr
import dask
import dask.array

shape = (149568, 66369, 3)  # 27.7G RGB image
chunks = (1, 66369, 3)  # one scanline
dtype = '|u1'
step = 512
repeat = 2
counter = 0  # global counter, not thread-safe

class ZarrStore(MutableMapping):
    def __init__(self, shape, chunks, dtype):
        self.chunk = numpy.zeros(chunks, dtype)
        self.store = {}
        self.store['.zattrs'] = json_dump({})
        self.store['.zarray'] = json_dump(
            {
                'zarr_format': 2,
                'shape': shape,
                'chunks': chunks,
                'dtype': dtype,
                'fill_value': 0,
                'order': 'C',
                'compressor': None,
                'filters': None,
            }
        )

    def flush(self):
        raise PermissionError

    def clear(self):
        raise PermissionError

    def keys(self):
        return self.store.keys()

    def items(self):
        return self.store.items()

    def values(self):
        return self.store.values()

    def __iter__(self):
        return iter(self.store.keys())

    def __len__(self):
        return len(self.store)

    def __delitem__(self, key):
        raise PermissionError

    def __setitem__(self, key, value):
        raise PermissionError

    def __contains__(self, key):
        return key in self.store

    def __getitem__(self, key):
        global counter
        if key in self.store:
            return self.store[key]
        counter += 1
        return self.chunk


def json_dump(obj):
    return json.dumps(
        obj,
        indent=1,
        sort_keys=True,
        ensure_ascii=True,
        separators=(',', ': '),
    ).encode('ascii')


if 1:
    store = ZarrStore(shape, chunks, dtype)
    z = zarr.open(store, mode='r')
else:
    # use zarr MemoryStore
    z = zarr.zeros(
        shape=shape,
        chunks=chunks,
        dtype=dtype,
        compressor=None,
    )

print('Python', sys.version)
print('Platform', sys.platform)
print('numpy', numpy.__version__)
print('zarr', zarr.__version__)
print('dask', dask.__version__)
print()
print(z.info)

print('z[:]  # slice all chunks')
for _ in range(repeat):
    counter = 0
    start = time.perf_counter()
    arr = z[:]
    print(f'read {counter} chunks in {time.perf_counter() - start:.3} s')
    del arr

print()
print('z[:293]  # slice first 293 chunks')
for _ in range(repeat):
    counter = 0
    start = time.perf_counter()
    arr = z[: shape[0] // step + 1]
    print(f'read {counter} chunks in {time.perf_counter() - start:.3} s')
    del arr

print()
print('z[::512]  # slice every 512th chunk, 293 total')
for _ in range(repeat):
    counter = 0
    start = time.perf_counter()
    arr = z[::step]
    print(f'read {counter} chunks in {time.perf_counter() - start:.3} s')
    del arr

print()
print('z[::149568]  # slice every 149568th chunk, 1 total')
for _ in range(repeat):
    counter = 0
    start = time.perf_counter()
    arr = z[:: shape[0]]
    print(f'read {counter} chunks in {time.perf_counter() - start:.3} s')
    del arr

print()
print('z.oindex[[0, 512 .. 149504]]  # index using list')
for _ in range(repeat):
    counter = 0
    start = time.perf_counter()
    arr = z.oindex[list(range(0, shape[0], step))]
    print(f'read {counter} chunks in {time.perf_counter() - start:.3} s')
    del arr

print()
print('dask.array.from_zarr(store)[::512].compute()  # slice using dask')
da = dask.array.from_zarr(store)
for _ in range(repeat):
    counter = 0
    start = time.perf_counter()
    arr = da[::step].compute()
    print(f'read {counter} chunks in {time.perf_counter() - start:.3} s')
    del arr

print()
print('dask.array.from_zarr(store).compute()  # slice all chunks using dask')
da = dask.array.from_zarr(store)
for _ in range(repeat):
    counter = 0
    start = time.perf_counter()
    arr = da.compute()
    print(f'read {counter} chunks in {time.perf_counter() - start:.3} s')
    del arr

print()
print('# baseline: manually copy all chunks from store to numpy array')
da = dask.array.from_zarr(store)
for _ in range(repeat):
    counter = 0
    start = time.perf_counter()
    arr = numpy.empty(shape, dtype)
    for i in range(arr.shape[0]):
        arr[i] = store[i]
    print(f'read {counter} chunks in {time.perf_counter() - start:.3} s')
    del arr