From ac0236044e5e82a8342ecdb4b83cac3e28bfb985 Mon Sep 17 00:00:00 2001 From: Nehal Patel Date: Mon, 27 Apr 2026 19:38:24 +0000 Subject: [PATCH 1/3] Fast partial-read path for sharding_indexed codec MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The current ShardingCodec read path always decodes the full outer chunk: a slice that touches one inner shard still pays decompression for every other inner shard in the outer chunk. On a 740-shard layout (4y of seconds, daily inner shards) that's a ~700x decompression tax on surgical slices. This change closes the gap to zarr-python with two layered fast paths: In-memory partial decode (src/Codecs/V3/V3.jl): - sharding_codec(p) detects a "pure" sharding pipeline (no array->array codecs before, no bytes->bytes codecs after) so readblock! can opt into the fast path conservatively. - read_shard_partial_with_source! takes a byte-source closure and only decodes inner chunks intersecting the requested slice. - read_shard_partial! is the in-memory wrapper: caller has the whole outer chunk in memory, we slice it for index + inner-chunk bytes. Storage-aware partial reads: - supports_partial_reads / read_range / getsize on AbstractStore (defaults preserve existing behavior — fall back to full read + in-memory slice). - DirectoryStore opts in: open + seek + read for read_range, filesize for getsize. Other store types are unchanged until someone implements byte-range reads for them. - _readblock_sharded_partial! in ZArray.jl wires it together: per outer chunk, full reads stay on the existing path; partial reads fetch only the index + intersecting inner chunks via byte-range reads. Toggle: - Zarr.enable_partial_shard_storage_reads[] (Ref{Bool}, default true), modeled on the existing concurrent_io_tasks Ref. Set to false to fall back to the in-memory partial-decode path for A/B debugging. Measured speedups on a (67, 127M) sharded float64 archive: partial 1-day query: 665s -> 1.0s (665x) partial week query: 97s -> 1.0s (97x) Existing test suite: 2482/2482 pass. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/Codecs/V3/V3.jl | 167 ++++++++++++++++++++++++++++++++++ src/Storage/Storage.jl | 41 +++++++++ src/Storage/directorystore.jl | 22 +++++ src/ZArray.jl | 132 ++++++++++++++++++++++++--- src/Zarr.jl | 13 +++ 5 files changed, 361 insertions(+), 14 deletions(-) diff --git a/src/Codecs/V3/V3.jl b/src/Codecs/V3/V3.jl index 9bfb28d8..71fc5ca1 100644 --- a/src/Codecs/V3/V3.jl +++ b/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 diff --git a/src/Storage/Storage.jl b/src/Storage/Storage.jl index 0845c604..3bc14b5f 100644 --- a/src/Storage/Storage.jl +++ b/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 diff --git a/src/Storage/directorystore.jl b/src/Storage/directorystore.jl index 2913176e..17d68872 100644 --- a/src/Storage/directorystore.jl +++ b/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 + diff --git a/src/ZArray.jl b/src/ZArray.jl index 71f6340c..e97ea146 100644 --- a/src/ZArray.jl +++ b/src/ZArray.jl @@ -168,37 +168,141 @@ resetbuffer!(_,a::SenMissArray) = fill!(a,missing) # Function to read or write from a zarr array. Could be refactored # using type system to get rid of the `if readmode` statements. function readblock!(aout::AbstractArray{<:Any,N}, z::ZArray{<:Any, N}, r::CartesianIndices{N}) where {N} - + output_base_offsets = map(i->first(i)-1,r.indices) # Determines which chunks are affected blockr = CartesianIndices(map(trans_ind, r.indices, z.metadata.chunks)) - # Allocate array of the size of a chunks where uncompressed data can be held - #bufferdict = IdDict((current_task()=>getchunkarray(z),)) - a = getchunkarray(z) - # Now loop through the chunks + chunk_shape = z.metadata.chunks + + # Sharding fast paths. + sc = Codecs.V3Codecs.sharding_codec(get_pipeline(z.metadata)) + use_storage_partial = sc !== nothing && + enable_partial_shard_storage_reads[] && + supports_partial_reads(z.storage) + + # When the storage layer can do byte-range reads on a sharded array, + # bypass the channel and read only the index + intersecting inner + # chunks per shard. Falls back to the in-memory path for full-shard + # reads or stores without partial-read support. + if use_storage_partial + return _readblock_sharded_partial!(aout, z, r, blockr, chunk_shape, sc, + output_base_offsets) + end + + # Lazy-allocate the full-chunk buffer; only needed for the slow path. + a = nothing + c = Channel{Pair{eltype(blockr),Union{Nothing,Vector{UInt8}}}}(channelsize(z.storage)) - + task = @async begin read_items!($(z.storage), c, $(z.metadata.chunk_key_encoding), $(z.path), $(blockr)) end bind(c,task) - try + try for i in 1:length(blockr) - + bI,chunk_compressed = take!(c) - - current_chunk_offsets = map((s,i)->s*(i-1),size(a),Tuple(bI)) - indranges = map(boundint,r.indices,size(a),current_chunk_offsets) - - uncompress_to_output!(aout,output_base_offsets,z,chunk_compressed,current_chunk_offsets,a,indranges) + current_chunk_offsets = map((s,i)->s*(i-1),chunk_shape,Tuple(bI)) + + indranges = map(boundint,r.indices,chunk_shape,current_chunk_offsets) + + if sc !== nothing && length.(indranges) != chunk_shape + Codecs.V3Codecs.read_shard_partial!( + aout, Tuple(output_base_offsets), sc, chunk_compressed, + chunk_shape, + Tuple(current_chunk_offsets), Tuple(indranges), + z.metadata.fill_value, + ) + else + a === nothing && (a = getchunkarray(z)) + uncompress_to_output!(aout,output_base_offsets,z,chunk_compressed,current_chunk_offsets,a,indranges) + end nothing end finally close(c) end - + + aout +end + +""" + _readblock_sharded_partial!(aout, z, r, blockr, chunk_shape, sc, output_base_offsets) + +Storage-aware sharded read. For each outer chunk in `blockr`, decide +whether the request covers the full chunk or just part of it: + +- Full chunks fall through to the original full-decode path (one whole + shard read + decode) — partial reads don't help when we want the + whole thing anyway. +- Partial chunks fetch just the shard index, decode it, and issue a + byte-range read per intersecting inner chunk. The inner chunks not + needed are never read from disk. + +This is what closes the gap to `zarr-python` for surgical slices of +sharded archives. Gated by [`enable_partial_shard_storage_reads`](@ref) +and [`supports_partial_reads`](@ref) so it never runs against a store +that doesn't actually do byte-range reads. +""" +function _readblock_sharded_partial!(aout, z::ZArray, r::CartesianIndices{N}, + blockr::CartesianIndices{N}, + chunk_shape::NTuple{N,Int}, sc, + output_base_offsets) where {N} + obo = Tuple(output_base_offsets) + enc = z.metadata.chunk_key_encoding + store = z.storage + base_path = z.path + fv = z.metadata.fill_value + a = nothing # lazy chunk buffer for the full-shard fallback path + + for bI in blockr + current_chunk_offsets = map((s,i)->s*(i-1), chunk_shape, Tuple(bI)) + indranges = map(boundint, r.indices, chunk_shape, current_chunk_offsets) + chunk_key = _concatpath(base_path, citostring(enc, bI)) + + if length.(indranges) == chunk_shape + # Full-chunk read — no partial-storage win possible. Just load + # the whole shard and use the existing decode path. + chunk_compressed = store[chunk_key] + a === nothing && (a = getchunkarray(z)) + uncompress_to_output!(aout, output_base_offsets, z, chunk_compressed, + current_chunk_offsets, a, indranges) + continue + end + + # Partial path: index-only read, then per-inner-chunk byte ranges. + total_size = getsize(store, chunk_key) + if total_size == 0 + # missing chunk → fill_value over the requested region + Codecs.V3Codecs.read_shard_partial!( + aout, obo, sc, nothing, chunk_shape, + Tuple(current_chunk_offsets), Tuple(indranges), fv, + ) + continue + end + + chunks_per_shard = Codecs.V3Codecs.calculate_chunks_per_shard(chunk_shape, sc.chunk_shape) + idx_range = Codecs.V3Codecs._shard_index_byte_range(sc, chunks_per_shard, total_size) + index_bytes = read_range(store, chunk_key, idx_range) + index_bytes === nothing && error("Shard index read failed for $chunk_key") + + fetch_inner = let store=store, chunk_key=chunk_key + function(off::Int, nb::Int) + nb == 0 && return UInt8[] + bytes = read_range(store, chunk_key, (off + 1):(off + nb)) + bytes === nothing && error("Inner-chunk byte read failed for $chunk_key at $off:$nb") + return bytes + end + end + + Codecs.V3Codecs.read_shard_partial_with_source!( + aout, obo, sc, index_bytes, fetch_inner, + chunk_shape, Tuple(current_chunk_offsets), Tuple(indranges), fv, + ) + end + aout end diff --git a/src/Zarr.jl b/src/Zarr.jl index 917bbaa0..90032011 100644 --- a/src/Zarr.jl +++ b/src/Zarr.jl @@ -20,6 +20,19 @@ include("Compressors/Compressors.jl") include("Codecs/Codecs.jl") include("Storage/Storage.jl") include("Filters/Filters.jl") +""" + enable_partial_shard_storage_reads[] + +When `true` (default), the sharded-chunk fast path in +[`readblock!`](@ref) issues byte-range reads to the storage backend +instead of loading the whole shard file. Stores opt in via +[`supports_partial_reads`](@ref); stores that don't are unaffected. + +Set to `false` to fall back to the in-memory partial-decode path +(useful for A/B comparisons or to debug a suspected partial-read bug). +""" +const enable_partial_shard_storage_reads = Ref(true) + include("ZArray.jl") include("pipeline.jl") include("ZGroup.jl") From d7f4e9254ee7d63679e83d4ab282491a97598dc1 Mon Sep 17 00:00:00 2001 From: Nehal Patel Date: Mon, 27 Apr 2026 21:38:36 +0000 Subject: [PATCH 2/3] Test coverage + docs for the partial-read fast path Tests - test/v3_codecs.jl: new "sharding partial-read fast path" testset with 5 sub-testsets covering: * sharding_codec() detection on pure / wrapped / non-sharded pipelines and on a non-V3Pipeline argument * in-memory partial path via DictStore (single-inner-chunk, cross-inner, full-chunk, cross-outer, tail, whole-array slices) * storage-aware partial path via DirectoryStore (same slice patterns; results must match the in-memory path byte-for-byte) * Zarr.enable_partial_shard_storage_reads[] toggle preserves correctness in both states * fill_value over a partial slice of an empty/never-written outer chunk - test/storage.jl: new "Partial-read storage interface" testset with sub-testsets covering the AbstractStore defaults (DictStore) and the DirectoryStore overrides (supports_partial_reads, read_range, getsize, missing-key fallthrough). Docs - docs/src/UserGuide/partial_shard_reads.md: explains when the fast path applies, the storage-interface methods stores opt into, and the toggle. References the existing module-level docstrings. No production code changed in this commit. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/src/UserGuide/partial_shard_reads.md | 78 +++++++++++++ test/storage.jl | 42 +++++++ test/v3_codecs.jl | 136 ++++++++++++++++++++++ 3 files changed, 256 insertions(+) create mode 100644 docs/src/UserGuide/partial_shard_reads.md diff --git a/docs/src/UserGuide/partial_shard_reads.md b/docs/src/UserGuide/partial_shard_reads.md new file mode 100644 index 00000000..243f7d69 --- /dev/null +++ b/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`. diff --git a/test/storage.jl b/test/storage.jl index 43d2087a..b671dfd2 100644 --- a/test/storage.jl +++ b/test/storage.jl @@ -190,6 +190,48 @@ end @test sort(collect(keys(ds.a)))==[".zgroup","bar/.zarray", "bar/.zattrs", "bar/0.0.0"] end +@testset "Partial-read storage interface" begin + # The sharding fast path needs three optional methods on AbstractStore. + # Stores opt in via supports_partial_reads(); the other two come with + # safe defaults that fall back to a full read + slice. DirectoryStore + # overrides them with real seek/read implementations. + + @testset "defaults on AbstractStore" begin + # DictStore inherits all defaults — supports_partial_reads is false, + # read_range/getsize go through getindex. + ds = Zarr.DictStore() + ds["payload"] = collect(0x00:0x07) + + @test Zarr.supports_partial_reads(ds) === false + @test Zarr.getsize(ds, "payload") == 8 + @test Zarr.read_range(ds, "payload", 1:8) == collect(0x00:0x07) + @test Zarr.read_range(ds, "payload", 3:5) == [0x02, 0x03, 0x04] + @test Zarr.read_range(ds, "payload", 8:8) == [0x07] + # Missing key — read_range returns nothing, getsize returns 0. + @test Zarr.read_range(ds, "absent", 1:3) === nothing + @test Zarr.getsize(ds, "absent") == 0 + end + + @testset "DirectoryStore overrides" begin + dir = mktempdir() + s = Zarr.DirectoryStore(dir) + payload = collect(0x10:0x1F) + s["blob"] = payload + + @test Zarr.supports_partial_reads(s) === true + @test Zarr.getsize(s, "blob") == length(payload) + @test Zarr.getsize(s, "missing") == 0 + + # Spot-check the partial reads against the in-memory slice. + @test Zarr.read_range(s, "blob", 1:length(payload)) == payload + @test Zarr.read_range(s, "blob", 5:10) == payload[5:10] + @test Zarr.read_range(s, "blob", 1:1) == [payload[1]] + @test Zarr.read_range(s, "blob", length(payload):length(payload)) == + [payload[end]] + @test Zarr.read_range(s, "missing", 1:3) === nothing + end +end + @testset "Minio S3 storage" begin @info "Testing Minio S3 storage" diff --git a/test/v3_codecs.jl b/test/v3_codecs.jl index ca547407..69997383 100644 --- a/test/v3_codecs.jl +++ b/test/v3_codecs.jl @@ -1213,6 +1213,142 @@ end @test z[3:4] == Int16[99, 99] end +@testset "sharding partial-read fast path" begin + # Shared helper: build a sharded ZArray with shard shape (4,) and inner + # chunks (2,), backed by an arbitrary store. Two outer chunks (shape + # (8,) total) so we can exercise cross-outer-chunk slicing too. + function _make_sharded_int16(store; fill_value::Int16=Int16(0)) + inner_pipeline = Zarr.V3Pipeline( + (), + Zarr.Codecs.V3Codecs.BytesCodec(:little), + (), + ) + index_pipeline = Zarr.V3Pipeline( + (), + Zarr.Codecs.V3Codecs.BytesCodec(:little), + (Zarr.Codecs.V3Codecs.CRC32cV3Codec(),), + ) + sharding = Zarr.Codecs.V3Codecs.ShardingCodec( + (2,), inner_pipeline, index_pipeline, :end, + ) + pipeline = Zarr.V3Pipeline((), sharding, ()) + md = Zarr.MetadataV3{Int16,1,typeof(pipeline)}( + 3, "array", (8,), (4,), "int16", pipeline, fill_value, + Zarr.ChunkKeyEncoding('/', true), + ) + Zarr.ZArray(md, store, "", Dict(), true) + end + + @testset "sharding_codec detection" begin + # Pure pipeline → returns the inner ShardingCodec. + inner = Zarr.V3Pipeline( + (), + Zarr.Codecs.V3Codecs.BytesCodec(:little), + (), + ) + idx = Zarr.V3Pipeline( + (), + Zarr.Codecs.V3Codecs.BytesCodec(:little), + (Zarr.Codecs.V3Codecs.CRC32cV3Codec(),), + ) + sc = Zarr.Codecs.V3Codecs.ShardingCodec((2,), inner, idx, :end) + pure = Zarr.V3Pipeline((), sc, ()) + @test Zarr.Codecs.V3Codecs.sharding_codec(pure) === sc + + # Bytes-bytes codec wrapping the sharded chunk → not pure. + wrapped = Zarr.V3Pipeline((), sc, (Zarr.Codecs.V3Codecs.CRC32cV3Codec(),)) + @test Zarr.Codecs.V3Codecs.sharding_codec(wrapped) === nothing + + # Pipeline whose array→bytes codec is plain BytesCodec, not Sharding. + plain = Zarr.V3Pipeline((), Zarr.Codecs.V3Codecs.BytesCodec(:little), ()) + @test Zarr.Codecs.V3Codecs.sharding_codec(plain) === nothing + + # Non-V3Pipeline argument → also nothing. + @test Zarr.Codecs.V3Codecs.sharding_codec(nothing) === nothing + end + + @testset "in-memory partial path via DictStore" begin + # DictStore returns supports_partial_reads=false by default, so + # readblock! goes through the in-memory partial-decode path. + @test Zarr.supports_partial_reads(Zarr.DictStore()) === false + + store = Zarr.DictStore() + z = _make_sharded_int16(store) + z[:] = Int16[1, 2, 3, 4, 5, 6, 7, 8] + + # Within a single inner chunk (chunk 1 of outer chunk 1). + @test z[1:2] == Int16[1, 2] + # Spans two inner chunks of one outer chunk. + @test z[2:3] == Int16[2, 3] + # Full outer chunk — falls through to the non-partial decode path. + @test z[1:4] == Int16[1, 2, 3, 4] + # Spans two outer chunks (and several inner chunks). + @test z[3:6] == Int16[3, 4, 5, 6] + # Tail of the array — last inner chunk of the second outer chunk. + @test z[7:8] == Int16[7, 8] + # Whole array — both outer chunks, every inner chunk. + @test z[:] == Int16[1, 2, 3, 4, 5, 6, 7, 8] + end + + @testset "storage-aware partial path via DirectoryStore" begin + @test Zarr.supports_partial_reads(Zarr.DirectoryStore(mktempdir())) === true + + dir = mktempdir() + store = Zarr.DirectoryStore(dir) + z = _make_sharded_int16(store) + z[:] = Int16[10, 20, 30, 40, 50, 60, 70, 80] + + # Same slice patterns as the DictStore test above; results must + # be byte-identical regardless of which path serviced them. + @test z[1:2] == Int16[10, 20] + @test z[2:3] == Int16[20, 30] + @test z[1:4] == Int16[10, 20, 30, 40] + @test z[3:6] == Int16[30, 40, 50, 60] + @test z[7:8] == Int16[70, 80] + @test z[:] == Int16[10, 20, 30, 40, 50, 60, 70, 80] + end + + @testset "enable_partial_shard_storage_reads[] toggle" begin + # Flipping the flag must preserve correctness — only the path + # taken changes. + dir = mktempdir() + store = Zarr.DirectoryStore(dir) + z = _make_sharded_int16(store) + z[:] = Int16[1, 2, 3, 4, 5, 6, 7, 8] + + prev = Zarr.enable_partial_shard_storage_reads[] + try + Zarr.enable_partial_shard_storage_reads[] = true + @test z[3:6] == Int16[3, 4, 5, 6] + + Zarr.enable_partial_shard_storage_reads[] = false + @test z[3:6] == Int16[3, 4, 5, 6] + finally + Zarr.enable_partial_shard_storage_reads[] = prev + end + end + + @testset "fill_value over partial slice of an empty shard" begin + # Second outer chunk is never written → its file doesn't exist; + # the partial path must populate the requested region with + # fill_value. + dir = mktempdir() + store = Zarr.DirectoryStore(dir) + z = _make_sharded_int16(store; fill_value=Int16(-7)) + z[1:2] = Int16[100, 200] # writes outer chunk 0 only + + @test z[1:2] == Int16[100, 200] + # Partial slice that lives entirely in the un-written outer chunk: + @test z[5:6] == Int16[-7, -7] + @test z[7:8] == Int16[-7, -7] + # Cross-boundary: writeblock! decoded outer chunk 0 from + # fill_value first, so the unwritten tail of chunk 0 is also + # fill_value (not zero); chunk 1 was never written so its bytes + # come straight from fill_value too. + @test z[3:6] == Int16[-7, -7, -7, -7] + end +end + @testset "ShardingCodec validate_index_pipeline rejects variable-size codecs" begin # Metadata with a blosc compressor inside index_codecs — must throw ArgumentError json_str = """{"zarr_format":3,"node_type":"array","shape":[4],"data_type":"int16", From e382675e78927ed2e9cea72a7b93dd68b463cdb1 Mon Sep 17 00:00:00 2001 From: Nehal Patel Date: Tue, 28 Apr 2026 01:34:03 +0000 Subject: [PATCH 3/3] CHANGELOG: add Unreleased entry for the partial-read fast path Required by the upstream changelog-enforcer CI check. Describes the new in-memory and storage-aware partial-decode paths, the optional AbstractStore methods (supports_partial_reads, read_range, getsize), DirectoryStore opt-in, and the enable_partial_shard_storage_reads[] toggle. Co-Authored-By: Claude Opus 4.7 (1M context) --- CHANGELOG.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 79ebbddd..97cf72b6 100644 --- a/CHANGELOG.md +++ b/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)