Custom Indexing API for Pipeline-level Parallelism

[ilan-gold]

Inspired by zarrs, where you can do a sort of jump-around-slicing, I broadly wish to be able to do the following disjoint slice indexing (although I am suggesting, motivated by this request, a more general API).

The current "public" solution would be to either asyncio.gather a bunch of disparate slice requests or to use numpy indexing if you are not using the async API.

I get better performance by batching together all of the slice requests into one indexer use the pipeline-level parallelism as opposed to top-level parallelism.

import zarr
import numpy as np

import asyncio
import platform
import subprocess

def clear_cache():
    if platform.system() == "Darwin":
        subprocess.call(['sync', '&&', 'sudo', 'purge'])
    elif platform.system() == "Linux":
        subprocess.call(['sudo', 'sh', '-c', "sync; echo 3 > /proc/sys/vm/drop_caches"])
    else:
        raise Exception("Unsupported platform")


class MultiBasicIndexer(zarr.core.indexing.Indexer):
    """An indexer that batches together multiple disparate slicing 1D slicing indexers together into a single indexer"""
    def __init__(self, indexers: list[zarr.core.indexing.Indexer]):
        self.shape = tuple(
            sum(i.shape[k] for i in indexers) for k in range(len(indexers[0].shape))
        )
        self.drop_axes = indexers[0].drop_axes  # maybe?
        self.indexers = indexers

    def __iter__(self):
        total = 0
        for i in self.indexers:
            for c in i:
                gap = c[2][0].stop - c[2][0].start
                yield type(c)(c[0], c[1], (slice(total, total + gap)), c[3])
                total += gap

z = zarr.create_array("foo.zarr", data=np.random.randn(10_000 * 10_000))
chunk_size = z.chunks[0]

chunks = list(range(0, 15, 2)) # disparate chunks i.e., skipping around by a factor of 2

slices = [zarr.core.indexing.BasicIndexer(slice(i * chunk_size, (i+1)*chunk_size), shape=z.metadata.shape, chunk_grid=z.metadata.chunk_grid) for i in range(0, 15, 2)]
np_indexer = np.concatenate([np.arange(i * chunk_size, (i+1)*chunk_size) for i in range(0, 15, 2)])

async def run():
    return np.concatenate(await asyncio.gather(*(z._async_array._get_selection(s, prototype=zarr.core.buffer.default_buffer_prototype()) for s in slices))).reshape(chunk_size * len(chunks))

assert (z[np_indexer] == zarr.core.sync.sync(run())).all()

clear_cache()
%timeit  z[np_indexer]
clear_cache()
%timeit zarr.core.sync.sync(run())
clear_cache()
%timeit zarr.core.sync.sync(z._async_array._get_selection(MultiBasicIndexer(slices), prototype=zarr.core.buffer.default_buffer_prototype()))

I'm uploading a juv runnable notebook as well with cache clearing built in:
indexing_benchmark.ipynb.zip

So I see about a 40% speedup for the batched indexing vs. asyncio.gather separated indexing (both of which are light years ahead of integer indexing). This is an upper bound - usually, the difference was ~20-30%, but zarrs gives another 20-30% which means we're now close to a 2X speedup.

Thus I see a potential route forward here as making the following public

• AsyncArray._get_selection (or something similar, just saying this because it exists)
• Indexer abc class + ChunkProjection for iteration
• An into_indexer function that returns Indexer-based classes for a simple selection (like tuple of slices or so). I think the Indexer return type would probably be enough to do useful things since the iterator for all Indexer subclasses returns ChunkProjection, which can be modified as needed (see my motivating example).

Open to other thoughts!

[ilan-gold]

I've been pondering this a bit more. While this solution is definitely faster, that's more a testament to just how slow integer-based indexing can be. For example, in the above example, the slices are perfectly aligned to chunks. But if you were to request, say, multiple slices from within a given chunk, then with the current MultiBasicIndexer solution, you end up fetching the same chunk multiple times because ChunkProjection goes in one at a time to the codec pipeline tied to both its chunk and the selection, which must be expressable via normal python indexing (i.e., not disparate, maybe uneven, slices).

When I account for this locally, making chunk_size smaller (i.e., the size of the slices fetched), the performance gap closes but the indexing via slices is still faster (albeit by less) despite the extra requests for the same on-disk chunk.

I added a class (cc @d-v-b for examples maybe! I should re-write using WrapperStore you mentioned on the call) that tracks requests to a store to check this

import zarr
from collections import Counter, defaultdict
from collections.abc import Iterable


class AccessTrackingStoreBase(zarr.storage.LocalStore):
    _access_count: Counter[str]
    _accessed: defaultdict[str, set]
    _accessed_keys: defaultdict[str, list[str]]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._access_count = Counter()
        self._accessed = defaultdict(set)
        self._accessed_keys = defaultdict(list)

        self._read_only = True

    def _check_and_track_key(self, key: str):
        for tracked in self._access_count:
            if tracked in key:
                self._access_count[tracked] += 1
                self._accessed[tracked].add(key)
                self._accessed_keys[tracked] += [key]

    def get_access_count(self, key: str) -> int:
        # access defaultdict when value is not there causes key to be there,
        # which causes it to be tracked
        if key not in self._access_count:
            msg = f"{key} not found among access count"
            raise KeyError(msg)
        return self._access_count[key]

    def get_subkeys_accessed(self, key: str) -> set[str]:
        if key not in self._accessed:
            msg = f"{key} not found among accessed"
            raise KeyError(msg)
        return self._accessed[key]

    def get_accessed_keys(self, key: str) -> list[str]:
        if key not in self._accessed_keys:
            msg = f"{key} not found among accessed keys"
            raise KeyError(msg)
        return self._accessed_keys[key]

    def initialize_key_trackers(self, keys_to_track: Iterable[str]) -> None:
        for k in keys_to_track:
            self._access_count[k] = 0
            self._accessed_keys[k] = []
            self._accessed[k] = set()

    def reset_key_trackers(self) -> None:
        self.initialize_key_trackers(self._access_count.keys())

    def assert_access_count(self, key: str, count: int):
        keys_accessed = self.get_subkeys_accessed(key)
        access_count = self.get_access_count(key)
        assert self.get_access_count(key) == count, (
            f"Found {access_count} accesses at {keys_accessed}"
        )

    async def get(
        self,
        key: str,
        prototype: zarr.core.buffer.BufferPrototype | None = None,
        byte_range: zarr.abc.store.ByteRequest | None = None,
    ) -> object:
        self._check_and_track_key(key)
        return await super().get(key, prototype=prototype, byte_range=byte_range)

and can see that the sliced-based indexing makes multiple requests for the same chunk (I didn't have any reason to believe this optimization existed in zarr-python but wanted to check).

This also brings up an issue that @d-v-b has mentioned before, of tying stores to arrays. At the moment, I don't think we can cache uncompressed chunks before indexing. We could cache the chunks themselves, compressed, via a cache-store though (however, this would still require repeated decompression).

That being said, this sort of thing does exist in zarrs via partial_decoder or with user-handled chunk caching, see: https://docs.rs/zarrs/latest/zarrs/array/struct.Array.html?search=cache#optimising-reads

I would venture a guess that this is a non-trivial portion of the speed improvement we see with zarrs (although likely not all).

So I think something concrete would be for zarr-python to somehow cache decompressed chunks within a single read_batch call from the codec pipeline.

At the moment, I am not sure how this sort of thing would actually change performance without a more general API like the one proposed above, but I think this is a nice diagnostic.

[d-v-b]

i don't see why the codec pipeline couldn't be responsible caching this stuff -- I'd say caching falls under the remit of "find the most efficient way to get me this chunk N times", which is arguably the job of the codec pipeline

[ilan-gold]

So I've been doing some further thinking about this issue. I think that whatever the solution ends up being, we would want some sort of pipeline-level cache similar to https://docs.rs/zarrs/latest/zarrs/array/struct.Array.html?search=cache#optimising-reads.

Creating a top-level multi-slice iterator-like API

One option discussed in the meeting on Friday was to create a top-level API that takes in slice-like selections (although, really, could be any sort of index selection) and the outputs one-array-per-selection. Pretty simple! This would be instead of making public the above requested APIs (i.e., Indexer etc.).

Downsides

• Caching Our current pipeline is tied to having an one out-buffer per-read call. This paradigm works for fine for the exemplar 1D case where the output buffer can be contiguous and written to with expected results, but breaks down in more dimensions (i.e., how do you write (2,2) selection followed by a (3,3) selection?). But implementing a cache that lives as long as a given BatchedCodecPipeline.read means that the cache is tied to having a single output buffer, which means we can't cache between different selections unless we extend the BatchedCodecPipeline.read and BatchedCodecPipeline.read_batch APIs to accept a cache argument (and then implement some sort of known API around the cache, which we wouldn't have to do if it were internal).
• We would be tied to whatever concrete argument types the new API accepts, which would limit customizability

Upsides

• Simplicity This would be super simple and concrete, although I'm not sure it's an improvement on what exists already i.e., asyncio.gather unless we implement internal caching (see above)

Using Awkward.Array

A more flexible option could be something like:

1. Do the requested public-API-ing of Indexer, _get_selection etc.
2. Make https://awkward-array.org a buffer option (potentially outside zarr-python)
3. Create a simple slices-only interface that outputs into an awkward array container

So my thought goes something like this:

import awkward as ak, numpy as np

# Encapsulated in some buffer class
builder = ak.ArrayBuilder()

np_arrays = [np.arange(16).reshape((4, 4)), np.arange(36).reshape((6, 6))] # i.e., two separate zarr collections
for arr in np_arrays: # i.e., within `read`
    builder.append(arr) # encapsulated within a `__setitem__` call

builder.snapshot() # <Array [[[0, 1, 2, 3], ..., [12, ..., 15]], ...] type='2 * var * var * int64'>

Downsides

• shape currently is required by both NDBuffer and Indexer, but of course does not really apply as a concept to Awkward.Array at least given the current typing (i.e., the current typing does not allow for ragged arrays).
• This sort of thing would never end up in zarr-python without accepting a new dependency

Upsides

• Using Awkward.Array would likely allow allow us to save on memory compared to the current pure-iteration approach because the output buffer would be an Awkward.Array that is built up instead of allocating a new buffer for every selection that is identical to the output of the selection itself in read i.e., the decompressed selection would fill the entire output buffer, which means the output buffer is redundant. This in theory could be solved by Codec pipeline performance #2904 (comment)
• The returned object might be able to be given semantic meaning or similar via the rich functionality offered by awkward.Array