Add case classes to enforce matching problem budget data assumptions (#42)

Author: borjaocejo

dualip/src/main/scala/com/linkedin/dualip/data/BudgetData.scala

@@ -0,0 +1,20 @@
+package com.linkedin.dualip.data
+
+/**
+  * Representation of the budget vector/budget per item id for Matching Problem.
+  *
+  * @param itemId - id of the item (i.e. product id)
+  * @param budget - the budget for itemId/the b(itemId) value in the budget (b) vector
+  */
+case class BudgetData(itemId: Int, budget: Double)
+
+/**
+  * Representation of the budget for a specific entity, given a constraint matrix.
+  * This is used to represent the budget of an entity that may have several constraints, as in the
+  * Multiple Matching Problem.
+  *
+  * @param entityIndex - index for the entity
+  * @param constraintIndex - index of the constraint
+  * @param budgetValue - budget for this entity-constraint combo, or b(entityIndex)(constraintIndex)
+  */
+case class MultipleMatchingBudgetData(entityIndex: Int, constraintIndex: Int, budgetValue: Double)

dualip/src/main/scala/com/linkedin/dualip/driver/LPSolverDriver.scala

@@ -345,7 +345,7 @@ object LPSolverDriver {
       val driverParams = LPSolverDriverParamsParser.parseArgs(args)
 
       println("-----------------------------------------------------------------------")
-      println("                      DuaLip v2.0     2022")
+      println("                      DuaLip v4.0     2023")
       println("            Dual Decomposition based Linear Program Solver")
       println("-----------------------------------------------------------------------")
       println("Settings:")

dualip/src/main/scala/com/linkedin/dualip/problem/MatchingSlateSolver.scala

@@ -1,7 +1,7 @@
 package com.linkedin.dualip.problem
 
 import breeze.linalg.{SparseVector => BSV}
-import com.linkedin.dualip.data.MatchingData
+import com.linkedin.dualip.data.{BudgetData, MatchingData}
 import com.linkedin.dualip.objective.distributedobjective.DistributedRegularizedObjective
 import com.linkedin.dualip.objective.{DualPrimalObjective, DualPrimalObjectiveLoader, PartialPrimalStats}
 import com.linkedin.dualip.projection.{BoxCutProjection, GreedyProjection, SimplexProjection, UnitBoxProjection}
@@ -12,7 +12,7 @@ import com.linkedin.optimization.util.VectorOperations.toBSV
 import com.twitter.algebird.{Max, Tuple5Semigroup}
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.sql.functions.{col, lit}
-import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}
+import org.apache.spark.sql.{DataFrame, Dataset, SparkSession}
 import org.apache.spark.storage.StorageLevel
 
 import scala.util.Try
@@ -160,8 +160,8 @@ object MatchingSolverDualObjectiveFunction extends DualPrimalObjectiveLoader {
     */
   def loadData(inputPaths: InputPaths)(implicit spark: SparkSession): (Dataset[MatchingData], BSV[Double]) = {
     import spark.implicits._
-    val budget = IOUtility.readDataFrame(inputPaths.vectorBPath, inputPaths.format)
-      .map { case Row(_c0: Number, _c1: Number) => (_c0.intValue(), _c1.doubleValue()) }
+    val budget = IOUtility.readDataFrame(inputPaths.vectorBPath, inputPaths.format).as[BudgetData]
+      .map { budgetRecord => (budgetRecord.itemId, budgetRecord.budget) }
       .collect
 
     val itemIds = budget.toMap.keySet

dualip/src/main/scala/com/linkedin/dualip/problem/MultipleMatchingSolver.scala

@@ -1,7 +1,7 @@
 package com.linkedin.dualip.problem
 
 import breeze.linalg.{SparseVector => BSV}
-import com.linkedin.dualip.data.MultipleMatchingData
+import com.linkedin.dualip.data.{MultipleMatchingBudgetData, MultipleMatchingData}
 import com.linkedin.dualip.objective.distributedobjective.DistributedRegularizedObjective
 import com.linkedin.dualip.objective.{DualPrimalObjective, DualPrimalObjectiveLoader, PartialPrimalStats}
 import com.linkedin.dualip.projection.{BoxCutProjection, GreedyProjection, SimplexProjection, UnitBoxProjection}
@@ -146,8 +146,8 @@ object MultipleMatchingSolverDualObjectiveFunction extends DualPrimalObjectiveLo
     // constraintIndex: index for the constraint; for a multiple-matching problem with three matching constraints per
     // entity, this column may assume values 0, 1 and 2
     // budgetValue: value of the budget
-    val budgetDF = IOUtility.readDataFrame(inputPaths.vectorBPath, inputPaths.format)
-      .toDF("entityIndex", "constraintIndex", "budgetValue")
+    val budgetDF = IOUtility.readDataFrame(inputPaths.vectorBPath, inputPaths.format).as[MultipleMatchingBudgetData]
+      .toDF()
 
     val matchingConstraintsPerIndex = budgetDF.select("constraintIndex").distinct().count().toInt
     val reindexBudgetDF = budgetDF

version.properties

@@ -1,2 +1,2 @@
-version=3.0.*
+version=4.0.*