Skip to content

Commit d33e05c

Browse files
authored
Add Tables.jl support to containers (#3104)
1 parent bcb803e commit d33e05c

File tree

6 files changed

+306
-0
lines changed

6 files changed

+306
-0
lines changed

docs/src/manual/containers.md

+100
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,42 @@ julia> swap.(x)
102102
(1, 2) (2, 2) (3, 2)
103103
```
104104

105+
### Tables
106+
107+
Use [`Containers.rowtable`](@ref) to convert the `Array` into a
108+
[Tables.jl](https://github.com/JuliaData/Tables.jl) compatible
109+
`Vector{<:NamedTuple}`:
110+
111+
```jldoctest containers_array
112+
julia> table = Containers.rowtable(x; header = [:I, :J, :value])
113+
6-element Vector{NamedTuple{(:I, :J, :value), Tuple{Int64, Int64, Tuple{Int64, Int64}}}}:
114+
(I = 1, J = 1, value = (1, 1))
115+
(I = 2, J = 1, value = (2, 1))
116+
(I = 1, J = 2, value = (1, 2))
117+
(I = 2, J = 2, value = (2, 2))
118+
(I = 1, J = 3, value = (1, 3))
119+
(I = 2, J = 3, value = (2, 3))
120+
```
121+
122+
Because it supports the [Tables.jl](https://github.com/JuliaData/Tables.jl)
123+
interface, you can pass it to any function which accepts a table as input:
124+
125+
```jldoctest containers_array
126+
julia> import DataFrames;
127+
128+
julia> DataFrames.DataFrame(table)
129+
6×3 DataFrame
130+
Row │ I J value
131+
│ Int64 Int64 Tuple…
132+
─────┼──────────────────────
133+
1 │ 1 1 (1, 1)
134+
2 │ 2 1 (2, 1)
135+
3 │ 1 2 (1, 2)
136+
4 │ 2 2 (2, 2)
137+
5 │ 1 3 (1, 3)
138+
6 │ 2 3 (2, 3)
139+
```
140+
105141
## DenseAxisArray
106142

107143
A [`Containers.DenseAxisArray`](@ref) is created when the index sets are
@@ -191,6 +227,38 @@ julia> x.data
191227
(2, :A) (2, :B)
192228
```
193229

230+
### Tables
231+
232+
Use [`Containers.rowtable`](@ref) to convert the `DenseAxisArray` into a
233+
[Tables.jl](https://github.com/JuliaData/Tables.jl) compatible
234+
`Vector{<:NamedTuple}`:
235+
236+
```jldoctest containers_dense
237+
julia> table = Containers.rowtable(x; header = [:I, :J, :value])
238+
4-element Vector{NamedTuple{(:I, :J, :value), Tuple{Int64, Symbol, Tuple{Int64, Symbol}}}}:
239+
(I = 1, J = :A, value = (1, :A))
240+
(I = 2, J = :A, value = (2, :A))
241+
(I = 1, J = :B, value = (1, :B))
242+
(I = 2, J = :B, value = (2, :B))
243+
```
244+
245+
Because it supports the [Tables.jl](https://github.com/JuliaData/Tables.jl)
246+
interface, you can pass it to any function which accepts a table as input:
247+
248+
```jldoctest containers_dense
249+
julia> import DataFrames;
250+
251+
julia> DataFrames.DataFrame(table)
252+
4×3 DataFrame
253+
Row │ I J value
254+
│ Int64 Symbol Tuple…
255+
─────┼────────────────────────
256+
1 │ 1 A (1, :A)
257+
2 │ 2 A (2, :A)
258+
3 │ 1 B (1, :B)
259+
4 │ 2 B (2, :B)
260+
```
261+
194262
## SparseAxisArray
195263

196264
A [`Containers.SparseAxisArray`](@ref) is created when the index sets are
@@ -252,6 +320,38 @@ JuMP.Containers.SparseAxisArray{Tuple{Symbol, Int64}, 1, Tuple{Int64}} with 2 en
252320
[3] = (:B, 3)
253321
```
254322

323+
### Tables
324+
325+
Use [`Containers.rowtable`](@ref) to convert the `SparseAxisArray` into a
326+
[Tables.jl](https://github.com/JuliaData/Tables.jl) compatible
327+
`Vector{<:NamedTuple}`:
328+
329+
```jldoctest containers_sparse
330+
julia> table = Containers.rowtable(x; header = [:I, :J, :value])
331+
4-element Vector{NamedTuple{(:I, :J, :value), Tuple{Int64, Symbol, Tuple{Int64, Symbol}}}}:
332+
(I = 3, J = :B, value = (3, :B))
333+
(I = 2, J = :A, value = (2, :A))
334+
(I = 2, J = :B, value = (2, :B))
335+
(I = 3, J = :A, value = (3, :A))
336+
```
337+
338+
Because it supports the [Tables.jl](https://github.com/JuliaData/Tables.jl)
339+
interface, you can pass it to any function which accepts a table as input:
340+
341+
```jldoctest containers_sparse
342+
julia> import DataFrames;
343+
344+
julia> DataFrames.DataFrame(table)
345+
4×3 DataFrame
346+
Row │ I J value
347+
│ Int64 Symbol Tuple…
348+
─────┼────────────────────────
349+
1 │ 3 B (3, :B)
350+
2 │ 2 A (2, :A)
351+
3 │ 2 B (2, :B)
352+
4 │ 3 A (3, :A)
353+
```
354+
255355
## Forcing the container type
256356

257357
Pass `container = T` to use `T` as the container. For example:

docs/src/reference/containers.md

+1
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ Containers
77
Containers.DenseAxisArray
88
Containers.SparseAxisArray
99
Containers.container
10+
Containers.rowtable
1011
Containers.default_container
1112
Containers.@container
1213
Containers.VectorizedProductIterator

docs/src/tutorials/linear/diet.jl

+10
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,16 @@ end
139139
# That's a lot of milk and ice cream! And sadly, we only get `0.6` of a
140140
# hamburger.
141141

142+
# We can also use the function [`Containers.rowtable`](@ref) to easily convert
143+
# the result into a DataFrame:
144+
145+
table = Containers.rowtable(value, x; header = [:food, :quantity])
146+
solution = DataFrames.DataFrame(table)
147+
148+
# This makes it easy to perform analyses our solution:
149+
150+
filter!(row -> row.quantity > 0.0, solution)
151+
142152
# ## Problem modification
143153

144154
# JuMP makes it easy to take an existing model and modify it by adding extra

src/Containers/Containers.jl

+1
Original file line numberDiff line numberDiff line change
@@ -54,5 +54,6 @@ include("nested_iterator.jl")
5454
include("no_duplicate_dict.jl")
5555
include("container.jl")
5656
include("macro.jl")
57+
include("tables.jl")
5758

5859
end

src/Containers/tables.jl

+68
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
# Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
2+
# This Source Code Form is subject to the terms of the Mozilla Public
3+
# License, v. 2.0. If a copy of the MPL was not distributed with this
4+
# file, You can obtain one at https://mozilla.org/MPL/2.0/.
5+
6+
_rows(x::Array) = zip(eachindex(x), Iterators.product(axes(x)...))
7+
8+
_rows(x::DenseAxisArray) = zip(vec(eachindex(x)), Iterators.product(axes(x)...))
9+
10+
_rows(x::SparseAxisArray) = zip(eachindex(x.data), keys(x.data))
11+
12+
"""
13+
rowtable([f::Function=identity,] x; [header::Vector{Symbol} = Symbol[]])
14+
15+
Applies the function `f` to all elements of the variable container `x`,
16+
returning the result as a `Vector` of `NamedTuple`s, where `header` is a vector
17+
containing the corresponding axis names.
18+
19+
If `x` is an `N`-dimensional array, there must be `N+1` names, so that the last
20+
name corresponds to the result of `f(x[i])`.
21+
22+
If `header` is left empty, then the default header is `[:x1, :x2, ..., :xN, :y]`.
23+
24+
!!! info
25+
A `Vector` of `NamedTuple`s implements the [Tables.jl](https://github.com/JuliaData/Tables.jl)
26+
interface, and so the result can be used as input for any function
27+
that consumes a 'Tables.jl' compatible source.
28+
29+
## Example
30+
31+
```jldoctest; setup=:(using JuMP)
32+
julia> model = Model();
33+
34+
julia> @variable(model, x[i=1:2, j=i:2] >= 0, start = i+j);
35+
36+
julia> Containers.rowtable(start_value, x; header = [:i, :j, :start])
37+
3-element Vector{NamedTuple{(:i, :j, :start), Tuple{Int64, Int64, Float64}}}:
38+
(i = 1, j = 2, start = 3.0)
39+
(i = 1, j = 1, start = 2.0)
40+
(i = 2, j = 2, start = 4.0)
41+
42+
julia> Containers.rowtable(x)
43+
3-element Vector{NamedTuple{(:x1, :x2, :y), Tuple{Int64, Int64, VariableRef}}}:
44+
(x1 = 1, x2 = 2, y = x[1,2])
45+
(x1 = 1, x2 = 1, y = x[1,1])
46+
(x1 = 2, x2 = 2, y = x[2,2])
47+
```
48+
"""
49+
function rowtable(
50+
f::Function,
51+
x::Union{Array,DenseAxisArray,SparseAxisArray};
52+
header::Vector{Symbol} = Symbol[],
53+
)
54+
if isempty(header)
55+
header = Symbol[Symbol("x$i") for i in 1:ndims(x)]
56+
push!(header, :y)
57+
end
58+
got, want = length(header), ndims(x) + 1
59+
if got != want
60+
error(
61+
"Invalid number of column names provided: Got $got, expected $want.",
62+
)
63+
end
64+
names = tuple(header...)
65+
return [NamedTuple{names}((args..., f(x[i]))) for (i, args) in _rows(x)]
66+
end
67+
68+
rowtable(x; kwargs...) = rowtable(identity, x; kwargs...)

test/Containers/tables.jl

+126
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
# Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
2+
# This Source Code Form is subject to the terms of the Mozilla Public
3+
# License, v. 2.0. If a copy of the MPL was not distributed with this
4+
# file, You can obtain one at https://mozilla.org/MPL/2.0/.
5+
6+
module TestTableInterface
7+
8+
using JuMP
9+
using Test
10+
11+
function runtests()
12+
for name in names(@__MODULE__; all = true)
13+
if startswith("$(name)", "test_")
14+
@testset "$(name)" begin
15+
getfield(@__MODULE__, name)()
16+
end
17+
end
18+
end
19+
return
20+
end
21+
22+
function test_denseaxisarray()
23+
model = Model()
24+
@variable(model, x[i = 4:10, j = 2002:2022] >= 0, start = 0.0)
25+
@test typeof(x) <: Containers.DenseAxisArray
26+
start_table = Containers.rowtable(start_value, x; header = [:i1, :i2, :i3])
27+
T = NamedTuple{(:i1, :i2, :i3),Tuple{Int,Int,Float64}}
28+
@test start_table isa Vector{T}
29+
@test length(start_table) == length(x)
30+
row = first(start_table)
31+
@test row == (i1 = 4, i2 = 2002, i3 = 0.0)
32+
x_table = Containers.rowtable(x; header = [:i1, :i2, :i3])
33+
@test x_table[1] == (i1 = 4, i2 = 2002, i3 = x[4, 2002])
34+
return
35+
end
36+
37+
function test_array()
38+
model = Model()
39+
@variable(model, x[1:10, 1:5] >= 0, start = 0.0)
40+
@test typeof(x) <: Array{VariableRef}
41+
start_table = Containers.rowtable(start_value, x; header = [:i1, :i2, :i3])
42+
T = NamedTuple{(:i1, :i2, :i3),Tuple{Int,Int,Float64}}
43+
@test start_table isa Vector{T}
44+
@test length(start_table) == length(x)
45+
row = first(start_table)
46+
@test row == (i1 = 1, i2 = 1, i3 = 0.0)
47+
x_table = Containers.rowtable(x; header = [:i1, :i2, :i3])
48+
@test x_table[1] == (i1 = 1, i2 = 1, i3 = x[1, 1])
49+
return
50+
end
51+
52+
function test_sparseaxisarray()
53+
model = Model()
54+
@variable(model, x[i = 1:10, j = 1:5; i + j <= 8] >= 0, start = 0)
55+
@test typeof(x) <: Containers.SparseAxisArray
56+
start_table = Containers.rowtable(start_value, x; header = [:i1, :i2, :i3])
57+
T = NamedTuple{(:i1, :i2, :i3),Tuple{Int,Int,Float64}}
58+
@test start_table isa Vector{T}
59+
@test length(start_table) == length(x)
60+
@test (i1 = 1, i2 = 1, i3 = 0.0) in start_table
61+
x_table = Containers.rowtable(x; header = [:i1, :i2, :i3])
62+
@test (i1 = 1, i2 = 1, i3 = x[1, 1]) in x_table
63+
return
64+
end
65+
66+
function test_col_name_error()
67+
model = Model()
68+
@variable(model, x[1:2, 1:2])
69+
@test_throws ErrorException Containers.rowtable(x; header = [:y, :a])
70+
@test_throws(
71+
ErrorException,
72+
Containers.rowtable(x; header = [:y, :a, :b, :c]),
73+
)
74+
@test Containers.rowtable(x; header = [:y, :a, :b]) isa Vector{<:NamedTuple}
75+
return
76+
end
77+
78+
# Mockup of custom variable type
79+
struct _MockVariable <: JuMP.AbstractVariable
80+
var::JuMP.ScalarVariable
81+
end
82+
83+
struct _MockVariableRef <: JuMP.AbstractVariableRef
84+
vref::VariableRef
85+
end
86+
87+
JuMP.name(v::_MockVariableRef) = JuMP.name(v.vref)
88+
89+
JuMP.owner_model(v::_MockVariableRef) = JuMP.owner_model(v.vref)
90+
91+
JuMP.start_value(v::_MockVariableRef) = JuMP.start_value(v.vref)
92+
93+
struct _Mock end
94+
95+
function JuMP.build_variable(::Function, info::JuMP.VariableInfo, _::_Mock)
96+
return _MockVariable(JuMP.ScalarVariable(info))
97+
end
98+
99+
function JuMP.add_variable(model::Model, x::_MockVariable, name::String)
100+
variable = JuMP.add_variable(model, x.var, name)
101+
return _MockVariableRef(variable)
102+
end
103+
104+
function test_custom_variable()
105+
model = Model()
106+
@variable(
107+
model,
108+
x[i = 1:3, j = 100:102] >= 0,
109+
_Mock(),
110+
container = Containers.DenseAxisArray,
111+
start = 0.0,
112+
)
113+
@test typeof(x) <: Containers.DenseAxisArray
114+
start_table = Containers.rowtable(start_value, x)
115+
T = NamedTuple{(:x1, :x2, :y),Tuple{Int,Int,Float64}}
116+
@test start_table isa Vector{T}
117+
@test length(start_table) == length(x)
118+
@test (x1 = 1, x2 = 100, y = 0.0) in start_table
119+
x_table = Containers.rowtable(x)
120+
@test (x1 = 1, x2 = 100, y = x[1, 100]) in x_table
121+
return
122+
end
123+
124+
end
125+
126+
TestTableInterface.runtests()

0 commit comments

Comments
 (0)