SNOW-1800374: adding support for options and partition_by to dataframewriter (#2841)

sfc-gh-batur · web-flow · commit 2f06485fec05 · 2025-01-14T14:02:02.000-08:00
1. Which Jira issue is this PR addressing? Make sure that there is an accompanying issue to your PR.  Fixes SNOW-1800374 2. Fill out the following pre-review checklist: - [x] I am adding a new automated test(s) to verify correctness of my new code - [ ] If this test skips Local Testing mode, I'm requesting review from @snowflakedb/local-testing - [ ] I am adding new logging messages - [ ] I am adding a new telemetry message - [ ] I am adding new credentials - [ ] I am adding a new dependency - [ ] If this is a new feature/behavior, I'm adding the Local Testing parity changes. - [x] I acknowledge that I have ensured my changes to be thread-safe. Follow the link for more information: [Thread-safe Developer Guidelines](https://github.com/snowflakedb/snowpark-python/blob/main/CONTRIBUTING.md#thread-safe-development) 3. Please describe how your code solves the related issue. Please write a short description of how your code change solves the related issue.
diff --git a/src/snowflake/snowpark/_internal/proto/ast.proto b/src/snowflake/snowpark/_internal/proto/ast.proto
@@ -89,7 +89,7 @@ message PythonTimeZone {
   int64 offset_seconds = 2;
 }
 
-// sp-type.ir:70
+// sp-type.ir:71
 message SpCallable {
   int64 id = 1;
   string name = 2;
@@ -136,71 +136,71 @@ message SpDataType {
   }
 }
 
-// sp-type.ir:27
+// sp-type.ir:28
 message SpArrayType {
   bool structured = 1;
   SpDataType ty = 2;
 }
 
-// sp-type.ir:31
+// sp-type.ir:32
 message SpColumnIdentifier {
   string name = 1;
 }
 
-// sp-type.ir:33
+// sp-type.ir:34
 message SpDecimalType {
   int64 precision = 1;
   int64 scale = 2;
 }
 
-// sp-type.ir:40
+// sp-type.ir:41
 message SpMapType {
   SpDataType key_ty = 1;
   bool structured = 2;
   SpDataType value_ty = 3;
 }
 
-// sp-type.ir:43
+// sp-type.ir:44
 message SpStringType {
   google.protobuf.Int64Value length = 1;
 }
 
-// sp-type.ir:44
+// sp-type.ir:45
 message SpStructField {
   SpColumnIdentifier column_identifier = 1;
   SpDataType data_type = 2;
   bool nullable = 3;
 }
 
-// sp-type.ir:45
+// sp-type.ir:46
 message SpStructType {
   repeated SpStructField fields = 1;
   bool structured = 2;
 }
 
-// sp-type.ir:47
+// sp-type.ir:48
 message SpTimestampType {
   SpTimestampTimeZone time_zone = 1;
 }
 
-// sp-type.ir:49
+// sp-type.ir:50
 message SpVectorType {
   int64 dimension = 1;
   SpDataType ty = 2;
 }
 
-// sp-type.ir:51
+// sp-type.ir:52
 message SpPandasSeriesType {
   SpDataType el_ty = 1;
 }
 
-// sp-type.ir:52
+// sp-type.ir:53
 message SpPandasDataFrameType {
   repeated string col_names = 1;
   repeated SpDataType col_types = 2;
 }
 
-// sp-type.ir:59
+// sp-type.ir:60
 message SpDataframeData {
   oneof sealed_value {
     SpDataframeData_List sp_dataframe_data__list = 1;
@@ -209,35 +209,35 @@ message SpDataframeData {
   }
 }
 
-// sp-type.ir:60
+// sp-type.ir:61
 message SpDataframeData_List {
   repeated Expr vs = 1;
 }
 
-// sp-type.ir:61
+// sp-type.ir:62
 message SpDataframeData_Tuple {
   repeated Expr vs = 1;
 }
 
-// sp-type.ir:62
+// sp-type.ir:63
 message SpDataframeData_Pandas {
   StagedPandasDataframe v = 1;
 }
 
-// sp-type.ir:65
+// sp-type.ir:66
 message SpDataframeSchema {
   oneof sealed_value {
     SpDataframeSchema_List sp_dataframe_schema__list = 1;
     SpDataframeSchema_Struct sp_dataframe_schema__struct = 2;
   }
 }
 
-// sp-type.ir:66
+// sp-type.ir:67
 message SpDataframeSchema_List {
   repeated string vs = 1;
 }
 
-// sp-type.ir:67
+// sp-type.ir:68
 message SpDataframeSchema_Struct {
   SpStructType v = 1;
 }
@@ -292,20 +292,20 @@ message SpNullOrder {
   }
 }
 
-// sp-type.ir:82
+// sp-type.ir:83
 message SpPivotValue {
   oneof sealed_value {
     SpPivotValue_Dataframe sp_pivot_value__dataframe = 1;
     SpPivotValue_Expr sp_pivot_value__expr = 2;
   }
 }
 
-// sp-type.ir:83
+// sp-type.ir:84
 message SpPivotValue_Expr {
   Expr v = 1;
 }
 
-// sp-type.ir:84
+// sp-type.ir:85
 message SpPivotValue_Dataframe {
   SpDataframeRef v = 1;
 }
@@ -363,7 +363,7 @@ message SrcPosition {
   int64 start_line = 5;
 }
 
-// sp-type.ir:55
+// sp-type.ir:56
 message StagedPandasDataframe {
   SpNameRef temp_table = 1;
 }
@@ -1585,7 +1585,7 @@ message SpDataframeApply {
   SrcPosition src = 3;
 }
 
-// sp-df-io.ir:185
+// sp-df-io.ir:187
 message SpDataframeCacheResult {
   SpDataframeExpr df = 1;
   SrcPosition src = 2;
@@ -1610,7 +1610,7 @@ message SpDataframeCollect {
   repeated Tuple_String_String statement_params = 7;
 }
 
-// sp-df-io.ir:167
+// sp-df-io.ir:169
 message SpDataframeCopyIntoTable {
   repeated Tuple_String_Expr copy_options = 1;
   SpDataframeExpr df = 2;
@@ -1634,7 +1634,7 @@ message SpDataframeCount {
   repeated Tuple_String_String statement_params = 4;
 }
 
-// sp-df-io.ir:151
+// sp-df-io.ir:153
 message SpDataframeCreateOrReplaceDynamicTable {
   List_Expr clustering_keys = 1;
   google.protobuf.StringValue comment = 2;
@@ -1652,7 +1652,7 @@ message SpDataframeCreateOrReplaceDynamicTable {
   string warehouse = 14;
 }
 
-// sp-df-io.ir:143
+// sp-df-io.ir:145
 message SpDataframeCreateOrReplaceView {
   google.protobuf.StringValue comment = 1;
   SpDataframeExpr df = 2;
@@ -2248,8 +2248,10 @@ message SpDataframeWithColumns {
 // sp-df-io.ir:84
 message SpDataframeWrite {
   SpDataframeExpr df = 1;
-  SpSaveMode save_mode = 2;
-  SrcPosition src = 3;
+  repeated Tuple_String_Expr options = 2;
+  Expr partition_by = 3;
+  SpSaveMode save_mode = 4;
+  SrcPosition src = 5;
 }
 
 message SpDataframeWriter {
@@ -2609,7 +2611,7 @@ message SpWindowSpecRowsBetween {
 message SpWindowType {
 }
 
-// sp-df-io.ir:133
+// sp-df-io.ir:135
 message SpWriteCopyIntoLocation {
   bool block = 1;
   repeated Tuple_String_Expr copy_options = 2;
@@ -2624,7 +2626,7 @@ message SpWriteCopyIntoLocation {
   repeated Tuple_String_String statement_params = 11;
 }
 
-// sp-df-io.ir:100
+// sp-df-io.ir:102
 message SpWriteCsv {
   bool block = 1;
   repeated Tuple_String_Expr copy_options = 2;
@@ -2646,7 +2648,7 @@ message SpWriteFile {
   }
 }
 
-// sp-df-io.ir:104
+// sp-df-io.ir:106
 message SpWriteJson {
   bool block = 1;
   repeated Tuple_String_Expr copy_options = 2;
@@ -2676,7 +2678,7 @@ message SpWritePandas {
   string table_type = 13;
 }
 
-// sp-df-io.ir:108
+// sp-df-io.ir:110
 message SpWriteParquet {
   bool block = 1;
   repeated Tuple_String_Expr copy_options = 2;
@@ -2689,7 +2691,7 @@ message SpWriteParquet {
   repeated Tuple_String_String statement_params = 9;
 }
 
-// sp-df-io.ir:112
+// sp-df-io.ir:114
 message SpWriteTable {
   bool block = 1;
   google.protobuf.BoolValue change_tracking = 2;
diff --git a/src/snowflake/snowpark/dataframe_writer.py b/src/snowflake/snowpark/dataframe_writer.py
@@ -16,6 +16,8 @@
 )
 from snowflake.snowpark._internal.ast.utils import (
     build_expr_from_snowpark_column_or_col_name,
+    build_expr_from_snowpark_column_or_sql_str,
+    build_expr_from_snowpark_column_or_python_val,
     debug_check_missing_ast,
     fill_sp_save_mode,
     fill_sp_write_file,
@@ -130,19 +132,37 @@ def mode(self, save_mode: str, _emit_ast: bool = True) -> "DataFrameWriter":
 
         return self
 
-    def partition_by(self, expr: ColumnOrSqlExpr) -> "DataFrameWriter":
+    def partition_by(
+        self, expr: ColumnOrSqlExpr, _emit_ast: bool = True
+    ) -> "DataFrameWriter":
         """Specifies an expression used to partition the unloaded table rows into separate files. It can be a
         :class:`Column`, a column name, or a SQL expression.
         """
         self._partition_by = expr
+
+        # Update AST if it exists.
+        if _emit_ast:
+            if self._ast_stmt is not None:
+                build_expr_from_snowpark_column_or_sql_str(
+                    self._ast_stmt.expr.sp_dataframe_write.partition_by, expr
+                )
+
         return self
 
-    def option(self, key: str, value: Any) -> "DataFrameWriter":
+    def option(self, key: str, value: Any, _emit_ast: bool = True) -> "DataFrameWriter":
         """Depending on the ``file_format_type`` specified, you can include more format specific options.
         Use the options documented in the `Format Type Options <https://docs.snowflake.com/en/sql-reference/sql/copy-into-location.html#format-type-options-formattypeoptions>`__.
         """
         aliased_key = get_aliased_option_name(key, WRITER_OPTIONS_ALIAS_MAP)
         self._cur_options[aliased_key] = value
+
+        # Update AST if it exists.
+        if _emit_ast:
+            if self._ast_stmt is not None:
+                t = self._ast_stmt.expr.sp_dataframe_write.options.add()
+                t._1 = aliased_key
+                build_expr_from_snowpark_column_or_python_val(t._2, value)
+
         return self
 
     def options(self, configs: Optional[Dict] = None, **kwargs) -> "DataFrameWriter":
diff --git a/tests/ast/data/DataFrame.write.test b/tests/ast/data/DataFrame.write.test
