SNOW-966358 Support TableFunction in DataFrame Select()

src/main/java/com/snowflake/snowpark_java/DataFrame.java

@@ -1312,6 +1312,81 @@ public DataFrame join(
             JavaUtils.columnArrayToSeq(Column.toScalaColumnArray(orderBy))));
   }
 
+  /**
+   * Joins the current DataFrame with the output of the specified table function `func`.
+   *
+   * <p>Pre-defined table functions can be found in `TableFunctions` class.
+   *
+   * <p>For example:
+   *
+   * <pre>{@code
+   * df.join(TableFunctions.flatten(
+   *   Functions.parse_json(df.col("col")),
+   *   "path", true, true, "both"
+   * ));
+   * }</pre>
+   *
+   * <p>Or load any Snowflake builtin table function via TableFunction Class.
+   *
+   * <pre>{@code
+   * Map<String, Column> args = new HashMap<>();
+   * args.put("input", Functions.parse_json(df.col("a")));
+   * df.join(new TableFunction("flatten").call(args));
+   * }</pre>
+   *
+   * @since 1.10.0
+   * @param func Column object, which can be one of the values in the TableFunctions class or an
+   *     object that you create from the `new TableFunction("name").call()`.
+   * @return The result DataFrame
+   */
+  public DataFrame join(Column func) {
+    return new DataFrame(this.df.join(func.toScalaColumn()));
+  }
+
+  /**
+   * Joins the current DataFrame with the output of the specified table function `func`.
+   *
+   * <p>To specify a PARTITION BY or ORDER BY clause, use the `partitionBy` and `orderBy` arguments.
+   *
+   * <p>Pre-defined table functions can be found in `TableFunctions` class.
+   *
+   * <p>For example:
+   *
+   * <pre>{@code
+   * df.join(TableFunctions.flatten(
+   *     Functions.parse_json(df.col("col1")),
+   *     "path", true, true, "both"
+   *   ),
+   *   new Column[] {df.col("col2")},
+   *   new Column[] {df.col("col1")}
+   * );
+   * }</pre>
+   *
+   * <p>Or load any Snowflake builtin table function via TableFunction Class.
+   *
+   * <pre>{@code
+   * Map<String, Column> args = new HashMap<>();
+   * args.put("input", Functions.parse_json(df.col("col1")));
+   * df.join(new TableFunction("flatten").call(args),
+   * new Column[] {df.col("col2")},
+   * new Column[] {df.col("col1")});
+   * }</pre>
+   *
+   * @since 1.10.0
+   * @param func Column object, which can be one of the values in the TableFunctions class or an
+   *     object that you create from the `new TableFunction("name").call()`.
+   * @param partitionBy An array of columns partitioned by.
+   * @param orderBy An array of columns ordered by.
+   * @return The result DataFrame
+   */
+  public DataFrame join(Column func, Column[] partitionBy, Column[] orderBy) {
+    return new DataFrame(
+        this.df.join(
+            func.toScalaColumn(),
+            JavaUtils.columnArrayToSeq(Column.toScalaColumnArray(partitionBy)),
+            JavaUtils.columnArrayToSeq(Column.toScalaColumnArray(orderBy))));
+  }
+
   com.snowflake.snowpark.DataFrame getScalaDataFrame() {
     return this.df;
   }

src/main/java/com/snowflake/snowpark_java/Session.java

@@ -566,6 +566,27 @@ public DataFrame tableFunction(TableFunction func, Map<String, Column> args) {
             func.getScalaTableFunction(), JavaUtils.javaStringColumnMapToScala(scalaArgs)));
   }
 
+  /**
+   * Creates a new DataFrame from the given table function and arguments.
+   *
+   * <p>Example
+   *
+   * <pre>{@code
+   * session.tableFunction(TableFunctions.flatten(
+   *   Functions.parse_json(df.col("col")),
+   *   "path", true, true, "both"
+   * ));
+   * }</pre>
+   *
+   * @since 1.10.0
+   * @param func Column object, which can be one of the values in the TableFunctions class or an
+   *     object that you create from the `new TableFunction("name").call()`.
+   * @return The result DataFrame
+   */
+  public DataFrame tableFunction(Column func) {
+    return new DataFrame(session.tableFunction(func.toScalaColumn()));
+  }
+
   /**
    * Returns a SProcRegistration object that you can use to register Stored Procedures.
    *

src/main/java/com/snowflake/snowpark_java/TableFunction.java

@@ -1,5 +1,9 @@
 package com.snowflake.snowpark_java;
 
+import com.snowflake.snowpark.internal.JavaUtils;
+import java.util.HashMap;
+import java.util.Map;
+
 /**
  * Looks up table functions by funcName and returns tableFunction object which can be used in {@code
  * DataFrame.join} and {@code Session.tableFunction} methods.
@@ -38,4 +42,31 @@ com.snowflake.snowpark.TableFunction getScalaTableFunction() {
   public String funcName() {
     return func.funcName();
   }
+
+  /**
+   * Create a Column reference by passing arguments in the TableFunction object.
+   *
+   * @param args A list of Column objects representing the arguments of the given table function
+   * @return A Column reference
+   * @since 1.10.0
+   */
+  public Column call(Column... args) {
+    return new Column(this.func.apply(JavaUtils.columnArrayToSeq(Column.toScalaColumnArray(args))));
+  }
+
+  /**
+   * Create a Column reference by passing arguments in the TableFunction object.
+   *
+   * @param args function arguments map of the given table function. Some functions, like flatten,
+   *     have named parameters. use this map to assign values to the corresponding parameters.
+   * @return A Column reference
+   * @since 1.10.0
+   */
+  public Column call(Map<String, Column> args) {
+    Map<String, com.snowflake.snowpark.Column> scalaArgs = new HashMap<>();
+    for (Map.Entry<String, Column> entry : args.entrySet()) {
+      scalaArgs.put(entry.getKey(), entry.getValue().toScalaColumn());
+    }
+    return new Column(this.func.apply(JavaUtils.javaStringColumnMapToScala(scalaArgs)));
+  }
 }

src/main/java/com/snowflake/snowpark_java/TableFunctions.java

@@ -40,6 +40,27 @@ public static TableFunction split_to_table() {
     return new TableFunction(com.snowflake.snowpark.tableFunctions.split_to_table());
   }
 
+  /**
+   * This table function splits a string (based on a specified delimiter) and flattens the results
+   * into rows.
+   *
+   * <p>Example
+   *
+   * <pre>{@code
+   * session.tableFunction(TableFunctions.split_to_table(,
+   *   Functions.lit("split by space"), Functions.lit(" ")));
+   * }</pre>
+   *
+   * @since 1.10.0
+   * @param str Text to be split.
+   * @param delimiter Text to split string by.
+   * @return The result Column reference
+   */
+  public static Column split_to_table(Column str, String delimiter) {
+    return new Column(
+        com.snowflake.snowpark.tableFunctions.split_to_table(str.toScalaColumn(), delimiter));
+  }
+
   /**
    * Flattens (explodes) compound values into multiple rows.
    *
@@ -77,4 +98,56 @@ public static TableFunction split_to_table() {
   public static TableFunction flatten() {
     return new TableFunction(com.snowflake.snowpark.tableFunctions.flatten());
   }
+
+  /**
+   * Flattens (explodes) compound values into multiple rows.
+   *
+   * <p>Example
+   *
+   * <pre>{@code
+   * df.join(TableFunctions.flatten(
+   *   Functions.parse_json(df.col("col")), "path", true, true, "both"));
+   * }</pre>
+   *
+   * @since 1.10.0
+   * @param input The expression that will be unseated into rows. The expression must be of data
+   *     type VariantType, MapType or ArrayType.
+   * @param path The path to the element within a VariantType data structure which needs to be
+   *     flattened. Can be a zero-length string (i.e. empty path) if the outermost element is to be
+   *     flattened. Default: Zero-length string (i.e. empty path)
+   * @param outer If FALSE, any input rows that cannot be expanded, either because they cannot be
+   *     accessed in the path or because they have zero fields or entries, are completely omitted
+   *     from the output. If TRUE, exactly one row is generated for zero-row expansions (with NULL
+   *     in the KEY, INDEX, and VALUE columns).
+   * @param recursive If FALSE, only the element referenced by PATH is expanded. If TRUE, the
+   *     expansion is performed for all sub-elements recursively. Default: FALSE
+   * @param mode ("object", "array", or "both") Specifies whether only objects, arrays, or both
+   *     should be flattened.
+   * @return The result Column reference
+   */
+  public static Column flatten(
+      Column input, String path, boolean outer, boolean recursive, String mode) {
+    return new Column(
+        com.snowflake.snowpark.tableFunctions.flatten(
+            input.toScalaColumn(), path, outer, recursive, mode));
+  }
+
+  /**
+   * Flattens (explodes) compound values into multiple rows.
+   *
+   * <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 of data
+   *     type VariantType, MapType or ArrayType.
+   * @return The result Column reference
+   */
+  public static Column flatten(Column input) {
+    return new Column(com.snowflake.snowpark.tableFunctions.flatten(input.toScalaColumn()));
+  }
 }