fix test

snowflakedb · Sep 27, 2024 · 1f4b581 · 1f4b581
1 parent c171397
commit 1f4b581
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 12 deletions.
diff --git a/src/main/scala/com/snowflake/snowpark/internal/JavaUtils.scala b/src/main/scala/com/snowflake/snowpark/internal/JavaUtils.scala
@@ -364,15 +364,18 @@ object JavaUtils {
       case (key, value) => key -> value
     }.toMap
 
-  def scalaMapToJavaWithVariantConversion(map: Map[_, _]): java.util.Map[Object, Object] =
-    map
-      .map[Object, Object] {
-        case (key, value: com.snowflake.snowpark.types.Variant) =>
-          key.asInstanceOf[Object] -> InternalUtils.createVariant(value).asInstanceOf[Object]
-        case (key, value) => key.asInstanceOf[Object] -> value.asInstanceOf[Object]
-      }
-      .asInstanceOf[Map[Object, Object]]
-      .asJava
+  def scalaMapToJavaWithVariantConversion(map: Map[_, _]): java.util.Map[Object, Object] = {
+    val result = new java.util.HashMap[Object, Object]()
+    map.foreach {
+      case (key, value: com.snowflake.snowpark.types.Variant) =>
+        result.put(
+          key.asInstanceOf[Object],
+          InternalUtils.createVariant(value).asInstanceOf[Object])
+      case (key, value) =>
+        result.put(key.asInstanceOf[Object], value.asInstanceOf[Object])
+    }
+    result
+  }
 
   def serialize(obj: Any): Array[Byte] = {
     val bos = new ByteArrayOutputStream()

diff --git a/src/main/scala/com/snowflake/snowpark/types/Variant.scala b/src/main/scala/com/snowflake/snowpark/types/Variant.scala
@@ -11,6 +11,7 @@ import Variant._
 import org.apache.commons.codec.binary.{Base64, Hex}
 
 import java.io.{IOException, UncheckedIOException}
+import java.util.function.Consumer
 import scala.util.hashing.MurmurHash3
 
 private[snowpark] object Variant {
@@ -248,9 +249,11 @@ class Variant private[snowpark] (
       {
         def mapToNode(map: JavaMap[Object, Object]): ObjectNode = {
           val result = MAPPER.createObjectNode()
-          map.asScala.foreach { case (key, value) =>
-            result.set(key.toString, objectToJsonNode(value))
-          }
+          map.keySet.forEach(new Consumer[Object] {
+            override def accept(key: Object): Unit = {
+              result.set(key.toString, objectToJsonNode(map.get(key)))
+            }
+          })
           result
         }
         obj match {