Optimize Arrow write paths and tighten Flight interop

Author: guangtao

Project.toml

@@ -46,6 +46,9 @@ gRPCServer = "608c6337-0d7d-447f-bb69-0f5674ee3959"
 [extensions]
 ArrowgRPCServerExt = "gRPCServer"
 
+[sources]
+ArrowTypes = { path = "src/ArrowTypes" }
+
 [compat]
 ArrowTypes = "1.1,2"
 BitIntegers = "0.2, 0.3"

README.md

@@ -40,6 +40,8 @@ The package can be installed by typing in the following in a Julia REPL:
 julia> using Pkg; Pkg.add("Arrow")
 ```
 
+Arrow.jl currently requires Julia `1.12+`.
+
 ## Local Development
 
 When developing on Arrow.jl it is recommended that you run the following to ensure that any

ext/arrowgrpcserverext/handlers.jl

@@ -35,8 +35,10 @@ function _server_streaming_handler(service::Flight.Service, method::Flight.Metho
         end
         try
             _drain_response!(stream, response)
-        finally
             _streaming_handler_result(task)
+            gRPCServer.close!(stream)
+        finally
+            istaskdone(task) || wait(task)
         end
     end
 end
@@ -86,9 +88,11 @@ function _bidi_streaming_handler(service::Flight.Service, method::Flight.MethodD
             for message in response
                 gRPCServer.send!(stream, message)
             end
+            _streaming_handler_result(task, producer)
             gRPCServer.close!(stream)
         finally
-            _streaming_handler_result(task, producer)
+            istaskdone(task) || wait(task)
+            isnothing(producer) || (istaskdone(producer) || wait(producer))
         end
         return nothing
     end

ext/arrowgrpcserverext/streams.jl

@@ -16,20 +16,27 @@
 # under the License.
 
 function _drain_response!(
-    stream::gRPCServer.ServerStream{T},
-    response::Channel{T},
-) where {T}
+    stream::gRPCServer.ServerStream,
+    response::Channel,
+)
+    # gRPCServer falls back to `ServerStream{Any}` when a descriptor only carries
+    # protobuf type names. Drain generically and let `send!` enforce compatibility.
     for message in response
         gRPCServer.send!(stream, message)
     end
-    gRPCServer.close!(stream)
     return nothing
 end
 
 function _streaming_handler_result(task::Task, producer::Union{Nothing,Task}=nothing)
     if !isnothing(producer)
+        if istaskfailed(producer)
+            throw(producer.exception)
+        end
         wait(producer)
     end
+    if istaskfailed(task)
+        throw(task.exception)
+    end
     wait(task)
     return nothing
 end

src/ArrowTypes/Project.toml

@@ -25,4 +25,4 @@ Sockets = "6462fe0b-24de-5631-8697-dd941f90decc"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [compat]
-julia = "1.0"
+julia = "1.12"

src/ArrowTypes/src/ArrowTypes.jl

@@ -405,26 +405,58 @@ function _toarroweltype(x)
     state = iterate(x)
     state === nothing && return Missing
     y, st = state
-    T = typeof(toarrow(y))
+    srcT = Union{}
+    stable = false
+    T = Missing
+    if y !== missing
+        srcT = typeof(y)
+        mapped = ArrowType(srcT)
+        stable = isconcretetype(mapped)
+        T = stable ? mapped : typeof(toarrow(y))
+    end
     while true
         state = iterate(x, st)
         state === nothing && return T
         y, st = state
-        S = typeof(toarrow(y))
+        if y === missing
+            S = Missing
+        elseif srcT === Union{}
+            srcT = typeof(y)
+            mapped = ArrowType(srcT)
+            stable = isconcretetype(mapped)
+            S = stable ? mapped : typeof(toarrow(y))
+        elseif stable && typeof(y) === srcT
+            continue
+        else
+            S = typeof(toarrow(y))
+            if stable && typeof(y) !== srcT
+                stable = false
+            end
+        end
         S === T && continue
         T = promoteunion(T, S)
     end
 end
 
+@inline _hasoffsetaxes(data) = Base.has_offset_axes(data)
+@inline _offsetshift(data) = _hasoffsetaxes(data) ? firstindex(data) - 1 : 0
+@inline _hasonebasedaxes(data) = !_hasoffsetaxes(data)
+
 # lazily call toarrow(x) on getindex for each x in data
 struct ToArrow{T,A} <: AbstractVector{T}
     data::A
     offset::Int
     needsconvert::Bool
 end
+@inline _sourcedata(x::ToArrow) = getfield(x, :data)
+@inline _sourceoffset(x::ToArrow) = getfield(x, :offset)
+@inline _needsconvert(x::ToArrow) = getfield(x, :needsconvert)
+@inline _sourcevalue(x::ToArrow, i::Integer) =
+    @inbounds getindex(_sourcedata(x), i + _sourceoffset(x))
+
 function ToArrow{T,A}(data::A) where {T,A}
     needsconvert = !(eltype(A) === T && concrete_or_concreteunion(T))
-    return ToArrow{T,A}(data, firstindex(data) - 1, needsconvert)
+    return ToArrow{T,A}(data, _offsetshift(data), needsconvert)
 end
 
 concrete_or_concreteunion(T) =
@@ -434,8 +466,7 @@ concrete_or_concreteunion(T) =
 function ToArrow(x::A) where {A}
     S = eltype(A)
     T = ArrowType(S)
-    fi = firstindex(x)
-    if S === T && concrete_or_concreteunion(S) && fi == 1
+    if S === T && concrete_or_concreteunion(S) && _hasonebasedaxes(x)
         return x
     elseif !concrete_or_concreteunion(T)
         # arrow needs concrete types, so try to find a concrete common type, preferring unions
@@ -472,12 +503,12 @@ function _convert(::Type{T}, x) where {T}
 end
 
 @inline function _toarrowvalue(x::ToArrow{T}, value) where {T}
-    x.needsconvert || return value
+    _needsconvert(x) || return value
     return _convert(T, toarrow(value))
 end
 
 Base.@propagate_inbounds function Base.getindex(x::ToArrow{T}, i::Int) where {T}
-    value = @inbounds getindex(x.data, i + x.offset)
+    value = _sourcevalue(x, i)
     return _toarrowvalue(x, value)
 end
 

src/ArrowTypes/test/tests.jl

@@ -185,14 +185,21 @@ end
     @test !ArrowTypes.concrete_or_concreteunion(Any)
 
     @testset "ToArrow" begin
+        @test !ArrowTypes._hasoffsetaxes([1, 2, 3])
+        @test ArrowTypes._offsetshift([1, 2, 3]) == 0
+
         x = ArrowTypes.ToArrow([1, 2, 3])
         @test x isa Vector{Int}
         @test x == [1, 2, 3]
 
+        baseview = @view [1, 2, 3][1:3]
+        x = ArrowTypes.ToArrow(baseview)
+        @test x === baseview
+
         x = ArrowTypes.ToArrow([:hey, :ho])
         @test x isa ArrowTypes.ToArrow{String,Vector{Symbol}}
         @test eltype(x) == String
-        @test getfield(x, :needsconvert)
+        @test ArrowTypes._needsconvert(x)
         @test x[1] == "hey"
         @test collect(x) == ["hey", "ho"]
         @test x == ["hey", "ho"]
@@ -203,32 +210,54 @@ end
         @test collect(x) == [1.0, 3.14]
         @test x == [1.0, 3.14]
 
+        x = ArrowTypes.ToArrow(Any[UUID(UInt128(1)), UUID(UInt128(2))])
+        @test x isa ArrowTypes.ToArrow{NTuple{16,UInt8},Vector{Any}}
+        @test eltype(x) == NTuple{16,UInt8}
+        @test collect(x) == [ArrowTypes.toarrow(UUID(UInt128(1))), ArrowTypes.toarrow(UUID(UInt128(2)))]
+
+        x = ArrowTypes.ToArrow(Any[missing, UUID(UInt128(1))])
+        @test x isa ArrowTypes.ToArrow{Union{Missing,NTuple{16,UInt8}},Vector{Any}}
+        @test eltype(x) == Union{Missing,NTuple{16,UInt8}}
+        @test isequal(collect(x), Union{Missing,NTuple{16,UInt8}}[missing, ArrowTypes.toarrow(UUID(UInt128(1)))])
+
         x = ArrowTypes.ToArrow(Any[1, 3.14, "hey"])
         @test x isa ArrowTypes.ToArrow{Union{Float64,String},Vector{Any}}
         @test eltype(x) == Union{Float64,String}
         @test collect(x) == Union{Float64,String}[1.0, 3.14, "hey"]
         @test x == [1.0, 3.14, "hey"]
 
+        x = ArrowTypes.ToArrow(Any[UUID(UInt128(1)), "tail"])
+        @test x isa ArrowTypes.ToArrow{Union{NTuple{16,UInt8},String},Vector{Any}}
+        @test eltype(x) == Union{NTuple{16,UInt8},String}
+        @test collect(x) == Union{NTuple{16,UInt8},String}[ArrowTypes.toarrow(UUID(UInt128(1))), "tail"]
+
         x = ArrowTypes.ToArrow(OffsetArray([1, 2, 3], -3:-1))
         @test x isa ArrowTypes.ToArrow{Int,OffsetVector{Int,Vector{Int}}}
+        @test ArrowTypes._hasoffsetaxes(getfield(x, :data))
+        @test getfield(x, :offset) == ArrowTypes._offsetshift(getfield(x, :data))
+        @test ArrowTypes._sourcedata(x) === getfield(x, :data)
+        @test ArrowTypes._sourceoffset(x) == getfield(x, :offset)
+        @test !ArrowTypes._needsconvert(x)
+        @test ArrowTypes._sourcevalue(x, 1) == 1
         @test eltype(x) == Int
-        @test !getfield(x, :needsconvert)
         @test x[1] == 1
         @test x[3] == 3
         @test collect(x) == [1, 2, 3]
         @test x == [1, 2, 3]
 
         x = ArrowTypes.ToArrow(OffsetArray(Union{Missing,Int}[1, missing], -3:-2))
         @test x isa ArrowTypes.ToArrow{Union{Missing,Int},OffsetVector{Union{Missing,Int},Vector{Union{Missing,Int}}}}
-        @test !getfield(x, :needsconvert)
+        @test !ArrowTypes._needsconvert(x)
         @test x[1] == 1
         @test x[2] === missing
         @test isequal(collect(x), Union{Missing,Int}[1, missing])
 
         x = ArrowTypes.ToArrow(OffsetArray(Any[1, 3.14], -3:-2))
         @test x isa ArrowTypes.ToArrow{Float64,OffsetVector{Any,Vector{Any}}}
+        @test getfield(x, :offset) == ArrowTypes._offsetshift(getfield(x, :data))
+        @test ArrowTypes._sourcevalue(x, 2) == 3.14
         @test eltype(x) == Float64
-        @test getfield(x, :needsconvert)
+        @test ArrowTypes._needsconvert(x)
         @test x[1] == 1
         @test x[2] == 3.14
         @test collect(x) == [1.0, 3.14]

src/arraytypes/arraytypes.jl

@@ -99,7 +99,7 @@ function arrowvector(
             dictencode=dictencode,
             kw...,
         )
-    elseif !(x isa DictEncode)
+    elseif !(x isa DictEncode) && !_keeprawmapvector(T, x)
         x = ToArrow(x)
     end
     S = maybemissing(eltype(x))
@@ -144,6 +144,64 @@ function _arrowtypemeta(meta, n, m)
     return toidict(dict)
 end
 
+@inline function _materializeconverted(x::ArrowTypes.ToArrow)
+    data = ArrowTypes._sourcedata(x)
+    if ArrowTypes._needsconvert(x) && !ArrowTypes.concrete_or_concreteunion(eltype(data))
+        return _materializeconverted(eltype(x), x)
+    end
+    return x
+end
+
+function _materializeconverted(::Type{T}, x::ArrowTypes.ToArrow{T,A}) where {T,A}
+    len = length(x)
+    data = Vector{T}(undef, len)
+    source = ArrowTypes._sourcedata(x)
+    i = 1
+    for value in source
+        @inbounds data[i] =
+            value isa T ? value : ArrowTypes._convert(T, ArrowTypes.toarrow(value))
+        i += 1
+    end
+    return data
+end
+
+@inline function _materializefixedbytes16(value)
+    if value isa ArrowTypes.UUID
+        return ArrowTypes._cast(NTuple{16,UInt8}, value.value)
+    elseif value isa NTuple{16,UInt8}
+        return value
+    else
+        return ArrowTypes._convert(NTuple{16,UInt8}, ArrowTypes.toarrow(value))
+    end
+end
+
+function _materializeconverted(
+    ::Type{NTuple{16,UInt8}},
+    x::ArrowTypes.ToArrow{NTuple{16,UInt8},A},
+) where {A}
+    len = length(x)
+    data = Vector{NTuple{16,UInt8}}(undef, len)
+    source = ArrowTypes._sourcedata(x)
+    i = 1
+    for value in source
+        @inbounds data[i] = _materializefixedbytes16(value)
+        i += 1
+    end
+    return data
+end
+
+@inline _toarrowvaliditysource(x::ArrowTypes.ToArrow) =
+    ArrowTypes._needsconvert(x) ? x : ArrowTypes._sourcedata(x)
+
+@inline _toarrowvalidity(x::ArrowTypes.ToArrow, data) =
+    data === x ? ValidityBitmap(x) : ValidityBitmap(data)
+
+@inline function _keeprawmapvector(::Type{T}, x) where {T}
+    return Base.has_offset_axes(x) &&
+           ArrowTypes.concrete_or_concreteunion(T) &&
+           ArrowKind(T) isa ArrowTypes.MapKind
+end
+
 # now we check for ArrowType converions and dispatch on ArrowKind
 function arrowvector(::Type{S}, x, i, nl, fi, de, ded, meta; kw...) where {S}
     meta = _normalizemeta(meta)
@@ -240,7 +298,7 @@ function ValidityBitmap(x::ArrowTypes.ToArrow)
     if !(T >: Missing)
         return ValidityBitmap(UInt8[], 1, length(x), 0)
     end
-    source = getfield(x, :needsconvert) ? x : getfield(x, :data)
+    source = _toarrowvaliditysource(x)
     return _validitybitmap(source, length(x))
 end
 

src/arraytypes/bool.jl

@@ -83,11 +83,12 @@ function arrowvector(::BoolKind, x, i, nl, fi, de, ded, meta; kw...)
 end
 
 function arrowvector(::BoolKind, x::ArrowTypes.ToArrow, i, nl, fi, de, ded, meta; kw...)
-    validity = ValidityBitmap(x)
-    len = length(x)
-    source = getfield(x, :needsconvert) ? x : getfield(x, :data)
+    data = _materializeconverted(x)
+    validity = _toarrowvalidity(x, data)
+    len = length(data)
+    source = data === x ? _toarrowvaliditysource(x) : data
     bytes = _packboolbytes(source, len)
-    return BoolVector{eltype(x)}(bytes, 1, validity, len, meta)
+    return BoolVector{eltype(data)}(bytes, 1, validity, len, meta)
 end
 
 function compress(Z::Meta.CompressionType.T, comp, p::P) where {P<:BoolVector}