SNOW-1552377: Add support for flatmap using udtf

CHANGELOG.md

@@ -15,6 +15,11 @@
   - `nullifzero`
   - `snowflake_cortex_sentiment`
 - Added `Catalog` class to manage snowflake objects. It can be accessed via `Session.catalog`.
+- Added new methods in class `DataFrame`:
+  - `col_regex`: Select columns that match with provided regex.
+  - `map` and its alias `foreach`: A method to apply user function on each row with 1-1 mapping.
+  - `flat_map`: A method to apply user function on each row with one to many mapping.
+  - `toJSON` and its alias `to_json`: Convert each row of dataframe into json string.
 
 #### Improvements
 
@@ -63,8 +68,6 @@
 #### New Features
 
 - Added support for property `version` and class method `get_active_session` for `Session` class.
-- Added new methods in class `DataFrame`:
-  - `col_regex`: Select columns that match with provided regex.
 - Added support for property `version` and class method `get_active_session` for `Session` class.
 - Added new methods and variables to enhance data type handling and JSON serialization/deserialization:
   - To `DataType`, its derived classes, and `StructField`:
@@ -84,7 +87,6 @@
   - `collect_list` an alias of `array_agg`.
   - `substring` makes `len` argument optional.
 - Added parameter `ast_enabled` to session for internal usage (default: `False`).
-- Added support for `Dataframe.toJSON`
 
 #### Improvements
 

docs/source/snowpark/dataframe.rst

@@ -50,6 +50,8 @@ DataFrame
     DataFrame.fillna
     DataFrame.filter
     DataFrame.first
+    DataFrame.flat_map
+    DataFrame.flatMap
     DataFrame.flatten
     DataFrame.groupBy
     DataFrame.group_by

src/snowflake/snowpark/_internal/analyzer/analyzer.py

@@ -325,7 +325,7 @@ def analyze(
         if isinstance(expr, WindowSpecDefinition):
             return window_spec_expression(
                 [
-                    self.analyze(
+                    self.to_sql_try_avoid_cast(
                         x, df_aliased_col_name_to_real_col_name, parse_local_name
                     )
                     for x in expr.partition_spec
@@ -456,7 +456,7 @@ def analyze(
             return table_function_partition_spec(
                 expr.over,
                 [
-                    self.analyze(
+                    self.to_sql_try_avoid_cast(
                         x, df_aliased_col_name_to_real_col_name, parse_local_name
                     )
                     for x in expr.partition_spec
@@ -623,7 +623,9 @@ def table_function_expression_extractor(
                 "NamedArgumentsTableFunction, GeneratorTableFunction, or FlattenFunction."
             )
         partition_spec_sql = (
-            self.analyze(expr.partition_spec, df_aliased_col_name_to_real_col_name)
+            self.to_sql_try_avoid_cast(
+                expr.partition_spec, df_aliased_col_name_to_real_col_name
+            )
             if expr.partition_spec
             else ""
         )