Merge pull request #35 from piever/pv/eltype

rename to collect_structarray

Author: Pietro Vertechi

NEWS.md

@@ -8,6 +8,6 @@
 
 ### New features
 
-- Added `collect_fieldarrays` function to collect an iterable of structs into a `StructArray` without having to allocate an array of structs
+- Added `collect_structarray` function to collect an iterable of structs into a `StructArray` without having to allocate an array of structs
 - `StructArray{T}(undef, dims)` and `StructArray(v::AbstractArray)` now support an `unwrap` keyword argument to specify on which types to do recursive unnesting of array of structs to struct of arrays
 

src/StructArrays.jl

@@ -2,6 +2,7 @@ module StructArrays
 
 import Requires
 export StructArray, StructVector
+export collect_structarray
 
 include("interface.jl")
 include("structarray.jl")

src/collect.jl

@@ -24,29 +24,30 @@ ArrayInitializer(unwrap = t->false) = ArrayInitializer(unwrap, arrayof)
 _reshape(v, itr, ::Base.HasShape) = reshape(v, axes(itr))
 _reshape(v, itr, ::Union{Base.HasLength, Base.SizeUnknown}) = v
 
-function collect_fieldarrays(itr; initializer = default_initializer)
+function collect_structarray(itr; initializer = default_initializer)
     sz = Base.IteratorSize(itr)
-    v = collect_fieldarrays(itr, sz, initializer = initializer)
+    v = collect_structarray(itr, sz, initializer = initializer)
     _reshape(v, itr, sz)
 end
 
-function collect_empty_fieldarrays(itr::T; initializer = default_initializer) where {T}
+function collect_empty_structarray(itr::T; initializer = default_initializer) where {T}
     S = Core.Compiler.return_type(first, Tuple{T})
     initializer(S, (0,))
 end
 
-function collect_fieldarrays(@nospecialize(itr), ::Union{Base.HasShape, Base.HasLength};
-    initializer = default_initializer)
+function collect_structarray(itr, ::Union{Base.HasShape, Base.HasLength};
+                             initializer = default_initializer) 
 
     st = iterate(itr)
-    st === nothing && return collect_empty_fieldarrays(itr, initializer = initializer)
+    st === nothing && return collect_empty_structarray(itr, initializer = initializer)
     el, i = st
-    dest = initializer(typeof(el), (length(itr),))
+    S = typeof(el)
+    dest = initializer(S, (length(itr),))
     dest[1] = el
-    collect_to_fieldarrays!(dest, itr, 2, i)
+    collect_to_structarray!(dest, itr, 2, i)
 end
 
-function collect_to_fieldarrays!(dest::AbstractArray{T}, itr, offs, st) where {T}
+function collect_to_structarray!(dest::AbstractArray{T}, itr, offs, st) where {T}
     # collect to dest array, checking the type of each result. if a result does not
     # match, widen the result type and re-dispatch.
     i = offs
@@ -58,24 +59,24 @@ function collect_to_fieldarrays!(dest::AbstractArray{T}, itr, offs, st) where {T
             @inbounds dest[i] = el
             i += 1
         else
-            new = widenfieldarrays(dest, i, el)
+            new = widenstructarray(dest, i, el)
             @inbounds new[i] = el
-            return collect_to_fieldarrays!(new, itr, i+1, st)
+            return collect_to_structarray!(new, itr, i+1, st)
         end
     end
     return dest
 end
 
-function collect_fieldarrays(itr, ::Base.SizeUnknown; initializer = default_initializer)
+function collect_structarray(itr, ::Base.SizeUnknown; initializer = default_initializer)
     elem = iterate(itr)
-    elem === nothing && return collect_empty_fieldarrays(itr; initializer = initializer)
+    elem === nothing && return collect_empty_structarray(itr; initializer = initializer)
     el, st = elem
     dest = initializer(typeof(el), (1,))
     dest[1] = el
-    grow_to_fieldarrays!(dest, itr, iterate(itr, st))
+    grow_to_structarray!(dest, itr, iterate(itr, st))
 end
 
-function grow_to_fieldarrays!(dest::AbstractArray{T}, itr, elem = iterate(itr)) where {T}
+function grow_to_structarray!(dest::AbstractArray{T}, itr, elem = iterate(itr)) where {T}
     # collect to dest array, checking the type of each result. if a result does not
     # match, widen the result type and re-dispatch.
     i = length(dest)+1
@@ -86,9 +87,9 @@ function grow_to_fieldarrays!(dest::AbstractArray{T}, itr, elem = iterate(itr))
             elem = iterate(itr, st)
             i += 1
         else
-            new = widenfieldarrays(dest, i, el)
+            new = widenstructarray(dest, i, el)
             push!(new, el)
-            return grow_to_fieldarrays!(new, itr, iterate(itr, st))
+            return grow_to_structarray!(new, itr, iterate(itr, st))
         end
     end
     return dest
@@ -99,22 +100,22 @@ function to_structarray(::Type{T}, nt::C) where {T, C}
     StructArray{S}(nt)
 end
 
-function widenfieldarrays(dest::StructArray{T}, i, el::S) where {T, S}
+function widenstructarray(dest::StructArray{T}, i, el::S) where {T, S}
     fs = fields(S)
     if fs === fields(T)
-        new_cols = (widenfieldarrays(fieldarrays(dest)[ind], i, getfieldindex(el, f, ind)) for (ind, f) in enumerate(fs))
+        new_cols = (widenstructarray(fieldarrays(dest)[ind], i, getfieldindex(el, f, ind)) for (ind, f) in enumerate(fs))
         nt = NamedTuple{fs}(Tuple(new_cols))
         v = to_structarray(T, nt)
     else
         widenarray(dest, i, el)
     end
 end
 
-widenfieldarrays(dest::AbstractArray, i, el) = widenarray(dest, i, el)
+widenstructarray(dest::AbstractArray, i, el) = widenarray(dest, i, el)
 
 function widenarray(dest::AbstractArray{T}, i, el::S) where {S, T}
     S <: T && return dest
-    new = similar(dest, promote_type(S, T), length(dest))
+    new = similar(dest, Base.promote_typejoin(S, T), length(dest))
     copyto!(new, 1, dest, 1, i-1)
     new
 end

test/runtests.jl

@@ -237,142 +237,142 @@ StructArrays.SkipConstructor(::Type{<:S}) = true
 end
 
 const initializer = StructArrays.ArrayInitializer(t -> t <: Union{Tuple, NamedTuple, Pair})
-collect_fieldarrays_rec(t) = StructArrays.collect_fieldarrays(t, initializer = initializer)
+collect_structarray_rec(t) = collect_structarray(t, initializer = initializer)
 
 @testset "collectnamedtuples" begin
     v = [(a = 1, b = 2), (a = 1, b = 3)]
-    collect_fieldarrays_rec(v) == StructArray((a = Int[1, 1], b = Int[2, 3]))
+    collect_structarray_rec(v) == StructArray((a = Int[1, 1], b = Int[2, 3]))
 
     # test inferrability with constant eltype
     itr = [(a = 1, b = 2), (a = 1, b = 2), (a = 1, b = 12)]
     el, st = iterate(itr)
     dest = initializer(typeof(el), (3,))
     dest[1] = el
-    @inferred StructArrays.collect_to_fieldarrays!(dest, itr, 2, st)
+    @inferred StructArrays.collect_to_structarray!(dest, itr, 2, st)
 
     v = [(a = 1, b = 2), (a = 1.2, b = 3)]
-    @test collect_fieldarrays_rec(v) == StructArray((a = [1, 1.2], b = Int[2, 3]))
-    @test typeof(collect_fieldarrays_rec(v)) == typeof(StructArray((a = [1, 1.2], b = Int[2, 3])))
+    @test collect_structarray_rec(v) == StructArray((a = [1, 1.2], b = Int[2, 3]))
+    @test typeof(collect_structarray_rec(v)) == typeof(StructArray((a = Real[1, 1.2], b = Int[2, 3])))
 
     s = StructArray(a = [1, 2], b  = [3, 4])
-    @test StructArrays.collect_fieldarrays(LazyRow(s, i) for i in eachindex(s)) == s
-    @test collect_fieldarrays_rec(LazyRow(s, i) for i in eachindex(s)) == s
+    @test collect_structarray(LazyRow(s, i) for i in eachindex(s)) == s
+    @test collect_structarray_rec(LazyRow(s, i) for i in eachindex(s)) == s
 
     v = [(a = 1, b = 2), (a = 1.2, b = "3")]
-    @test collect_fieldarrays_rec(v) == StructArray((a = [1, 1.2], b = Any[2, "3"]))
-    @test typeof(collect_fieldarrays_rec(v)) == typeof(StructArray((a = [1, 1.2], b = Any[2, "3"])))
+    @test collect_structarray_rec(v) == StructArray((a = [1, 1.2], b = Any[2, "3"]))
+    @test typeof(collect_structarray_rec(v)) == typeof(StructArray((a = Real[1, 1.2], b = Any[2, "3"])))
 
     v = [(a = 1, b = 2), (a = 1.2, b = 2), (a = 1, b = "3")]
-    @test collect_fieldarrays_rec(v) == StructArray((a = [1, 1.2, 1], b = Any[2, 2, "3"]))
-    @test typeof(collect_fieldarrays_rec(v)) == typeof(StructArray((a = [1, 1.2, 1], b = Any[2, 2, "3"])))
+    @test collect_structarray_rec(v) == StructArray((a = Real[1, 1.2, 1], b = Any[2, 2, "3"]))
+    @test typeof(collect_structarray_rec(v)) == typeof(StructArray((a = Real[1, 1.2, 1], b = Any[2, 2, "3"])))
 
     # length unknown
     itr = Iterators.filter(isodd, 1:8)
     tuple_itr = ((a = i+1, b = i-1) for i in itr)
-    @test collect_fieldarrays_rec(tuple_itr) == StructArray((a = [2, 4, 6, 8], b = [0, 2, 4, 6]))
+    @test collect_structarray_rec(tuple_itr) == StructArray((a = [2, 4, 6, 8], b = [0, 2, 4, 6]))
     tuple_itr_real = (i == 1 ? (a = 1.2, b =i-1) : (a = i+1, b = i-1) for i in itr)
-    @test collect_fieldarrays_rec(tuple_itr_real) == StructArray((a = Real[1.2, 4, 6, 8], b = [0, 2, 4, 6]))
+    @test collect_structarray_rec(tuple_itr_real) == StructArray((a = Real[1.2, 4, 6, 8], b = [0, 2, 4, 6]))
 
     # empty
     itr = Iterators.filter(t -> t > 10, 1:8)
     tuple_itr = ((a = i+1, b = i-1) for i in itr)
-    @test collect_fieldarrays_rec(tuple_itr) == StructArray((a = Int[], b = Int[]))
+    @test collect_structarray_rec(tuple_itr) == StructArray((a = Int[], b = Int[]))
 
     itr = (i for i in 0:-1)
     tuple_itr = ((a = i+1, b = i-1) for i in itr)
-    @test collect_fieldarrays_rec(tuple_itr) == StructArray((a = Int[], b = Int[]))
+    @test collect_structarray_rec(tuple_itr) == StructArray((a = Int[], b = Int[]))
 end
 
 @testset "collecttuples" begin
     v = [(1, 2), (1, 3)]
-    @test collect_fieldarrays_rec(v) == StructArray((Int[1, 1], Int[2, 3]))
-    @inferred collect_fieldarrays_rec(v)
+    @test collect_structarray_rec(v) == StructArray((Int[1, 1], Int[2, 3]))
+    @inferred collect_structarray_rec(v)
 
-    @test StructArrays.collect_fieldarrays(v) == StructArray((Int[1, 1], Int[2, 3]))
-    @inferred StructArrays.collect_fieldarrays(v)
+    @test collect_structarray(v) == StructArray((Int[1, 1], Int[2, 3]))
+    @inferred collect_structarray(v)
 
     v = [(1, 2), (1.2, 3)]
-    @test collect_fieldarrays_rec(v) == StructArray(([1, 1.2], Int[2, 3]))
+    @test collect_structarray_rec(v) == StructArray((Real[1, 1.2], Int[2, 3]))
 
     v = [(1, 2), (1.2, "3")]
-    @test collect_fieldarrays_rec(v) == StructArray(([1, 1.2], Any[2, "3"]))
-    @test typeof(collect_fieldarrays_rec(v)) == typeof(StructArray(([1, 1.2], Any[2, "3"])))
+    @test collect_structarray_rec(v) == StructArray((Real[1, 1.2], Any[2, "3"]))
+    @test typeof(collect_structarray_rec(v)) == typeof(StructArray((Real[1, 1.2], Any[2, "3"])))
 
     v = [(1, 2), (1.2, 2), (1, "3")]
-    @test collect_fieldarrays_rec(v) == StructArray(([1, 1.2, 1], Any[2, 2, "3"]))
+    @test collect_structarray_rec(v) == StructArray((Real[1, 1.2, 1], Any[2, 2, "3"]))
     # length unknown
     itr = Iterators.filter(isodd, 1:8)
     tuple_itr = ((i+1, i-1) for i in itr)
-    @test collect_fieldarrays_rec(tuple_itr) == StructArray(([2, 4, 6, 8], [0, 2, 4, 6]))
+    @test collect_structarray_rec(tuple_itr) == StructArray(([2, 4, 6, 8], [0, 2, 4, 6]))
     tuple_itr_real = (i == 1 ? (1.2, i-1) : (i+1, i-1) for i in itr)
-    @test collect_fieldarrays_rec(tuple_itr_real) == StructArray(([1.2, 4, 6, 8], [0, 2, 4, 6]))
-    @test typeof(collect_fieldarrays_rec(tuple_itr_real)) == typeof(StructArray(([1.2, 4, 6, 8], [0, 2, 4, 6])))
+    @test collect_structarray_rec(tuple_itr_real) == StructArray((Real[1.2, 4, 6, 8], [0, 2, 4, 6]))
+    @test typeof(collect_structarray_rec(tuple_itr_real)) == typeof(StructArray((Real[1.2, 4, 6, 8], [0, 2, 4, 6])))
 
     # empty
     itr = Iterators.filter(t -> t > 10, 1:8)
     tuple_itr = ((i+1, i-1) for i in itr)
-    @test collect_fieldarrays_rec(tuple_itr) == StructArray((Int[], Int[]))
+    @test collect_structarray_rec(tuple_itr) == StructArray((Int[], Int[]))
 
     itr = (i for i in 0:-1)
     tuple_itr = ((i+1, i-1) for i in itr)
-    @test collect_fieldarrays_rec(tuple_itr) == StructArray((Int[], Int[]))
+    @test collect_structarray_rec(tuple_itr) == StructArray((Int[], Int[]))
 end
 
 @testset "collectscalars" begin
     v = (i for i in 1:3)
-    @test collect_fieldarrays_rec(v) == [1,2,3]
-    @inferred collect_fieldarrays_rec(v)
+    @test collect_structarray_rec(v) == [1,2,3]
+    @inferred collect_structarray_rec(v)
 
     v = (i == 1 ? 1.2 : i for i in 1:3)
-    @test collect_fieldarrays_rec(v) == collect(v)
+    @test collect_structarray_rec(v) == collect(v)
 
     itr = Iterators.filter(isodd, 1:100)
-    @test collect_fieldarrays_rec(itr) == collect(itr)
+    @test collect_structarray_rec(itr) == collect(itr)
     real_itr = (i == 1 ? 1.5 : i for i in itr)
-    @test collect_fieldarrays_rec(real_itr) == collect(real_itr)
-    @test eltype(collect_fieldarrays_rec(real_itr)) == Float64
+    @test collect_structarray_rec(real_itr) == collect(real_itr)
+    @test eltype(collect_structarray_rec(real_itr)) == Real
 
     #empty
     itr = Iterators.filter(t -> t > 10, 1:8)
     tuple_itr = (exp(i) for i in itr)
-    @test collect_fieldarrays_rec(tuple_itr) == Float64[]
+    @test collect_structarray_rec(tuple_itr) == Float64[]
 
     itr = (i for i in 0:-1)
     tuple_itr = (exp(i) for i in itr)
-    @test collect_fieldarrays_rec(tuple_itr) == Float64[]
+    @test collect_structarray_rec(tuple_itr) == Float64[]
 
-    t = collect_fieldarrays_rec((a = i,) for i in (1, missing, 3))
+    t = collect_structarray_rec((a = i,) for i in (1, missing, 3))
     @test StructArrays.fieldarrays(t)[1] isa Array{Union{Int, Missing}}
     @test isequal(StructArrays.fieldarrays(t)[1], [1, missing, 3])
 end
 
 @testset "collectpairs" begin
     v = (i=>i+1 for i in 1:3)
-    @test collect_fieldarrays_rec(v) == StructArray{Pair{Int, Int}}([1,2,3], [2,3,4])
-    @test eltype(collect_fieldarrays_rec(v)) == Pair{Int, Int}
+    @test collect_structarray_rec(v) == StructArray{Pair{Int, Int}}([1,2,3], [2,3,4])
+    @test eltype(collect_structarray_rec(v)) == Pair{Int, Int}
 
     v = (i == 1 ? (1.2 => i+1) : (i => i+1) for i in 1:3)
-    @test collect_fieldarrays_rec(v) == StructArray{Pair{Float64, Int}}([1.2,2,3], [2,3,4])
-    @test eltype(collect_fieldarrays_rec(v)) == Pair{Float64, Int}
+    @test collect_structarray_rec(v) == StructArray{Pair{Real, Int}}([1.2,2,3], [2,3,4])
+    @test eltype(collect_structarray_rec(v)) == Pair{Real, Int}
 
     v = ((a=i,) => (b="a$i",) for i in 1:3)
-    @test collect_fieldarrays_rec(v) == StructArray{Pair{NamedTuple{(:a,),Tuple{Int64}},NamedTuple{(:b,),Tuple{String}}}}(StructArray((a = [1,2,3],)), StructArray((b = ["a1","a2","a3"],)))
-    @test eltype(collect_fieldarrays_rec(v)) == Pair{NamedTuple{(:a,), Tuple{Int64}}, NamedTuple{(:b,), Tuple{String}}}
+    @test collect_structarray_rec(v) == StructArray{Pair{NamedTuple{(:a,),Tuple{Int64}},NamedTuple{(:b,),Tuple{String}}}}(StructArray((a = [1,2,3],)), StructArray((b = ["a1","a2","a3"],)))
+    @test eltype(collect_structarray_rec(v)) == Pair{NamedTuple{(:a,), Tuple{Int64}}, NamedTuple{(:b,), Tuple{String}}}
 
     v = (i == 1 ? (a="1",) => (b="a$i",) : (a=i,) => (b="a$i",) for i in 1:3)
-    @test collect_fieldarrays_rec(v) == StructArray{Pair{NamedTuple{(:a,),Tuple{Any}},NamedTuple{(:b,),Tuple{String}}}}(StructArray((a = ["1",2,3],)), StructArray((b = ["a1","a2","a3"],)))
-    @test eltype(collect_fieldarrays_rec(v)) == Pair{NamedTuple{(:a,), Tuple{Any}}, NamedTuple{(:b,), Tuple{String}}}
+    @test collect_structarray_rec(v) == StructArray{Pair{NamedTuple{(:a,),Tuple{Any}},NamedTuple{(:b,),Tuple{String}}}}(StructArray((a = ["1",2,3],)), StructArray((b = ["a1","a2","a3"],)))
+    @test eltype(collect_structarray_rec(v)) == Pair{NamedTuple{(:a,), Tuple{Any}}, NamedTuple{(:b,), Tuple{String}}}
 
     # empty
     v = ((a=i,) => (b="a$i",) for i in 0:-1)
-    @test collect_fieldarrays_rec(v) == StructArray{Pair{NamedTuple{(:a,),Tuple{Int64}},NamedTuple{(:b,),Tuple{String}}}}(StructArray((a = Int[],)), StructArray((b = String[],)))
-    @test eltype(collect_fieldarrays_rec(v)) == Pair{NamedTuple{(:a,), Tuple{Int}}, NamedTuple{(:b,), Tuple{String}}}
+    @test collect_structarray_rec(v) == StructArray{Pair{NamedTuple{(:a,),Tuple{Int64}},NamedTuple{(:b,),Tuple{String}}}}(StructArray((a = Int[],)), StructArray((b = String[],)))
+    @test eltype(collect_structarray_rec(v)) == Pair{NamedTuple{(:a,), Tuple{Int}}, NamedTuple{(:b,), Tuple{String}}}
 
     v = Iterators.filter(t -> t.first.a == 4, ((a=i,) => (b="a$i",) for i in 1:3))
-    @test collect_fieldarrays_rec(v) == StructArray{Pair{NamedTuple{(:a,),Tuple{Int64}},NamedTuple{(:b,),Tuple{String}}}}(StructArray((a = Int[],)), StructArray((b = String[],)))
-    @test eltype(collect_fieldarrays_rec(v)) == Pair{NamedTuple{(:a,), Tuple{Int}}, NamedTuple{(:b,), Tuple{String}}}
+    @test collect_structarray_rec(v) == StructArray{Pair{NamedTuple{(:a,),Tuple{Int64}},NamedTuple{(:b,),Tuple{String}}}}(StructArray((a = Int[],)), StructArray((b = String[],)))
+    @test eltype(collect_structarray_rec(v)) == Pair{NamedTuple{(:a,), Tuple{Int}}, NamedTuple{(:b,), Tuple{String}}}
 
-    t = collect_fieldarrays_rec((b = 1,) => (a = i,) for i in (2, missing, 3))
+    t = collect_structarray_rec((b = 1,) => (a = i,) for i in (2, missing, 3))
     s = StructArray{Pair{NamedTuple{(:b,),Tuple{Int64}},NamedTuple{(:a,),Tuple{Union{Missing, Int64}}}}}(StructArray(b = [1,1,1]), StructArray(a = [2, missing, 3]))
     @test s[1] == t[1]
     @test ismissing(t[2].second.a)
@@ -381,7 +381,7 @@ end
 
 @testset "collect2D" begin
     s = (l for l in [(a=i, b=j) for i in 1:3, j in 1:4])
-    v = StructArrays.collect_fieldarrays(s)
+    v = collect_structarray(s)
     @test size(v) == (3, 4)
     @test v.a == [i for i in 1:3, j in 1:4]
     @test v.b == [j for i in 1:3, j in 1:4]