implement three-argument similar method (from #94) (#218)

* Implement three-argument `similar` methods

Array wrappers such as `OffsetArray` assume that the three argument
version exists for its interior arrays.

* Apply suggestions from code review

Co-authored-by: Lucas C Wilcox <lucas@swirlee.com>

* simplify similar signatures

* support and test similar and reshape with offsets

Co-authored-by: Keno Fischer <keno@alumni.harvard.edu>
Co-authored-by: Lucas C Wilcox <lucas@swirlee.com>

Author: 3 people

src/structarray.jl

@@ -167,22 +167,26 @@ function Base.IndexStyle(::Type{S}) where {S<:StructArray}
     index_type(S) === Int ? IndexLinear() : IndexCartesian()
 end
 
-function _undef_array(::Type{T}, sz; unwrap::F = alwaysfalse) where {T, F}
+function undef_array(::Type{T}, sz; unwrap::F = alwaysfalse) where {T, F}
     if unwrap(T)
         return StructArray{T}(undef, sz; unwrap = unwrap)
     else
         return Array{T}(undef, sz)
     end
 end
 
-function _similar(v::AbstractArray, ::Type{Z}; unwrap::F = alwaysfalse) where {Z, F}
+function similar_array(v::AbstractArray, ::Type{Z}; unwrap::F = alwaysfalse) where {Z, F}
     if unwrap(Z)
-        return buildfromschema(typ -> _similar(v, typ; unwrap = unwrap), Z)
+        return buildfromschema(typ -> similar_array(v, typ; unwrap = unwrap), Z)
     else
         return similar(v, Z)
     end
 end
 
+function similar_structarray(v::AbstractArray, ::Type{Z}; unwrap::F = alwaysfalse) where {Z, F}
+    buildfromschema(typ -> similar_array(v, typ; unwrap = unwrap), Z)
+end
+
 """
     StructArray{T}(undef, dims; unwrap=T->false)
 
@@ -204,14 +208,10 @@ julia> StructArray{ComplexF64}(undef, (2,3))
 StructArray(::Base.UndefInitializer, sz::Dims)
 
 function StructArray{T}(::Base.UndefInitializer, sz::Dims; unwrap::F = alwaysfalse) where {T, F}
-    buildfromschema(typ -> _undef_array(typ, sz; unwrap = unwrap), T)
+    buildfromschema(typ -> undef_array(typ, sz; unwrap = unwrap), T)
 end
 StructArray{T}(u::Base.UndefInitializer, d::Integer...; unwrap::F = alwaysfalse) where {T, F} = StructArray{T}(u, convert(Dims, d); unwrap = unwrap)
 
-function similar_structarray(v::AbstractArray, ::Type{Z}; unwrap::F = alwaysfalse) where {Z, F}
-    buildfromschema(typ -> _similar(v, typ; unwrap = unwrap), Z)
-end
-
 """
     StructArray(A; unwrap = T->false)
 
@@ -276,14 +276,34 @@ Base.convert(::Type{StructArray}, v::StructArray) = v
 Base.convert(::Type{StructVector}, v::AbstractVector) = StructVector(v)
 Base.convert(::Type{StructVector}, v::StructVector) = v
 
-function Base.similar(::Type{<:StructArray{T, <:Any, C}}, sz::Dims) where {T, C}
-    buildfromschema(typ -> similar(typ, sz), T, C)
+# Mimic OffsetArrays signatures
+const OffsetAxisKnownLength = Union{Integer, AbstractUnitRange}
+const OffsetAxis = Union{OffsetAxisKnownLength, Colon}
+
+const OffsetShapeKnownLength = Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}}
+const OffsetShape = Tuple{OffsetAxis,Vararg{OffsetAxis}}
+
+# Helper function to avoid adding too many dispatches to `Base.similar`
+function _similar(s::StructArray{T}, ::Type{T}, sz) where {T}
+    return StructArray{T}(map(typ -> similar(typ, sz), components(s)))
 end
 
-Base.similar(s::StructArray, sz::Base.DimOrInd...) = similar(s, Base.to_shape(sz))
-Base.similar(s::StructArray) = similar(s, Base.to_shape(axes(s)))
-function Base.similar(s::StructArray{T}, sz::Tuple) where {T}
-    StructArray{T}(map(typ -> similar(typ, sz), components(s)))
+function _similar(s::StructArray{T}, S::Type, sz) where {T}
+    # If not specified, we don't really know what kind of array to use for each
+    # interior type, so we just pick the first one arbitrarily. If users need
+    # something else, they need to be more specific.
+    c1 = first(components(s))
+    return isnonemptystructtype(S) ? buildfromschema(typ -> similar(c1, typ, sz), S) : similar(c1, S, sz)
+end
+
+for type in (:Dims, :OffsetShapeKnownLength)
+    @eval function Base.similar(::Type{<:StructArray{T, N, C}}, sz::$(type)) where {T, N, C}
+        return buildfromschema(typ -> similar(typ, sz), T, C)
+    end
+
+    @eval function Base.similar(s::StructArray, S::Type, sz::$(type))
+        return _similar(s, S, sz)
+    end
 end
 
 @deprecate fieldarrays(x) StructArrays.components(x)
@@ -437,8 +457,10 @@ end
 
 Base.copy(s::StructArray{T}) where {T} = StructArray{T}(map(copy, components(s)))
 
-function Base.reshape(s::StructArray{T}, d::Dims) where {T}
-    StructArray{T}(map(x -> reshape(x, d), components(s)))
+for type in (:Dims, :OffsetShape)
+    @eval function Base.reshape(s::StructArray{T}, d::$(type)) where {T}
+        StructArray{T}(map(x -> reshape(x, d), components(s)))
+    end
 end
 
 function showfields(io::IO, fields::NTuple{N, Any}) where N

test/runtests.jl

@@ -1,6 +1,6 @@
 using StructArrays
 using StructArrays: staticschema, iscompatible, _promote_typejoin, append!!
-using OffsetArrays: OffsetArray
+using OffsetArrays: OffsetArray, OffsetVector, OffsetMatrix
 using StaticArrays
 import Tables, PooledArrays, WeakRefStrings
 using TypedTables: Table
@@ -318,13 +318,57 @@ end
     s = similar(t)
     @test eltype(s) == NamedTuple{(:a, :b), Tuple{Float64, Bool}}
     @test size(s) == (10,)
+    @test s isa StructArray
+
     t = StructArray(a = rand(10, 2), b = rand(Bool, 10, 2))
     s = similar(t, 3, 5)
     @test eltype(s) == NamedTuple{(:a, :b), Tuple{Float64, Bool}}
     @test size(s) == (3, 5)
+    @test s isa StructArray
+
     s = similar(t, (3, 5))
     @test eltype(s) == NamedTuple{(:a, :b), Tuple{Float64, Bool}}
     @test size(s) == (3, 5)
+    @test s isa StructArray
+
+    s = similar(t, (0:2, 5))
+    @test eltype(s) == NamedTuple{(:a, :b), Tuple{Float64, Bool}}
+    @test axes(s) == (0:2, 1:5)
+    @test s isa StructArray
+    @test s.a isa OffsetArray
+    @test s.b isa OffsetArray
+
+    s = similar(t, ComplexF64, 10)
+    @test s isa StructArray{ComplexF64, 1, NamedTuple{(:re, :im), Tuple{Vector{Float64}, Vector{Float64}}}}
+    @test size(s) == (10,)
+
+    s = similar(t, ComplexF64, 0:9)
+    VectorType = OffsetVector{Float64, Vector{Float64}}
+    @test s isa StructArray{ComplexF64, 1, NamedTuple{(:re, :im), Tuple{VectorType, VectorType}}}
+    @test axes(s) == (0:9,)
+
+    s = similar(t, Float32, 2, 2)
+    @test s isa Matrix{Float32}
+    @test size(s) == (2, 2)
+
+    s = similar(t, Float32, 0:1, 2)
+    @test s isa OffsetMatrix{Float32, Matrix{Float32}}
+    @test axes(s) == (0:1, 1:2)
+end
+
+@testset "similar type" begin
+    t = StructArray(a = rand(10), b = rand(10))
+    T = typeof(t)
+    s = similar(T, 3)
+    @test typeof(s) == typeof(t)
+    @test size(s) == (3,)
+
+    s = similar(T, 0:2)
+    @test axes(s) == (0:2,)
+    @test s isa StructArray{NamedTuple{(:a, :b), Tuple{Float64, Float64}}}
+    VectorType = OffsetVector{Float64, Vector{Float64}}
+    @test s.a isa VectorType
+    @test s.b isa VectorType
 end
 
 @testset "empty" begin
@@ -803,6 +847,10 @@ end
     rs = reshape(s, (2, 2))
     @test rs.a == [1 3; 2 4]
     @test rs.b == ["a" "c"; "b" "d"]
+
+    rs = reshape(s, (0:1, :))
+    @test rs.a == OffsetArray([1 3; 2 4], (-1, 0))
+    @test rs.b == OffsetArray(["a" "c"; "b" "d"], (-1, 0))
 end
 
 @testset "lazy" begin
@@ -1091,3 +1139,16 @@ end
     C = map(zero, NamedTuple{(:a, :b, :c)}(map(zero, fieldtypes(types))))
     @test A === C
 end
+
+@testset "OffsetArray zero" begin
+    s = StructArray{ComplexF64}((rand(2), rand(2)))
+    soff = OffsetArray(s, 0:1)
+    @test isa(parent(zero(soff)), StructArray)
+end
+
+# issue #230
+@testset "StaticArray zero" begin
+    u = StructArray([SVector(1.0)])
+    @test zero(u) == StructArray([SVector(0.0)])
+    @test typeof(zero(u)) == typeof(StructArray([SVector(0.0)]))
+end