samtalki commented on code in PR #607: URL: https://github.com/apache/arrow-julia/pull/607#discussion_r3962711823
########## src/cdata.jl: ########## @@ -0,0 +1,647 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +struct ArrowSchema + format::Cstring + name::Cstring + metadata::Cstring + flags::Int64 + n_children::Int64 + children::Ptr{Ptr{ArrowSchema}} + dictionary::Ptr{ArrowSchema} + release::Ptr{Cvoid} + private_data::Ptr{Cvoid} +end + +struct ArrowArray + length::Int64 + null_count::Int64 + offset::Int64 + n_buffers::Int64 + n_children::Int64 + buffers::Ptr{Ptr{Cvoid}} + children::Ptr{Ptr{ArrowArray}} + dictionary::Ptr{ArrowArray} + release::Ptr{Cvoid} + private_data::Ptr{Cvoid} +end + +const _CDATA_PTR_SIZE = sizeof(Ptr{Cvoid}) +@assert isbitstype(ArrowSchema) +@assert isbitstype(ArrowArray) +# On 32 bit ABIs with 8 byte Int64 alignment the C structs contain padding, so +# the packed size formulas hold only where pointer and Int64 sizes agree. +@static if Sys.WORD_SIZE == 64 + @assert sizeof(ArrowSchema) == 7 * _CDATA_PTR_SIZE + 2 * sizeof(Int64) + @assert sizeof(ArrowArray) == 5 * _CDATA_PTR_SIZE + 5 * sizeof(Int64) +end + +const ARROW_FLAG_DICTIONARY_ORDERED = Int64(1) +const ARROW_FLAG_NULLABLE = Int64(2) +const ARROW_FLAG_MAP_KEYS_SORTED = Int64(4) + +const _CDATA_MAX_FORMAT_BYTES = 4096 + +abstract type CDataFormat end + +struct CDataNullFormat <: CDataFormat end +struct CDataPrimitiveFormat <: CDataFormat + storage::Type +end + +abstract type CDataVector{T} <: ArrowVector{T} end + +mutable struct CDataOwner + schema::Base.RefValue{ArrowSchema} + array::Base.RefValue{ArrowArray} + released::Bool + lock::ReentrantLock +end + +function CDataOwner(schema_ptr::Ptr{ArrowSchema}, array_ptr::Ptr{ArrowArray}) + # Per the Arrow C Data Interface spec, move the base structures into Julia + # owned storage and mark the sources released without calling their release + # callbacks, like arrow-rs `from_raw` and nanoarrow `ArrowArrayMove`. + owner = CDataOwner( + Ref(unsafe_load(schema_ptr)), + Ref(unsafe_load(array_ptr)), + false, + ReentrantLock(), + ) + _clear_schema_release!(schema_ptr) + _clear_array_release!(array_ptr) + finalizer(_finalize_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 +end + +struct CDataPrimitive{T,S,A<:AbstractVector{S}} <: CDataVector{T} + owner::CDataOwner + validity::CDataValidity + data::A +end + +struct CDataNode + schema::ArrowSchema + array::ArrowArray + format::CDataFormat + buffers::Vector{Ptr{Cvoid}} + 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) + +_owner(x::CDataVector) = getfield(x, :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)) + +validitybitmap(x::CDataNull) = nothing +nullcount(x::CDataNull) = x.len +nullcount(x::CDataVector) = validitybitmap(x).null_count + +@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(bytes::Vector{UInt8}, bitoffset::Int, len::Int) + len == 0 && return 0 + firstbit = bitoffset + lastbit = bitoffset + len - 1 + firstbyte = firstbit >>> 3 + lastbyte = lastbit >>> 3 + firstmask = 0xff << (firstbit & 7) + lastmask = 0xff >>> (7 - (lastbit & 7)) + set = 0 + if firstbyte == lastbyte + set = count_ones(@inbounds(bytes[firstbyte + 1]) & firstmask & lastmask) + else + set = count_ones(@inbounds(bytes[firstbyte + 1]) & firstmask) + @inbounds for b = (firstbyte + 2):lastbyte + set += count_ones(bytes[b]) + end + set += count_ones(@inbounds(bytes[lastbyte + 1]) & lastmask) + end + return len - set +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 + +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 Review Comment: Yes, exactly. Users can index and iterate the imported array. These are accessor functions that protect buffer reads against concurrent release. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
