fix(indexing): make sparse selections O(npoints) instead of O(nchunks) - #262
Open
d-v-b wants to merge 2 commits into
Open
fix(indexing): make sparse selections O(npoints) instead of O(nchunks)#262d-v-b wants to merge 2 commits into
d-v-b wants to merge 2 commits into
Conversation
CoordinateIndexer (general path and sorted-1D fast path) and IntArrayDimIndexer built a dense per-chunk histogram via np.bincount(..., minlength=nchunks) / np.zeros(nchunks) plus a full-length cumsum. For arrays with very many chunks this allocates memory proportional to the total chunk count regardless of how few points are selected — a 3-point write to an array with 1.4e11 chunks tried to allocate 1 TiB and raised MemoryError. Replace the dense histogram with run boundaries computed directly on the sorted raveled chunk ids (sorted_run_ends), storing a compressed cumsum aligned with the occupied chunks, and index it positionally in __iter__. Memory and time now scale with the number of selected points. Fixes zarr-developers#4174 Assisted-by: ClaudeCode:claude-fable-5
…rties The compressed per-occupied-chunk cumsum introduced for zarr-developersgh-4174 changed the observable semantics of chunk_nitems_cumsum (and removed chunk_nitems on IntArrayDimIndexer). Although zarr.core is documented as private API, external code is known to introspect these indexers, so be conservative: store the compressed offsets under a new name (chunk_run_ends, aligned with chunk_rixs / dim_chunk_ixs) and restore chunk_nitems / chunk_nitems_cumsum as properties that lazily rebuild the original dense arrays, warning with ZarrDeprecationWarning. The O(nchunks) cost is now only paid if someone actually accesses them — which was the status quo before the fix. Assisted-by: ClaudeCode:claude-fable-5
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
🤖 AI text below 🤖
Fixes zarr-developers#4174
Problem
CoordinateIndexer(both the general path and the sorted-1Dsearchsortedfast path) andIntArrayDimIndexerbuilt a dense per-chunk histogram withnp.bincount(..., minlength=nchunks)/np.zeros(nchunks), followed by a full-lengthnp.cumsum. Memory and time were proportional to the total number of chunks in the array, regardless of how few points were selected. The reporter's 3-pointset_coordinate_selectionon a(10_000_000, 13_854)array with(1, 1)chunks (~1.4e11 chunks) tried to allocate 1 TiB and raisedMemoryError.get_coordinate_selection,vindex, mask selections (viaMaskIndexer), and orthogonal integer-array selections (per-dimension) were all affected.Fix
The dense histogram was only ever consumed by
__iter__to find each occupied chunk's start/stop range in the chunk-sorted point list — and the points are already sorted (or verified pre-sorted) by raveled chunk id. The newsorted_run_endshelper computes the occupied chunk ids and a compressed cumsum (one entry per occupied chunk) directly from run boundaries in the sorted ids;__iter__indexes it by position instead of by raveled chunk id. Cost is now O(npoints).Measured with the scaled-down repro from the investigation (3-point selection, 1e8 chunks): peak allocation drops from 1.6 GB to 0.1 MB for both get and set, on both the general and fast paths.
Tests
(2**22, 2**22)array with(1, 1)chunks (2**44 chunks — any dense per-chunk allocation fails outright).CoordinateIndexer(general sorted, general unsorted, and dense-duplicates fast path) andIntArrayDimIndexer(increasing/decreasing/unordered) on a grid with 2**62 chunks.Out of scope, noted during investigation
normalize_chunks_1dmaterializesnp.full(n_chunks_along_dim, size)at array-creation time, so creating an array with an astronomically chunked single dimension is still slow/impossible (opening one is fine —ChunkGrid.from_metadatabuilds O(1)FixedDimensions). That is a separate creation-time issue and is not addressed here.API compatibility
Although
zarr.coreis documented private, the indexer attributes are introspectable, so the change is conservative: the compressed offsets live under a new name (chunk_run_ends, aligned withchunk_rixs/dim_chunk_ixs), and the pre-existing dense attributes —CoordinateIndexer.chunk_nitems_cumsum,IntArrayDimIndexer.chunk_nitems,IntArrayDimIndexer.chunk_nitems_cumsum— remain as deprecated properties that lazily rebuild the original dense arrays (withZarrDeprecationWarning). External code that read them sees identical values; the O(nchunks) allocation is only paid on access. A GitHub-wide code search found no importers of these attributes (only vendored copies and ports), so removal after a deprecation cycle should be safe.🤖 Generated with Claude Code