Custom Indexing API for Pipeline-level Parallelism #3175

Activity

1. dstansby

   added

   enhancementNew features or improvementsNew features or improvements

   on Jul 6, 2025
2. ilan-gold commented on Jul 11, 2025

   

   ilan-gold

   on Jul 11, 2025

   Last edited by ilan-gold

   ContributorAuthor

   More actions

   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.

   👀React with 👀1Reacted by Davis Bennett
3. d-v-b commented on Jul 11, 2025

   

   d-v-b

   on Jul 11, 2025

   Contributor

   More actions

   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
4. ilan-gold commented on Jul 28, 2025

   

   ilan-gold

   on Jul 28, 2025

   ContributorAuthor

   More actions

   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