onda-arrow-nested-single-table-approach.jl

using UUIDs, Dates, Arrow, Tables

#####
##### Signals/Annotations/Recordings
#####

struct Signal
    file_uri::String
    file_metadata::Union{Nothing,Dict{String,String}}
    channel_names::Vector{String}
    start_nanosecond::Nanosecond
    stop_nanosecond::Nanosecond
    sample_unit::String
    sample_resolution_in_unit::Float64
    sample_offset_in_unit::Float64
    sample_type::String
    sample_rate::Float64
end

Arrow.ArrowTypes.registertype!(Signal, Signal)

struct Annotation{V}
    value::V
    start_nanosecond::Nanosecond
    stop_nanosecond::Nanosecond
end

_annotation_value_type(::Type{Annotation{V}}) where {V} = V

Arrow.ArrowTypes.registertype!(Annotation, Annotation)

mutable struct Recording{V}
    _row::Tables.ColumnsRow{Tables.CopiedColumns{Arrow.Table}}
    _signals::Dict{String,Signal}
    _annotations::Dict{UUID,Annotation{V}}
    function Recording{V}(_row) where {V}
        recording = new{V}()
        recording._row = _row
        return recording
    end
    function Recording{V}(_signals, _annotations) where {V}
        recording = new{V}()
        recording._signals = _signals
        recording._annotations = _annotations
        return recording
    end
end

function Base.getproperty(recording::Recording, field::Symbol)
    if field === :signals
        if !isdefined(recording, :_signals)
            recording._signals = recording._row.signals
        end
        field = :_signals
    elseif field === :annotations
        if !isdefined(recording, :_annotations)
            recording._annotations = recording._row.annotations
        end
        field = :_annotations
    end
    return getfield(recording, field)
end

function read_recordings(tbl)
    schema = Tables.schema(tbl)
    schema.names === (:uuid, :signals, :annotations) || error("")
    V = _annotation_value_type(valtype(schema.types[3]))
    schema.types === (UUID, Dict{String,Signal}, Dict{UUID,Annotation{V}}) || error("")
    allunique(tbl.uuid) || error("")
    recordings = (Recording{V}(row) for row in Tables.rows(tbl))
    return Dict{UUID,Recording{V}}(zip(tbl.uuid, recordings))
end

#####
##### experiments
#####

write_table(tbl) = (io = IOBuffer(); Arrow.write(io, tbl); seekstart(io); Arrow.Table(io))

signals = Dict(
    "eeg" => Signal(
        "file://joe/dave/eeg.lpcm.zst",
        nothing,
        ["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z"],
        Nanosecond(0),
        Nanosecond(100000),
        "microvolt",
        1.0,
        1.0,
        "float64",
        256.0
    ),
    "ecg" => Signal(
        "file://joe/dave/ecg.lpcm.zst",
        Dict("a" => "absidufbaid", "b" => "adjfhaudi"),
        ["a", "b", "c", "d"],
        Nanosecond(0),
        Nanosecond(100000),
        "microvolt",
        0.134,
        1.9875,
        "int32",
        128.0
    )
)

annotations = Dict(uuid4() => Annotation((x = "asdf", y = 3), Nanosecond(0), Nanosecond(10)),
                   uuid4() => Annotation((x = "gtior", y = 43), Nanosecond(0), Nanosecond(10)))

n = 10000
tbl = write_table((uuid=[uuid4() for _ in 1:n], signals=[signals for _ in 1:n], annotations=[annotations for _ in 1:n]))

@time recordings = read_recordings(tbl); # 0.001641 seconds (10.06 k allocations: 1.643 MiB)
example = first(values(recordings));
@time example.signals; # 0.000054 seconds (150 allocations: 9.688 KiB)
@time example.signals; # 0.000004 seconds
@time example.annotations; # 0.000058 seconds (77 allocations: 5.688 KiB)
@time example.annotations; # 0.000005 seconds

function load_all(recordings)
    for r in values(recordings)
        r.signals
        r.annotations
    end
end

recordings = read_recordings(tbl);
@time load_all(recordings); # 0.293740 seconds (2.26 M allocations: 150.022 MiB, 29.98% gc time)