java api

snowflakedb · Nov 29, 2023 · 4304f72 · 4304f72
1 parent eb96b32
commit 4304f72
Show file tree

Hide file tree

Showing 3 changed files with 82 additions and 1 deletion.
diff --git a/src/main/java/com/snowflake/snowpark_java/TableFunctions.java b/src/main/java/com/snowflake/snowpark_java/TableFunctions.java
@@ -1,5 +1,6 @@
 package com.snowflake.snowpark_java;
 
+
 /**
  * Provides utility functions that generate table function expressions that can be passed to
  * DataFrame join method and Session tableFunction method.
@@ -150,4 +151,24 @@ public static Column flatten(
   public static Column flatten(Column input) {
     return new Column(com.snowflake.snowpark.tableFunctions.flatten(input.toScalaColumn()));
   }
+
+  /**
+   * Flattens a given array or map type column into individual rows. The output column(s) in case of
+   * array input column is `VALUE`, and are `KEY` and `VALUE` in case of amp input column.
+   *
+   * <p>Example
+   *
+   * <pre>{@code
+   * df.join(TableFunctions.flatten(
+   *   Functions.parse_json(df.col("col"))));
+   * }</pre>
+   *
+   * @since 1.10.0
+   * @param input The expression that will be unseated into rows. The expression must be MapType or
+   *     ArrayType data.
+   * @return The result Column reference
+   */
+  public static Column explode(Column input) {
+    return new Column(com.snowflake.snowpark.tableFunctions.explode(input.toScalaColumn()));
+  }
 }
diff --git a/src/main/scala/com/snowflake/snowpark/tableFunctions.scala b/src/main/scala/com/snowflake/snowpark/tableFunctions.scala
@@ -216,7 +216,8 @@ object tableFunctions {
    * }}}
    *
    * @since 1.10.0
-   * @param input an array or map column
+   * @param input The expression that will be unseated into rows.
+   *              The expression must be MapType or ArrayType data.
    * @return The result Column reference
    */
   def explode(input: Column): Column = TableFunction("explode").apply(input)

diff --git a/src/test/java/com/snowflake/snowpark_test/JavaTableFunctionSuite.java b/src/test/java/com/snowflake/snowpark_test/JavaTableFunctionSuite.java
@@ -121,4 +121,63 @@ public void argumentInFlatten() {
             .select("value"),
         new Row[] {Row.create("1"), Row.create("2")});
   }
+
+  @Test
+  public void explodeWithDataFrame() {
+    // select
+    DataFrame df =
+        getSession()
+            .createDataFrame(
+                new Row[] {Row.create("{\"a\":1, \"b\":2}")},
+                StructType.create(new StructField("col", DataTypes.StringType)));
+    DataFrame df1 =
+        df.select(
+            Functions.parse_json(df.col("col"))
+                .cast(DataTypes.createMapType(DataTypes.StringType, DataTypes.IntegerType))
+                .as("col"));
+    checkAnswer(
+        df1.select(TableFunctions.explode(df1.col("col"))),
+        new Row[] {Row.create("a", "1"), Row.create("b", "2")});
+    // join
+    df =
+        getSession()
+            .createDataFrame(
+                new Row[] {Row.create("[1, 2]")},
+                StructType.create(new StructField("col", DataTypes.StringType)));
+    df1 =
+        df.select(
+            Functions.parse_json(df.col("col"))
+                .cast(DataTypes.createArrayType(DataTypes.IntegerType))
+                .as("col"));
+    checkAnswer(
+        df1.join(TableFunctions.explode(df1.col("col"))).select("VALUE"),
+        new Row[] {Row.create("1"), Row.create("2")});
+  }
+
+  @Test
+  public void explodeWithSession() {
+    DataFrame df =
+        getSession()
+            .createDataFrame(
+                new Row[] {Row.create("{\"a\":1, \"b\":2}")},
+                StructType.create(new StructField("col", DataTypes.StringType)));
+    DataFrame df1 =
+        df.select(
+            Functions.parse_json(df.col("col"))
+                .cast(DataTypes.createMapType(DataTypes.StringType, DataTypes.IntegerType))
+                .as("col"));
+    checkAnswer(
+        getSession().tableFunction(TableFunctions.explode(df1.col("col"))).select("KEY", "VALUE"),
+        new Row[] {Row.create("a", "1"), Row.create("b", "2")});
+
+    checkAnswer(
+        getSession()
+            .tableFunction(
+                TableFunctions.explode(
+                    Functions.parse_json(Functions.lit("{\"a\":1, \"b\":2}"))
+                        .cast(
+                            DataTypes.createMapType(DataTypes.StringType, DataTypes.IntegerType))))
+            .select("KEY", "VALUE"),
+        new Row[] {Row.create("a", "1"), Row.create("b", "2")});
+  }
 }