feat: add lazy attribute loading using ResumableFunctions (#16)

Author: Beforerr

Project.toml

@@ -1,7 +1,7 @@
 name = "CommonDataFormat"
 uuid = "a9737db6-c05c-4e48-868b-6bc41491d9d9"
-version = "0.1.2"
 authors = ["Beforerr <[email protected]> and contributors"]
+version = "0.1.2"
 
 [deps]
 CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193"
@@ -10,6 +10,7 @@ Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 DiskArrays = "3c3547ce-8d99-4f5e-a174-61eb10b00ae3"
 LibDeflate = "9255714d-24a7-4b30-8ea3-d46a97f7e13b"
 Mmap = "a63ad114-7e13-5084-954f-fe012c677804"
+ResumableFunctions = "c5292f4c-5179-55e1-98c5-05642aab7184"
 StaticStrings = "4db0a0c5-418a-4e1d-8806-cb305fe13294"
 UnixTimes = "ab1a18e7-b408-4913-896c-624bb82ed7f4"
 
@@ -27,6 +28,7 @@ Dictionaries = "0.4"
 DiskArrays = "0.4"
 LibDeflate = "0.4.3"
 Mmap = "1"
+ResumableFunctions = "1.0.4"
 StaticStrings = "0.2.6"
 UnixTimes = "1.7.2"
 julia = "1.10"

src/CommonDataFormat.jl

@@ -3,6 +3,7 @@ module CommonDataFormat
 using Dates, UnixTimes
 using Mmap
 using Dictionaries
+using ResumableFunctions
 using DiskArrays
 using StaticStrings
 using Base.Threads

src/loading/attribute.jl

@@ -19,13 +19,13 @@ end
 
 Load all attributes from the CDF file.
 """
-function attrib(cdf::CDFDataset; predicate = is_global)
+function attrib(cdf::CDFDataset; predicate=is_global)
     RecordSizeType = recordsize_type(cdf)
     buffer = cdf.buffer
     cdf_encoding = cdf.cdr.encoding
     offsets = get_offsets(buffer, cdf.gdr.ADRhead, RecordSizeType)
     adrs = map(of -> ADR(buffer, of, RecordSizeType), offsets)
-    adrs = filter(predicate, adrs)
+    adrs = filter!(predicate, adrs)
     names = map(adr -> String(adr.Name), adrs)
     aedrs = map(adrs) do adr
         load_attribute_entries(buffer, adr, RecordSizeType, cdf_encoding)
@@ -42,7 +42,7 @@ function attrib(cdf::CDFDataset, name::String)
     RecordSizeType = recordsize_type(cdf)
     buffer = cdf.buffer
     cdf_encoding = cdf.cdr.encoding
-    offsets = get_offsets(buffer, cdf.gdr.ADRhead, RecordSizeType)
+    offsets = get_offsets_lazy(buffer, cdf.gdr.ADRhead, RecordSizeType)
     for offset in offsets
         adr = ADR(buffer, offset, RecordSizeType)
         name == String(adr.Name) && return load_attribute_entries(buffer, adr, RecordSizeType, cdf_encoding)
@@ -59,11 +59,11 @@ function vattrib(cdf::CDFDataset, varnum::Integer)
     RecordSizeType = recordsize_type(cdf)
     buffer = cdf.buffer
     cdf_encoding = cdf.cdr.encoding
-    attributes = Dict{String, Union{String, Vector}}()
-    offsets = get_offsets(buffer, cdf.gdr.ADRhead, RecordSizeType)
+    attributes = Dict{String,Union{String,Vector}}()
+    offsets = get_offsets_lazy(buffer, cdf.gdr.ADRhead, RecordSizeType)
     for offset in offsets
+        is_global(buffer, offset, RecordSizeType) && continue
         adr = ADR(buffer, offset, RecordSizeType)
-        is_global(adr) && continue
         @assert min(adr.AgrEDRhead, adr.AzEDRhead) == 0
         head = max(adr.AgrEDRhead, adr.AzEDRhead)
         found = _search_aedr_entries(buffer, head, RecordSizeType, cdf_encoding, varnum)
@@ -95,11 +95,11 @@ function vattrib(cdf, varnum, name)
     cdf_encoding = cdf.cdr.encoding
 
     # Search for the specific attribute by name first
-    offsets = get_offsets(buffer, cdf.gdr.ADRhead, RecordSizeType)
+    offsets = get_offsets_lazy(buffer, cdf.gdr.ADRhead, RecordSizeType)
     name_bytes = codeunits(name)
     for offset in offsets
+        is_global(buffer, offset, RecordSizeType) && continue
         adr = ADR(buffer, offset, RecordSizeType)
-        is_global(adr) && continue
         adr.Name != name_bytes && continue
         @assert min(adr.AgrEDRhead, adr.AzEDRhead) == 0
         head = max(adr.AgrEDRhead, adr.AzEDRhead)
@@ -127,15 +127,14 @@ end
 
 Return a list of attribute names in the CDF file.
 """
-function attribnames(cdf::CDFDataset; filter = is_global)
+function attribnames(cdf::CDFDataset; predicate=is_global)
     names = String[]
     buffer = cdf.buffer
     RecordSizeType = recordsize_type(cdf)
-    offsets = get_offsets(buffer, cdf.gdr.ADRhead, RecordSizeType)
-    sizehint!(names, length(offsets))
+    offsets = get_offsets_lazy(buffer, cdf.gdr.ADRhead, RecordSizeType)
     for offset in offsets
         adr = ADR(buffer, offset, RecordSizeType)
-        filter(adr) && push!(names, String(adr.Name))
+        predicate(adr) && push!(names, String(adr.Name))
     end
     return names
 end

src/parsing.jl

@@ -125,6 +125,12 @@ function readname(buf::Vector{UInt8}, offset::Int)
     return @views buf[offset:(offset + 255)]
 end
 
+@resumable function get_offsets_lazy(buffer::Vector{UInt8}, pos::Int64, ::Type{RecordSizeType}) where {RecordSizeType}
+    while pos != 0
+        @yield pos
+        pos = read_be(buffer, pos + 1 + sizeof(RecordSizeType) + 4, RecordSizeType)
+    end
+end
 
 function get_offsets!(offsets, buffer::Vector{UInt8}, pos::Int64, RecordSizeType)
     while pos != 0

src/records/adr.jl

@@ -20,6 +20,8 @@ struct ADR{FSZ, S} <: Record
 end
 
 is_global(adr) = adr.Scope == 1
+is_global(buffer, offset, ::Type{Int32}) = read_be(buffer, offset + 17, Int32) == 1
+is_global(buffer, offset, ::Type{Int64}) = read_be(buffer, offset + 29, Int32) == 1
 
 
 """