snowpark-connect 0.27.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

snowflake/snowpark_connect/expression/map_unresolved_star.py

@@ -9,8 +9,11 @@ import snowflake.snowpark.functions as snowpark_fn
 from snowflake.snowpark._internal.analyzer.analyzer_utils import (
     quote_name_without_upper_casing,
 )
-from snowflake.snowpark.types import StructType
+from snowflake.snowpark.types import StringType, StructType, VariantType
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
+from snowflake.snowpark_connect.column_qualifier import ColumnQualifier
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.typed_column import TypedColumn
 from snowflake.snowpark_connect.utils.context import get_outer_dataframes
@@ -26,7 +29,7 @@ def check_struct_and_get_field_datatype(field_name, schema):
         else:
             return None
     else:
-        None
+        return None
 
 
 def map_unresolved_star(
@@ -34,7 +37,6 @@ def map_unresolved_star(
     column_mapping: ColumnNameMap,
     typer: ExpressionTyper,
 ) -> tuple[list[str], TypedColumn]:
-
     if exp.unresolved_star.HasField("unparsed_target"):
         unparsed_target = exp.unresolved_star.unparsed_target
         name_parts = split_fully_qualified_spark_name(unparsed_target)
@@ -54,16 +56,17 @@ def map_unresolved_star(
             return spark_names, typed_column
 
         # scenario where it is expanding * to mulitple columns
-        spark_names = []
-        snowpark_names = []
-        qualifiers = []
+        spark_names: list[str] = []
+        snowpark_names: list[str] = []
+        qualifiers: list[set[ColumnQualifier]] = []
 
+        target_qualifier = ColumnQualifier(tuple(name_parts[:-1]))
         (
             spark_names,
             snowpark_names,
             qualifiers,
         ) = column_mapping.get_spark_and_snowpark_columns_with_qualifier_for_qualifier(
-            name_parts[:-1]
+            target_qualifier
         )
 
         if len(spark_names) == 0:
@@ -74,7 +77,7 @@ def map_unresolved_star(
                     snowpark_names,
                     qualifiers,
                 ) = column_mapping_for_outer_df.get_spark_and_snowpark_columns_with_qualifier_for_qualifier(
-                    name_parts[:-1]
+                    target_qualifier
                 )
                 if len(spark_names) > 0:
                     break
@@ -103,7 +106,7 @@ def map_unresolved_star(
                 prefix_candidate_str = f"{prefix_candidate_str}.{name_parts[i]}"
             prefix_candidate = (
                 column_mapping.get_snowpark_column_name_from_spark_column_name(
-                    prefix_candidate_str, allow_non_exists=True, is_qualified=(i > 0)
+                    prefix_candidate_str, allow_non_exists=True
                 )
             )
             if prefix_candidate is None:
@@ -140,29 +143,86 @@ def map_unresolved_star(
                     final_sql_expr,
                     lambda final_sql_expr=final_sql_expr: typer.type(final_sql_expr),
                 )
-                typed_column.set_multi_col_qualifiers([[] for _ in spark_names])
+                typed_column.set_multi_col_qualifiers([set() for _ in spark_names])
                 return spark_names, typed_column
     else:
-        result_exp = snowpark_fn.sql_expr(
-            ", ".join(column_mapping.get_snowpark_columns())
-        )
+        snowpark_columns = column_mapping.get_snowpark_columns()
+        result_exp = snowpark_fn.sql_expr(", ".join(snowpark_columns))
         spark_names = column_mapping.get_spark_columns()
         typed_column = TypedColumn(
-            result_exp, lambda: [f.datatype for f in typer.df.schema]
+            result_exp,
+            lambda: [f.datatype for f in typer.df.schema if f.name in snowpark_columns],
         )
         typed_column.set_multi_col_qualifiers(column_mapping.get_qualifiers())
         return spark_names, typed_column
 
-    raise AnalysisException(
+    exception = AnalysisException(
         f"[UNRESOLVED_STAR] The unresolved star expression {exp} is not supported."
     )
+    attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+    raise exception
+
+
+def map_unresolved_star_as_single_column(
+    exp: expressions_proto.Expression,
+    column_mapping: ColumnNameMap,
+    typer: ExpressionTyper,
+) -> tuple[str, TypedColumn]:
+    """
+    Similar to map_unresolved_star but returns a single tuple containing
+    a combined spark column name and a TypedColumn representing a struct instead of many columns. .
+    If star resolves to single column, it works the same.
+    """
+    if exp.unresolved_star.HasField("unparsed_target"):
+        names, tc = map_unresolved_star(exp, column_mapping, typer)
+        assert len(names) == 1, "Expected single column"
+        return names[0], tc
+    else:
+        snowpark_columns = column_mapping.get_snowpark_columns()
+        spark_names = column_mapping.get_spark_columns()
+
+        if len(spark_names) == 1:
+            names, tc = map_unresolved_star(exp, column_mapping, typer)
+            return names[0], tc
+
+        fields_cols = [
+            (
+                spark_name,
+                TypedColumn(
+                    snowpark_fn.sql_expr(snowpark_name),
+                    lambda snowpark_name=snowpark_name: typer.type(
+                        snowpark_fn.sql_expr(snowpark_name)
+                    ),
+                ),
+            )
+            for spark_name, snowpark_name in zip(spark_names, snowpark_columns)
+        ]
+
+        result_exp = snowpark_fn.object_construct_keep_null(
+            *[
+                name_with_col
+                for name, typed_col in fields_cols
+                for name_with_col in (
+                    snowpark_fn.lit(name),
+                    typed_col.column(to_semi_structure=True),
+                )
+            ]
+        ).cast(VariantType())
+
+        combined_spark_name = "value"
+        typed_column = TypedColumn(
+            result_exp,
+            lambda: [VariantType()],
+        )
+        typed_column.set_multi_col_qualifiers([set() for _ in spark_names])
+        return combined_spark_name, typed_column
 
 
 def map_unresolved_star_struct(
     exp: expressions_proto.Expression,
     column_mapping: ColumnNameMap,
     typer: ExpressionTyper,
-) -> tuple[list[str], list]:
+) -> tuple[list[str], list[TypedColumn]]:
     unparsed_target = exp.unresolved_star.unparsed_target
     name_parts = split_fully_qualified_spark_name(unparsed_target)
 
@@ -170,7 +230,7 @@ def map_unresolved_star_struct(
         len(name_parts) > 1 and name_parts[-1] == "*"
     ), f"Unable to parse unparsed_target {unparsed_target}"
 
-    expanded_args = []
+    expanded_args: list[TypedColumn] = []
     for i in range(0, len(name_parts) - 1):
         if i == 0:
             prefix_candidate_str = name_parts[i]
@@ -181,7 +241,7 @@ def map_unresolved_star_struct(
             prefix_candidate_str = f"{prefix_candidate_str}.{name_parts[i]}"
         prefix_candidate = (
             column_mapping.get_snowpark_column_name_from_spark_column_name(
-                prefix_candidate_str, allow_non_exists=True, is_qualified=(i > 0)
+                prefix_candidate_str, allow_non_exists=True
             )
         )
         if prefix_candidate is None:
@@ -207,13 +267,17 @@ def map_unresolved_star_struct(
             if prefix_candidate is None:
                 continue
 
-            spark_names = candidate_leaf_field.names
+            spark_names: list[str] = candidate_leaf_field.names
             prefix_candidate = ":".join(fields)
 
-            for spark_name in spark_names:
-                expanded_args.append(snowpark_fn.lit(spark_name))
-                field_snowpark_name = f"{prefix_candidate}:{spark_name}"
+            for struct_field in candidate_leaf_field.fields:
+                lit_col = snowpark_fn.lit(struct_field.name)
+                expanded_args.append(TypedColumn(lit_col, lambda: [StringType()]))
+                field_snowpark_name = f"{prefix_candidate}:{struct_field.name}"
                 field_col = snowpark_fn.sql_expr(field_snowpark_name)
-                expanded_args.append(field_col)
+                field_type = struct_field.datatype
+                expanded_args.append(
+                    TypedColumn(field_col, lambda field_type=field_type: [field_type])
+                )
 
     return spark_names, expanded_args

snowflake/snowpark_connect/expression/map_update_fields.py

@@ -6,8 +6,17 @@ import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 from pyspark.errors.exceptions.base import AnalysisException
 
 import snowflake.snowpark.functions as snowpark_fn
-from snowflake.snowpark.types import DataType, StringType, StructField, StructType
+from snowflake.snowpark.types import (
+    DataType,
+    MapType,
+    StringType,
+    StructField,
+    StructType,
+    VariantType,
+)
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.typed_column import TypedColumn
 from snowflake.snowpark_connect.utils.identifiers import (
@@ -39,9 +48,11 @@ def update_field_in_schema(
                         field.name, updated_subschema, field.nullable, _is_column=False
                     )
                 else:
-                    raise AnalysisException(
+                    exception = AnalysisException(
                         message=f"[FIELD_NOT_FOUND] No such struct field `{field_str}` in `{field}`"
                     )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+                    raise exception
             field_updated = True
         else:
             new_field = field  # leave unchanged
@@ -59,9 +70,11 @@ def update_field_in_schema(
             # if the value type is None that means we want to drop the field and spark does not throw an error if the field does not exists
             # but if the value type is not None, it means we should add or update this field which has already been covered above
             # if we reach this code, it means the field should have existed
-            raise AnalysisException(
+            exception = AnalysisException(
                 message=f"[FIELD_NOT_FOUND] No such struct field `{field_str}`"
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+            raise exception
     return StructType(new_fields)
 
 
@@ -99,13 +112,16 @@ def map_update_fields(
     )
 
     if not isinstance(struct_typed_column.typ, StructType):
-        raise AnalysisException(
+        exception = AnalysisException(
             f'[DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE] Cannot resolve "update_fields({struct_name}, ...)" due to data type mismatch: Parameter 1 requires the "STRUCT" type'
         )
+        attach_custom_error_code(exception, ErrorCodes.TYPE_MISMATCH)
+        raise exception
 
     final_schema = struct_typed_column.typ
     value_column_list = []
-    input_types_to_the_udf = [struct_typed_column.typ]
+    # Snowflake UDFs don't support StructType/MapType, convert to VariantType
+    input_types_to_the_udf = [VariantType()]
     update_operation_strs = []
     array_of_named_parts = []
     for field_expression, value_expression in zip(field_expressions, value_expressions):
@@ -129,7 +145,11 @@ def map_update_fields(
             )
             update_operation_strs.append(f"WithField({value_spark_name})")
             value_column_list.append(value_typed_column.col)
-            input_types_to_the_udf.append(value_typed_column.typ)
+            # Convert StructType/MapType to VariantType for Snowflake UDFs (ArrayType is supported)
+            if isinstance(value_typed_column.typ, (StructType, MapType)):
+                input_types_to_the_udf.append(VariantType())
+            else:
+                input_types_to_the_udf.append(value_typed_column.typ)
 
         array_of_named_parts.append(name_parts)
 
@@ -137,28 +157,60 @@ def map_update_fields(
     final_name = f"update_fields({struct_name}, {update_operations_str})"
 
     if len(final_schema.fields) == 0:
-        raise AnalysisException(
+        exception = AnalysisException(
             f'[DATATYPE_MISMATCH.CANNOT_DROP_ALL_FIELDS] Cannot resolve "{final_name}" due to data type mismatch: Cannot drop all fields in struct.'
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+        raise exception
 
+    # Snowflake UDFs don't support StructType, so we use VariantType
+    # The result will be automatically cast back to the struct type
     @snowpark_fn.udf(
         input_types=input_types_to_the_udf,
-        return_type=final_schema,
+        return_type=VariantType(),
     )
     def _update(dictionary, *array_of_value):
+        if dictionary is None:
+            return None
+
+        # Recursively copy to create mutable dict from Snowflake's VARIANT objects
+        def make_mutable_copy(obj):
+            if obj is None:
+                return None
+            elif isinstance(obj, dict):
+                return {k: make_mutable_copy(v) for k, v in obj.items()}
+            elif isinstance(obj, (list, tuple)):
+                return [make_mutable_copy(item) for item in obj]
+            else:
+                return obj
+
+        result = make_mutable_copy(dictionary)
+
         for fields_array, value in zip(array_of_named_parts, array_of_value):
-            current = dictionary
+            current = result
             for k in fields_array[:-1]:
                 current = current.get(k)
-            if value == "_SNOWPARK_CONNECT_UPDATE_FIELD_DROP_":
-                current.pop(fields_array[-1], None)
-            else:
-                current[fields_array[-1]] = value
-        return dictionary
+                if current is None:
+                    break
 
-    final_exp = _update(
-        struct_typed_column.col,
-        *value_column_list,
-    )
+            if current is not None and isinstance(current, dict):
+                if value == "_SNOWPARK_CONNECT_UPDATE_FIELD_DROP_":
+                    current.pop(fields_array[-1], None)
+                else:
+                    current[fields_array[-1]] = value
+
+        return result
+
+    # Cast inputs to VARIANT (Snowflake UDFs don't support complex types directly)
+    struct_as_variant = struct_typed_column.col.cast(VariantType())
+    variant_value_list = [
+        col.cast(VariantType()) if isinstance(udf_type, VariantType) else col
+        for col, udf_type in zip(value_column_list, input_types_to_the_udf[1:])
+    ]
+
+    udf_result = _update(struct_as_variant, *variant_value_list)
+
+    # Cast the VariantType result back to the target StructType
+    final_exp = udf_result.cast(final_schema)
 
     return [final_name], TypedColumn(final_exp, lambda: typer.type(final_exp))

snowflake/snowpark_connect/expression/map_window_function.py

@@ -6,7 +6,11 @@ import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 
 from snowflake import snowpark
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
-from snowflake.snowpark_connect.error.error_utils import SparkException
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import (
+    SparkException,
+    attach_custom_error_code,
+)
 from snowflake.snowpark_connect.expression.literal import get_literal_field_and_name
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.typed_column import TypedColumn
@@ -29,6 +33,8 @@ SPARK_RANKING_FUNCTIONS = frozenset(
     ]
 )
 
+RANGE_BASED_WINDOW_FRAME_ONLY_SNOWFLAKE_FUNCTIONS = frozenset(["percent_rank"])
+
 CAPITAL_FUNCTION_NAMES = frozenset(["rank()", "dense_rank()", "percent_rank()"])
 
 
@@ -128,6 +134,11 @@ def map_window_function(
             case expressions_proto.Expression.Window.WindowFrame.FrameType.FRAME_TYPE_ROW:
                 frame_name.append("ROWS BETWEEN")
                 frame_type_func_string = "rows_between"
+                if proto_func_name in RANGE_BASED_WINDOW_FRAME_ONLY_SNOWFLAKE_FUNCTIONS:
+                    # Seems like Snowflake and Spark have different understanding of some functions. For those,
+                    # Spark only allows rows_between while Snowflake only allows range_between. To be compatible
+                    # with Spark, we have to use range_between here.
+                    frame_type_func_string = "range_between"
                 lower_name, lower = parse_frame_boundary(
                     exp.window.frame_spec.lower, is_upper=False
                 )
@@ -138,9 +149,11 @@ def map_window_function(
                     lower != snowpark.Window.UNBOUNDED_PRECEDING
                     or upper != snowpark.Window.CURRENT_ROW
                 ):
-                    raise SparkException.invalid_ranking_function_window_frame(
+                    exception = SparkException.invalid_ranking_function_window_frame(
                         window_frame=f"specifiedwindowframe(RowFrame, {lower_name}, {upper_name})"
                     )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+                    raise exception
 
                 is_unbounded = (
                     lower == snowpark.Window.UNBOUNDED_PRECEDING
@@ -165,9 +178,11 @@ def map_window_function(
                     orders = orders[:1]
 
                 if proto_func_name in SPARK_RANKING_FUNCTIONS:
-                    raise SparkException.invalid_ranking_function_window_frame(
+                    exception = SparkException.invalid_ranking_function_window_frame(
                         window_frame=f"specifiedwindowframe(RangeFrame, {lower_name}, {upper_name})"
                     )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+                    raise exception
 
                 is_unbounded = (
                     lower == snowpark.Window.UNBOUNDED_PRECEDING

snowflake/snowpark_connect/includes/python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py

@@ -49,7 +49,7 @@ def main(infile: IO, outfile: IO) -> None:
     )
 
     spark_connect_session = SparkSession.builder.remote(connect_url).getOrCreate()
-    spark_connect_session._client._session_id = session_id  # type: ignore[attr-defined]
+    spark_connect_session._client.session_id = session_id  # type: ignore[attr-defined]
 
     # TODO(SPARK-44460): Pass credentials.
     # TODO(SPARK-44461): Enable Process Isolation

snowflake/snowpark_connect/includes/python/pyspark/sql/connect/streaming/worker/listener_worker.py

@@ -57,7 +57,7 @@ def main(infile: IO, outfile: IO) -> None:
     )
 
     spark_connect_session = SparkSession.builder.remote(connect_url).getOrCreate()
-    spark_connect_session._client._session_id = session_id  # type: ignore[attr-defined]
+    spark_connect_session._client.session_id = session_id  # type: ignore[attr-defined]
 
     # TODO(SPARK-44460): Pass credentials.
     # TODO(SPARK-44461): Enable Process Isolation

snowflake/snowpark_connect/proto/snowflake_expression_ext_pb2.py

@@ -16,7 +16,7 @@ from pyspark.sql.connect.proto import expressions_pb2 as spark_dot_connect_dot_e
 from pyspark.sql.connect.proto import relations_pb2 as spark_dot_connect_dot_relations__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1esnowflake_expression_ext.proto\x12\rsnowflake.ext\x1a\x1fspark/connect/expressions.proto\x1a\x1dspark/connect/relations.proto\"\x98\x01\n\x0c\x45xpExtension\x12@\n\x0enamed_argument\x18\x01 \x01(\x0b\x32&.snowflake.ext.NamedArgumentExpressionH\x00\x12@\n\x13subquery_expression\x18\x02 \x01(\x0b\x32!.snowflake.ext.SubqueryExpressionH\x00\x42\x04\n\x02op\"P\n\x17NamedArgumentExpression\x12\x0b\n\x03key\x18\x01 \x01(\t\x12(\n\x05value\x18\x02 \x01(\x0b\x32\x19.spark.connect.Expression\"\xf4\x04\n\x12SubqueryExpression\x12&\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x45\n\rsubquery_type\x18\x02 \x01(\x0e\x32..snowflake.ext.SubqueryExpression.SubqueryType\x12Q\n\x11table_arg_options\x18\x03 \x01(\x0b\x32\x31.snowflake.ext.SubqueryExpression.TableArgOptionsH\x00\x88\x01\x01\x12\x35\n\x12in_subquery_values\x18\x04 \x03(\x0b\x32\x19.spark.connect.Expression\x1a\xbb\x01\n\x0fTableArgOptions\x12\x31\n\x0epartition_spec\x18\x01 \x03(\x0b\x32\x19.spark.connect.Expression\x12\x37\n\norder_spec\x18\x02 \x03(\x0b\x32#.spark.connect.Expression.SortOrder\x12\"\n\x15with_single_partition\x18\x03 \x01(\x08H\x00\x88\x01\x01\x42\x18\n\x16_with_single_partition\"\x90\x01\n\x0cSubqueryType\x12\x19\n\x15SUBQUERY_TYPE_UNKNOWN\x10\x00\x12\x18\n\x14SUBQUERY_TYPE_SCALAR\x10\x01\x12\x18\n\x14SUBQUERY_TYPE_EXISTS\x10\x02\x12\x1b\n\x17SUBQUERY_TYPE_TABLE_ARG\x10\x03\x12\x14\n\x10SUBQUERY_TYPE_IN\x10\x04\x42\x14\n\x12_table_arg_optionsb\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1esnowflake_expression_ext.proto\x12\rsnowflake.ext\x1a\x1fspark/connect/expressions.proto\x1a\x1dspark/connect/relations.proto\"\xde\x01\n\x0c\x45xpExtension\x12@\n\x0enamed_argument\x18\x01 \x01(\x0b\x32&.snowflake.ext.NamedArgumentExpressionH\x00\x12@\n\x13subquery_expression\x18\x02 \x01(\x0b\x32!.snowflake.ext.SubqueryExpressionH\x00\x12\x44\n\x10interval_literal\x18\x03 \x01(\x0b\x32(.snowflake.ext.IntervalLiteralExpressionH\x00\x42\x04\n\x02op\"P\n\x17NamedArgumentExpression\x12\x0b\n\x03key\x18\x01 \x01(\t\x12(\n\x05value\x18\x02 \x01(\x0b\x32\x19.spark.connect.Expression\"\xf4\x04\n\x12SubqueryExpression\x12&\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x45\n\rsubquery_type\x18\x02 \x01(\x0e\x32..snowflake.ext.SubqueryExpression.SubqueryType\x12Q\n\x11table_arg_options\x18\x03 \x01(\x0b\x32\x31.snowflake.ext.SubqueryExpression.TableArgOptionsH\x00\x88\x01\x01\x12\x35\n\x12in_subquery_values\x18\x04 \x03(\x0b\x32\x19.spark.connect.Expression\x1a\xbb\x01\n\x0fTableArgOptions\x12\x31\n\x0epartition_spec\x18\x01 \x03(\x0b\x32\x19.spark.connect.Expression\x12\x37\n\norder_spec\x18\x02 \x03(\x0b\x32#.spark.connect.Expression.SortOrder\x12\"\n\x15with_single_partition\x18\x03 \x01(\x08H\x00\x88\x01\x01\x42\x18\n\x16_with_single_partition\"\x90\x01\n\x0cSubqueryType\x12\x19\n\x15SUBQUERY_TYPE_UNKNOWN\x10\x00\x12\x18\n\x14SUBQUERY_TYPE_SCALAR\x10\x01\x12\x18\n\x14SUBQUERY_TYPE_EXISTS\x10\x02\x12\x1b\n\x17SUBQUERY_TYPE_TABLE_ARG\x10\x03\x12\x14\n\x10SUBQUERY_TYPE_IN\x10\x04\x42\x14\n\x12_table_arg_options\"\x9f\x01\n\x19IntervalLiteralExpression\x12\x32\n\x07literal\x18\x01 \x01(\x0b\x32!.spark.connect.Expression.Literal\x12\x18\n\x0bstart_field\x18\x02 \x01(\x05H\x00\x88\x01\x01\x12\x16\n\tend_field\x18\x03 \x01(\x05H\x01\x88\x01\x01\x42\x0e\n\x0c_start_fieldB\x0c\n\n_end_fieldb\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -24,13 +24,15 @@ _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'snowflake_expression_ext_pb
 if _descriptor._USE_C_DESCRIPTORS == False:
   DESCRIPTOR._options = None
   _globals['_EXPEXTENSION']._serialized_start=114
-  _globals['_EXPEXTENSION']._serialized_end=266
-  _globals['_NAMEDARGUMENTEXPRESSION']._serialized_start=268
-  _globals['_NAMEDARGUMENTEXPRESSION']._serialized_end=348
-  _globals['_SUBQUERYEXPRESSION']._serialized_start=351
-  _globals['_SUBQUERYEXPRESSION']._serialized_end=979
-  _globals['_SUBQUERYEXPRESSION_TABLEARGOPTIONS']._serialized_start=623
-  _globals['_SUBQUERYEXPRESSION_TABLEARGOPTIONS']._serialized_end=810
-  _globals['_SUBQUERYEXPRESSION_SUBQUERYTYPE']._serialized_start=813
-  _globals['_SUBQUERYEXPRESSION_SUBQUERYTYPE']._serialized_end=957
+  _globals['_EXPEXTENSION']._serialized_end=336
+  _globals['_NAMEDARGUMENTEXPRESSION']._serialized_start=338
+  _globals['_NAMEDARGUMENTEXPRESSION']._serialized_end=418
+  _globals['_SUBQUERYEXPRESSION']._serialized_start=421
+  _globals['_SUBQUERYEXPRESSION']._serialized_end=1049
+  _globals['_SUBQUERYEXPRESSION_TABLEARGOPTIONS']._serialized_start=693
+  _globals['_SUBQUERYEXPRESSION_TABLEARGOPTIONS']._serialized_end=880
+  _globals['_SUBQUERYEXPRESSION_SUBQUERYTYPE']._serialized_start=883
+  _globals['_SUBQUERYEXPRESSION_SUBQUERYTYPE']._serialized_end=1027
+  _globals['_INTERVALLITERALEXPRESSION']._serialized_start=1052
+  _globals['_INTERVALLITERALEXPRESSION']._serialized_end=1211
 # @@protoc_insertion_point(module_scope)

snowflake/snowpark_connect/proto/snowflake_expression_ext_pb2.pyi

@@ -9,12 +9,14 @@ from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Map
 DESCRIPTOR: _descriptor.FileDescriptor
 
 class ExpExtension(_message.Message):
-    __slots__ = ("named_argument", "subquery_expression")
+    __slots__ = ("named_argument", "subquery_expression", "interval_literal")
     NAMED_ARGUMENT_FIELD_NUMBER: _ClassVar[int]
     SUBQUERY_EXPRESSION_FIELD_NUMBER: _ClassVar[int]
+    INTERVAL_LITERAL_FIELD_NUMBER: _ClassVar[int]
     named_argument: NamedArgumentExpression
     subquery_expression: SubqueryExpression
-    def __init__(self, named_argument: _Optional[_Union[NamedArgumentExpression, _Mapping]] = ..., subquery_expression: _Optional[_Union[SubqueryExpression, _Mapping]] = ...) -> None: ...
+    interval_literal: IntervalLiteralExpression
+    def __init__(self, named_argument: _Optional[_Union[NamedArgumentExpression, _Mapping]] = ..., subquery_expression: _Optional[_Union[SubqueryExpression, _Mapping]] = ..., interval_literal: _Optional[_Union[IntervalLiteralExpression, _Mapping]] = ...) -> None: ...
 
 class NamedArgumentExpression(_message.Message):
     __slots__ = ("key", "value")
@@ -56,3 +58,13 @@ class SubqueryExpression(_message.Message):
     table_arg_options: SubqueryExpression.TableArgOptions
     in_subquery_values: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression]
     def __init__(self, input: _Optional[_Union[_relations_pb2.Relation, _Mapping]] = ..., subquery_type: _Optional[_Union[SubqueryExpression.SubqueryType, str]] = ..., table_arg_options: _Optional[_Union[SubqueryExpression.TableArgOptions, _Mapping]] = ..., in_subquery_values: _Optional[_Iterable[_Union[_expressions_pb2.Expression, _Mapping]]] = ...) -> None: ...
+
+class IntervalLiteralExpression(_message.Message):
+    __slots__ = ("literal", "start_field", "end_field")
+    LITERAL_FIELD_NUMBER: _ClassVar[int]
+    START_FIELD_FIELD_NUMBER: _ClassVar[int]
+    END_FIELD_FIELD_NUMBER: _ClassVar[int]
+    literal: _expressions_pb2.Expression.Literal
+    start_field: int
+    end_field: int
+    def __init__(self, literal: _Optional[_Union[_expressions_pb2.Expression.Literal, _Mapping]] = ..., start_field: _Optional[int] = ..., end_field: _Optional[int] = ...) -> None: ...

snowflake/snowpark_connect/proto/snowflake_relation_ext_pb2.py

@@ -16,7 +16,7 @@ from pyspark.sql.connect.proto import relations_pb2 as spark_dot_connect_dot_rel
 from pyspark.sql.connect.proto import expressions_pb2 as spark_dot_connect_dot_expressions__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1csnowflake_relation_ext.proto\x12\rsnowflake.ext\x1a\x1dspark/connect/relations.proto\x1a\x1fspark/connect/expressions.proto\"\xe3\x02\n\tExtension\x12(\n\x07rdd_map\x18\x01 \x01(\x0b\x32\x15.snowflake.ext.RddMapH\x00\x12.\n\nrdd_reduce\x18\x02 \x01(\x0b\x32\x18.snowflake.ext.RddReduceH\x00\x12G\n\x17subquery_column_aliases\x18\x03 \x01(\x0b\x32$.snowflake.ext.SubqueryColumnAliasesH\x00\x12\x32\n\x0clateral_join\x18\x04 \x01(\x0b\x32\x1a.snowflake.ext.LateralJoinH\x00\x12J\n\x19udtf_with_table_arguments\x18\x05 \x01(\x0b\x32%.snowflake.ext.UDTFWithTableArgumentsH\x00\x12-\n\taggregate\x18\x06 \x01(\x0b\x32\x18.snowflake.ext.AggregateH\x00\x42\x04\n\x02op\">\n\x06RddMap\x12&\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x0c\n\x04\x66unc\x18\x02 \x01(\x0c\"A\n\tRddReduce\x12&\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x0c\n\x04\x66unc\x18\x02 \x01(\x0c\"P\n\x15SubqueryColumnAliases\x12&\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x0f\n\x07\x61liases\x18\x02 \x03(\t\"\\\n\x0bLateralJoin\x12%\n\x04left\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12&\n\x05right\x18\x02 \x01(\x0b\x32\x17.spark.connect.Relation\"\x98\x01\n\x16UDTFWithTableArguments\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12,\n\targuments\x18\x02 \x03(\x0b\x32\x19.spark.connect.Expression\x12\x39\n\x0ftable_arguments\x18\x03 \x03(\x0b\x32 .snowflake.ext.TableArgumentInfo\"`\n\x11TableArgumentInfo\x12/\n\x0etable_argument\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x1a\n\x12table_argument_idx\x18\x02 \x01(\x05\"\xc7\x05\n\tAggregate\x12&\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x36\n\ngroup_type\x18\x02 \x01(\x0e\x32\".snowflake.ext.Aggregate.GroupType\x12\x37\n\x14grouping_expressions\x18\x03 \x03(\x0b\x32\x19.spark.connect.Expression\x12\x38\n\x15\x61ggregate_expressions\x18\x04 \x03(\x0b\x32\x19.spark.connect.Expression\x12-\n\x05pivot\x18\x05 \x01(\x0b\x32\x1e.snowflake.ext.Aggregate.Pivot\x12<\n\rgrouping_sets\x18\x06 \x03(\x0b\x32%.snowflake.ext.Aggregate.GroupingSets\x12\x33\n\x10having_condition\x18\x07 \x01(\x0b\x32\x19.spark.connect.Expression\x1a\x62\n\x05Pivot\x12&\n\x03\x63ol\x18\x01 \x01(\x0b\x32\x19.spark.connect.Expression\x12\x31\n\x06values\x18\x02 \x03(\x0b\x32!.spark.connect.Expression.Literal\x1a?\n\x0cGroupingSets\x12/\n\x0cgrouping_set\x18\x01 \x03(\x0b\x32\x19.spark.connect.Expression\"\x9f\x01\n\tGroupType\x12\x1a\n\x16GROUP_TYPE_UNSPECIFIED\x10\x00\x12\x16\n\x12GROUP_TYPE_GROUPBY\x10\x01\x12\x15\n\x11GROUP_TYPE_ROLLUP\x10\x02\x12\x13\n\x0fGROUP_TYPE_CUBE\x10\x03\x12\x14\n\x10GROUP_TYPE_PIVOT\x10\x04\x12\x1c\n\x18GROUP_TYPE_GROUPING_SETS\x10\x05\x62\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1csnowflake_relation_ext.proto\x12\rsnowflake.ext\x1a\x1dspark/connect/relations.proto\x1a\x1fspark/connect/expressions.proto\"\xe3\x02\n\tExtension\x12(\n\x07rdd_map\x18\x01 \x01(\x0b\x32\x15.snowflake.ext.RddMapH\x00\x12.\n\nrdd_reduce\x18\x02 \x01(\x0b\x32\x18.snowflake.ext.RddReduceH\x00\x12G\n\x17subquery_column_aliases\x18\x03 \x01(\x0b\x32$.snowflake.ext.SubqueryColumnAliasesH\x00\x12\x32\n\x0clateral_join\x18\x04 \x01(\x0b\x32\x1a.snowflake.ext.LateralJoinH\x00\x12J\n\x19udtf_with_table_arguments\x18\x05 \x01(\x0b\x32%.snowflake.ext.UDTFWithTableArgumentsH\x00\x12-\n\taggregate\x18\x06 \x01(\x0b\x32\x18.snowflake.ext.AggregateH\x00\x42\x04\n\x02op\">\n\x06RddMap\x12&\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x0c\n\x04\x66unc\x18\x02 \x01(\x0c\"A\n\tRddReduce\x12&\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x0c\n\x04\x66unc\x18\x02 \x01(\x0c\"P\n\x15SubqueryColumnAliases\x12&\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x0f\n\x07\x61liases\x18\x02 \x03(\t\"\\\n\x0bLateralJoin\x12%\n\x04left\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12&\n\x05right\x18\x02 \x01(\x0b\x32\x17.spark.connect.Relation\"\x98\x01\n\x16UDTFWithTableArguments\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12,\n\targuments\x18\x02 \x03(\x0b\x32\x19.spark.connect.Expression\x12\x39\n\x0ftable_arguments\x18\x03 \x03(\x0b\x32 .snowflake.ext.TableArgumentInfo\"`\n\x11TableArgumentInfo\x12/\n\x0etable_argument\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x1a\n\x12table_argument_idx\x18\x02 \x01(\x05\"\xbf\x06\n\tAggregate\x12&\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x36\n\ngroup_type\x18\x02 \x01(\x0e\x32\".snowflake.ext.Aggregate.GroupType\x12\x37\n\x14grouping_expressions\x18\x03 \x03(\x0b\x32\x19.spark.connect.Expression\x12\x38\n\x15\x61ggregate_expressions\x18\x04 \x03(\x0b\x32\x19.spark.connect.Expression\x12-\n\x05pivot\x18\x05 \x01(\x0b\x32\x1e.snowflake.ext.Aggregate.Pivot\x12<\n\rgrouping_sets\x18\x06 \x03(\x0b\x32%.snowflake.ext.Aggregate.GroupingSets\x12\x33\n\x10having_condition\x18\x07 \x01(\x0b\x32\x19.spark.connect.Expression\x1a\xd9\x01\n\x05Pivot\x12\x30\n\rpivot_columns\x18\x01 \x03(\x0b\x32\x19.spark.connect.Expression\x12?\n\x0cpivot_values\x18\x02 \x03(\x0b\x32).snowflake.ext.Aggregate.Pivot.PivotValue\x1a]\n\nPivotValue\x12\x31\n\x06values\x18\x01 \x03(\x0b\x32!.spark.connect.Expression.Literal\x12\x12\n\x05\x61lias\x18\x02 \x01(\tH\x00\x88\x01\x01\x42\x08\n\x06_alias\x1a?\n\x0cGroupingSets\x12/\n\x0cgrouping_set\x18\x01 \x03(\x0b\x32\x19.spark.connect.Expression\"\x9f\x01\n\tGroupType\x12\x1a\n\x16GROUP_TYPE_UNSPECIFIED\x10\x00\x12\x16\n\x12GROUP_TYPE_GROUPBY\x10\x01\x12\x15\n\x11GROUP_TYPE_ROLLUP\x10\x02\x12\x13\n\x0fGROUP_TYPE_CUBE\x10\x03\x12\x14\n\x10GROUP_TYPE_PIVOT\x10\x04\x12\x1c\n\x18GROUP_TYPE_GROUPING_SETS\x10\x05\x62\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -38,11 +38,13 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _globals['_TABLEARGUMENTINFO']._serialized_start=931
   _globals['_TABLEARGUMENTINFO']._serialized_end=1027
   _globals['_AGGREGATE']._serialized_start=1030
-  _globals['_AGGREGATE']._serialized_end=1741
-  _globals['_AGGREGATE_PIVOT']._serialized_start=1416
-  _globals['_AGGREGATE_PIVOT']._serialized_end=1514
-  _globals['_AGGREGATE_GROUPINGSETS']._serialized_start=1516
-  _globals['_AGGREGATE_GROUPINGSETS']._serialized_end=1579
-  _globals['_AGGREGATE_GROUPTYPE']._serialized_start=1582
-  _globals['_AGGREGATE_GROUPTYPE']._serialized_end=1741
+  _globals['_AGGREGATE']._serialized_end=1861
+  _globals['_AGGREGATE_PIVOT']._serialized_start=1417
+  _globals['_AGGREGATE_PIVOT']._serialized_end=1634
+  _globals['_AGGREGATE_PIVOT_PIVOTVALUE']._serialized_start=1541
+  _globals['_AGGREGATE_PIVOT_PIVOTVALUE']._serialized_end=1634
+  _globals['_AGGREGATE_GROUPINGSETS']._serialized_start=1636
+  _globals['_AGGREGATE_GROUPINGSETS']._serialized_end=1699
+  _globals['_AGGREGATE_GROUPTYPE']._serialized_start=1702
+  _globals['_AGGREGATE_GROUPTYPE']._serialized_end=1861
 # @@protoc_insertion_point(module_scope)

snowflake/snowpark_connect/proto/snowflake_relation_ext_pb2.pyi

@@ -91,12 +91,19 @@ class Aggregate(_message.Message):
     GROUP_TYPE_PIVOT: Aggregate.GroupType
     GROUP_TYPE_GROUPING_SETS: Aggregate.GroupType
     class Pivot(_message.Message):
-        __slots__ = ("col", "values")
-        COL_FIELD_NUMBER: _ClassVar[int]
-        VALUES_FIELD_NUMBER: _ClassVar[int]
-        col: _expressions_pb2.Expression
-        values: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression.Literal]
-        def __init__(self, col: _Optional[_Union[_expressions_pb2.Expression, _Mapping]] = ..., values: _Optional[_Iterable[_Union[_expressions_pb2.Expression.Literal, _Mapping]]] = ...) -> None: ...
+        __slots__ = ("pivot_columns", "pivot_values")
+        class PivotValue(_message.Message):
+            __slots__ = ("values", "alias")
+            VALUES_FIELD_NUMBER: _ClassVar[int]
+            ALIAS_FIELD_NUMBER: _ClassVar[int]
+            values: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression.Literal]
+            alias: str
+            def __init__(self, values: _Optional[_Iterable[_Union[_expressions_pb2.Expression.Literal, _Mapping]]] = ..., alias: _Optional[str] = ...) -> None: ...
+        PIVOT_COLUMNS_FIELD_NUMBER: _ClassVar[int]
+        PIVOT_VALUES_FIELD_NUMBER: _ClassVar[int]
+        pivot_columns: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression]
+        pivot_values: _containers.RepeatedCompositeFieldContainer[Aggregate.Pivot.PivotValue]
+        def __init__(self, pivot_columns: _Optional[_Iterable[_Union[_expressions_pb2.Expression, _Mapping]]] = ..., pivot_values: _Optional[_Iterable[_Union[Aggregate.Pivot.PivotValue, _Mapping]]] = ...) -> None: ...
     class GroupingSets(_message.Message):
         __slots__ = ("grouping_set",)
         GROUPING_SET_FIELD_NUMBER: _ClassVar[int]