Fast partial-read path for sharding_indexed codec

CHANGELOG.md

@@ -2,6 +2,13 @@
 
 ## Unreleased
 
+- Fast partial-read path for the `sharding_indexed` codec [#264](https://github.com/JuliaIO/Zarr.jl/pull/264)
+  - in-memory partial decode in `Codecs.V3Codecs.read_shard_partial!` and `read_shard_partial_with_source!` — only inner chunks intersecting the requested slice are decompressed; the rest are skipped
+  - storage-aware partial reads via three new optional `AbstractStore` methods (`supports_partial_reads`, `read_range`, `getsize`) — stores opt in to byte-range reads; safe defaults preserve correctness for backends that don't
+  - `DirectoryStore` opts in (using `seek` + `readbytes!` and `filesize`); other backends inherit the defaults
+  - new `Zarr.enable_partial_shard_storage_reads[]` `Ref{Bool}` flag (default `true`); flip to `false` to fall back to the in-memory partial-decode path for A/B comparisons
+  - applies only when the codec pipeline is "pure" sharding (no array→array codecs before, no bytes→bytes codecs after); compound pipelines run on the existing path unchanged
+
 ## v0.10.0 - 2026-04-24
 
 - Enable `sharding_indexed` codec for Zarr v3 [#241](https://github.com/JuliaIO/Zarr.jl/pull/241)

docs/src/UserGuide/partial_shard_reads.md

@@ -0,0 +1,78 @@
+# Partial reads of sharded arrays
+
+Zarr v3's `sharding_indexed` codec packs many small inner chunks into one
+larger outer chunk file, with an index that records each inner chunk's
+byte offset and length inside the file. The index lets a reader fetch
+just the inner chunks intersecting a user's slice instead of decoding
+the whole outer chunk.
+
+`Zarr.jl` exploits this in two layers:
+
+1. **In-memory partial decode** — when an outer chunk has been read into
+   memory (the existing path), only decode the inner chunks the request
+   actually touches. Skips decompression for unrelated inner chunks.
+2. **Storage-aware partial read** — when the storage backend supports
+   byte-range reads, fetch only the shard index plus the bytes of the
+   intersecting inner chunks. Skips both the I/O and the decode for
+   everything else.
+
+Both paths are transparent to user code — `arr[a:b, c]` is the same call
+whether the array is sharded or not. The fast paths kick in
+automatically when the codec pipeline is "pure" sharding (no
+array→array codecs before the sharding codec, no bytes→bytes codecs
+wrapping it).
+
+## When the fast path applies
+
+The fast path is taken when:
+
+- The codec pipeline is exactly one `ShardingCodec` (no surrounding
+  codecs). This is the common shape that `zarr-python` produces and
+  that practical sharded archives use.
+- The requested slice is *partial* — i.e. it covers fewer elements than
+  the outer chunk shape. For full-chunk reads the existing path is
+  already optimal (the whole shard's bytes need to come off disk and
+  be decoded anyway).
+
+When the pipeline has additional codecs (e.g. transpose, blosc-around-
+sharding) or when the request happens to align exactly with an outer
+chunk, the existing decode path runs unchanged.
+
+## Storage backends
+
+Stores opt into the storage-aware path by implementing three optional
+methods. The defaults provided by `AbstractStore` keep correctness for
+backends that don't implement them — any such store falls back to the
+in-memory partial-decode path automatically.
+
+```@docs
+Zarr.supports_partial_reads
+Zarr.read_range
+Zarr.getsize
+```
+
+`DirectoryStore` opts in (using `seek` + `readbytes!` for `read_range`
+and `filesize` for `getsize`). Other built-in backends inherit the
+defaults; adding native byte-range support to e.g. `S3Store` or
+`HTTPStore` is a one-method change for each, since both wire protocols
+support byte ranges natively.
+
+## Toggle
+
+```@docs
+Zarr.enable_partial_shard_storage_reads
+```
+
+Set the flag to `false` to fall back to the in-memory partial-decode
+path even on stores that support byte-range reads. Useful for A/B
+performance comparisons or to debug a suspected partial-read bug.
+
+## Reference
+
+The storage-aware path lives in `Zarr._readblock_sharded_partial!`
+(in `src/ZArray.jl`). The shared decode loop used by both paths is
+`Zarr.Codecs.V3Codecs.read_shard_partial_with_source!` (in
+`src/Codecs/V3/V3.jl`); the in-memory wrapper is
+`Zarr.Codecs.V3Codecs.read_shard_partial!`. The pipeline detection
+helper that distinguishes "pure" sharding from compound pipelines is
+`Zarr.Codecs.V3Codecs.sharding_codec`.

src/Codecs/V3/V3.jl

@@ -592,6 +592,173 @@ function zdecode!(data::AbstractArray, encoded::Vector{UInt8}, c::ShardingCodec{
     return data
 end
 
+"""
+    sharding_codec(p) -> Union{Nothing, ShardingCodec}
+
+If `p` is a "pure" sharding pipeline — no array→array codecs before
+sharding and no bytes→bytes codecs after — return its inner
+`ShardingCodec`. Otherwise return `nothing`.
+
+Used by `ZArray.readblock!` to dispatch partial reads to a fast path
+that decodes only the inner chunks intersecting the requested slice
+instead of the whole shard. Pre/post codecs would have to be
+reapplied per inner chunk, which is the general path and not yet
+optimized.
+"""
+function sharding_codec(p::V3Pipeline)
+    isempty(p.array_array) || return nothing
+    isempty(p.bytes_bytes) || return nothing
+    return p.array_bytes isa ShardingCodec ? p.array_bytes : nothing
+end
+sharding_codec(_) = nothing
+
+
+"""
+    _shard_index_byte_range(sc, total_size) -> UnitRange{Int}
+
+Byte range, 1-based and inclusive, where the shard index lives inside a
+shard file of size `total_size`.
+"""
+function _shard_index_byte_range(sc::ShardingCodec{N}, chunks_per_shard::NTuple{N,Int}, total_size::Int) where N
+    index_size = compute_encoded_index_size(chunks_per_shard, sc)
+    if sc.index_location == :start
+        return 1:index_size
+    else
+        return (total_size - index_size + 1):total_size
+    end
+end
+
+
+"""
+    read_shard_partial_with_source!(aout, output_base_offsets, sc, index_bytes,
+                                    fetch_inner_chunk_bytes,
+                                    chunk_shape, current_chunk_offsets,
+                                    indranges, fill_value)
+
+Decode just the inner chunks of one outer (sharded) chunk that intersect
+`indranges`, writing each intersection straight into the right slice of
+`aout`. The caller supplies the already-read shard `index_bytes` and a
+closure `fetch_inner_chunk_bytes(offset_from_shard_start, nbytes) ->
+Vector{UInt8}` that returns the raw bytes for an inner chunk. This lets
+the same loop drive both the in-memory case (slice a `chunk_compressed`
+blob) and the byte-range storage case (`seek` + `read` the file).
+
+Coordinate conventions:
+
+- `indranges`             — global 1-based ranges, intersection of the
+                            user request with this outer chunk
+- `current_chunk_offsets` — 0-based global offsets of this outer chunk's
+                            origin
+- `output_base_offsets`   — 0-based global offsets of the user's request
+                            (so `aout` indices = global −
+                            `output_base_offsets`)
+"""
+function read_shard_partial_with_source!(
+    aout::AbstractArray{T,N},
+    output_base_offsets::NTuple{N,Int},
+    sc::ShardingCodec{N},
+    index_bytes::Vector{UInt8},
+    fetch_inner_chunk_bytes::F,
+    chunk_shape::NTuple{N,Int},
+    current_chunk_offsets::NTuple{N,Int},
+    indranges::NTuple{N,UnitRange{Int}},
+    fill_value,
+) where {T, N, F}
+    fv = fill_value === nothing ? zero(T) : T(fill_value)
+
+    local_indranges = ntuple(N) do i
+        lo = first(indranges[i]) - current_chunk_offsets[i]
+        hi = last(indranges[i])  - current_chunk_offsets[i]
+        lo:hi
+    end
+
+    chunks_per_shard = calculate_chunks_per_shard(chunk_shape, sc.chunk_shape)
+    index = decode_shard_index(index_bytes, chunks_per_shard, sc)
+
+    inner_buf = Array{T}(undef, sc.chunk_shape)
+
+    for cart_idx in CartesianIndices(chunks_per_shard)
+        ic_coords      = Tuple(cart_idx)
+        ic_array_slice = get_chunk_slice_in_shard(ic_coords, sc.chunk_shape, chunk_shape)
+        intersection = ntuple(N) do i
+            lo = max(first(ic_array_slice[i]), first(local_indranges[i]))
+            hi = min(last(ic_array_slice[i]),  last(local_indranges[i]))
+            lo:hi
+        end
+        any(isempty, intersection) && continue
+
+        ic_local = ntuple(N) do i
+            lo = first(intersection[i]) - first(ic_array_slice[i]) + 1
+            hi = last(intersection[i])  - first(ic_array_slice[i]) + 1
+            lo:hi
+        end
+        aout_dest = ntuple(N) do i
+            lo = first(intersection[i]) + current_chunk_offsets[i] - output_base_offsets[i]
+            hi = last(intersection[i])  + current_chunk_offsets[i] - output_base_offsets[i]
+            lo:hi
+        end
+
+        chunk_slice = get_chunk_slice(index, ic_coords)
+        if chunk_slice === nothing
+            view(aout, aout_dest...) .= fv
+            continue
+        end
+
+        offset_start, offset_end = chunk_slice
+        encoded_chunk = fetch_inner_chunk_bytes(Int(offset_start), Int(offset_end - offset_start))
+        pipeline_decode!(sc.codecs, inner_buf, encoded_chunk)
+
+        copyto!(view(aout, aout_dest...), view(inner_buf, ic_local...))
+    end
+    return aout
+end
+
+
+"""
+    read_shard_partial!(aout, output_base_offsets, sc, chunk_compressed,
+                        chunk_shape, current_chunk_offsets, indranges,
+                        fill_value)
+
+In-memory variant: caller has already loaded the entire shard into
+`chunk_compressed`. Pulls the index out of that blob and forwards to
+`read_shard_partial_with_source!` with a closure that slices the blob
+for each inner chunk's bytes. `nothing`/empty `chunk_compressed`
+fills the request region with `fill_value` and returns.
+"""
+function read_shard_partial!(
+    aout::AbstractArray{T,N},
+    output_base_offsets::NTuple{N,Int},
+    sc::ShardingCodec{N},
+    chunk_compressed::Union{Vector{UInt8}, Nothing},
+    chunk_shape::NTuple{N,Int},
+    current_chunk_offsets::NTuple{N,Int},
+    indranges::NTuple{N,UnitRange{Int}},
+    fill_value,
+) where {T, N}
+    if chunk_compressed === nothing || isempty(chunk_compressed)
+        fv = fill_value === nothing ? zero(T) : T(fill_value)
+        aout_dest_outer = ntuple(N) do i
+            lo = first(indranges[i]) - output_base_offsets[i]
+            hi = last(indranges[i])  - output_base_offsets[i]
+            lo:hi
+        end
+        view(aout, aout_dest_outer...) .= fv
+        return aout
+    end
+
+    chunks_per_shard = calculate_chunks_per_shard(chunk_shape, sc.chunk_shape)
+    idx_range = _shard_index_byte_range(sc, chunks_per_shard, length(chunk_compressed))
+    index_bytes = chunk_compressed[idx_range]
+
+    fetch = (off::Int, nb::Int) -> chunk_compressed[(off + 1):(off + nb)]
+
+    return read_shard_partial_with_source!(
+        aout, output_base_offsets, sc, index_bytes, fetch,
+        chunk_shape, current_chunk_offsets, indranges, fill_value,
+    )
+end
+
+
 struct TransposeCodec{N} <: V3Codec{:array, :array}
     order::NTuple{N, Int}  # permutation (1-based Julia indexing)
 end

src/Storage/Storage.jl

@@ -214,6 +214,47 @@ end
 
 
 
+"""
+    supports_partial_reads(::AbstractStore) -> Bool
+
+Whether this store implements efficient byte-range reads via
+[`read_range`](@ref) and [`getsize`](@ref). Defaults to `false`; stores
+override to opt in (e.g. `DirectoryStore`).
+
+The sharding partial-read fast path uses this to skip loading whole
+shard files when only a few inner chunks are needed.
+"""
+supports_partial_reads(::AbstractStore) = false
+
+"""
+    read_range(s::AbstractStore, key::AbstractString, byte_range::UnitRange{Int})
+        -> Union{Vector{UInt8}, Nothing}
+
+Read just `byte_range` (1-based, inclusive) from the value stored under
+`key`. Returns `nothing` if the key doesn't exist. Default implementation
+falls back to `s[key][byte_range]`; stores that opt into
+[`supports_partial_reads`](@ref) should override with a real partial read.
+"""
+function read_range(s::AbstractStore, key::AbstractString, byte_range::UnitRange{Int})
+    bytes = s[key]
+    bytes === nothing && return nothing
+    return bytes[byte_range]
+end
+
+"""
+    getsize(s::AbstractStore, key::AbstractString) -> Int
+
+Size in bytes of the value stored under `key`, or 0 if it doesn't exist.
+Default falls back to `length(s[key])`; partial-read-capable stores
+should override with a cheap size lookup (e.g. `filesize`).
+"""
+function getsize(s::AbstractStore, key::AbstractString)
+    bytes = s[key]
+    bytes === nothing && return 0
+    return length(bytes)
+end
+
+
 ## Handling sequential vs parallel IO
 struct SequentialRead end
 struct ConcurrentRead

src/Storage/directorystore.jl

@@ -57,3 +57,25 @@ function subkeys(s::DirectoryStore,p)
 end
 Base.delete!(s::DirectoryStore, k::String) = isfile(joinpath(s.folder, k)) && rm(joinpath(s.folder, k))
 
+# Partial-read support. seek+read into a fresh buffer is much cheaper
+# than reading the whole file when the caller only wants a few KB.
+supports_partial_reads(::DirectoryStore) = true
+
+function read_range(d::DirectoryStore, i::AbstractString, byte_range::UnitRange{Int})
+    fname = joinpath(d.folder, i)
+    isfile(fname) || return nothing
+    n = length(byte_range)
+    n == 0 && return UInt8[]
+    buf = Vector{UInt8}(undef, n)
+    open(fname, "r") do io
+        seek(io, first(byte_range) - 1)
+        readbytes!(io, buf, n)
+    end
+    return buf
+end
+
+function getsize(d::DirectoryStore, i::AbstractString)
+    fname = joinpath(d.folder, i)
+    isfile(fname) ? filesize(fname) : 0
+end
+