From 6c0b4b89bc0a15636d0a58de29c67bb3f0515657 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/2] Add C Data Interface import foundation Add ArrowSchema and ArrowArray ABI structs plus C Data import for null and primitive arrays. Root struct arrays are exposed as Tables.jl column tables with DataAPI metadata. Add validation for pointer layout, child counts, buffers, offset arithmetic, metadata bounds, dictionary pointer consistency, release idempotency, and reads after release. Keep aligned imported buffers rooted behind a shared owner, and mirror arrow-rs by copying misaligned fixed width buffers into aligned storage before typed access. The copy and collect methods materialize Julia owned arrays. 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 Generated-by: OpenAI Codex --- src/Arrow.jl | 4 +- src/arraytypes/dictencoding.jl | 27 +- src/cdata.jl | 859 +++++++++++++++++++++++++++++ test/cdata.jl | 977 +++++++++++++++++++++++++++++++++ test/runtests.jl | 30 + 5 files changed, 1893 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..451decb1 --- /dev/null +++ b/src/cdata.jl @@ -0,0 +1,859 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +mutable struct ArrowSchema + format::Cstring + name::Cstring + metadata::Cstring + flags::Int64 + n_children::Int64 + children::Ptr{Ptr{ArrowSchema}} + dictionary::Ptr{ArrowSchema} + release::Ptr{Cvoid} + private_data::Ptr{Cvoid} +end + +mutable struct ArrowArray + length::Int64 + null_count::Int64 + offset::Int64 + n_buffers::Int64 + n_children::Int64 + buffers::Ptr{Ptr{Cvoid}} + children::Ptr{Ptr{ArrowArray}} + dictionary::Ptr{ArrowArray} + release::Ptr{Cvoid} + private_data::Ptr{Cvoid} +end + +const _CDATA_PTR_SIZE = sizeof(Ptr{Cvoid}) +@assert sizeof(ArrowSchema) == 7 * _CDATA_PTR_SIZE + 2 * sizeof(Int64) +@assert sizeof(ArrowArray) == 5 * _CDATA_PTR_SIZE + 5 * sizeof(Int64) + +const ARROW_FLAG_DICTIONARY_ORDERED = Int64(1) +const ARROW_FLAG_NULLABLE = Int64(2) +const ARROW_FLAG_MAP_KEYS_SORTED = Int64(4) + +const _CDATA_MAX_CHILDREN = 10_000 +const _CDATA_MAX_DEPTH = 128 +const _CDATA_MAX_BUFFERS = 64 +const _CDATA_MAX_FORMAT_BYTES = 4096 +const _CDATA_MAX_NAME_BYTES = 1 << 16 +const _CDATA_MAX_METADATA_PAIRS = 4096 +const _CDATA_MAX_METADATA_BYTES = 1 << 20 +const _CDATA_MAX_METADATA_FIELD_BYTES = 1 << 20 + +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_ptr::Ptr{ArrowSchema} + array_ptr::Ptr{ArrowArray} + released::Bool + lock::ReentrantLock +end + +function CDataOwner(schema_ptr::Ptr{ArrowSchema}, array_ptr::Ptr{ArrowArray}) + owner = CDataOwner(schema_ptr, array_ptr, false, ReentrantLock()) + finalizer(release_c_data, owner) + return owner +end + +struct CDataValidity + bytes::Vector{UInt8} + bitoffset::Int + len::Int + null_count::Int +end + +struct CDataNull{T} <: CDataVector{T} + owner::CDataOwner + len::Int + metadata::Union{Nothing,Base.ImmutableDict{String,String}} +end + +struct CDataPrimitive{T,S,A<:AbstractVector{S}} <: CDataVector{T} + owner::CDataOwner + validity::CDataValidity + data::A + metadata::Union{Nothing,Base.ImmutableDict{String,String}} +end + +struct 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::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::CDataSlice, i::Integer) + return _with_live(x) do + @boundscheck checkbounds(x, i) + return @inbounds x.parent[x.first + Int(i) - 1] + end +end + +@propagate_inbounds function Base.getindex(x::CDataMasked, i::Integer) + return _with_live(x) do + @boundscheck checkbounds(x, i) + if !_valid(x.parent_validity, i) + return missing + end + return @inbounds x.parent[i] + end +end + +function Base.collect(x::CDataVector{T}) where {T} + return _with_live(x) do + out = Vector{T}(undef, length(x)) + for i in eachindex(x) + @inbounds out[i] = x[i] + end + return out + end +end + +Base.copy(x::CDataVector) = collect(x) + +function Base.copy(t::CDataTable) + return _with_live(t) do + names = getfield(t, :names) + columns = getfield(t, :columns) + return NamedTuple{Tuple(names)}(Tuple(copy(col) for col in columns)) + end +end + +Tables.istable(::Type{CDataTable}) = true +Tables.columnaccess(::Type{CDataTable}) = true +Tables.columns(t::CDataTable) = t +Tables.columnnames(t::CDataTable) = getfield(t, :names) +Tables.schema(t::CDataTable) = Tables.Schema(getfield(t, :names), getfield(t, :types)) +Tables.getcolumn(t::CDataTable, i::Int) = (_check_live(t); getfield(t, :columns)[i]) +Tables.getcolumn(t::CDataTable, nm::Symbol) = (_check_live(t); getfield(t, :lookup)[nm]) +Tables.rowcount(t::CDataTable) = getfield(t, :rowcount) + +Base.getindex(t::CDataTable, i::Int) = Tables.getcolumn(t, i) +Base.getindex(t::CDataTable, nm::Symbol) = Tables.getcolumn(t, nm) +function Base.getproperty(t::CDataTable, nm::Symbol) + lookup = getfield(t, :lookup) + haskey(lookup, nm) && return Tables.getcolumn(t, nm) + return getfield(t, nm) +end +Base.propertynames(t::CDataTable, private::Bool=false) = + private ? fieldnames(typeof(t)) : Tuple(getfield(t, :names)) + +DataAPI.metadatasupport(::Type{CDataTable}) = (read=true, write=false) +DataAPI.colmetadatasupport(::Type{CDataTable}) = (read=true, write=false) + +function DataAPI.metadata(t::CDataTable; style::Bool=false) + meta = getmetadata(t) + meta === nothing && return Dict{String,Any}() + if style + return Dict(k => (v, :default) for (k, v) in meta) + else + return Dict(meta) + end +end + +function DataAPI.metadata(t::CDataTable, key::AbstractString; style::Bool=false) + meta = getmetadata(t) + meta === nothing && throw(KeyError(key)) + val = meta[key] + return style ? (val, :default) : val +end + +function DataAPI.metadata(t::CDataTable, key::AbstractString, default; style::Bool=false) + meta = getmetadata(t) + if meta !== nothing && haskey(meta, key) + val = meta[key] + return style ? (val, :default) : val + end + return style ? (default, :default) : default +end + +function DataAPI.metadatakeys(t::CDataTable) + meta = getmetadata(t) + meta === nothing && return () + return keys(meta) +end + +function DataAPI.colmetadata(t::CDataTable; style::Bool=false) + pairs = Pair{Symbol,Any}[] + for col in getfield(t, :names) + meta = getmetadata(t[col]) + if meta !== nothing && !isempty(meta) + push!(pairs, col => DataAPI.colmetadata(t, col; style=style)) + end + end + return Dict(pairs) +end + +function DataAPI.colmetadata(t::CDataTable, col; style::Bool=false) + meta = getmetadata(t[col]) + meta === nothing && return Dict{String,Any}() + if style + return Dict(k => (v, :default) for (k, v) in meta) + else + return Dict(meta) + end +end + +function DataAPI.colmetadata(t::CDataTable, col, key::AbstractString; style::Bool=false) + meta = getmetadata(t[col]) + meta === nothing && throw(KeyError(key)) + val = meta[key] + return style ? (val, :default) : val +end + +function DataAPI.colmetadata( + t::CDataTable, + col, + key::AbstractString, + default; + style::Bool=false, +) + meta = getmetadata(t[col]) + if meta !== nothing && haskey(meta, key) + val = meta[key] + return style ? (val, :default) : val + end + return style ? (default, :default) : default +end + +function DataAPI.colmetadatakeys(t::CDataTable, col) + meta = getmetadata(t[col]) + meta === nothing && return () + return keys(meta) +end + +function DataAPI.colmetadatakeys(t::CDataTable) + return ( + col => DataAPI.colmetadatakeys(t, col) for + col in Tables.columnnames(t) if getmetadata(t[col]) !== nothing + ) +end + +function release_c_data(owner::CDataOwner) + array_ptr = Ptr{ArrowArray}(C_NULL) + schema_ptr = Ptr{ArrowSchema}(C_NULL) + array_release = Ptr{Cvoid}(C_NULL) + schema_release = Ptr{Cvoid}(C_NULL) + lock(owner.lock) + try + owner.released && return + owner.released = true + array_ptr = owner.array_ptr + schema_ptr = owner.schema_ptr + if array_ptr != C_NULL + array = unsafe_load(array_ptr) + array_release = array.release + end + if schema_ptr != C_NULL + schema = unsafe_load(schema_ptr) + schema_release = schema.release + end + finally + unlock(owner.lock) + end + # Per the Arrow C Data Interface spec, consumers release only the base structures. + if array_ptr != C_NULL && array_release != C_NULL + ccall(array_release, Cvoid, (Ptr{ArrowArray},), array_ptr) + end + if schema_ptr != C_NULL && schema_release != C_NULL + ccall(schema_release, Cvoid, (Ptr{ArrowSchema},), schema_ptr) + end + return +end + +""" + Arrow.release_c_data(x) + +Release C Data resources owned by an imported array or table. The call is +idempotent. Reads through imported arrays throw after release. +""" +release_c_data(x::CDataVector) = release_c_data(_owner(x)) +release_c_data(t::CDataTable) = release_c_data(getfield(t, :owner)) + +function _to_int(x::Int64, name) + x > typemax(Int) && throw(ArgumentError("$name exceeds the Julia Int range")) + return Int(x) +end + +function _checked_nonnegative(x::Int64, name) + x < 0 && throw(ArgumentError("$name must be nonnegative")) + return _to_int(x, name) +end + +function _checked_add(a::Int, b::Int, name) + b > typemax(Int) - a && throw(ArgumentError("$name overflows")) + return a + b +end + +function _checked_mul(a::Int, b::Int, name) + a != 0 && b > typemax(Int) รท a && throw(ArgumentError("$name overflows")) + return a * b +end + +function _parse_c_data_format(format::AbstractString) + format == "n" && return CDataNullFormat() + format == "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 has fewer buffers than required")) + if expected > 0 && array.buffers == C_NULL + throw(ArgumentError("ArrowArray.buffers is NULL")) + end + buffers = Ptr{Cvoid}[] + for i = 1:expected + push!(buffers, unsafe_load(array.buffers, i)) + end + return buffers +end + +function _load_child_ptrs(schema::ArrowSchema, array::ArrowArray, count::Int) + count == 0 && return Ptr{ArrowSchema}[], Ptr{ArrowArray}[] + schema.children == C_NULL && throw(ArgumentError("ArrowSchema.children is NULL")) + array.children == C_NULL && throw(ArgumentError("ArrowArray.children is NULL")) + schema_children = Ptr{ArrowSchema}[] + array_children = Ptr{ArrowArray}[] + for i = 1:count + schema_child = unsafe_load(schema.children, i) + array_child = unsafe_load(array.children, i) + schema_child == C_NULL && throw(ArgumentError("ArrowSchema child pointer is NULL")) + array_child == C_NULL && throw(ArgumentError("ArrowArray child pointer is NULL")) + push!(schema_children, schema_child) + push!(array_children, array_child) + end + return schema_children, array_children +end + +function _validate_common(schema::ArrowSchema, array::ArrowArray, top_level::Bool) + schema.format == C_NULL && throw(ArgumentError("ArrowSchema.format is NULL")) + if top_level + schema.release == C_NULL && throw(ArgumentError("ArrowSchema.release is NULL")) + array.release == C_NULL && throw(ArgumentError("ArrowArray.release is NULL")) + elseif schema.release == C_NULL || array.release == C_NULL + throw(ArgumentError("released Arrow C Data child structure")) + end + len = _checked_nonnegative(array.length, "ArrowArray.length") + offset = _checked_nonnegative(array.offset, "ArrowArray.offset") + _checked_nonnegative(array.n_buffers, "ArrowArray.n_buffers") + n_children = _checked_nonnegative(array.n_children, "ArrowArray.n_children") + n_children > _CDATA_MAX_CHILDREN && + throw(ArgumentError("ArrowArray.n_children exceeds the import limit")) + schema_children = _checked_nonnegative(schema.n_children, "ArrowSchema.n_children") + schema_children > _CDATA_MAX_CHILDREN && + throw(ArgumentError("ArrowSchema.n_children exceeds the import limit")) + schema_children == n_children || + throw(ArgumentError("ArrowSchema and ArrowArray child counts differ")) + null_count = array.null_count + if !(null_count == -1 || 0 <= null_count <= array.length) + throw(ArgumentError("ArrowArray.null_count is out of range")) + end + if (schema.dictionary == C_NULL) != (array.dictionary == C_NULL) + throw(ArgumentError("ArrowSchema and ArrowArray dictionary pointers differ")) + end + schema.dictionary == C_NULL || + throw(ArgumentError("dictionary encoded Arrow C Data import is not supported")) + return len, offset, Int(null_count), n_children +end + +function _validate_node( + schema_ptr::Ptr{ArrowSchema}, + array_ptr::Ptr{ArrowArray}; + top_level::Bool=false, + depth::Int=1, +) + depth > _CDATA_MAX_DEPTH && + throw(ArgumentError("Arrow C Data nesting exceeds the import limit")) + schema_ptr == C_NULL && throw(ArgumentError("ArrowSchema pointer is NULL")) + array_ptr == C_NULL && throw(ArgumentError("ArrowArray pointer is NULL")) + schema = unsafe_load(schema_ptr) + array = unsafe_load(array_ptr) + len, offset, null_count, n_children = _validate_common(schema, array, top_level) + format = _parse_c_data_format( + _unsafe_string_bounded( + schema.format, + _CDATA_MAX_FORMAT_BYTES, + "ArrowSchema.format", + ), + ) + expected_children = _expected_children(format) + if expected_children !== nothing && n_children != expected_children + throw(ArgumentError("Arrow C Data child count does not match the format")) + end + buffers = _load_buffers(array, _expected_buffers(format)) + schema_child_ptrs, array_child_ptrs = _load_child_ptrs(schema, array, n_children) + children = CDataNode[] + for i in eachindex(schema_child_ptrs) + push!( + children, + _validate_node( + schema_child_ptrs[i], + array_child_ptrs[i]; + top_level=false, + depth=depth + 1, + ), + ) + end + node = CDataNode( + schema_ptr, + array_ptr, + schema, + array, + format, + _load_name(schema.name), + _parse_c_metadata(schema.metadata), + buffers, + children, + len, + offset, + null_count, + ) + _validate_layout(node) + return node +end + +function _validate_layout(node::CDataNode) + total = _checked_add(node.offset, node.len, "ArrowArray offset plus length") + validity_bytes = cld(total, 8) + if _expected_buffers(node.format) > 0 + validity = node.buffers[1] + if node.null_count == -1 + if node.len > 0 && validity == C_NULL + throw(ArgumentError("unknown null count requires a validity bitmap")) + end + elseif node.null_count > 0 && validity == C_NULL + throw(ArgumentError("null values require a validity bitmap")) + end + validity_bytes >= 0 || throw(ArgumentError("invalid validity bitmap size")) + end + _validate_data_layout(node.format, node, total) + _validate_value_layout(node) + return +end + +function _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 _slice_for_table(child::CDataVector, offset::Int, len::Int) + if offset == 0 && length(child) == len + return child + else + first = offset + 1 + return CDataSlice{eltype(child),typeof(child)}(child, first, len) + end +end + +function _mask_for_struct(child::CDataVector, parent_validity::CDataValidity) + if parent_validity.null_count == 0 + return child + else + # Per the Arrow struct validity spec, parent and child validity are independent. + T = Union{eltype(child),Missing} + return CDataMasked{T,typeof(child)}(child, parent_validity) + end +end + +function _table_from_struct(node::CDataNode, owner::CDataOwner, convert::Bool) + parent_validity = _make_validity(node) + columns = AbstractVector[] + names = Symbol[] + types = Type[] + for (i, child_node) in enumerate(node.children) + child = _import_node(child_node, owner, convert) + col = _mask_for_struct( + _slice_for_table(child, node.offset, node.len), + parent_validity, + ) + name = Symbol(child_node.name === nothing ? "f$(i)" : child_node.name) + push!(names, name) + push!(types, eltype(col)) + push!(columns, col) + end + lookup = Dict{Symbol,AbstractVector}(names[i] => columns[i] for i in eachindex(names)) + return CDataTable(names, types, columns, lookup, node.metadata, owner, node.len) +end + +""" + Arrow.from_c_data(schema_ptr, array_ptr; convert=true) + +Import an Arrow C Data Interface schema and array pair. + +Aligned imported buffers are viewed without copying and are released by +`release_c_data` or finalization. Misaligned fixed width data buffers are copied +into aligned Julia storage before typed access. Use `copy` or `collect` on +imported arrays to make Julia owned arrays. + +Per the Arrow C Data Interface spec, producers describe buffer sizes through the +schema, length, and offset fields. The importer validates those layout facts and +does not inspect allocator metadata for foreign pointers. + +A top level struct array is returned as a Tables.jl column table. +""" +function from_c_data( + schema_ptr::Ptr{ArrowSchema}, + array_ptr::Ptr{ArrowArray}; + convert::Bool=true, +) + owner = CDataOwner(schema_ptr, array_ptr) + try + node = _validate_node(schema_ptr, array_ptr; top_level=true) + if node.format isa CDataStructFormat + return _table_from_struct(node, owner, convert) + else + return _import_node(node, owner, convert) + end + catch + release_c_data(owner) + rethrow() + end +end + +from_c_data(schema_ptr::Ptr{Cvoid}, array_ptr::Ptr{Cvoid}; kw...) = + from_c_data(Ptr{ArrowSchema}(schema_ptr), Ptr{ArrowArray}(array_ptr); kw...) diff --git a/test/cdata.jl b/test/cdata.jl new file mode 100644 index 00000000..dd77361e --- /dev/null +++ b/test/cdata.jl @@ -0,0 +1,977 @@ +# 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) + + cc = Sys.which("cc") + if cc === nothing + @test_skip "C compiler not available" + else + c_src = """ + #include + #include + #include + struct ArrowSchema { + const char* format; + const char* name; + const char* metadata; + int64_t flags; + int64_t n_children; + struct ArrowSchema** children; + struct ArrowSchema* dictionary; + void (*release)(struct ArrowSchema*); + void* private_data; + }; + struct ArrowArray { + int64_t length; + int64_t null_count; + int64_t offset; + int64_t n_buffers; + int64_t n_children; + const void** buffers; + struct ArrowArray** children; + struct ArrowArray* dictionary; + void (*release)(struct ArrowArray*); + void* private_data; + }; + int main(void) { + printf("%zu %zu %zu %zu %zu %zu %zu %zu %zu\\n", + offsetof(struct ArrowSchema, format), + offsetof(struct ArrowSchema, name), + offsetof(struct ArrowSchema, metadata), + offsetof(struct ArrowSchema, flags), + offsetof(struct ArrowSchema, n_children), + offsetof(struct ArrowSchema, children), + offsetof(struct ArrowSchema, dictionary), + offsetof(struct ArrowSchema, release), + offsetof(struct ArrowSchema, private_data)); + printf("%zu %zu %zu %zu %zu %zu %zu %zu %zu %zu\\n", + offsetof(struct ArrowArray, length), + offsetof(struct ArrowArray, null_count), + offsetof(struct ArrowArray, offset), + offsetof(struct ArrowArray, n_buffers), + offsetof(struct ArrowArray, n_children), + offsetof(struct ArrowArray, buffers), + offsetof(struct ArrowArray, children), + offsetof(struct ArrowArray, dictionary), + offsetof(struct ArrowArray, release), + offsetof(struct ArrowArray, private_data)); + return 0; + } + """ + mktempdir() do dir + src = joinpath(dir, "layout.c") + bin = joinpath(dir, "layout") + write(src, c_src) + run(`$cc -o $bin $src`) + lines = split(readchomp(`$bin`), '\n') + schema_offsets = parse.(Int, split(lines[1])) + array_offsets = parse.(Int, split(lines[2])) + for (i, offset) in enumerate(schema_offsets) + @test fieldoffset(Arrow.ArrowSchema, i) == offset + end + for (i, offset) in enumerate(array_offsets) + @test fieldoffset(Arrow.ArrowArray, i) == offset + end + end + end + end + + @testset "primitive arrays" begin + data = Int32[1, 2, 3] + f = _primitive_fixture("i", data) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test x isa Arrow.CDataVector + @test collect(x) == Int32[1, 2, 3] + @test copy(x) == Int32[1, 2, 3] + data[1] = 99 + @test collect(x) == Int32[99, 2, 3] + + f = _unaligned_primitive_fixture("i", Int32[7, 8, 9]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test collect(x) == Int32[7, 8, 9] + fill!(f.roots[end], 0x00) + @test collect(x) == Int32[7, 8, 9] + f = _unaligned_primitive_fixture("i", Int32[6, 7, 8, 9]; offset=Int64(1), len=2) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test collect(x) == Int32[7, 8] + fill!(f.roots[end], 0x00) + @test collect(x) == Int32[7, 8] + + validity = UInt8[0b00000101] + f = _primitive_fixture( + "i", + Int32[10, 20, 30]; + validity=validity, + null_count=Int64(1), + flags=Arrow.ARROW_FLAG_NULLABLE, + ) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test isequal(collect(x), Union{Int32,Missing}[10, missing, 30]) + + f = _primitive_fixture( + "i", + Int32[10, 20, 30]; + validity=validity, + null_count=Int64(-1), + flags=Arrow.ARROW_FLAG_NULLABLE, + ) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test Arrow.nullcount(x) == 1 + @test isequal(collect(x), Union{Int32,Missing}[10, missing, 30]) + end + + @testset "null arrays" begin + f = _cdata_fixture("n", 3, Ptr{Cvoid}[]; null_count=Int64(3)) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test x isa Arrow.CDataVector{Missing} + @test size(x) == (3,) + @test Arrow.validitybitmap(x) === nothing + @test Arrow.nullcount(x) == 3 + @test x[2] === missing + @test isequal(collect(x), [missing, missing, missing]) + node = Arrow._validate_node(_schema_ptr(f), _array_ptr(f); top_level=true) + validity = Arrow._make_validity(node) + @test validity.bytes == UInt8[] + @test validity.null_count == 0 + end + + @testset "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 DataAPI.metadata(tbl) == Dict("source" => "cdata") + @test DataAPI.metadata(tbl; style=true) == Dict("source" => ("cdata", :default)) + @test DataAPI.metadata(tbl, "source") == "cdata" + @test DataAPI.metadata(tbl, "source"; style=true) == ("cdata", :default) + @test DataAPI.metadata(tbl, "source", "fallback") == "cdata" + @test DataAPI.metadata(tbl, "source", "fallback"; style=true) == ("cdata", :default) + @test DataAPI.metadata(tbl, "missing", "fallback") == "fallback" + @test DataAPI.metadata(tbl, "missing", "fallback"; style=true) == + ("fallback", :default) + @test Set(DataAPI.metadatakeys(tbl)) == Set(["source"]) + @test DataAPI.colmetadata(tbl, :x) == Dict("unit" => "id") + @test DataAPI.colmetadata(tbl, :x; style=true) == Dict("unit" => ("id", :default)) + @test DataAPI.colmetadata(tbl, :x, "unit") == "id" + @test DataAPI.colmetadata(tbl, :x, "unit"; style=true) == ("id", :default) + @test DataAPI.colmetadata(tbl, :x, "unit", "fallback") == "id" + @test DataAPI.colmetadata(tbl, :x, "unit", "fallback"; style=true) == + ("id", :default) + @test DataAPI.colmetadata(tbl, :x, "missing", "fallback") == "fallback" + @test DataAPI.colmetadata(tbl, :x, "missing", "fallback"; style=true) == + ("fallback", :default) + @test Set(DataAPI.colmetadatakeys(tbl, :x)) == Set(["unit"]) + colkeys = collect(DataAPI.colmetadatakeys(tbl)) + @test length(colkeys) == 1 + @test first(colkeys[1]) == :x + @test Set(last(colkeys[1])) == Set(["unit"]) + @test DataAPI.colmetadata(tbl) == Dict(:x => Dict("unit" => "id")) + + shortmeta_child = + _primitive_fixture("i", Int32[1]; name="x", metadata=Dict("k" => "v")) + shortmeta_root = _cdata_fixture( + "+s", + 1, + Ptr{Cvoid}[C_NULL]; + children=[shortmeta_child], + metadata=Dict("x" => "yz"), + ) + shortmeta_tbl = + Arrow.from_c_data(_schema_ptr(shortmeta_root), _array_ptr(shortmeta_root)) + @test DataAPI.metadata(shortmeta_tbl, "x") == "yz" + @test DataAPI.colmetadata(shortmeta_tbl, :x, "k") == "v" + + col = Tables.getcolumn(tbl, :x) + GC.gc(true) + @test collect(col) == Int32[1, 2, 3] + end + + @testset "struct table offsets and owner roots" begin + child = _primitive_fixture("i", Int32[10, 20, 30]; name="x") + root = + _cdata_fixture("+s", 2, Ptr{Cvoid}[C_NULL]; offset=Int64(1), children=[child]) + tbl = Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + @test collect(Tables.getcolumn(tbl, :x)) == Int32[20, 30] + + col = let + tbl2 = Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + Tables.getcolumn(tbl2, :x) + end + GC.gc(true) + GC.gc(true) + @test collect(col) == Int32[20, 30] + Arrow.release_c_data(col) + + child = _primitive_fixture("i", Int32[10, 20, 30]; offset=Int64(1), len=2, name="x") + root = + _cdata_fixture("+s", 2, Ptr{Cvoid}[C_NULL]; offset=Int64(1), children=[child]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + end + + @testset "struct parent validity masks children" begin + child_validity = UInt8[0b00001011] + parent_validity = UInt8[0b00001101] + child = _primitive_fixture( + "i", + Int32[1, 2, 3, 4]; + validity=child_validity, + null_count=Int64(1), + flags=Arrow.ARROW_FLAG_NULLABLE, + name="x", + ) + root = _cdata_fixture( + "+s", + 4, + Ptr{Cvoid}[Ptr{Cvoid}(pointer(parent_validity))]; + children=[child], + null_count=Int64(1), + flags=Arrow.ARROW_FLAG_NULLABLE, + ) + push!(root.roots, parent_validity) + tbl = Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + col = Tables.getcolumn(tbl, :x) + @test Tables.schema(tbl).types == (Union{Int32,Missing},) + @test Arrow.nullcount(col) == 2 + @test isequal(collect(col), Union{Int32,Missing}[1, missing, missing, 4]) + @test isequal(copy(tbl), (x=Union{Int32,Missing}[1, missing, missing, 4],)) + + parent_validity = UInt8[0b00000010] + child = _primitive_fixture("i", Int32[10, 20, 30]; name="x") + root = _cdata_fixture( + "+s", + 2, + Ptr{Cvoid}[Ptr{Cvoid}(pointer(parent_validity))]; + offset=Int64(1), + children=[child], + null_count=Int64(1), + flags=Arrow.ARROW_FLAG_NULLABLE, + ) + push!(root.roots, parent_validity) + tbl = Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + @test isequal(collect(Tables.getcolumn(tbl, :x)), Union{Int32,Missing}[20, missing]) + end + + @testset "release behavior" begin + f = _primitive_fixture("i", Int32[1, 2, 3]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test x[1] == 1 + Arrow.release_c_data(x) + @test f.array[].release == C_NULL + @test f.schema[].release == C_NULL + @test_nowarn Arrow.release_c_data(x) + @test_throws ArgumentError x[1] + + f = _primitive_fixture("i", Int32[1, 2, 3]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + owner = Arrow._owner(x) + owner_lock = getfield(owner, :lock) + started = Channel{Nothing}(1) + locked = false + lock(owner_lock) + locked = true + try + task = Threads.@spawn begin + put!(started, nothing) + Arrow.release_c_data(x) + end + take!(started) + sleep(0.05) + @test f.array[].release != C_NULL + unlock(owner_lock) + locked = false + wait(task) + finally + locked && unlock(owner_lock) + end + @test f.array[].release == C_NULL + @test f.schema[].release == C_NULL + + f = _primitive_fixture("i", Int32[1, 2, 3]) + _set_array!(f; release=_CDATA_RELEASE_ARRAY_REENTRANT) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + _CDATA_REENTRANT_OBJECT[] = x + _CDATA_REENTRANT_RELEASES[] = 0 + try + @test_nowarn Arrow.release_c_data(x) + @test _CDATA_REENTRANT_RELEASES[] == 1 + @test f.array[].release == C_NULL + @test f.schema[].release == C_NULL + @test_throws ArgumentError x[1] + finally + _CDATA_REENTRANT_OBJECT[] = nothing + end + end + + @testset "table release behavior" begin + child = _primitive_fixture("i", Int32[1, 2, 3]; name="x") + root = _cdata_fixture("+s", 3, Ptr{Cvoid}[C_NULL]; children=[child]) + tbl = Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + col = Tables.getcolumn(tbl, :x) + @test col[1] == 1 + Arrow.release_c_data(tbl) + @test_nowarn Arrow.release_c_data(tbl) + @test_throws ArgumentError Tables.getcolumn(tbl, :x) + @test_throws ArgumentError col[1] + end + + @testset "copy and collect own the result" begin + f = _primitive_fixture("i", Int32[1, 2, 3]) + x = Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + a = collect(x) + b = copy(x) + Arrow.release_c_data(x) + @test a == Int32[1, 2, 3] + @test b == Int32[1, 2, 3] + end + + @testset "malformed inputs" begin + f = _primitive_fixture("i", Int32[1]) + @test_throws ArgumentError Arrow.from_c_data( + Ptr{Arrow.ArrowSchema}(C_NULL), + _array_ptr(f), + ) + @test_throws ArgumentError Arrow.from_c_data( + _schema_ptr(f), + Ptr{Arrow.ArrowArray}(C_NULL), + ) + + f = _cdata_fixture("?", 0, Ptr{Cvoid}[]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + @test f.array[].release == C_NULL + @test f.schema[].release == C_NULL + + f = _primitive_fixture("i", Int32[1]) + f.schema[] = Arrow.ArrowSchema( + f.schema[].format, + f.schema[].name, + f.schema[].metadata, + f.schema[].flags, + f.schema[].n_children, + f.schema[].children, + f.schema[].dictionary, + C_NULL, + f.schema[].private_data, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + f.array[] = Arrow.ArrowArray( + f.array[].length, + f.array[].null_count, + f.array[].offset, + f.array[].n_buffers, + f.array[].n_children, + f.array[].buffers, + f.array[].children, + f.array[].dictionary, + C_NULL, + f.array[].private_data, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + f.array[] = Arrow.ArrowArray( + -1, + 0, + 0, + 2, + 0, + f.array[].buffers, + f.array[].children, + f.array[].dictionary, + f.array[].release, + f.array[].private_data, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + f.array[] = Arrow.ArrowArray( + f.array[].length, + f.array[].null_count, + -1, + f.array[].n_buffers, + f.array[].n_children, + f.array[].buffers, + f.array[].children, + f.array[].dictionary, + f.array[].release, + f.array[].private_data, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + f.array[] = Arrow.ArrowArray( + f.array[].length, + f.array[].null_count, + f.array[].offset, + -1, + f.array[].n_children, + f.array[].buffers, + f.array[].children, + f.array[].dictionary, + f.array[].release, + f.array[].private_data, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + f.array[] = Arrow.ArrowArray( + f.array[].length, + f.array[].null_count, + f.array[].offset, + f.array[].n_buffers, + -1, + f.array[].buffers, + f.array[].children, + f.array[].dictionary, + f.array[].release, + f.array[].private_data, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + f.schema[] = Arrow.ArrowSchema( + f.schema[].format, + f.schema[].name, + f.schema[].metadata, + f.schema[].flags, + -1, + f.schema[].children, + f.schema[].dictionary, + f.schema[].release, + f.schema[].private_data, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]; null_count=Int64(2)) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]; null_count=Int64(-1)) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture( + "i", + Int32[1]; + null_count=Int64(1), + flags=Arrow.ARROW_FLAG_NULLABLE, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture( + "i", + Int32[1, 2, 3]; + validity=UInt8[0b00000111], + null_count=Int64(1), + flags=Arrow.ARROW_FLAG_NULLABLE, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture( + "i", + Int32[1, 2, 3]; + validity=UInt8[0b00000101], + null_count=Int64(0), + flags=Arrow.ARROW_FLAG_NULLABLE, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + _set_array!(f; buffers=Ptr{Ptr{Cvoid}}(C_NULL)) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _cdata_fixture("i", 1, Ptr{Cvoid}[C_NULL, C_NULL]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _cdata_fixture("i", 1, Ptr{Cvoid}[C_NULL]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + badmeta = reinterpret(UInt8, Int32[1, -1]) + f = _primitive_fixture("i", Int32[1]) + f.schema[] = Arrow.ArrowSchema( + f.schema[].format, + f.schema[].name, + Cstring(pointer(badmeta)), + f.schema[].flags, + f.schema[].n_children, + f.schema[].children, + f.schema[].dictionary, + f.schema[].release, + f.schema[].private_data, + ) + push!(f.roots, badmeta) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + badmeta = reinterpret(UInt8, Int32[1, 0, -1]) + f = _primitive_fixture("i", Int32[1]) + _set_schema!(f; metadata=Cstring(pointer(badmeta))) + push!(f.roots, badmeta) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + f.schema[] = Arrow.ArrowSchema( + f.schema[].format, + f.schema[].name, + f.schema[].metadata, + f.schema[].flags, + f.schema[].n_children, + Ptr{Arrow.ArrowSchema}(UInt(0x01)), + f.schema[].dictionary, + f.schema[].release, + f.schema[].private_data, + ) + f.array[] = Arrow.ArrowArray( + f.array[].length, + f.array[].null_count, + f.array[].offset, + f.array[].n_buffers, + 1, + f.array[].buffers, + Ptr{Ptr{Arrow.ArrowArray}}(C_NULL), + f.array[].dictionary, + f.array[].release, + f.array[].private_data, + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + child = _primitive_fixture("i", Int32[1]) + schema_children = Ptr{Arrow.ArrowSchema}[Ptr{Arrow.ArrowSchema}(C_NULL)] + array_children = Ptr{Arrow.ArrowArray}[_array_ptr(child)] + f = _cdata_fixture("+s", 1, Ptr{Cvoid}[C_NULL]) + _set_schema!( + f; + n_children=1, + children=Ptr{Ptr{Arrow.ArrowSchema}}(pointer(schema_children)), + ) + _set_array!( + f; + n_children=1, + children=Ptr{Ptr{Arrow.ArrowArray}}(pointer(array_children)), + ) + append!(f.roots, Any[child, schema_children, array_children]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + child = _primitive_fixture("i", Int32[1]) + schema_children = Ptr{Arrow.ArrowSchema}[_schema_ptr(child)] + array_children = Ptr{Arrow.ArrowArray}[Ptr{Arrow.ArrowArray}(C_NULL)] + f = _cdata_fixture("+s", 1, Ptr{Cvoid}[C_NULL]) + f.schema[] = Arrow.ArrowSchema( + f.schema[].format, + f.schema[].name, + f.schema[].metadata, + f.schema[].flags, + 1, + Ptr{Ptr{Arrow.ArrowSchema}}(pointer(schema_children)), + f.schema[].dictionary, + f.schema[].release, + f.schema[].private_data, + ) + f.array[] = Arrow.ArrowArray( + f.array[].length, + f.array[].null_count, + f.array[].offset, + f.array[].n_buffers, + 1, + f.array[].buffers, + Ptr{Ptr{Arrow.ArrowArray}}(pointer(array_children)), + f.array[].dictionary, + f.array[].release, + f.array[].private_data, + ) + append!(f.roots, Any[child, schema_children, array_children]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + child = _primitive_fixture("i", Int32[1]) + _set_schema!(child; release=C_NULL) + schema_children = Ptr{Arrow.ArrowSchema}[_schema_ptr(child)] + array_children = Ptr{Arrow.ArrowArray}[_array_ptr(child)] + root = _cdata_fixture("+s", 1, Ptr{Cvoid}[C_NULL]) + _set_schema!( + root; + n_children=1, + children=Ptr{Ptr{Arrow.ArrowSchema}}(pointer(schema_children)), + ) + _set_array!( + root; + n_children=1, + children=Ptr{Ptr{Arrow.ArrowArray}}(pointer(array_children)), + ) + append!(root.roots, Any[child, schema_children, array_children]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(root), _array_ptr(root)) + + child = _primitive_fixture("i", Int32[1]) + f = _primitive_fixture("i", Int32[1]) + schema_children = Ptr{Arrow.ArrowSchema}[_schema_ptr(child)] + array_children = Ptr{Arrow.ArrowArray}[_array_ptr(child)] + _set_schema!( + f; + n_children=1, + children=Ptr{Ptr{Arrow.ArrowSchema}}(pointer(schema_children)), + ) + _set_array!( + f; + n_children=1, + children=Ptr{Ptr{Arrow.ArrowArray}}(pointer(array_children)), + ) + append!(f.roots, Any[child, schema_children, array_children]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + dict_schema = Ref(f.schema[]) + f.schema[] = Arrow.ArrowSchema( + f.schema[].format, + f.schema[].name, + f.schema[].metadata, + f.schema[].flags, + f.schema[].n_children, + f.schema[].children, + Base.unsafe_convert(Ptr{Arrow.ArrowSchema}, dict_schema), + f.schema[].release, + f.schema[].private_data, + ) + push!(f.roots, dict_schema) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + dict_array = Ref(f.array[]) + f.array[] = Arrow.ArrowArray( + f.array[].length, + f.array[].null_count, + f.array[].offset, + f.array[].n_buffers, + f.array[].n_children, + f.array[].buffers, + f.array[].children, + Base.unsafe_convert(Ptr{Arrow.ArrowArray}, dict_array), + f.array[].release, + f.array[].private_data, + ) + push!(f.roots, dict_array) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + dict_schema = Ref(f.schema[]) + dict_array = Ref(f.array[]) + _set_schema!(f; dictionary=Base.unsafe_convert(Ptr{Arrow.ArrowSchema}, dict_schema)) + _set_array!(f; dictionary=Base.unsafe_convert(Ptr{Arrow.ArrowArray}, dict_array)) + append!(f.roots, Any[dict_schema, dict_array]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + f = _primitive_fixture("i", Int32[1]) + fmt = fill(UInt8('i'), Arrow._CDATA_MAX_FORMAT_BYTES + 1) + f.schema[] = Arrow.ArrowSchema( + Cstring(pointer(fmt)), + f.schema[].name, + f.schema[].metadata, + f.schema[].flags, + f.schema[].n_children, + f.schema[].children, + f.schema[].dictionary, + f.schema[].release, + f.schema[].private_data, + ) + push!(f.roots, fmt) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + child = _primitive_fixture("i", Int32[]) + for _ = 1:Arrow._CDATA_MAX_DEPTH + child = _cdata_fixture("+s", 0, Ptr{Cvoid}[C_NULL]; children=[child]) + end + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(child), _array_ptr(child)) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 315d1b60..c591dcf7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -39,6 +39,7 @@ include(joinpath(@__DIR__, "testtables.jl")) include(joinpath(@__DIR__, "testappend.jl")) include(joinpath(@__DIR__, "integrationtest.jl")) include(joinpath(@__DIR__, "dates.jl")) +include(joinpath(@__DIR__, "cdata.jl")) struct CustomStruct x::Int @@ -50,6 +51,25 @@ struct CustomStruct2{sym} x::Int end +struct OffsetPool{T} <: AbstractVector{T} + data::Vector{T} + first::Int +end +Base.size(x::OffsetPool) = size(x.data) +Base.axes(x::OffsetPool) = (x.first:(x.first + length(x.data) - 1),) +Base.IndexStyle(::Type{<:OffsetPool}) = IndexLinear() +Base.getindex(x::OffsetPool, i::Int) = x.data[i - x.first + 1] + +struct OffsetRefVector{T,P,R} <: AbstractVector{T} + pool::P + refs::R +end +Base.size(x::OffsetRefVector) = size(x.refs) +Base.IndexStyle(::Type{<:OffsetRefVector}) = IndexLinear() +Base.getindex(x::OffsetRefVector, i::Int) = x.pool[x.refs[i]] +DataAPI.refpool(x::OffsetRefVector) = x.pool +DataAPI.refarray(x::OffsetRefVector) = x.refs + @testset ExtendedTestSet "Arrow" begin @testset "table roundtrips" begin for case in testtables @@ -353,6 +373,16 @@ end @test isa(first(av.indices), Signed) @test length(av) == 3 @test eltype(av) == String + + pool = OffsetPool(["a", "b"], 1000) + x = OffsetRefVector{String,typeof(pool),Vector{Int}}(pool, [1000, 1001, 1000]) + av = Arrow.toarrowvector(x) + @test av.indices == Int8[0, 1, 0] + @test collect(av) == ["a", "b", "a"] + @test Arrow.Table(Arrow.tobuffer((a=x,))).a == ["a", "b", "a"] + + bad = OffsetRefVector{String,typeof(pool),Vector{Int}}(pool, [1000, 1002]) + @test_throws ArgumentError Arrow.toarrowvector(bad) end @testset "# 120" begin From e306a2bb50654a80251ee9e615717731b1459872 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/2] 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 Generated-by: OpenAI Codex --- src/cdata.jl | 539 +++++++++++++++++++++++++++++++++++++++++++++++++- test/cdata.jl | 366 +++++++++++++++++++++++++++++++++- 2 files changed, 899 insertions(+), 6 deletions(-) diff --git a/src/cdata.jl b/src/cdata.jl index 451decb1..66d36b68 100644 --- a/src/cdata.jl +++ b/src/cdata.jl @@ -55,13 +55,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 @@ -99,6 +111,61 @@ struct CDataPrimitive{T,S,A<:AbstractVector{S}} <: CDataVector{T} metadata::Union{Nothing,Base.ImmutableDict{String,String}} end +struct CDataBool{T} <: CDataVector{T} + owner::CDataOwner + validity::CDataValidity + data::Vector{UInt8} + bitoffset::Int + len::Int + metadata::Union{Nothing,Base.ImmutableDict{String,String}} +end + +struct CDataBinary{T,O,A<:AbstractVector{O}} <: CDataVector{T} + owner::CDataOwner + validity::CDataValidity + offsets::A + data::Vector{UInt8} + len::Int + metadata::Union{Nothing,Base.ImmutableDict{String,String}} +end + +struct CDataFixedSizeBinary{T} <: CDataVector{T} + owner::CDataOwner + validity::CDataValidity + data::Vector{UInt8} + bytewidth::Int + len::Int + metadata::Union{Nothing,Base.ImmutableDict{String,String}} +end + +struct CDataList{T,O,A<:AbstractVector{O},C<:AbstractVector} <: CDataVector{T} + owner::CDataOwner + validity::CDataValidity + offsets::A + data::C + len::Int + metadata::Union{Nothing,Base.ImmutableDict{String,String}} +end + +struct CDataFixedSizeList{T,C<:AbstractVector} <: CDataVector{T} + owner::CDataOwner + validity::CDataValidity + data::C + listsize::Int + offset::Int + len::Int + metadata::Union{Nothing,Base.ImmutableDict{String,String}} +end + +struct CDataStruct{T,S,fnames} <: CDataVector{T} + owner::CDataOwner + validity::CDataValidity + data::S + offset::Int + len::Int + metadata::Union{Nothing,Base.ImmutableDict{String,String}} +end + struct CDataSlice{T,V<:AbstractVector{T}} <: CDataVector{T} parent::V first::Int @@ -139,6 +206,12 @@ Base.IndexStyle(::Type{<:CDataVector}) = Base.IndexLinear() Base.size(x::CDataNull) = (x.len,) Base.size(x::CDataPrimitive) = size(x.data) +Base.size(x::CDataBool) = (x.len,) +Base.size(x::CDataBinary) = (x.len,) +Base.size(x::CDataFixedSizeBinary) = (x.len,) +Base.size(x::CDataList) = (x.len,) +Base.size(x::CDataFixedSizeList) = (x.len,) +Base.size(x::CDataStruct) = (x.len,) Base.size(x::CDataSlice) = (x.len,) Base.size(x::CDataMasked) = size(x.parent) Base.length(t::CDataTable) = length(getfield(t, :columns)) @@ -219,6 +292,94 @@ 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,fnames}, + i::Integer, +) where {T,S,fnames} + return _with_live(x) do + @boundscheck checkbounds(x, i) + if !_valid(x.validity, i) + return missing + end + j = x.offset + Int(i) + vals = ntuple(k -> @inbounds(x.data[k][j]), fieldcount(S)) + NT = Base.nonmissingtype(T) + if isnamedtuple(NT) || istuple(NT) + return ArrowTypes.fromarrow(T, NT(vals)) + else + return ArrowTypes.fromarrow(T, _fromarrowstruct(NT, Val{fnames}(), vals...)) + end + end +end + @propagate_inbounds function Base.getindex(x::CDataSlice, i::Integer) return _with_live(x) do @boundscheck checkbounds(x, i) @@ -429,6 +590,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) @@ -440,16 +602,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) @@ -676,12 +934,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) @@ -692,6 +958,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 @@ -720,6 +1042,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) @@ -755,6 +1151,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 @@ -781,6 +1181,141 @@ function _import_node( ) end +function _import_node(::CDataBoolFormat, node::CDataNode, owner::CDataOwner, convert::Bool) + validity = _make_validity(node) + nullable = _nullable(node.schema, validity.null_count) + T = nullable ? Union{Bool,Missing} : Bool + nbytes = cld(_checked_add(node.offset, node.len, "boolean bitmap length"), 8) + data = + nbytes == 0 ? UInt8[] : + unsafe_wrap(Array, Ptr{UInt8}(node.buffers[2]), nbytes; own=false) + return CDataBool{T}(owner, validity, data, node.offset, node.len, node.metadata) +end + +function _import_node( + format::CDataBinaryFormat{O}, + node::CDataNode, + owner::CDataOwner, + convert::Bool, +) where {O} + validity = _make_validity(node) + nullable = _nullable(node.schema, validity.null_count) + T = nullable ? Union{format.juliatype,Missing} : format.juliatype + offsets = _wrap_offsets(node.buffers[2], O, node.offset, node.len + 1) + data_len = offsets[end] == offsets[1] ? 0 : Int(offsets[end]) + data = + data_len == 0 ? UInt8[] : + unsafe_wrap(Array, Ptr{UInt8}(node.buffers[3]), data_len; own=false) + return CDataBinary{T,O,typeof(offsets)}( + owner, + validity, + offsets, + data, + node.len, + node.metadata, + ) +end + +function _import_node( + format::CDataFixedSizeBinaryFormat, + node::CDataNode, + owner::CDataOwner, + convert::Bool, +) + validity = _make_validity(node) + nullable = _nullable(node.schema, validity.null_count) + storage = NTuple{format.bytewidth,UInt8} + T = nullable ? Union{storage,Missing} : storage + nbytes = _checked_mul(node.len, format.bytewidth, "fixed size binary byte count") + data = + nbytes == 0 ? UInt8[] : + unsafe_wrap( + Array, + Ptr{UInt8}(node.buffers[2]) + + _checked_mul(node.offset, format.bytewidth, "fixed size binary byte offset"), + nbytes; + own=false, + ) + return CDataFixedSizeBinary{T}( + owner, + validity, + data, + format.bytewidth, + node.len, + node.metadata, + ) +end + +function _import_node( + format::CDataListFormat{O}, + node::CDataNode, + owner::CDataOwner, + convert::Bool, +) where {O} + validity = _make_validity(node) + child = _import_node(node.children[1], owner, convert) + nullable = _nullable(node.schema, validity.null_count) + storage = Vector{eltype(child)} + T = nullable ? Union{storage,Missing} : storage + offsets = _wrap_offsets(node.buffers[2], O, node.offset, node.len + 1) + return CDataList{T,O,typeof(offsets),typeof(child)}( + owner, + validity, + offsets, + child, + node.len, + node.metadata, + ) +end + +function _import_node( + format::CDataFixedSizeListFormat, + node::CDataNode, + owner::CDataOwner, + convert::Bool, +) + validity = _make_validity(node) + child = _import_node(node.children[1], owner, convert) + nullable = _nullable(node.schema, validity.null_count) + storage = NTuple{format.listsize,eltype(child)} + T = nullable ? Union{storage,Missing} : storage + return CDataFixedSizeList{T,typeof(child)}( + owner, + validity, + child, + format.listsize, + node.offset, + node.len, + node.metadata, + ) +end + +function _import_node( + ::CDataStructFormat, + node::CDataNode, + owner::CDataOwner, + convert::Bool, +) + validity = _make_validity(node) + children = Tuple(_import_node(child, owner, convert) for child in node.children) + names = Tuple( + Symbol(child.name === nothing ? "f$(i)" : child.name) for + (i, child) in enumerate(node.children) + ) + types = Tuple(eltype(child) for child in children) + storage = NamedTuple{names,Tuple{types...}} + nullable = _nullable(node.schema, validity.null_count) + T = nullable ? Union{storage,Missing} : storage + return CDataStruct{T,typeof(children),names}( + owner, + validity, + children, + node.offset, + node.len, + node.metadata, + ) +end + function _slice_for_table(child::CDataVector, offset::Int, len::Int) if offset == 0 && length(child) == len return child diff --git a/test/cdata.jl b/test/cdata.jl index dd77361e..1d49c102 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 = @@ -407,10 +409,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, @@ -420,7 +655,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)) @@ -432,8 +667,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 DataAPI.metadata(tbl) == Dict("source" => "cdata") @@ -463,6 +698,13 @@ 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( @@ -547,6 +789,46 @@ 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 + 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)) @@ -781,6 +1063,57 @@ end f = _cdata_fixture("i", 1, Ptr{Cvoid}[C_NULL]) @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + child = _primitive_fixture("i", Int32[1]) + f = _cdata_fixture("+l", 1, Ptr{Cvoid}[C_NULL, C_NULL]; children=[child]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + child = _primitive_fixture("i", Int32[1, 2]) + offsets = Int32[0, 3] + f = _cdata_fixture( + "+l", + 1, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets))]; + children=[child], + ) + push!(f.roots, offsets) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + child = _primitive_fixture("i", Int32[1, 2, 3, 4, 5]) + f = _cdata_fixture("+w:3", 2, Ptr{Cvoid}[C_NULL]; children=[child]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + offsets = Int32[0, 3, 2] + bytes = UInt8[1, 2, 3] + f = _cdata_fixture( + "z", + 2, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets)), Ptr{Cvoid}(pointer(bytes))], + ) + append!(f.roots, Any[offsets, bytes]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + offsets = Int32[0, -1] + bytes = UInt8[] + f = _cdata_fixture( + "z", + 1, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets)), Ptr{Cvoid}(C_NULL)], + ) + append!(f.roots, Any[offsets, bytes]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + + offsets = Int32[0, 1] + bytes = UInt8[0x01] + f = _cdata_fixture( + "u", + 1, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets)), Ptr{Cvoid}(pointer(bytes))]; + flags=Arrow.ARROW_FLAG_NULLABLE, + null_count=Int64(-1), + ) + append!(f.roots, Any[offsets, bytes]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + badmeta = reinterpret(UInt8, Int32[1, -1]) f = _primitive_fixture("i", Int32[1]) f.schema[] = Arrow.ArrowSchema( @@ -968,10 +1301,35 @@ end push!(f.roots, fmt) @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + for fmt in ("tsq:", "d:x,2", "d:10,2,64", "w:0", "+w:x") + f = _cdata_fixture(fmt, 0, Ptr{Cvoid}[C_NULL, C_NULL]) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + end + + f = _cdata_fixture( + "w:$(Arrow._CDATA_MAX_FIXED_SIZE + 1)", + 0, + Ptr{Cvoid}[C_NULL, C_NULL], + ) + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(f), _array_ptr(f)) + child = _primitive_fixture("i", Int32[]) for _ = 1:Arrow._CDATA_MAX_DEPTH child = _cdata_fixture("+s", 0, Ptr{Cvoid}[C_NULL]; children=[child]) end @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(child), _array_ptr(child)) + + child = _primitive_fixture("i", Int32[]) + for _ = 1:Arrow._CDATA_MAX_DEPTH + offsets = Int32[0] + child = _cdata_fixture( + "+l", + 0, + Ptr{Cvoid}[C_NULL, Ptr{Cvoid}(pointer(offsets))]; + children=[child], + ) + push!(child.roots, offsets) + end + @test_throws ArgumentError Arrow.from_c_data(_schema_ptr(child), _array_ptr(child)) end end