From 1b9428156b9a0429b8ac2e12a195334c4109f0b2 Mon Sep 17 00:00:00 2001 From: samtalki <10187005+samtalki@users.noreply.github.com> Date: Sun, 5 Jul 2026 22:04:43 -0400 Subject: [PATCH 1/4] GH-184: Add C Data Interface import foundation Add immutable ArrowSchema and ArrowArray ABI structs plus C Data import for null, primitive, and root struct arrays. Root structs are exposed as Tables.jl column tables with DataAPI metadata. Validate pointer layout, child counts, buffer counts, offsets, metadata bounds, flag bits, dictionary pointer consistency, release idempotency, and reads after release. Move the base ArrowSchema and ArrowArray structures into Julia owned storage at import and mark the sources released without calling their callbacks, following the C Data Interface move semantics used by arrow-rs from_raw and nanoarrow ArrowArrayMove. Callers may free or reuse the passed structures once from_c_data returns; release_c_data or finalization releases the moved copies through the producer callbacks. Keep imported buffers rooted behind a shared release owner. Copy misaligned fixed width buffers into aligned Julia storage before typed access, and materialize Julia owned arrays from copy and collect. Adjust dictionary encoding offsets for non 1-based reference pools so categorical pools with missing values keep correct dictionary indices. Co-authored-by: Robert Buessow Co-authored-by: Olle Martensson Co-authored-by: Claude Fable 5 Generated-by: OpenAI Codex --- src/Arrow.jl | 4 +- src/arraytypes/dictencoding.jl | 27 +- src/cdata.jl | 928 +++++++++++++++++++++++++++++++++ test/cdata.jl | 819 +++++++++++++++++++++++++++++ test/runtests.jl | 1 + 5 files changed, 1775 insertions(+), 4 deletions(-) create mode 100644 src/cdata.jl create mode 100644 test/cdata.jl diff --git a/src/Arrow.jl b/src/Arrow.jl index 6f3ccdf8..5fb38ddc 100644 --- a/src/Arrow.jl +++ b/src/Arrow.jl @@ -26,11 +26,12 @@ This implementation supports the 1.0 version of the specification, including sup * Extension types * Streaming, file, record batch, and replacement and isdelta dictionary messages * Buffer compression/decompression via the standard LZ4 frame and Zstd formats + * C Data Interface import It currently doesn't include support for: * Tensors or sparse tensors * Flight RPC - * C data interface + * C Data Interface export Third-party data formats: * csv and parquet support via the existing [CSV.jl](https://github.com/JuliaData/CSV.jl) and [Parquet.jl](https://github.com/JuliaIO/Parquet.jl) packages @@ -76,6 +77,7 @@ include("utils.jl") include("arraytypes/arraytypes.jl") include("eltypes.jl") include("table.jl") +include("cdata.jl") include("write.jl") include("append.jl") include("show.jl") diff --git a/src/arraytypes/dictencoding.jl b/src/arraytypes/dictencoding.jl index 3e3576c5..71f3c23e 100644 --- a/src/arraytypes/dictencoding.jl +++ b/src/arraytypes/dictencoding.jl @@ -119,6 +119,27 @@ signedtype(::Type{UInt32}) = Int32 signedtype(::Type{UInt64}) = Int64 signedtype(::Type{T}) where {T<:Signed} = T +function _dict_indices(refa, pool, validity::ValidityBitmap) + ET = signedtype(length(pool)) + inds = Vector{ET}(undef, length(refa)) + lo = firstindex(pool) + hi = lastindex(pool) + j = 1 + @inbounds for idx in eachindex(refa) + ref = refa[idx] + # Per the Arrow columnar format spec, dictionary data stores 0-based indices. + if lo <= ref <= hi + inds[j] = ET(ref - lo) + elseif validity[j] + throw(ArgumentError("dictionary reference index is out of range")) + else + inds[j] = zero(ET) + end + j += 1 + end + return inds +end + indtype(d::DictEncoded{T,S,A}) where {T,S,A} = S indtype(c::Compressed{Z,A}) where {Z,A<:DictEncoded} = indtype(c.data) @@ -229,12 +250,12 @@ function arrowvector( else pool = DataAPI.refpool(x) refa = DataAPI.refarray(x) - inds = copyto!(similar(Vector{signedtype(length(pool))}, length(refa)), refa) end # adjust to "offset" instead of index - inds .-= firstindex(refa) + inds = _dict_indices(refa, pool, validity) + pooldata = firstindex(pool) == 1 ? pool : collect(pool) data = arrowvector( - pool, + pooldata, i, nl, fi, diff --git a/src/cdata.jl b/src/cdata.jl new file mode 100644 index 00000000..efd18a1b --- /dev/null +++ b/src/cdata.jl @@ -0,0 +1,928 @@ +# 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 + +abstract type CDataFormat end + +struct CDataNullFormat <: CDataFormat end +struct CDataPrimitiveFormat <: CDataFormat + storage::Type +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 CDataStruct{T,S,names} <: CDataVector{T} + owner::CDataOwner + validity::CDataValidity + data::S + 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::CDataStruct) = (x.validity.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::CDataStruct{T,S,names}, + i::Integer, +) where {T,S,names} + return _with_live(x) do + @boundscheck checkbounds(x, i) + !_valid(x.validity, i) && return missing + NT = Base.nonmissingtype(T) + vals = ntuple(j -> @inbounds(x.data[j][i]), fieldcount(S)) + return NT(vals) + 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 == "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 == "+s" && return CDataStructFormat() + throw(ArgumentError("unsupported Arrow C Data format string: $format")) +end + +_expected_buffers(::CDataNullFormat) = 0 +_expected_buffers(::CDataPrimitiveFormat) = 2 +_expected_buffers(::CDataStructFormat) = 1 + +_expected_children(::CDataNullFormat) = 0 +_expected_children(::CDataPrimitiveFormat) = 0 +_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 _validate_data_layout(::CDataNullFormat, node::CDataNode, total::Int) + return +end + +function _aligned(ptr::Ptr{Cvoid}, ::Type{T}) where {T} + return UInt(ptr) % Base.datatype_alignment(T) == 0 +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(::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 _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 _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 _struct_type(schema::ArrowSchema, validity::CDataValidity, data, names) + NT = NamedTuple{names,Tuple{(eltype(x) for x in data)...}} + return _nullable(schema, validity.null_count) ? Union{NT,Missing} : NT +end + +function _import_node( + ::CDataStructFormat, + node::CDataNode, + owner::CDataOwner, + convert::Bool, +) + validity = _make_validity(node) + names, data = _struct_columns(node, owner, convert) + T = _struct_type(node.schema, validity, data, names) + return CDataStruct{T,typeof(data),names}(owner, validity, data, 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...) diff --git a/test/cdata.jl b/test/cdata.jl new file mode 100644 index 00000000..6eaf4d30 --- /dev/null +++ b/test/cdata.jl @@ -0,0 +1,819 @@ +# 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 + +const _CDATA_RELEASE_SCHEMA = + @cfunction(_cdata_release_schema, Cvoid, (Ptr{Arrow.ArrowSchema},)) +const _CDATA_RELEASE_ARRAY = + @cfunction(_cdata_release_array, Cvoid, (Ptr{Arrow.ArrowArray},)) +const _CDATA_REENTRANT_OBJECT = Ref{Any}(nothing) +const _CDATA_REENTRANT_RELEASES = Ref(0) + +function _cdata_release_array_reentrant(ptr::Ptr{Arrow.ArrowArray}) + _CDATA_REENTRANT_RELEASES[] += 1 + x = _CDATA_REENTRANT_OBJECT[] + x === nothing || Arrow.release_c_data(x) + _cdata_release_array(ptr) + return +end + +const _CDATA_RELEASE_ARRAY_REENTRANT = + @cfunction(_cdata_release_array_reentrant, Cvoid, (Ptr{Arrow.ArrowArray},)) + +mutable struct CDataFixture + schema::Ref{Arrow.ArrowSchema} + array::Ref{Arrow.ArrowArray} + roots::Vector{Any} +end + +const _CDATA_FIXTURE_ROOTS = Any[] + +_schema_ptr(x::CDataFixture) = Base.unsafe_convert(Ptr{Arrow.ArrowSchema}, x.schema) +_array_ptr(x::CDataFixture) = Base.unsafe_convert(Ptr{Arrow.ArrowArray}, x.array) + +function _cstring_root(s::Union{Nothing,String}, roots) + s === nothing && return Cstring(C_NULL) + bytes = Vector{UInt8}(s * "\0") + push!(roots, bytes) + return Cstring(pointer(bytes)) +end + +function _metadata_root(meta::Union{Nothing,AbstractDict}, roots) + meta === nothing && return Cstring(C_NULL) + io = IOBuffer() + write(io, Int32(length(meta))) + for (k, v) in meta + kb = codeunits(String(k)) + vb = codeunits(String(v)) + write(io, Int32(length(kb))) + write(io, kb) + write(io, Int32(length(vb))) + write(io, vb) + end + bytes = take!(io) + push!(roots, bytes) + return Cstring(pointer(bytes)) +end + +function _cdata_fixture( + fmt::String, + len::Integer, + buffers::Vector{Ptr{Cvoid}}; + name=nothing, + metadata=nothing, + flags::Int64=0, + null_count::Int64=0, + offset::Int64=0, + children::Vector{CDataFixture}=CDataFixture[], +) + roots = Any[] + append!(roots, children) + schema_ptrs = + [Base.unsafe_convert(Ptr{Arrow.ArrowSchema}, child.schema) for child in children] + array_ptrs = + [Base.unsafe_convert(Ptr{Arrow.ArrowArray}, child.array) for child in children] + if !isempty(children) + push!(roots, schema_ptrs) + push!(roots, array_ptrs) + end + buffer_ptrs = copy(buffers) + if !isempty(buffer_ptrs) + push!(roots, buffer_ptrs) + end + schema = Ref( + Arrow.ArrowSchema( + _cstring_root(fmt, roots), + _cstring_root(name, roots), + _metadata_root(metadata, roots), + flags, + Int64(length(children)), + isempty(children) ? Ptr{Ptr{Arrow.ArrowSchema}}(C_NULL) : + Ptr{Ptr{Arrow.ArrowSchema}}(pointer(schema_ptrs)), + Ptr{Arrow.ArrowSchema}(C_NULL), + _CDATA_RELEASE_SCHEMA, + Ptr{Cvoid}(C_NULL), + ), + ) + array = Ref( + Arrow.ArrowArray( + Int64(len), + null_count, + offset, + Int64(length(buffer_ptrs)), + Int64(length(children)), + isempty(buffer_ptrs) ? Ptr{Ptr{Cvoid}}(C_NULL) : + Ptr{Ptr{Cvoid}}(pointer(buffer_ptrs)), + isempty(children) ? Ptr{Ptr{Arrow.ArrowArray}}(C_NULL) : + Ptr{Ptr{Arrow.ArrowArray}}(pointer(array_ptrs)), + Ptr{Arrow.ArrowArray}(C_NULL), + _CDATA_RELEASE_ARRAY, + Ptr{Cvoid}(C_NULL), + ), + ) + push!(roots, schema) + push!(roots, array) + fixture = CDataFixture(schema, array, roots) + push!(_CDATA_FIXTURE_ROOTS, fixture) + return fixture +end + +function _primitive_fixture( + fmt, + data::Vector{T}; + validity=nothing, + null_count::Int64=0, + flags::Int64=0, + offset::Int64=0, + len::Int=length(data) - Int(offset), + name=nothing, + metadata=nothing, +) where {T} + roots = Any[data] + buffers = Ptr{Cvoid}[ + validity === nothing ? Ptr{Cvoid}(C_NULL) : Ptr{Cvoid}(pointer(validity)), + isempty(data) ? Ptr{Cvoid}(C_NULL) : Ptr{Cvoid}(pointer(data)), + ] + validity !== nothing && push!(roots, validity) + fixture = _cdata_fixture( + fmt, + len, + buffers; + name=name, + metadata=metadata, + flags=flags, + null_count=null_count, + offset=offset, + ) + append!(fixture.roots, roots) + return fixture +end + +function _unaligned_primitive_fixture( + fmt, + data::Vector{T}; + offset::Int64=Int64(0), + len::Int64=Int64(length(data)) - offset, +) where {T} + nbytes = length(data) * sizeof(T) + bytes = Vector{UInt8}(undef, nbytes + 1) + unsafe_copyto!(pointer(bytes, 2), Ptr{UInt8}(pointer(data)), nbytes) + buffers = Ptr{Cvoid}[Ptr{Cvoid}(C_NULL), Ptr{Cvoid}(pointer(bytes, 2))] + fixture = _cdata_fixture(fmt, len, buffers; offset=offset) + append!(fixture.roots, Any[data, bytes]) + return fixture +end + +function _set_schema!( + f::CDataFixture; + format=f.schema[].format, + name=f.schema[].name, + metadata=f.schema[].metadata, + flags=f.schema[].flags, + n_children=f.schema[].n_children, + children=f.schema[].children, + dictionary=f.schema[].dictionary, + release=f.schema[].release, + private_data=f.schema[].private_data, +) + f.schema[] = Arrow.ArrowSchema( + format, + name, + metadata, + flags, + n_children, + children, + dictionary, + release, + private_data, + ) + return f +end + +function _set_array!( + f::CDataFixture; + length=f.array[].length, + null_count=f.array[].null_count, + offset=f.array[].offset, + n_buffers=f.array[].n_buffers, + buffers=f.array[].buffers, + n_children=f.array[].n_children, + children=f.array[].children, + dictionary=f.array[].dictionary, + release=f.array[].release, + private_data=f.array[].private_data, +) + f.array[] = Arrow.ArrowArray( + length, + null_count, + offset, + n_buffers, + n_children, + buffers, + children, + dictionary, + release, + private_data, + ) + return f +end + +@testset "Arrow C Data Interface import" begin + @testset "ABI layout" begin + ptr = sizeof(Ptr{Cvoid}) + @test sizeof(Arrow.ArrowSchema) == 7 * ptr + 2 * sizeof(Int64) + @test sizeof(Arrow.ArrowArray) == 5 * ptr + 5 * sizeof(Int64) + @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 "struct root table with names and metadata" begin + xchild = + _primitive_fixture("i", Int32[1, 2, 3]; name="x", metadata=Dict("unit" => "id")) + ychild = _primitive_fixture("g", Float64[1.5, 2.5, 3.5]; name="y") + 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, Float64) + @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) == [1.5, 2.5, 3.5] + @test copy(tbl) == (x=Int32[1, 2, 3], y=[1.5, 2.5, 3.5]) + @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")) + + 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 "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 + + 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) + + 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) + 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 From e84c39d67a5954960af1c5eb48a85fbbb583147f Mon Sep 17 00:00:00 2001 From: samtalki <10187005+samtalki@users.noreply.github.com> Date: Mon, 6 Jul 2026 14:03:27 -0400 Subject: [PATCH 2/4] GH-184: Expand C Data Interface import coverage Add import support for boolean, string, binary, list, fixed size, temporal, timestamp, duration, interval, and decimal C Data formats. Extend malformed input, offset, null count, release, copy, collect, metadata, and GC rooting coverage for the broader importer. Co-authored-by: Robert Buessow Co-authored-by: Olle Martensson Co-authored-by: Claude Fable 5 Generated-by: OpenAI Codex --- src/cdata.jl | 527 ++++++++++++++++++++++++++++++++++++++++++++++++-- test/cdata.jl | 397 ++++++++++++++++++++++++++++++++++++- 2 files changed, 905 insertions(+), 19 deletions(-) diff --git a/src/cdata.jl b/src/cdata.jl index efd18a1b..aac8b419 100644 --- a/src/cdata.jl +++ b/src/cdata.jl @@ -57,13 +57,25 @@ 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 @@ -111,10 +123,58 @@ struct CDataPrimitive{T,S,A<:AbstractVector{S}} <: CDataVector{T} metadata::Union{Nothing,Base.ImmutableDict{String,String}} end -struct CDataStruct{T,S,names} <: CDataVector{T} +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 @@ -158,7 +218,12 @@ Base.IndexStyle(::Type{<:CDataVector}) = Base.IndexLinear() Base.size(x::CDataNull) = (x.len,) Base.size(x::CDataPrimitive) = size(x.data) -Base.size(x::CDataStruct) = (x.validity.len,) +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)) @@ -239,16 +304,91 @@ end 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,names}, + x::CDataStruct{T,S,fnames}, i::Integer, -) where {T,S,names} +) where {T,S,fnames} return _with_live(x) do @boundscheck checkbounds(x, i) - !_valid(x.validity, i) && return missing + 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) - vals = ntuple(j -> @inbounds(x.data[j][i]), fieldcount(S)) - return NT(vals) + if isnamedtuple(NT) || istuple(NT) + return ArrowTypes.fromarrow(T, NT(vals)) + else + return ArrowTypes.fromarrow(T, _fromarrowstruct(NT, Val{fnames}(), vals...)) + end end end @@ -449,6 +589,7 @@ 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) @@ -460,16 +601,112 @@ function _parse_c_data_format(format::AbstractString) 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) @@ -710,12 +947,20 @@ function _validate_layout(node::CDataNode) 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 _aligned(ptr::Ptr{Cvoid}, ::Type{T}) where {T} - return UInt(ptr) % Base.datatype_alignment(T) == 0 +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) @@ -726,6 +971,62 @@ function _validate_data_layout(format::CDataPrimitiveFormat, node::CDataNode, to 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 @@ -754,6 +1055,80 @@ function _validate_value_layout(node::CDataNode) 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) @@ -789,6 +1164,10 @@ function _wrap_data(ptr::Ptr{Cvoid}, ::Type{T}, offset::Int, len::Int) where {T} 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 @@ -830,9 +1209,113 @@ function _struct_columns(node::CDataNode, owner::CDataOwner, convert::Bool) return Tuple(names), Tuple(columns) end -function _struct_type(schema::ArrowSchema, validity::CDataValidity, data, names) - NT = NamedTuple{names,Tuple{(eltype(x) for x in data)...}} - return _nullable(schema, validity.null_count) ? Union{NT,Missing} : NT +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( @@ -842,9 +1325,23 @@ function _import_node( convert::Bool, ) validity = _make_validity(node) - names, data = _struct_columns(node, owner, convert) - T = _struct_type(node.schema, validity, data, names) - return CDataStruct{T,typeof(data),names}(owner, validity, data, node.metadata) + 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) diff --git a/test/cdata.jl b/test/cdata.jl index 6eaf4d30..0cceef0f 100644 --- a/test/cdata.jl +++ b/test/cdata.jl @@ -53,6 +53,8 @@ function _cdata_release_array(ptr::Ptr{Arrow.ArrowArray}) return end +using Dates + const _CDATA_RELEASE_SCHEMA = @cfunction(_cdata_release_schema, Cvoid, (Ptr{Arrow.ArrowSchema},)) const _CDATA_RELEASE_ARRAY = @@ -337,10 +339,243 @@ end @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")) - ychild = _primitive_fixture("g", Float64[1.5, 2.5, 3.5]; name="y") + 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, @@ -350,7 +585,7 @@ end ) tbl = Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) @test Tables.columnnames(tbl) == [:x, :y] - @test Tables.schema(tbl).types == (Int32, Float64) + @test Tables.schema(tbl).types == (Int32, String) @test length(tbl) == 2 @test Tables.rowcount(tbl) == 3 @test Tables.istable(typeof(tbl)) @@ -362,8 +597,8 @@ end @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) == [1.5, 2.5, 3.5] - @test copy(tbl) == (x=Int32[1, 2, 3], y=[1.5, 2.5, 3.5]) + @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") @@ -383,6 +618,27 @@ end @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] @@ -524,6 +780,47 @@ end @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)) @@ -721,6 +1018,73 @@ end 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, @@ -810,10 +1174,35 @@ end 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 end From c8d7f5531b9f0541812c73159329778489606c19 Mon Sep 17 00:00:00 2001 From: samtalki <10187005+samtalki@users.noreply.github.com> Date: Mon, 6 Jul 2026 02:50:21 -0400 Subject: [PATCH 3/4] GH-184: Export Arrow arrays through the C Data Interface Add array and table export entry points with independent schema and array owners. Release children that were already exported when a later child fails to build, so partially constructed exports do not leak owner registry entries. Cover primitive and struct round trips, release order, moved child ownership, GC, metadata, malformed layout, unsupported export cases, and exported buffer layout checks for nested, string, and fixed size arrays. Co-authored-by: Robert Buessow Co-authored-by: Olle Martensson Co-authored-by: Claude Fable 5 Generated-by: OpenAI Codex --- src/Arrow.jl | 4 +- src/cdata.jl | 636 ++++++++++++++++++++++++++++++++++++++++++++++++++ test/cdata.jl | 473 +++++++++++++++++++++++++++++++++++++ 3 files changed, 1111 insertions(+), 2 deletions(-) diff --git a/src/Arrow.jl b/src/Arrow.jl index 5fb38ddc..0aed012f 100644 --- a/src/Arrow.jl +++ b/src/Arrow.jl @@ -26,12 +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 + * C Data Interface import and export It currently doesn't include support for: * Tensors or sparse tensors * Flight RPC - * C Data Interface export Third-party data formats: * csv and parquet support via the existing [CSV.jl](https://github.com/JuliaData/CSV.jl) and [Parquet.jl](https://github.com/JuliaIO/Parquet.jl) packages @@ -140,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/cdata.jl b/src/cdata.jl index aac8b419..7190e511 100644 --- a/src/cdata.jl +++ b/src/cdata.jl @@ -1423,3 +1423,639 @@ 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/cdata.jl b/test/cdata.jl index 0cceef0f..d8e1e553 100644 --- a/test/cdata.jl +++ b/test/cdata.jl @@ -1206,3 +1206,476 @@ end bad(child) 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 +end From 5a4a7eaaf49d7bb1a2b44a9ad8eb9cc84610b87f Mon Sep 17 00:00:00 2001 From: samtalki <10187005+samtalki@users.noreply.github.com> Date: Mon, 6 Jul 2026 14:09:53 -0400 Subject: [PATCH 4/4] GH-184: Harden C Data Interface tests Add deterministic malformed import fuzzing, nested malformed input and export layout checks, a C producer smoke test, repeated GC stress, and an optional PyArrow C Data smoke test with skips when the external dependency is unavailable. Co-authored-by: Robert Buessow Co-authored-by: Olle Martensson Generated-by: OpenAI Codex --- test/Project.toml | 1 + test/cdata.jl | 309 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 310 insertions(+) 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 index d8e1e553..c77ce876 100644 --- a/test/cdata.jl +++ b/test/cdata.jl @@ -54,6 +54,8 @@ function _cdata_release_array(ptr::Ptr{Arrow.ArrowArray}) end using Dates +using Libdl +using Random const _CDATA_RELEASE_SCHEMA = @cfunction(_cdata_release_schema, Cvoid, (Ptr{Arrow.ArrowSchema},)) @@ -242,6 +244,10 @@ function _set_schema!( return f end +function _replace_schema!(f::CDataFixture; kwargs...) + return _set_schema!(f; kwargs...) +end + function _set_array!( f::CDataFixture; length=f.array[].length, @@ -270,6 +276,10 @@ function _set_array!( 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}) @@ -1205,6 +1215,214 @@ end 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}) = @@ -1678,4 +1896,95 @@ end @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