From e8dc10c1649c395a41589c6599450f3b28067527 Mon Sep 17 00:00:00 2001 From: Beforerr Date: Sun, 30 Aug 2026 11:23:04 +0900 Subject: [PATCH] perf: improve TTFX by type-erasing data movement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit What changed 1. Type-erased data movement (src/loading/variable.jl). readblock! is now a thin wrapper; everything below it takes Ptr{UInt8} + element size + Vector{Int} dims/strides: - _read_records! — record range → bytes, threaded decompression compiled once per file format instead of per shape. - _copy_block! — one strided N-d gather kernel serving both the row→column majority swap and sub-block extraction. - Big-endian byte swap keyed on word size, not T. 2. Array(var), collect(var), and show read via readblock! directly — DiskArrays' generic indexing costs ~60 ms compile per shape. Base.summary(::CDFVariable) is now cheap (Base's was ~190 ms/type). 3. Workload (src/precompile.jl): seeds one variable per common shape (floats ≤3-D, TT2000/Epoch 1-D, strings ≤2-D) with the ops people actually run — show, attribs, Array, var[1], record slice, DateTime.. Rare shapes skipped. refactor: VVREntry records per-leaf whether it's a CVVR The old code took the last leaf's type for the whole variable, which is wrong when a compressed variable stores a poorly-compressing block as a plain VVR. --- src/dataset.jl | 2 +- src/loading/variable.jl | 293 ++++++++++++++++++++++++---------------- src/parsing.jl | 21 ++- src/precompile.jl | 28 ++-- src/records/vdr.jl | 7 +- src/records/vvr.jl | 15 -- src/variable.jl | 50 ++++--- test/Project.toml | 1 + test/runtests.jl | 24 +++- 9 files changed, 275 insertions(+), 166 deletions(-) diff --git a/src/dataset.jl b/src/dataset.jl index da9a9d4..5ee1c80 100644 --- a/src/dataset.jl +++ b/src/dataset.jl @@ -145,7 +145,7 @@ function Base.show(io::IO, m::MIME"text/plain", cdf::CDFDataset) print(io, " ", var.name, " : ", size(var), " ") printstyled(io, variable_type(var); bold = true) print(io, " ", CDFDataType(var.vdr.data_type)) - !isempty(var) && print(io, " [", var[1], " … ", var[end], "]") + !isempty(var) && print(io, " [", first(_record(var, 1)), " … ", last(_record(var, size(var)[end])), "]") println(io) end println(io, cdf.cdr) diff --git a/src/loading/variable.jl b/src/loading/variable.jl index 97d596c..d50ab89 100644 --- a/src/loading/variable.jl +++ b/src/loading/variable.jl @@ -6,7 +6,7 @@ return chnl end else - const _decompressors = Ref{Union{Channel{Decompressor}, Nothing}}(nothing) + const _decompressors = Ref{Union{Channel{Decompressor},Nothing}}(nothing) function decompressors() if _decompressors[] === nothing n_ch = nthreads() @@ -16,26 +16,104 @@ else end return _decompressors[] end + __init__() = _decompressors[] = nothing end -""" - majority_swap!(data, dims_without_record) - -Convert row-major data layout to column-major (Julia's native layout) in-place. -For row-major CDF files, this reverses the dimension order (except record dimension). -""" -function majority_swap!(data::AbstractArray{T, N}, dims_without_record) where {T, N} - N <= 2 && return data - perm = ntuple(i -> N - i, N - 1) - reversed_dims = reverse(dims_without_record) - temp = similar(data, dims_without_record) - for slc in eachslice(data; dims = N) - permutedims!(temp, reshape(slc, reversed_dims), perm) - copyto!(slc, temp) +# ---- Type-erased block moves: dims/strides are runtime vectors ---- +# Element movers, chosen once per block by element size; `ByteMove` covers odd sizes (CHAR data). +struct WordMove{U} end +struct ByteMove end +@inline (::WordMove{U})(d::Ptr{UInt8}, s::Ptr{UInt8}, ::Int) where {U} = unsafe_store!(Ptr{U}(d), unsafe_load(Ptr{U}(s))) +@inline (::ByteMove)(d::Ptr{UInt8}, s::Ptr{UInt8}, esz::Int) = unsafe_copyto!(d, s, esz) + +# Gather a column-major block of `dims` elements (`esz` bytes each) whose source byte +# strides are `sstrides` into contiguous `dst`. +function _copy_block!(dst::Ptr{UInt8}, src::Ptr{UInt8}, esz::Int, dims::Vector{Int}, sstrides::Vector{Int}) + esz == 8 && return _copy_block!(dst, src, esz, dims, sstrides, WordMove{UInt64}()) + esz == 4 && return _copy_block!(dst, src, esz, dims, sstrides, WordMove{UInt32}()) + esz == 2 && return _copy_block!(dst, src, esz, dims, sstrides, WordMove{UInt16}()) + esz == 1 && return _copy_block!(dst, src, esz, dims, sstrides, WordMove{UInt8}()) + esz == 16 && return _copy_block!(dst, src, esz, dims, sstrides, WordMove{UInt128}()) + return _copy_block!(dst, src, esz, dims, sstrides, ByteMove()) +end + +function _copy_block!(dst::Ptr{UInt8}, src::Ptr{UInt8}, esz::Int, dims::Vector{Int}, sstrides::Vector{Int}, move) + n = length(dims) + n0 = dims[1] + s0 = sstrides[1] + run_bytes = n0 * esz + memcpy_run = s0 == esz && run_bytes > 64 # a memmove call per short run costs more than it saves + nruns = 1 + for k in 2:n + nruns *= dims[k] end - return data + idx = zeros(Int, n) + soff = 0 + doff = 0 + for _ in 1:nruns + if memcpy_run + unsafe_copyto!(dst + doff, src + soff, run_bytes) + else + for i in 0:(n0-1) + move(dst + doff + i * esz, src + soff + i * s0, esz) + end + end + doff += run_bytes + k = 2 + while k <= n + idx[k] += 1 + soff += sstrides[k] + idx[k] < dims[k] && break + soff -= dims[k] * sstrides[k] + idx[k] = 0 + k += 1 + end + end + return end +# Row-major records → column-major in place. Each record of `rdims` is stored in +# `reverse(rdims)` order; swapped through a bounded temp, a chunk of records at a time. +function _majority_swap!(data::Ptr{UInt8}, esz::Int, rdims::Vector{Int}, nrec::Int) + k = length(rdims) + k <= 1 && return + record_bytes = prod(rdims) * esz + sstrides = Vector{Int}(undef, k + 1) + s = esz + for i in k:-1:1 + sstrides[i] = s + s *= rdims[i] + end + sstrides[k+1] = record_bytes + chunk = clamp((1 << 20) ÷ record_bytes, 1, nrec) + dims = [rdims; chunk] + temp = Vector{UInt8}(undef, chunk * record_bytes) + GC.@preserve temp for r0 in 0:chunk:(nrec-1) + n = min(chunk, nrec - r0) + dims[end] = n + base = data + r0 * record_bytes + _copy_block!(pointer(temp), base, esz, dims, sstrides) + unsafe_copyto!(base, pointer(temp), n * record_bytes) + end + return +end + +# Extract `ranges` of the record dims (all records) from contiguous column-major `src`. +function _copy_subblock!(dst::Ptr{UInt8}, src::Ptr{UInt8}, esz::Int, rdims::Vector{Int}, ranges::Vector{UnitRange{Int}}, nrec::Int) + k = length(rdims) + sstrides = Vector{Int}(undef, k + 1) + s = esz + off = 0 + for i in 1:k + sstrides[i] = s + off += (first(ranges[i]) - 1) * s + s *= rdims[i] + end + sstrides[k+1] = s + dims = [length.(ranges); nrec] + _copy_block!(dst, src + off, esz, dims, sstrides) + return +end function variable(cdf::CDFDataset, name) vdr = find_vdr(cdf, name) @@ -48,7 +126,7 @@ function _variable(cdf, name, vdr) M = num_record_dims(vdr, cdf) return Base.Cartesian.@nif 12 d -> (M == d - 1) d -> ( d == 12 ? throw(ArgumentError("variable has $M dimensions; the CDF format allows at most 10")) : - _variable(cdf, name, vdr, Val(d - 1)) + _variable(cdf, name, vdr, Val(d - 1)) ) end @@ -56,8 +134,8 @@ function _variable(cdf, name, vdr, ::Val{M}) where {M} dims = (map(Int, record_sizes(vdr, cdf, Val(M)))..., Int(vdr.max_rec) + 1) code = Int(vdr.data_type) if code == 51 || code == 52 # CHAR/UCHAR: eltype depends on runtime num_elems - T = StaticString{Int(vdr.num_elems), UInt8} - return CDFVariable{T, M + 1, typeof(vdr), typeof(cdf)}(name, vdr, cdf, dims) + T = StaticString{Int(vdr.num_elems),UInt8} + return CDFVariable{T,M + 1,typeof(vdr),typeof(cdf)}(name, vdr, cdf, dims) end # Branch to static constructor per element type return Base.Cartesian.@nif( @@ -68,132 +146,113 @@ function _variable(cdf, name, vdr, ::Val{M}) where {M} ) end -@inline _construct(cdf, name, vdr, dims::NTuple{N, Int}, ::Type{T}) where {N, T} = - CDFVariable{T, N, typeof(vdr), typeof(cdf)}(name, vdr, cdf, dims) - -function _read_full!(dest::AbstractArray{T, N}, ds, name, vdr) where {T, N} - Base.require_one_based_indexing(dest) - Tfile = julia_type(vdr.data_type, vdr.num_elems) - T === Tfile || throw(ArgumentError("element type mismatch for \"$name\": file has $Tfile, destination has $T")) - dims = (map(Int, record_sizes(vdr, ds, Val(N - 1)))..., Int(vdr.max_rec) + 1) - size(dest) == dims || throw(DimensionMismatch("variable \"$name\" has size $dims, destination has size $(size(dest))")) - var = CDFVariable{T, N, typeof(vdr), typeof(ds)}(name, vdr, ds, dims) - DiskArrays.readblock!(var, dest, axes(dest)...) - return dest -end - -@inline _record_view(A::AbstractArray{<:Any, M}, r) where {M} = - view(A, ntuple(_ -> Colon(), M - 1)..., r) +@inline _construct(cdf, name, vdr, dims::NTuple{N,Int}, ::Type{T}) where {N,T} = + CDFVariable{T,N,typeof(vdr),typeof(cdf)}(name, vdr, cdf, dims) -function DiskArrays.readblock!(var::CDFVariable{T, N}, dest::AbstractArray{T}, ranges::Vararg{AbstractUnitRange{<:Integer}, N}; nbuffers = nthreads()) where {T, N} +function DiskArrays.readblock!(var::CDFVariable{T,N}, dest::AbstractArray{T}, ranges::Vararg{AbstractUnitRange{<:Integer},N}) where {T,N} N > 0 && @boundscheck checkbounds(var, ranges...) isempty(dest) && return dest + if !(dest isa Array) + copyto!(dest, DiskArrays.readblock!(var, Array{T}(undef, size(dest)), ranges...)) + return dest + end + rdims = Int[Base.front(var.dims)...] + other_ranges = UnitRange{Int}[Int(first(r)):Int(last(r)) for r in Base.front(ranges)] + rec_range = Int(first(ranges[end])):Int(last(ranges[end])) + GC.@preserve dest _readblock!(Ptr{UInt8}(pointer(dest)), sizeof(T), rdims, other_ranges, rec_range, var.vdr, var.parentdataset) + is_big_endian_encoding(var) && _byte_swap!(dest) + return dest +end - buffer = parent(var.parentdataset) - RecordSizeType = recordsize_type(var.parentdataset) - entries, vvr_type = read_vvrs(var.vdr, var.parentdataset) - isempty(entries) && return dest - compression = if !isempty(entries) # # vvr records is the ultimative source - vvr_type == VVR_ ? NoCompression : variable_compression(var.vdr, var.parentdataset) +function _readblock!(dest::Ptr{UInt8}, esz::Int, rdims::Vector{Int}, other_ranges::Vector{UnitRange{Int}}, rec_range::UnitRange{Int}, vdr, ds) + entries = read_vvrs(vdr, ds) + isempty(entries) && return + compression = any(e -> e.compressed, entries) ? variable_compression(vdr, ds) : NoCompression + buffer = parent(ds) + RST = recordsize_type(ds) + swap = majority(ds) == Row + nrec = length(rec_range) + record_bytes = prod(rdims) * esz + if length.(other_ranges) == rdims + dest_vec = unsafe_wrap(Vector{UInt8}, dest, nrec * record_bytes) + _read_records!(dest_vec, buffer, entries, first(rec_range), last(rec_range), record_bytes, compression, RST) + swap && _majority_swap!(dest, esz, rdims, nrec) else - NoCompression + scratch = Vector{UInt8}(undef, nrec * record_bytes) + GC.@preserve scratch begin + _read_records!(scratch, buffer, entries, first(rec_range), last(rec_range), record_bytes, compression, RST) + swap && _majority_swap!(pointer(scratch), esz, rdims, nrec) + _copy_subblock!(dest, pointer(scratch), esz, rdims, other_ranges, nrec) + end end + return +end + +# Whole-variable reads skip DiskArrays' indexing layer, which is costly to compile +function Base.Array(var::CDFVariable{T,N}) where {T,N} + return DiskArrays.readblock!(var, Array{T,N}(undef, size(var)), axes(var)...) +end - record_range = ranges[end] - other_ranges = ranges[1:(N - 1)] - dims_without_record = var.dims[1:(N - 1)] - - is_full_record = length.(other_ranges) == dims_without_record - is_no_compression = compression == NoCompression - - first_rec = first(record_range) - last_rec = last(record_range) - start_idx = findfirst(entry -> entry.first <= first_rec <= entry.last, entries)::Int - end_idx = findfirst(entry -> entry.first <= last_rec <= entry.last, entries)::Int - record_size = prod(dims_without_record) - - is_row_major = majority(var) == Row - # If the variable is not compressed and the other dimension ranges are the same as the variable range - # we can directly read the data into dest - if is_no_compression && is_full_record - record_bytes = record_size * sizeof(T) - doffs = 1 - header_skip = sizeof(RecordSizeType) + sizeof(Int32) +function _record(var::CDFVariable{T,N}, r::Int) where {T,N} + rdims = Base.front(var.dims) + dest = Array{T,N}(undef, rdims..., 1) + return DiskArrays.readblock!(var, dest, ntuple(i -> 1:rdims[i], N - 1)..., r:r) +end + +# Fill `dest` (exactly `nrec * record_bytes` bytes) with records `rec_first:rec_last`. +# Element-type agnostic so the (threaded) decompression loop compiles once. +function _read_records!(dest::Vector{UInt8}, buffer::Vector{UInt8}, entries::Vector{VVREntry}, rec_first::Int, rec_last::Int, record_bytes::Int, compression::CompressionType, ::Type{RST}) where {RST} + start_idx = findfirst(e -> e.first <= rec_first <= e.last, entries)::Int + end_idx = findfirst(e -> e.first <= rec_last <= e.last, entries)::Int + if compression == NoCompression for i in start_idx:end_idx - entry = entries[i] - overlap_first = max(first_rec, entry.first) - overlap_last = min(last_rec, entry.last) - N_elems = (overlap_last - overlap_first + 1) * record_size - data_start = entry.offset + 1 + header_skip - byte_offset = data_start + (overlap_first - entry.first) * record_bytes - _copy_to!(dest, doffs, buffer, byte_offset, N_elems) - doffs += N_elems + _read_entry!(dest, buffer, entries[i], rec_first, rec_last, record_bytes, compression, RST, nothing) end - @assert doffs == length(dest) + 1 - is_row_major && majority_swap!(dest, dims_without_record) else - Base.@inbounds Threads.@threads for i in start_idx:end_idx - entry = entries[i] - if is_full_record && entry.first >= first_rec && entry.last <= last_rec - # full entry - dest_range = dst_src_ranges(first_rec, last_rec, entry)[1] - dest_view = _record_view(dest, dest_range) - total_elems = record_size * length(entry) - decompressor = take!(decompressors()) - try - load_cvvr_data!(dest_view, 1, buffer, entry.offset, total_elems, RecordSizeType, compression; decompressor) - is_row_major && majority_swap!(dest_view, dims_without_record) - finally - put!(decompressors(), decompressor) - end - else - # partial entry - (dest_range, local_range) = dst_src_ranges(first_rec, last_rec, entry) - dest_view = _record_view(dest, dest_range) - n_records = length(entry) - total_elems = record_size * n_records - chunk = Vector{T}(undef, total_elems) - - if is_no_compression - load_vvr_data!(chunk, 1, buffer, entry.offset, total_elems, RecordSizeType) - else - decompressor = take!(decompressors()) - try - load_cvvr_data!(chunk, 1, buffer, entry.offset, total_elems, RecordSizeType, compression; decompressor) - finally - put!(decompressors(), decompressor) - end - end - - # chunk_data = _load_entry_chunk(var, entry, RecordSizeType, buffer, var.compression; decompressor) - chunk_array = reshape(chunk, dims_without_record..., :) - is_row_major && majority_swap!(chunk_array, dims_without_record) - src_view = view(chunk_array, other_ranges..., local_range) - dest_view .= src_view + Threads.@threads for i in start_idx:end_idx + decompressor = take!(decompressors()) + try + _read_entry!(dest, buffer, entries[i], rec_first, rec_last, record_bytes, compression, RST, decompressor) + finally + put!(decompressors(), decompressor) end end end - is_big_endian_encoding(var) && _byte_swap!(dest) return dest end +function _read_entry!(dest, buffer, e::VVREntry, rec_first, rec_last, record_bytes, compression, ::Type{RST}, decompressor) where {RST} + lo, hi = max(rec_first, e.first), min(rec_last, e.last) + doffs = (lo - rec_first) * record_bytes + 1 + nbytes = (hi - lo + 1) * record_bytes + skip = (lo - e.first) * record_bytes + if !e.compressed + load_vvr_data!(dest, doffs, buffer, e.offset + skip, nbytes, RST) + elseif skip == 0 && hi == e.last + load_cvvr_data!(dest, doffs, buffer, e.offset, nbytes, RST, compression; decompressor) + else + scratch = Vector{UInt8}(undef, length(e) * record_bytes) + load_cvvr_data!(scratch, 1, buffer, e.offset, length(scratch), RST, compression; decompressor) + copyto!(dest, doffs, scratch, skip + 1, nbytes) + end + return +end + function collect_vxr_entries!(entries::Vector{VVREntry}, src, offset, ::Type{FieldSizeT}) where {FieldSizeT} - vvr_type = nothing while offset != 0 vxr = VXR{FieldSizeT}(src, offset) for (first, last, leaf_offset) in vxr record_type = get_record_type(src, leaf_offset, FieldSizeT) @assert record_type in (VVR_, CVVR_, VXR_) if record_type == VXR_ - vvr_type = collect_vxr_entries!(entries, src, leaf_offset, FieldSizeT) + collect_vxr_entries!(entries, src, leaf_offset, FieldSizeT) else - push!(entries, VVREntry(Int(first) + 1, Int(last) + 1, leaf_offset)) - vvr_type = record_type + push!(entries, VVREntry(Int(first) + 1, Int(last) + 1, leaf_offset, record_type == CVVR_)) end end offset = vxr.vxr_next end - return vvr_type + return entries end function variable_compression(vdr::AbstractVDR{FieldSizeT}, cdf) where {FieldSizeT} diff --git a/src/parsing.jl b/src/parsing.jl index 4cecf6d..b21aaae 100644 --- a/src/parsing.jl +++ b/src/parsing.jl @@ -106,5 +106,22 @@ function validate_cdf_magic(magic_bytes) return magic_bytes in cdf_magic_bytes end -_byte_swap!(data) = map!(ntoh, data, data) -_byte_swap!(data::AbstractArray{<:StaticString{N}}) where {N} = data +# Big-endian files: swap each `unit`-byte word in place (0 = raw bytes, no swap). +_swap_unit(::Type{T}) where {T} = sizeof(T) +_swap_unit(::Type{<:StaticString}) = 0 +_swap_unit(::Type{Epoch16}) = 8 # two Float64s +function _byte_swap!(p::Ptr{UInt8}, nbytes::Int, unit::Int) + unit <= 1 && return + unit == 2 && return _swap_words!(Ptr{UInt16}(p), nbytes ÷ 2) + unit == 4 && return _swap_words!(Ptr{UInt32}(p), nbytes ÷ 4) + unit == 8 && return _swap_words!(Ptr{UInt64}(p), nbytes ÷ 8) + throw(ArgumentError("unsupported byte-swap unit $unit")) +end +function _swap_words!(p::Ptr{U}, n::Int) where {U} + for i in 1:n + unsafe_store!(p, ntoh(unsafe_load(p, i)), i) + end + return +end +_byte_swap!(data::Array{T}) where {T} = + GC.@preserve data _byte_swap!(Ptr{UInt8}(pointer(data)), sizeof(data), _swap_unit(T)) diff --git a/src/precompile.jl b/src/precompile.jl index c33ce89..bf007a8 100644 --- a/src/precompile.jl +++ b/src/precompile.jl @@ -1,16 +1,26 @@ -precompile(Array, (CDFVariable{TT2000, 1, VDR{Int64}, CDFDataset{Int64}},)) -for T in (Float32, Float64), i in 1:3 - precompile(Array, (CDFVariable{T, i, VDR{Int64}, CDFDataset{Int64}},)) -end +# Seed one variable per common (eltype, ndims) shape; every new pair costs DiskArrays +# indexing compiles, so rarer shapes (integers, Epoch16, >3-D, extra string widths) are left out. +_workload_key(var) = (eltype(var) <: StaticString ? StaticString : eltype(var), ndims(var)) +_workload_maxdims(T) = T <: AbstractFloat ? 3 : T <: StaticString ? 2 : T in (TT2000, Epoch) ? 1 : 0 PrecompileTools.@setup_workload begin - elx_file = joinpath(@__DIR__, "../data/elb_l2_epdef_20210914_v01.cdf") + files = joinpath.(@__DIR__, "../data", ("elb_l2_epdef_20210914_v01.cdf", "a_cdf_with_compressed_vars.cdf")) + seen = Set{Tuple{Any, Int}}() PrecompileTools.@compile_workload begin - ds = CDFDataset(elx_file) - for name in keys(ds) - var = ds[name] - isempty(var) || var[1] + for (i, file) in enumerate(files) + ds = CDFDataset(file) + ds.attrib + for var in ds + key = _workload_key(var) + (isempty(var) || key[2] > _workload_maxdims(key[1]) || key in seen) && continue + push!(seen, key) + get(var.attrib, "FILLVAL", nothing) + Array(var) + var[1] + var[ntuple(_ -> Colon(), ndims(var) - 1)..., 1:1] + eltype(var) <: CDFDateTime && DateTime.(var[1:1]) + end end end end diff --git a/src/records/vdr.jl b/src/records/vdr.jl index 36f8652..9a98e3b 100644 --- a/src/records/vdr.jl +++ b/src/records/vdr.jl @@ -118,13 +118,10 @@ end # 2 Whether or not a compression method might be applied to this variable data. Set indicates that a compression is chosen by the user and the data might be compressed, depending on the data size and content. If the compressed data becomes larger than its uncompressed data, no compression is applied and the data are stored as uncompressed, even the compression bit is set. The compressed data is stored in Compressed Variable Value Record (CVVR) while uncompressed data go into Variable Value Record (VVR). Clear indicates that a compression will not be used. The CPRorSPRoffset field provides the offset of the Compressed Parameters Record if this compression bit is set and the compression used. function read_vvrs(vdr::AbstractVDR{FieldSizeT}, cdf) where {FieldSizeT} - vxr_head = vdr.vxr_head entries = Vector{VVREntry}() - src = parent(cdf) sizehint!(entries, 1) - vvr_type = collect_vxr_entries!(entries, src, Int(vxr_head), FieldSizeT) - vvr_type = @something vvr_type VVR_ - return entries, vvr_type + collect_vxr_entries!(entries, parent(cdf), Int(vdr.vxr_head), FieldSizeT) + return entries end is_record_varying(vdr) = !is_nrv(vdr) diff --git a/src/records/vvr.jl b/src/records/vvr.jl index 82bf0a7..a00cbd1 100644 --- a/src/records/vvr.jl +++ b/src/records/vvr.jl @@ -1,18 +1,3 @@ -""" -Variable Value Record (VVR) - contains actual variable data -""" -struct VVR{T} - header::Header - data::Vector{T} # Raw variable data -end - -@inline function VVR(buffer::Vector{UInt8}, offset, ::Type{RecordSizeType}, data) where {RecordSizeType} - pos = offset + 1 - header = Header(buffer, pos, RecordSizeType) - @assert header.record_type == 7 "Invalid VVR record type" - return VVR(header, data) -end - function _copy_to!(dest, doffs, src, soffs, N) T = eltype(dest) GC.@preserve dest src begin diff --git a/src/variable.jl b/src/variable.jl index bb6f05a..3bc1b7f 100644 --- a/src/variable.jl +++ b/src/variable.jl @@ -1,18 +1,19 @@ -abstract type AbstractVariable{T, N} <: DiskArrays.AbstractDiskArray{T, N} end +abstract type AbstractVariable{T,N} <: DiskArrays.AbstractDiskArray{T,N} end struct VVREntry first::Int last::Int offset::Int + compressed::Bool # leaf is a CVVR (a compressed variable may still store poorly-compressing blocks as plain VVRs) end Base.length(entry::VVREntry) = entry.last - entry.first + 1 -struct CDFVariable{T, N, V, P} <: AbstractVariable{T, N} +struct CDFVariable{T,N,V,P} <: AbstractVariable{T,N} name::String vdr::V parentdataset::P - dims::NTuple{N, Int} + dims::NTuple{N,Int} end Base.size(var::CDFVariable) = var.dims @@ -20,16 +21,6 @@ Base.size(var::CDFVariable) = var.dims @inline majority(var::CDFVariable) = majority(var.parentdataset) @inline is_big_endian_encoding(var::CDFVariable) = is_big_endian_encoding(var.parentdataset) -function dst_src_ranges(first, last, entry) - overlap_first = max(first, entry.first) - overlap_last = min(last, entry.last) - local_first = overlap_first - entry.first + 1 - local_last = overlap_last - entry.first + 1 - dest_first = overlap_first - first + 1 - dest_last = overlap_last - first + 1 - return (dest_first:dest_last, local_first:local_last) -end - # Codes seem to be faster if we disable chunking DiskArrays.haschunks(::CDFVariable) = DiskArrays.Unchunked() # DiskArrays.haschunks(::CDFVariable) = DiskArrays.Chunked() @@ -46,7 +37,7 @@ function _eachchunk(var::CDFVariable) end function _eachchunk_vvrs(var::CDFVariable) - vvrs, _ = read_vvrs(var.vdr, var.parentdataset) + vvrs = read_vvrs(var.vdr, var.parentdataset) N = ndims(var) chunks = ntuple(N) do i if i != N @@ -54,9 +45,9 @@ function _eachchunk_vvrs(var::CDFVariable) else chunksizes = length.(vvrs) if length(vvrs) > 0 - chunksizes[end] = @views var.dims[N] - sum(chunksizes[1:(end - 1)]) + chunksizes[end] = @views var.dims[N] - sum(chunksizes[1:(end-1)]) end - DiskArrays.IrregularChunks(chunksizes = chunksizes) + DiskArrays.IrregularChunks(chunksizes=chunksizes) end end return DiskArrays.GridChunks(chunks) @@ -83,6 +74,33 @@ attrib(var::CDFVariable, name::String) = get(attrib(var), name) is_record_varying(v::CDFVariable) = is_record_varying(v.vdr) variable_type(v::CDFVariable) = get(v.attrib, "VAR_TYPE", "unknown") +# Base's summary/showarg specializes on the full type; keep it cheap per element type. +Base.summary(io::IO, var::CDFVariable{T}) where {T} = + print(io, Base.dims2string(size(var)), " CDFVariable{", T, "}") + +# DiskArrays' reductions read chunk-by-chunk through its generic indexing, +# which compiles per shape (~40-260 ms). +# With Unchunked chunking that already means one full read +for f in (:sum, :prod, :maximum, :minimum, :extrema, :all, :any, :count) + @eval Base.$f(var::CDFVariable; kw...) = Base.$f(Array(var); kw...) + @eval Base.$f(g::Function, var::CDFVariable; kw...) = Base.$f(g, Array(var); kw...) +end + +Base.collect(var::CDFVariable) = Array(var) + +function Base.iterate(var::CDFVariable) + isempty(var) && return nothing + A = Array(var) + x, st = iterate(A)::Tuple + return x, (A, st) +end +function Base.iterate(::CDFVariable, state::Tuple{Array,Int}) + A, st = state + r = iterate(A, st) + r === nothing && return nothing + return r[1], (A, r[2]) +end + function Base.show(io::IO, m::MIME"text/plain", var::CDFVariable) summary(io, var) println(io) diff --git a/test/Project.toml b/test/Project.toml index 2b76fbb..22e9537 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -4,6 +4,7 @@ CommonDataFormat = "a9737db6-c05c-4e48-868b-6bc41491d9d9" CommonDataModel = "1fbeeb36-5f17-413c-809b-666fb144f157" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" +JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" JuliaC = "acedd4c2-ced6-4a15-accc-2607eb759ba2" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/test/runtests.jl b/test/runtests.jl index 0b499d1..4a5db80 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -4,7 +4,9 @@ import CommonDataFormat as CDF include("utils.jl") include("epochs_test.jl") -include("comprehensive_test.jl") +@testset "Comprehensive Test" begin + include("comprehensive_test.jl") +end include("cdf2_test.jl") include("CommonDataModelExt_test.jl") include("decompress_test.jl") @@ -32,6 +34,26 @@ end include("trim.jl") end +# Reference the type-erased block kernels against permutedims/view on random shapes. +@testset "block copy kernels" begin + for esz in (1, 4, 8, 9), k in 1:4 + rdims = rand(1:5, k) + nrec = rand(1:4) + raw = rand(UInt8, prod(rdims) * esz * nrec) + stored = reshape(raw, esz, reverse(rdims)..., nrec) + expected = permutedims(stored, (1, (k + 1):-1:2..., k + 2)) + swapped = copy(raw) + GC.@preserve swapped CDF._majority_swap!(pointer(swapped), esz, rdims, nrec) + @test swapped == vec(expected) + + ranges = [let a = rand(1:d); a:rand(a:d) end for d in rdims] + src = reshape(raw, esz, rdims..., nrec) + dst = Vector{UInt8}(undef, esz * prod(length, ranges) * nrec) + GC.@preserve dst raw CDF._copy_subblock!(pointer(dst), pointer(raw), esz, rdims, ranges, nrec) + @test dst == vec(src[:, ranges..., :]) + end +end + @testset "Fill Value" begin for T in (Int8, Int16, Int32, Int64, Float32, Float64, UInt8, UInt16, UInt32) @test CDF.fillvalue(T) isa T