snowpark-connect 0.22.1__py3-none-any.whl → 0.23.0__py3-none-any.whl

snowflake/snowpark_connect/config.py

@@ -4,10 +4,8 @@
 
 # Proto source for reference:
 # https://github.com/apache/spark/blob/branch-3.5/connector/connect/common/src/main/protobuf/spark/connect/base.proto#L420
-import os
 import re
 import sys
-import time
 from collections import defaultdict
 from copy import copy, deepcopy
 from typing import Any
@@ -533,17 +531,8 @@ def set_snowflake_parameters(
                 snowpark_session.sql(
                     f"ALTER SESSION SET TIMEZONE = '{value}'"
                 ).collect()
-                set_jvm_timezone(value)
-                if hasattr(time, "tzset"):
-                    os.environ["TZ"] = value
-                    time.tzset()
             else:
                 snowpark_session.sql("ALTER SESSION UNSET TIMEZONE").collect()
-                reset_jvm_timezone_to_system_default()
-                if hasattr(time, "tzset") and False:
-                    if "TZ" in os.environ:
-                        del os.environ["TZ"]
-                    time.tzset()
         case "spark.sql.globalTempDatabase":
             if not value:
                 value = global_config.default_static_global_config.get(key)

snowflake/snowpark_connect/error/error_utils.py

@@ -75,6 +75,10 @@ terminate_multi_args_exception_pattern = (
 snowpark_connect_exception_pattern = re.compile(
     r"\[snowpark-connect-exception(?::(\w+))?\]\s*(.+?)'\s*is not recognized"
 )
+invalid_bit_pattern = re.compile(
+    r"Invalid bit position: \d+ exceeds the bit (?:upper|lower) limit",
+    re.IGNORECASE,
+)
 
 
 def contains_udtf_select(sql_string):
@@ -107,6 +111,9 @@ def _get_converted_known_sql_or_custom_exception(
         return SparkRuntimeException(
             message="Unexpected value for start in function slice: SQL array indices start at 1."
         )
+    invalid_bit = invalid_bit_pattern.search(msg)
+    if invalid_bit:
+        return IllegalArgumentException(message=invalid_bit.group(0))
     match = snowpark_connect_exception_pattern.search(
         ex.message if hasattr(ex, "message") else str(ex)
     )

snowflake/snowpark_connect/error/exceptions.py

@@ -22,3 +22,7 @@ class MissingSchema(SnowparkConnectException):
         super().__init__(
             "No default schema found in session",
         )
+
+
+class MaxRetryExceeded(SnowparkConnectException):
+    ...

snowflake/snowpark_connect/expression/hybrid_column_map.py

@@ -0,0 +1,192 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+"""
+Hybrid column mapping for HAVING clause resolution.
+
+This module provides a special column mapping that can resolve expressions
+in the context of both the input DataFrame (for base columns) and the
+aggregated DataFrame (for aggregate expressions and aliases).
+"""
+
+from typing import Dict, List
+
+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.expression.typer import ExpressionTyper
+from snowflake.snowpark_connect.typed_column import TypedColumn
+
+
+class HybridColumnMap:
+    """
+    A column mapping that can resolve expressions in both input and aggregated contexts.
+
+    This is specifically designed for HAVING clause resolution where expressions may reference:
+    1. Base columns from the input DataFrame (to build new aggregates)
+    2. Existing aggregate expressions and their aliases
+    3. Grouping columns
+    """
+
+    def __init__(
+        self,
+        input_column_map: ColumnNameMap,
+        input_typer: ExpressionTyper,
+        aggregated_column_map: ColumnNameMap,
+        aggregated_typer: ExpressionTyper,
+        aggregate_expressions: List[expressions_proto.Expression],
+        grouping_expressions: List[expressions_proto.Expression],
+        aggregate_aliases: Dict[str, expressions_proto.Expression],
+    ) -> None:
+        self.input_column_map = input_column_map
+        self.input_typer = input_typer
+        self.aggregated_column_map = aggregated_column_map
+        self.aggregated_typer = aggregated_typer
+        self.aggregate_expressions = aggregate_expressions
+        self.grouping_expressions = grouping_expressions
+        self.aggregate_aliases = aggregate_aliases
+
+    def is_aggregate_function(self, exp: expressions_proto.Expression) -> bool:
+        """Check if an expression is an aggregate function."""
+        if exp.WhichOneof("expr_type") == "unresolved_function":
+            func_name = exp.unresolved_function.function_name.lower()
+            # Common aggregate functions - expand this list as needed
+            aggregate_functions = {
+                "avg",
+                "average",
+                "sum",
+                "count",
+                "min",
+                "max",
+                "stddev",
+                "stddev_pop",
+                "stddev_samp",
+                "variance",
+                "var_pop",
+                "var_samp",
+                "collect_list",
+                "collect_set",
+                "first",
+                "last",
+                "any_value",
+                "bool_and",
+                "bool_or",
+                "corr",
+                "covar_pop",
+                "covar_samp",
+                "kurtosis",
+                "skewness",
+                "percentile_cont",
+                "percentile_disc",
+                "approx_count_distinct",
+            }
+            return func_name in aggregate_functions
+        return False
+
+    def is_grouping_column(self, column_name: str) -> bool:
+        """Check if a column name refers to a grouping column."""
+        for group_exp in self.grouping_expressions:
+            if (
+                group_exp.WhichOneof("expr_type") == "unresolved_attribute"
+                and group_exp.unresolved_attribute.unparsed_identifier == column_name
+            ):
+                return True
+        return False
+
+    def resolve_expression(
+        self, exp: expressions_proto.Expression
+    ) -> tuple[list[str], TypedColumn]:
+        """
+        Resolve an expression in the hybrid context.
+
+        Strategy:
+        1. If it's an aggregate function -> create new aggregate using input context
+        2. If it's an alias to existing aggregate -> use aggregated context
+        3. If it's a grouping column -> try aggregated context first, fall back to input context
+           (handles exclude_grouping_columns=True case)
+        4. Otherwise -> try input context first, then aggregated context
+        """
+        from snowflake.snowpark_connect.expression.map_expression import map_expression
+
+        expr_type = exp.WhichOneof("expr_type")
+
+        # Handle aggregate functions - need to evaluate against input DataFrame
+        if self.is_aggregate_function(exp):
+            return map_expression(exp, self.input_column_map, self.input_typer)
+
+        # Handle column references
+        if expr_type == "unresolved_attribute":
+            column_name = exp.unresolved_attribute.unparsed_identifier
+
+            # Check if it's an alias to an existing aggregate expression
+            if column_name in self.aggregate_aliases:
+                # Use the aggregated context to get the alias
+                return map_expression(
+                    exp, self.aggregated_column_map, self.aggregated_typer
+                )
+
+            # Check if it's a grouping column
+            if self.is_grouping_column(column_name):
+                # Try aggregated context first (for cases where grouping columns are included)
+                try:
+                    return map_expression(
+                        exp, self.aggregated_column_map, self.aggregated_typer
+                    )
+                except Exception:
+                    # Fall back to input context if grouping columns were excluded
+                    # This handles the exclude_grouping_columns=True case
+                    return map_expression(exp, self.input_column_map, self.input_typer)
+
+            # Try input context first (for base columns used in new aggregates)
+            try:
+                return map_expression(exp, self.input_column_map, self.input_typer)
+            except Exception:
+                # Fall back to aggregated context
+                return map_expression(
+                    exp, self.aggregated_column_map, self.aggregated_typer
+                )
+
+        # For other expression types, try aggregated context first (likely references to computed values)
+        try:
+            return map_expression(
+                exp, self.aggregated_column_map, self.aggregated_typer
+            )
+        except Exception:
+            # Fall back to input context
+            return map_expression(exp, self.input_column_map, self.input_typer)
+
+
+def create_hybrid_column_map_for_having(
+    input_df: snowpark.DataFrame,
+    input_column_map: ColumnNameMap,
+    aggregated_df: snowpark.DataFrame,
+    aggregated_column_map: ColumnNameMap,
+    aggregate_expressions: List[expressions_proto.Expression],
+    grouping_expressions: List[expressions_proto.Expression],
+    spark_columns: List[str],
+    raw_aggregations: List[tuple[str, TypedColumn]],
+) -> HybridColumnMap:
+    """
+    Create a HybridColumnMap instance for HAVING clause resolution.
+    """
+    # Create typers for both contexts
+    input_typer = ExpressionTyper(input_df)
+    aggregated_typer = ExpressionTyper(aggregated_df)
+
+    # Build alias mapping from spark column names to aggregate expressions
+    aggregate_aliases = {}
+    for i, (spark_name, _) in enumerate(raw_aggregations):
+        if i < len(aggregate_expressions):
+            aggregate_aliases[spark_name] = aggregate_expressions[i]
+
+    return HybridColumnMap(
+        input_column_map=input_column_map,
+        input_typer=input_typer,
+        aggregated_column_map=aggregated_column_map,
+        aggregated_typer=aggregated_typer,
+        aggregate_expressions=aggregate_expressions,
+        grouping_expressions=grouping_expressions,
+        aggregate_aliases=aggregate_aliases,
+    )

snowflake/snowpark_connect/expression/literal.py

@@ -10,6 +10,7 @@ import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 from tzlocal import get_localzone
 
 from snowflake.snowpark_connect.config import global_config
+from snowflake.snowpark_connect.utils.context import get_is_evaluating_sql
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
@@ -47,17 +48,7 @@ def get_literal_field_and_name(literal: expressions_proto.Expression.Literal):
             ).date()
             return date, f"DATE '{date}'"
         case "timestamp" | "timestamp_ntz" as t:
-            # Note - Clients need to ensure local_timezone is the same as spark_sql_session_timeZone config.
-            # No need to apply timezone for lit datetime, because we set the TIMEZONE parameter in snowpark session,
-            # the snowflake backend would convert the lit datetime correctly. However, for returned column name, the
-            # timezone needs to be added. Pyspark has a weird behavior that datetime.datetime always gets converted
-            # to local timezone before printing according to spark_sql_session_timeZone setting. Haven't found
-            # official doc about it, but this behavior is based on my testings.
-            tz = (
-                ZoneInfo(global_config.spark_sql_session_timeZone)
-                if hasattr(global_config, "spark_sql_session_timeZone")
-                else get_localzone()
-            )
+            local_tz = get_localzone()
             if t == "timestamp":
                 microseconds = literal.timestamp
             else:
@@ -66,11 +57,17 @@ def get_literal_field_and_name(literal: expressions_proto.Expression.Literal):
                 microseconds // 1_000_000
             ) + datetime.timedelta(microseconds=microseconds % 1_000_000)
             tz_dt = datetime.datetime.fromtimestamp(
-                microseconds // 1_000_000, tz=tz
+                microseconds // 1_000_000, tz=local_tz
             ) + datetime.timedelta(microseconds=microseconds % 1_000_000)
             if t == "timestamp_ntz":
                 lit_dt = lit_dt.astimezone(datetime.timezone.utc)
                 tz_dt = tz_dt.astimezone(datetime.timezone.utc)
+            elif not get_is_evaluating_sql():
+                config_tz = global_config.spark_sql_session_timeZone
+                config_tz = ZoneInfo(config_tz) if config_tz else local_tz
+                tz_dt = tz_dt.astimezone(config_tz)
+                lit_dt = lit_dt.astimezone(local_tz)
+
             return lit_dt, f"{t.upper()} '{tz_dt.strftime('%Y-%m-%d %H:%M:%S')}'"
         case "day_time_interval":
             # TODO(SNOW-1920942): Snowflake SQL is missing an "interval" type.

snowflake/snowpark_connect/expression/map_cast.py

@@ -127,10 +127,11 @@ def map_cast(
         from_type = StringType()
     if isinstance(to_type, StringType):
         to_type = StringType()
+
+    # todo - verify if that's correct SNOW-2248680
     if isinstance(from_type, TimestampType):
         from_type = TimestampType()
-    if isinstance(to_type, TimestampType):
-        to_type = TimestampType()
+
     match (from_type, to_type):
         case (_, _) if (from_type == to_type):
             result_exp = col
@@ -185,6 +186,17 @@ def map_cast(
         case (DateType(), TimestampType()):
             result_exp = snowpark_fn.to_timestamp(col)
             result_exp = result_exp.cast(TimestampType(TimestampTimeZone.NTZ))
+        case (TimestampType() as f, TimestampType() as t) if f.tzinfo == t.tzinfo:
+            result_exp = col
+        case (
+            TimestampType(),
+            TimestampType() as t,
+        ) if t.tzinfo == TimestampTimeZone.NTZ:
+            zone = global_config.spark_sql_session_timeZone
+            result_exp = snowpark_fn.convert_timezone(snowpark_fn.lit(zone), col).cast(
+                TimestampType(TimestampTimeZone.NTZ)
+            )
+        # todo: verify if more support for LTZ and TZ is needed - SNOW-2248680
         case (TimestampType(), TimestampType()):
             result_exp = col
         case (_, TimestampType()) if isinstance(from_type, _NumericType):
@@ -259,8 +271,12 @@ def map_cast(
         case (_, _) if isinstance(from_type, _FractionalType) and isinstance(
             to_type, _IntegralType
         ):
-            result_exp = snowpark_fn.when(col < 0, snowpark_fn.ceil(col)).otherwise(
-                snowpark_fn.floor(col)
+            result_exp = (
+                snowpark_fn.when(
+                    col == snowpark_fn.lit(float("nan")), snowpark_fn.lit(0)
+                )
+                .when(col < 0, snowpark_fn.ceil(col))
+                .otherwise(snowpark_fn.floor(col))
             )
             result_exp = result_exp.cast(to_type)
         case (StringType(), _) if (isinstance(to_type, _IntegralType)):

snowflake/snowpark_connect/expression/map_expression.py

@@ -11,6 +11,7 @@ import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
 from snowflake.snowpark import Session
 from snowflake.snowpark._internal.analyzer.expression import UnresolvedAttribute
+from snowflake.snowpark.types import TimestampTimeZone, TimestampType
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
 from snowflake.snowpark_connect.expression import (
     map_extension,
@@ -190,9 +191,15 @@ def map_expression(
                 return [lit_name], TypedColumn(
                     snowpark_fn.lit(lit_value, return_type), lambda: [return_type]
                 )
+            result_exp = snowpark_fn.lit(lit_value)
+
+            if lit_type_str == "timestamp_ntz" and isinstance(
+                lit_value, datetime.datetime
+            ):
+                result_exp = result_exp.cast(TimestampType(TimestampTimeZone.NTZ))
 
             return [lit_name], TypedColumn(
-                snowpark_fn.lit(lit_value), lambda: [map_simple_types(lit_type_str)]
+                result_exp, lambda: [map_simple_types(lit_type_str)]
             )
         case "sort_order":
             child_name, child_column = map_single_column_expression(

snowflake/snowpark_connect/expression/map_unresolved_extract_value.py

@@ -5,6 +5,7 @@
 import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 
 import snowflake.snowpark.functions as snowpark_fn
+from snowflake.snowpark._internal.analyzer.expression import Literal
 from snowflake.snowpark.types import ArrayType, MapType, StructType, _IntegralType
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
 from snowflake.snowpark_connect.config import global_config
@@ -57,7 +58,8 @@ def map_unresolved_extract_value(
         extract_fn = snowpark_fn.get_ignore_case
     # Set index to a dummy value before we use it later in the ansi mode check.
     index = snowpark_fn.lit(1)
-    if _check_if_array_type(extract_typed_column, child_typed_column):
+    is_array = _check_if_array_type(extract_typed_column, child_typed_column)
+    if is_array:
         # Set all non-valid array indices to NULL.
         # This is done because both conditions of a CASE WHEN statement are executed regardless of if the condition is true or not.
         # Getting a negative index in Snowflake throws an error; thus, we convert all non-valid array indices to NULL before getting the index.
@@ -74,12 +76,37 @@ def map_unresolved_extract_value(
 
     spark_sql_ansi_enabled = global_config.spark_sql_ansi_enabled
 
-    if spark_sql_ansi_enabled and _check_if_array_type(
-        extract_typed_column, child_typed_column
-    ):
+    if spark_sql_ansi_enabled and is_array:
         result_exp = snowpark_fn.when(
             index.isNull(),
             child_typed_column.col.getItem("[snowpark_connect::INVALID_ARRAY_INDEX]"),
         ).otherwise(result_exp)
 
-    return spark_function_name, TypedColumn(result_exp, lambda: typer.type(result_exp))
+    def _get_extracted_value_type():
+        if is_array:
+            return [child_typed_column.typ.element_type]
+        elif isinstance(child_typed_column.typ, MapType):
+            return [child_typed_column.typ.value_type]
+        elif (
+            isinstance(child_typed_column.typ, StructType)
+            and isinstance(extract_typed_column.col._expr1, Literal)
+            and isinstance(extract_typed_column.col._expr1.value, str)
+        ):
+            struct = dict(
+                {
+                    (
+                        f.name
+                        if global_config.spark_sql_caseSensitive
+                        else f.name.lower(),
+                        f.datatype,
+                    )
+                    for f in child_typed_column.typ.fields
+                }
+            )
+            key = extract_typed_column.col._expr1.value
+            key = key if global_config.spark_sql_caseSensitive else key.lower()
+
+            return [struct[key]] if key in struct else typer.type(result_exp)
+        return typer.type(result_exp)
+
+    return spark_function_name, TypedColumn(result_exp, _get_extracted_value_type)

snowflake/snowpark_connect/expression/map_unresolved_function.py

@@ -1107,7 +1107,7 @@ def map_unresolved_function(
             result_exp = TypedColumn(
                 result_exp, lambda: [ArrayType(snowpark_typed_args[0].typ)]
             )
-        case "array_size" | "cardinality":
+        case "array_size":
             array_type = snowpark_typed_args[0].typ
             if not isinstance(array_type, ArrayType):
                 raise AnalysisException(
@@ -1116,6 +1116,16 @@ def map_unresolved_function(
             result_exp = TypedColumn(
                 snowpark_fn.array_size(*snowpark_args), lambda: [LongType()]
             )
+        case "cardinality":
+            arg_type = snowpark_typed_args[0].typ
+            if isinstance(arg_type, (ArrayType, MapType)):
+                result_exp = TypedColumn(
+                    snowpark_fn.size(*snowpark_args), lambda: [LongType()]
+                )
+            else:
+                raise AnalysisException(
+                    f"Expected argument '{snowpark_arg_names[0]}' to have an ArrayType or MapType, but got {arg_type.simpleString()}."
+                )
         case "array_sort":
             result_exp = TypedColumn(
                 snowpark_fn.array_sort(*snowpark_args),
@@ -1295,10 +1305,35 @@ def map_unresolved_function(
             )
             result_exp = TypedColumn(result_exp, lambda: [LongType()])
         case "bit_get" | "getbit":
-            bit_get_function = snowpark_fn.function("GETBIT")
-            result_exp = TypedColumn(
-                bit_get_function(*snowpark_args), lambda: [LongType()]
+            snowflake_compat = get_boolean_session_config_param(
+                "enable_snowflake_extension_behavior"
             )
+            col, pos = snowpark_args
+            if snowflake_compat:
+                bit_get_function = snowpark_fn.function("GETBIT")(col, pos)
+            else:
+                raise_error = _raise_error_helper(LongType())
+                bit_get_function = snowpark_fn.when(
+                    (snowpark_fn.lit(0) <= pos) & (pos <= snowpark_fn.lit(63))
+                    | snowpark_fn.is_null(pos),
+                    snowpark_fn.function("GETBIT")(col, pos),
+                ).otherwise(
+                    raise_error(
+                        snowpark_fn.concat(
+                            snowpark_fn.lit(
+                                "Invalid bit position: ",
+                            ),
+                            snowpark_fn.cast(
+                                pos,
+                                StringType(),
+                            ),
+                            snowpark_fn.lit(
+                                " exceeds the bit upper limit",
+                            ),
+                        )
+                    )
+                )
+            result_exp = TypedColumn(bit_get_function, lambda: [LongType()])
         case "bit_length":
             bit_length_function = snowpark_fn.function("bit_length")
             result_exp = TypedColumn(
@@ -3787,7 +3822,13 @@ def map_unresolved_function(
         case "locate":
             substr = unwrap_literal(exp.unresolved_function.arguments[0])
             value = snowpark_args[1]
-            start_pos = unwrap_literal(exp.unresolved_function.arguments[2])
+            if len(exp.unresolved_function.arguments) == 3:
+                start_pos = unwrap_literal(exp.unresolved_function.arguments[2])
+            else:
+                # start_pos is an optional argument and if not provided we should default to 1.
+                # This path will only be reached by spark connect scala clients.
+                start_pos = 1
+                spark_function_name = f"locate({', '.join(snowpark_arg_names)}, 1)"
 
             if start_pos > 0:
                 result_exp = snowpark_fn.locate(substr, value, start_pos)
@@ -5496,9 +5537,27 @@ def map_unresolved_function(
             ):
                 result_exp = snowpark_fn.lit(None)
             else:
+                right_expr = snowpark_fn.right(*snowpark_args)
+                if isinstance(snowpark_typed_args[0].typ, TimestampType):
+                    # Spark format is always displayed as YYY-MM-DD HH:mm:ss.FF6
+                    # When microseconds are equal to 0 .FF6 part is removed
+                    # When microseconds are equal to 0 at the end, they are removed i.e. .123000 -> .123 when displayed
+
+                    formated_timestamp = snowpark_fn.to_varchar(
+                        snowpark_args[0], "YYYY-MM-DD HH:MI:SS.FF6"
+                    )
+                    right_expr = snowpark_fn.right(
+                        snowpark_fn.regexp_replace(
+                            snowpark_fn.regexp_replace(formated_timestamp, "0+$", ""),
+                            "\\.$",
+                            "",
+                        ),
+                        snowpark_args[1],
+                    )
+
                 result_exp = snowpark_fn.when(
                     snowpark_args[1] <= 0, snowpark_fn.lit("")
-                ).otherwise(snowpark_fn.right(*snowpark_args))
+                ).otherwise(right_expr)
             result_type = StringType()
         case "rint":
             result_exp = snowpark_fn.cast(
@@ -6729,6 +6788,7 @@ def map_unresolved_function(
                         if value == "" or any(
                             c in value for c in [",", "\n", "\r", '"', "'"]
                         ):
+                            value = value.replace("\\", "\\\\").replace('"', '\\"')
                             result.append(f'"{value}"')
                         else:
                             result.append(value)

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\"\x92\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\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\"\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')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -38,11 +38,11 @@ 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=1688
-  _globals['_AGGREGATE_PIVOT']._serialized_start=1363
-  _globals['_AGGREGATE_PIVOT']._serialized_end=1461
-  _globals['_AGGREGATE_GROUPINGSETS']._serialized_start=1463
-  _globals['_AGGREGATE_GROUPINGSETS']._serialized_end=1526
-  _globals['_AGGREGATE_GROUPTYPE']._serialized_start=1529
-  _globals['_AGGREGATE_GROUPTYPE']._serialized_end=1688
+  _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
 # @@protoc_insertion_point(module_scope)

snowflake/snowpark_connect/proto/snowflake_relation_ext_pb2.pyi

@@ -75,7 +75,7 @@ class TableArgumentInfo(_message.Message):
     def __init__(self, table_argument: _Optional[_Union[_relations_pb2.Relation, _Mapping]] = ..., table_argument_idx: _Optional[int] = ...) -> None: ...
 
 class Aggregate(_message.Message):
-    __slots__ = ("input", "group_type", "grouping_expressions", "aggregate_expressions", "pivot", "grouping_sets")
+    __slots__ = ("input", "group_type", "grouping_expressions", "aggregate_expressions", "pivot", "grouping_sets", "having_condition")
     class GroupType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
         __slots__ = ()
         GROUP_TYPE_UNSPECIFIED: _ClassVar[Aggregate.GroupType]
@@ -108,10 +108,12 @@ class Aggregate(_message.Message):
     AGGREGATE_EXPRESSIONS_FIELD_NUMBER: _ClassVar[int]
     PIVOT_FIELD_NUMBER: _ClassVar[int]
     GROUPING_SETS_FIELD_NUMBER: _ClassVar[int]
+    HAVING_CONDITION_FIELD_NUMBER: _ClassVar[int]
     input: _relations_pb2.Relation
     group_type: Aggregate.GroupType
     grouping_expressions: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression]
     aggregate_expressions: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression]
     pivot: Aggregate.Pivot
     grouping_sets: _containers.RepeatedCompositeFieldContainer[Aggregate.GroupingSets]
-    def __init__(self, input: _Optional[_Union[_relations_pb2.Relation, _Mapping]] = ..., group_type: _Optional[_Union[Aggregate.GroupType, str]] = ..., grouping_expressions: _Optional[_Iterable[_Union[_expressions_pb2.Expression, _Mapping]]] = ..., aggregate_expressions: _Optional[_Iterable[_Union[_expressions_pb2.Expression, _Mapping]]] = ..., pivot: _Optional[_Union[Aggregate.Pivot, _Mapping]] = ..., grouping_sets: _Optional[_Iterable[_Union[Aggregate.GroupingSets, _Mapping]]] = ...) -> None: ...
+    having_condition: _expressions_pb2.Expression
+    def __init__(self, input: _Optional[_Union[_relations_pb2.Relation, _Mapping]] = ..., group_type: _Optional[_Union[Aggregate.GroupType, str]] = ..., grouping_expressions: _Optional[_Iterable[_Union[_expressions_pb2.Expression, _Mapping]]] = ..., aggregate_expressions: _Optional[_Iterable[_Union[_expressions_pb2.Expression, _Mapping]]] = ..., pivot: _Optional[_Union[Aggregate.Pivot, _Mapping]] = ..., grouping_sets: _Optional[_Iterable[_Union[Aggregate.GroupingSets, _Mapping]]] = ..., having_condition: _Optional[_Union[_expressions_pb2.Expression, _Mapping]] = ...) -> None: ...