From d365f729d531f8fb89cda07133f99791a20ddfbc Mon Sep 17 00:00:00 2001
From: Ruslan Sadykov <41117501+rrsadykov@users.noreply.github.com>
Date: Tue, 28 Nov 2023 15:43:19 +0100
Subject: [PATCH] Get custom data of a column in a disaggregated solution
 (#1110)

* Get custom data of a column in a disaggregated solution

* Add tests for BD.customdata()
---
 .gitignore                                               | 3 ++-
 Project.toml                                             | 2 +-
 src/MOIwrapper.jl                                        | 5 +++++
 test/MathOptInterface/MOI_wrapper.jl                     | 2 +-
 test/e2e_extra/advanced_colgen/improve_relax_callback.jl | 1 +
 test/runtests.jl                                         | 1 +
 6 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/.gitignore b/.gitignore
index e5b85bc7f..2953547c9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -14,4 +14,5 @@ test/Manifest.toml
 *.dot
 *.mps
 *DS_Store
-test/.222-revise-exit-code
\ No newline at end of file
+test/.222-revise-exit-code
+test/playgap.json
diff --git a/Project.toml b/Project.toml
index b392c8be7..94520128b 100644
--- a/Project.toml
+++ b/Project.toml
@@ -23,7 +23,7 @@ TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [compat]
-BlockDecomposition = "1.13"
+BlockDecomposition = "1.14"
 Crayons = "4.1"
 DataStructures = "0.17, 0.18"
 DynamicSparseArrays = "0.7"
diff --git a/src/MOIwrapper.jl b/src/MOIwrapper.jl
index 16202038b..40d9d0d88 100644
--- a/src/MOIwrapper.jl
+++ b/src/MOIwrapper.jl
@@ -1108,6 +1108,11 @@ end
 
 BD.value(info::ColumnInfo) = info.column_val
 
+function BD.customdata(info::ColumnInfo) 
+    masterform = getmaster(info.optimizer.inner.re_formulation)
+    return getcustomdata(masterform, info.column_var_id)
+end
+
 function BD.value(info::ColumnInfo, index::MOI.VariableIndex)
     varid = info.optimizer.env.varids[index]
     origin_form_uid = getoriginformuid(info.column_var_id)
diff --git a/test/MathOptInterface/MOI_wrapper.jl b/test/MathOptInterface/MOI_wrapper.jl
index aa3abd7e0..b260b3633 100644
--- a/test/MathOptInterface/MOI_wrapper.jl
+++ b/test/MathOptInterface/MOI_wrapper.jl
@@ -78,7 +78,7 @@ function test_runtests()
             "test_linear_integration_Interval",
             "test_linear_integration_delete_variables",
             "test_linear_transform",
-            # Can we just exclude them? 
+            # To see later if we need to support SOS2 integration
             "test_linear_SOS2_integration",
         ],
         # This argument is useful to prevent tests from failing on future
diff --git a/test/e2e_extra/advanced_colgen/improve_relax_callback.jl b/test/e2e_extra/advanced_colgen/improve_relax_callback.jl
index d2c499212..4e1b9c3a2 100644
--- a/test/e2e_extra/advanced_colgen/improve_relax_callback.jl
+++ b/test/e2e_extra/advanced_colgen/improve_relax_callback.jl
@@ -216,6 +216,7 @@ function test_improve_relaxation(; do_improve::Bool)
         sets = BD.getsolutions(model, b)
         for s in sets
             @test BD.value(s) == 1.0 # value of the master column variable
+            @test BD.customdata(s).items == [1] || BD.customdata(s).items == [2, 3] # either [1] or [2, 3] 
             @test BD.value(s, x[b, 1]) != BD.value(s, x[b, 2]) # only x[1,1] in its set
             @test BD.value(s, x[b, 1]) != BD.value(s, x[b, 3]) # only x[1,1] in its set
             @test BD.value(s, x[b, 2]) == BD.value(s, x[b, 3]) # x[1,2] and x[1,3] in the same set
diff --git a/test/runtests.jl b/test/runtests.jl
index 2331a8de0..952951c7c 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -62,6 +62,7 @@ else
     end
 
     run_e2e_tests()
+    run_e2e_extra_tests()
 end
 
 @testset "Version" begin