Performance improvements (v2 only for now)

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)

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

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))
 

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