implement lazy row (JuliaArrays#24)

Author: Pietro Vertechi

src/StructArrays.jl

@@ -8,6 +8,7 @@ include("structarray.jl")
 include("utils.jl")
 include("collect.jl")
 include("sort.jl")
+include("lazy.jl")
 
 function __init__()
     Requires.@require Tables="bd369af6-aec1-5ad0-b16a-f7cc5008161c" include("tables.jl")
@@ -16,6 +17,7 @@ function __init__()
     end
     Requires.@require WeakRefStrings="ea10d353-3f73-51f8-a26c-33c1cb351aa5" begin
         isstringarray(::WeakRefStrings.StringArray) = true
+        arrayof(::Type{T}, d) where {T<:AbstractString} = WeakRefStrings.StringArray{T}(d)
     end
 end
 

src/collect.jl

@@ -1,18 +1,25 @@
-struct StructArrayInitializer{F}
+arrayof(S, d) = Array{S}(undef, d)
+
+struct StructArrayInitializer{F, G}
     unwrap::F
+    arrayof::G
 end
-StructArrayInitializer() = StructArrayInitializer(t -> false)
+StructArrayInitializer(unwrap = t->false) = StructArrayInitializer(unwrap, arrayof)
 
 const default_initializer = StructArrayInitializer()
 
-(s::StructArrayInitializer)(S, d) = StructArray{S}(undef, d; unwrap = s.unwrap)
+function (s::StructArrayInitializer)(S, d)
+    ai = ArrayInitializer(s.unwrap, s.arrayof)
+    buildfromschema(typ -> ai(typ, d), S)
+end
 
-struct ArrayInitializer{F}
+struct ArrayInitializer{F, G}
     unwrap::F
+    arrayof::G
 end
-ArrayInitializer() = ArrayInitializer(t -> false)
+ArrayInitializer(unwrap = t->false) = ArrayInitializer(unwrap, arrayof)
 
-(s::ArrayInitializer)(S, d) = _undef_array(S, d; unwrap = s.unwrap)
+(s::ArrayInitializer)(S, d) = s.unwrap(S) ? buildfromschema(typ -> s(typ, d), S) : s.arrayof(S, d)
 
 _reshape(v, itr, ::Base.HasShape) = reshape(v, axes(itr))
 _reshape(v, itr, ::Union{Base.HasLength, Base.SizeUnknown}) = v

src/lazy.jl

@@ -0,0 +1,16 @@
+struct LazyRow{T, N, C, I}
+    columns::StructArray{T, N, C} # a `Columns` object
+    index::I
+end
+
+Base.getproperty(c::LazyRow, nm::Symbol) = getproperty(getfield(c, 1), nm)[getfield(c, 2)]
+Base.setproperty!(c::LazyRow, nm::Symbol, val) = (getproperty(getfield(c, 1), nm)[getfield(c, 2)] = val; nothing)
+Base.propertynames(c::LazyRow) = propertynames(getfield(c, 1))
+
+staticschema(::Type{<:LazyRow{T}}) where {T} = staticschema(T)
+buildfromschema(f, ::Type{<:LazyRow{T}}) where {T} = buildfromschema(f, T)
+
+iscompatible(::Type{<:LazyRow{S}}, ::Type{StructArray{T, N, C}}) where {S, T, N, C} =
+    iscompatible(S, StructArray{T, N, C})
+
+(s::ArrayInitializer)(::Type{<:LazyRow{T}}, d) where {T} = buildfromschema(typ -> s(typ, d), T)

test/runtests.jl

@@ -1,4 +1,5 @@
 using StructArrays
+using StructArrays: LazyRow
 import Tables, PooledArrays, WeakRefStrings
 using Test
 
@@ -251,6 +252,10 @@ collect_fieldarrays_rec(t) = StructArrays.collect_fieldarrays(t, initializer = i
     @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])))
 
+    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
+
     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"])))