Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sfc-gh-batur merged 4 commits into from
Jan 14, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
bb7b106
changes
sfc-gh-batur Jan 13, 2025
a171b24
updating ast.proto
sfc-gh-batur Jan 13, 2025
66feaf7
moar changes
sfc-gh-batur Jan 13, 2025
5b9564e
Merge remote-tracking branch 'origin/main' into SNOW-1800374
sfc-gh-batur Jan 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 36 additions & 34 deletions src/snowflake/snowpark/_internal/proto/ast.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -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;
}
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -363,7 +363,7 @@ message SrcPosition {
int64 start_line = 5;
}

// sp-type.ir:55
// sp-type.ir:56
message StagedPandasDataframe {
SpNameRef temp_table = 1;
}
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand Down
24 changes: 22 additions & 2 deletions src/snowflake/snowpark/dataframe_writer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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":
Expand Down
Loading
Loading