setindex!: convert to eltype (fixes #216) (#227)

This introduces the internal utility `maybe_convert_elt` used to
convert values to the element type of the StructArray before
assignment. This fixes errors that are caused by assigning a value
that does not have the same fields as the struct but which can be
converted to such a type (e.g., `T<:Real` -> `Complex{T}`). Rather
than calling `convert` directly, `maybe_convert_elt` can be
specialized for particular types, like LazyRow, which should not be
converted.

Fixes #131
Fixes #216
Closes #184

Co-authored-by: Gustavo Goretkin <[email protected]>

Author: timholy

src/lazy.jl

@@ -57,6 +57,8 @@ iscompatible(::Type{<:LazyRow{R}}, ::Type{S}) where {R, S<:StructArray} = iscomp
 
 (s::ArrayInitializer)(::Type{<:LazyRow{T}}, d) where {T} = buildfromschema(typ -> s(typ, d), T)
 
+maybe_convert_elt(::Type{T}, vals::LazyRow) where T = vals
+
 """
     LazyRows(s::StructArray)
 

src/structarray.jl

@@ -351,34 +351,40 @@ end
     StructArray{T}(map(v -> @inbounds(view(v, I...)), components(s)))
 end
 
-Base.@propagate_inbounds function Base.setindex!(s::StructArray{<:Any, <:Any, <:Any, CartesianIndex{N}}, vals, I::Vararg{Int, N}) where {N}
+Base.@propagate_inbounds function Base.setindex!(s::StructArray{T, <:Any, <:Any, CartesianIndex{N}}, vals, I::Vararg{Int, N}) where {T,N}
     @boundscheck checkbounds(s, I...)
-    foreachfield((col, val) -> (@inbounds col[I...] = val), s, vals)
-    s
+    valsT = maybe_convert_elt(T, vals)
+    foreachfield((col, val) -> (@inbounds col[I...] = val), s, valsT)
+    return s
 end
 
-Base.@propagate_inbounds function Base.setindex!(s::StructArray{<:Any, <:Any, <:Any, Int}, vals, I::Int)
+Base.@propagate_inbounds function Base.setindex!(s::StructArray{T, <:Any, <:Any, Int}, vals, I::Int) where T
     @boundscheck checkbounds(s, I)
-    foreachfield((col, val) -> (@inbounds col[I] = val), s, vals)
-    s
+    valsT = maybe_convert_elt(T, vals)
+    foreachfield((col, val) -> (@inbounds col[I] = val), s, valsT)
+    return s
 end
 
 for f in (:push!, :pushfirst!)
-    @eval function Base.$f(s::StructVector, vals)
-        foreachfield($f, s, vals)
+    @eval function Base.$f(s::StructVector{T}, vals) where T
+        valsT = maybe_convert_elt(T, vals)
+        foreachfield($f, s, valsT)
         return s
     end
 end
 
 for f in (:append!, :prepend!)
-    @eval function Base.$f(s::StructVector, vals::StructVector)
+    @eval function Base.$f(s::StructVector{T}, vals::StructVector{T}) where T
+        # If these aren't the same type, there's no guarantee that x.a "means" the same thing as y.a,
+        # even when all the field names match.
         foreachfield($f, s, vals)
         return s
     end
 end
 
-function Base.insert!(s::StructVector, i::Integer, vals)
-    foreachfield((v, val) -> insert!(v, i, val), s, vals)
+function Base.insert!(s::StructVector{T}, i::Integer, vals) where T
+    valsT = maybe_convert_elt(T, vals)
+    foreachfield((v, val) -> insert!(v, i, val), s, valsT)
     return s
 end
 

src/tables.jl

@@ -20,6 +20,8 @@ function try_compatible_columns(rows::R, s::StructArray) where {R}
     _schema(NT) == Tables.schema(rows) || return nothing
     return Tables.columntable(rows)
 end
+try_compatible_columns(rows::StructArray{T}, s::StructArray{T}) where {T} = Tables.columntable(rows)
+try_compatible_columns(rows::StructArray{R}, s::StructArray{S}) where {R,S} = nothing
 
 for (f, g) in zip((:append!, :prepend!), (:push!, :pushfirst!))
     @eval function Base.$f(s::StructVector, rows)

src/utils.jl

@@ -186,3 +186,14 @@ possible internal constructors. `T` should be a concrete type.
     construct = Expr(:new, :T, vars...)
     Expr(:block, assign..., construct)
 end
+
+# Specialize this for types like LazyRow that shouldn't convert
+"""
+    StructArrays.maybe_convert_elt(T, x)
+
+Element conversion before assignment in a StructArray.
+By default, this calls `convert(T, x)`; however, you can specialize it for other types.
+"""
+maybe_convert_elt(::Type{T}, vals) where T = convert(T, vals)
+maybe_convert_elt(::Type{T}, vals::Tuple) where T = T <: Tuple ? convert(T, vals) : vals  # assignment of fields by position
+maybe_convert_elt(::Type{T}, vals::NamedTuple) where T = T<:NamedTuple ? convert(T, vals) : vals # assignment of fields by name

test/runtests.jl

@@ -9,17 +9,72 @@ using Adapt: adapt, Adapt
 using Test
 
 using Documenter: doctest
-if Base.VERSION >= v"1.6"
+if Base.VERSION >= v"1.6" && Int === Int64
     doctest(StructArrays)
 end
 
+# Most types should not be viewed as equivalent merely
+# because they have the same field names. (Exception:
+# NamedTuples are distinguished only by field names, so they
+# are treated as equivalent to any struct with the same
+# field names.) To test proper behavior, define two types
+# that are "structurally" equivalent...
+struct Meters
+    x::Float64
+end
+struct Millimeters
+    x::Float64
+end
+# ...but not naively transferrable
+Base.convert(::Type{Meters}, x::Millimeters) = Meters(x.x/1000)
+Base.convert(::Type{Millimeters}, x::Meters) = Millimeters(x.x*1000)
+
+
 @testset "index" begin
     a, b = [1 2; 3 4], [4 5; 6 7]
     t = StructArray((a = a, b = b))
     @test (@inferred t[2,2]) == (a = 4, b = 7)
     @test (@inferred t[2,1:2]) == StructArray((a = [3, 4], b = [6, 7]))
     @test_throws BoundsError t[3,3]
     @test (@inferred view(t, 2, 1:2)) == StructArray((a = view(a, 2, 1:2), b = view(b, 2, 1:2)))
+
+    # Element type conversion (issue #216)
+    x = StructArray{Complex{Int}}((a, b))
+    x[1,1] = 10
+    x[2,2] = 20im
+    @test x[1,1] === 10 + 0im
+    @test x[2,2] === 0 + 20im
+
+    # Test that explicit `setindex!` returns the entire array
+    # (Julia's parser ensures that chained assignment returns the value)
+    @test setindex!(x, 22, 3) === x
+end
+
+@testset "eltype conversion" begin
+    v = StructArray{Complex{Int}}(([1,2,3], [4,5,6]))
+    @test append!(v, [7, 8]) == [1+4im, 2+5im, 3+6im, 7+0im, 8+0im]
+    push!(v, (im=12, re=11))  # NamedTuples support field assignment by name
+    @test v[end] === 11 + 12im
+    v[end] = (re=9, im=10)
+    @test v[end] === 9 + 10im
+
+    # For some eltypes, the structarray is "nameless" and we can use regular Tuples
+    v = StructArray([SVector(true, false), SVector(false, false)])
+    v[end] = (true, true)
+    @test v[end] === SVector(true, true)
+    push!(v, (false, false))
+    @test v[end] === SVector(false, false)
+
+    z = StructArray{Meters}(undef, 0)
+    push!(z, Millimeters(1100))
+    @test length(z) == 1
+    @test z[1] === Meters(1.1)
+    append!(z, [Millimeters(1200)])
+    @test z[2] === Meters(1.2)
+    append!(z, StructArray{Millimeters}(([1500.0],)))
+    @test z[3] === Meters(1.5)
+    insert!(z, 3, Millimeters(2000))
+    @test z[3] === Meters(2.0)
 end
 
 @testset "components" begin