Add AST support

Author: sfc-gh-mayliu

src/snowflake/snowpark/_internal/proto/ast.proto

@@ -1,5 +1,5 @@
 // N.B. This file is generated by `//ir-dsl-c`. DO NOT EDIT!
-// Generated from `{[email protected]:snowflakedb/snowflake.git}/Snowpark/ast`.
+// Generated from `{[email protected]:snowflake-eng/snowflake.git}/Snowpark/ast`.
 
 syntax = "proto3";
 
@@ -987,7 +987,7 @@ message DataframeCollect {
   repeated Tuple_String_String statement_params = 7;
 }
 
-// dataframe-io.ir:165
+// dataframe-io.ir:167
 message DataframeCopyIntoTable {
   repeated Tuple_String_Expr copy_options = 1;
   Expr df = 2;
@@ -1011,7 +1011,7 @@ message DataframeCount {
   repeated Tuple_String_String statement_params = 4;
 }
 
-// dataframe-io.ir:148
+// dataframe-io.ir:150
 message DataframeCreateOrReplaceDynamicTable {
   repeated Expr clustering_keys = 1;
   google.protobuf.StringValue comment = 2;
@@ -1030,7 +1030,7 @@ message DataframeCreateOrReplaceDynamicTable {
   string warehouse = 15;
 }
 
-// dataframe-io.ir:139
+// dataframe-io.ir:141
 message DataframeCreateOrReplaceView {
   google.protobuf.StringValue comment = 1;
   bool copy_grants = 2;
@@ -2685,7 +2685,7 @@ message WindowSpecRowsBetween {
   WindowSpecExpr wnd = 4;
 }
 
-// dataframe-io.ir:116
+// dataframe-io.ir:118
 message WriteCopyIntoLocation {
   bool block = 1;
   repeated Tuple_String_Expr copy_options = 2;
@@ -2704,7 +2704,7 @@ message WriteCopyIntoLocation {
   Expr writer = 15;
 }
 
-// dataframe-io.ir:123
+// dataframe-io.ir:125
 message WriteCsv {
   bool block = 1;
   repeated Tuple_String_Expr copy_options = 2;
@@ -2731,15 +2731,15 @@ message WriteFile {
   }
 }
 
-// dataframe-io.ir:129
+// dataframe-io.ir:131
 message WriteInsertInto {
   bool overwrite = 1;
   SrcPosition src = 2;
   NameRef table_name = 3;
   Expr writer = 4;
 }
 
-// dataframe-io.ir:125
+// dataframe-io.ir:127
 message WriteJson {
   bool block = 1;
   repeated Tuple_String_Expr copy_options = 2;
@@ -2773,7 +2773,7 @@ message WritePandas {
   string table_type = 13;
 }
 
-// dataframe-io.ir:127
+// dataframe-io.ir:129
 message WriteParquet {
   bool block = 1;
   repeated Tuple_String_Expr copy_options = 2;
@@ -2790,7 +2790,7 @@ message WriteParquet {
   Expr writer = 13;
 }
 
-// dataframe-io.ir:121
+// dataframe-io.ir:123
 message WriteSave {
   bool block = 1;
   repeated Tuple_String_Expr copy_options = 2;
@@ -2821,9 +2821,11 @@ message WriteTable {
   repeated Tuple_String_Expr iceberg_config = 10;
   google.protobuf.Int64Value max_data_extension_time = 11;
   SaveMode mode = 12;
-  SrcPosition src = 13;
-  repeated Tuple_String_String statement_params = 14;
-  NameRef table_name = 15;
-  string table_type = 16;
-  Expr writer = 17;
+  Expr overwrite_condition = 13;
+  SrcPosition src = 14;
+  repeated Tuple_String_String statement_params = 15;
+  google.protobuf.BoolValue table_exists = 16;
+  NameRef table_name = 17;
+  string table_type = 18;
+  Expr writer = 19;
 }

src/snowflake/snowpark/dataframe_writer.py

@@ -411,6 +411,8 @@ def save_as_table(
             # change_tracking: Optional[bool] = None,
             # copy_grants: bool = False,
             # iceberg_config: Optional[dict] = None,
+            # table_exists: Optional[bool] = None,
+            # overwrite_condition: Optional[ColumnOrSqlExpr] = None,
 
             build_table_name(expr.table_name, table_name)
 
@@ -452,6 +454,12 @@ def save_as_table(
                     t = expr.iceberg_config.add()
                     t._1 = k
                     build_expr_from_python_val(t._2, v)
+            if table_exists is not None:
+                expr.table_exists.value = table_exists
+            if overwrite_condition is not None:
+                build_expr_from_snowpark_column_or_sql_str(
+                    expr.overwrite_condition, overwrite_condition
+                )
 
             self._dataframe._session._ast_batch.eval(stmt)