perf: O(n) dataset iteration

src/dataset.jl

@@ -68,6 +68,16 @@ majority(cdf::CDFDataset) = majority(cdf.cdr)
     throw(ArgumentError("Unknown property $name"))
 end
 
+# Load the (z or r) VDR at a known offset
+function _vdr_at(cdf::CDFDataset, offset::Int)
+    buffer = parent(cdf)
+    RecordSizeType = recordsize_type(cdf)
+    record_type = read_be(buffer, offset + 1 + sizeof(RecordSizeType), Int32)
+    @assert record_type in (8, 3)
+    return record_type == 8 ? VDR(buffer, offset, RecordSizeType) :
+        rVDR(buffer, offset, GDR(cdf), RecordSizeType)
+end
+
 function find_vdr(cdf::CDFDataset, var_name::String)
     gdr = GDR(cdf)
     RecordSizeType = recordsize_type(cdf)
@@ -77,13 +87,7 @@ function find_vdr(cdf::CDFDataset, var_name::String)
     for current_offset in (gdr.rVDRhead, gdr.zVDRhead)
         while current_offset != 0
             if readname(buffer, current_offset + vdr_name_offset) == var_name_bytes
-                record_type = read_be(buffer, current_offset + 1 + sizeof(RecordSizeType), Int32)
-                @assert record_type in (8, 3)
-                if record_type == 8
-                    return VDR(buffer, current_offset, RecordSizeType)
-                else
-                    return rVDR(buffer, current_offset, gdr, RecordSizeType)
-                end
+                return _vdr_at(cdf, Int(current_offset))
             end
             current_offset = read_be(buffer, current_offset + 5 + sizeof(RecordSizeType), RecordSizeType)
         end
@@ -96,7 +100,7 @@ function Base.getindex(cdf::CDFDataset, var_name::String)
     return variable(cdf, var_name)
 end
 
-Base.length(cdf::CDFDataset) = length(keys(cdf))
+Base.length(cdf::CDFDataset) = Int(GDR(cdf).NrVars + GDR(cdf).NzVars)
 
 function Base.keys(cdf::CDFDataset)
     RecordSizeType = recordsize_type(cdf)
@@ -116,17 +120,30 @@ function Base.keys(cdf::CDFDataset)
     return varnames
 end
 
-Base.haskey(cdf::CDFDataset, var_name::String) = var_name in keys(cdf)
+Base.haskey(cdf::CDFDataset, var_name::String) = !isnothing(find_vdr(cdf, var_name))
 
-Base.iterate(cdf::CDFDataset, state = 1) = state > length(cdf) ? nothing : (cdf[keys(cdf)[state]], state + 1)
+# Walk the rVDR then zVDR chain directly; state = (offset, still_in_r_chain)
+function Base.iterate(cdf::CDFDataset, state = (Int(GDR(cdf).rVDRhead), true))
+    offset, in_rchain = state
+    if offset == 0
+        in_rchain || return nothing
+        offset, in_rchain = Int(GDR(cdf).zVDRhead), false
+        offset == 0 && return nothing
+    end
+    RecordSizeType = recordsize_type(cdf)
+    buffer = parent(cdf)
+    name = String(readname(buffer, offset + 45 + 5 * sizeof(RecordSizeType)))
+    var = _variable(cdf, name, _vdr_at(cdf, offset))
+    next_offset = Int(read_be(buffer, offset + 5 + sizeof(RecordSizeType), RecordSizeType))
+    return (var, (next_offset, in_rchain))
+end
 
 function Base.show(io::IO, m::MIME"text/plain", cdf::CDFDataset)
     println(io, typeof(cdf))
     println(io, "path: ", cdf.filename)
     println(io, "variables:")
-    for key in keys(cdf)
-        var = cdf[key]
-        print(io, "  ", key, " : ", size(var), " ")
+    for var in cdf
+        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], "]")

test/runtests.jl

@@ -38,6 +38,16 @@ end
     @test string(Epoch(-1.0e31)) == "FILLVAL"
 end
 
+@testset "iteration" begin
+    for file in ("a_cdf.cdf", "ac_h0_mfi_20230102_v07.cdf") # z-only and r+z chains
+        ds = CDFDataset(data_path(file))
+        @test length(ds) == length(keys(ds))
+        @test [v.name for v in ds] == keys(ds)
+    end
+    ds = CDFDataset(data_path("a_cdf.cdf"))
+    @test haskey(ds, "var") && !haskey(ds, "nonexistent")
+end
+
 @testset "close" begin
     ds = CDFDataset(data_path("a_cdf.cdf"))
     v = ds["var"][:]