Merge pull request #78 from piever/pv/cart

Add type parameter for index style in StructArray

Author: Pietro Vertechi

src/lazy.jl

@@ -1,5 +1,5 @@
 struct LazyRow{T, N, C, I}
-    columns::StructArray{T, N, C} # a `Columns` object
+    columns::StructArray{T, N, C, I} # a `Columns` object
     index::I
 end
 
@@ -30,19 +30,17 @@ iscompatible(::Type{<:LazyRow{S}}, ::Type{StructArray{T, N, C}}) where {S, T, N,
 (s::ArrayInitializer)(::Type{<:LazyRow{T}}, d) where {T} = buildfromschema(typ -> s(typ, d), T)
 
 struct LazyRows{T, N, C, I} <: AbstractArray{LazyRow{T, N, C, I}, N}
-    columns::StructArray{T, N, C}
+    columns::StructArray{T, N, C, I}
 end
-LazyRows(s::S) where {S<:StructArray} = LazyRows(IndexStyle(S), s)
-LazyRows(::IndexLinear, s::StructArray{T, N, C}) where {T, N, C} = LazyRows{T, N, C, Int}(s)
-LazyRows(::IndexCartesian, s::StructArray{T, N, C}) where {T, N, C} = LazyRows{T, N, C, CartesianIndex{N}}(s)
 Base.parent(v::LazyRows) = getfield(v, 1)
 fieldarrays(v::LazyRows) = fieldarrays(parent(v))
 
 Base.size(v::LazyRows) = size(parent(v))
-Base.getindex(v::LazyRows{<:Any, <:Any, <:Any, <:Integer}, i::Integer) = LazyRow(parent(v), i)
-Base.getindex(v::LazyRows{<:Any, <:Any, <:Any, <:CartesianIndex}, i::Integer...) = LazyRow(parent(v), CartesianIndex(i))
+Base.getindex(v::LazyRows{<:Any, <:Any, <:Any, Int}, i::Int) = LazyRow(parent(v), i)
+Base.getindex(v::LazyRows{<:Any, <:Any, <:Any, CartesianIndex{N}}, i::Vararg{Int, N}) where {N} = LazyRow(parent(v), CartesianIndex(i))
 
-Base.IndexStyle(::Type{<:LazyRows{<:Any, <:Any, <:Any, <:Integer}}) = IndexLinear()
+_best_index(::Type{LazyRows{T, N, C, I}}) where {T, N, C, I} = I
+Base.IndexStyle(::Type{L}) where {L<:LazyRows} = _indexstyle(_best_index(L))
 
 function Base.showarg(io::IO, s::LazyRows{T}, toplevel) where T
     print(io, "LazyRows")

src/structarray.jl

@@ -3,7 +3,7 @@ A type that stores an array of structures as a structure of arrays.
 # Fields:
 - `fieldarrays`: a (named) tuple of arrays. Also `fieldarrays(x)`
 """
-struct StructArray{T, N, C<:Tup} <: AbstractArray{T, N}
+struct StructArray{T, N, C<:Tup, I} <: AbstractArray{T, N}
     fieldarrays::C
 
     function StructArray{T, N, C}(c) where {T, N, C<:Tup}
@@ -14,10 +14,18 @@ struct StructArray{T, N, C<:Tup} <: AbstractArray{T, N}
                 axes(c[i]) == ax || error("all field arrays must have same shape")
             end
         end
-        new{T, N, C}(c)
+        new{T, N, C, _best_index(c...)}(c)
     end
 end
 
+_best_index() = Int
+_best_index(col::AbstractArray, cols::AbstractArray...) = _best_index(IndexStyle(col, cols...), col)
+_best_index(::IndexLinear, ::AbstractArray) = Int
+_best_index(::IndexCartesian, ::AbstractArray{T, N}) where {T, N} = CartesianIndex{N}
+_best_index(::Type{StructArray{T, N, C, I}}) where {T, N, C, I} = I
+_indexstyle(::Type{Int}) = IndexLinear()
+_indexstyle(::Type{CartesianIndex{N}}) where {N} = IndexCartesian()
+
 _dims(c::Tup) = length(axes(c[1]))
 _dims(c::EmptyTup) = 1
 
@@ -37,17 +45,11 @@ _structarray(args::Tuple, names) = _structarray(args, Tuple(names))
 _structarray(args::Tuple, ::Tuple) = _structarray(args, nothing)
 _structarray(args::NTuple{N, Any}, names::NTuple{N, Symbol}) where {N} = StructArray(NamedTuple{names}(args))
 
-const StructVector{T, C<:Tup} = StructArray{T, 1, C}
+const StructVector{T, C<:Tup, I} = StructArray{T, 1, C, I}
 StructVector{T}(args...; kwargs...) where {T} = StructArray{T}(args...; kwargs...)
 StructVector(args...; kwargs...) = StructArray(args...; kwargs...)
 
-_indexstyle(::Type{Tuple{}}) = IndexStyle(Union{})
-_indexstyle(::Type{T}) where {T<:Tuple} = IndexStyle(IndexStyle(tuple_type_head(T)), _indexstyle(tuple_type_tail(T)))
-_indexstyle(::Type{NamedTuple{names, types}}) where {names, types} = _indexstyle(types)
-
-function Base.IndexStyle(::Type{StructArray{T, N, C}}) where {T, N, C}
-    _indexstyle(C)
-end
+Base.IndexStyle(::Type{S}) where {S<:StructArray} = _indexstyle(_best_index(S))
 
 _undef_array(::Type{T}, sz; unwrap = t -> false) where {T} = unwrap(T) ? StructArray{T}(undef, sz; unwrap = unwrap) : Array{T}(undef, sz)
 
@@ -124,22 +126,34 @@ function get_ith(cols::NTuple{N, Any}, I...) where N
     end
 end
 
-Base.@propagate_inbounds function Base.getindex(x::StructArray{T, N, C}, I::Int...) where {T, N, C}
+Base.@propagate_inbounds function Base.getindex(x::StructArray{T, <:Any, <:Any, CartesianIndex{N}}, I::Vararg{Int, N}) where {T, N}
     cols = fieldarrays(x)
     @boundscheck checkbounds(x, I...)
     return createinstance(T, get_ith(cols, I...)...)
 end
 
+Base.@propagate_inbounds function Base.getindex(x::StructArray{T, <:Any, <:Any, Int}, I::Int) where {T}
+    cols = fieldarrays(x)
+    @boundscheck checkbounds(x, I)
+    return createinstance(T, get_ith(cols, I)...)
+end
+
 function Base.view(s::StructArray{T, N, C}, I...) where {T, N, C}
     StructArray{T}(map(v -> view(v, I...), fieldarrays(s)))
 end
 
-Base.@propagate_inbounds function Base.setindex!(s::StructArray, vals, I::Int...)
+Base.@propagate_inbounds function Base.setindex!(s::StructArray{<:Any, <:Any, <:Any, CartesianIndex{N}}, vals, I::Vararg{Int, N}) where {N}
     @boundscheck checkbounds(s, I...)
     foreachfield((col, val) -> (@inbounds col[I...] = val), s, vals)
     s
 end
 
+Base.@propagate_inbounds function Base.setindex!(s::StructArray{<:Any, <:Any, <:Any, Int}, vals, I::Int)
+    @boundscheck checkbounds(s, I)
+    foreachfield((col, val) -> (@inbounds col[I] = val), s, vals)
+    s
+end
+
 function Base.push!(s::StructArray, vals)
     foreachfield(push!, s, vals)
     return s