perf: improve TTFX by type-erasing data movement

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)

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

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

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)

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

src/variable.jl

@@ -1,35 +1,26 @@
-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
 
 @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,17 +37,17 @@ 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
             DiskArrays.RegularChunks(var.dims[i], 0, var.dims[i])
         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)

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"

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