diff --git a/CHANGELOG.md b/CHANGELOG.md index 8a09e2f9..bd2ec858 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,13 @@ ## Unreleased +- V2 performance improvements [#280](https://github.com/JuliaIO/Zarr.jl/pull/280) + - `readblock!` fast path for single-chunk full-reads that bypasses the readtask channel and the chunk-shaped scratch buffer, decoding straight into the output array + - `writeblock!` fast path for single-chunk full-overwrites that bypasses the readtask/writetask channels and the chunk-shaped scratch buffer, encoding straight from the input array into the store + - Zero-copy chunk write for `NoCompressor` + no filters: the single-chunk fastpath hands a `reinterpret(UInt8, ain)` view straight to the store, skipping the chunk-sized `Vector{UInt8}` allocation + memcpy + - `NoCompressor` writes: replace `append!` over the reinterpret view with bulk `resize!` + `copyto!` in the generic `zcompress!` fallback + - `NoCompressor` reads: bulk-copy `zuncompress!` method dispatched on `::NoCompressor` bypasses `copyto!(::Array, ::ReinterpretArray)`'s element-by-element walk + - `getchunkarray_undef` skips the dead zero-fill of the chunk-shaped scratch buffer on full-overwrite paths - Added manual pagination in order to go beyond the default 1k [#282](https://github.com/JuliaIO/Zarr.jl/pull/282) - Added `wait` to writetask in `writeblock!` [#281](https://github.com/JuliaIO/Zarr.jl/pull/281) - Fix getattrs for v3 [#277](https://github.com/JuliaIO/Zarr.jl/pull/277) diff --git a/src/Compressors/Compressors.jl b/src/Compressors/Compressors.jl index 60b52558..d2f631e2 100644 --- a/src/Compressors/Compressors.jl +++ b/src/Compressors/Compressors.jl @@ -29,10 +29,11 @@ getCompressor(::Nothing) = NoCompressor() zcompress!(compressed,data,c,::Nothing) = zcompress!(compressed,data,c) zuncompress!(data,compressed,c,::Nothing) = zuncompress!(data,compressed,c) -# Fallback definition of mutating form of compress and uncompress -function zcompress!(compressed, data, c) - empty!(compressed) - append!(compressed,zcompress(data, c)) +# Bulk `resize!` + `copyto!` (not `append!`): avoids elementwise growth over `NoCompressor`'s view. +function zcompress!(compressed, data, c) + src = zcompress(data, c) + resize!(compressed, length(src)) + copyto!(compressed, src) end zuncompress!(data, compressed, c) = copyto!(data, zuncompress(compressed, c, eltype(data))) @@ -72,6 +73,18 @@ function zcompress(a, ::NoCompressor) _reinterpret(UInt8,a) end +# Fast path: bulk `unsafe_copyto!` avoids the elementwise `ReinterpretArray` copy of the fallback. +function zuncompress!(data::Array{T}, compressed::Vector{UInt8}, ::NoCompressor) where {T} + isbitstype(T) || return copyto!(data, _reinterpret(T, compressed)) + n = sizeof(data) + n == length(compressed) || throw(DimensionMismatch( + "Encoded byte length $(length(compressed)) does not match output byte size $n" + )) + GC.@preserve data compressed unsafe_copyto!(Ptr{UInt8}(pointer(data)), + pointer(compressed), n) + return data +end + JSON.lower(::NoCompressor) = nothing compressortypes[nothing] = NoCompressor diff --git a/src/ZArray.jl b/src/ZArray.jl index b1aeb4e1..857b9a6b 100644 --- a/src/ZArray.jl +++ b/src/ZArray.jl @@ -160,21 +160,95 @@ _zero(::Type{<:Vector{T}}) where T = T[] _zero(::Type{Char}) = Char(0) getchunkarray(z::ZArray) = fill(_zero(eltype(z)), z.metadata.chunks) +# Same as `getchunkarray` but skips the zero/fill_value-fill. Use only when +# the caller guarantees the buffer will be fully overwritten before any read +# (e.g. decode-into or full-chunk-write paths). +# +# Falls back to the standard `getchunkarray` for `>:Missing` element types: +# the codec pipeline requires the underlying buffer to be the `isbits` inner +# of a `SenMissArray`, not an `Array{Union{Missing,T}}` directly (Blosc and +# friends reject non-isbits eltypes). +function getchunkarray_undef(z::ZArray{T}) where {T} + Missing <: T && return getchunkarray(z) + return Array{T}(undef, z.metadata.chunks) +end + maybeinner(a::Array) = a maybeinner(a::SenMissArray) = a.x resetbuffer!(fv,a::Array) = fv === nothing || fill!(a,fv) resetbuffer!(_,a::SenMissArray) = fill!(a,missing) +# Returns the chunk index when the call qualifies for the single-chunk fast +# path: a plain `Array{T,N}` matching the chunk shape, `Missing <: T` false, +# and only one chunk touched. `size(arr) == chunks` with `length(blockr) == 1` +# implies the slice aligns with that chunk's full range, so no extra +# coverage check is needed. +function singlechunk_fastpath(arr, z::ZArray{T,N}, blockr::CartesianIndices{N}) where {T,N} + arr isa Array{T,N} && !(Missing <: T) && + size(arr) == z.metadata.chunks && length(blockr) == 1 || return nothing + return first(blockr) +end + +# Single-chunk full-overwrite write. Encodes `ain` and pushes the chunk to +# the store. Split out of `writeblock!` so we can specialize on metadata +# type below. +function write_singlechunk_fastpath!(z::ZArray, ain::Array, bI::CartesianIndex) + data_encoded = compress_raw(ain, z) + if data_encoded === nothing + if store_isinitialized(z.storage, z.path, bI, z.metadata.chunk_key_encoding) + store_deletechunk(z.storage, z.path, bI, z.metadata.chunk_key_encoding) + end + else + store_writechunk(z.storage, data_encoded, z.path, bI, z.metadata.chunk_key_encoding) + end + return nothing +end + +# Zero-copy write for V2 with no compressor and no filters: the on-disk +# bytes are exactly `ain`'s in-memory representation (Zarr V2 default +# little-endian matches Julia on little-endian hosts — i.e. all currently +# supported platforms), so pass a `reinterpret(UInt8, ain)` view straight +# to the store and skip the chunk-sized `Vector{UInt8}` allocation + +# memcpy that the generic pipeline performs. +# +# Safe to hand out a view (not an owned copy) because the singlechunk +# fastpath guarantees `ain` is the caller's input array, not the shared +# scratch buffer the multi-chunk path reuses across iterations. +function write_singlechunk_fastpath!( + z::ZArray{T,N,<:AbstractStore,<:MetadataV2{T,N,NoCompressor,Nothing}}, + ain::Array{T,N}, bI::CartesianIndex, +) where {T,N} + fv = z.metadata.fill_value + if fv !== nothing && all(isequal(fv), ain) + if store_isinitialized(z.storage, z.path, bI, z.metadata.chunk_key_encoding) + store_deletechunk(z.storage, z.path, bI, z.metadata.chunk_key_encoding) + end + return nothing + end + store_writechunk(z.storage, _reinterpret(UInt8, ain), + z.path, bI, z.metadata.chunk_key_encoding) + return nothing +end + # 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) + # Fast path: single-chunk full-read decodes directly into `aout`, skipping the readtask channel and scratch buffer. + bI = singlechunk_fastpath(aout, z, blockr) + if bI !== nothing + chunk_compressed = store_readchunk(z.storage, z.path, bI, z.metadata.chunk_key_encoding) + uncompress_raw!(aout, z, chunk_compressed) + return aout + end + # Allocate the chunk-shaped scratch buffer. Reads always either fill it + # from a decode (which writes every element) or fall through to the + # fill-value path (which calls `fill!` itself), so we don't need to + # pre-zero it. + a = getchunkarray_undef(z) # Now loop through the chunks c = Channel{Pair{eltype(blockr),Union{Nothing,Vector{UInt8}}}}(channelsize(z.storage)) @@ -203,14 +277,22 @@ function readblock!(aout::AbstractArray{<:Any,N}, z::ZArray{<:Any, N}, r::Cartes end function writeblock!(ain::AbstractArray{<:Any,N}, z::ZArray{<:Any, N}, r::CartesianIndices{N}) where {N} - + z.writeable || error("Can not write to read-only ZArray") input_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) + # Fast path: single-chunk full-overwrite skips the readtask/writetask channels and the scratch buffer. + bI = singlechunk_fastpath(ain, z, blockr) + if bI !== nothing + write_singlechunk_fastpath!(z, ain, bI) + return ain + end + # If `fill_value === nothing`, the legacy behaviour is that un-written + # cells in a partial-write to a new chunk default to zero (via the + # zero-fill of `getchunkarray`). Preserve that. Otherwise `resetbuffer!` + # handles initialisation explicitly and we save the dead memset. + a = z.metadata.fill_value === nothing ? getchunkarray(z) : getchunkarray_undef(z) # Now loop through the chunks readchannel = Channel{Pair{eltype(blockr),Union{Nothing,Vector{UInt8}}}}(channelsize(z.storage)) diff --git a/test/runtests.jl b/test/runtests.jl index d1ee7c0c..6bb66083 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -83,6 +83,93 @@ using Dates GC.gc() end end + + @testset "NoCompressor single-chunk zero-copy fastpath" begin + # The V2+NoCompressor+no-filters specialization of + # `write_singlechunk_fastpath!` hands a `reinterpret(UInt8, ain)` view + # straight to the store instead of allocating a chunk-sized + # `Vector{UInt8}` and memcpy'ing into it. The on-disk bytes must + # still match the array's bit pattern exactly. + + # Dispatch: the specialized method is selected for V2 + NoCompressor + no filters. + mktempdir(@__DIR__) do dir + z3 = zcreate(Float32, 4, 4, 2; path=joinpath(dir, "disp"), + chunks=(4, 4, 2), compressor=Zarr.NoCompressor()) + ain = rand(Float32, 4, 4, 2) + m = which(Zarr.write_singlechunk_fastpath!, (typeof(z3), typeof(ain), CartesianIndex{3})) + @test occursin("MetadataV2", string(m.sig)) + @test occursin("NoCompressor", string(m.sig)) + GC.gc() + end + + # Round-trip across rank 0, 1, 3, 4 and a mix of bitstypes. + mktempdir(@__DIR__) do dir + for (T, shape) in ((Float32, (8, 8, 4)), + (Int64, (5, 3)), + (UInt16, (16,)), + (ComplexF32, (4, 4)), + (Float64, ())) + name = "rt_$(T)_$(length(shape))d" + # Match chunks to full shape so writes hit the single-chunk fastpath. + chunks = shape == () ? () : shape + z = zcreate(T, shape...; path=joinpath(dir, name), + chunks=chunks, compressor=Zarr.NoCompressor()) + # Drive `ain` through a controlled byte pattern so round-trip + # failure can't be masked by zero-init. + ain = T === ComplexF32 ? + reshape(ComplexF32[ComplexF32(i, -i) for i in 1:prod(shape == () ? (1,) : shape)], + shape == () ? () : shape) : + reshape(T[T(i) for i in 1:prod(shape == () ? (1,) : shape)], + shape == () ? () : shape) + if shape == () + z[] = ain[] + else + z[(Colon() for _ in shape)...] = ain + end + # On-disk chunk file equals reinterpret(UInt8, ain). + chunk_path = shape == () ? joinpath(dir, name, "0") : + joinpath(dir, name, join(fill("0", length(shape)), ".")) + @test isfile(chunk_path) + on_disk = read(chunk_path) + expected = Vector{UInt8}(undef, sizeof(ain)) + GC.@preserve ain expected unsafe_copyto!(pointer(expected), + Ptr{UInt8}(pointer(ain)), sizeof(ain)) + @test on_disk == expected + # Reads round-trip via the regular Zarr.jl path. + @test Array(zopen(joinpath(dir, name))) == ain + GC.gc() + end + end + + # All-fill-value chunks are elided (no on-disk file written) — same + # semantics as the generic V2 pipeline_encode. + mktempdir(@__DIR__) do dir + z = zcreate(Float32, 4, 4; path=joinpath(dir, "elide"), + chunks=(4, 4), compressor=Zarr.NoCompressor(), + fill_value=Float32(7)) + z[:, :] = fill(Float32(7), 4, 4) + @test !ispath(joinpath(dir, "elide", "0.0")) + # Overwriting an existing chunk with all-fill removes the file. + z[:, :] = rand(Float32, 4, 4) + @test ispath(joinpath(dir, "elide", "0.0")) + z[:, :] = fill(Float32(7), 4, 4) + @test !ispath(joinpath(dir, "elide", "0.0")) + GC.gc() + end + + # The store must fully consume the reinterpret view before returning: + # mutating `ain` after the write must not affect on-disk content. + mktempdir(@__DIR__) do dir + z = zcreate(Float32, 8, 8; path=joinpath(dir, "alias"), + chunks=(8, 8), compressor=Zarr.NoCompressor()) + ain = rand(Float32, 8, 8) + snapshot = copy(ain) + z[:, :] = ain + fill!(ain, Float32(NaN)) # clobber after the write returns + @test zopen(joinpath(dir, "alias"))[:, :] == snapshot + GC.gc() + end + end end @testset "Groups" begin