index_labels method and IndexLabel type indexing

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "6.0.22"
+version = "6.0.23"
 
 [deps]
 ArrayInterfaceCore = "30b0a656-2188-435a-8636-2ec0e6a096e2"

lib/ArrayInterfaceCore/Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterfaceCore"
 uuid = "30b0a656-2188-435a-8636-2ec0e6a096e2"
-version = "0.1.17"
+version = "0.1.18"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

lib/ArrayInterfaceCore/src/ArrayInterfaceCore.jl

@@ -79,7 +79,7 @@ _fl1(x::Tuple, x1) = ((length(x) - 1, 0),)
 Returns the parent array that type `T` wraps.
 """
 parent_type(x) = parent_type(typeof(x))
-parent_type(::Type{Symmetric{T,S}}) where {T,S} = S
+parent_type(@nospecialize T::Type{<:Union{Symmetric,Hermitian}}) = fieldtype(T, :data)
 parent_type(::Type{<:AbstractTriangular{T,S}}) where {T,S} = S
 parent_type(@nospecialize T::Type{<:PermutedDimsArray}) = fieldtype(T, :parent)
 parent_type(@nospecialize T::Type{<:Adjoint}) = fieldtype(T, :parent)
@@ -616,6 +616,15 @@ Base.@propagate_inbounds function Base.getindex(ind::TridiagonalIndex, i::Int)
     end
 end
 
+"""
+    Label(label)
+
+A type that clearly communicates that `label` refers to a key-index mapping.
+"""
+struct Label{L} <: ArrayIndex{0}
+    label::L
+end
+
 _cartesian_index(i::Tuple{Vararg{Int}}) = CartesianIndex(i)
 _cartesian_index(::Any) = nothing
 
@@ -715,6 +724,7 @@ julia> ArrayInterfaceCore.ndims_index([CartesianIndex(1, 2), CartesianIndex(1, 3
 ndims_index(::Type{<:Base.AbstractCartesianIndex{N}}) where {N} = N
 # preserve CartesianIndices{0} as they consume a dimension.
 ndims_index(::Type{CartesianIndices{0,Tuple{}}}) = 1
+ndims_index(@nospecialize T::Type{<:Union{Number,Label,Symbol,AbstractString,AbstractChar}}) = 1
 ndims_index(@nospecialize T::Type{<:AbstractArray{Bool}}) = ndims(T)
 ndims_index(@nospecialize T::Type{<:AbstractArray}) = ndims_index(eltype(T))
 ndims_index(@nospecialize T::Type{<:Base.LogicalIndex}) = ndims(fieldtype(T, :mask))
@@ -742,7 +752,7 @@ julia> ndims(CartesianIndices((2,2))[[CartesianIndex(1, 1), CartesianIndex(1, 2)
 ndims_shape(T::DataType) = ndims_index(T)
 ndims_shape(::Type{Colon}) = 1
 ndims_shape(@nospecialize T::Type{<:CartesianIndices}) = ndims(T)
-ndims_shape(@nospecialize T::Type{<:Union{Number,Base.AbstractCartesianIndex}}) = 0
+ndims_shape(@nospecialize T::Type{<:Union{Number,Base.AbstractCartesianIndex,Label,Symbol,AbstractString,AbstractChar}}) = 0
 ndims_shape(@nospecialize T::Type{<:AbstractArray{Bool}}) = 1
 ndims_shape(@nospecialize T::Type{<:AbstractArray}) = ndims(T)
 ndims_shape(x) = ndims_shape(typeof(x))

src/ArrayInterface.jl

@@ -6,7 +6,7 @@ import ArrayInterfaceCore: allowed_getindex, allowed_setindex!, aos_to_soa, buff
     issingular, isstructured, matrix_colors, restructure, lu_instance,
     safevec, zeromatrix, ColoringAlgorithm, fast_scalar_indexing, parameterless_type,
     ndims_index, ndims_shape, is_splat_index, is_forwarding_wrapper, IndicesInfo,
-    map_tuple_type, flatten_tuples, GetIndex, SetIndex!, defines_strides,
+    map_tuple_type, flatten_tuples, GetIndex, SetIndex!, Label, defines_strides,
     stride_preserving_index
 
 # ArrayIndex subtypes and methods

src/axes.jl

@@ -90,7 +90,6 @@ function axes_types(::Type{A}) where {T,N,S,A<:Base.ReshapedReinterpretArray{T,N
     end
 end
 
-
 # FUTURE NOTE: we avoid  `SOneTo(1)` when `axis(A, dim::Int)``. This is inended to decreases
 # breaking changes for this adopting this method to situations where they clearly benefit
 # from the propagation of static axes. This creates the somewhat awkward situation of
@@ -113,7 +112,7 @@ axes(A::ReshapedArray) = Base.axes(A)
 @inline function axes(x::Union{MatAdjTrans,PermutedDimsArray})
     map(GetIndex{false}(axes(parent(x))), to_parent_dims(x))
 end
-axes(A::VecAdjTrans) = (SOneTo{1}(), axes(parent(A), 1))
+axes(A::VecAdjTrans) = (SOneTo{1}(), getfield(axes(parent(A)), 1))
 
 @inline axes(x::SubArray) = flatten_tuples(map(Base.Fix1(_sub_axes, x), sub_axes_map(typeof(x))))
 @inline _sub_axes(x::SubArray, axis::SOneTo) = axis
@@ -248,3 +247,64 @@ lazy_axes(x::AbstractRange, ::StaticInt{1}) = Base.axes1(x)
 lazy_axes(x, ::Colon) = LazyAxis{:}(x)
 lazy_axes(x, ::StaticInt{dim}) where {dim} = ndims(x) < dim ? SOneTo{1}() : LazyAxis{dim}(x)
 @inline lazy_axes(x, dims::Tuple) = map(Base.Fix1(lazy_axes, x), dims)
+
+"""
+    index_labels(x)
+    index_labels(x, dim)
+
+Returns a tuple of labels assigned to each axis or a collection of labels corresponding to
+axis `dim` of `x`. Default is to simply return `map(keys, axes(x))`.
+"""
+index_labels(x, dim)  = index_labels(x, to_dims(x, dim))
+index_labels(@nospecialize x::Number) = ()
+@inline index_labels(x, d::CanonicalInt) = d > ndims(x) ? keys(axes(x, d)) : index_labels(x)[d]
+@inline index_labels(x) = is_forwarding_wrapper(x) ? index_labels(parent(x)) : map(keys, axes(x))
+function index_labels(x::Union{MatAdjTrans,PermutedDimsArray})
+    map(GetIndex{false}(index_labels(parent(x))), to_parent_dims(x))
+end
+index_labels(x::VecAdjTrans) = (keys(SOneTo{1}()), getfield(index_labels(parent(x)), 1))
+index_labels(x::SubArray) = _sub_index_labels(parent(x), x.indices, IndicesInfo(x))
+function _sub_index_labels(x::AbstractArray, inds::Tuple, ::IndicesInfo{N,pdims,cdims}) where {N,pdims,cdims}
+    labels = index_labels(x)
+    flatten_tuples(ntuple(Val{nfields(pdims)}()) do i
+        pdim_i = getfield(pdims, i)
+        cdim_i = getfield(cdims, i)
+        index = getfield(inds, i)
+        if pdim_i isa Tuple || cdim_i isa Tuple # no direct mapping to parent axes
+            index_labels(index)
+        elseif cdim_i === 0  # integer indexing drops axes
+            ()
+        elseif pdim_i === 0  # trailing dimension
+            LinearIndices((SOneTo{1}(),))
+        elseif index isa Base.Slice  # index into labels where there is direct mapping to parent axis
+            (getfield(labels, pdim_i),)
+        else
+            (@inbounds(getfield(labels, pdim_i)[index]),)
+        end
+    end)
+end
+index_labels(x::Union{LinearIndices,CartesianIndices}) = map(first ∘ index_labels, axes(x))
+index_labels(x::Union{Symmetric,Hermitian}) = index_labels(parent(x))
+index_labels(x::LazyAxis{N,P}) where {N,P} = (index_labels(getfield(x, :parent), StaticInt(N)),)
+@inline @inline function index_labels(x::Base.NonReshapedReinterpretArray{T,N,S}) where {T,N,S}
+    if sizeof(T) === sizeof(S)
+        return index_labels(parent(x))
+    else
+        return flatten_tuples((keys(StaticInt(1):size(x, 1)), Base.tail(index_labels(parent(x)))))
+    end
+end
+function index_labels(x::Base.ReshapedReinterpretArray{T,N,S}) where {T,N,S}
+    _reinterpret_index_labels(div(StaticInt(sizeof(S)), StaticInt(sizeof(T))), x)
+end
+@inline function _reinterpreted_fieldnames(@nospecialize T::Type{<:Base.ReshapedReinterpretArray})
+    S = eltype(parent_type(T))
+    isstructtype(S) ? fieldnames(S) : ()
+end
+function _reinterpret_index_labels(s::StaticInt{N}, x::Base.ReshapedReinterpretArray) where {N}
+    __reinterpret_index_labels(s, _reinterpreted_fieldnames(typeof(x)), index_labels(parent(x)))
+end
+@inline function __reinterpret_index_labels(::StaticInt{N}, fields::NTuple{M,Symbol}, ks::Tuple) where {N,M}
+    N === M ? (fields, ks...,) : (LinearIndices((SOneTo{N}(),)), ks...,)
+end
+_reinterpret_index_labels(::StaticInt{1}, x::Base.ReshapedReinterpretArray) = index_labels(parent(x))
+_reinterpret_index_labels(::StaticInt{0}, x::Base.ReshapedReinterpretArray) = tail(index_labels(parent(x)))

src/dimensions.jl

@@ -142,6 +142,9 @@ end
         return ntuple(Compat.Returns(:_), StaticInt(ndims(T)))
     end
 end
+known_dimnames(::Type{<:LazyAxis{:,P}}) where {P} = (first(known_dimnames(P)),)
+known_dimnames(::Type{<:LazyAxis{N,P}}) where {N,P} = (getfield(known_dimnames(P), N),)
+
 @inline function known_dimnames(::Type{T}) where {T}
     if is_forwarding_wrapper(T)
         return known_dimnames(parent_type(T))
@@ -201,6 +204,8 @@ end
         return ntuple(Compat.Returns(static(:_)), StaticInt(ndims(x)))
     end
 end
+dimnames(x::LazyAxis{:,P}) where {P} = (first(dimnames(getfield(x, :parent))),)
+dimnames(x::LazyAxis{N,P}) where {N,P} = (getfield(dimnames(getfield(x, :parent)), N),)
 @inline function dimnames(x::X) where {X}
     if is_forwarding_wrapper(X)
         return dimnames(parent(x))

src/indexing.jl

@@ -170,14 +170,33 @@ end
 @inline function to_index(x, i::Base.Fix2{typeof(>=),<:Union{Base.BitInteger,StaticInt}})
     max(canonicalize(i.x), static_first(x)):static_last(x)
 end
+@inline function to_index(x, i::Base.Fix2{<:Union{typeof(>),typeof(>=),typeof(<=),typeof(<),typeof(isless)},<:Label})
+    findall(i.f(i.x.label), first(index_labels(x)))
+end
 @inline function to_index(x, i::Base.Fix2{typeof(>),<:Union{Base.BitInteger,StaticInt}})
     max(_add1(canonicalize(i.x)), static_first(x)):static_last(x)
 end
-# integer indexing
-to_index(x, i::AbstractArray{<:Integer}) = i
+to_index(x, i::AbstractArray{<:Union{Base.BitInteger,StaticInt}}) = i
 to_index(x, @nospecialize(i::StaticInt)) = i
 to_index(x, i::Integer) = Int(i)
 @inline to_index(x, i) = to_index(IndexStyle(x), x, i)
+# key indexing
+function to_index(x, i::Label)
+    index = findfirst(==(getfield(i, :label)), first(index_labels(x)))
+    # delay throwing bounds-error if we didn't find label
+    index === nothing ? offset1(x) - 1 : index
+end
+function to_index(x, i::Union{Symbol,AbstractString,AbstractChar,Number})
+    index = findfirst(==(i), getfield(index_labels(x), 1))
+    index === nothing ? offset1(x) - 1 : index
+end
+# TODO there's probably a more efficient way of doing this
+to_index(x, ks::AbstractArray{<:Label}) = [to_index(x, k) for k in ks]
+function to_index(x, ks::AbstractArray{<:Union{Symbol,AbstractString,AbstractChar,Number}})
+    [to_index(x, k) for k in ks]
+end
+
+# integer indexing
 function to_index(S::IndexStyle, x, i)
     throw(ArgumentError(
         "invalid index: $S does not support indices of type $(typeof(i)) for instances of type $(typeof(x))."

test/axes.jl

@@ -104,3 +104,38 @@ if isdefined(Base, :ReshapedReinterpretArray)
     @inferred(ArrayInterface.axes(fa)) isa ArrayInterface.axes_types(fa)
   end
 end
+
+@testset "index_labels" begin
+    colors = LabelledArray([(R = rand(), G = rand(), B = rand()) for i ∈ 1:100], (range(-10, 10, length=100),));
+    caxis = ArrayInterface.LazyAxis{1}(colors);
+    colormat = reinterpret(reshape, Float64, colors);
+    cmat_view1 = view(colormat, :, 4);
+    cmat_view2 = view(colormat, :, 4:7);
+    cmat_view3 = view(colormat, 2:3,:);
+    absym_abstr = LabelledArray(rand(Int64, 2,2), ([:a, :b], ["a", "b"],));
+
+    @test @inferred(ArrayInterface.index_labels(colors)) == (range(-10, 10, length=100),)
+    @test @inferred(ArrayInterface.index_labels(caxis)) == (range(-10, 10, length=100),)
+    @test ArrayInterface.index_labels(view(colors, :, :), 2) == keys(static(1):static(1))
+    @test @inferred(ArrayInterface.index_labels(LinearIndices((caxis,)))) == (range(-10, 10, length=100),)
+    @test @inferred(ArrayInterface.index_labels(colormat)) == ((:R, :G, :B), range(-10, 10, length=100))
+    @test @inferred(ArrayInterface.index_labels(colormat')) == (range(-10, 10, length=100), (:R, :G, :B))
+    @test @inferred(ArrayInterface.index_labels(cmat_view1)) == ((:R, :G, :B),)
+    @test @inferred((ArrayInterface.index_labels(cmat_view2))) == ((:R, :G, :B), -9.393939393939394:0.20202020202020202:-8.787878787878787)
+    @test @inferred((ArrayInterface.index_labels(view(colormat, 1, :)'))) == (keys(static(1):static(1)), range(-10, 10, length=100))
+    # can't infer this b/c tuple is being indexed by range
+    @test ArrayInterface.index_labels(cmat_view3) == ((:G, :B), -10.0:0.20202020202020202:10.0)
+    @test @inferred(ArrayInterface.index_labels(Symmetric(view(colormat, :, 1:3)))) == ((:R, :G, :B), -10.0:0.20202020202020202:-9.595959595959595)
+
+    @test @inferred(ArrayInterface.index_labels(reinterpret(Int8, absym_abstr))) == (keys(Base.OneTo(16)), ["a", "b"])
+    @test @inferred(ArrayInterface.index_labels(reinterpret(reshape, Int8, absym_abstr))) == (keys(static(1):static(8)), [:a, :b], ["a", "b"])
+    @test @inferred(ArrayInterface.index_labels(reinterpret(reshape, Int64, LabelledArray(rand(Int32, 2,2), ([:a, :b], ["a", "b"],))))) == (["a", "b"],)
+    @test @inferred(ArrayInterface.index_labels(reinterpret(reshape, Float64, LabelledArray(rand(Int64, 2,2), ([:a, :b], ["a", "b"],))))) == ([:a, :b], ["a", "b"],)
+    @test @inferred(ArrayInterface.index_labels(reinterpret(Float64, absym_abstr))) == ([:a, :b], ["a", "b"],)
+
+    @test @inferred(ArrayInterface.getindex(colormat, :R, :)) == colormat[1, :]
+    @test @inferred(ArrayInterface.getindex(cmat_view1, :R)) == cmat_view1[1]
+    @test @inferred(ArrayInterface.getindex(colormat, :,ArrayInterface.Label(-9.595959595959595))) == colormat[:, 3]
+    @test @inferred(ArrayInterface.getindex(colormat, :,<=(ArrayInterface.Label(-9.595959595959595)))) == colormat[:, 1:3]
+    @test @inferred(ArrayInterface.getindex(absym_abstr, :, ["a"])) == absym_abstr[:,[1]]
+end

test/dimensions.jl

@@ -29,13 +29,18 @@ end
     r4 = reinterpret(reshape, Float64, x)
     w = Wrapper(x)
     dnums = ntuple(+, length(d))
+    lz2 = ArrayInterface.lazy_axes(x)[2]
+    lzslice = ArrayInterface.LazyAxis{:}(x)
+
     @test @inferred(ArrayInterface.has_dimnames(x)) == true
     @test @inferred(ArrayInterface.has_dimnames(z)) == true
     @test @inferred(ArrayInterface.has_dimnames(ones(2, 2))) == false
     @test @inferred(ArrayInterface.has_dimnames(Array{Int,2})) == false
     @test @inferred(ArrayInterface.has_dimnames(typeof(x))) == true
     @test @inferred(ArrayInterface.has_dimnames(typeof(view(x, :, 1, :)))) == true
     @test @inferred(ArrayInterface.dimnames(x)) === d
+    @test @inferred(ArrayInterface.dimnames(lz2)) === (static(:y),)
+    @test @inferred(ArrayInterface.dimnames(lzslice)) === (static(:x),)
     @test @inferred(ArrayInterface.dimnames(w)) === d
     @test @inferred(ArrayInterface.dimnames(r1)) === d
     @test @inferred(ArrayInterface.dimnames(r2)) === (static(:_), d...)
@@ -64,6 +69,8 @@ end
     # multidmensional indices
     @test @inferred(ArrayInterface.known_dimnames(view(x, ones(Int, 2, 2), 1))) === (:_, :_)
     @test @inferred(ArrayInterface.known_dimnames(view(x, [CartesianIndex(1,1), CartesianIndex(1,1)]))) === (:_,)
+    @test @inferred(ArrayInterface.known_dimnames(lz2)) === (:y,)
+    @test @inferred(ArrayInterface.known_dimnames(lzslice)) === (:x,)
 
     @test @inferred(ArrayInterface.known_dimnames(z)) === (nothing, :y)
     @test @inferred(ArrayInterface.known_dimnames(reshape(x, (1, 4)))) === (:x, :y)

test/setup.jl

@@ -34,7 +34,16 @@ function ArrayInterface.known_dimnames(::Type{T}) where {L,T<:NamedDimsWrapper{L
     ArrayInterface.Static.known(L)
 end
 
-Base.parent(x::NamedDimsWrapper) = x.parent
+struct LabelledArray{T,N,P<:AbstractArray{T,N},L} <: ArrayInterface.AbstractArray2{T,N}
+    parent::P
+    labels::L
+
+    LabelledArray(p::P, labels::L) where {P,L} = new{eltype(P),ndims(p),P,L}(p, labels)
+end
+ArrayInterface.is_forwarding_wrapper(::Type{<:LabelledArray}) = true
+Base.parent(x::LabelledArray) = getfield(x, :parent)
+ArrayInterface.parent_type(::Type{T}) where {P,T<:LabelledArray{<:Any,<:Any,P}} = P
+ArrayInterface.index_labels(x::LabelledArray) = getfield(x, :labels)
 
 # Dummy array type with undetermined contiguity properties
 struct DummyZeros{T,N} <: AbstractArray{T,N}