From ac0236044e5e82a8342ecdb4b83cac3e28bfb985 Mon Sep 17 00:00:00 2001 From: Nehal Patel Date: Mon, 27 Apr 2026 19:38:24 +0000 Subject: [PATCH 1/8] 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/8] 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 42050c56fe1e4490ff5c3c5b5bbea350dfc37e9c Mon Sep 17 00:00:00 2001 From: Nehal Patel Date: Mon, 27 Apr 2026 19:51:28 +0000 Subject: [PATCH 3/8] Parallelize sharded partial reads with @spawn MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sharding partial-read fast paths were single-threaded: a request that touched many inner chunks of one shard, or many outer chunks across the array, decoded them serially. Adds two layers of internal threading so one readblock! call scales with available cores. read_shard_partial_with_source! (Codecs/V3/V3.jl) - precomputes the work list of intersecting inner chunks (each writes to a disjoint region of `aout`, so safe to parallelize) - dispatches the decode loop with @sync / Threads.@spawn - bounded buffer pool (Channel of nthreads buffers) caps memory instead of allocating one full-shard buffer per task _readblock_sharded_partial! (ZArray.jl) - splits incoming blockr into full vs partial outer-chunk reads - full-chunk reads keep the existing serial path with one shared chunk buffer (already saturates one core) - partial-chunk reads dispatch with @sync / Threads.@spawn — each task reads its shard index + intersecting inner chunks - inner-chunk threading inside each task still applies, so multi-symbol multi-day queries get both axes of parallelism Toggle: - Zarr.enable_threaded_shard_decode[] (Ref{Bool}, default true). Falls back to sequential when nthreads()==1, |work|==1, or the flag is off. Threading is opt-out via the flag, so single-threaded callers see no behavior change. Both flag declarations moved above include() so nested modules can import them. Measured on a (129 syms × 1.48 yr) symbol-major archive, BTC × full quotes history, all 6 vars: user-serial path: 23.9s → 10.3s (2.3x from internal threading) user-level @threads over 6 vars: 12.7s → 5.2s Python (zarr-python 3, internal parallelism) baseline: 4.0s. Existing test suite: 2482/2482 pass. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/Codecs/V3/V3.jl | 70 ++++++++++++++++++++++++++++++++++++++------- src/ZArray.jl | 63 +++++++++++++++++++++++++++++++--------- src/Zarr.jl | 31 ++++++++++++++------ 3 files changed, 133 insertions(+), 31 deletions(-) diff --git a/src/Codecs/V3/V3.jl b/src/Codecs/V3/V3.jl index 71fc5ca1..3d67752a 100644 --- a/src/Codecs/V3/V3.jl +++ b/src/Codecs/V3/V3.jl @@ -5,6 +5,7 @@ import ..Codecs: zencode, zdecode, zencode!, zdecode! import ...Zarr: ZlibCompressor, ZstdCompressor, zcompress, zuncompress import ...Zarr: BloscCompressor as ZarrBloscCompressor import ...Zarr: AbstractCodecPipeline, V3Pipeline, pipeline_encode, pipeline_decode! +import ...Zarr: enable_threaded_shard_decode using CRC32c: CRC32c using JSON: JSON using ChunkCodecLibZlib: GzipCodec as LibZGzipCodec, GzipEncodeOptions @@ -675,8 +676,14 @@ function read_shard_partial_with_source!( 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) - + # Build the work list of intersecting inner chunks up front, then + # decode them. Each entry's output region is disjoint by + # construction (different ic_coords → non-overlapping aout_dest), so + # the decode loop is safe to parallelize when threads are available. + Work = Tuple{NTuple{N,UnitRange{Int}}, + NTuple{N,UnitRange{Int}}, + Union{Nothing,Tuple{Int,Int}}} + work = Work[] 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) @@ -698,17 +705,60 @@ function read_shard_partial_with_source!( lo:hi end - chunk_slice = get_chunk_slice(index, ic_coords) - if chunk_slice === nothing - view(aout, aout_dest...) .= fv - continue + cs = get_chunk_slice(index, ic_coords) + push!(work, (ic_local, aout_dest, cs === nothing ? nothing : (Int(cs[1]), Int(cs[2])))) + end + + if isempty(work) + return aout + end + + # Run sequentially when threading wouldn't help (single thread, one + # task, or the user has explicitly turned the knob off). + use_threads = Threads.nthreads() > 1 && + length(work) > 1 && + enable_threaded_shard_decode[] + + if !use_threads + inner_buf = Array{T}(undef, sc.chunk_shape) + for (ic_local, aout_dest, cs) in work + if cs === nothing + view(aout, aout_dest...) .= fv + else + offset_start, offset_end = cs + encoded_chunk = fetch_inner_chunk_bytes(offset_start, offset_end - offset_start) + pipeline_decode!(sc.codecs, inner_buf, encoded_chunk) + copyto!(view(aout, aout_dest...), view(inner_buf, ic_local...)) + end end + return aout + 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) + # Threaded path: bounded buffer pool so we don't allocate one + # full-shard buffer per task. min(nthreads, |work|) buffers is + # enough — each task takes one off the channel, decodes, returns it. + n_workers = min(Threads.nthreads(), length(work)) + pool = Channel{Array{T,N}}(n_workers) + for _ in 1:n_workers + put!(pool, Array{T,N}(undef, sc.chunk_shape)) + end - copyto!(view(aout, aout_dest...), view(inner_buf, ic_local...)) + @sync for (ic_local, aout_dest, cs) in work + Threads.@spawn begin + if cs === nothing + view(aout, aout_dest...) .= fv + else + offset_start, offset_end = cs + buf = take!(pool) + try + encoded_chunk = fetch_inner_chunk_bytes(offset_start, offset_end - offset_start) + pipeline_decode!(sc.codecs, buf, encoded_chunk) + copyto!(view(aout, aout_dest...), view(buf, ic_local...)) + finally + put!(pool, buf) + end + end + end end return aout end diff --git a/src/ZArray.jl b/src/ZArray.jl index e97ea146..5c4c3f82 100644 --- a/src/ZArray.jl +++ b/src/ZArray.jl @@ -255,34 +255,60 @@ function _readblock_sharded_partial!(aout, z::ZArray, r::CartesianIndices{N}, store = z.storage base_path = z.path fv = z.metadata.fill_value - a = nothing # lazy chunk buffer for the full-shard fallback path + # Two task lists. Full-chunk reads still go through the legacy + # decode path (one big shared chunk buffer); partial reads dispatch + # to the storage-partial path. Splitting this way keeps the buffer + # for the full-chunk path single-instance. + full_tasks = Tuple{NTuple{N,Int}, NTuple{N,UnitRange{Int}}, String}[] + partial_tasks = Tuple{NTuple{N,Int}, NTuple{N,UnitRange{Int}}, String}[] 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. + push!(full_tasks, (Tuple(current_chunk_offsets), Tuple(indranges), chunk_key)) + else + push!(partial_tasks, (Tuple(current_chunk_offsets), Tuple(indranges), chunk_key)) + end + end + + # Full-chunk reads keep the existing serial path with a shared + # chunk buffer — these are intrinsically large reads that decompress + # the whole shard, so the existing path already saturates one core + # per call. + if !isempty(full_tasks) + a = getchunkarray(z) + for (current_chunk_offsets, indranges, chunk_key) in full_tasks 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 + collect(current_chunk_offsets), a, collect(indranges)) end + end - # Partial path: index-only read, then per-inner-chunk byte ranges. + isempty(partial_tasks) && return aout + + # Run partial reads in parallel across outer chunks. Each task is + # one outer chunk: it reads the shard index, then either fills + # missing data with fill_value or hands off to + # read_shard_partial_with_source! which itself may spawn per-inner- + # chunk tasks. Tasks write to disjoint regions of `aout` (different + # outer chunks = different time windows by construction), so no + # locking needed. + use_threads = Threads.nthreads() > 1 && + length(partial_tasks) > 1 && + enable_threaded_shard_decode[] + + do_one = function(task) + current_chunk_offsets, indranges, chunk_key = task 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, + current_chunk_offsets, indranges, fv, ) - continue + return 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) @@ -299,8 +325,19 @@ function _readblock_sharded_partial!(aout, z::ZArray, r::CartesianIndices{N}, Codecs.V3Codecs.read_shard_partial_with_source!( aout, obo, sc, index_bytes, fetch_inner, - chunk_shape, Tuple(current_chunk_offsets), Tuple(indranges), fv, + chunk_shape, current_chunk_offsets, indranges, fv, ) + return + end + + if use_threads + @sync for task in partial_tasks + Threads.@spawn do_one(task) + end + else + for task in partial_tasks + do_one(task) + end end aout diff --git a/src/Zarr.jl b/src/Zarr.jl index 90032011..492e053f 100644 --- a/src/Zarr.jl +++ b/src/Zarr.jl @@ -12,14 +12,6 @@ ZarrFormat(v::ZarrFormat) = v #Default Zarr Version const DV = ZarrFormat(Val(2)) -include("types.jl") -include("chunkkeyencoding.jl") -include("metadata.jl") -include("metadata3.jl") -include("Compressors/Compressors.jl") -include("Codecs/Codecs.jl") -include("Storage/Storage.jl") -include("Filters/Filters.jl") """ enable_partial_shard_storage_reads[] @@ -33,6 +25,29 @@ Set to `false` to fall back to the in-memory partial-decode path """ const enable_partial_shard_storage_reads = Ref(true) +""" + enable_threaded_shard_decode[] + +When `true` (default) and Julia is running with more than one thread, +inner-chunk decompresses in `read_shard_partial_with_source!` are +dispatched to `Threads.@spawn` so that one shard read scales with +available cores. Falls back to a sequential loop on single-threaded +runs or when the work list has fewer than two entries. + +Set to `false` to force the sequential path even with `-t > 1` +(useful for debugging or for callers that already parallelize at a +higher level). +""" +const enable_threaded_shard_decode = Ref(true) + +include("types.jl") +include("chunkkeyencoding.jl") +include("metadata.jl") +include("metadata3.jl") +include("Compressors/Compressors.jl") +include("Codecs/Codecs.jl") +include("Storage/Storage.jl") +include("Filters/Filters.jl") include("ZArray.jl") include("pipeline.jl") include("ZGroup.jl") From ca47c6d73258ef33cf2c4a7f57350b3f40bf9960 Mon Sep 17 00:00:00 2001 From: Nehal Patel Date: Mon, 27 Apr 2026 20:38:02 +0000 Subject: [PATCH 4/8] In-place codec_decode! + cap inner-decode pool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two chunk-sized transient allocations on every inner-chunk decode in the V3 partial-shard path: one in zstd's `cc_decode` (returns a fresh Vector{UInt8}) and one in BytesCodec's `collect(reinterpret(...))`. pipeline_decode! then did a final `copyto!(output, arr)` — a third copy. For a 187 MB inner chunk that's GBs of waste per query. Patch A: in-place codec_decode! API (Codecs/V3/V3.jl) - Generic fallbacks dispatched on V3Codec{In,Out} so any codec gets a working method (allocates + copies, but no caller-visible API change). - Specialized BytesCodec — reinterpret bytes as UInt8 view of the output array and `copyto!` straight in. No fresh Vector{T}. - Specialized ZstdV3Codec — calls ChunkCodecCore.decode! into the caller's buffer directly. No fresh decoded Vector{UInt8}. Patch B: pipeline_decode! threads in-place dispatch (pipeline.jl) - For the common case (no array_array codecs, ≥1 bytes_bytes codec), sizes a single intermediate buffer to the array-bytes step's expected output and reuses it across the chain via codec_decode!. - Final array-bytes step writes straight into `output` — eliminates the `arr` allocation + redundant copy. - Falls back to the old path when transpose codecs or other array-array steps are present. Patch C: cap inner-decode buffer pool (Zarr.jl + V3.jl) - New Zarr.max_concurrent_inner_decodes Ref{Int} (default 8), modeled on zarr-python's async.concurrency = 10. - read_shard_partial_with_source! pool sized to min(nthreads, |work|, max_concurrent_inner_decodes[]). - On a `-t 32` run with ~187 MB chunks the pool no longer reserves ~6 GB upfront for buffers most tasks never use. Measured (HL BTC quotes full history × 6 vars on hyperliquid_1s_symmajor.zarr, -t 32): user-serial: 9.6s → 7.2s (-25%) user @threads-over-vars: 5.1s → 3.3s (-35%, beats Python's 4.0s) Allocations dropped 35-40% across partial-read benches. Existing test suite: 2482/2482 pass. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/Codecs/V3/V3.jl | 91 ++++++++++++++++++++++++++++++++++++++++++--- src/Zarr.jl | 16 ++++++++ src/pipeline.jl | 55 +++++++++++++++++++++++---- 3 files changed, 149 insertions(+), 13 deletions(-) diff --git a/src/Codecs/V3/V3.jl b/src/Codecs/V3/V3.jl index 3d67752a..2ef0031d 100644 --- a/src/Codecs/V3/V3.jl +++ b/src/Codecs/V3/V3.jl @@ -5,11 +5,11 @@ import ..Codecs: zencode, zdecode, zencode!, zdecode! import ...Zarr: ZlibCompressor, ZstdCompressor, zcompress, zuncompress import ...Zarr: BloscCompressor as ZarrBloscCompressor import ...Zarr: AbstractCodecPipeline, V3Pipeline, pipeline_encode, pipeline_decode! -import ...Zarr: enable_threaded_shard_decode +import ...Zarr: enable_threaded_shard_decode, max_concurrent_inner_decodes using CRC32c: CRC32c using JSON: JSON using ChunkCodecLibZlib: GzipCodec as LibZGzipCodec, GzipEncodeOptions -using ChunkCodecCore: encode as cc_encode, decode as cc_decode +using ChunkCodecCore: encode as cc_encode, decode as cc_decode, decode! as cc_decode! abstract type V3Codec{In,Out} end @@ -735,9 +735,11 @@ function read_shard_partial_with_source!( end # Threaded path: bounded buffer pool so we don't allocate one - # full-shard buffer per task. min(nthreads, |work|) buffers is - # enough — each task takes one off the channel, decodes, returns it. - n_workers = min(Threads.nthreads(), length(work)) + # full-shard buffer per task. Cap by `max_concurrent_inner_decodes` + # (mirrors zarr-python's `async.concurrency = 10`) — each chunk + # buffer is ~prod(chunk_shape)*sizeof(T) bytes, so allocating one + # per OS thread on a `-t 32` run wastes GBs. + n_workers = min(Threads.nthreads(), length(work), max_concurrent_inner_decodes[]) pool = Channel{Array{T,N}}(n_workers) for _ in 1:n_workers put!(pool, Array{T,N}(undef, sc.chunk_shape)) @@ -839,6 +841,47 @@ function JSON.lower(c::TransposeCodec) end # codec_encode / codec_decode methods for V3 codecs +# +# Each codec also exposes a `codec_decode!(c, output, encoded[; …])` variant +# that writes into a caller-owned `output` buffer instead of returning a +# fresh allocation. Used by `pipeline_decode!` on the partial-shard read +# path to avoid chunk-sized transient allocations on every inner-chunk +# decode. Generic fallbacks below dispatch on the codec's In/Out tags so +# any codec that doesn't override gets a working (allocating) default. + +# Generic in-place fallbacks. These exist so callers can always write +# `codec_decode!(c, out, enc)` without worrying about whether the codec +# was specialized; the fallback just does the allocating decode and +# copies. Specialized methods follow per-codec. +function codec_decode!(c::V3Codec{:bytes,:bytes}, + output::AbstractVector{UInt8}, + encoded::AbstractVector{UInt8}) + bytes = codec_decode(c, encoded isa Vector{UInt8} ? encoded : collect(encoded)) + length(bytes) == length(output) || + throw(DimensionMismatch("codec_decode! output sized $(length(output)) " * + "but $(name(c)) returned $(length(bytes)) bytes")) + copyto!(output, bytes) + return output +end + +function codec_decode!(c::V3Codec{:array,:bytes}, + output::AbstractArray{T,N}, + encoded::AbstractVector{UInt8}; + fill_value=nothing) where {T, N} + arr = codec_decode(c, encoded isa Vector{UInt8} ? encoded : collect(encoded), + T, size(output); fill_value) + copyto!(output, arr) + return output +end + +function codec_decode!(c::V3Codec{:array,:array}, + output::AbstractArray, + encoded::AbstractArray) + arr = codec_decode(c, encoded) + copyto!(output, arr) + return output +end + function codec_encode(c::BytesCodec, data::AbstractArray) if _needs_bswap(c.endian) @@ -856,6 +899,30 @@ function codec_decode(c::BytesCodec, encoded::Vector{UInt8}, ::Type{T}, shape::N return reshape(arr, shape) end +# In-place: copy bytes straight into output, skipping the +# `collect(reinterpret(...))` materialization. For little-endian → little- +# endian (the common case), this is a single bulk byte copy. For mixed +# endian we still do the bswap but at least there's no fresh Vector{T} +# allocation per inner-chunk decode. +function codec_decode!(c::BytesCodec, + output::AbstractArray{T,N}, + encoded::AbstractVector{UInt8}; + fill_value=nothing) where {T, N} + expected = length(output) * sizeof(T) + length(encoded) == expected || + throw(DimensionMismatch("BytesCodec output buffer expects $expected " * + "bytes but encoded length is $(length(encoded))")) + out_bytes = reinterpret(UInt8, vec(output)) # zero-copy view + copyto!(out_bytes, encoded) + if _needs_bswap(c.endian) + out_vec = vec(output) + @inbounds for i in eachindex(out_vec) + out_vec[i] = bswap(out_vec[i]) + end + end + return output +end + function codec_encode(c::ShardingCodec, data::AbstractArray) encoded = UInt8[] zencode!(encoded, data, c) @@ -977,6 +1044,20 @@ function codec_decode(c::ZstdV3Codec, encoded::Vector{UInt8}) return collect(zuncompress(encoded, comp, UInt8)) end +# In-place: decompress directly into the caller's buffer via +# ChunkCodecCore.decode!. The buffer must already be sized to the +# decompressed length (callers know it: it's the inner-chunk byte size +# = prod(chunk_shape) * sizeof(T) for the sharded V3 read path). +function codec_decode!(c::ZstdV3Codec, + output::AbstractVector{UInt8}, + encoded::AbstractVector{UInt8}) + comp = ZstdCompressor(level=c.level) + cc_decode!(comp.config.codec, + output isa Vector{UInt8} ? output : @view(output[1:end]), + encoded isa Vector{UInt8} ? encoded : collect(encoded)) + return output +end + struct CRC32cV3Codec <: V3Codec{:bytes, :bytes} end name(::CRC32cV3Codec) = "crc32c" diff --git a/src/Zarr.jl b/src/Zarr.jl index 492e053f..e14f6b82 100644 --- a/src/Zarr.jl +++ b/src/Zarr.jl @@ -40,6 +40,22 @@ higher level). """ const enable_threaded_shard_decode = Ref(true) +""" + max_concurrent_inner_decodes[] + +Upper bound on the number of inner-chunk decode tasks dispatched in +parallel by `read_shard_partial_with_source!`. Default `8`, modeled on +zarr-python's `async.concurrency = 10`. Tasks beyond this queue on +the buffer-pool channel. + +Capping matters because each in-flight task holds a chunk-sized buffer +(e.g. 187 MB for our 1s archive shape) and `Threads.nthreads()` is +often much larger than the number of inner chunks per shard. Without +the cap, a `-t 32` run pre-allocates ~6 GB of decode buffers even when +the actual work is a handful of chunks. +""" +const max_concurrent_inner_decodes = Ref(8) + include("types.jl") include("chunkkeyencoding.jl") include("metadata.jl") diff --git a/src/pipeline.jl b/src/pipeline.jl index d172f9f2..e896317a 100644 --- a/src/pipeline.jl +++ b/src/pipeline.jl @@ -31,19 +31,58 @@ function pipeline_encode(p::V3Pipeline, data::AbstractArray, fill_value) end function pipeline_decode!(p::V3Pipeline, output::AbstractArray, compressed::Vector{UInt8}; fill_value=nothing) - # Phase 3 reverse: bytes->bytes codecs (reverse order) - bytes = compressed - for codec in reverse(collect(p.bytes_bytes)) - bytes = Codecs.V3Codecs.codec_decode(codec, bytes) - end - # Phase 2 reverse: bytes->array codec - # Compute the intermediate shape — the shape data has after array_array encoding + # Compute the byte size of the array-bytes step's output — equal to + # the (post-array_array) decoded element count × sizeof(eltype). For + # the common sharded inner-chunk path this is exactly the inner chunk + # in bytes, which lets us route bytes-bytes codecs through a single + # reusable buffer and the final array-bytes step straight into + # `output` with no intermediate Array{T} allocation. intermediate_shape = foldl( (sz, codec) -> Codecs.V3Codecs.encoded_shape(codec, sz), p.array_array; init=size(output) ) + decoded_byte_size = prod(intermediate_shape) * sizeof(eltype(output)) + + # Walk the bytes-bytes chain in reverse. The last step's output must + # be exactly `decoded_byte_size`; intermediate steps run via the + # allocating `codec_decode` (their output sizes vary per codec and + # aren't worth pre-sizing for the rare multi-bytes-codec case). + n_bb = length(p.bytes_bytes) + if isempty(p.array_array) && n_bb >= 1 + # Common case: pipeline ends in [array_bytes] and starts with + # one or more bytes-bytes codecs. The first reverse step (= last + # forward step) gets the final-sized output buffer; downstream + # steps cascade through `codec_decode`. For the dominant case + # `[BytesCodec, ZstdCompressor]` (n_bb == 1), this collapses to + # one in-place zstd decode + one in-place BytesCodec decode and + # zero chunk-sized intermediate allocations. + bytes_buf = Vector{UInt8}(undef, decoded_byte_size) + bytes = compressed + bb = collect(p.bytes_bytes) + for i in length(bb):-1:1 + codec = bb[i] + if i == 1 + Codecs.V3Codecs.codec_decode!(codec, bytes_buf, bytes) + else + bytes = Codecs.V3Codecs.codec_decode(codec, bytes) + end + end + Codecs.V3Codecs.codec_decode!(p.array_bytes, output, bytes_buf; fill_value) + return output + end + + # Fallback path: array_array codecs present (e.g. transpose) or no + # bytes-bytes codecs. Same logic as before, allocating where + # necessary. + bytes = compressed + for codec in reverse(collect(p.bytes_bytes)) + bytes = Codecs.V3Codecs.codec_decode(codec, bytes) + end + if isempty(p.array_array) + Codecs.V3Codecs.codec_decode!(p.array_bytes, output, bytes; fill_value) + return output + end arr = Codecs.V3Codecs.codec_decode(p.array_bytes, bytes, eltype(output), intermediate_shape; fill_value) - # Phase 1 reverse: array->array codecs (reverse order) for codec in reverse(collect(p.array_array)) arr = Codecs.V3Codecs.codec_decode(codec, arr) end From 390bd58b29c2af09016d8c88e28b4f160a6e86a6 Mon Sep 17 00:00:00 2001 From: Nehal Patel Date: Mon, 27 Apr 2026 20:54:04 +0000 Subject: [PATCH 5/8] Skip the bytes scratch buffer when array_bytes is BytesCodec MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pipeline_decode! was still allocating a chunk-sized Vector{UInt8} on every inner-chunk decode for the bytes-bytes step's output. For our prod shape that's ~187 MB per inner chunk × ~7 outer chunks per query = ~1.3 GB of transient allocations per arr[:, si] call. The dominant pipeline shape is `[BytesCodec, ZstdV3Codec]` — BytesCodec just reinterprets bytes as the array's element type, so the bytes-bytes output IS the byte view of the typed output array. Decoding zstd straight into `reinterpret(UInt8, vec(output))` eliminates the scratch entirely. Two new fast paths added: - matching-endian common case: zstd into output's byte view, return. - mismatched-endian variant: zstd into output's byte view, then in-place bswap via codec_decode!(::BytesCodec, ...). For multi-step bytes-bytes chains or pipelines with array_array codecs (transpose etc.), keep the existing scratch-buffer / fallback paths unchanged. Measured (HL BTC quotes full history × 6 vars, -t 32): user-serial: 7.2s → 3.2s (Python 4.0s) user @threads-over-vars: 3.3s → 1.1s per-call alloc on f64 vars: 2.9 GB → 1.7 GB Existing test suite: 2482/2482 pass. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/pipeline.jl | 68 ++++++++++++++++++++++++++++++------------------- 1 file changed, 42 insertions(+), 26 deletions(-) diff --git a/src/pipeline.jl b/src/pipeline.jl index e896317a..ed3862c4 100644 --- a/src/pipeline.jl +++ b/src/pipeline.jl @@ -31,31 +31,44 @@ function pipeline_encode(p::V3Pipeline, data::AbstractArray, fill_value) end function pipeline_decode!(p::V3Pipeline, output::AbstractArray, compressed::Vector{UInt8}; fill_value=nothing) - # Compute the byte size of the array-bytes step's output — equal to - # the (post-array_array) decoded element count × sizeof(eltype). For - # the common sharded inner-chunk path this is exactly the inner chunk - # in bytes, which lets us route bytes-bytes codecs through a single - # reusable buffer and the final array-bytes step straight into - # `output` with no intermediate Array{T} allocation. intermediate_shape = foldl( (sz, codec) -> Codecs.V3Codecs.encoded_shape(codec, sz), p.array_array; init=size(output) ) decoded_byte_size = prod(intermediate_shape) * sizeof(eltype(output)) - # Walk the bytes-bytes chain in reverse. The last step's output must - # be exactly `decoded_byte_size`; intermediate steps run via the - # allocating `codec_decode` (their output sizes vary per codec and - # aren't worth pre-sizing for the rare multi-bytes-codec case). n_bb = length(p.bytes_bytes) + ab = p.array_bytes + + # Hot path: pipeline is exactly `[BytesCodec, one bytes_bytes codec]` + # with matching endian. The bytes-bytes step's output IS the byte + # view of `output`'s underlying memory, so we decompress straight + # into it — no intermediate Vector{UInt8} allocation, no second + # copy from a scratch buffer into the typed array. This is what + # this archive (and most well-configured sharded archives) exercise + # on every inner-chunk decode. + if isempty(p.array_array) && n_bb == 1 && ab isa Codecs.V3Codecs.BytesCodec && + !Codecs.V3Codecs._needs_bswap(ab.endian) + bytes_view = reinterpret(UInt8, vec(output)) + Codecs.V3Codecs.codec_decode!(only(p.bytes_bytes), bytes_view, compressed) + return output + end + + # Endian-mismatch variant: still avoid the extra buffer by + # decoding into the byte view, then byte-swapping in place via the + # array_bytes codec's in-place dispatch. + if isempty(p.array_array) && n_bb == 1 && ab isa Codecs.V3Codecs.BytesCodec + bytes_view = reinterpret(UInt8, vec(output)) + Codecs.V3Codecs.codec_decode!(only(p.bytes_bytes), bytes_view, compressed) + # codec_decode!(::BytesCodec) handles bswap when needed. + Codecs.V3Codecs.codec_decode!(ab, output, bytes_view; fill_value) + return output + end + + # Multi bytes-bytes step (rare): need one chunk-sized scratch buffer + # to chain through. Final step writes into the buffer; array_bytes + # then writes into output. if isempty(p.array_array) && n_bb >= 1 - # Common case: pipeline ends in [array_bytes] and starts with - # one or more bytes-bytes codecs. The first reverse step (= last - # forward step) gets the final-sized output buffer; downstream - # steps cascade through `codec_decode`. For the dominant case - # `[BytesCodec, ZstdCompressor]` (n_bb == 1), this collapses to - # one in-place zstd decode + one in-place BytesCodec decode and - # zero chunk-sized intermediate allocations. bytes_buf = Vector{UInt8}(undef, decoded_byte_size) bytes = compressed bb = collect(p.bytes_bytes) @@ -67,22 +80,25 @@ function pipeline_decode!(p::V3Pipeline, output::AbstractArray, compressed::Vect bytes = Codecs.V3Codecs.codec_decode(codec, bytes) end end - Codecs.V3Codecs.codec_decode!(p.array_bytes, output, bytes_buf; fill_value) + Codecs.V3Codecs.codec_decode!(ab, output, bytes_buf; fill_value) return output end - # Fallback path: array_array codecs present (e.g. transpose) or no - # bytes-bytes codecs. Same logic as before, allocating where - # necessary. + # No bytes-bytes step (uncompressed): array_bytes from the encoded + # input directly into `output`. + if isempty(p.array_array) + Codecs.V3Codecs.codec_decode!(ab, output, compressed; fill_value) + return output + end + + # Fallback for pipelines with array_array codecs (e.g. transpose). + # Allocate as before — these are uncommon enough that further + # tuning isn't worth the case-analysis. bytes = compressed for codec in reverse(collect(p.bytes_bytes)) bytes = Codecs.V3Codecs.codec_decode(codec, bytes) end - if isempty(p.array_array) - Codecs.V3Codecs.codec_decode!(p.array_bytes, output, bytes; fill_value) - return output - end - arr = Codecs.V3Codecs.codec_decode(p.array_bytes, bytes, eltype(output), intermediate_shape; fill_value) + arr = Codecs.V3Codecs.codec_decode(ab, bytes, eltype(output), intermediate_shape; fill_value) for codec in reverse(collect(p.array_array)) arr = Codecs.V3Codecs.codec_decode(codec, arr) end From 3161866c882e96fbd18ea75438b476d1aaff2188 Mon Sep 17 00:00:00 2001 From: Nehal Patel Date: Mon, 27 Apr 2026 22:21:47 +0000 Subject: [PATCH 6/8] Test coverage + docs for the threading & in-place codec work MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tests - test/v3_codecs.jl: three new testsets, 34 cases: * "codec_decode! in-place API" — BytesCodec little-endian and big-endian (bswap) round-trips, dimension-mismatch error, ZstdV3Codec in-place decode, and the generic V3Codec{:bytes, :bytes} fallback via CRC32cV3Codec. * "pipeline_decode! V3 paths" — exercises each branch: matching-endian fast path (BytesCodec :little + Zstd), endian-mismatch variant (BytesCodec :big), no-bytes-bytes path (BytesCodec only), multi bytes-bytes scratch-buffer branch (Zstd + CRC32c), and the array_array fallback (TransposeCodec). * "threading flags preserve correctness" — flips enable_threaded_shard_decode[] and max_concurrent_inner_decodes[] across realistic combinations, verifies reads still match. Docs - docs/src/UserGuide/partial_shard_reads.md: new "Threading" and "In-place codec API" sections describing the user-transparent parallelism, the two new toggles, and how downstream codecs can opt into in-place dispatch. No production code changed in this commit. Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/src/UserGuide/partial_shard_reads.md | 40 +++++ test/v3_codecs.jl | 170 ++++++++++++++++++++++ 2 files changed, 210 insertions(+) diff --git a/docs/src/UserGuide/partial_shard_reads.md b/docs/src/UserGuide/partial_shard_reads.md index 243f7d69..8973a217 100644 --- a/docs/src/UserGuide/partial_shard_reads.md +++ b/docs/src/UserGuide/partial_shard_reads.md @@ -67,6 +67,46 @@ 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. +## Threading + +When Julia is started with more than one thread (`julia -t N`), the +sharded partial-read path internally dispatches inner-chunk decodes to +`Threads.@spawn` so a single `arr[a:b, c]` call scales with available +cores — the same way `zarr-python` parallelizes inner-chunk decodes +inside one `__getitem__`. User code is unchanged; the threading is +transparent. + +Two layers of parallelism kick in automatically: + +1. Within one outer chunk, the inner chunks intersecting the request + decode in parallel (bounded by `max_concurrent_inner_decodes[]`). +2. Across outer chunks (when the request spans more than one), the + per-chunk reads dispatch in parallel. + +Both fall back to a sequential loop on single-threaded runs or when +the work list has fewer than two entries. + +```@docs +Zarr.enable_threaded_shard_decode +Zarr.max_concurrent_inner_decodes +``` + +## In-place codec API + +The decode pipeline avoids per-inner-chunk transient allocations by +threading the caller's output buffer through each codec via +`Zarr.Codecs.V3Codecs.codec_decode!`. For the dominant pipeline shape +`[BytesCodec, ZstdV3Codec]` (matching system endian), the inner chunk +decompresses straight into `reinterpret(UInt8, vec(output))` — no +intermediate `Vector{UInt8}` and no second copy. + +`codec_decode!` is dispatched on the codec's `In/Out` tag pair so any +new codec written for `Zarr.jl` can opt in by adding a specialized +method. A generic fallback (allocate + `copyto!`) is provided for +`V3Codec{:bytes,:bytes}`, `V3Codec{:array,:bytes}`, and +`V3Codec{:array,:array}`, so codecs without a specialization remain +correct, just less alloc-friendly. + ## Reference The storage-aware path lives in `Zarr._readblock_sharded_partial!` diff --git a/test/v3_codecs.jl b/test/v3_codecs.jl index 69997383..3e6b5c29 100644 --- a/test/v3_codecs.jl +++ b/test/v3_codecs.jl @@ -1349,6 +1349,176 @@ end end end +@testset "codec_decode! in-place API" begin + @testset "BytesCodec little-endian copies straight into output" begin + bc = Zarr.Codecs.V3Codecs.BytesCodec(:little) + data = Float64[1.5, -2.5, 3.5, -4.5] + encoded = collect(reinterpret(UInt8, data)) + output = Vector{Float64}(undef, length(data)) + Zarr.Codecs.V3Codecs.codec_decode!(bc, output, encoded) + @test output == data + end + + @testset "BytesCodec big-endian bswaps in place" begin + # On a little-endian host, BytesCodec(:big) needs bswap. Encode + # by bswapping the data bytes; the in-place decode must undo it. + bc = Zarr.Codecs.V3Codecs.BytesCodec(:big) + data = Int32[100, -200, 300, -400] + encoded = collect(reinterpret(UInt8, bswap.(data))) + output = Vector{Int32}(undef, length(data)) + Zarr.Codecs.V3Codecs.codec_decode!(bc, output, encoded) + @test output == data + end + + @testset "BytesCodec rejects size mismatch" begin + bc = Zarr.Codecs.V3Codecs.BytesCodec(:little) + output = Vector{Float64}(undef, 3) + @test_throws DimensionMismatch Zarr.Codecs.V3Codecs.codec_decode!( + bc, output, UInt8[0, 0, 0]) + end + + @testset "ZstdV3Codec round-trip via in-place decode" begin + zc = Zarr.Codecs.V3Codecs.ZstdV3Codec(3) + raw = collect(0x10:0x4F) # 64 bytes + encoded = Zarr.Codecs.V3Codecs.codec_encode(zc, raw) + output = Vector{UInt8}(undef, length(raw)) + Zarr.Codecs.V3Codecs.codec_decode!(zc, output, encoded) + @test output == raw + end + + @testset "Generic fallback for bytes→bytes codec without specialization" begin + # CRC32cV3Codec doesn't have its own codec_decode! method; the + # generic V3Codec{:bytes,:bytes} fallback should still work. + cc = Zarr.Codecs.V3Codecs.CRC32cV3Codec() + raw = collect(0x01:0x10) + encoded = Zarr.Codecs.V3Codecs.codec_encode(cc, raw) + output = Vector{UInt8}(undef, length(raw)) + Zarr.Codecs.V3Codecs.codec_decode!(cc, output, encoded) + @test output == raw + end +end + +@testset "pipeline_decode! V3 paths" begin + @testset "matching-endian fast path: [BytesCodec :little, Zstd]" begin + # The hot path. Decompresses zstd directly into reinterpret(UInt8, + # vec(output)) — no scratch buffer. + pipeline = Zarr.V3Pipeline( + (), + Zarr.Codecs.V3Codecs.BytesCodec(:little), + (Zarr.Codecs.V3Codecs.ZstdV3Codec(3),), + ) + data = Float64[1.0, 2.0, 3.0, 4.0, 5.0] + encoded = Zarr.pipeline_encode(pipeline, data, nothing) + output = Vector{Float64}(undef, length(data)) + Zarr.pipeline_decode!(pipeline, output, encoded) + @test output == data + end + + @testset "endian-mismatch variant: BytesCodec :big triggers in-place bswap" begin + pipeline = Zarr.V3Pipeline( + (), + Zarr.Codecs.V3Codecs.BytesCodec(:big), + (Zarr.Codecs.V3Codecs.ZstdV3Codec(3),), + ) + data = Int32[1, 2, 3, 4, 5, 6, 7, 8] + encoded = Zarr.pipeline_encode(pipeline, data, nothing) + output = Vector{Int32}(undef, length(data)) + Zarr.pipeline_decode!(pipeline, output, encoded) + @test output == data + end + + @testset "no bytes-bytes step: BytesCodec only" begin + # Uncompressed array; pipeline_decode! routes straight through + # codec_decode!(::BytesCodec, ...). + pipeline = Zarr.V3Pipeline( + (), + Zarr.Codecs.V3Codecs.BytesCodec(:little), + (), + ) + data = Float64[1.5, 2.5, 3.5, 4.5] + encoded = Zarr.pipeline_encode(pipeline, data, nothing) + output = Vector{Float64}(undef, length(data)) + Zarr.pipeline_decode!(pipeline, output, encoded) + @test output == data + end + + @testset "multi bytes-bytes step uses scratch buffer" begin + # Two bytes-bytes codecs forces the scratch-buffer branch (the + # zero-buf fast path only fires for a single bytes-bytes codec). + pipeline = Zarr.V3Pipeline( + (), + Zarr.Codecs.V3Codecs.BytesCodec(:little), + (Zarr.Codecs.V3Codecs.ZstdV3Codec(3), + Zarr.Codecs.V3Codecs.CRC32cV3Codec()), + ) + data = Float64[10.0, 20.0, 30.0] + encoded = Zarr.pipeline_encode(pipeline, data, nothing) + output = Vector{Float64}(undef, length(data)) + Zarr.pipeline_decode!(pipeline, output, encoded) + @test output == data + end + + @testset "with array_array codec: TransposeCodec falls through to legacy path" begin + # TransposeCodec((2,1)) on a (2,3) matrix → encoded shape (3,2). + pipeline = Zarr.V3Pipeline( + (Zarr.Codecs.V3Codecs.TransposeCodec((2, 1)),), + Zarr.Codecs.V3Codecs.BytesCodec(:little), + (Zarr.Codecs.V3Codecs.ZstdV3Codec(3),), + ) + data = reshape(Float64.(1:6), 2, 3) + encoded = Zarr.pipeline_encode(pipeline, data, nothing) + output = Matrix{Float64}(undef, 2, 3) + Zarr.pipeline_decode!(pipeline, output, encoded) + @test output == data + end +end + +@testset "threading flags preserve correctness" begin + # Sharded ZArray; flip enable_threaded_shard_decode[] and + # max_concurrent_inner_decodes[] across realistic combinations and + # confirm reads still match the written data. Catches any divergence + # between the threaded and sequential paths. + inner_pipeline = Zarr.V3Pipeline( + (), + Zarr.Codecs.V3Codecs.BytesCodec(:little), + (Zarr.Codecs.V3Codecs.ZstdV3Codec(3),), + ) + 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, Int16(0), + Zarr.ChunkKeyEncoding('/', true), + ) + + dir = mktempdir() + store = Zarr.DirectoryStore(dir) + z = Zarr.ZArray(md, store, "", Dict(), true) + z[:] = Int16[1, 2, 3, 4, 5, 6, 7, 8] + + prev_threaded = Zarr.enable_threaded_shard_decode[] + prev_max = Zarr.max_concurrent_inner_decodes[] + try + for threaded in (true, false), max_workers in (1, 2, 8) + Zarr.enable_threaded_shard_decode[] = threaded + Zarr.max_concurrent_inner_decodes[] = max_workers + @test z[2:5] == Int16[2, 3, 4, 5] # cross outer-chunk boundary + @test z[1:4] == Int16[1, 2, 3, 4] # full outer chunk + @test z[3:3] == Int16[3] # single-element slice + @test z[:] == Int16[1, 2, 3, 4, 5, 6, 7, 8] + end + finally + Zarr.enable_threaded_shard_decode[] = prev_threaded + Zarr.max_concurrent_inner_decodes[] = prev_max + 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 83900bd000a73e763dc02c703e56bf055d6935cd Mon Sep 17 00:00:00 2001 From: Nehal Patel Date: Tue, 28 Apr 2026 01:34:03 +0000 Subject: [PATCH 7/8] 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) From 92aea90abe06aa9855ef0d9f761f93cb2157b85c Mon Sep 17 00:00:00 2001 From: Nehal Patel Date: Tue, 28 Apr 2026 01:34:47 +0000 Subject: [PATCH 8/8] CHANGELOG: add Unreleased entry for the threading + in-place codec work Describes the internal Threads.@spawn dispatch in read_shard_partial_with_source! and _readblock_sharded_partial!, the new enable_threaded_shard_decode[] and max_concurrent_inner_decodes[] toggles, the in-place codec_decode! API, and the rewritten V3 pipeline_decode! that eliminates per-inner-chunk scratch allocations. Co-Authored-By: Claude Opus 4.7 (1M context) --- CHANGELOG.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 97cf72b6..2ef81a64 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,12 @@ ## Unreleased +- Internal threading and in-place codec API for sharded reads [#265](https://github.com/JuliaIO/Zarr.jl/pull/265) + - inner-chunk decodes within one shard and outer-chunk reads across one request now dispatch to `Threads.@spawn` when Julia is started with `-t > 1` — single `arr[a:b, c]` calls scale with available cores, transparent to user code + - new `Zarr.enable_threaded_shard_decode[]` `Ref{Bool}` flag (default `true`); set `false` to force the sequential decode path + - new `Zarr.max_concurrent_inner_decodes[]` `Ref{Int}` flag (default `8`, modeled on zarr-python's `async.concurrency = 10`) caps the buffer pool size independently of `Threads.nthreads()` + - new `Codecs.V3Codecs.codec_decode!(c, output, encoded; …)` API that writes into a caller-owned buffer; specialized for `BytesCodec` (zero-copy reinterpret + bulk byte copy) and `ZstdV3Codec` (`ChunkCodecCore.decode!` straight into the caller's buffer); generic fallbacks dispatched on the codec's `In/Out` tag pair + - V3 `pipeline_decode!` rewritten to thread the caller's output array through the decode chain — for the dominant `[BytesCodec, ZstdV3Codec]` shape it decodes zstd directly into `reinterpret(UInt8, vec(output))`, eliminating per-inner-chunk transient allocations - 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