Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 4 commits into from
Jun 17, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 38 additions & 16 deletions src/structarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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)

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
63 changes: 62 additions & 1 deletion test/runtests.jl
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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