switch to getcolumn and support tables in ND StructArray (#119)

* switch to getcolumn and support tables in ND StructArray

* lazy cleanup

* more tests

* revert changing columns

Author: Pietro Vertechi

src/StructArrays.jl

@@ -1,6 +1,7 @@
 module StructArrays
 
 using Base: tuple_type_cons, tuple_type_head, tuple_type_tail, tail
+using Tables: getcolumn, Tables
 
 export StructArray, StructVector, LazyRow, LazyRows
 export collect_structarray, fieldarrays

src/lazy.jl

@@ -18,7 +18,9 @@ Base.propertynames(c::LazyRow) = propertynames(getfield(c, 1))
 
 function Base.show(io::IO, c::LazyRow)
     print(io, "LazyRow")
-    show(io, to_tup(c))
+    columns, index = getfield(c, 1), getfield(c, 2)
+    tup = StructArray(fieldarrays(columns))[index]
+    show(io, tup)
 end
 
 staticschema(::Type{<:LazyRow{T}}) where {T} = staticschema(T)

src/tables.jl

@@ -1,11 +1,9 @@
-using Tables: Tables
+Tables.isrowtable(::Type{<:StructArray}) = true
 
-Tables.isrowtable(::Type{<:StructVector}) = true
-
-Tables.columnaccess(::Type{<:StructVector}) = true
-Tables.columns(s::StructVector) = fieldarrays(s)
-
-Tables.schema(s::StructVector) = Tables.Schema(staticschema(eltype(s)))
+Tables.columnaccess(::Type{<:StructArray}) = true
+Tables.columns(s::StructArray) = fieldarrays(s)
+Tables.getcolumn(s::StructArray, i::Int) = getproperty(s, i)
+Tables.schema(s::StructArray) = Tables.Schema(staticschema(eltype(s)))
 
 function Base.append!(s::StructVector, rows)
     if Tables.isrowtable(rows) && Tables.columnaccess(rows)
@@ -14,9 +12,7 @@ function Base.append!(s::StructVector, rows)
         table = Tables.columns(rows)
         isempty(_setdiff(propertynames(s), Tables.columnnames(rows))) ||
             _invalid_columns_error(s, rows)
-        _foreach(propertynames(s)) do name
-            append!(getproperty(s, name), Tables.getcolumn(table, name))
-        end
+        foreachfield(append!, s, table)
         return s
     else
         # Otherwise, fallback to a generic implementation expecting

src/utils.jl

@@ -21,13 +21,6 @@ function buildfromschema(initializer, ::Type{T}, ::Type{NT}) where {T, NT<:Tup}
     StructArray{T}(nt)
 end
 
-@static if VERSION < v"1.2.0"
-    @inline _getproperty(v::Tuple, field) = getfield(v, field)
-    @inline _getproperty(v, field) = getproperty(v, field)
-else
-    const _getproperty = getproperty
-end
-
 function _foreachfield(names, L)
     vars = ntuple(i -> gensym(), L)
     exprs = Expr[]
@@ -36,7 +29,7 @@ function _foreachfield(names, L)
     end
     for field in names
         sym = QuoteNode(field)
-        args = [Expr(:call, :_getproperty, var, sym) for var in vars]
+        args = [Expr(:call, :getcolumn, var, sym) for var in vars]
         push!(exprs, Expr(:call, :f, args...))
     end
     push!(exprs, :(return nothing))
@@ -122,13 +115,6 @@ function replace_storage(f, s::StructArray{T}) where T
     StructArray{T}(newcols)
 end
 
-to_tup(c::T) where {T} = to_tup(c, fieldnames(staticschema(T)))
-function to_tup(c, fields::NTuple{N, Symbol}) where N
-    t = ntuple(i -> getproperty(c, fields[i]), N)
-    return NamedTuple{fields}(t)
-end
-to_tup(c, fields::NTuple{N, Int}) where {N} = ntuple(i -> _getproperty(c, fields[i]), N)
-
 astuple(::Type{NamedTuple{names, types}}) where {names, types} = types
 astuple(::Type{T}) where {T<:Tuple} = T
 
@@ -148,8 +134,3 @@ _setdiff(a, b) = setdiff(a, b)
 @inline _setdiff(a::Tuple, ::Tuple{}) = a
 @inline _setdiff(a::Tuple, b::Tuple) = _setdiff(_exclude(a, b[1]), Base.tail(b))
 @inline _exclude(a, b) = foldl((ys, x) -> x == b ? ys : (ys..., x), a; init = ())
-
-# _foreach(f, xs) = foreach(f, xs)
-_foreach(f, xs::Tuple) = foldl((_, x) -> (f(x); nothing), xs; init = nothing)
-# Note `foreach` is not optimized for tuples yet.
-# See: https://github.com/JuliaLang/julia/pull/31901

test/runtests.jl

@@ -318,11 +318,6 @@ cols_infer() = StructArray(([1, 2], [1.2, 2.3]))
     @inferred cols_infer()
 end
 
-@testset "to_(named)tuple" begin
-    @test StructArrays.to_tup((1, 2, 3)) == (1, 2, 3)
-    @test StructArrays.to_tup(2 + 3im) == (re = 2, im = 3)
-end
-
 @testset "propertynames" begin
     a = StructArray{ComplexF64}((Float64[], Float64[]))
     @test sort(collect(propertynames(a))) == [:im, :re]
@@ -341,6 +336,10 @@ end
     @test Tables.rowaccess(typeof(s))
     @test Tables.columnaccess(s)
     @test Tables.columnaccess(typeof(s))
+    @test Tables.getcolumn(s, 1) == [1]
+    @test Tables.getcolumn(s, :a) == [1]
+    @test Tables.getcolumn(s, 2) == ["test"]
+    @test Tables.getcolumn(s, :b) == ["test"]
     @test append!(StructArray([1im]), [(re = 111, im = 222)]) ==
         StructArray([1im, 111 + 222im])
     @test append!(StructArray([1im]), (x for x in [(re = 111, im = 222)])) ==
@@ -626,11 +625,11 @@ end
     rows = LazyRows(s)
     @test IndexStyle(rows) isa IndexLinear
     @test IndexStyle(typeof(rows)) isa IndexLinear
-    @test all(t -> StructArrays._getproperty(t, 1) >= 0, s)
+    @test all(t -> Tables.getcolumn(t, 1) >= 0, s)
     @test all(t -> getproperty(t, 1) >= 0, rows)
     setproperty!(rows[13], 1, -12)
     setproperty!(rows[13], 2, 0)
-    @test !all(t -> StructArrays._getproperty(t, 1) >= 0, s)
+    @test !all(t -> Tables.getcolumn(t, 1) >= 0, s)
     @test !all(t -> getproperty(t, 1) >= 0, rows)
 
     io = IOBuffer()