Skip to content

Commit db0b32c

Browse files
sfc-gh-heshahsfc-gh-heshah
authored
SNOW-1916224: Ensure that all APIs have AST support (#3170)
1 parent 61982e0 commit db0b32c

27 files changed

+10792
-554
lines changed

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

Lines changed: 45 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -223,7 +223,7 @@ message FlattenMode {
223223
}
224224
}
225225

226-
// dataframe.ir:230
226+
// dataframe.ir:232
227227
message JoinType {
228228
oneof variant {
229229
bool join_type__asof = 1;
@@ -741,14 +741,14 @@ message CreateDataframe {
741741
SrcPosition src = 3;
742742
}
743743

744-
// dataframe.ir:163
744+
// dataframe.ir:165
745745
message DataframeAgg {
746746
Expr df = 1;
747747
ExprArgList exprs = 2;
748748
SrcPosition src = 3;
749749
}
750750

751-
// dataframe.ir:168
751+
// dataframe.ir:170
752752
message DataframeAlias {
753753
Expr df = 1;
754754
string name = 2;
@@ -888,7 +888,7 @@ message DataframeCreateOrReplaceView {
888888
repeated Tuple_String_String statement_params = 6;
889889
}
890890

891-
// dataframe.ir:173
891+
// dataframe.ir:175
892892
message DataframeCrossJoin {
893893
Expr lhs = 1;
894894
google.protobuf.StringValue lsuffix = 2;
@@ -904,48 +904,48 @@ message DataframeCube {
904904
SrcPosition src = 3;
905905
}
906906

907-
// dataframe.ir:180
907+
// dataframe.ir:182
908908
message DataframeDescribe {
909909
ExprArgList cols = 1;
910910
Expr df = 2;
911911
SrcPosition src = 3;
912912
}
913913

914-
// dataframe.ir:185
914+
// dataframe.ir:187
915915
message DataframeDistinct {
916916
Expr df = 1;
917917
SrcPosition src = 2;
918918
}
919919

920-
// dataframe.ir:189
920+
// dataframe.ir:191
921921
message DataframeDrop {
922922
ExprArgList cols = 1;
923923
Expr df = 2;
924924
SrcPosition src = 3;
925925
}
926926

927-
// dataframe.ir:194
927+
// dataframe.ir:196
928928
message DataframeDropDuplicates {
929929
ExprArgList cols = 1;
930930
Expr df = 2;
931931
SrcPosition src = 3;
932932
}
933933

934-
// dataframe.ir:199
934+
// dataframe.ir:201
935935
message DataframeExcept {
936936
Expr df = 1;
937937
Expr other = 2;
938938
SrcPosition src = 3;
939939
}
940940

941-
// dataframe.ir:204
941+
// dataframe.ir:206
942942
message DataframeFilter {
943943
Expr condition = 1;
944944
Expr df = 2;
945945
SrcPosition src = 3;
946946
}
947947

948-
// dataframe.ir:218
948+
// dataframe.ir:220
949949
message DataframeFirst {
950950
bool block = 1;
951951
Expr df = 2;
@@ -954,7 +954,7 @@ message DataframeFirst {
954954
repeated Tuple_String_String statement_params = 5;
955955
}
956956

957-
// dataframe.ir:209
957+
// dataframe.ir:211
958958
message DataframeFlatten {
959959
Expr df = 1;
960960
Expr input = 2;
@@ -979,14 +979,14 @@ message DataframeGroupByGroupingSets {
979979
SrcPosition src = 3;
980980
}
981981

982-
// dataframe.ir:225
982+
// dataframe.ir:227
983983
message DataframeIntersect {
984984
Expr df = 1;
985985
Expr other = 2;
986986
SrcPosition src = 3;
987987
}
988988

989-
// dataframe.ir:241
989+
// dataframe.ir:243
990990
message DataframeJoin {
991991
Expr join_expr = 1;
992992
JoinType join_type = 2;
@@ -998,14 +998,14 @@ message DataframeJoin {
998998
SrcPosition src = 8;
999999
}
10001000

1001-
// dataframe.ir:251
1001+
// dataframe.ir:253
10021002
message DataframeJoinTableFunction {
10031003
Expr fn = 1;
10041004
Expr lhs = 2;
10051005
SrcPosition src = 3;
10061006
}
10071007

1008-
// dataframe.ir:256
1008+
// dataframe.ir:258
10091009
message DataframeLimit {
10101010
Expr df = 1;
10111011
int64 n = 2;
@@ -1033,25 +1033,27 @@ message DataframeNaDrop_Scala {
10331033
// dataframe.ir:142
10341034
message DataframeNaFill {
10351035
Expr df = 1;
1036-
SrcPosition src = 2;
1037-
ExprArgList subset = 3;
1038-
Expr value = 4;
1039-
repeated Tuple_String_Expr value_map = 5;
1036+
bool include_decimal = 2;
1037+
SrcPosition src = 3;
1038+
ExprArgList subset = 4;
1039+
Expr value = 5;
1040+
repeated Tuple_String_Expr value_map = 6;
10401041
}
10411042

1042-
// dataframe.ir:149
1043+
// dataframe.ir:150
10431044
message DataframeNaReplace {
10441045
Expr df = 1;
1045-
repeated Tuple_Expr_Expr replacement_map = 2;
1046-
SrcPosition src = 3;
1047-
ExprArgList subset = 4;
1048-
repeated Expr to_replace_list = 5;
1049-
Expr to_replace_value = 6;
1050-
Expr value = 7;
1051-
repeated Expr values = 8;
1046+
bool include_decimal = 2;
1047+
repeated Tuple_Expr_Expr replacement_map = 3;
1048+
SrcPosition src = 4;
1049+
ExprArgList subset = 5;
1050+
repeated Expr to_replace_list = 6;
1051+
Expr to_replace_value = 7;
1052+
Expr value = 8;
1053+
repeated Expr values = 9;
10521054
}
10531055

1054-
// dataframe.ir:262
1056+
// dataframe.ir:264
10551057
message DataframeNaturalJoin {
10561058
JoinType join_type = 1;
10571059
Expr lhs = 2;
@@ -1068,7 +1070,7 @@ message DataframePivot {
10681070
PivotValue values = 5;
10691071
}
10701072

1071-
// dataframe.ir:276
1073+
// dataframe.ir:278
10721074
message DataframeRandomSplit {
10731075
Expr df = 1;
10741076
google.protobuf.Int64Value seed = 2;
@@ -1117,7 +1119,7 @@ message DataframeRef {
11171119
SrcPosition src = 2;
11181120
}
11191121

1120-
// dataframe.ir:283
1122+
// dataframe.ir:285
11211123
message DataframeRename {
11221124
Expr col_or_mapper = 1;
11231125
Expr df = 2;
@@ -1132,7 +1134,7 @@ message DataframeRollup {
11321134
SrcPosition src = 3;
11331135
}
11341136

1135-
// dataframe.ir:289
1137+
// dataframe.ir:291
11361138
message DataframeSample {
11371139
Expr df = 1;
11381140
google.protobuf.Int64Value num = 2;
@@ -1141,7 +1143,7 @@ message DataframeSample {
11411143
SrcPosition src = 5;
11421144
}
11431145

1144-
// dataframe.ir:296
1146+
// dataframe.ir:298
11451147
message DataframeSelect {
11461148
ExprArgList cols = 1;
11471149
Expr df = 2;
@@ -1156,7 +1158,7 @@ message DataframeShow {
11561158
SrcPosition src = 3;
11571159
}
11581160

1159-
// dataframe.ir:303
1161+
// dataframe.ir:305
11601162
message DataframeSort {
11611163
Expr ascending = 1;
11621164
ExprArgList cols = 2;
@@ -1240,7 +1242,7 @@ message DataframeToPandasBatches {
12401242
repeated Tuple_String_String statement_params = 4;
12411243
}
12421244

1243-
// dataframe.ir:309
1245+
// dataframe.ir:311
12441246
message DataframeUnion {
12451247
bool all = 1;
12461248
bool allow_missing_columns = 2;
@@ -1250,7 +1252,7 @@ message DataframeUnion {
12501252
SrcPosition src = 6;
12511253
}
12521254

1253-
// dataframe.ir:268
1255+
// dataframe.ir:270
12541256
message DataframeUnpivot {
12551257
repeated Expr column_list = 1;
12561258
Expr df = 2;
@@ -1260,23 +1262,23 @@ message DataframeUnpivot {
12601262
string value_column = 6;
12611263
}
12621264

1263-
// dataframe.ir:317
1265+
// dataframe.ir:319
12641266
message DataframeWithColumn {
12651267
Expr col = 1;
12661268
string col_name = 2;
12671269
Expr df = 3;
12681270
SrcPosition src = 4;
12691271
}
12701272

1271-
// dataframe.ir:323
1273+
// dataframe.ir:325
12721274
message DataframeWithColumnRenamed {
12731275
Expr col = 1;
12741276
Expr df = 2;
12751277
string new_name = 3;
12761278
SrcPosition src = 4;
12771279
}
12781280

1279-
// dataframe.ir:329
1281+
// dataframe.ir:331
12801282
message DataframeWithColumns {
12811283
repeated string col_names = 1;
12821284
Expr df = 2;
@@ -1629,7 +1631,7 @@ message Geq {
16291631
SrcPosition src = 3;
16301632
}
16311633

1632-
// dataframe.ir:335
1634+
// dataframe.ir:337
16331635
message GroupingSets {
16341636
ExprArgList sets = 1;
16351637
SrcPosition src = 2;
@@ -2246,8 +2248,8 @@ message TableFnCallAlias {
22462248
// fn.ir:151
22472249
message TableFnCallOver {
22482250
Expr lhs = 1;
2249-
repeated Expr order_by = 2;
2250-
repeated Expr partition_by = 3;
2251+
ExprArgList order_by = 2;
2252+
ExprArgList partition_by = 3;
22512253
SrcPosition src = 4;
22522254
}
22532255

@@ -2283,7 +2285,7 @@ message TableUpdate {
22832285
repeated Tuple_String_String statement_params = 7;
22842286
}
22852287

2286-
// dataframe.ir:339
2288+
// dataframe.ir:341
22872289
message ToSnowparkPandas {
22882290
repeated string columns = 1;
22892291
Expr df = 2;

‎src/snowflake/snowpark/dataframe.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2213,7 +2213,7 @@ def group_by(
22132213
for c in col_list:
22142214
build_expr_from_snowpark_column_or_col_name(expr.cols.args.add(), c)
22152215

2216-
expr.df.dataframe_ref.id.bitfield1 = self._ast_id
2216+
self._set_ast_ref(expr.df)
22172217
else:
22182218
stmt = _ast_stmt
22192219

‎src/snowflake/snowpark/dataframe_na_functions.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -410,6 +410,7 @@ def fill(
410410
ast.subset.variadic = False
411411
for col in subset:
412412
build_expr_from_python_val(ast.subset.args.add(), col)
413+
ast.include_decimal = include_decimal
413414

414415
if subset is None:
415416
subset = self._dataframe.columns
@@ -633,6 +634,7 @@ def replace(
633634
ast.subset.variadic = False
634635
for col in subset:
635636
build_expr_from_python_val(ast.subset.args.add(), col)
637+
ast.include_decimal = include_decimal
636638

637639
# Modify subset.
638640
if subset is None:

0 commit comments

Comments
 (0)