Optimize wrapped-column write paths and fix RAT

Author: guangtao

dev/release/rat_exclude_files.txt

@@ -16,6 +16,7 @@
 # under the License.
 
 Manifest.toml
+*/Manifest.toml
 dev/release/apache-rat-*.jar
 dev/release/filtered_rat.txt
 dev/release/rat.xml

src/ArrowTypes/src/ArrowTypes.jl

@@ -419,6 +419,12 @@ end
 # lazily call toarrow(x) on getindex for each x in data
 struct ToArrow{T,A} <: AbstractVector{T}
     data::A
+    offset::Int
+    needsconvert::Bool
+end
+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)
 end
 
 concrete_or_concreteunion(T) =
@@ -464,7 +470,29 @@ function _convert(::Type{T}, x) where {T}
         return convert(T, x)
     end
 end
-Base.getindex(x::ToArrow{T}, i::Int) where {T} =
-    _convert(T, toarrow(getindex(x.data, i + firstindex(x.data) - 1)))
+
+@inline function _toarrowvalue(x::ToArrow{T}, value) where {T}
+    x.needsconvert || 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)
+    return _toarrowvalue(x, value)
+end
+
+function Base.iterate(x::ToArrow)
+    state = iterate(x.data)
+    state === nothing && return nothing
+    value, st = state
+    return _toarrowvalue(x, value), st
+end
+
+function Base.iterate(x::ToArrow, st)
+    state = iterate(x.data, st)
+    state === nothing && return nothing
+    value, st = state
+    return _toarrowvalue(x, value), st
+end
 
 end # module ArrowTypes

src/ArrowTypes/test/tests.jl

@@ -192,26 +192,46 @@ end
         x = ArrowTypes.ToArrow([:hey, :ho])
         @test x isa ArrowTypes.ToArrow{String,Vector{Symbol}}
         @test eltype(x) == String
+        @test getfield(x, :needsconvert)
+        @test x[1] == "hey"
+        @test collect(x) == ["hey", "ho"]
         @test x == ["hey", "ho"]
 
         x = ArrowTypes.ToArrow(Any[1, 3.14])
         @test x isa ArrowTypes.ToArrow{Float64,Vector{Any}}
         @test eltype(x) == Float64
+        @test collect(x) == [1.0, 3.14]
         @test x == [1.0, 3.14]
 
         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(OffsetArray([1, 2, 3], -3:-1))
         @test x isa ArrowTypes.ToArrow{Int,OffsetVector{Int,Vector{Int}}}
         @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 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 eltype(x) == Float64
+        @test getfield(x, :needsconvert)
+        @test x[1] == 1
+        @test x[2] == 3.14
+        @test collect(x) == [1.0, 3.14]
         @test x == [1, 3.14]
 
         @testset "respect non-missing concrete type" begin

src/arraytypes/arraytypes.jl

@@ -201,15 +201,10 @@ end
 Base.size(p::ValidityBitmap) = (p.ℓ,)
 nullcount(x::ValidityBitmap) = x.nc
 
-function ValidityBitmap(x)
-    T = eltype(x)
-    if !(T >: Missing)
-        return ValidityBitmap(UInt8[], 1, length(x), 0)
-    end
+function _validitybitmap(x, len)
     len = length(x)
     blen = cld(len, 8)
     bytes = Vector{UInt8}(undef, blen)
-    st = iterate(x)
     nc = 0
     b = 0xff
     j = k = 1
@@ -232,6 +227,23 @@ function ValidityBitmap(x)
     return ValidityBitmap(nc == 0 ? UInt8[] : bytes, 1, nc == 0 ? 0 : len, nc)
 end
 
+function ValidityBitmap(x)
+    T = eltype(x)
+    if !(T >: Missing)
+        return ValidityBitmap(UInt8[], 1, length(x), 0)
+    end
+    return _validitybitmap(x, length(x))
+end
+
+function ValidityBitmap(x::ArrowTypes.ToArrow)
+    T = eltype(x)
+    if !(T >: Missing)
+        return ValidityBitmap(UInt8[], 1, length(x), 0)
+    end
+    source = getfield(x, :needsconvert) ? x : getfield(x, :data)
+    return _validitybitmap(source, length(x))
+end
+
 @propagate_inbounds function Base.getindex(p::ValidityBitmap, i::Integer)
     # no boundscheck because parent array should do it
     # if a validity bitmap is empty, it either means:

src/arraytypes/bool.jl

@@ -52,9 +52,7 @@ end
 
 arrowvector(::BoolKind, x::BoolVector, i, nl, fi, de, ded, meta; kw...) = x
 
-function arrowvector(::BoolKind, x, i, nl, fi, de, ded, meta; kw...)
-    validity = ValidityBitmap(x)
-    len = length(x)
+function _packboolbytes(x, len)
     blen = cld(len, 8)
     bytes = Vector{UInt8}(undef, blen)
     b = 0xff
@@ -74,6 +72,21 @@ function arrowvector(::BoolKind, x, i, nl, fi, de, ded, meta; kw...)
     if j > 1
         bytes[k] = b
     end
+    return bytes
+end
+
+function arrowvector(::BoolKind, x, i, nl, fi, de, ded, meta; kw...)
+    validity = ValidityBitmap(x)
+    len = length(x)
+    bytes = _packboolbytes(x, len)
+    return BoolVector{eltype(x)}(bytes, 1, validity, len, meta)
+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)
+    bytes = _packboolbytes(source, len)
     return BoolVector{eltype(x)}(bytes, 1, validity, len, meta)
 end
 

src/arraytypes/list.jl

@@ -219,6 +219,22 @@ end
     return x, (i, chunk, chunk_i, chunk_len, len)
 end
 
+@inline function _writeuint8chunk(io::IO, bytes)
+    GC.@preserve bytes begin
+        return Base.unsafe_write(io, pointer(bytes), length(bytes))
+    end
+end
+
+function writearray(io::IO, ::Type{UInt8}, col::ToList{UInt8,stringtype}) where {stringtype}
+    n = 0
+    for chunk in col.data
+        chunk === missing && continue
+        bytes = stringtype ? _codeunits(chunk) : chunk
+        n += _writeuint8chunk(io, bytes)
+    end
+    return n
+end
+
 arrowvector(::ListKind, x::List, i, nl, fi, de, ded, meta; kw...) = x
 
 function arrowvector(::ListKind, x, i, nl, fi, de, ded, meta; largelists::Bool=false, kw...)

src/utils.jl

@@ -37,6 +37,33 @@ end
 # efficient writing of arrays
 writearray(io, col) = writearray(io, maybemissing(eltype(col)), col)
 
+function _writearrayfallback(io::IO, ::Type{T}, col) where {T}
+    n = 0
+    data = Vector{UInt8}(undef, sizeof(col))
+    buf = IOBuffer(data; write=true)
+    for x in col
+        n += Base.write(buf, coalesce(x, ArrowTypes.default(T)))
+    end
+    n = Base.write(io, take!(buf))
+    return n
+end
+
+@inline function _writearraycontiguous(io::IO, ::Type{T}, data) where {T}
+    return Base.unsafe_write(io, pointer(data), sizeof(T) * length(data))
+end
+
+@inline function _contiguoustoarrowdata(::Type{T}, col::ArrowTypes.ToArrow) where {T}
+    getfield(col, :needsconvert) && return nothing
+    data = getfield(col, :data)
+    strides(data) == (1,) || return nothing
+    if data isa AbstractVector{T}
+        return isbitstype(T) ? data : nothing
+    elseif isbitstype(T) && data isa AbstractVector{Union{T,Missing}}
+        return data
+    end
+    return nothing
+end
+
 function writearray(io::IO, ::Type{T}, col) where {T}
     if col isa Vector{T}
         n = Base.write(io, col)
@@ -51,17 +78,17 @@ function writearray(io::IO, ::Type{T}, col) where {T}
             n += writearray(io, T, A)
         end
     else
-        n = 0
-        data = Vector{UInt8}(undef, sizeof(col))
-        buf = IOBuffer(data; write=true)
-        for x in col
-            n += Base.write(buf, coalesce(x, ArrowTypes.default(T)))
-        end
-        n = Base.write(io, take!(buf))
+        n = _writearrayfallback(io, T, col)
     end
     return n
 end
 
+function writearray(io::IO, ::Type{T}, col::ArrowTypes.ToArrow) where {T}
+    data = _contiguoustoarrowdata(T, col)
+    isnothing(data) || return _writearraycontiguous(io, T, data)
+    return _writearrayfallback(io, T, col)
+end
+
 getbit(v::UInt8, n::Integer) = (v & (1 << (n - 1))) > 0x00
 
 function setbit(v::UInt8, b::Bool, n::Integer)

test/Project.toml

@@ -1,3 +1,20 @@
+# 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.
+
 [deps]
 Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"
 ArrowTypes = "31f734f8-188a-4ce0-8406-c8a06bd891cd"

test/runtests.jl

@@ -295,6 +295,34 @@ end
             @test isequal(tt.a, ['a', missing])
         end
 
+        @testset "# offset bool write paths" begin
+            t = (
+                a=OffsetArray(Bool[true, false, true], -1:1),
+                b=OffsetArray(Union{Missing,Bool}[true, missing, false], -1:1),
+                c=OffsetArray(Any[true, false, true], -1:1),
+                d=OffsetArray(Any[true, missing, false], -1:1),
+            )
+            tt = Arrow.Table(Arrow.tobuffer(t))
+            @test tt.a == Bool[true, false, true]
+            @test isequal(tt.b, Union{Missing,Bool}[true, missing, false])
+            @test tt.c == Bool[true, false, true]
+            @test isequal(tt.d, Union{Missing,Bool}[true, missing, false])
+        end
+
+        @testset "# offset primitive write paths" begin
+            t = (
+                a=OffsetArray(Int64[1, 2, 3], -1:1),
+                b=OffsetArray(Union{Missing,Int64}[1, missing, 3], -1:1),
+                c=OffsetArray(Any[1, 2, 3], -1:1),
+                d=OffsetArray(Any[1, missing, 3], -1:1),
+            )
+            tt = Arrow.Table(Arrow.tobuffer(t))
+            @test tt.a == Int64[1, 2, 3]
+            @test isequal(tt.b, Union{Missing,Int64}[1, missing, 3])
+            @test tt.c == Int64[1, 2, 3]
+            @test isequal(tt.d, Union{Missing,Int64}[1, missing, 3])
+        end
+
         @testset "# automatic custom struct serialization/deserialization" begin
             t = (col1=[CustomStruct(1, 2.3, "hey"), CustomStruct(4, 5.6, "there")],)
 
@@ -974,6 +1002,18 @@ end
             @test isequal(t1.bm, t2.bm)
             @test isequal(t1.c, t2.c)
             @test isequal(t1.cm, t2.cm)
+
+            toffset = (
+                b=OffsetArray([b"01", b"", b"3"], -1:1),
+                bm=OffsetArray(Union{Missing,Base.CodeUnits{UInt8,String}}[b"01", b"3", missing], -1:1),
+                c=OffsetArray(["a", "b", "c"], -1:1),
+                cm=OffsetArray(Union{Missing,String}["a", "c", missing], -1:1),
+            )
+            ttoffset = Arrow.Table(Arrow.tobuffer(toffset))
+            @test collect(toffset.b) == ttoffset.b
+            @test isequal(collect(toffset.bm), ttoffset.bm)
+            @test collect(toffset.c) == ttoffset.c
+            @test isequal(collect(toffset.cm), ttoffset.cm)
         end
 
         @testset "# 435" begin