Skip to content

Commit

Permalink
Add reading metadata from Arrow.Table (#481)
Browse files Browse the repository at this point in the history
Partially solves #337 for
`Arrow.Table`.
  • Loading branch information
bkamins authored Dec 1, 2024
1 parent fc8b899 commit 9041cbf
Show file tree
Hide file tree
Showing 2 changed files with 87 additions and 0 deletions.
47 changes: 47 additions & 0 deletions src/table.jl
Original file line number Diff line number Diff line change
Expand Up @@ -363,6 +363,53 @@ See [the official Arrow documentation for more details on custom application met
getmetadata(t::Table) = getfield(t, :metadata)[]
getmetadata(::Any) = nothing

DataAPI.metadatasupport(::Type{Table}) = (read=true, write=false)
DataAPI.colmetadatasupport(::Type{Table}) = (read=true, write=false)

function DataAPI.metadata(t::Table, key::AbstractString; style::Bool=false)
meta = getmetadata(t)[key]
return style ? (meta, :default) : meta
end

function DataAPI.metadata(t::Table, key::AbstractString, default; style::Bool=false)
meta = getmetadata(t)
if meta !== nothing
haskey(meta, key) && return style ? meta[key] : (meta[key], :default)
end
return style ? (default, :default) : default
end

function DataAPI.metadatakeys(t::Table)
meta = getmetadata(t)
meta === nothing && return ()
return keys(meta)
end

function DataAPI.colmetadata(t::Table, col, key::AbstractString; style::Bool=false)
meta = getmetadata(t[col])[key]
return style ? (meta, :default) : meta
end

function DataAPI.colmetadata(t::Table, col, key::AbstractString, default; style::Bool=false)
meta = getmetadata(t[col])
if meta !== nothing
haskey(meta, key) && return style ? (meta[key], :default) : meta[key]
end
return style ? (default, :default) : default
end

function DataAPI.colmetadatakeys(t::Table, col)
meta = getmetadata(t[col])
meta === nothing && return ()
return keys(meta)
end

function DataAPI.colmetadatakeys(t::Table)
return (col => DataAPI.colmetadatakeys(t, col) for
col in Tables.columnnames(t) if
getmetadata(t[col]) !== nothing)
end

Tables.istable(::Table) = true
Tables.columnaccess(::Table) = true
Tables.columns(t::Table) = Tables.CopiedColumns(t)
Expand Down
40 changes: 40 additions & 0 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1070,4 +1070,44 @@ end
@test tbl.col16[1] == Dates.Time(0, 0, 0)
end
end # @testset "misc"

@testset "DataAPI.metadata" begin
df = DataFrame(a=1, b=2, c=3)
for i in 1:2
io = IOBuffer()
if i == 1 # skip writing metadata in the first iteration
Arrow.write(io, df)
else
Arrow.write(io, df, metadata=metadata(df), colmetadata=colmetadata(df))
end
seekstart(io)
tbl = Arrow.Table(io)

@test DataAPI.metadatasupport(typeof(tbl)) == (read=true, write=false)
@test metadata(tbl) == metadata(df)
@test metadata(tbl; style=true) == metadata(df; style=true)
@test_throws Exception metadata(tbl, "xyz")
@test metadata(tbl, "xyz", "something") == "something"
@test metadata(tbl, "xyz", "something"; style=true) == ("something", :default)
@test metadatakeys(tbl) == metadatakeys(df)

@test DataAPI.colmetadatasupport(typeof(tbl)) == (read=true, write=false)
@test colmetadata(tbl) == colmetadata(df)
@test colmetadata(tbl; style=true) == colmetadata(df; style=true)
@test_throws MethodError colmetadata(tbl, "xyz")
@test_throws KeyError colmetadata(tbl, :xyz)
@test colmetadata(tbl, :b) == colmetadata(df, :b)
@test_throws MethodError colmetadata(tbl, :b, "xyz")
@test colmetadata(tbl, :b, "xyz", "something") == "something"
@test colmetadata(tbl, :b, "xyz", "something"; style=true) == ("something", :default)
@test Set(colmetadatakeys(tbl)) == Set(colmetadatakeys(df))

# add metadata for the second iteration
metadata!(df, "tkey", "tvalue")
metadata!(df, "tkey2", "tvalue2")
colmetadata!(df, :a, "ackey", "acvalue")
colmetadata!(df, :a, "ackey2", "acvalue2")
colmetadata!(df, :c, "cckey", "ccvalue")
end
end # @testset "DataAPI.metadata"
end

0 comments on commit 9041cbf

Please sign in to comment.