[Containers] add slicing for SparseAxisArray (#3031)

Author: odow

docs/src/manual/containers.md

@@ -207,33 +207,32 @@ JuMP.Containers.SparseAxisArray{Tuple{Int64, Int64}, 2, Tuple{Int64, Int64}} wit
 
 The `[indices; condition]` syntax is used:
 ```jldoctest containers_sparse
-julia> x = Containers.@container([i = 1:3, j = [:A, :B]; i > 1 && j == :B], (i, j))
-JuMP.Containers.SparseAxisArray{Tuple{Int64, Symbol}, 2, Tuple{Int64, Symbol}} with 2 entries:
+julia> x = Containers.@container([i = 1:3, j = [:A, :B]; i > 1], (i, j))
+JuMP.Containers.SparseAxisArray{Tuple{Int64, Symbol}, 2, Tuple{Int64, Symbol}} with 4 entries:
+  [2, A]  =  (2, :A)
   [2, B]  =  (2, :B)
+  [3, A]  =  (3, :A)
   [3, B]  =  (3, :B)
 ```
 Here we have the index sets `i = 1:3, j = [:A, :B]`, followed by `;`, and then a
-condition, which evaluates to `true` or `false`: `i > 1 && j == :B`.
+condition, which evaluates to `true` or `false`: `i > 1`.
 
 ### Slicing
 
-```@meta
-# TODO: This is included so we know to update the documentation when this is fixed.
-```
-
-Slicing is not supported.
+Slicing is supported:
 ```jldoctest containers_sparse
-julia> x[:, :B]
-ERROR: ArgumentError: Indexing with `:` is not supported by Containers.SparseAxisArray
-[...]
+julia> y = x[:, :B]
+JuMP.Containers.SparseAxisArray{Tuple{Int64, Symbol}, 1, Tuple{Int64}} with 2 entries:
+  [2]  =  (2, :B)
+  [3]  =  (3, :B)
 ```
 
 ### Looping
 
 Use `eachindex` to loop over the elements:
 ```jldoctest containers_sparse
-julia> for key in eachindex(x)
-           println(x[key])
+julia> for key in eachindex(y)
+           println(y[key])
        end
 (2, :B)
 (3, :B)
@@ -247,10 +246,10 @@ Broadcasting over a SparseAxisArray returns a SparseAxisArray
 julia> swap(x::Tuple) = (last(x), first(x))
 swap (generic function with 1 method)
 
-julia> swap.(x)
-JuMP.Containers.SparseAxisArray{Tuple{Symbol, Int64}, 2, Tuple{Int64, Symbol}} with 2 entries:
-  [2, B]  =  (:B, 2)
-  [3, B]  =  (:B, 3)
+julia> swap.(y)
+JuMP.Containers.SparseAxisArray{Tuple{Symbol, Int64}, 1, Tuple{Int64}} with 2 entries:
+  [2]  =  (:B, 2)
+  [3]  =  (:B, 3)
 ```
 
 ## Forcing the container type

src/Containers/SparseAxisArray.jl

@@ -85,20 +85,6 @@ function Base.haskey(sa::SparseAxisArray{T,1,Tuple{I}}, idx::I) where {T,I}
     return haskey(sa.data, (idx,))
 end
 
-# Error for sa[..., :, ...]
-_colon_error() = nothing
-
-_colon_error(::Any, args...) = _colon_error(args...)
-
-function _colon_error(::Colon, args...)
-    return throw(
-        ArgumentError(
-            "Indexing with `:` is not supported by" *
-            " Containers.SparseAxisArray",
-        ),
-    )
-end
-
 function Base.setindex!(
     d::SparseAxisArray{T,N,K},
     value,
@@ -107,11 +93,16 @@ function Base.setindex!(
     return setindex!(d, value, idx...)
 end
 
-function Base.setindex!(d::SparseAxisArray{T,N}, value, idx...) where {T,N}
+function Base.setindex!(d::SparseAxisArray{T,N,K}, value, idx...) where {T,N,K}
     if length(idx) < N
         throw(BoundsError(d, idx))
+    elseif _sliced_key_type(K, idx...) !== nothing
+        throw(
+            ArgumentError(
+                "Slicing is not when calling setindex! on a SparseAxisArray",
+            ),
+        )
     end
-    _colon_error(idx...)
     return setindex!(d.data, value, idx)
 end
 
@@ -122,14 +113,52 @@ function Base.getindex(
     return getindex(d, idx...)
 end
 
-function Base.getindex(d::SparseAxisArray{T,N}, idx...) where {T,N}
+function Base.getindex(d::SparseAxisArray{T,N,K}, idx...) where {T,N,K}
     if length(idx) < N
         throw(BoundsError(d, idx))
     end
-    _colon_error(idx...)
+    K2 = _sliced_key_type(K, idx...)
+    if K2 !== nothing
+        new_data = Dict{K2,T}(
+            _sliced_key(k, idx) => v for (k, v) in d.data if _filter(k, idx)
+        )
+        return SparseAxisArray(new_data)
+    end
     return getindex(d.data, idx)
 end
 
+# Method to check whether an index is an attempt at a slice.
+@generated function _sliced_key_type(::Type{K}, idx...) where {K<:Tuple}
+    expr = Expr(:curly, :Tuple)
+    for i in 1:length(idx)
+        Ki = K.parameters[i]
+        if idx[i] <: Colon || idx[i] <: AbstractVector{<:Ki}
+            push!(expr.args, Ki)
+        end
+    end
+    return length(expr.args) == 1 ? :(nothing) : expr
+end
+
+# Methods to check whether a key `k` is a valid subset of `idx`.
+_filter(::Any, ::Colon) = true
+_filter(ki::Any, i::Any) = ki == i
+_filter(ki::K, i::AbstractVector{<:K}) where {K} = ki in i
+_filter(::Tuple{}, ::Tuple{}) = true
+function _filter(k::Tuple, idx::Tuple)
+    return _filter(k[1], idx[1]) && _filter(Base.tail(k), Base.tail(idx))
+end
+
+# Methods to subset the key into a new key, dropping all singleton axes.
+_sliced_key(k, ::Any) = (k,)
+_sliced_key(::K, ::K) where {K} = ()
+_sliced_key(::Tuple{}, ::Tuple{}) = ()
+function _sliced_key(k::Tuple, idx::Tuple)
+    return tuple(
+        _sliced_key(k[1], idx[1])...,
+        _sliced_key(Base.tail(k), Base.tail(idx))...,
+    )
+end
+
 Base.eachindex(d::SparseAxisArray) = keys(d.data)
 
 ################

test/Containers/SparseAxisArray.jl

@@ -14,14 +14,6 @@ using Test
 @testset "SparseAxisArray" begin
     function sparse_test(d, sum_d, d2, d3, dsqr, d_bads)
         sqr(x) = x^2
-        @testset "Colon indexing" begin
-            err = ArgumentError(
-                "Indexing with `:` is not supported by" *
-                " Containers.SparseAxisArray",
-            )
-            @test_throws err d[:, ntuple(one, ndims(d) - 1)...]
-            @test_throws err d[ntuple(i -> :a, ndims(d) - 1)..., :]
-        end
         @testset "Map" begin
             @test d == @inferred map(identity, d)
             @test dsqr == @inferred map(sqr, d)
@@ -195,4 +187,37 @@ $(SparseAxisArray{Float64,2,Tuple{Symbol,Char}}) with 2 entries"""
         @test y isa SparseAxisArray{Any,2,Tuple{Any,Int}}
         @test isempty(y)
     end
+    @testset "Slicing" begin
+        Containers.@container(x[i = 1:4, j = 1:2; isodd(i + j)], i + j)
+        @test x[:, :] == x
+        @test x[1, :] == Containers.@container(y[j = 1:2; isodd(1 + j)], 1 + j)
+        @test x[:, 1] == Containers.@container(z[i = 1:4; isodd(i + 1)], i + 1)
+        @test isempty(x[[1, 3], [1, 3]])
+        @test typeof(x[[1, 3], [1, 3]]) == typeof(x)
+        @test typeof(x[[1, 3], 1]) ==
+              Containers.SparseAxisArray{Int,1,Tuple{Int}}
+        @test isempty(x[[1, 3], 1])
+        Containers.@container(y[i = 1:4; isodd(i)], i)
+        @test y[:] == y
+        Containers.@container(y[i = 1:4; isodd(i)], i)
+        @test y[[1, 3]] == y
+        z = Containers.@container([i = 1:3, j = [:A, :B]; i > 1], (i, j))
+        @test z[2, :] == Containers.@container([j = [:A, :B]; true], (2, j))
+        @test z[:, :A] == Containers.@container([i = 2:3; true], (i, :A))
+        @test z[:, :] == z
+        @test z[1:2, :A] == Containers.@container([i = 2:2; true], (i, :A))
+        @test z[2, [:A, :B]] ==
+              Containers.@container([j = [:A, :B]; true], (2, j))
+        @test z[1:2, [:A, :B]] ==
+              Containers.@container([i = 2:2, j = [:A, :B]; true], (i, j))
+    end
+    @testset "Slicing on set" begin
+        Containers.@container(x[i = 1:4, j = 1:2; isodd(i + j)], i + j)
+        err = ArgumentError(
+            "Slicing is not when calling setindex! on a SparseAxisArray",
+        )
+        @test_throws(err, x[:, :] = 1)
+        @test_throws(err, x[1, :] = 1)
+        @test_throws(err, x[1, 1:2] = 1)
+    end
 end