Add DataFrames.jl extension

Project.toml

@@ -9,32 +9,36 @@ MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 MutableArithmetics = "d8a4904e-b15c-11e9-3269-09a3773c0cb0"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
-Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [weakdeps]
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
 
 [extensions]
+JuMPDataFramesExt = "DataFrames"
 JuMPDimensionalDataExt = "DimensionalData"
 
 [compat]
+DataFrames = "1"
 DimensionalData = "0.24, 0.25, 0.26.2, 0.27"
 LinearAlgebra = "<0.0.1, 1.6"
 MacroTools = "0.5"
 MathOptInterface = "1.25.2"
 MutableArithmetics = "1.1"
 OrderedCollections = "1"
-Printf = "<0.0.1, 1.6"
 PrecompileTools = "1"
+Printf = "<0.0.1, 1.6"
 SparseArrays = "<0.0.1, 1.6"
 Test = "<0.0.1, 1.6"
 julia = "1.6"
 
 [extras]
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["DimensionalData", "Test"]
+test = ["DataFrames", "DimensionalData", "Test"]

ext/JuMPDataFramesExt.jl

@@ -0,0 +1,23 @@
+#  Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+module JuMPDataFramesExt
+
+import DataFrames
+import JuMP
+
+function JuMP.Containers.container(
+    f::Function,
+    indices,
+    ::Type{DataFrames.DataFrame},
+    names::AbstractVector,
+)
+    rows = vec(collect(indices))
+    df = DataFrames.DataFrame(NamedTuple{tuple(names...)}(arg) for arg in rows)
+    df.value = [f(arg...) for arg in rows]
+    return df
+end
+
+end #module

ext/test_DataFrames.jl

@@ -0,0 +1,83 @@
+#  Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+module TestContainersDataFrames
+
+using Test
+
+using DataFrames
+using JuMP
+
+function test_dimension_data_vector()
+    model = Model()
+    @variable(model, x[i = 2:4], container = DataFrame)
+    @test x isa DataFrame
+    @test size(x) == (3, 2)
+    @test names(x) == ["i", "value"]
+    return
+end
+
+function test_dimension_data_matrix()
+    model = Model()
+    @variable(model, x[i = 2:4, j = ["a", "b"]], container = DataFrame)
+    @test x isa DataFrame
+    @test size(x) == (6, 3)
+    @test names(x) == ["i", "j", "value"]
+    @test sum(x[x.j .== "a", :value]) isa AffExpr
+    return
+end
+
+function test_dimension_data_triangle()
+    model = Model()
+    @variable(model, x[i = 2:4, j in i:4], container = DataFrame)
+    @test x isa DataFrame
+    @test size(x) == (6, 3)
+    @test names(x) == ["i", "j", "value"]
+    return
+end
+
+function test_dimension_data_sparse()
+    model = Model()
+    @variable(model, x[i in 1:4, j in 1:4; isodd(i + j)], container = DataFrame)
+    @test x isa DataFrame
+    @test size(x) == (8, 3)
+    @test x.i == [1, 1, 2, 2, 3, 3, 4, 4]
+    @test x.j == [2, 4, 1, 3, 2, 4, 1, 3]
+    @test names(x) == ["i", "j", "value"]
+    return
+end
+
+function test_dataframes_expression()
+    model = Model()
+    B = ["a", "b"]
+    @variable(model, x[i = 2:4, j = B], container = DataFrame)
+    @expression(
+        model,
+        expr[j = B],
+        sum(x[x.j .== j, :value]),
+        container = DataFrame,
+    )
+    @test expr isa DataFrame
+    @test expr.j == ["a", "b"]
+    expr2 = DataFrames.combine(
+        DataFrames.groupby(x, :j),
+        :value => sum => :value,
+    )
+    @test expr == expr2
+    return
+end
+
+function test_data_frames_missing_names()
+    model = Model()
+    x = @variable(model, [1:3, 1:2], container = DataFrame)
+    @test all(startswith.(names(x), ["##", "##", "value"]))
+    x = @variable(model, [i in 1:3, 1:2], container = DataFrame)
+    @test all(startswith.(names(x), ["i", "##", "value"]))
+    x = @variable(model, [1:3, j in 1:2], container = DataFrame)
+    @test all(startswith.(names(x), ["##", "j", "value"]))
+    return
+end
+
+end