diff --git a/CHANGELOG.md b/CHANGELOG.md index 79ebbddd..2ef81a64 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,19 @@ ## 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 + - `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) diff --git a/docs/src/UserGuide/partial_shard_reads.md b/docs/src/UserGuide/partial_shard_reads.md new file mode 100644 index 00000000..8973a217 --- /dev/null +++ b/docs/src/UserGuide/partial_shard_reads.md @@ -0,0 +1,118 @@ +# 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. + +## 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!` +(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/src/Codecs/V3/V3.jl b/src/Codecs/V3/V3.jl index 9bfb28d8..2ef0031d 100644 --- a/src/Codecs/V3/V3.jl +++ b/src/Codecs/V3/V3.jl @@ -5,10 +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, 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 @@ -592,6 +593,224 @@ 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) + + # 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) + 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 + + 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 + + # Threaded path: bounded buffer pool so we don't allocate one + # 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)) + end + + @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 + + +""" + 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 @@ -622,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) @@ -639,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) @@ -760,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/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..5c4c3f82 100644 --- a/src/ZArray.jl +++ b/src/ZArray.jl @@ -168,37 +168,178 @@ 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 + + # 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 + 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] + uncompress_to_output!(aout, output_base_offsets, z, chunk_compressed, + collect(current_chunk_offsets), a, collect(indranges)) + end + end + + 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 + Codecs.V3Codecs.read_shard_partial!( + aout, obo, sc, nothing, chunk_shape, + current_chunk_offsets, indranges, fv, + ) + 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) + 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, 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 end diff --git a/src/Zarr.jl b/src/Zarr.jl index 917bbaa0..e14f6b82 100644 --- a/src/Zarr.jl +++ b/src/Zarr.jl @@ -12,6 +12,50 @@ ZarrFormat(v::ZarrFormat) = v #Default Zarr Version const DV = ZarrFormat(Val(2)) +""" + 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) + +""" + 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) + +""" + 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..ed3862c4 100644 --- a/src/pipeline.jl +++ b/src/pipeline.jl @@ -31,19 +31,74 @@ 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 intermediate_shape = foldl( (sz, codec) -> Codecs.V3Codecs.encoded_shape(codec, sz), p.array_array; init=size(output) ) - arr = Codecs.V3Codecs.codec_decode(p.array_bytes, bytes, eltype(output), intermediate_shape; fill_value) - # Phase 1 reverse: array->array codecs (reverse order) + decoded_byte_size = prod(intermediate_shape) * sizeof(eltype(output)) + + 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 + 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!(ab, output, bytes_buf; fill_value) + return output + end + + # 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 + 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 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..3e6b5c29 100644 --- a/test/v3_codecs.jl +++ b/test/v3_codecs.jl @@ -1213,6 +1213,312 @@ 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 "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",