diff --git a/src/Arrow.jl b/src/Arrow.jl index 6f3ccdf8..0aed012f 100644 --- a/src/Arrow.jl +++ b/src/Arrow.jl @@ -26,11 +26,11 @@ 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 and export It currently doesn't include support for: * Tensors or sparse tensors * Flight RPC - * C data interface 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 +76,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") @@ -138,6 +139,7 @@ function __init__() resize!(empty!(ZSTD_COMPRESSOR), nt) resize!(empty!(LZ4_FRAME_DECOMPRESSOR), nt) resize!(empty!(ZSTD_DECOMPRESSOR), nt) + _init_c_data_export_callbacks!() return end 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..7190e511 --- /dev/null +++ b/src/cdata.jl @@ -0,0 +1,2061 @@ +# 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. + +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 + +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 _SUPPORTED_CDATA_FLAGS = + ARROW_FLAG_DICTIONARY_ORDERED | ARROW_FLAG_NULLABLE | ARROW_FLAG_MAP_KEYS_SORTED + +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::Base.RefValue{ArrowSchema} + array::Base.RefValue{ArrowArray} + released::Bool + lock::ReentrantLock +end + +function CDataOwner(schema_ptr::Ptr{ArrowSchema}, array_ptr::Ptr{ArrowArray}) + # Per the Arrow C Data Interface spec, move the base structures into Julia + # owned storage and mark the sources released without calling their release + # callbacks, like arrow-rs `from_raw` and nanoarrow `ArrowArrayMove`. + owner = CDataOwner( + Ref(unsafe_load(schema_ptr)), + Ref(unsafe_load(array_ptr)), + false, + ReentrantLock(), + ) + _clear_schema_release!(schema_ptr) + _clear_array_release!(array_ptr) + 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(meta, key::AbstractString, style::Bool) + val = meta[key] + return style ? (val, :default) : val +end + +function _dataapi_metadata(meta, key::AbstractString, default, style::Bool) + if meta !== nothing && haskey(meta, key) + val = meta[key] + return style ? (val, :default) : val + end + return style ? (default, :default) : default +end + +DataAPI.metadata(t::CDataTable, key::AbstractString; style::Bool=false) = + _dataapi_metadata(getmetadata(t), key, style) +DataAPI.metadata(t::CDataTable, key::AbstractString, default; style::Bool=false) = + _dataapi_metadata(getmetadata(t), key, default, style) + +function DataAPI.metadatakeys(t::CDataTable) + meta = getmetadata(t) + meta === nothing && return () + return keys(meta) +end + +DataAPI.colmetadata(t::CDataTable, col, key::AbstractString; style::Bool=false) = + _dataapi_metadata(getmetadata(t[col]), key, style) +DataAPI.colmetadata(t::CDataTable, col, key::AbstractString, default; style::Bool=false) = + _dataapi_metadata(getmetadata(t[col]), key, default, style) + +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 _clear_schema_release!(ptr::Ptr{ArrowSchema}) + ptr == C_NULL && return + schema = unsafe_load(ptr) + schema.release == C_NULL && return + unsafe_store!( + ptr, + ArrowSchema( + schema.format, + schema.name, + schema.metadata, + schema.flags, + schema.n_children, + schema.children, + schema.dictionary, + C_NULL, + schema.private_data, + ), + ) + return +end + +function _clear_array_release!(ptr::Ptr{ArrowArray}) + ptr == C_NULL && return + array = unsafe_load(ptr) + array.release == C_NULL && return + unsafe_store!( + ptr, + 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 + +function release_c_data(owner::CDataOwner) + array_release = Ptr{Cvoid}(C_NULL) + schema_release = Ptr{Cvoid}(C_NULL) + lock(owner.lock) + try + owner.released && return + owner.released = true + array_release = owner.array[].release + schema_release = owner.schema[].release + finally + unlock(owner.lock) + end + # Per the Arrow C Data Interface spec, consumers release only the base structures. + if array_release != C_NULL + ccall(array_release, Cvoid, (Ptr{ArrowArray},), owner.array) + end + if schema_release != C_NULL + ccall(schema_release, Cvoid, (Ptr{ArrowSchema},), owner.schema) + 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.n_buffers does not match the format")) + 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_flags(schema::ArrowSchema, format::CDataFormat) + unknown = schema.flags & ~_SUPPORTED_CDATA_FLAGS + unknown == 0 || throw(ArgumentError("ArrowSchema.flags contains unknown bits")) + if (schema.flags & ARROW_FLAG_DICTIONARY_ORDERED) != 0 && schema.dictionary == C_NULL + throw(ArgumentError("dictionary ordered flag requires a dictionary schema")) + end + if (schema.flags & ARROW_FLAG_MAP_KEYS_SORTED) != 0 + throw(ArgumentError("map keys sorted flag requires a map schema")) + end + return +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_add(offset, len, "ArrowArray offset plus length") + _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 <= len) + 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", + ), + ) + _validate_flags(schema, 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 validity_bytes > 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 _struct_name(node::CDataNode, i::Int) + return Symbol(node.name === nothing ? "f$(i)" : node.name) +end + +function _struct_columns(node::CDataNode, owner::CDataOwner, convert::Bool) + columns = AbstractVector[] + names = Symbol[] + for (i, child_node) in enumerate(node.children) + child = _import_node(child_node, owner, convert) + push!(columns, _slice_for_table(child, node.offset, node.len)) + push!(names, _struct_name(child_node, i)) + end + return Tuple(names), Tuple(columns) +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) + names_tuple, data = _struct_columns(node, owner, convert) + names = collect(names_tuple) + columns = AbstractVector[_mask_for_struct(col, parent_validity) for col in data] + types = Type[eltype(col) for col in columns] + 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. + +The importer moves the base `ArrowSchema` and `ArrowArray` structures into +Julia owned storage and marks the passed structures released +(`release = C_NULL`) without calling their release callbacks, following the +C Data Interface move semantics. Callers may free or reuse the passed +structures as soon as `from_c_data` returns; the moved copies are released +through the producer callbacks by `release_c_data` or finalization. + +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, +) + schema_ptr == C_NULL && throw(ArgumentError("ArrowSchema pointer is NULL")) + array_ptr == C_NULL && throw(ArgumentError("ArrowArray pointer is NULL")) + owner = CDataOwner(schema_ptr, array_ptr) + try + node = GC.@preserve owner _validate_node( + Base.unsafe_convert(Ptr{ArrowSchema}, owner.schema), + Base.unsafe_convert(Ptr{ArrowArray}, owner.array); + 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...) + +mutable struct CDataExportSchemaOwner + refs::Vector{Ref{ArrowSchema}} + roots::Vector{Any} +end + +mutable struct CDataExportArrayOwner + refs::Vector{Ref{ArrowArray}} + roots::Vector{Any} +end + +const _CDATA_EXPORT_LOCK = ReentrantLock() +const _CDATA_EXPORT_NEXT_TOKEN = Ref{UInt}(0) +const _CDATA_EXPORT_SCHEMA_OWNERS = Dict{UInt,CDataExportSchemaOwner}() +const _CDATA_EXPORT_ARRAY_OWNERS = Dict{UInt,CDataExportArrayOwner}() + +function _next_c_data_export_token() + lock(_CDATA_EXPORT_LOCK) + try + token = _CDATA_EXPORT_NEXT_TOKEN[] + UInt(1) + token == 0 && throw(ArgumentError("Arrow C Data export token overflowed")) + _CDATA_EXPORT_NEXT_TOKEN[] = token + return token + finally + unlock(_CDATA_EXPORT_LOCK) + end +end + +function _release_exported_schema(ptr::Ptr{ArrowSchema}) + ptr == C_NULL && return + schema = unsafe_load(ptr) + schema.release == C_NULL && return + for i = 1:schema.n_children + child = unsafe_load(schema.children, i) + if child != C_NULL + child_schema = unsafe_load(child) + if child_schema.release != C_NULL + ccall(child_schema.release, Cvoid, (Ptr{ArrowSchema},), child) + end + end + end + if schema.dictionary != C_NULL + dictionary = unsafe_load(schema.dictionary) + if dictionary.release != C_NULL + ccall(dictionary.release, Cvoid, (Ptr{ArrowSchema},), schema.dictionary) + end + end + token = UInt(schema.private_data) + owner = lock(_CDATA_EXPORT_LOCK) do + pop!(_CDATA_EXPORT_SCHEMA_OWNERS, token, nothing) + end + if owner === nothing + _clear_schema_release!(ptr) + return + end + _clear_schema_release!(ptr) + for ref in owner.refs + _clear_schema_release!(Base.unsafe_convert(Ptr{ArrowSchema}, ref)) + end + empty!(owner.roots) + empty!(owner.refs) + return +end + +function _release_exported_array(ptr::Ptr{ArrowArray}) + ptr == C_NULL && return + array = unsafe_load(ptr) + array.release == C_NULL && return + for i = 1:array.n_children + child = unsafe_load(array.children, i) + if child != C_NULL + child_array = unsafe_load(child) + if child_array.release != C_NULL + ccall(child_array.release, Cvoid, (Ptr{ArrowArray},), child) + end + end + end + if array.dictionary != C_NULL + dictionary = unsafe_load(array.dictionary) + if dictionary.release != C_NULL + ccall(dictionary.release, Cvoid, (Ptr{ArrowArray},), array.dictionary) + end + end + token = UInt(array.private_data) + owner = lock(_CDATA_EXPORT_LOCK) do + pop!(_CDATA_EXPORT_ARRAY_OWNERS, token, nothing) + end + if owner === nothing + _clear_array_release!(ptr) + return + end + _clear_array_release!(ptr) + for ref in owner.refs + _clear_array_release!(Base.unsafe_convert(Ptr{ArrowArray}, ref)) + end + empty!(owner.roots) + empty!(owner.refs) + return +end + +global _CDATA_EXPORT_SCHEMA_RELEASE::Ptr{Cvoid} = C_NULL +global _CDATA_EXPORT_ARRAY_RELEASE::Ptr{Cvoid} = C_NULL + +function _init_c_data_export_callbacks!() + global _CDATA_EXPORT_SCHEMA_RELEASE = + @cfunction(_release_exported_schema, Cvoid, (Ptr{ArrowSchema},)) + global _CDATA_EXPORT_ARRAY_RELEASE = + @cfunction(_release_exported_array, Cvoid, (Ptr{ArrowArray},)) + return +end + +function _checked_int64(x::Integer, name) + x < 0 && throw(ArgumentError("$name must be nonnegative")) + x > typemax(Int64) && throw(ArgumentError("$name exceeds the Int64 range")) + return Int64(x) +end + +function _primitive_c_data_format(::Type{T}) where {T} + T === Missing && return "n" + T === Bool && return "b" + T === Int8 && return "c" + T === UInt8 && return "C" + T === Int16 && return "s" + T === UInt16 && return "S" + T === Int32 && return "i" + T === UInt32 && return "I" + T === Int64 && return "l" + T === UInt64 && return "L" + T === Float16 && return "e" + T === Float32 && return "f" + T === Float64 && return "g" + T === Date{Meta.DateUnit.DAY,Int32} && return "tdD" + T === Date{Meta.DateUnit.MILLISECOND,Int64} && return "tdm" + T === Time{Meta.TimeUnit.SECOND,Int32} && return "tts" + T === Time{Meta.TimeUnit.MILLISECOND,Int32} && return "ttm" + T === Time{Meta.TimeUnit.MICROSECOND,Int64} && return "ttu" + T === Time{Meta.TimeUnit.NANOSECOND,Int64} && return "ttn" + T === Duration{Meta.TimeUnit.SECOND} && return "tDs" + T === Duration{Meta.TimeUnit.MILLISECOND} && return "tDm" + T === Duration{Meta.TimeUnit.MICROSECOND} && return "tDu" + T === Duration{Meta.TimeUnit.NANOSECOND} && return "tDn" + T === Interval{Meta.IntervalUnit.YEAR_MONTH,Int32} && return "tiM" + T === Interval{Meta.IntervalUnit.DAY_TIME,Int64} && return "tiD" + if T <: Timestamp + U = T.parameters[1] + TZ = T.parameters[2] + unit = + U === Meta.TimeUnit.SECOND ? "s" : + U === Meta.TimeUnit.MILLISECOND ? "m" : + U === Meta.TimeUnit.MICROSECOND ? "u" : + U === Meta.TimeUnit.NANOSECOND ? "n" : + throw(ArgumentError("unsupported Arrow timestamp unit for C Data export")) + tz = TZ === nothing ? "" : String(TZ) + return "ts$(unit):$(tz)" + elseif T <: Decimal + P = T.parameters[1] + S = T.parameters[2] + I = T.parameters[3] + I === Int128 && return "d:$(P),$(S),128" + I === Int256 && return "d:$(P),$(S),256" + end + throw(ArgumentError("unsupported Arrow C Data export type: $T")) +end + +_c_data_format(::NullVector) = "n" +_c_data_format(::BoolVector) = "b" +_c_data_format(v::Primitive) = _primitive_c_data_format(Base.nonmissingtype(eltype(v))) + +function _c_data_format(v::List{T,Int32}) where {T} + S = Base.nonmissingtype(T) + liststringtype(v) && return S <: AbstractString ? "u" : "z" + return "+l" +end + +function _c_data_format(v::List{T,Int64}) where {T} + S = Base.nonmissingtype(T) + liststringtype(v) && return S <: AbstractString ? "U" : "Z" + return "+L" +end + +function _fixed_size_width(v::FixedSizeList{T}) where {T} + S = Base.nonmissingtype(T) + K = ArrowTypes.ArrowKind(ArrowTypes.ArrowType(S)) + return ArrowTypes.getsize(K) +end + +_is_fixed_size_binary(v::FixedSizeList) = eltype(v.data) === UInt8 + +function _c_data_format(v::FixedSizeList) + n = _fixed_size_width(v) + return _is_fixed_size_binary(v) ? "w:$(n)" : "+w:$(n)" +end + +_c_data_format(::Struct) = "+s" + +function _c_data_format(v::ArrowVector) + throw(ArgumentError("unsupported Arrow C Data export array type: $(typeof(v))")) +end + +_c_data_children(::ArrowVector) = () +_c_data_children(v::List) = liststringtype(v) ? () : (v.data,) +_c_data_children(v::FixedSizeList) = _is_fixed_size_binary(v) ? () : (v.data,) +_c_data_children(v::Struct) = v.data + +_c_data_child_name(::ArrowVector, i::Integer) = "" +function _c_data_child_name(v::Struct, i::Integer) + names = fieldnames(Base.nonmissingtype(eltype(v))) + return i <= length(names) ? String(names[i]) : "f$(i)" +end + +function _assert_c_data_export_supported(v::ArrowVector) + if v isa Compressed + throw( + ArgumentError( + "compressed Arrow arrays cannot be exported through the C Data Interface", + ), + ) + elseif v isa DictEncoded + throw(ArgumentError("dictionary encoded Arrow C Data export is not supported")) + elseif v isa Map + throw(ArgumentError("map Arrow C Data export is not supported")) + elseif v isa Union{DenseUnion,SparseUnion} + throw(ArgumentError("union Arrow C Data export is not supported")) + end + _c_data_format(v) + for child in _c_data_children(v) + _assert_c_data_export_supported(child) + end + return +end + +function _metadata_to_c_data(meta) + meta === nothing && return UInt8[] + isempty(meta) && return UInt8[] + length(meta) > _CDATA_MAX_METADATA_PAIRS && + throw(ArgumentError("Arrow C Data metadata has too many pairs")) + total = 4 + io = IOBuffer() + Base.write(io, Int32(length(meta))) + for (k, v) in meta + key = codeunits(String(k)) + val = codeunits(String(v)) + length(key) > _CDATA_MAX_METADATA_FIELD_BYTES && + throw(ArgumentError("Arrow C Data metadata key is too large")) + total = _checked_add(total, 4, "metadata byte count") + total = _checked_add(total, length(key), "metadata byte count") + total > _CDATA_MAX_METADATA_BYTES && + throw(ArgumentError("Arrow C Data metadata byte count exceeds the limit")) + length(val) > _CDATA_MAX_METADATA_FIELD_BYTES && + throw(ArgumentError("Arrow C Data metadata value is too large")) + total = _checked_add(total, 4, "metadata byte count") + total = _checked_add(total, length(val), "metadata byte count") + total > _CDATA_MAX_METADATA_BYTES && + throw(ArgumentError("Arrow C Data metadata byte count exceeds the limit")) + Base.write(io, Int32(length(key))) + Base.write(io, key) + Base.write(io, Int32(length(val))) + Base.write(io, val) + end + return take!(io) +end + +function _export_cstring(s::AbstractString, roots::Vector{Any}) + bytes = Vector{UInt8}(s * "\0") + push!(roots, bytes) + return GC.@preserve bytes begin + Cstring(pointer(bytes)) + end +end + +function _export_metadata(meta, roots::Vector{Any}) + bytes = _metadata_to_c_data(meta) + isempty(bytes) && return Cstring(C_NULL) + push!(roots, bytes) + return GC.@preserve bytes begin + Cstring(pointer(bytes)) + end +end + +function _schema_flags(v::ArrowVector) + flags = Int64(0) + if eltype(v) >: Missing || v isa NullVector + flags |= ARROW_FLAG_NULLABLE + end + return flags +end + +function _make_c_data_child_schemas!(v::ArrowVector, owner::CDataExportSchemaOwner) + children = _c_data_children(v) + isempty(children) && return Int64(0), Ptr{Ptr{ArrowSchema}}(C_NULL) + ptrs = Ptr{ArrowSchema}[] + try + for (i, child) in enumerate(children) + ref = _make_c_data_schema(child, _c_data_child_name(v, i)) + push!(ptrs, Base.unsafe_convert(Ptr{ArrowSchema}, ref)) + end + catch + # Release children built before the failure so their owners do not leak. + foreach(_release_exported_schema, ptrs) + rethrow() + end + push!(owner.roots, ptrs) + return GC.@preserve ptrs begin + Int64(length(ptrs)), Ptr{Ptr{ArrowSchema}}(pointer(ptrs)) + end +end + +function _fill_c_data_schema!( + ref::Ref{ArrowSchema}, + v::ArrowVector, + name::AbstractString, + token::UInt, + owner::CDataExportSchemaOwner, +) + push!(owner.refs, ref) + format = _export_cstring(_c_data_format(v), owner.roots) + field_name = _export_cstring(String(name), owner.roots) + metadata = _export_metadata(getmetadata(v), owner.roots) + n_children, children = _make_c_data_child_schemas!(v, owner) + ref[] = ArrowSchema( + format, + field_name, + metadata, + _schema_flags(v), + n_children, + children, + Ptr{ArrowSchema}(C_NULL), + _CDATA_EXPORT_SCHEMA_RELEASE, + Ptr{Cvoid}(token), + ) + return ref +end + +function _make_c_data_schema(v::ArrowVector, name::AbstractString) + token = _next_c_data_export_token() + owner = CDataExportSchemaOwner(Ref{ArrowSchema}[], Any[]) + ref = Ref{ArrowSchema}() + _fill_c_data_schema!(ref, v, name, token, owner) + _store_c_data_schema_owner!(token, owner) + return ref +end + +function _store_c_data_schema_owner!(token::UInt, owner::CDataExportSchemaOwner) + lock(_CDATA_EXPORT_LOCK) + try + _CDATA_EXPORT_SCHEMA_OWNERS[token] = owner + return + finally + unlock(_CDATA_EXPORT_LOCK) + end +end + +function _validity_ptr(v::ArrowVector, roots::Vector{Any}) + validity = validitybitmap(v) + validity.nc == 0 && return Ptr{Cvoid}(C_NULL) + validity.nc > 0 || throw(ArgumentError("validity null count is invalid")) + len = length(v) + validity.nc <= len || throw(ArgumentError("validity null count exceeds array length")) + validity.โ„“ >= len || throw(ArgumentError("validity bitmap length is too short")) + isempty(validity.bytes) && throw(ArgumentError("validity bitmap is empty")) + nbytes = cld(len, 8) + nbytes > 0 || throw(ArgumentError("validity bitmap is empty")) + bytes = validity.bytes + _check_export_range(validity.pos, nbytes, length(bytes), "validity bitmap") + push!(roots, bytes) + return GC.@preserve bytes begin + Ptr{Cvoid}(pointer(bytes, validity.pos)) + end +end + +_validity_ptr(::NullVector, roots::Vector{Any}) = Ptr{Cvoid}(C_NULL) + +function _materialized_vector(x, ::Type{T}) where {T} + x isa Vector{T} && return x + return collect(T, x) +end + +function _optional_data_ptr(x::AbstractVector, roots::Vector{Any}) + isempty(x) && return Ptr{Cvoid}(C_NULL) + push!(roots, x) + return GC.@preserve x begin + Ptr{Cvoid}(pointer(x)) + end +end + +function _required_data_ptr(x::AbstractVector, roots::Vector{Any}, name) + isempty(x) && throw(ArgumentError("$name is empty")) + push!(roots, x) + return GC.@preserve x begin + Ptr{Cvoid}(pointer(x)) + end +end + +function _buffer_array_ptr(buffers::Vector{Ptr{Cvoid}}, roots::Vector{Any}) + isempty(buffers) && return Int64(0), Ptr{Ptr{Cvoid}}(C_NULL) + push!(roots, buffers) + return GC.@preserve buffers begin + Int64(length(buffers)), Ptr{Ptr{Cvoid}}(pointer(buffers)) + end +end + +function _check_export_range(pos::Int, n::Int, len::Int, name) + pos > 0 || throw(ArgumentError("$name position is invalid")) + n >= 0 || throw(ArgumentError("$name length is invalid")) + pos <= len || throw(ArgumentError("$name is too short")) + n <= len - pos + 1 || throw(ArgumentError("$name is too short")) + return +end + +function _validate_offsets_for_export(offsets::AbstractVector, len::Int, name) + len >= 0 || throw(ArgumentError("$name length is invalid")) + count = _checked_add(len, 1, "$name offset count") + length(offsets) >= count || + throw(ArgumentError("$name must contain length + 1 offsets")) + prev = offsets[1] + prev < 0 && throw(ArgumentError("$name contains a negative offset")) + prev > typemax(Int) && throw(ArgumentError("$name offset exceeds the Julia Int range")) + for i = 2:count + cur = offsets[i] + cur < 0 && throw(ArgumentError("$name contains a negative offset")) + cur > typemax(Int) && + throw(ArgumentError("$name offset exceeds the Julia Int range")) + cur < prev && throw(ArgumentError("$name is not monotonic")) + prev = cur + end + return Int(prev) +end + +function _primitive_buffers(v::Primitive, roots::Vector{Any}) + T = Base.nonmissingtype(eltype(v)) + data = _materialized_vector(v.data, T) + length(data) >= length(v) || throw(ArgumentError("primitive data buffer is too short")) + buffers = Ptr{Cvoid}[_validity_ptr(v, roots), _optional_data_ptr(data, roots)] + return _buffer_array_ptr(buffers, roots) +end + +function _bool_buffers(v::BoolVector, roots::Vector{Any}) + data_bytes = cld(length(v), 8) + if data_bytes > 0 + _check_export_range(v.pos, data_bytes, length(v.arrow), "boolean data buffer") + end + data_ptr = if data_bytes == 0 + Ptr{Cvoid}(C_NULL) + else + bytes = v.arrow + push!(roots, bytes) + GC.@preserve bytes begin + Ptr{Cvoid}(pointer(bytes, v.pos)) + end + end + buffers = Ptr{Cvoid}[_validity_ptr(v, roots), data_ptr] + return _buffer_array_ptr(buffers, roots) +end + +function _list_buffers(v::List{T,O}, roots::Vector{Any}) where {T,O} + offsets = _materialized_vector(v.offsets.offsets, O) + last = _validate_offsets_for_export(offsets, length(v), "list offset buffer") + offset_ptr = _required_data_ptr(offsets, roots, "list offset buffer") + if liststringtype(v) + data = _materialized_vector(v.data, UInt8) + last <= length(data) || + throw(ArgumentError("list offset exceeds data buffer length")) + buffers = + Ptr{Cvoid}[_validity_ptr(v, roots), offset_ptr, _optional_data_ptr(data, roots)] + else + child = v.data + last <= length(child) || + throw(ArgumentError("list offset exceeds child array length")) + buffers = Ptr{Cvoid}[_validity_ptr(v, roots), offset_ptr] + end + return _buffer_array_ptr(buffers, roots) +end + +function _fixed_size_list_buffers(v::FixedSizeList, roots::Vector{Any}) + if _is_fixed_size_binary(v) + n = _fixed_size_width(v) + nbytes = _checked_mul(length(v), n, "fixed size binary byte count") + data = _materialized_vector(v.data, UInt8) + length(data) >= nbytes || + throw(ArgumentError("fixed size binary data buffer is too short")) + buffers = Ptr{Cvoid}[_validity_ptr(v, roots), _optional_data_ptr(data, roots)] + else + child_len = + _checked_mul(length(v), _fixed_size_width(v), "fixed size list child length") + length(v.data) >= child_len || + throw(ArgumentError("fixed size list child array is too short")) + buffers = Ptr{Cvoid}[_validity_ptr(v, roots)] + end + return _buffer_array_ptr(buffers, roots) +end + +function _struct_buffers(v::Struct, roots::Vector{Any}) + for child in v.data + length(child) >= length(v) || + throw(ArgumentError("struct child array is too short")) + end + return _buffer_array_ptr(Ptr{Cvoid}[_validity_ptr(v, roots)], roots) +end + +_c_data_buffers(v::NullVector, roots::Vector{Any}) = Int64(0), Ptr{Ptr{Cvoid}}(C_NULL) +_c_data_buffers(v::Primitive, roots::Vector{Any}) = _primitive_buffers(v, roots) +_c_data_buffers(v::BoolVector, roots::Vector{Any}) = _bool_buffers(v, roots) +_c_data_buffers(v::List, roots::Vector{Any}) = _list_buffers(v, roots) +_c_data_buffers(v::FixedSizeList, roots::Vector{Any}) = _fixed_size_list_buffers(v, roots) +_c_data_buffers(v::Struct, roots::Vector{Any}) = _struct_buffers(v, roots) + +function _make_c_data_child_arrays!(v::ArrowVector, owner::CDataExportArrayOwner) + children = _c_data_children(v) + isempty(children) && return Int64(0), Ptr{Ptr{ArrowArray}}(C_NULL) + ptrs = Ptr{ArrowArray}[] + try + for child in children + ref = _make_c_data_array(child) + push!(ptrs, Base.unsafe_convert(Ptr{ArrowArray}, ref)) + end + catch + # Release children built before the failure so their owners do not leak. + foreach(_release_exported_array, ptrs) + rethrow() + end + push!(owner.roots, ptrs) + return GC.@preserve ptrs begin + Int64(length(ptrs)), Ptr{Ptr{ArrowArray}}(pointer(ptrs)) + end +end + +function _c_data_null_count(v::ArrowVector) + nc = nullcount(v) + nc <= length(v) || throw(ArgumentError("ArrowArray.null_count exceeds length")) + return _checked_int64(nc, "ArrowArray.null_count") +end +_c_data_null_count(v::NullVector) = _checked_int64(length(v), "ArrowArray.null_count") + +function _fill_c_data_array!( + ref::Ref{ArrowArray}, + v::ArrowVector, + token::UInt, + owner::CDataExportArrayOwner, +) + push!(owner.refs, ref) + n_buffers, buffers = _c_data_buffers(v, owner.roots) + n_children, children = _make_c_data_child_arrays!(v, owner) + ref[] = ArrowArray( + _checked_int64(length(v), "ArrowArray.length"), + _c_data_null_count(v), + Int64(0), + n_buffers, + n_children, + buffers, + children, + Ptr{ArrowArray}(C_NULL), + _CDATA_EXPORT_ARRAY_RELEASE, + Ptr{Cvoid}(token), + ) + return ref +end + +function _make_c_data_array(v::ArrowVector) + token = _next_c_data_export_token() + owner = CDataExportArrayOwner(Ref{ArrowArray}[], Any[v]) + ref = Ref{ArrowArray}() + _fill_c_data_array!(ref, v, token, owner) + _store_c_data_array_owner!(token, owner) + return ref +end + +function _store_c_data_array_owner!(token::UInt, owner::CDataExportArrayOwner) + lock(_CDATA_EXPORT_LOCK) + try + _CDATA_EXPORT_ARRAY_OWNERS[token] = owner + return + finally + unlock(_CDATA_EXPORT_LOCK) + end +end + +function _to_c_data_refs(col::ArrowVector, name::AbstractString) + _assert_c_data_export_supported(col) + schema_ref = _make_c_data_schema(col, name) + try + return schema_ref, _make_c_data_array(col) + catch + _release_exported_schema(Base.unsafe_convert(Ptr{ArrowSchema}, schema_ref)) + rethrow() + end +end + +""" + Arrow.to_c_data(col::ArrowVector; name="") -> (Ref{ArrowSchema}, Ref{ArrowArray}) + +Export an Arrow array through the Arrow C Data Interface. The returned schema +and array have independent release callbacks; releasing the schema does not +release array buffers. +""" +to_c_data(col::ArrowVector; name::AbstractString="") = _to_c_data_refs(col, name) + +function _table_to_c_data_struct(tbl, names) + cols = Tables.columns(tbl) + name_strings = String.(collect(names)) + arrow_tbl = toarrowtable( + cols, + Dict{Int64,Any}(), + false, + nothing, + true, + false, + false, + DEFAULT_MAX_DEPTH, + getmetadata(tbl), + nothing, + ) + length(name_strings) == length(arrow_tbl.cols) || + throw(ArgumentError("names length must match the number of table columns")) + syms = Tuple(Symbol.(name_strings)) + data = Tuple(arrow_tbl.cols) + types = Tuple(eltype(col) for col in data) + T = NamedTuple{syms,Tuple{types...}} + validity = ValidityBitmap(UInt8[], 1, Tables.rowcount(arrow_tbl), 0) + return Struct{T,typeof(data),syms}( + validity, + data, + Tables.rowcount(arrow_tbl), + arrow_tbl.metadata, + ) +end + +""" + Arrow.to_c_data(tbl; names=String.(Tables.columnnames(tbl))) + -> (Ref{ArrowSchema}, Ref{ArrowArray}) + +Export a Tables.jl column table as a root C Data struct array. +""" +function to_c_data(tbl; names=String.(Tables.columnnames(Tables.columns(tbl)))) + root = _table_to_c_data_struct(tbl, names) + return _to_c_data_refs(root, "") +end diff --git a/test/Project.toml b/test/Project.toml index c2e02aa8..fd674490 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -22,6 +22,7 @@ DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" FilePathsBase = "48062228-2e41-5def-b9a4-89aafe57970f" JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" +Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb" Mmap = "a63ad114-7e13-5084-954f-fe012c677804" OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" PooledArrays = "2dfb63ee-cc39-5dd5-95bd-886bf059d720" diff --git a/test/cdata.jl b/test/cdata.jl new file mode 100644 index 00000000..c77ce876 --- /dev/null +++ b/test/cdata.jl @@ -0,0 +1,1990 @@ +# 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 +using Libdl +using Random + +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 _replace_schema!(f::CDataFixture; kwargs...) + return _set_schema!(f; kwargs...) +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 + +function _replace_array!(f::CDataFixture; len=f.array[].length, kwargs...) + return _set_array!(f; length=len, kwargs...) +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) + @test fieldcount(Arrow.ArrowSchema) == 9 + @test fieldcount(Arrow.ArrowArray) == 10 + 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]) + f = _cdata_fixture("n", 3, Ptr{Cvoid}[]; null_count=Int64(3)) + 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 Dict(Arrow.getmetadata(tbl)) == Dict("source" => "cdata") + @test DataAPI.metadata(tbl) == Dict("source" => "cdata") + @test DataAPI.metadata(tbl, "source") == "cdata" + @test DataAPI.metadata(tbl, "source", "fallback"; style=true) == ("cdata", :default) + @test DataAPI.metadata(tbl, "missing", "fallback") == "fallback" + @test Set(DataAPI.metadatakeys(tbl)) == Set(["source"]) + @test DataAPI.colmetadata(tbl, :x, "unit") == "id" + @test DataAPI.colmetadata(tbl, :x, "unit", "fallback"; style=true) == + ("id", :default) + @test DataAPI.colmetadata(tbl, :x, "missing", "fallback") == "fallback" + @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 "nested struct columns" begin + child_validity = UInt8[0b00011101] + xchild = _primitive_fixture( + "i", + Int32[10, 20, 30, 40, 50]; + validity=child_validity, + null_count=Int64(1), + flags=Arrow.ARROW_FLAG_NULLABLE, + name="x", + ) + struct_validity = UInt8[0b00011011] + nested = _cdata_fixture( + "+s", + 5, + Ptr{Cvoid}[Ptr{Cvoid}(pointer(struct_validity))]; + name="point", + children=[xchild], + metadata=Dict("shape" => "point"), + null_count=Int64(1), + flags=Arrow.ARROW_FLAG_NULLABLE, + ) + push!(nested.roots, struct_validity) + T = NamedTuple{(:x,),Tuple{Union{Int32,Missing}}} + + parent_validity = UInt8[0b00001111] + root = _cdata_fixture( + "+s", + 5, + Ptr{Cvoid}[Ptr{Cvoid}(pointer(parent_validity))]; + children=[nested], + 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, :point) + @test eltype(col) == Union{T,Missing} + @test Arrow.nullcount(col) == 2 + @test Dict(Arrow.getmetadata(col)) == Dict("shape" => "point") + @test isequal( + collect(col), + Union{T,Missing}[(x=10,), (x=missing,), missing, (x=40,), missing], + ) + + root = + _cdata_fixture("+s", 3, Ptr{Cvoid}[C_NULL]; offset=Int64(1), children=[nested]) + tbl = Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + col = Tables.getcolumn(tbl, :point) + @test Tables.schema(tbl).types == (Union{T,Missing},) + @test isequal(collect(col), Union{T,Missing}[(x=missing,), missing, (x=40,)]) + + collected = collect(col) + copied = copy(col) + Arrow.release_c_data(col) + @test root.array[].release == C_NULL + @test root.schema[].release == C_NULL + @test nested.array[].release != C_NULL + @test nested.schema[].release != C_NULL + @test_throws ArgumentError col[1] + @test isequal(collected, Union{T,Missing}[(x=missing,), missing, (x=40,)]) + @test isequal(copied, Union{T,Missing}[(x=missing,), missing, (x=40,)]) + 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 + child2 = _primitive_fixture("i", Int32[10, 20, 30]; name="x") + root2 = _cdata_fixture( + "+s", + 2, + Ptr{Cvoid}[C_NULL]; + offset=Int64(1), + children=[child2], + ) + tbl2 = Arrow.from_c_data(_schema_ptr(root2), _array_ptr(root2)) + 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 + f = _primitive_fixture("tdD", 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)) + owner = Arrow._owner(x) + @test x[1] == 1 + Arrow.release_c_data(x) + @test owner.array[].release == C_NULL + @test owner.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 owner.array[].release != C_NULL + unlock(owner_lock) + locked = false + wait(task) + finally + locked && unlock(owner_lock) + end + @test owner.array[].release == C_NULL + @test owner.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)) + owner = Arrow._owner(x) + _CDATA_REENTRANT_OBJECT[] = x + _CDATA_REENTRANT_RELEASES[] = 0 + try + @test_nowarn Arrow.release_c_data(x) + @test _CDATA_REENTRANT_RELEASES[] == 1 + @test owner.array[].release == C_NULL + @test owner.schema[].release == C_NULL + @test_throws ArgumentError x[1] + finally + _CDATA_REENTRANT_OBJECT[] = nothing + end + end + + @testset "import moves the base structures" begin + f = _primitive_fixture("i", Int32[1, 2, 3]) + _set_array!(f; release=_CDATA_RELEASE_ARRAY_REENTRANT) + _CDATA_REENTRANT_OBJECT[] = nothing + _CDATA_REENTRANT_RELEASES[] = 0 + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + # The sources are marked released without calling their callbacks. + @test f.schema[].release == C_NULL + @test f.array[].release == C_NULL + @test _CDATA_REENTRANT_RELEASES[] == 0 + # The caller may reuse the source structures immediately after import. + f.schema[] = Arrow.ArrowSchema( + Cstring(C_NULL), + Cstring(C_NULL), + Cstring(C_NULL), + -1, + -1, + C_NULL, + C_NULL, + C_NULL, + C_NULL, + ) + f.array[] = + Arrow.ArrowArray(-1, -1, -1, -1, -1, C_NULL, C_NULL, C_NULL, C_NULL, C_NULL) + @test x == Int32[1, 2, 3] + Arrow.release_c_data(x) + @test _CDATA_REENTRANT_RELEASES[] == 1 + @test_nowarn Arrow.release_c_data(x) + @test _CDATA_REENTRANT_RELEASES[] == 1 + @test_throws ArgumentError x[1] + + # A moved-from source cannot be imported again. + f = _primitive_fixture("i", Int32[1, 2]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + Arrow.release_c_data(x) + + @test_throws ArgumentError Arrow.from_c_data( + Ptr{Arrow.ArrowSchema}(C_NULL), + Ptr{Arrow.ArrowArray}(C_NULL), + ) + 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 + bad(f) = @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + bad_primitive(mutator) = begin + f = _primitive_fixture("i", Int32[1]) + mutator(f) + bad(f) + end + + 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}[]) + bad(f) + @test f.array[].release == C_NULL + @test f.schema[].release == C_NULL + + for mutator in ( + f -> _set_schema!(f; release=C_NULL), + f -> _set_array!(f; release=C_NULL), + f -> _set_array!(f; length=Int64(-1)), + f -> _set_array!(f; offset=Int64(-1)), + f -> _set_array!(f; length=typemax(Int64), offset=Int64(1)), + f -> _set_array!(f; n_buffers=Int64(-1)), + f -> _set_array!(f; n_children=Int64(-1)), + f -> _set_schema!(f; n_children=Int64(-1)), + f -> _set_array!(f; n_children=Int64(1)), + f -> _set_schema!(f; n_children=Int64(1)), + ) + bad_primitive(mutator) + end + + for null_count in (Int64(2), Int64(-1), Int64(1)) + bad( + _primitive_fixture( + "i", + Int32[1]; + null_count=null_count, + flags=Arrow.ARROW_FLAG_NULLABLE, + ), + ) + end + + for (validity, null_count) in + ((UInt8[0b00000111], Int64(1)), (UInt8[0b00000101], Int64(0))) + bad( + _primitive_fixture( + "i", + Int32[1, 2, 3]; + validity=validity, + null_count=null_count, + flags=Arrow.ARROW_FLAG_NULLABLE, + ), + ) + end + + f = _primitive_fixture("i", Int32[1]) + _set_array!(f; buffers=Ptr{Ptr{Cvoid}}(C_NULL)) + bad(f) + + data = Int32[1] + f = _cdata_fixture("i", 1, Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(data)), C_NULL]) + push!(f.roots, data) + bad(f) + + for f in ( + _cdata_fixture("n", 0, Ptr{Cvoid}[C_NULL]), + _cdata_fixture("i", 1, Ptr{Cvoid}[C_NULL, C_NULL]), + _cdata_fixture("i", 1, Ptr{Cvoid}[C_NULL]), + ) + bad(f) + end + + child = _primitive_fixture("i", Int32[1]) + f = _cdata_fixture("+l", 1, Ptr{Cvoid}[C_NULL, C_NULL]; children=[child]) + bad(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) + bad(f) + + child = _primitive_fixture("i", Int32[1, 2, 3, 4, 5]) + f = _cdata_fixture("+w:3", 2, Ptr{Cvoid}[C_NULL]; children=[child]) + bad(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]) + bad(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]) + bad(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]) + bad(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) + bad(f) + + for flags in ( + Int64(8), + Arrow.ARROW_FLAG_DICTIONARY_ORDERED, + Arrow.ARROW_FLAG_MAP_KEYS_SORTED, + ) + bad_primitive(f -> _set_schema!(f; flags=flags)) + end + + for words in (Int32[1, -1], Int32[1, 0, -1]) + bytes = reinterpret(UInt8, words) + f = _primitive_fixture("i", Int32[1]) + _set_schema!(f; metadata=Cstring(pointer(bytes))) + append!(f.roots, Any[words, bytes]) + bad(f) + end + + function bad_child(schema_ptr, array_ptr, child) + schema_children = Ptr{Arrow.ArrowSchema}[schema_ptr] + array_children = Ptr{Arrow.ArrowArray}[array_ptr] + 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]) + bad(root) + end + + child = _primitive_fixture("i", Int32[1]) + bad_child(Ptr{Arrow.ArrowSchema}(C_NULL), _array_ptr(child), child) + bad_child(_schema_ptr(child), Ptr{Arrow.ArrowArray}(C_NULL), child) + + child = _primitive_fixture("i", Int32[1]) + _set_schema!(child; release=C_NULL) + bad_child(_schema_ptr(child), _array_ptr(child), child) + + child = _primitive_fixture("i", Int32[1]) + _set_array!(child; release=C_NULL) + bad_child(_schema_ptr(child), _array_ptr(child), child) + + 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]) + bad(f) + + f = _primitive_fixture("i", Int32[1]) + dict_schema = Ref(f.schema[]) + _set_schema!(f; dictionary=Base.unsafe_convert(Ptr{Arrow.ArrowSchema}, dict_schema)) + push!(f.roots, dict_schema) + bad(f) + + f = _primitive_fixture("i", Int32[1]) + dict_array = Ref(f.array[]) + _set_array!(f; dictionary=Base.unsafe_convert(Ptr{Arrow.ArrowArray}, dict_array)) + push!(f.roots, dict_array) + bad(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]) + bad(f) + + f = _primitive_fixture("i", Int32[1]) + fmt = fill(UInt8('i'), Arrow._CDATA_MAX_FORMAT_BYTES + 1) + _set_schema!(f; format=Cstring(pointer(fmt))) + push!(f.roots, fmt) + bad(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 + bad(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 + bad(child) + end + + @testset "deterministic malformed import fuzz" begin + iters = something(tryparse(Int, get(ENV, "ARROW_CDATA_FUZZ_ITERS", "64")), 64) + iters = clamp(iters, 0, 10_000) + rng = Random.MersenneTwister(0x0cda7a) + for _ = 1:iters + f = _primitive_fixture("i", Int32[10, 20, 30]) + case = rand(rng, 1:11) + if case == 1 + _replace_array!(f; len=Int64(-1)) + elseif case == 2 + _replace_array!(f; offset=Int64(-1)) + elseif case == 3 + _replace_array!(f; null_count=Int64(4)) + elseif case == 4 + _replace_array!(f; n_buffers=Int64(1)) + elseif case == 5 + _replace_array!(f; buffers=Ptr{Ptr{Cvoid}}(C_NULL)) + elseif case == 6 + _replace_schema!(f; format=Cstring(C_NULL)) + elseif case == 7 + _replace_schema!(f; release=Ptr{Cvoid}(C_NULL)) + elseif case == 8 + _replace_array!(f; release=Ptr{Cvoid}(C_NULL)) + elseif case == 9 + fmt = Vector{UInt8}("+l\0") + push!(f.roots, fmt) + _replace_schema!(f; format=Cstring(pointer(fmt))) + elseif case == 10 + offsets = Int32[0, 2, 1, 3] + bytes = UInt8[0x01, 0x02, 0x03] + buffers = Ptr{Cvoid}[ + C_NULL, + Ptr{Cvoid}(pointer(offsets)), + Ptr{Cvoid}(pointer(bytes)), + ] + push!(f.roots, offsets) + push!(f.roots, bytes) + push!(f.roots, buffers) + fmt = Vector{UInt8}("z\0") + push!(f.roots, fmt) + _replace_schema!(f; format=Cstring(pointer(fmt))) + _replace_array!( + f; + buffers=Ptr{Ptr{Cvoid}}(pointer(buffers)), + n_buffers=Int64(3), + ) + else + fmt = Vector{UInt8}("w:$(Arrow._CDATA_MAX_FIXED_SIZE + rand(rng, 1:8))\0") + push!(f.roots, fmt) + _replace_schema!(f; format=Cstring(pointer(fmt))) + end + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + end + end + + @testset "malformed nested inputs" begin + child = _primitive_fixture("i", Int32[1, 2, 3]) + _replace_schema!(child; release=Ptr{Cvoid}(C_NULL)) + offsets = Int32[0, 1] + 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]) + root = _cdata_fixture("+s", 2, Ptr{Cvoid}[C_NULL]; children=[child]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + + child = _primitive_fixture("i", Int32[]) + f = _cdata_fixture( + "+w:2", + 1, + Ptr{Cvoid}[C_NULL]; + offset=Int64(typemax(Int)), + children=[child], + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + end + + @testset "C producer smoke" begin + cc = Sys.which("cc") + if cc === nothing || Sys.iswindows() + @test_skip "C producer smoke requires a C compiler" + else + c_src = """ + #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; + }; + static const char format[] = "i"; + static const uint8_t validity[] = {0x05}; + static const int32_t data[] = {10, 20, 30}; + static const void* buffers[] = {validity, data}; + static void release_schema(struct ArrowSchema* schema) { + schema->release = 0; + } + static void release_array(struct ArrowArray* array) { + array->release = 0; + } + int make_nullable_int32(struct ArrowSchema* schema, struct ArrowArray* array) { + if (!schema || !array) return -1; + schema->format = format; + schema->name = ""; + schema->metadata = 0; + schema->flags = 2; + schema->n_children = 0; + schema->children = 0; + schema->dictionary = 0; + schema->release = release_schema; + schema->private_data = 0; + array->length = 3; + array->null_count = 1; + array->offset = 0; + array->n_buffers = 2; + array->n_children = 0; + array->buffers = buffers; + array->children = 0; + array->dictionary = 0; + array->release = release_array; + array->private_data = 0; + return 0; + } + """ + mktempdir() do dir + src = joinpath(dir, "producer.c") + lib = joinpath(dir, "producer.$(Libdl.dlext)") + write(src, c_src) + if Sys.isapple() + run(`$cc -dynamiclib -o $lib $src`) + else + run(`$cc -shared -fPIC -o $lib $src`) + end + handle = Libdl.dlopen(lib) + try + make = Libdl.dlsym(handle, :make_nullable_int32) + schema = Ref( + Arrow.ArrowSchema( + Cstring(C_NULL), + Cstring(C_NULL), + Cstring(C_NULL), + 0, + 0, + Ptr{Ptr{Arrow.ArrowSchema}}(C_NULL), + Ptr{Arrow.ArrowSchema}(C_NULL), + Ptr{Cvoid}(C_NULL), + Ptr{Cvoid}(C_NULL), + ), + ) + array = Ref( + Arrow.ArrowArray( + 0, + 0, + 0, + 0, + 0, + Ptr{Ptr{Cvoid}}(C_NULL), + Ptr{Ptr{Arrow.ArrowArray}}(C_NULL), + Ptr{Arrow.ArrowArray}(C_NULL), + Ptr{Cvoid}(C_NULL), + Ptr{Cvoid}(C_NULL), + ), + ) + @test ccall( + make, + Cint, + (Ptr{Arrow.ArrowSchema}, Ptr{Arrow.ArrowArray}), + Base.unsafe_convert(Ptr{Arrow.ArrowSchema}, schema), + Base.unsafe_convert(Ptr{Arrow.ArrowArray}, array), + ) == 0 + x = Arrow.from_c_data( + Base.unsafe_convert(Ptr{Arrow.ArrowSchema}, schema), + Base.unsafe_convert(Ptr{Arrow.ArrowArray}, array), + ) + @test isequal(collect(x), Union{Int32,Missing}[10, missing, 30]) + Arrow.release_c_data(x) + @test schema[].release == C_NULL + @test array[].release == C_NULL + finally + Libdl.dlclose(handle) + end + end + end + end +end + +_schema_ref_ptr(ref::Ref{Arrow.ArrowSchema}) = + Base.unsafe_convert(Ptr{Arrow.ArrowSchema}, ref) +_array_ref_ptr(ref::Ref{Arrow.ArrowArray}) = Base.unsafe_convert(Ptr{Arrow.ArrowArray}, ref) + +function _release_exported_schema!(ref::Ref{Arrow.ArrowSchema}) + ptr = _schema_ref_ptr(ref) + release = unsafe_load(ptr).release + release == C_NULL && return + ccall(release, Cvoid, (Ptr{Arrow.ArrowSchema},), ptr) + return +end + +function _release_exported_array!(ref::Ref{Arrow.ArrowArray}) + ptr = _array_ref_ptr(ref) + release = unsafe_load(ptr).release + release == C_NULL && return + ccall(release, Cvoid, (Ptr{Arrow.ArrowArray},), ptr) + return +end + +function _mark_exported_schema_released!(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 _mark_exported_array_released!(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 + +function _import_exported(schema_ref, array_ref; convert=true) + return Arrow.from_c_data( + _schema_ref_ptr(schema_ref), + _array_ref_ptr(array_ref); + convert=convert, + ) +end + +function _exported_buffers(array_ref) + array = array_ref[] + array.n_buffers == 0 && return Ptr{Cvoid}[] + return collect(unsafe_wrap(Array, array.buffers, Int(array.n_buffers); own=false)) +end + +function _exported_children(array_ref) + array = array_ref[] + array.n_children == 0 && return Ptr{Arrow.ArrowArray}[] + return collect(unsafe_wrap(Array, array.children, Int(array.n_children); own=false)) +end + +function _exported_child_schemas(schema_ref) + schema = schema_ref[] + schema.n_children == 0 && return Ptr{Arrow.ArrowSchema}[] + return collect(unsafe_wrap(Array, schema.children, Int(schema.n_children); own=false)) +end + +function _exported_vector(::Type{T}, ptr::Ptr{Cvoid}, n::Integer) where {T} + n == 0 && return T[] + return unsafe_wrap(Array, Ptr{T}(ptr), Int(n); own=false) +end + +function _exported_bits(ptr::Ptr{Cvoid}, len::Integer) + len == 0 && return Bool[] + bytes = _exported_vector(UInt8, ptr, cld(Int(len), 8)) + return [ + Arrow.getbit(bytes[((i - 1) >>> 3) + 1], ((i - 1) & 0x07) + 1) for i = 1:Int(len) + ] +end + +@testset "Arrow C Data Interface export" begin + @testset "primitive arrays" begin + col = Arrow.toarrowvector(Int32[1, 2, 3]) + schema, array = Arrow.to_c_data(col; name="ints") + @test unsafe_string(schema[].format) == "i" + @test unsafe_string(schema[].name) == "ints" + @test array[].n_buffers == 2 + @test schema[].private_data != C_NULL + @test array[].private_data != C_NULL + imported = _import_exported(schema, array) + @test collect(imported) == Int32[1, 2, 3] + Arrow.release_c_data(imported) + @test schema[].release == C_NULL + @test array[].release == C_NULL + end + + @testset "boolean arrays" begin + col = Arrow.toarrowvector(Union{Bool,Missing}[true, missing, false, true]) + schema, array = Arrow.to_c_data(col) + @test unsafe_string(schema[].format) == "b" + @test array[].n_buffers == 2 + @test array[].null_count == 1 + buffers = _exported_buffers(array) + @test _exported_bits(buffers[1], array[].length) == [true, false, true, true] + values = _exported_bits(buffers[2], array[].length) + @test values[[1, 3, 4]] == [true, false, true] + _release_exported_array!(array) + _release_exported_schema!(schema) + end + + @testset "string and binary arrays" begin + strings = Union{String,Missing}["ab", "", missing, "cd"] + schema, array = Arrow.to_c_data(Arrow.toarrowvector(strings)) + @test unsafe_string(schema[].format) == "u" + @test array[].n_buffers == 3 + @test array[].null_count == 1 + buffers = _exported_buffers(array) + @test _exported_bits(buffers[1], array[].length) == [true, true, false, true] + offsets = _exported_vector(Int32, buffers[2], array[].length + 1) + @test collect(offsets) == Int32[0, 2, 2, 2, 4] + data = _exported_vector(UInt8, buffers[3], offsets[end]) + @test String(copy(data)) == "abcd" + _release_exported_array!(array) + _release_exported_schema!(schema) + + bytes = [b"ab", b"", b"cd"] + schema, array = Arrow.to_c_data(Arrow.toarrowvector(bytes)) + @test unsafe_string(schema[].format) == "z" + @test array[].n_buffers == 3 + buffers = _exported_buffers(array) + offsets = _exported_vector(Int32, buffers[2], array[].length + 1) + @test collect(offsets) == Int32[0, 2, 2, 4] + data = _exported_vector(UInt8, buffers[3], offsets[end]) + @test collect(data) == UInt8[0x61, 0x62, 0x63, 0x64] + _release_exported_array!(array) + _release_exported_schema!(schema) + end + + @testset "nested arrays" begin + lists = [Int32[1, 2], Int32[], Int32[3]] + schema, array = Arrow.to_c_data(Arrow.toarrowvector(lists)) + @test unsafe_string(schema[].format) == "+l" + @test schema[].n_children == 1 + @test array[].n_children == 1 + child_schema = unsafe_load(_exported_child_schemas(schema)[1]) + @test unsafe_string(child_schema.format) == "i" + buffers = _exported_buffers(array) + offsets = _exported_vector(Int32, buffers[2], array[].length + 1) + @test collect(offsets) == Int32[0, 2, 2, 3] + child_array = unsafe_load(_exported_children(array)[1]) + @test child_array.length == 3 + child_buffers = collect( + unsafe_wrap(Array, child_array.buffers, Int(child_array.n_buffers); own=false), + ) + child_data = _exported_vector(Int32, child_buffers[2], child_array.length) + @test collect(child_data) == Int32[1, 2, 3] + _release_exported_array!(array) + _release_exported_schema!(schema) + + fixed = [(0x01, 0x02), (0x03, 0x04)] + schema, array = Arrow.to_c_data(Arrow.toarrowvector(fixed)) + @test unsafe_string(schema[].format) == "w:2" + buffers = _exported_buffers(array) + data = _exported_vector(UInt8, buffers[2], 4) + @test collect(data) == UInt8[0x01, 0x02, 0x03, 0x04] + _release_exported_array!(array) + _release_exported_schema!(schema) + + structs = [(a=Int32(1), b=1.5), (a=Int32(2), b=2.5)] + schema, array = Arrow.to_c_data(Arrow.toarrowvector(structs)) + @test unsafe_string(schema[].format) == "+s" + imported = _import_exported(schema, array) + @test Tables.columnnames(imported) == [:a, :b] + @test collect(imported.a) == Int32[1, 2] + @test collect(imported.b) == [1.5, 2.5] + Arrow.release_c_data(imported) + end + + @testset "table names and metadata" begin + tbl = (col1=Int32[1, 2], col2=Float64[1.5, 2.5]) + meta = Dict("source" => "export") + colmeta = Dict("unit" => "id") + arrow_tbl = Arrow.Table( + Arrow.tobuffer(tbl; metadata=meta, colmetadata=Dict(:col1 => colmeta)), + ) + schema, array = Arrow.to_c_data(arrow_tbl; names=["left", "right"]) + imported = _import_exported(schema, array) + @test Tables.columnnames(imported) == [:left, :right] + @test collect(imported.left) == Int32[1, 2] + @test collect(imported.right) == [1.5, 2.5] + @test DataAPI.metadata(imported, "source") == "export" + @test DataAPI.colmetadata(imported, :left, "unit") == "id" + Arrow.release_c_data(imported) + end + + @testset "null and empty arrays" begin + nulls = Arrow.toarrowvector([missing, missing]) + schema, array = Arrow.to_c_data(nulls) + @test unsafe_string(schema[].format) == "n" + @test array[].n_buffers == 0 + imported = _import_exported(schema, array) + @test isequal(collect(imported), [missing, missing]) + Arrow.release_c_data(imported) + + empty = Arrow.toarrowvector(Int32[]) + schema, array = Arrow.to_c_data(empty) + @test array[].length == 0 + @test array[].n_buffers == 2 + imported = _import_exported(schema, array) + @test collect(imported) == Int32[] + Arrow.release_c_data(imported) + end + + @testset "release ordering and GC roots" begin + schema, array = Arrow.to_c_data(Arrow.toarrowvector(Int32[1, 2, 3])) + schema_token = UInt(schema[].private_data) + array_token = UInt(array[].private_data) + imported = _import_exported(schema, array) + # The import moves the exported base structures into the importer. + @test schema[].release == C_NULL + @test array[].release == C_NULL + @test haskey(Arrow._CDATA_EXPORT_SCHEMA_OWNERS, schema_token) + @test haskey(Arrow._CDATA_EXPORT_ARRAY_OWNERS, array_token) + scratch = [Vector{UInt8}(undef, 4096) for _ = 1:128] + @test sum(length, scratch) > 0 + GC.gc(true) + @test collect(imported) == Int32[1, 2, 3] + Arrow.release_c_data(imported) + @test !haskey(Arrow._CDATA_EXPORT_SCHEMA_OWNERS, schema_token) + @test !haskey(Arrow._CDATA_EXPORT_ARRAY_OWNERS, array_token) + + nested = Arrow.toarrowvector([(a=Int32(4), b="x"), (a=Int32(5), b="y")]) + schema, array = Arrow.to_c_data(nested) + array_release = array[].release + schema_release = schema[].release + schema_token = UInt(schema[].private_data) + ccall(array_release, Cvoid, (Ptr{Arrow.ArrowArray},), _array_ref_ptr(array)) + @test array[].release == C_NULL + @test schema[].release == schema_release + @test haskey(Arrow._CDATA_EXPORT_SCHEMA_OWNERS, schema_token) + scratch = [Vector{UInt8}(undef, 4096) for _ = 1:128] + @test sum(length, scratch) > 0 + GC.gc(true) + @test unsafe_string(schema[].format) == "+s" + first_child = unsafe_load(schema[].children, 1) + second_child = unsafe_load(schema[].children, 2) + @test unsafe_string(unsafe_load(first_child).name) == "a" + @test unsafe_string(unsafe_load(first_child).format) == "i" + @test unsafe_string(unsafe_load(second_child).name) == "b" + @test unsafe_string(unsafe_load(second_child).format) == "u" + ccall(schema_release, Cvoid, (Ptr{Arrow.ArrowSchema},), _schema_ref_ptr(schema)) + @test schema[].release == C_NULL + ccall(array_release, Cvoid, (Ptr{Arrow.ArrowArray},), _array_ref_ptr(array)) + ccall(schema_release, Cvoid, (Ptr{Arrow.ArrowSchema},), _schema_ref_ptr(schema)) + @test array[].release == C_NULL + @test schema[].release == C_NULL + + schema, array = Arrow.to_c_data(Arrow.toarrowvector(Int32[9, 10])) + schema_copy = Ref( + Arrow.ArrowSchema( + schema[].format, + schema[].name, + schema[].metadata, + schema[].flags, + schema[].n_children, + schema[].children, + schema[].dictionary, + schema[].release, + schema[].private_data, + ), + ) + array_copy = Ref( + Arrow.ArrowArray( + array[].length, + array[].null_count, + array[].offset, + array[].n_buffers, + array[].n_children, + array[].buffers, + array[].children, + array[].dictionary, + array[].release, + array[].private_data, + ), + ) + ccall( + array_copy[].release, + Cvoid, + (Ptr{Arrow.ArrowArray},), + _array_ref_ptr(array_copy), + ) + @test array_copy[].release == C_NULL + @test array[].release == C_NULL + ccall( + schema_copy[].release, + Cvoid, + (Ptr{Arrow.ArrowSchema},), + _schema_ref_ptr(schema_copy), + ) + @test schema_copy[].release == C_NULL + @test schema[].release == C_NULL + + schema, array = let + col = Arrow.toarrowvector(Int32[6, 7, 8]) + Arrow.to_c_data(col) + end + GC.gc(true) + GC.gc(true) + imported = _import_exported(schema, array) + @test collect(imported) == Int32[6, 7, 8] + Arrow.release_c_data(imported) + + schema, array = Arrow.to_c_data( + Arrow.toarrowvector([(a=Int32(1), b=Int32(2)), (a=Int32(3), b=Int32(4))]), + ) + child_schema_ptr = _exported_child_schemas(schema)[1] + child_array_ptr = _exported_children(array)[1] + child_schema_copy = Ref(unsafe_load(child_schema_ptr)) + child_array_copy = Ref(unsafe_load(child_array_ptr)) + _mark_exported_schema_released!(child_schema_ptr) + _mark_exported_array_released!(child_array_ptr) + _release_exported_schema!(schema) + _release_exported_array!(array) + @test schema[].release == C_NULL + @test array[].release == C_NULL + @test child_schema_copy[].release != C_NULL + @test child_array_copy[].release != C_NULL + scratch = [Vector{UInt8}(undef, 4096) for _ = 1:128] + @test sum(length, scratch) > 0 + GC.gc(true) + @test unsafe_string(child_schema_copy[].name) == "a" + @test unsafe_string(child_schema_copy[].format) == "i" + child_buffers = collect( + unsafe_wrap( + Array, + child_array_copy[].buffers, + Int(child_array_copy[].n_buffers); + own=false, + ), + ) + child_data = _exported_vector(Int32, child_buffers[2], child_array_copy[].length) + @test collect(child_data) == Int32[1, 3] + ccall( + child_array_copy[].release, + Cvoid, + (Ptr{Arrow.ArrowArray},), + _array_ref_ptr(child_array_copy), + ) + ccall( + child_schema_copy[].release, + Cvoid, + (Ptr{Arrow.ArrowSchema},), + _schema_ref_ptr(child_schema_copy), + ) + @test child_array_copy[].release == C_NULL + @test child_schema_copy[].release == C_NULL + end + + @testset "unsupported arrays" begin + schema_count = length(Arrow._CDATA_EXPORT_SCHEMA_OWNERS) + array_count = length(Arrow._CDATA_EXPORT_ARRAY_OWNERS) + map_col = Arrow.toarrowvector([Dict(Int32(1) => Float32(2))]) + @test_throws ArgumentError Arrow.to_c_data(map_col) + @test length(Arrow._CDATA_EXPORT_SCHEMA_OWNERS) == schema_count + @test length(Arrow._CDATA_EXPORT_ARRAY_OWNERS) == array_count + + dict_col = Arrow.toarrowvector(Arrow.DictEncode(["a", "b"])) + @test dict_col isa Arrow.DictEncoded + @test_throws ArgumentError Arrow.to_c_data(dict_col) + @test length(Arrow._CDATA_EXPORT_SCHEMA_OWNERS) == schema_count + @test length(Arrow._CDATA_EXPORT_ARRAY_OWNERS) == array_count + end + + @testset "malformed export layouts" begin + schema_count = length(Arrow._CDATA_EXPORT_SCHEMA_OWNERS) + array_count = length(Arrow._CDATA_EXPORT_ARRAY_OWNERS) + + short_validity = Arrow.ValidityBitmap(UInt8[0xff], 1, 1, 1) + bad_validity = Arrow.Primitive{Union{Missing,Int32},Vector{Int32}}( + UInt8[], + short_validity, + Int32[1, 2, 3, 4, 5, 6, 7, 8, 9], + 9, + nothing, + ) + @test_throws ArgumentError Arrow.to_c_data(bad_validity) + + bad_null_count = Arrow.Primitive{Union{Missing,Int32},Vector{Int32}}( + UInt8[], + Arrow.ValidityBitmap(UInt8[0x00], 1, 1, 2), + Int32[1], + 1, + nothing, + ) + @test_throws ArgumentError Arrow.to_c_data(bad_null_count) + + wrapped_validity = Arrow.Primitive{Union{Missing,Int32},Vector{Int32}}( + UInt8[], + Arrow.ValidityBitmap(UInt8[0xff], typemax(Int), typemax(Int), 1), + Int32[1, 2, 3, 4, 5, 6, 7, 8, 9], + 9, + nothing, + ) + @test_throws ArgumentError Arrow.to_c_data(wrapped_validity) + + bad_bool = Arrow.BoolVector{Bool}( + UInt8[0xff], + typemax(Int), + Arrow.ValidityBitmap(UInt8[], 1, 0, 0), + 9, + nothing, + ) + @test_throws ArgumentError Arrow.to_c_data(bad_bool) + + child = Arrow.toarrowvector(Int32[]) + bad_offsets = Arrow.List{Vector{Int32},Int32,typeof(child)}( + UInt8[], + Arrow.ValidityBitmap(UInt8[], 1, 0, 0), + Arrow.Offsets(UInt8[], Int32[0]), + child, + typemax(Int), + nothing, + ) + @test_throws ArgumentError Arrow.to_c_data(bad_offsets) + + # A failure on a later child must release the children already built. + good_child = Arrow.Primitive{Int32,Vector{Int32}}( + UInt8[], + Arrow.ValidityBitmap(UInt8[], 1, 3, 0), + Int32[1, 2, 3], + 3, + nothing, + ) + bad_child = Arrow.Primitive{Union{Missing,Int32},Vector{Int32}}( + UInt8[], + Arrow.ValidityBitmap(UInt8[0xff], 1, 1, 1), + Int32[1, 2, 3], + 3, + nothing, + ) + T = NamedTuple{(:a, :b),Tuple{Int32,Union{Missing,Int32}}} + data = (good_child, bad_child) + bad_struct = Arrow.Struct{T,typeof(data),(:a, :b)}( + Arrow.ValidityBitmap(UInt8[], 1, 3, 0), + data, + 3, + nothing, + ) + @test_throws ArgumentError Arrow.to_c_data(bad_struct) + + @test length(Arrow._CDATA_EXPORT_SCHEMA_OWNERS) == schema_count + @test length(Arrow._CDATA_EXPORT_ARRAY_OWNERS) == array_count + end + + @testset "malformed nested export layouts" begin + schema_count = length(Arrow._CDATA_EXPORT_SCHEMA_OWNERS) + array_count = length(Arrow._CDATA_EXPORT_ARRAY_OWNERS) + + child = Arrow.toarrowvector(Int32[1]) + bad_struct = + Arrow.Struct{NamedTuple{(:a,),Tuple{Int32}},Tuple{typeof(child)},(:a,)}( + Arrow.ValidityBitmap(UInt8[], 1, 0, 0), + (child,), + 2, + nothing, + ) + @test_throws ArgumentError Arrow.to_c_data(bad_struct) + + fixed_child = Arrow.toarrowvector(Int32[1, 2, 3, 4, 5]) + bad_fixed = Arrow.FixedSizeList{NTuple{3,Int32},typeof(fixed_child)}( + UInt8[], + Arrow.ValidityBitmap(UInt8[], 1, 0, 0), + fixed_child, + 2, + nothing, + ) + @test_throws ArgumentError Arrow.to_c_data(bad_fixed) + + list_child = Arrow.toarrowvector(Int32[1]) + bad_list = Arrow.List{Vector{Int32},Int32,typeof(list_child)}( + UInt8[], + Arrow.ValidityBitmap(UInt8[], 1, 0, 0), + Arrow.Offsets(UInt8[], Int32[0, 2]), + list_child, + 1, + nothing, + ) + @test_throws ArgumentError Arrow.to_c_data(bad_list) + + @test length(Arrow._CDATA_EXPORT_SCHEMA_OWNERS) == schema_count + @test length(Arrow._CDATA_EXPORT_ARRAY_OWNERS) == array_count + end + + @testset "repeated GC stress" begin + for i = 1:24 + tbl = ( + id=Int32[i, i + 1, i + 2], + label=["a$(i)", "b$(i)", "c$(i)"], + flags=Union{Bool,Missing}[true, missing, isodd(i)], + ) + schema, array = Arrow.to_c_data(tbl) + imported = _import_exported(schema, array) + GC.gc(true) + GC.gc(true) + @test collect(imported.id) == tbl.id + @test collect(imported.label) == tbl.label + @test isequal(collect(imported.flags), tbl.flags) + Arrow.release_c_data(imported) + @test schema[].release == C_NULL + @test array[].release == C_NULL + end + end + + @testset "optional PyArrow C Data smoke" begin + python = Sys.which("python3") + if python === nothing + @test_skip "python3 not available" + else + script = """ + try: + import pyarrow as pa + except Exception: + print("skip: pyarrow unavailable") + raise SystemExit(0) + arr = pa.array([1, None, 3], type=pa.int32()) + if not hasattr(arr, "__arrow_c_array__"): + print("skip: pyarrow C array export unavailable") + raise SystemExit(0) + if not hasattr(pa.Array, "_import_from_c_capsule"): + print("skip: pyarrow C capsule import unavailable") + raise SystemExit(0) + capsules = arr.__arrow_c_array__() + out = pa.Array._import_from_c_capsule(*capsules) + assert out.to_pylist() == [1, None, 3] + print("ok") + """ + out = readchomp(`$python -c $script`) + if startswith(out, "skip:") + @test_skip out + else + @test out == "ok" + end + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 315d1b60..784ff562 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