Alter jvm_scala_object to better handle nested scala classes. (#838)

combust · Dec 22, 2022 · fb43440 · fb43440
1 parent 96eaaf9
commit fb43440
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 8 deletions.
diff --git a/python/mleap/pyspark/feature/math_binary.py b/python/mleap/pyspark/feature/math_binary.py
@@ -90,8 +90,8 @@ def __init__(
         # if operation is not None, we can proceed to instantiate the scala classes
         if operation:
             scalaBinaryOperation = jvm_scala_object(
-                _jvm().ml.combust.mleap.core.feature.BinaryOperation,
-                operation.name
+                _jvm().ml.combust.mleap.core.feature,
+                f"BinaryOperation${operation.name}$"
             )
 
             scalaMathBinaryModel = _jvm().ml.combust.mleap.core.feature.MathBinaryModel(

diff --git a/python/mleap/pyspark/feature/math_unary.py b/python/mleap/pyspark/feature/math_unary.py
@@ -46,8 +46,8 @@ def __init__(self, operation=None, inputCol=None, outputCol=None):
         # if operation is not None, we can proceed to instantiate the scala classes
         if operation:
             scalaUnaryOperation = jvm_scala_object(
-                _jvm().ml.combust.mleap.core.feature.UnaryOperation,
-                operation.name
+                _jvm().ml.combust.mleap.core.feature,
+                f"UnaryOperation${operation.name}$",
             )
 
             scalaMathUnaryModel = _jvm().ml.combust.mleap.core.feature.MathUnaryModel(scalaUnaryOperation)

diff --git a/python/mleap/pyspark/feature/string_map.py b/python/mleap/pyspark/feature/string_map.py
@@ -45,8 +45,8 @@ def validate_args():
             .toMap(_jvm().scala.Predef.conforms())
 
         handle_invalid_jvm = jvm_scala_object(
-            _jvm().ml.combust.mleap.core.feature.HandleInvalid,
-            handleInvalid.capitalize(),
+            _jvm().ml.combust.mleap.core.feature,
+            f"HandleInvalid${handleInvalid.capitalize()}$",
         )
 
         string_map_model = _jvm().ml.combust.mleap.core.feature.StringMapModel(

diff --git a/python/mleap/pyspark/py2scala.py b/python/mleap/pyspark/py2scala.py
@@ -13,7 +13,7 @@ def jvm_scala_object(jpkg, obj):
         (for reference see file ml.combust.mleap.core.feature.MathUnaryModel)
     """
     return getattr(
-        getattr(jpkg, obj + "$"),  # JavaClass
+        getattr(jpkg, obj),  # JavaClass
         "MODULE$",  # JavaObject
     )
 
@@ -25,4 +25,4 @@ def Some(value):
     return _jvm().scala.Some(value)
 
 def ScalaNone():
-    return jvm_scala_object(_jvm().scala, "None")
+    return jvm_scala_object(_jvm().scala, "None$")