Tables support for IndexedVarArray (#30)

* Tables support for IndexedVarArray

* Implement Tables support along recent JuMP changes

* Format fix

* Move rowtable out of JuMP.Containers namespace

* Test with generic column names until new JuMP ver

Author: trulsf

Project.toml

@@ -7,12 +7,8 @@ version = "0.6.2"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
-Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
 [compat]
 Dictionaries = "0.3"
 JuMP = "1"
-Requires = "1.3"
-Tables = "1.7"
 julia = "1.6"

src/SparseVariables.jl

@@ -3,28 +3,19 @@ module SparseVariables
 using Dictionaries
 using JuMP
 using LinearAlgebra
-using Tables
-using Requires
 
 include("sparsearray.jl")
 include("sparsevararray.jl")
 include("macros.jl")
 include("dictionaries.jl")
-include("tables.jl")
 include("indexedarray.jl")
+include("tables.jl")
 
 export SparseArray
 export SparseVarArray
 export IndexedVarArray
 export @sparsevariable
 export insertvar!
 export unsafe_insertvar!
-export table
-
-function __init__()
-    @require DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" include(
-        "dataframes.jl",
-    )
-end
 
 end # module

src/dataframes.jl

@@ -1,107 +1,36 @@
-abstract type SolutionTable end
-
-Tables.istable(::Type{<:SolutionTable}) = true
-Tables.rowaccess(::Type{<:SolutionTable}) = true
-
-rows(t::SolutionTable) = t
-names(t::SolutionTable) = getfield(t, :names)
-lookup(t::SolutionTable) = getfield(t, :lookup)
-
-Base.eltype(::SolutionTable) = SolutionRow
-Base.length(t::SolutionTable) = length(t.var)
-
-struct SolutionRow <: Tables.AbstractRow
-    index_vals::Any
-    sol_val::Float64
-    source::SolutionTable
-end
-
-function Tables.getcolumn(s::SolutionRow, i::Int)
-    if i > length(getfield(s, :index_vals))
-        return getfield(s, :sol_val)
-    end
-    return getfield(s, :index_vals)[i]
-end
-
-function Tables.getcolumn(s::SolutionRow, nm::Symbol)
-    i = lookup(getfield(s, :source))[nm]
-    if i > length(getfield(s, :index_vals))
-        return getfield(s, :sol_val)
-    end
-    return getfield(s, :index_vals)[i]
-end
-
-Tables.columnnames(s::SolutionRow) = names(getfield(s, :source))
-
-struct SolutionTableSparse <: SolutionTable
-    names::Vector{Symbol}
-    lookup::Dict{Symbol,Int}
-    var::SparseVarArray
-end
-
-SolutionTableSparse(v::SparseVarArray) = SolutionTableSparse(v, Symbol(v.name))
-
-function SolutionTableSparse(v::SparseVarArray, name)
-    if length(v) > 0 && !has_values(first(v.data).model)
-        error("No solution values available for variable")
-    end
-    names = vcat(v.index_names, name)
-    lookup = Dict(nm => i for (i, nm) in enumerate(names))
-    return SolutionTableSparse(names, lookup, v)
-end
-
-function Base.iterate(t::SolutionTableSparse, state = nothing)
-    next =
-        isnothing(state) ? iterate(keys(t.var.data)) :
-        iterate(keys(t.var.data), state)
-    next === nothing && return nothing
-    return SolutionRow(next[1], JuMP.value(t.var[next[1]]), t), next[2]
-end
-
-table(var::SparseVarArray) = SolutionTableSparse(var)
-table(var::SparseVarArray, name) = SolutionTableSparse(var, name)
-
-struct SolutionTableDense <: SolutionTable
-    names::Vector{Symbol}
-    lookup::Dict{Symbol,Int}
-    index_lookup::Dict
-    var::Containers.DenseAxisArray
-end
-
-function SolutionTableDense(
-    v::Containers.DenseAxisArray{VariableRef,N,Ax,L},
-    name,
-    colnames...,
-) where {N,Ax,L}
-    if length(colnames) < length(axes(v))
-        error("Not enough column names provided")
-    end
-    if length(v) > 0 && !has_values(first(v).model)
-        error("No solution values available for variable")
+function _rows(x::Union{SparseArray,SparseVarArray,IndexedVarArray})
+    return zip(eachindex(x.data), keys(x.data))
+end
+
+# The rowtable functions should be moved to the JuMP.Containers namespace 
+# when Tables support is available in JuMP 
+function rowtable(
+    f::Function,
+    x::AbstractSparseArray;
+    header::Vector{Symbol} = Symbol[],
+)
+    if isempty(header)
+        header = Symbol[Symbol("x$i") for i in 1:ndims(x)]
+        push!(header, :y)
     end
-    names = vcat(colnames..., name)
-    lookup = Dict(nm => i for (i, nm) in enumerate(names))
-    index_lookup = Dict()
-    for (i, ax) in enumerate(v.axes)
-        index_lookup[i] = collect(ax)
+    got, want = length(header), ndims(x) + 1
+    if got != want
+        error(
+            "Invalid number of column names provided: Got $got, expected $want.",
+        )
     end
-    return SolutionTableDense(names, lookup, index_lookup, v)
+    names = tuple(header...)
+    return [NamedTuple{names}((args..., f(x[i]))) for (i, args) in _rows(x)]
 end
 
-function Base.iterate(t::SolutionTableDense, state = nothing)
-    next =
-        isnothing(state) ? iterate(eachindex(t.var)) :
-        iterate(eachindex(t.var), state)
-    next === nothing && return nothing
-    index = next[1]
-    index_vals = [t.index_lookup[i][index[i]] for i in 1:length(index)]
-    return SolutionRow(index_vals, JuMP.value(t.var[next[1]]), t), next[2]
+function rowtable(f::Function, x::IndexedVarArray, col_header::Symbol)
+    header = Symbol[k for k in keys(x.index_names)]
+    push!(header, col_header)
+    return rowtable(f, x; header = header)
 end
 
-function table(
-    var::Containers.DenseAxisArray{VariableRef,N,Ax,L},
-    name,
-    colnames...,
-) where {N,Ax,L}
-    return SolutionTableDense(var, name, colnames...)
+function rowtable(f::Function, x::IndexedVarArray)
+    header = Symbol[k for k in keys(x.index_names)]
+    push!(header, Symbol(f))
+    return rowtable(f, x; header = header)
 end

src/tables.jl

@@ -1,5 +1,4 @@
 [deps]
-DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 HiGHS = "87dc4568-4c63-4d18-b0c0-bb2238e4078b"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"

test/Project.toml

@@ -1,5 +1,4 @@
 using Base: product
-using DataFrames
 using Dictionaries
 using HiGHS
 using JuMP
@@ -195,23 +194,11 @@ end
     set_optimizer_attribute(m, MOI.Silent(), true)
     optimize!(m)
 
-    tab = table(y)
-    @test typeof(tab) == SV.SolutionTableSparse
-
+    tab = SparseVariables.rowtable(value, y; header = [:car, :year, :value])
     @test length(tab) == 5
 
-    r = first(tab)
-    @test typeof(r) == SV.SolutionRow
+    r = tab[1]
     @test r.car == "bmw"
-
-    t2 = table(u, :u, :car, :year)
-    @test typeof(t2) == SV.SolutionTableDense
-    @test length(t2) == 12
-    rows = collect(t2)
-    @test rows[11].year == 2003
-
-    df = dataframe(u, :u, :car, :year)
-    @test first(df.car) == "ford"
 end
 
 @testset "IndexedVarArray" begin
@@ -300,6 +287,37 @@ end
     @test length(z3.index_cache[4]) == 0
 end
 
+@testset "Tables IndexedVarArray" begin
+    m = Model()
+    @variable(m, y[car = cars, year = year] >= 0; container = IndexedVarArray)
+    for c in cars
+        insertvar!(y, c, 2002)
+    end
+    @constraint(m, sum(y[:, :]) <= 300)
+    @constraint(
+        m,
+        [i in year],
+        sum(car_cost[c, i] * y[c, i] for (c, i) in SV.select(y, :, i)) <= 200
+    )
+
+    @objective(m, Max, sum(y[c, i] for c in cars, i in year))
+
+    set_optimizer(m, HiGHS.Optimizer)
+    set_optimizer_attribute(m, MOI.Silent(), true)
+    optimize!(m)
+
+    tab = SparseVariables.rowtable(value, y)
+
+    T = NamedTuple{(:i1, :i2, :value),Tuple{String,Int,Float64}}
+    @test tab isa Vector{T}
+
+    @test length(tab) == 3
+    r = tab[1]
+    @test r.i1 == "ford"
+    @test r.i2 == 2002
+    @test r.value == 300.0
+end
+
 @testset "JuMP extension" begin
 
     # Test JuMP Extension