diff --git a/src/Arrow.jl b/src/Arrow.jl index 6f3ccdf8..5fb38ddc 100644 --- a/src/Arrow.jl +++ b/src/Arrow.jl @@ -26,11 +26,12 @@ This implementation supports the 1.0 version of the specification, including sup * Extension types * Streaming, file, record batch, and replacement and isdelta dictionary messages * Buffer compression/decompression via the standard LZ4 frame and Zstd formats + * C Data Interface import It currently doesn't include support for: * Tensors or sparse tensors * Flight RPC - * C data interface + * C Data Interface export Third-party data formats: * csv and parquet support via the existing [CSV.jl](https://github.com/JuliaData/CSV.jl) and [Parquet.jl](https://github.com/JuliaIO/Parquet.jl) packages @@ -76,6 +77,7 @@ include("utils.jl") include("arraytypes/arraytypes.jl") include("eltypes.jl") include("table.jl") +include("cdata.jl") include("write.jl") include("append.jl") include("show.jl") diff --git a/src/arraytypes/dictencoding.jl b/src/arraytypes/dictencoding.jl index 3e3576c5..71f3c23e 100644 --- a/src/arraytypes/dictencoding.jl +++ b/src/arraytypes/dictencoding.jl @@ -119,6 +119,27 @@ signedtype(::Type{UInt32}) = Int32 signedtype(::Type{UInt64}) = Int64 signedtype(::Type{T}) where {T<:Signed} = T +function _dict_indices(refa, pool, validity::ValidityBitmap) + ET = signedtype(length(pool)) + inds = Vector{ET}(undef, length(refa)) + lo = firstindex(pool) + hi = lastindex(pool) + j = 1 + @inbounds for idx in eachindex(refa) + ref = refa[idx] + # Per the Arrow columnar format spec, dictionary data stores 0-based indices. + if lo <= ref <= hi + inds[j] = ET(ref - lo) + elseif validity[j] + throw(ArgumentError("dictionary reference index is out of range")) + else + inds[j] = zero(ET) + end + j += 1 + end + return inds +end + indtype(d::DictEncoded{T,S,A}) where {T,S,A} = S indtype(c::Compressed{Z,A}) where {Z,A<:DictEncoded} = indtype(c.data) @@ -229,12 +250,12 @@ function arrowvector( else pool = DataAPI.refpool(x) refa = DataAPI.refarray(x) - inds = copyto!(similar(Vector{signedtype(length(pool))}, length(refa)), refa) end # adjust to "offset" instead of index - inds .-= firstindex(refa) + inds = _dict_indices(refa, pool, validity) + pooldata = firstindex(pool) == 1 ? pool : collect(pool) data = arrowvector( - pool, + pooldata, i, nl, fi, diff --git a/src/cdata.jl b/src/cdata.jl new file mode 100644 index 00000000..66d36b68 --- /dev/null +++ b/src/cdata.jl @@ -0,0 +1,1394 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +mutable struct ArrowSchema + format::Cstring + name::Cstring + metadata::Cstring + flags::Int64 + n_children::Int64 + children::Ptr{Ptr{ArrowSchema}} + dictionary::Ptr{ArrowSchema} + release::Ptr{Cvoid} + private_data::Ptr{Cvoid} +end + +mutable struct ArrowArray + length::Int64 + null_count::Int64 + offset::Int64 + n_buffers::Int64 + n_children::Int64 + buffers::Ptr{Ptr{Cvoid}} + children::Ptr{Ptr{ArrowArray}} + dictionary::Ptr{ArrowArray} + release::Ptr{Cvoid} + private_data::Ptr{Cvoid} +end + +const _CDATA_PTR_SIZE = sizeof(Ptr{Cvoid}) +@assert sizeof(ArrowSchema) == 7 * _CDATA_PTR_SIZE + 2 * sizeof(Int64) +@assert sizeof(ArrowArray) == 5 * _CDATA_PTR_SIZE + 5 * sizeof(Int64) + +const ARROW_FLAG_DICTIONARY_ORDERED = Int64(1) +const ARROW_FLAG_NULLABLE = Int64(2) +const ARROW_FLAG_MAP_KEYS_SORTED = Int64(4) + +const _CDATA_MAX_CHILDREN = 10_000 +const _CDATA_MAX_DEPTH = 128 +const _CDATA_MAX_BUFFERS = 64 +const _CDATA_MAX_FORMAT_BYTES = 4096 +const _CDATA_MAX_NAME_BYTES = 1 << 16 +const _CDATA_MAX_METADATA_PAIRS = 4096 +const _CDATA_MAX_METADATA_BYTES = 1 << 20 +const _CDATA_MAX_METADATA_FIELD_BYTES = 1 << 20 +const _CDATA_MAX_FIXED_SIZE = 4096 + +abstract type CDataFormat end + +struct CDataNullFormat <: CDataFormat end +struct CDataBoolFormat <: CDataFormat end +struct CDataPrimitiveFormat <: CDataFormat + storage::Type +end +struct CDataBinaryFormat{O} <: CDataFormat + juliatype::Type +end +struct CDataFixedSizeBinaryFormat <: CDataFormat + bytewidth::Int +end +struct CDataListFormat{O} <: CDataFormat end +struct CDataFixedSizeListFormat <: CDataFormat + listsize::Int +end +struct CDataStructFormat <: CDataFormat end + +abstract type CDataVector{T} <: ArrowVector{T} end + +mutable struct CDataOwner + schema_ptr::Ptr{ArrowSchema} + array_ptr::Ptr{ArrowArray} + released::Bool + lock::ReentrantLock +end + +function CDataOwner(schema_ptr::Ptr{ArrowSchema}, array_ptr::Ptr{ArrowArray}) + owner = CDataOwner(schema_ptr, array_ptr, false, ReentrantLock()) + finalizer(release_c_data, owner) + return owner +end + +struct CDataValidity + bytes::Vector{UInt8} + bitoffset::Int + len::Int + null_count::Int +end + +struct CDataNull{T} <: CDataVector{T} + owner::CDataOwner + len::Int + metadata::Union{Nothing,Base.ImmutableDict{String,String}} +end + +struct CDataPrimitive{T,S,A<:AbstractVector{S}} <: CDataVector{T} + owner::CDataOwner + validity::CDataValidity + data::A + metadata::Union{Nothing,Base.ImmutableDict{String,String}} +end + +struct CDataBool{T} <: CDataVector{T} + owner::CDataOwner + validity::CDataValidity + data::Vector{UInt8} + bitoffset::Int + len::Int + metadata::Union{Nothing,Base.ImmutableDict{String,String}} +end + +struct CDataBinary{T,O,A<:AbstractVector{O}} <: CDataVector{T} + owner::CDataOwner + validity::CDataValidity + offsets::A + data::Vector{UInt8} + len::Int + metadata::Union{Nothing,Base.ImmutableDict{String,String}} +end + +struct CDataFixedSizeBinary{T} <: CDataVector{T} + owner::CDataOwner + validity::CDataValidity + data::Vector{UInt8} + bytewidth::Int + len::Int + metadata::Union{Nothing,Base.ImmutableDict{String,String}} +end + +struct CDataList{T,O,A<:AbstractVector{O},C<:AbstractVector} <: CDataVector{T} + owner::CDataOwner + validity::CDataValidity + offsets::A + data::C + len::Int + metadata::Union{Nothing,Base.ImmutableDict{String,String}} +end + +struct CDataFixedSizeList{T,C<:AbstractVector} <: CDataVector{T} + owner::CDataOwner + validity::CDataValidity + data::C + listsize::Int + offset::Int + len::Int + metadata::Union{Nothing,Base.ImmutableDict{String,String}} +end + +struct CDataStruct{T,S,fnames} <: CDataVector{T} + owner::CDataOwner + validity::CDataValidity + data::S + offset::Int + len::Int + metadata::Union{Nothing,Base.ImmutableDict{String,String}} +end + +struct CDataSlice{T,V<:AbstractVector{T}} <: CDataVector{T} + parent::V + first::Int + len::Int +end + +struct CDataMasked{T,V<:CDataVector} <: CDataVector{T} + parent::V + parent_validity::CDataValidity +end + +struct CDataTable <: Tables.AbstractColumns + names::Vector{Symbol} + types::Vector{Type} + columns::Vector{AbstractVector} + lookup::Dict{Symbol,AbstractVector} + metadata::Union{Nothing,Base.ImmutableDict{String,String}} + owner::CDataOwner + rowcount::Int +end + +struct CDataNode + schema_ptr::Ptr{ArrowSchema} + array_ptr::Ptr{ArrowArray} + schema::ArrowSchema + array::ArrowArray + format::CDataFormat + name::Union{Nothing,String} + metadata::Union{Nothing,Base.ImmutableDict{String,String}} + buffers::Vector{Ptr{Cvoid}} + children::Vector{CDataNode} + len::Int + offset::Int + null_count::Int +end + +Base.IndexStyle(::Type{<:CDataVector}) = Base.IndexLinear() + +Base.size(x::CDataNull) = (x.len,) +Base.size(x::CDataPrimitive) = size(x.data) +Base.size(x::CDataBool) = (x.len,) +Base.size(x::CDataBinary) = (x.len,) +Base.size(x::CDataFixedSizeBinary) = (x.len,) +Base.size(x::CDataList) = (x.len,) +Base.size(x::CDataFixedSizeList) = (x.len,) +Base.size(x::CDataStruct) = (x.len,) +Base.size(x::CDataSlice) = (x.len,) +Base.size(x::CDataMasked) = size(x.parent) +Base.length(t::CDataTable) = length(getfield(t, :columns)) + +_owner(x::CDataVector) = getfield(x, :owner) +_owner(x::CDataSlice) = _owner(x.parent) +_owner(x::CDataMasked) = _owner(x.parent) + +function _check_live(owner::CDataOwner) + lock(owner.lock) + try + owner.released && throw(ArgumentError("Arrow C Data object has been released")) + return + finally + unlock(owner.lock) + end +end + +_check_live(x::CDataVector) = _check_live(_owner(x)) +_check_live(t::CDataTable) = _check_live(getfield(t, :owner)) + +function _with_live(f::F, owner::CDataOwner) where {F} + lock(owner.lock) + try + owner.released && throw(ArgumentError("Arrow C Data object has been released")) + return f() + finally + unlock(owner.lock) + end +end + +_with_live(f::F, x::CDataVector) where {F} = _with_live(f, _owner(x)) +_with_live(f::F, t::CDataTable) where {F} = _with_live(f, getfield(t, :owner)) + +validitybitmap(x::CDataNull) = nothing +nullcount(x::CDataNull) = x.len +nullcount(x::CDataVector) = validitybitmap(x).null_count +nullcount(x::CDataSlice) = count(i -> ismissing(x[i]), eachindex(x)) +nullcount(x::CDataMasked) = count(i -> ismissing(x[i]), eachindex(x)) +getmetadata(x::CDataSlice) = getmetadata(x.parent) +getmetadata(x::CDataMasked) = getmetadata(x.parent) +getmetadata(t::CDataTable) = getfield(t, :metadata) +validitybitmap(::CDataMasked) = nothing + +@inline function _valid_bit(bytes::Vector{UInt8}, bitoffset::Int, i::Integer) + pos = bitoffset + Int(i) - 1 + byte = @inbounds bytes[(pos >>> 3) + 1] + return getbit(byte, (pos & 0x07) + 1) +end + +@inline function _valid(v::CDataValidity, i::Integer) + v.null_count == 0 && return true + return _valid_bit(v.bytes, v.bitoffset, i) +end + +function _count_nulls(v::CDataValidity) + n = 0 + for i = 1:v.len + n += _valid_bit(v.bytes, v.bitoffset, i) ? 0 : 1 + end + return n +end + +@propagate_inbounds function Base.getindex(x::CDataNull, i::Integer) + return _with_live(x) do + @boundscheck checkbounds(x, i) + return missing + end +end + +@propagate_inbounds function Base.getindex(x::CDataPrimitive{T}, i::Integer) where {T} + return _with_live(x) do + @boundscheck checkbounds(x, i) + if !_valid(x.validity, i) + return missing + end + return @inbounds ArrowTypes.fromarrow(T, x.data[i]) + end +end + +@propagate_inbounds function Base.getindex(x::CDataBool{T}, i::Integer) where {T} + return _with_live(x) do + @boundscheck checkbounds(x, i) + if !_valid(x.validity, i) + return missing + end + pos = x.bitoffset + Int(i) - 1 + byte = @inbounds x.data[(pos >>> 3) + 1] + return ArrowTypes.fromarrow(T, getbit(byte, (pos & 0x07) + 1)) + end +end + +@propagate_inbounds function Base.getindex(x::CDataBinary{T}, i::Integer) where {T} + return _with_live(x) do + @boundscheck checkbounds(x, i) + if !_valid(x.validity, i) + return missing + end + lo = Int(@inbounds x.offsets[i]) + 1 + hi = Int(@inbounds x.offsets[i + 1]) + n = hi - lo + 1 + if n == 0 + return ArrowTypes.fromarrow(T, "") + end + data = x.data + owner = _owner(x) + GC.@preserve x data owner begin + return ArrowTypes.fromarrow(T, pointer(data, lo), n) + end + end +end + +@propagate_inbounds function Base.getindex(x::CDataFixedSizeBinary{T}, i::Integer) where {T} + return _with_live(x) do + @boundscheck checkbounds(x, i) + if !_valid(x.validity, i) + return missing + end + offset = (Int(i) - 1) * x.bytewidth + tup = ntuple(j -> @inbounds(x.data[offset + j]), x.bytewidth) + return ArrowTypes.fromarrow(T, tup) + end +end + +@propagate_inbounds function Base.getindex(x::CDataList{T}, i::Integer) where {T} + return _with_live(x) do + @boundscheck checkbounds(x, i) + if !_valid(x.validity, i) + return missing + end + lo = Int(@inbounds x.offsets[i]) + 1 + hi = Int(@inbounds x.offsets[i + 1]) + return ArrowTypes.fromarrow(T, @view x.data[lo:hi]) + end +end + +@propagate_inbounds function Base.getindex(x::CDataFixedSizeList{T}, i::Integer) where {T} + return _with_live(x) do + @boundscheck checkbounds(x, i) + if !_valid(x.validity, i) + return missing + end + offset = (x.offset + Int(i) - 1) * x.listsize + tup = ntuple(j -> @inbounds(x.data[offset + j]), x.listsize) + return ArrowTypes.fromarrow(T, tup) + end +end + +@propagate_inbounds function Base.getindex( + x::CDataStruct{T,S,fnames}, + i::Integer, +) where {T,S,fnames} + return _with_live(x) do + @boundscheck checkbounds(x, i) + if !_valid(x.validity, i) + return missing + end + j = x.offset + Int(i) + vals = ntuple(k -> @inbounds(x.data[k][j]), fieldcount(S)) + NT = Base.nonmissingtype(T) + if isnamedtuple(NT) || istuple(NT) + return ArrowTypes.fromarrow(T, NT(vals)) + else + return ArrowTypes.fromarrow(T, _fromarrowstruct(NT, Val{fnames}(), vals...)) + end + end +end + +@propagate_inbounds function Base.getindex(x::CDataSlice, i::Integer) + return _with_live(x) do + @boundscheck checkbounds(x, i) + return @inbounds x.parent[x.first + Int(i) - 1] + end +end + +@propagate_inbounds function Base.getindex(x::CDataMasked, i::Integer) + return _with_live(x) do + @boundscheck checkbounds(x, i) + if !_valid(x.parent_validity, i) + return missing + end + return @inbounds x.parent[i] + end +end + +function Base.collect(x::CDataVector{T}) where {T} + return _with_live(x) do + out = Vector{T}(undef, length(x)) + for i in eachindex(x) + @inbounds out[i] = x[i] + end + return out + end +end + +Base.copy(x::CDataVector) = collect(x) + +function Base.copy(t::CDataTable) + return _with_live(t) do + names = getfield(t, :names) + columns = getfield(t, :columns) + return NamedTuple{Tuple(names)}(Tuple(copy(col) for col in columns)) + end +end + +Tables.istable(::Type{CDataTable}) = true +Tables.columnaccess(::Type{CDataTable}) = true +Tables.columns(t::CDataTable) = t +Tables.columnnames(t::CDataTable) = getfield(t, :names) +Tables.schema(t::CDataTable) = Tables.Schema(getfield(t, :names), getfield(t, :types)) +Tables.getcolumn(t::CDataTable, i::Int) = (_check_live(t); getfield(t, :columns)[i]) +Tables.getcolumn(t::CDataTable, nm::Symbol) = (_check_live(t); getfield(t, :lookup)[nm]) +Tables.rowcount(t::CDataTable) = getfield(t, :rowcount) + +Base.getindex(t::CDataTable, i::Int) = Tables.getcolumn(t, i) +Base.getindex(t::CDataTable, nm::Symbol) = Tables.getcolumn(t, nm) +function Base.getproperty(t::CDataTable, nm::Symbol) + lookup = getfield(t, :lookup) + haskey(lookup, nm) && return Tables.getcolumn(t, nm) + return getfield(t, nm) +end +Base.propertynames(t::CDataTable, private::Bool=false) = + private ? fieldnames(typeof(t)) : Tuple(getfield(t, :names)) + +DataAPI.metadatasupport(::Type{CDataTable}) = (read=true, write=false) +DataAPI.colmetadatasupport(::Type{CDataTable}) = (read=true, write=false) + +function DataAPI.metadata(t::CDataTable; style::Bool=false) + meta = getmetadata(t) + meta === nothing && return Dict{String,Any}() + if style + return Dict(k => (v, :default) for (k, v) in meta) + else + return Dict(meta) + end +end + +function DataAPI.metadata(t::CDataTable, key::AbstractString; style::Bool=false) + meta = getmetadata(t) + meta === nothing && throw(KeyError(key)) + val = meta[key] + return style ? (val, :default) : val +end + +function DataAPI.metadata(t::CDataTable, key::AbstractString, default; style::Bool=false) + meta = getmetadata(t) + if meta !== nothing && haskey(meta, key) + val = meta[key] + return style ? (val, :default) : val + end + return style ? (default, :default) : default +end + +function DataAPI.metadatakeys(t::CDataTable) + meta = getmetadata(t) + meta === nothing && return () + return keys(meta) +end + +function DataAPI.colmetadata(t::CDataTable; style::Bool=false) + pairs = Pair{Symbol,Any}[] + for col in getfield(t, :names) + meta = getmetadata(t[col]) + if meta !== nothing && !isempty(meta) + push!(pairs, col => DataAPI.colmetadata(t, col; style=style)) + end + end + return Dict(pairs) +end + +function DataAPI.colmetadata(t::CDataTable, col; style::Bool=false) + meta = getmetadata(t[col]) + meta === nothing && return Dict{String,Any}() + if style + return Dict(k => (v, :default) for (k, v) in meta) + else + return Dict(meta) + end +end + +function DataAPI.colmetadata(t::CDataTable, col, key::AbstractString; style::Bool=false) + meta = getmetadata(t[col]) + meta === nothing && throw(KeyError(key)) + val = meta[key] + return style ? (val, :default) : val +end + +function DataAPI.colmetadata( + t::CDataTable, + col, + key::AbstractString, + default; + style::Bool=false, +) + meta = getmetadata(t[col]) + if meta !== nothing && haskey(meta, key) + val = meta[key] + return style ? (val, :default) : val + end + return style ? (default, :default) : default +end + +function DataAPI.colmetadatakeys(t::CDataTable, col) + meta = getmetadata(t[col]) + meta === nothing && return () + return keys(meta) +end + +function DataAPI.colmetadatakeys(t::CDataTable) + return ( + col => DataAPI.colmetadatakeys(t, col) for + col in Tables.columnnames(t) if getmetadata(t[col]) !== nothing + ) +end + +function release_c_data(owner::CDataOwner) + array_ptr = Ptr{ArrowArray}(C_NULL) + schema_ptr = Ptr{ArrowSchema}(C_NULL) + array_release = Ptr{Cvoid}(C_NULL) + schema_release = Ptr{Cvoid}(C_NULL) + lock(owner.lock) + try + owner.released && return + owner.released = true + array_ptr = owner.array_ptr + schema_ptr = owner.schema_ptr + if array_ptr != C_NULL + array = unsafe_load(array_ptr) + array_release = array.release + end + if schema_ptr != C_NULL + schema = unsafe_load(schema_ptr) + schema_release = schema.release + end + finally + unlock(owner.lock) + end + # Per the Arrow C Data Interface spec, consumers release only the base structures. + if array_ptr != C_NULL && array_release != C_NULL + ccall(array_release, Cvoid, (Ptr{ArrowArray},), array_ptr) + end + if schema_ptr != C_NULL && schema_release != C_NULL + ccall(schema_release, Cvoid, (Ptr{ArrowSchema},), schema_ptr) + end + return +end + +""" + Arrow.release_c_data(x) + +Release C Data resources owned by an imported array or table. The call is +idempotent. Reads through imported arrays throw after release. +""" +release_c_data(x::CDataVector) = release_c_data(_owner(x)) +release_c_data(t::CDataTable) = release_c_data(getfield(t, :owner)) + +function _to_int(x::Int64, name) + x > typemax(Int) && throw(ArgumentError("$name exceeds the Julia Int range")) + return Int(x) +end + +function _checked_nonnegative(x::Int64, name) + x < 0 && throw(ArgumentError("$name must be nonnegative")) + return _to_int(x, name) +end + +function _checked_add(a::Int, b::Int, name) + b > typemax(Int) - a && throw(ArgumentError("$name overflows")) + return a + b +end + +function _checked_mul(a::Int, b::Int, name) + a != 0 && b > typemax(Int) รท a && throw(ArgumentError("$name overflows")) + return a * b +end + +function _parse_c_data_format(format::AbstractString) + format == "n" && return CDataNullFormat() + format == "b" && return CDataBoolFormat() + format == "c" && return CDataPrimitiveFormat(Int8) + format == "C" && return CDataPrimitiveFormat(UInt8) + format == "s" && return CDataPrimitiveFormat(Int16) + format == "S" && return CDataPrimitiveFormat(UInt16) + format == "i" && return CDataPrimitiveFormat(Int32) + format == "I" && return CDataPrimitiveFormat(UInt32) + format == "l" && return CDataPrimitiveFormat(Int64) + format == "L" && return CDataPrimitiveFormat(UInt64) + format == "e" && return CDataPrimitiveFormat(Float16) + format == "f" && return CDataPrimitiveFormat(Float32) + format == "g" && return CDataPrimitiveFormat(Float64) + format == "z" && return CDataBinaryFormat{Int32}(Base.CodeUnits) + format == "Z" && return CDataBinaryFormat{Int64}(Base.CodeUnits) + format == "u" && return CDataBinaryFormat{Int32}(String) + format == "U" && return CDataBinaryFormat{Int64}(String) + format == "+l" && return CDataListFormat{Int32}() + format == "+L" && return CDataListFormat{Int64}() + format == "+s" && return CDataStructFormat() + format == "tdD" && return CDataPrimitiveFormat(Date{Meta.DateUnit.DAY,Int32}) + format == "tdm" && return CDataPrimitiveFormat(Date{Meta.DateUnit.MILLISECOND,Int64}) + format == "tts" && return CDataPrimitiveFormat(Time{Meta.TimeUnit.SECOND,Int32}) + format == "ttm" && return CDataPrimitiveFormat(Time{Meta.TimeUnit.MILLISECOND,Int32}) + format == "ttu" && return CDataPrimitiveFormat(Time{Meta.TimeUnit.MICROSECOND,Int64}) + format == "ttn" && return CDataPrimitiveFormat(Time{Meta.TimeUnit.NANOSECOND,Int64}) + format == "tDs" && return CDataPrimitiveFormat(Duration{Meta.TimeUnit.SECOND}) + format == "tDm" && return CDataPrimitiveFormat(Duration{Meta.TimeUnit.MILLISECOND}) + format == "tDu" && return CDataPrimitiveFormat(Duration{Meta.TimeUnit.MICROSECOND}) + format == "tDn" && return CDataPrimitiveFormat(Duration{Meta.TimeUnit.NANOSECOND}) + format == "tiM" && + return CDataPrimitiveFormat(Interval{Meta.IntervalUnit.YEAR_MONTH,Int32}) + format == "tiD" && + return CDataPrimitiveFormat(Interval{Meta.IntervalUnit.DAY_TIME,Int64}) + if startswith(format, "ts") + return CDataPrimitiveFormat(_parse_timestamp_format(format)) + elseif startswith(format, "d:") + return CDataPrimitiveFormat(_parse_decimal_format(format)) + elseif startswith(format, "w:") + return CDataFixedSizeBinaryFormat( + _parse_positive_int(format[3:end], format, _CDATA_MAX_FIXED_SIZE), + ) + elseif startswith(format, "+w:") + return CDataFixedSizeListFormat( + _parse_positive_int(format[4:end], format, _CDATA_MAX_FIXED_SIZE), + ) + end + throw(ArgumentError("unsupported Arrow C Data format string: $format")) +end + +function _parse_positive_int(s, format, max_value::Int=typemax(Int)) + n = try + parse(Int, s) + catch + throw(ArgumentError("invalid Arrow C Data format string: $format")) + end + n <= 0 && throw(ArgumentError("invalid Arrow C Data format string: $format")) + n > max_value && + throw(ArgumentError("Arrow C Data format size exceeds the import limit: $format")) + return n +end + +function _parse_timestamp_format(format) + length(format) >= 4 || + throw(ArgumentError("invalid Arrow C Data format string: $format")) + format[4] == ':' || throw(ArgumentError("invalid Arrow C Data format string: $format")) + unit = format[3] + U = + unit == 's' ? Meta.TimeUnit.SECOND : + unit == 'm' ? Meta.TimeUnit.MILLISECOND : + unit == 'u' ? Meta.TimeUnit.MICROSECOND : + unit == 'n' ? Meta.TimeUnit.NANOSECOND : + throw(ArgumentError("invalid Arrow C Data timestamp unit: $format")) + tz = length(format) == 4 ? nothing : Symbol(format[5:end]) + return Timestamp{U,tz} +end + +function _parse_decimal_format(format) + parts = split(format[3:end], ',') + 2 <= length(parts) <= 3 || + throw(ArgumentError("invalid Arrow C Data decimal format: $format")) + precision = _parse_int_field(parts[1], "decimal precision", format) + scale = _parse_int_field(parts[2], "decimal scale", format) + bitwidth = + length(parts) == 3 ? _parse_int_field(parts[3], "decimal bit width", format) : 128 + precision > 0 || throw(ArgumentError("decimal precision must be positive")) + if bitwidth == 128 + return Decimal{precision,scale,Int128} + elseif bitwidth == 256 + return Decimal{precision,scale,Int256} + else + throw(ArgumentError("unsupported decimal bit width: $bitwidth")) + end +end + +function _parse_int_field(s, field, format) + try + return parse(Int, s) + catch + throw(ArgumentError("invalid Arrow C Data $field: $format")) + end +end + +_expected_buffers(::CDataNullFormat) = 0 +_expected_buffers(::CDataBoolFormat) = 2 +_expected_buffers(::CDataPrimitiveFormat) = 2 +_expected_buffers(::CDataBinaryFormat) = 3 +_expected_buffers(::CDataFixedSizeBinaryFormat) = 2 +_expected_buffers(::CDataListFormat) = 2 +_expected_buffers(::CDataFixedSizeListFormat) = 1 +_expected_buffers(::CDataStructFormat) = 1 + +_expected_children(::CDataNullFormat) = 0 +_expected_children(::CDataBoolFormat) = 0 +_expected_children(::CDataPrimitiveFormat) = 0 +_expected_children(::CDataBinaryFormat) = 0 +_expected_children(::CDataFixedSizeBinaryFormat) = 0 +_expected_children(::CDataListFormat) = 1 +_expected_children(::CDataFixedSizeListFormat) = 1 +_expected_children(::CDataStructFormat) = nothing + +function _nullable(schema::ArrowSchema, null_count::Int) + return (schema.flags & ARROW_FLAG_NULLABLE) != 0 || null_count != 0 +end + +function _julia_type(storage::Type, nullable::Bool, convert::Bool) + T = convert ? finaljuliatype(storage) : storage + return nullable ? Union{T,Missing} : T +end + +function _load_name(ptr::Cstring) + ptr == C_NULL && return nothing + name = _unsafe_string_bounded(ptr, _CDATA_MAX_NAME_BYTES, "ArrowSchema.name") + return isempty(name) ? nothing : name +end + +function _unsafe_string_bounded(ptr::Cstring, maxbytes::Int, name) + bytes = UInt8[] + sizehint!(bytes, min(maxbytes, 128)) + p = Ptr{UInt8}(ptr) + for i = 1:maxbytes + byte = unsafe_load(p, i) + byte == 0x00 && return String(bytes) + push!(bytes, byte) + end + throw(ArgumentError("$name exceeds the import limit")) +end + +function _metadata_dict(pairs) + isempty(pairs) && return Base.ImmutableDict{String,String}() + return toidict(pairs) +end + +@inline function _unsafe_load_int32(p::Ptr{UInt8}) + b1 = UInt32(unsafe_load(p, 1)) + b2 = UInt32(unsafe_load(p, 2)) + b3 = UInt32(unsafe_load(p, 3)) + b4 = UInt32(unsafe_load(p, 4)) + u = + ENDIAN_BOM == 0x04030201 ? b1 | (b2 << 8) | (b3 << 16) | (b4 << 24) : + ENDIAN_BOM == 0x01020304 ? (b1 << 24) | (b2 << 16) | (b3 << 8) | b4 : + error("unsupported host byte order") + return reinterpret(Int32, u) +end + +function _parse_c_metadata(ptr::Cstring) + ptr == C_NULL && return nothing + # Per the Arrow C Data Interface spec, metadata is length encoded and not null terminated. + p = Ptr{UInt8}(ptr) + count = Int(_unsafe_load_int32(p)) + count < 0 && throw(ArgumentError("Arrow C Data metadata pair count is negative")) + count > _CDATA_MAX_METADATA_PAIRS && + throw(ArgumentError("Arrow C Data metadata pair count exceeds the limit")) + pos = 4 + total = 4 + pairs = Pair{String,String}[] + for _ = 1:count + key_len = Int(_unsafe_load_int32(p + pos)) + pos += 4 + total += 4 + key_len < 0 && throw(ArgumentError("Arrow C Data metadata key length is negative")) + key_len > _CDATA_MAX_METADATA_FIELD_BYTES && + throw(ArgumentError("Arrow C Data metadata key length exceeds the limit")) + total = _checked_add(total, key_len, "metadata byte count") + total > _CDATA_MAX_METADATA_BYTES && + throw(ArgumentError("Arrow C Data metadata byte count exceeds the limit")) + key = unsafe_string(p + pos, key_len) + pos += key_len + + value_len = Int(_unsafe_load_int32(p + pos)) + pos += 4 + total += 4 + value_len < 0 && + throw(ArgumentError("Arrow C Data metadata value length is negative")) + value_len > _CDATA_MAX_METADATA_FIELD_BYTES && + throw(ArgumentError("Arrow C Data metadata value length exceeds the limit")) + total = _checked_add(total, value_len, "metadata byte count") + total > _CDATA_MAX_METADATA_BYTES && + throw(ArgumentError("Arrow C Data metadata byte count exceeds the limit")) + value = unsafe_string(p + pos, value_len) + pos += value_len + push!(pairs, key => value) + end + return _metadata_dict(pairs) +end + +function _load_buffers(array::ArrowArray, expected::Int) + n_buffers = _checked_nonnegative(array.n_buffers, "ArrowArray.n_buffers") + n_buffers > _CDATA_MAX_BUFFERS && + throw(ArgumentError("ArrowArray.n_buffers exceeds the import limit")) + n_buffers >= expected || + throw(ArgumentError("ArrowArray has fewer buffers than required")) + if expected > 0 && array.buffers == C_NULL + throw(ArgumentError("ArrowArray.buffers is NULL")) + end + buffers = Ptr{Cvoid}[] + for i = 1:expected + push!(buffers, unsafe_load(array.buffers, i)) + end + return buffers +end + +function _load_child_ptrs(schema::ArrowSchema, array::ArrowArray, count::Int) + count == 0 && return Ptr{ArrowSchema}[], Ptr{ArrowArray}[] + schema.children == C_NULL && throw(ArgumentError("ArrowSchema.children is NULL")) + array.children == C_NULL && throw(ArgumentError("ArrowArray.children is NULL")) + schema_children = Ptr{ArrowSchema}[] + array_children = Ptr{ArrowArray}[] + for i = 1:count + schema_child = unsafe_load(schema.children, i) + array_child = unsafe_load(array.children, i) + schema_child == C_NULL && throw(ArgumentError("ArrowSchema child pointer is NULL")) + array_child == C_NULL && throw(ArgumentError("ArrowArray child pointer is NULL")) + push!(schema_children, schema_child) + push!(array_children, array_child) + end + return schema_children, array_children +end + +function _validate_common(schema::ArrowSchema, array::ArrowArray, top_level::Bool) + schema.format == C_NULL && throw(ArgumentError("ArrowSchema.format is NULL")) + if top_level + schema.release == C_NULL && throw(ArgumentError("ArrowSchema.release is NULL")) + array.release == C_NULL && throw(ArgumentError("ArrowArray.release is NULL")) + elseif schema.release == C_NULL || array.release == C_NULL + throw(ArgumentError("released Arrow C Data child structure")) + end + len = _checked_nonnegative(array.length, "ArrowArray.length") + offset = _checked_nonnegative(array.offset, "ArrowArray.offset") + _checked_nonnegative(array.n_buffers, "ArrowArray.n_buffers") + n_children = _checked_nonnegative(array.n_children, "ArrowArray.n_children") + n_children > _CDATA_MAX_CHILDREN && + throw(ArgumentError("ArrowArray.n_children exceeds the import limit")) + schema_children = _checked_nonnegative(schema.n_children, "ArrowSchema.n_children") + schema_children > _CDATA_MAX_CHILDREN && + throw(ArgumentError("ArrowSchema.n_children exceeds the import limit")) + schema_children == n_children || + throw(ArgumentError("ArrowSchema and ArrowArray child counts differ")) + null_count = array.null_count + if !(null_count == -1 || 0 <= null_count <= array.length) + throw(ArgumentError("ArrowArray.null_count is out of range")) + end + if (schema.dictionary == C_NULL) != (array.dictionary == C_NULL) + throw(ArgumentError("ArrowSchema and ArrowArray dictionary pointers differ")) + end + schema.dictionary == C_NULL || + throw(ArgumentError("dictionary encoded Arrow C Data import is not supported")) + return len, offset, Int(null_count), n_children +end + +function _validate_node( + schema_ptr::Ptr{ArrowSchema}, + array_ptr::Ptr{ArrowArray}; + top_level::Bool=false, + depth::Int=1, +) + depth > _CDATA_MAX_DEPTH && + throw(ArgumentError("Arrow C Data nesting exceeds the import limit")) + schema_ptr == C_NULL && throw(ArgumentError("ArrowSchema pointer is NULL")) + array_ptr == C_NULL && throw(ArgumentError("ArrowArray pointer is NULL")) + schema = unsafe_load(schema_ptr) + array = unsafe_load(array_ptr) + len, offset, null_count, n_children = _validate_common(schema, array, top_level) + format = _parse_c_data_format( + _unsafe_string_bounded( + schema.format, + _CDATA_MAX_FORMAT_BYTES, + "ArrowSchema.format", + ), + ) + expected_children = _expected_children(format) + if expected_children !== nothing && n_children != expected_children + throw(ArgumentError("Arrow C Data child count does not match the format")) + end + buffers = _load_buffers(array, _expected_buffers(format)) + schema_child_ptrs, array_child_ptrs = _load_child_ptrs(schema, array, n_children) + children = CDataNode[] + for i in eachindex(schema_child_ptrs) + push!( + children, + _validate_node( + schema_child_ptrs[i], + array_child_ptrs[i]; + top_level=false, + depth=depth + 1, + ), + ) + end + node = CDataNode( + schema_ptr, + array_ptr, + schema, + array, + format, + _load_name(schema.name), + _parse_c_metadata(schema.metadata), + buffers, + children, + len, + offset, + null_count, + ) + _validate_layout(node) + return node +end + +function _validate_layout(node::CDataNode) + total = _checked_add(node.offset, node.len, "ArrowArray offset plus length") + validity_bytes = cld(total, 8) + if _expected_buffers(node.format) > 0 + validity = node.buffers[1] + if node.null_count == -1 + if node.len > 0 && validity == C_NULL + throw(ArgumentError("unknown null count requires a validity bitmap")) + end + elseif node.null_count > 0 && validity == C_NULL + throw(ArgumentError("null values require a validity bitmap")) + end + validity_bytes >= 0 || throw(ArgumentError("invalid validity bitmap size")) + end + _validate_data_layout(node.format, node, total) + _validate_value_layout(node) + return +end + +function _aligned(ptr::Ptr{Cvoid}, ::Type{T}) where {T} + return UInt(ptr) % Base.datatype_alignment(T) == 0 +end + +function _validate_data_layout(::CDataNullFormat, node::CDataNode, total::Int) + return +end + +function _validate_data_layout(::CDataBoolFormat, node::CDataNode, total::Int) + nbytes = cld(total, 8) + nbytes > 0 && + node.buffers[2] == C_NULL && + throw(ArgumentError("boolean data buffer is NULL")) + return +end + +function _validate_data_layout(format::CDataPrimitiveFormat, node::CDataNode, total::Int) + nbytes = _checked_mul(total, sizeof(format.storage), "primitive data byte count") + if nbytes > 0 + node.buffers[2] == C_NULL && throw(ArgumentError("primitive data buffer is NULL")) + end + return +end + +function _validate_data_layout( + format::CDataFixedSizeBinaryFormat, + node::CDataNode, + total::Int, +) + nbytes = _checked_mul(total, format.bytewidth, "fixed size binary byte count") + nbytes > 0 && + node.buffers[2] == C_NULL && + throw(ArgumentError("fixed size binary data buffer is NULL")) + return +end + +function _validate_data_layout( + format::CDataBinaryFormat{O}, + node::CDataNode, + total::Int, +) where {O} + node.buffers[2] == C_NULL && throw(ArgumentError("offset buffer is NULL")) + first, last = _validate_offsets(Ptr{O}(node.buffers[2]), node.offset, node.len) + last < first && throw(ArgumentError("offsets are not monotonic")) + last > first && node.buffers[3] == C_NULL && throw(ArgumentError("data buffer is NULL")) + format.juliatype === String && _validate_utf8_offsets( + Ptr{O}(node.buffers[2]), + node.offset, + node.len, + node.buffers[3], + first, + last, + ) + return +end + +function _validate_data_layout( + format::CDataListFormat{O}, + node::CDataNode, + total::Int, +) where {O} + node.buffers[2] == C_NULL && throw(ArgumentError("offset buffer is NULL")) + first, last = _validate_offsets(Ptr{O}(node.buffers[2]), node.offset, node.len) + last < first && throw(ArgumentError("offsets are not monotonic")) + child = node.children[1] + last <= child.len || throw(ArgumentError("list offset exceeds child length")) + return +end + +function _validate_data_layout( + format::CDataFixedSizeListFormat, + node::CDataNode, + total::Int, +) + required = _checked_mul(total, format.listsize, "fixed size list child length") + node.children[1].len >= required || + throw(ArgumentError("fixed size list child is too short")) + return +end + +function _validate_data_layout(::CDataStructFormat, node::CDataNode, total::Int) + # Per the Arrow C Data Interface spec, struct children must cover length + offset. + for child in node.children + child.len >= total || throw(ArgumentError("struct child is too short")) + end + return +end + +function _count_bitmap_nulls(node::CDataNode) + if _expected_buffers(node.format) == 0 || node.len == 0 || node.buffers[1] == C_NULL + return 0 + end + nbytes = cld(_checked_add(node.offset, node.len, "validity bitmap length"), 8) + bytes = unsafe_wrap(Array, Ptr{UInt8}(node.buffers[1]), nbytes; own=false) + return _count_nulls(CDataValidity(bytes, node.offset, node.len, 0)) +end + +function _validate_value_layout(node::CDataNode) + if _expected_buffers(node.format) > 0 && + node.null_count >= 0 && + node.buffers[1] != C_NULL + actual = _count_bitmap_nulls(node) + actual == node.null_count || + throw(ArgumentError("ArrowArray.null_count does not match the validity bitmap")) + end + return +end + +function _validate_offsets(ptr::Ptr{O}, offset::Int, len::Int) where {O} + last_index = _checked_add(_checked_add(offset, len, "offset index"), 1, "offset index") + _checked_mul(last_index, sizeof(O), "offset buffer byte count") + _checked_mul(offset, sizeof(O), "offset buffer byte offset") + first = _offset_to_int(_unsafe_load_offset(ptr, offset + 1), "offset") + prev = first + for i = (offset + 2):last_index + cur = _offset_to_int(_unsafe_load_offset(ptr, i), "offset") + cur < prev && throw(ArgumentError("offsets are not monotonic")) + prev = cur + end + return first, prev +end + +@inline function _unsafe_load_int64(p::Ptr{UInt8}) + b1 = UInt64(unsafe_load(p, 1)) + b2 = UInt64(unsafe_load(p, 2)) + b3 = UInt64(unsafe_load(p, 3)) + b4 = UInt64(unsafe_load(p, 4)) + b5 = UInt64(unsafe_load(p, 5)) + b6 = UInt64(unsafe_load(p, 6)) + b7 = UInt64(unsafe_load(p, 7)) + b8 = UInt64(unsafe_load(p, 8)) + u = + ENDIAN_BOM == 0x04030201 ? + b1 | (b2 << 8) | (b3 << 16) | (b4 << 24) | (b5 << 32) | (b6 << 40) | (b7 << 48) | + (b8 << 56) : + ENDIAN_BOM == 0x01020304 ? + (b1 << 56) | (b2 << 48) | (b3 << 40) | (b4 << 32) | (b5 << 24) | (b6 << 16) | + (b7 << 8) | b8 : error("unsupported host byte order") + return reinterpret(Int64, u) +end + +@inline function _unsafe_load_offset(ptr::Ptr{Int32}, i::Int) + return _unsafe_load_int32( + Ptr{UInt8}(ptr) + _checked_mul(i - 1, sizeof(Int32), "offset byte index"), + ) +end + +@inline function _unsafe_load_offset(ptr::Ptr{Int64}, i::Int) + return _unsafe_load_int64( + Ptr{UInt8}(ptr) + _checked_mul(i - 1, sizeof(Int64), "offset byte index"), + ) +end + +function _validate_utf8_offsets( + ptr::Ptr{O}, + offset::Int, + len::Int, + data_ptr::Ptr{Cvoid}, + first::Int, + last::Int, +) where {O} + last == first && return + bytes = unsafe_wrap(Array, Ptr{UInt8}(data_ptr), last; own=false) + prev = _offset_to_int(_unsafe_load_offset(ptr, offset + 1), "offset") + for i = (offset + 2):(offset + len + 1) + cur = _offset_to_int(_unsafe_load_offset(ptr, i), "offset") + if cur > prev + if !isvalid(String, @view bytes[(prev + 1):cur]) + throw(ArgumentError("UTF-8 data is invalid")) + end + end + prev = cur + end + return +end + +function _offset_to_int(x, name) + x < 0 && throw(ArgumentError("$name is negative")) + x > typemax(Int) && throw(ArgumentError("$name exceeds the Julia Int range")) + return Int(x) +end + +function _make_validity(node::CDataNode) + if _expected_buffers(node.format) == 0 + return CDataValidity(UInt8[], 0, node.len, 0) + end + ptr = node.buffers[1] + if node.null_count == 0 || node.len == 0 + return CDataValidity(UInt8[], 0, node.len, 0) + end + nbytes = cld(_checked_add(node.offset, node.len, "validity bitmap length"), 8) + bytes = unsafe_wrap(Array, Ptr{UInt8}(ptr), nbytes; own=false) + validity = CDataValidity(bytes, node.offset, node.len, max(node.null_count, 0)) + if node.null_count == -1 + return CDataValidity(bytes, node.offset, node.len, _count_nulls(validity)) + else + return validity + end +end + +function _copy_aligned_data(ptr::Ptr{Cvoid}, ::Type{T}, offset::Int, len::Int) where {T} + len == 0 && return T[] + nbytes = _checked_mul(len, sizeof(T), "data buffer byte count") + out = Vector{T}(undef, len) + src = Ptr{UInt8}(ptr) + _checked_mul(offset, sizeof(T), "data buffer byte offset") + GC.@preserve out unsafe_copyto!(Ptr{UInt8}(pointer(out)), src, nbytes) + return out +end + +function _wrap_data(ptr::Ptr{Cvoid}, ::Type{T}, offset::Int, len::Int) where {T} + len == 0 && return T[] + # Mirror arrow-rs: copy only misaligned fixed width buffers into aligned storage. + !_aligned(ptr, T) && return _copy_aligned_data(ptr, T, offset, len) + p = Ptr{T}(ptr) + _checked_mul(offset, sizeof(T), "data buffer byte offset") + return unsafe_wrap(Array, p, len; own=false) +end + +function _wrap_offsets(ptr::Ptr{Cvoid}, ::Type{T}, offset::Int, len::Int) where {T} + return _wrap_data(ptr, T, offset, len) +end + +function _import_node(node::CDataNode, owner::CDataOwner, convert::Bool) + return _import_node(node.format, node, owner, convert) +end + +function _import_node(::CDataNullFormat, node::CDataNode, owner::CDataOwner, convert::Bool) + return CDataNull{Missing}(owner, node.len, node.metadata) +end + +function _import_node( + format::CDataPrimitiveFormat, + node::CDataNode, + owner::CDataOwner, + convert::Bool, +) + validity = _make_validity(node) + nullable = _nullable(node.schema, validity.null_count) + T = _julia_type(format.storage, nullable, convert) + data = _wrap_data(node.buffers[2], format.storage, node.offset, node.len) + return CDataPrimitive{T,format.storage,typeof(data)}( + owner, + validity, + data, + node.metadata, + ) +end + +function _import_node(::CDataBoolFormat, node::CDataNode, owner::CDataOwner, convert::Bool) + validity = _make_validity(node) + nullable = _nullable(node.schema, validity.null_count) + T = nullable ? Union{Bool,Missing} : Bool + nbytes = cld(_checked_add(node.offset, node.len, "boolean bitmap length"), 8) + data = + nbytes == 0 ? UInt8[] : + unsafe_wrap(Array, Ptr{UInt8}(node.buffers[2]), nbytes; own=false) + return CDataBool{T}(owner, validity, data, node.offset, node.len, node.metadata) +end + +function _import_node( + format::CDataBinaryFormat{O}, + node::CDataNode, + owner::CDataOwner, + convert::Bool, +) where {O} + validity = _make_validity(node) + nullable = _nullable(node.schema, validity.null_count) + T = nullable ? Union{format.juliatype,Missing} : format.juliatype + offsets = _wrap_offsets(node.buffers[2], O, node.offset, node.len + 1) + data_len = offsets[end] == offsets[1] ? 0 : Int(offsets[end]) + data = + data_len == 0 ? UInt8[] : + unsafe_wrap(Array, Ptr{UInt8}(node.buffers[3]), data_len; own=false) + return CDataBinary{T,O,typeof(offsets)}( + owner, + validity, + offsets, + data, + node.len, + node.metadata, + ) +end + +function _import_node( + format::CDataFixedSizeBinaryFormat, + node::CDataNode, + owner::CDataOwner, + convert::Bool, +) + validity = _make_validity(node) + nullable = _nullable(node.schema, validity.null_count) + storage = NTuple{format.bytewidth,UInt8} + T = nullable ? Union{storage,Missing} : storage + nbytes = _checked_mul(node.len, format.bytewidth, "fixed size binary byte count") + data = + nbytes == 0 ? UInt8[] : + unsafe_wrap( + Array, + Ptr{UInt8}(node.buffers[2]) + + _checked_mul(node.offset, format.bytewidth, "fixed size binary byte offset"), + nbytes; + own=false, + ) + return CDataFixedSizeBinary{T}( + owner, + validity, + data, + format.bytewidth, + node.len, + node.metadata, + ) +end + +function _import_node( + format::CDataListFormat{O}, + node::CDataNode, + owner::CDataOwner, + convert::Bool, +) where {O} + validity = _make_validity(node) + child = _import_node(node.children[1], owner, convert) + nullable = _nullable(node.schema, validity.null_count) + storage = Vector{eltype(child)} + T = nullable ? Union{storage,Missing} : storage + offsets = _wrap_offsets(node.buffers[2], O, node.offset, node.len + 1) + return CDataList{T,O,typeof(offsets),typeof(child)}( + owner, + validity, + offsets, + child, + node.len, + node.metadata, + ) +end + +function _import_node( + format::CDataFixedSizeListFormat, + node::CDataNode, + owner::CDataOwner, + convert::Bool, +) + validity = _make_validity(node) + child = _import_node(node.children[1], owner, convert) + nullable = _nullable(node.schema, validity.null_count) + storage = NTuple{format.listsize,eltype(child)} + T = nullable ? Union{storage,Missing} : storage + return CDataFixedSizeList{T,typeof(child)}( + owner, + validity, + child, + format.listsize, + node.offset, + node.len, + node.metadata, + ) +end + +function _import_node( + ::CDataStructFormat, + node::CDataNode, + owner::CDataOwner, + convert::Bool, +) + validity = _make_validity(node) + children = Tuple(_import_node(child, owner, convert) for child in node.children) + names = Tuple( + Symbol(child.name === nothing ? "f$(i)" : child.name) for + (i, child) in enumerate(node.children) + ) + types = Tuple(eltype(child) for child in children) + storage = NamedTuple{names,Tuple{types...}} + nullable = _nullable(node.schema, validity.null_count) + T = nullable ? Union{storage,Missing} : storage + return CDataStruct{T,typeof(children),names}( + owner, + validity, + children, + node.offset, + node.len, + node.metadata, + ) +end + +function _slice_for_table(child::CDataVector, offset::Int, len::Int) + if offset == 0 && length(child) == len + return child + else + first = offset + 1 + return CDataSlice{eltype(child),typeof(child)}(child, first, len) + end +end + +function _mask_for_struct(child::CDataVector, parent_validity::CDataValidity) + if parent_validity.null_count == 0 + return child + else + # Per the Arrow struct validity spec, parent and child validity are independent. + T = Union{eltype(child),Missing} + return CDataMasked{T,typeof(child)}(child, parent_validity) + end +end + +function _table_from_struct(node::CDataNode, owner::CDataOwner, convert::Bool) + parent_validity = _make_validity(node) + columns = AbstractVector[] + names = Symbol[] + types = Type[] + for (i, child_node) in enumerate(node.children) + child = _import_node(child_node, owner, convert) + col = _mask_for_struct( + _slice_for_table(child, node.offset, node.len), + parent_validity, + ) + name = Symbol(child_node.name === nothing ? "f$(i)" : child_node.name) + push!(names, name) + push!(types, eltype(col)) + push!(columns, col) + end + lookup = Dict{Symbol,AbstractVector}(names[i] => columns[i] for i in eachindex(names)) + return CDataTable(names, types, columns, lookup, node.metadata, owner, node.len) +end + +""" + Arrow.from_c_data(schema_ptr, array_ptr; convert=true) + +Import an Arrow C Data Interface schema and array pair. + +Aligned imported buffers are viewed without copying and are released by +`release_c_data` or finalization. Misaligned fixed width data buffers are copied +into aligned Julia storage before typed access. Use `copy` or `collect` on +imported arrays to make Julia owned arrays. + +Per the Arrow C Data Interface spec, producers describe buffer sizes through the +schema, length, and offset fields. The importer validates those layout facts and +does not inspect allocator metadata for foreign pointers. + +A top level struct array is returned as a Tables.jl column table. +""" +function from_c_data( + schema_ptr::Ptr{ArrowSchema}, + array_ptr::Ptr{ArrowArray}; + convert::Bool=true, +) + owner = CDataOwner(schema_ptr, array_ptr) + try + node = _validate_node(schema_ptr, array_ptr; top_level=true) + if node.format isa CDataStructFormat + return _table_from_struct(node, owner, convert) + else + return _import_node(node, owner, convert) + end + catch + release_c_data(owner) + rethrow() + end +end + +from_c_data(schema_ptr::Ptr{Cvoid}, array_ptr::Ptr{Cvoid}; kw...) = + from_c_data(Ptr{ArrowSchema}(schema_ptr), Ptr{ArrowArray}(array_ptr); kw...) diff --git a/test/cdata.jl b/test/cdata.jl new file mode 100644 index 00000000..1d49c102 --- /dev/null +++ b/test/cdata.jl @@ -0,0 +1,1335 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +function _cdata_release_schema(ptr::Ptr{Arrow.ArrowSchema}) + schema = unsafe_load(ptr) + unsafe_store!( + ptr, + Arrow.ArrowSchema( + schema.format, + schema.name, + schema.metadata, + schema.flags, + schema.n_children, + schema.children, + schema.dictionary, + C_NULL, + schema.private_data, + ), + ) + return +end + +function _cdata_release_array(ptr::Ptr{Arrow.ArrowArray}) + array = unsafe_load(ptr) + unsafe_store!( + ptr, + Arrow.ArrowArray( + array.length, + array.null_count, + array.offset, + array.n_buffers, + array.n_children, + array.buffers, + array.children, + array.dictionary, + C_NULL, + array.private_data, + ), + ) + return +end + +using Dates + +const _CDATA_RELEASE_SCHEMA = + @cfunction(_cdata_release_schema, Cvoid, (Ptr{Arrow.ArrowSchema},)) +const _CDATA_RELEASE_ARRAY = + @cfunction(_cdata_release_array, Cvoid, (Ptr{Arrow.ArrowArray},)) +const _CDATA_REENTRANT_OBJECT = Ref{Any}(nothing) +const _CDATA_REENTRANT_RELEASES = Ref(0) + +function _cdata_release_array_reentrant(ptr::Ptr{Arrow.ArrowArray}) + _CDATA_REENTRANT_RELEASES[] += 1 + x = _CDATA_REENTRANT_OBJECT[] + x === nothing || Arrow.release_c_data(x) + _cdata_release_array(ptr) + return +end + +const _CDATA_RELEASE_ARRAY_REENTRANT = + @cfunction(_cdata_release_array_reentrant, Cvoid, (Ptr{Arrow.ArrowArray},)) + +mutable struct CDataFixture + schema::Ref{Arrow.ArrowSchema} + array::Ref{Arrow.ArrowArray} + roots::Vector{Any} +end + +const _CDATA_FIXTURE_ROOTS = Any[] + +_schema_ptr(x::CDataFixture) = Base.unsafe_convert(Ptr{Arrow.ArrowSchema}, x.schema) +_array_ptr(x::CDataFixture) = Base.unsafe_convert(Ptr{Arrow.ArrowArray}, x.array) + +function _cstring_root(s::Union{Nothing,String}, roots) + s === nothing && return Cstring(C_NULL) + bytes = Vector{UInt8}(s * "\0") + push!(roots, bytes) + return Cstring(pointer(bytes)) +end + +function _metadata_root(meta::Union{Nothing,AbstractDict}, roots) + meta === nothing && return Cstring(C_NULL) + io = IOBuffer() + write(io, Int32(length(meta))) + for (k, v) in meta + kb = codeunits(String(k)) + vb = codeunits(String(v)) + write(io, Int32(length(kb))) + write(io, kb) + write(io, Int32(length(vb))) + write(io, vb) + end + bytes = take!(io) + push!(roots, bytes) + return Cstring(pointer(bytes)) +end + +function _cdata_fixture( + fmt::String, + len::Integer, + buffers::Vector{Ptr{Cvoid}}; + name=nothing, + metadata=nothing, + flags::Int64=0, + null_count::Int64=0, + offset::Int64=0, + children::Vector{CDataFixture}=CDataFixture[], +) + roots = Any[] + append!(roots, children) + schema_ptrs = + [Base.unsafe_convert(Ptr{Arrow.ArrowSchema}, child.schema) for child in children] + array_ptrs = + [Base.unsafe_convert(Ptr{Arrow.ArrowArray}, child.array) for child in children] + if !isempty(children) + push!(roots, schema_ptrs) + push!(roots, array_ptrs) + end + buffer_ptrs = copy(buffers) + if !isempty(buffer_ptrs) + push!(roots, buffer_ptrs) + end + schema = Ref( + Arrow.ArrowSchema( + _cstring_root(fmt, roots), + _cstring_root(name, roots), + _metadata_root(metadata, roots), + flags, + Int64(length(children)), + isempty(children) ? Ptr{Ptr{Arrow.ArrowSchema}}(C_NULL) : + Ptr{Ptr{Arrow.ArrowSchema}}(pointer(schema_ptrs)), + Ptr{Arrow.ArrowSchema}(C_NULL), + _CDATA_RELEASE_SCHEMA, + Ptr{Cvoid}(C_NULL), + ), + ) + array = Ref( + Arrow.ArrowArray( + Int64(len), + null_count, + offset, + Int64(length(buffer_ptrs)), + Int64(length(children)), + isempty(buffer_ptrs) ? Ptr{Ptr{Cvoid}}(C_NULL) : + Ptr{Ptr{Cvoid}}(pointer(buffer_ptrs)), + isempty(children) ? Ptr{Ptr{Arrow.ArrowArray}}(C_NULL) : + Ptr{Ptr{Arrow.ArrowArray}}(pointer(array_ptrs)), + Ptr{Arrow.ArrowArray}(C_NULL), + _CDATA_RELEASE_ARRAY, + Ptr{Cvoid}(C_NULL), + ), + ) + push!(roots, schema) + push!(roots, array) + fixture = CDataFixture(schema, array, roots) + push!(_CDATA_FIXTURE_ROOTS, fixture) + return fixture +end + +function _primitive_fixture( + fmt, + data::Vector{T}; + validity=nothing, + null_count::Int64=0, + flags::Int64=0, + offset::Int64=0, + len::Int=length(data) - Int(offset), + name=nothing, + metadata=nothing, +) where {T} + roots = Any[data] + buffers = Ptr{Cvoid}[ + validity === nothing ? Ptr{Cvoid}(C_NULL) : Ptr{Cvoid}(pointer(validity)), + isempty(data) ? Ptr{Cvoid}(C_NULL) : Ptr{Cvoid}(pointer(data)), + ] + validity !== nothing && push!(roots, validity) + fixture = _cdata_fixture( + fmt, + len, + buffers; + name=name, + metadata=metadata, + flags=flags, + null_count=null_count, + offset=offset, + ) + append!(fixture.roots, roots) + return fixture +end + +function _unaligned_primitive_fixture( + fmt, + data::Vector{T}; + offset::Int64=Int64(0), + len::Int64=Int64(length(data)) - offset, +) where {T} + nbytes = length(data) * sizeof(T) + bytes = Vector{UInt8}(undef, nbytes + 1) + unsafe_copyto!(pointer(bytes, 2), Ptr{UInt8}(pointer(data)), nbytes) + buffers = Ptr{Cvoid}[Ptr{Cvoid}(C_NULL), Ptr{Cvoid}(pointer(bytes, 2))] + fixture = _cdata_fixture(fmt, len, buffers; offset=offset) + append!(fixture.roots, Any[data, bytes]) + return fixture +end + +function _set_schema!( + f::CDataFixture; + format=f.schema[].format, + name=f.schema[].name, + metadata=f.schema[].metadata, + flags=f.schema[].flags, + n_children=f.schema[].n_children, + children=f.schema[].children, + dictionary=f.schema[].dictionary, + release=f.schema[].release, + private_data=f.schema[].private_data, +) + f.schema[] = Arrow.ArrowSchema( + format, + name, + metadata, + flags, + n_children, + children, + dictionary, + release, + private_data, + ) + return f +end + +function _set_array!( + f::CDataFixture; + length=f.array[].length, + null_count=f.array[].null_count, + offset=f.array[].offset, + n_buffers=f.array[].n_buffers, + buffers=f.array[].buffers, + n_children=f.array[].n_children, + children=f.array[].children, + dictionary=f.array[].dictionary, + release=f.array[].release, + private_data=f.array[].private_data, +) + f.array[] = Arrow.ArrowArray( + length, + null_count, + offset, + n_buffers, + n_children, + buffers, + children, + dictionary, + release, + private_data, + ) + return f +end + +@testset "Arrow C Data Interface import" begin + @testset "ABI layout" begin + ptr = sizeof(Ptr{Cvoid}) + @test sizeof(Arrow.ArrowSchema) == 7 * ptr + 2 * sizeof(Int64) + @test sizeof(Arrow.ArrowArray) == 5 * ptr + 5 * sizeof(Int64) + + cc = Sys.which("cc") + if cc === nothing + @test_skip "C compiler not available" + else + c_src = """ + #include + #include + #include + struct ArrowSchema { + const char* format; + const char* name; + const char* metadata; + int64_t flags; + int64_t n_children; + struct ArrowSchema** children; + struct ArrowSchema* dictionary; + void (*release)(struct ArrowSchema*); + void* private_data; + }; + struct ArrowArray { + int64_t length; + int64_t null_count; + int64_t offset; + int64_t n_buffers; + int64_t n_children; + const void** buffers; + struct ArrowArray** children; + struct ArrowArray* dictionary; + void (*release)(struct ArrowArray*); + void* private_data; + }; + int main(void) { + printf("%zu %zu %zu %zu %zu %zu %zu %zu %zu\\n", + offsetof(struct ArrowSchema, format), + offsetof(struct ArrowSchema, name), + offsetof(struct ArrowSchema, metadata), + offsetof(struct ArrowSchema, flags), + offsetof(struct ArrowSchema, n_children), + offsetof(struct ArrowSchema, children), + offsetof(struct ArrowSchema, dictionary), + offsetof(struct ArrowSchema, release), + offsetof(struct ArrowSchema, private_data)); + printf("%zu %zu %zu %zu %zu %zu %zu %zu %zu %zu\\n", + offsetof(struct ArrowArray, length), + offsetof(struct ArrowArray, null_count), + offsetof(struct ArrowArray, offset), + offsetof(struct ArrowArray, n_buffers), + offsetof(struct ArrowArray, n_children), + offsetof(struct ArrowArray, buffers), + offsetof(struct ArrowArray, children), + offsetof(struct ArrowArray, dictionary), + offsetof(struct ArrowArray, release), + offsetof(struct ArrowArray, private_data)); + return 0; + } + """ + mktempdir() do dir + src = joinpath(dir, "layout.c") + bin = joinpath(dir, "layout") + write(src, c_src) + run(`$cc -o $bin $src`) + lines = split(readchomp(`$bin`), '\n') + schema_offsets = parse.(Int, split(lines[1])) + array_offsets = parse.(Int, split(lines[2])) + for (i, offset) in enumerate(schema_offsets) + @test fieldoffset(Arrow.ArrowSchema, i) == offset + end + for (i, offset) in enumerate(array_offsets) + @test fieldoffset(Arrow.ArrowArray, i) == offset + end + end + end + end + + @testset "primitive arrays" begin + data = Int32[1, 2, 3] + f = _primitive_fixture("i", data) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test x isa Arrow.CDataVector + @test collect(x) == Int32[1, 2, 3] + @test copy(x) == Int32[1, 2, 3] + data[1] = 99 + @test collect(x) == Int32[99, 2, 3] + + f = _unaligned_primitive_fixture("i", Int32[7, 8, 9]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test collect(x) == Int32[7, 8, 9] + fill!(f.roots[end], 0x00) + @test collect(x) == Int32[7, 8, 9] + f = _unaligned_primitive_fixture("i", Int32[6, 7, 8, 9]; offset=Int64(1), len=2) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test collect(x) == Int32[7, 8] + fill!(f.roots[end], 0x00) + @test collect(x) == Int32[7, 8] + + validity = UInt8[0b00000101] + f = _primitive_fixture( + "i", + Int32[10, 20, 30]; + validity=validity, + null_count=Int64(1), + flags=Arrow.ARROW_FLAG_NULLABLE, + ) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test isequal(collect(x), Union{Int32,Missing}[10, missing, 30]) + + f = _primitive_fixture( + "i", + Int32[10, 20, 30]; + validity=validity, + null_count=Int64(-1), + flags=Arrow.ARROW_FLAG_NULLABLE, + ) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test Arrow.nullcount(x) == 1 + @test isequal(collect(x), Union{Int32,Missing}[10, missing, 30]) + end + + @testset "null arrays" begin + f = _cdata_fixture("n", 3, Ptr{Cvoid}[]; null_count=Int64(3)) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test x isa Arrow.CDataVector{Missing} + @test size(x) == (3,) + @test Arrow.validitybitmap(x) === nothing + @test Arrow.nullcount(x) == 3 + @test x[2] === missing + @test isequal(collect(x), [missing, missing, missing]) + node = Arrow._validate_node(_schema_ptr(f), _array_ptr(f); top_level=true) + validity = Arrow._make_validity(node) + @test validity.bytes == UInt8[] + @test validity.null_count == 0 + end + + @testset "bool bit offsets" begin + data = UInt8[0b10110110] + validity = UInt8[0xff] + f = _cdata_fixture( + "b", + 5, + Ptr{Cvoid}[Ptr{Cvoid}(pointer(validity)), Ptr{Cvoid}(pointer(data))]; + offset=Int64(3), + ) + append!(f.roots, Any[data, validity]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test collect(x) == [false, true, true, false, true] + + validity = UInt8[0b00001101] + f = _cdata_fixture( + "b", + 4, + Ptr{Cvoid}[Ptr{Cvoid}(pointer(validity)), Ptr{Cvoid}(pointer(data))]; + flags=Arrow.ARROW_FLAG_NULLABLE, + null_count=Int64(-1), + ) + append!(f.roots, Any[data, validity]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test Arrow.nullcount(x) == 1 + @test isequal(collect(x), Union{Bool,Missing}[false, missing, true, false]) + end + + @testset "string and binary arrays" begin + offsets = Int32[0, 3, 3, 6] + bytes = Vector{UInt8}(codeunits("abcdef")) + f = _cdata_fixture( + "u", + 3, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets)), Ptr{Cvoid}(pointer(bytes))], + ) + append!(f.roots, Any[offsets, bytes]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test collect(x) == ["abc", "", "def"] + offsets[2] = 2 + @test collect(x) == ["ab", "c", "def"] + + offset_values = Int32[0, 1, 3] + offset_nbytes = length(offset_values) * sizeof(Int32) + offset_bytes = Vector{UInt8}(undef, offset_nbytes + 1) + unsafe_copyto!( + pointer(offset_bytes, 2), + Ptr{UInt8}(pointer(offset_values)), + offset_nbytes, + ) + bytes = Vector{UInt8}(codeunits("abc")) + f = _cdata_fixture( + "u", + 2, + Ptr{Cvoid}[ + C_NULL, + Ptr{Cvoid}(pointer(offset_bytes, 2)), + Ptr{Cvoid}(pointer(bytes)), + ], + ) + append!(f.roots, Any[offset_bytes, bytes]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + fill!(offset_bytes, 0x00) + @test collect(x) == ["a", "bc"] + + offsets64 = Int64[0, 2, 5] + bytes2 = Vector{UInt8}(codeunits("hello")) + f = _cdata_fixture( + "U", + 2, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets64)), Ptr{Cvoid}(pointer(bytes2))], + ) + append!(f.roots, Any[offsets64, bytes2]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test collect(x) == ["he", "llo"] + + bad_utf8 = UInt8[0xff] + offsets = Int32[0, 1] + f = _cdata_fixture( + "u", + 1, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets)), Ptr{Cvoid}(pointer(bad_utf8))], + ) + append!(f.roots, Any[offsets, bad_utf8]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + offsets = Int32[0, 2, 3] + bytes = UInt8[0x01, 0x02, 0xff] + f = _cdata_fixture( + "z", + 2, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets)), Ptr{Cvoid}(pointer(bytes))], + ) + append!(f.roots, Any[offsets, bytes]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test collect(x) == [b"\x01\x02", b"\xff"] + + offsets64 = Int64[0, 1, 3] + f = _cdata_fixture( + "Z", + 2, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets64)), Ptr{Cvoid}(pointer(bytes))], + ) + append!(f.roots, Any[offsets64, bytes]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test collect(x) == [b"\x01", b"\x02\xff"] + + offsets = Int32[123] + f = _cdata_fixture("u", 0, Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets)), C_NULL]) + push!(f.roots, offsets) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test collect(x) == String[] + + offsets = Int32[123, 123, 123] + f = _cdata_fixture("u", 2, Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets)), C_NULL]) + push!(f.roots, offsets) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test collect(x) == ["", ""] + + offsets = Int32[0, 1, 2, 3] + bytes = Vector{UInt8}(codeunits("abc")) + validity = UInt8[0b00000101] + f = _cdata_fixture( + "u", + 3, + Ptr{Cvoid}[ + Ptr{Cvoid}(pointer(validity)), + Ptr{Cvoid}(pointer(offsets)), + Ptr{Cvoid}(pointer(bytes)), + ]; + flags=Arrow.ARROW_FLAG_NULLABLE, + null_count=Int64(-1), + ) + append!(f.roots, Any[offsets, bytes, validity]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + y = copy(x) + Arrow.release_c_data(x) + @test isequal(y, Union{String,Missing}["a", missing, "c"]) + @test_throws ArgumentError x[1] + end + + @testset "fixed size binary" begin + data = UInt8[0x01, 0x02, 0x03, 0x04, 0x05, 0x06] + f = _cdata_fixture("w:3", 2, Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(data))]) + push!(f.roots, data) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test collect(x) == [(0x01, 0x02, 0x03), (0x04, 0x05, 0x06)] + end + + @testset "list of primitives" begin + child = _primitive_fixture("i", Int32[1, 2, 3, 4, 5]) + offsets = Int32[0, 2, 5] + f = _cdata_fixture( + "+l", + 2, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets))]; + children=[child], + ) + push!(f.roots, offsets) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test collect(x) == [Int32[1, 2], Int32[3, 4, 5]] + + child = _primitive_fixture("i", Int32[10, 20, 30, 40, 50, 60]) + offsets = Int64[0, 1, 3, 6] + f = _cdata_fixture( + "+L", + 2, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets))]; + offset=Int64(1), + children=[child], + ) + push!(f.roots, offsets) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test collect(x) == [Int32[20, 30], Int32[40, 50, 60]] + + child = _primitive_fixture("i", Int32[10, 20, 30]) + offset_values = Int64[0, 1, 3] + offset_nbytes = length(offset_values) * sizeof(Int64) + offset_bytes = Vector{UInt8}(undef, offset_nbytes + 1) + unsafe_copyto!( + pointer(offset_bytes, 2), + Ptr{UInt8}(pointer(offset_values)), + offset_nbytes, + ) + f = _cdata_fixture( + "+L", + 2, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offset_bytes, 2))]; + children=[child], + ) + push!(f.roots, offset_bytes) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + fill!(offset_bytes, 0x00) + @test collect(x) == [Int32[10], Int32[20, 30]] + + child = _primitive_fixture("i", Int32[1, 2, 3, 4, 5]) + offsets = Int32[0, 2, 2, 5] + validity = UInt8[0b00000101] + f = _cdata_fixture( + "+l", + 3, + Ptr{Cvoid}[Ptr{Cvoid}(pointer(validity)), Ptr{Cvoid}(pointer(offsets))]; + flags=Arrow.ARROW_FLAG_NULLABLE, + null_count=Int64(-1), + children=[child], + ) + append!(f.roots, Any[offsets, validity]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test isequal( + collect(x), + Union{Vector{Int32},Missing}[Int32[1, 2], missing, Int32[3, 4, 5]], + ) + end + + @testset "fixed size list" begin + child = _primitive_fixture("f", Float32[1, 2, 3, 4, 5, 6]) + f = _cdata_fixture("+w:3", 2, Ptr{Cvoid}[C_NULL]; children=[child]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f); convert=false) + @test collect(x) == [(1.0f0, 2.0f0, 3.0f0), (4.0f0, 5.0f0, 6.0f0)] + + child = _primitive_fixture("i", Int32[1, 2, 3, 4, 5, 6, 7, 8, 9]) + f = _cdata_fixture("+w:3", 2, Ptr{Cvoid}[C_NULL]; offset=Int64(1), children=[child]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f); convert=false) + @test collect(x) == [(Int32(4), Int32(5), Int32(6)), (Int32(7), Int32(8), Int32(9))] + end + + @testset "struct root table with names and metadata" begin + xchild = + _primitive_fixture("i", Int32[1, 2, 3]; name="x", metadata=Dict("unit" => "id")) + yoffsets = Int32[0, 1, 2, 3] + ybytes = Vector{UInt8}(codeunits("abc")) + ychild = _cdata_fixture( + "u", + 3, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(yoffsets)), Ptr{Cvoid}(pointer(ybytes))]; + name="y", + ) + append!(ychild.roots, Any[yoffsets, ybytes]) + root = _cdata_fixture( + "+s", + 3, + Ptr{Cvoid}[C_NULL]; + children=[xchild, ychild], + metadata=Dict("source" => "cdata"), + ) + tbl = Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + @test Tables.columnnames(tbl) == [:x, :y] + @test Tables.schema(tbl).types == (Int32, String) + @test length(tbl) == 2 + @test Tables.rowcount(tbl) == 3 + @test Tables.istable(typeof(tbl)) + @test Tables.columnaccess(typeof(tbl)) + @test Tables.columns(tbl) === tbl + @test propertynames(tbl) == (:x, :y) + @test :rowcount in propertynames(tbl, true) + @test tbl.rowcount == 3 + @test collect(Tables.getcolumn(tbl, :x)) == Int32[1, 2, 3] + @test collect(Tables.getcolumn(tbl, 1)) == Int32[1, 2, 3] + @test collect(tbl[1]) == Int32[1, 2, 3] + @test collect(tbl.y) == ["a", "b", "c"] + @test copy(tbl) == (x=Int32[1, 2, 3], y=["a", "b", "c"]) + @test DataAPI.metadatasupport(typeof(tbl)) == (read=true, write=false) + @test DataAPI.colmetadatasupport(typeof(tbl)) == (read=true, write=false) + @test DataAPI.metadata(tbl) == Dict("source" => "cdata") + @test DataAPI.metadata(tbl; style=true) == Dict("source" => ("cdata", :default)) + @test DataAPI.metadata(tbl, "source") == "cdata" + @test DataAPI.metadata(tbl, "source"; style=true) == ("cdata", :default) + @test DataAPI.metadata(tbl, "source", "fallback") == "cdata" + @test DataAPI.metadata(tbl, "source", "fallback"; style=true) == ("cdata", :default) + @test DataAPI.metadata(tbl, "missing", "fallback") == "fallback" + @test DataAPI.metadata(tbl, "missing", "fallback"; style=true) == + ("fallback", :default) + @test Set(DataAPI.metadatakeys(tbl)) == Set(["source"]) + @test DataAPI.colmetadata(tbl, :x) == Dict("unit" => "id") + @test DataAPI.colmetadata(tbl, :x; style=true) == Dict("unit" => ("id", :default)) + @test DataAPI.colmetadata(tbl, :x, "unit") == "id" + @test DataAPI.colmetadata(tbl, :x, "unit"; style=true) == ("id", :default) + @test DataAPI.colmetadata(tbl, :x, "unit", "fallback") == "id" + @test DataAPI.colmetadata(tbl, :x, "unit", "fallback"; style=true) == + ("id", :default) + @test DataAPI.colmetadata(tbl, :x, "missing", "fallback") == "fallback" + @test DataAPI.colmetadata(tbl, :x, "missing", "fallback"; style=true) == + ("fallback", :default) + @test Set(DataAPI.colmetadatakeys(tbl, :x)) == Set(["unit"]) + colkeys = collect(DataAPI.colmetadatakeys(tbl)) + @test length(colkeys) == 1 + @test first(colkeys[1]) == :x + @test Set(last(colkeys[1])) == Set(["unit"]) + @test DataAPI.colmetadata(tbl) == Dict(:x => Dict("unit" => "id")) + + a = _primitive_fixture("i", Int32[1, 2, 3]; name="a") + b = _primitive_fixture("i", Int32[4, 5, 6]; name="b") + pair = _cdata_fixture("+s", 3, Ptr{Cvoid}[C_NULL]; name="pair", children=[a, b]) + root = _cdata_fixture("+s", 3, Ptr{Cvoid}[C_NULL]; children=[pair]) + pair_tbl = Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + @test collect(pair_tbl.pair) == [(a=1, b=4), (a=2, b=5), (a=3, b=6)] + + shortmeta_child = + _primitive_fixture("i", Int32[1]; name="x", metadata=Dict("k" => "v")) + shortmeta_root = _cdata_fixture( + "+s", + 1, + Ptr{Cvoid}[C_NULL]; + children=[shortmeta_child], + metadata=Dict("x" => "yz"), + ) + shortmeta_tbl = + Arrow.from_c_data(_schema_ptr(shortmeta_root), _array_ptr(shortmeta_root)) + @test DataAPI.metadata(shortmeta_tbl, "x") == "yz" + @test DataAPI.colmetadata(shortmeta_tbl, :x, "k") == "v" + + col = Tables.getcolumn(tbl, :x) + GC.gc(true) + @test collect(col) == Int32[1, 2, 3] + end + + @testset "struct table offsets and owner roots" begin + child = _primitive_fixture("i", Int32[10, 20, 30]; name="x") + root = + _cdata_fixture("+s", 2, Ptr{Cvoid}[C_NULL]; offset=Int64(1), children=[child]) + tbl = Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + @test collect(Tables.getcolumn(tbl, :x)) == Int32[20, 30] + + col = let + tbl2 = Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + Tables.getcolumn(tbl2, :x) + end + GC.gc(true) + GC.gc(true) + @test collect(col) == Int32[20, 30] + Arrow.release_c_data(col) + + child = _primitive_fixture("i", Int32[10, 20, 30]; offset=Int64(1), len=2, name="x") + root = + _cdata_fixture("+s", 2, Ptr{Cvoid}[C_NULL]; offset=Int64(1), children=[child]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + end + + @testset "struct parent validity masks children" begin + child_validity = UInt8[0b00001011] + parent_validity = UInt8[0b00001101] + child = _primitive_fixture( + "i", + Int32[1, 2, 3, 4]; + validity=child_validity, + null_count=Int64(1), + flags=Arrow.ARROW_FLAG_NULLABLE, + name="x", + ) + root = _cdata_fixture( + "+s", + 4, + Ptr{Cvoid}[Ptr{Cvoid}(pointer(parent_validity))]; + children=[child], + null_count=Int64(1), + flags=Arrow.ARROW_FLAG_NULLABLE, + ) + push!(root.roots, parent_validity) + tbl = Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + col = Tables.getcolumn(tbl, :x) + @test Tables.schema(tbl).types == (Union{Int32,Missing},) + @test Arrow.nullcount(col) == 2 + @test isequal(collect(col), Union{Int32,Missing}[1, missing, missing, 4]) + @test isequal(copy(tbl), (x=Union{Int32,Missing}[1, missing, missing, 4],)) + + parent_validity = UInt8[0b00000010] + child = _primitive_fixture("i", Int32[10, 20, 30]; name="x") + root = _cdata_fixture( + "+s", + 2, + Ptr{Cvoid}[Ptr{Cvoid}(pointer(parent_validity))]; + offset=Int64(1), + children=[child], + null_count=Int64(1), + flags=Arrow.ARROW_FLAG_NULLABLE, + ) + push!(root.roots, parent_validity) + tbl = Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + @test isequal(collect(Tables.getcolumn(tbl, :x)), Union{Int32,Missing}[20, missing]) + end + + @testset "temporal and decimal formats" begin + dates = Arrow.DATE[Arrow.DATE(1), Arrow.DATE(2)] + f = _primitive_fixture("tdD", dates) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f); convert=false) + @test collect(x) == dates + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f); convert=true) + @test collect(x) == convert.(Dates.Date, dates) + + D = Arrow.Decimal{10,2,Int128} + decimals = D[D(123), D(-45)] + f = _primitive_fixture("d:10,2", decimals) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f); convert=false) + @test collect(x) == decimals + + timestamps = Arrow.Timestamp{Arrow.Meta.TimeUnit.MILLISECOND,nothing}[ + Arrow.Timestamp{Arrow.Meta.TimeUnit.MILLISECOND,nothing}(0), + Arrow.Timestamp{Arrow.Meta.TimeUnit.MILLISECOND,nothing}(1), + ] + f = _primitive_fixture("tsm:", timestamps) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f); convert=true) + @test collect(x) == + Dates.DateTime[Dates.DateTime(1970), Dates.DateTime(1970, 1, 1, 0, 0, 0, 1)] + + durations = Arrow.Duration{Arrow.Meta.TimeUnit.SECOND}[ + Arrow.Duration{Arrow.Meta.TimeUnit.SECOND}(1), + Arrow.Duration{Arrow.Meta.TimeUnit.SECOND}(2), + ] + f = _primitive_fixture("tDs", durations) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f); convert=true) + @test collect(x) == Dates.Second[Dates.Second(1), Dates.Second(2)] + + intervals = Arrow.Interval{Arrow.Meta.IntervalUnit.YEAR_MONTH,Int32}[ + Arrow.Interval{Arrow.Meta.IntervalUnit.YEAR_MONTH}(12), + Arrow.Interval{Arrow.Meta.IntervalUnit.YEAR_MONTH}(18), + ] + f = _primitive_fixture("tiM", intervals) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f); convert=false) + @test collect(x) == intervals + end + + @testset "release behavior" begin + f = _primitive_fixture("i", Int32[1, 2, 3]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test x[1] == 1 + Arrow.release_c_data(x) + @test f.array[].release == C_NULL + @test f.schema[].release == C_NULL + @test_nowarn Arrow.release_c_data(x) + @test_throws ArgumentError x[1] + + f = _primitive_fixture("i", Int32[1, 2, 3]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + owner = Arrow._owner(x) + owner_lock = getfield(owner, :lock) + started = Channel{Nothing}(1) + locked = false + lock(owner_lock) + locked = true + try + task = Threads.@spawn begin + put!(started, nothing) + Arrow.release_c_data(x) + end + take!(started) + sleep(0.05) + @test f.array[].release != C_NULL + unlock(owner_lock) + locked = false + wait(task) + finally + locked && unlock(owner_lock) + end + @test f.array[].release == C_NULL + @test f.schema[].release == C_NULL + + f = _primitive_fixture("i", Int32[1, 2, 3]) + _set_array!(f; release=_CDATA_RELEASE_ARRAY_REENTRANT) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + _CDATA_REENTRANT_OBJECT[] = x + _CDATA_REENTRANT_RELEASES[] = 0 + try + @test_nowarn Arrow.release_c_data(x) + @test _CDATA_REENTRANT_RELEASES[] == 1 + @test f.array[].release == C_NULL + @test f.schema[].release == C_NULL + @test_throws ArgumentError x[1] + finally + _CDATA_REENTRANT_OBJECT[] = nothing + end + end + + @testset "table release behavior" begin + child = _primitive_fixture("i", Int32[1, 2, 3]; name="x") + root = _cdata_fixture("+s", 3, Ptr{Cvoid}[C_NULL]; children=[child]) + tbl = Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + col = Tables.getcolumn(tbl, :x) + @test col[1] == 1 + Arrow.release_c_data(tbl) + @test_nowarn Arrow.release_c_data(tbl) + @test_throws ArgumentError Tables.getcolumn(tbl, :x) + @test_throws ArgumentError col[1] + end + + @testset "copy and collect own the result" begin + f = _primitive_fixture("i", Int32[1, 2, 3]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + a = collect(x) + b = copy(x) + Arrow.release_c_data(x) + @test a == Int32[1, 2, 3] + @test b == Int32[1, 2, 3] + end + + @testset "malformed inputs" begin + f = _primitive_fixture("i", Int32[1]) + @test_throws ArgumentError Arrow.from_c_data( + Ptr{Arrow.ArrowSchema}(C_NULL), + _array_ptr(f), + ) + @test_throws ArgumentError Arrow.from_c_data( + _schema_ptr(f), + Ptr{Arrow.ArrowArray}(C_NULL), + ) + + f = _cdata_fixture("?", 0, Ptr{Cvoid}[]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test f.array[].release == C_NULL + @test f.schema[].release == C_NULL + + f = _primitive_fixture("i", Int32[1]) + f.schema[] = Arrow.ArrowSchema( + f.schema[].format, + f.schema[].name, + f.schema[].metadata, + f.schema[].flags, + f.schema[].n_children, + f.schema[].children, + f.schema[].dictionary, + C_NULL, + f.schema[].private_data, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + f.array[] = Arrow.ArrowArray( + f.array[].length, + f.array[].null_count, + f.array[].offset, + f.array[].n_buffers, + f.array[].n_children, + f.array[].buffers, + f.array[].children, + f.array[].dictionary, + C_NULL, + f.array[].private_data, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + f.array[] = Arrow.ArrowArray( + -1, + 0, + 0, + 2, + 0, + f.array[].buffers, + f.array[].children, + f.array[].dictionary, + f.array[].release, + f.array[].private_data, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + f.array[] = Arrow.ArrowArray( + f.array[].length, + f.array[].null_count, + -1, + f.array[].n_buffers, + f.array[].n_children, + f.array[].buffers, + f.array[].children, + f.array[].dictionary, + f.array[].release, + f.array[].private_data, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + f.array[] = Arrow.ArrowArray( + f.array[].length, + f.array[].null_count, + f.array[].offset, + -1, + f.array[].n_children, + f.array[].buffers, + f.array[].children, + f.array[].dictionary, + f.array[].release, + f.array[].private_data, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + f.array[] = Arrow.ArrowArray( + f.array[].length, + f.array[].null_count, + f.array[].offset, + f.array[].n_buffers, + -1, + f.array[].buffers, + f.array[].children, + f.array[].dictionary, + f.array[].release, + f.array[].private_data, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + f.schema[] = Arrow.ArrowSchema( + f.schema[].format, + f.schema[].name, + f.schema[].metadata, + f.schema[].flags, + -1, + f.schema[].children, + f.schema[].dictionary, + f.schema[].release, + f.schema[].private_data, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]; null_count=Int64(2)) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]; null_count=Int64(-1)) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture( + "i", + Int32[1]; + null_count=Int64(1), + flags=Arrow.ARROW_FLAG_NULLABLE, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture( + "i", + Int32[1, 2, 3]; + validity=UInt8[0b00000111], + null_count=Int64(1), + flags=Arrow.ARROW_FLAG_NULLABLE, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture( + "i", + Int32[1, 2, 3]; + validity=UInt8[0b00000101], + null_count=Int64(0), + flags=Arrow.ARROW_FLAG_NULLABLE, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + _set_array!(f; buffers=Ptr{Ptr{Cvoid}}(C_NULL)) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _cdata_fixture("i", 1, Ptr{Cvoid}[C_NULL, C_NULL]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _cdata_fixture("i", 1, Ptr{Cvoid}[C_NULL]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + child = _primitive_fixture("i", Int32[1]) + f = _cdata_fixture("+l", 1, Ptr{Cvoid}[C_NULL, C_NULL]; children=[child]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + child = _primitive_fixture("i", Int32[1, 2]) + offsets = Int32[0, 3] + f = _cdata_fixture( + "+l", + 1, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets))]; + children=[child], + ) + push!(f.roots, offsets) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + child = _primitive_fixture("i", Int32[1, 2, 3, 4, 5]) + f = _cdata_fixture("+w:3", 2, Ptr{Cvoid}[C_NULL]; children=[child]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + offsets = Int32[0, 3, 2] + bytes = UInt8[1, 2, 3] + f = _cdata_fixture( + "z", + 2, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets)), Ptr{Cvoid}(pointer(bytes))], + ) + append!(f.roots, Any[offsets, bytes]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + offsets = Int32[0, -1] + bytes = UInt8[] + f = _cdata_fixture( + "z", + 1, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets)), Ptr{Cvoid}(C_NULL)], + ) + append!(f.roots, Any[offsets, bytes]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + offsets = Int32[0, 1] + bytes = UInt8[0x01] + f = _cdata_fixture( + "u", + 1, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets)), Ptr{Cvoid}(pointer(bytes))]; + flags=Arrow.ARROW_FLAG_NULLABLE, + null_count=Int64(-1), + ) + append!(f.roots, Any[offsets, bytes]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + badmeta = reinterpret(UInt8, Int32[1, -1]) + f = _primitive_fixture("i", Int32[1]) + f.schema[] = Arrow.ArrowSchema( + f.schema[].format, + f.schema[].name, + Cstring(pointer(badmeta)), + f.schema[].flags, + f.schema[].n_children, + f.schema[].children, + f.schema[].dictionary, + f.schema[].release, + f.schema[].private_data, + ) + push!(f.roots, badmeta) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + badmeta = reinterpret(UInt8, Int32[1, 0, -1]) + f = _primitive_fixture("i", Int32[1]) + _set_schema!(f; metadata=Cstring(pointer(badmeta))) + push!(f.roots, badmeta) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + f.schema[] = Arrow.ArrowSchema( + f.schema[].format, + f.schema[].name, + f.schema[].metadata, + f.schema[].flags, + f.schema[].n_children, + Ptr{Arrow.ArrowSchema}(UInt(0x01)), + f.schema[].dictionary, + f.schema[].release, + f.schema[].private_data, + ) + f.array[] = Arrow.ArrowArray( + f.array[].length, + f.array[].null_count, + f.array[].offset, + f.array[].n_buffers, + 1, + f.array[].buffers, + Ptr{Ptr{Arrow.ArrowArray}}(C_NULL), + f.array[].dictionary, + f.array[].release, + f.array[].private_data, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + child = _primitive_fixture("i", Int32[1]) + schema_children = Ptr{Arrow.ArrowSchema}[Ptr{Arrow.ArrowSchema}(C_NULL)] + array_children = Ptr{Arrow.ArrowArray}[_array_ptr(child)] + f = _cdata_fixture("+s", 1, Ptr{Cvoid}[C_NULL]) + _set_schema!( + f; + n_children=1, + children=Ptr{Ptr{Arrow.ArrowSchema}}(pointer(schema_children)), + ) + _set_array!( + f; + n_children=1, + children=Ptr{Ptr{Arrow.ArrowArray}}(pointer(array_children)), + ) + append!(f.roots, Any[child, schema_children, array_children]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + child = _primitive_fixture("i", Int32[1]) + schema_children = Ptr{Arrow.ArrowSchema}[_schema_ptr(child)] + array_children = Ptr{Arrow.ArrowArray}[Ptr{Arrow.ArrowArray}(C_NULL)] + f = _cdata_fixture("+s", 1, Ptr{Cvoid}[C_NULL]) + f.schema[] = Arrow.ArrowSchema( + f.schema[].format, + f.schema[].name, + f.schema[].metadata, + f.schema[].flags, + 1, + Ptr{Ptr{Arrow.ArrowSchema}}(pointer(schema_children)), + f.schema[].dictionary, + f.schema[].release, + f.schema[].private_data, + ) + f.array[] = Arrow.ArrowArray( + f.array[].length, + f.array[].null_count, + f.array[].offset, + f.array[].n_buffers, + 1, + f.array[].buffers, + Ptr{Ptr{Arrow.ArrowArray}}(pointer(array_children)), + f.array[].dictionary, + f.array[].release, + f.array[].private_data, + ) + append!(f.roots, Any[child, schema_children, array_children]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + child = _primitive_fixture("i", Int32[1]) + _set_schema!(child; release=C_NULL) + schema_children = Ptr{Arrow.ArrowSchema}[_schema_ptr(child)] + array_children = Ptr{Arrow.ArrowArray}[_array_ptr(child)] + root = _cdata_fixture("+s", 1, Ptr{Cvoid}[C_NULL]) + _set_schema!( + root; + n_children=1, + children=Ptr{Ptr{Arrow.ArrowSchema}}(pointer(schema_children)), + ) + _set_array!( + root; + n_children=1, + children=Ptr{Ptr{Arrow.ArrowArray}}(pointer(array_children)), + ) + append!(root.roots, Any[child, schema_children, array_children]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + + child = _primitive_fixture("i", Int32[1]) + f = _primitive_fixture("i", Int32[1]) + schema_children = Ptr{Arrow.ArrowSchema}[_schema_ptr(child)] + array_children = Ptr{Arrow.ArrowArray}[_array_ptr(child)] + _set_schema!( + f; + n_children=1, + children=Ptr{Ptr{Arrow.ArrowSchema}}(pointer(schema_children)), + ) + _set_array!( + f; + n_children=1, + children=Ptr{Ptr{Arrow.ArrowArray}}(pointer(array_children)), + ) + append!(f.roots, Any[child, schema_children, array_children]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + dict_schema = Ref(f.schema[]) + f.schema[] = Arrow.ArrowSchema( + f.schema[].format, + f.schema[].name, + f.schema[].metadata, + f.schema[].flags, + f.schema[].n_children, + f.schema[].children, + Base.unsafe_convert(Ptr{Arrow.ArrowSchema}, dict_schema), + f.schema[].release, + f.schema[].private_data, + ) + push!(f.roots, dict_schema) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + dict_array = Ref(f.array[]) + f.array[] = Arrow.ArrowArray( + f.array[].length, + f.array[].null_count, + f.array[].offset, + f.array[].n_buffers, + f.array[].n_children, + f.array[].buffers, + f.array[].children, + Base.unsafe_convert(Ptr{Arrow.ArrowArray}, dict_array), + f.array[].release, + f.array[].private_data, + ) + push!(f.roots, dict_array) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + dict_schema = Ref(f.schema[]) + dict_array = Ref(f.array[]) + _set_schema!(f; dictionary=Base.unsafe_convert(Ptr{Arrow.ArrowSchema}, dict_schema)) + _set_array!(f; dictionary=Base.unsafe_convert(Ptr{Arrow.ArrowArray}, dict_array)) + append!(f.roots, Any[dict_schema, dict_array]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + fmt = fill(UInt8('i'), Arrow._CDATA_MAX_FORMAT_BYTES + 1) + f.schema[] = Arrow.ArrowSchema( + Cstring(pointer(fmt)), + f.schema[].name, + f.schema[].metadata, + f.schema[].flags, + f.schema[].n_children, + f.schema[].children, + f.schema[].dictionary, + f.schema[].release, + f.schema[].private_data, + ) + push!(f.roots, fmt) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + for fmt in ("tsq:", "d:x,2", "d:10,2,64", "w:0", "+w:x") + f = _cdata_fixture(fmt, 0, Ptr{Cvoid}[C_NULL, C_NULL]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + end + + f = _cdata_fixture( + "w:$(Arrow._CDATA_MAX_FIXED_SIZE + 1)", + 0, + Ptr{Cvoid}[C_NULL, C_NULL], + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + child = _primitive_fixture("i", Int32[]) + for _ = 1:Arrow._CDATA_MAX_DEPTH + child = _cdata_fixture("+s", 0, Ptr{Cvoid}[C_NULL]; children=[child]) + end + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(child), _array_ptr(child)) + + child = _primitive_fixture("i", Int32[]) + for _ = 1:Arrow._CDATA_MAX_DEPTH + offsets = Int32[0] + child = _cdata_fixture( + "+l", + 0, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets))]; + children=[child], + ) + push!(child.roots, offsets) + end + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(child), _array_ptr(child)) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 315d1b60..c591dcf7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -39,6 +39,7 @@ include(joinpath(@__DIR__, "testtables.jl")) include(joinpath(@__DIR__, "testappend.jl")) include(joinpath(@__DIR__, "integrationtest.jl")) include(joinpath(@__DIR__, "dates.jl")) +include(joinpath(@__DIR__, "cdata.jl")) struct CustomStruct x::Int @@ -50,6 +51,25 @@ struct CustomStruct2{sym} x::Int end +struct OffsetPool{T} <: AbstractVector{T} + data::Vector{T} + first::Int +end +Base.size(x::OffsetPool) = size(x.data) +Base.axes(x::OffsetPool) = (x.first:(x.first + length(x.data) - 1),) +Base.IndexStyle(::Type{<:OffsetPool}) = IndexLinear() +Base.getindex(x::OffsetPool, i::Int) = x.data[i - x.first + 1] + +struct OffsetRefVector{T,P,R} <: AbstractVector{T} + pool::P + refs::R +end +Base.size(x::OffsetRefVector) = size(x.refs) +Base.IndexStyle(::Type{<:OffsetRefVector}) = IndexLinear() +Base.getindex(x::OffsetRefVector, i::Int) = x.pool[x.refs[i]] +DataAPI.refpool(x::OffsetRefVector) = x.pool +DataAPI.refarray(x::OffsetRefVector) = x.refs + @testset ExtendedTestSet "Arrow" begin @testset "table roundtrips" begin for case in testtables @@ -353,6 +373,16 @@ end @test isa(first(av.indices), Signed) @test length(av) == 3 @test eltype(av) == String + + pool = OffsetPool(["a", "b"], 1000) + x = OffsetRefVector{String,typeof(pool),Vector{Int}}(pool, [1000, 1001, 1000]) + av = Arrow.toarrowvector(x) + @test av.indices == Int8[0, 1, 0] + @test collect(av) == ["a", "b", "a"] + @test Arrow.Table(Arrow.tobuffer((a=x,))).a == ["a", "b", "a"] + + bad = OffsetRefVector{String,typeof(pool),Vector{Int}}(pool, [1000, 1002]) + @test_throws ArgumentError Arrow.toarrowvector(bad) end @testset "# 120" begin