add JsonNode function for java and scala with test case

snowflakedb · Aug 27, 2024 · 6036ed3 · 6036ed3
1 parent 0a83ef9
commit 6036ed3
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 2 deletions.
diff --git a/src/main/java/com/snowflake/snowpark_java/types/Variant.java b/src/main/java/com/snowflake/snowpark_java/types/Variant.java
@@ -383,7 +383,6 @@ public String asJsonString() {
    * @since 1.14.0
    */
   public JsonNode asJsonNode() {
-    verify(VariantTypes.Object);
     return objectToJsonNode(value);
   }
 

diff --git a/src/main/scala/com/snowflake/snowpark/types/Variant.scala b/src/main/scala/com/snowflake/snowpark/types/Variant.scala
@@ -44,6 +44,8 @@ private[snowpark] object Variant {
 
     object Object extends VariantType
 
+    object Jsonnode extends VariantType
+
     // internal used when converting from Java
     def getType(name: String): VariantType = name match {
       case "RealNumber" => RealNumber
@@ -381,6 +383,15 @@ class Variant private[snowpark] (
     }
   }
 
+  /**
+   * Converts the variant as valid JsonNode
+   *
+   * @since 1.14.0
+   */
+  def asJsonNode(): JsonNode = {
+    objectToJsonNode(value)
+  }
+
   /**
    * Converts the variant as binary value
    * @since 0.2.0

diff --git a/src/test/java/com/snowflake/snowpark_test/JavaVariantSuite.java b/src/test/java/com/snowflake/snowpark_test/JavaVariantSuite.java
@@ -308,7 +308,7 @@ public void javaVariantConversion() {
     assert (result.asMap().get("a").get().asString().equals("[1,2]"));
     assert (result.asMap().get("a").get().asJsonString().equals("[1,2]"));
 
-    // Scala variant to Java Variant
+//     Scala variant to Java Variant
     com.snowflake.snowpark.types.Variant sv =
         new com.snowflake.snowpark.types.Variant("{\"a\": [1, 2], \"b\": \"c\"}");
     assert InternalUtils.createVariant(null) == null;

diff --git a/src/test/scala/com/snowflake/snowpark_test/ScalaVariantSuite.scala b/src/test/scala/com/snowflake/snowpark_test/ScalaVariantSuite.scala
@@ -1,6 +1,7 @@
 package com.snowflake.snowpark_test
 
 import com.snowflake.snowpark.types.{Geography, Variant}
+import com.snowflake.snowpark_java.types.InternalUtils
 import org.scalatest.FunSuite
 
 import java.io.UncheckedIOException
@@ -311,4 +312,11 @@ class ScalaVariantSuite extends FunSuite {
 
     assert(v1.toString() == "123")
   }
+
+  test("JsonNode") {
+    val sv = new Variant("{\"a\": [1, 2], \"b\": 3, \"c\": \"xyz\"}")
+    assert(sv.asMap.get("a").get.asJsonNode().isArray)
+    assert(sv.asMap.get("b").get.asJsonNode().asInt().equals(3))
+    assert(sv.asMap.get("c").get.asJsonNode().asText().equals("xyz"))
+  }
 }