[SW-2573] Fix interactionConstraints on H2OXGBoostMOJOModel in Python…

… API (#2554)
h2oai · Jul 26, 2021 · dc35e1f · dc35e1f
1 parent fff5a05
commit dc35e1f
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 1 deletion.
diff --git a/api-generation/src/main/scala/ai/h2o/sparkling/api/generation/python/MOJOModelTemplate.scala b/api-generation/src/main/scala/ai/h2o/sparkling/api/generation/python/MOJOModelTemplate.scala
@@ -78,6 +78,7 @@ object MOJOModelTemplate
   }
 
   private def generateValueConversion(parameter: Parameter): String = parameter.dataType match {
+    case x if x.isArray && x.getComponentType.isArray() => "H2OTypeConverters.scala2DArrayToPython2DArray(value)"
     case x if x.isArray => "H2OTypeConverters.scalaArrayToPythonArray(value)"
     case _ => "value"
   }

diff --git a/py-scoring/src/ai/h2o/sparkling/ml/params/H2OTypeConverters.py b/py-scoring/src/ai/h2o/sparkling/ml/params/H2OTypeConverters.py
@@ -501,6 +501,15 @@ def scalaArrayToPythonArray(array):
         else:
             raise TypeError("Invalid type.")
 
+    @staticmethod
+    def scala2DArrayToPython2DArray(array):
+        if array is None:
+            return None
+        elif isinstance(array, JavaObject):
+            return [H2OTypeConverters.scalaArrayToPythonArray(v) for v in array]
+        else:
+            raise TypeError("Invalid type.")
+
     @staticmethod
     def scalaToPythonDataFrame(jdf):
         if jdf is None:

diff --git a/py/tests/unit/with_runtime_sparkling/test_mojo_parameters.py b/py/tests/unit/with_runtime_sparkling/test_mojo_parameters.py
@@ -39,7 +39,7 @@ def testXGBoostParameters(prostateDataset):
                            monotoneConstraints={'AGE': 1, 'RACE': -1},
                            interactionConstraints=[['AGE', 'RACE', 'DPROS'], ['DCAPS', 'PSA']])
     model = algorithm.fit(prostateDataset)
-    ignored=["getInteractionConstraints", "getMonotoneConstraints"]  # Will be fixed by SW-2573 and SW-2572
+    ignored=["getMonotoneConstraints"]  # Will be fixed by SW-2572
     compareParameterValues(algorithm, model, ignored)