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

snowflake/snowpark_connect/dataframe_container.py

@@ -4,14 +4,40 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Callable
 
 from snowflake import snowpark
 from snowflake.snowpark.types import StructField, StructType
-from snowflake.snowpark_connect.hidden_column import HiddenColumn
+from snowflake.snowpark_connect.column_qualifier import ColumnQualifier
 
 if TYPE_CHECKING:
+    import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
+
     from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
+    from snowflake.snowpark_connect.typed_column import TypedColumn
+
+
+@dataclass
+class AggregateMetadata:
+    """
+    Metadata about aggregation for resolving expressions in ORDER BY.
+
+    When a Sort operation follows an Aggregate operation, ORDER BY expressions
+    may reference:
+    1. Grouping columns from the GROUP BY clause
+    2. Aggregate result columns (aliases)
+    3. Expressions on pre-aggregation columns (e.g., year(date) where date existed before GROUP BY)
+
+    This metadata enables hybrid resolution similar to HAVING clause.
+    """
+
+    input_column_map: ColumnNameMap
+    input_dataframe: snowpark.DataFrame
+    grouping_expressions: list[expressions_proto.Expression]
+    aggregate_expressions: list[expressions_proto.Expression]
+    spark_columns: list[str]
+    raw_aggregations: list[tuple[str, TypedColumn]]
 
 
 class DataFrameContainer:
@@ -30,6 +56,9 @@ class DataFrameContainer:
         alias: str | None = None,
         cached_schema_getter: Callable[[], StructType] | None = None,
         partition_hint: int | None = None,
+        can_be_cached: bool = True,
+        can_be_materialized: bool = True,
+        aggregate_metadata: AggregateMetadata | None = None,
     ) -> None:
         """
         Initialize a new DataFrameContainer.
@@ -41,12 +70,16 @@ class DataFrameContainer:
             alias: Optional alias for the DataFrame
             cached_schema_getter: Optional function to get cached schema
             partition_hint: Optional partition count from repartition() operations
+            aggregate_metadata: Optional metadata about aggregation for ORDER BY resolution
         """
         self._dataframe = dataframe
         self._column_map = self._create_default_column_map(column_map)
         self._table_name = table_name
         self._alias = alias
         self._partition_hint = partition_hint
+        self._can_be_cached = can_be_cached
+        self._can_be_materialized = can_be_materialized
+        self._aggregate_metadata = aggregate_metadata
 
         if cached_schema_getter is not None:
             self._apply_cached_schema_getter(cached_schema_getter)
@@ -59,13 +92,16 @@ class DataFrameContainer:
         snowpark_column_names: list[str],
         snowpark_column_types: list | None = None,
         column_metadata: dict | None = None,
-        column_qualifiers: list[list[str]] | None = None,
+        column_qualifiers: list[set[ColumnQualifier]] | None = None,
         parent_column_name_map: ColumnNameMap | None = None,
-        hidden_columns: set[HiddenColumn] | None = None,
         table_name: str | None = None,
         alias: str | None = None,
         cached_schema_getter: Callable[[], StructType] | None = None,
         partition_hint: int | None = None,
+        equivalent_snowpark_names: list[set[str]] | None = None,
+        column_is_hidden: list[bool] | None = None,
+        can_be_cached: bool = True,
+        aggregate_metadata: AggregateMetadata | None = None,
     ) -> DataFrameContainer:
         """
         Create a new container with complete column mapping configuration.
@@ -78,11 +114,14 @@ class DataFrameContainer:
             column_metadata: Optional metadata dictionary
             column_qualifiers: Optional column qualifiers
             parent_column_name_map: Optional parent column name map
-            hidden_columns: Optional list of hidden column names
             table_name: Optional table name
             alias: Optional alias
             cached_schema_getter: Optional function to get cached schema
             partition_hint: Optional partition count from repartition() operations
+            equivalent_snowpark_names: list of sets with old snowpark names that can be resolved with an existing column
+            column_is_hidden: Optional list of booleans indicating whether each column is hidden
+            can_be_cached: Optional boolean indicating if the dataframe can be cached
+            aggregate_metadata: Optional metadata about aggregation for ORDER BY resolution
 
         Returns:
             A new DataFrameContainer instance
@@ -101,7 +140,8 @@ class DataFrameContainer:
             column_metadata,
             column_qualifiers,
             parent_column_name_map,
-            hidden_columns,
+            equivalent_snowpark_names,
+            column_is_hidden,
         )
 
         # Determine the schema getter to use
@@ -129,8 +169,25 @@ class DataFrameContainer:
             alias=alias,
             cached_schema_getter=final_schema_getter,
             partition_hint=partition_hint,
+            can_be_cached=can_be_cached,
+            aggregate_metadata=aggregate_metadata,
         )
 
+    @property
+    def can_be_cached(self) -> bool:
+        """Indicate if the DataFrame can be cached in df_cache"""
+        return self._can_be_cached
+
+    @property
+    def can_be_materialized(self) -> bool:
+        """Indicate if the DataFrame can be materialized in df_cache"""
+        return self._can_be_materialized
+
+    def without_materialization(self):
+        """Prevent the DataFrame from being materialized in df_cache"""
+        self._can_be_materialized = False
+        return self
+
     @property
     def dataframe(self) -> snowpark.DataFrame:
         """Get the underlying Snowpark DataFrame."""
@@ -224,9 +281,10 @@ class DataFrameContainer:
         spark_column_names: list[str],
         snowpark_column_names: list[str],
         column_metadata: dict | None = None,
-        column_qualifiers: list[list[str]] | None = None,
+        column_qualifiers: list[set[ColumnQualifier]] | None = None,
         parent_column_name_map: ColumnNameMap | None = None,
-        hidden_columns: set[HiddenColumn] | None = None,
+        equivalent_snowpark_names: list[set[str]] | None = None,
+        column_is_hidden: list[bool] | None = None,
     ) -> ColumnNameMap:
         """Create a ColumnNameMap with the provided configuration."""
         from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
@@ -237,7 +295,8 @@ class DataFrameContainer:
             column_metadata=column_metadata,
             column_qualifiers=column_qualifiers,
             parent_column_name_map=parent_column_name_map,
-            hidden_columns=hidden_columns,
+            equivalent_snowpark_names=equivalent_snowpark_names,
+            column_is_hidden=column_is_hidden,
         )
 
     @staticmethod
@@ -262,3 +321,38 @@ class DataFrameContainer:
                 )
             ]
         )
+
+    def without_hidden_columns(self) -> DataFrameContainer:
+        from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
+
+        if not any(c.is_hidden for c in self._column_map.columns):
+            return self
+
+        hidden_column_names = [
+            c.snowpark_name for c in self._column_map.columns if c.is_hidden
+        ]
+        visible_columns = [c for c in self._column_map.columns if not c.is_hidden]
+
+        filtered_df = self._dataframe.drop(hidden_column_names)
+        filtered_column_map = ColumnNameMap(
+            spark_column_names=[c.spark_name for c in visible_columns],
+            snowpark_column_names=[c.snowpark_name for c in visible_columns],
+            column_metadata=self._column_map.column_metadata,
+            column_qualifiers=[c.qualifiers for c in visible_columns],
+            parent_column_name_map=self._column_map._parent_column_name_map,
+        )
+
+        return DataFrameContainer(
+            dataframe=filtered_df,
+            column_map=filtered_column_map,
+            table_name=self._table_name,
+            alias=self._alias,
+            cached_schema_getter=lambda: StructType(
+                [
+                    field
+                    for field in self._dataframe.schema.fields
+                    if field.name not in hidden_column_names
+                ]
+            ),
+            partition_hint=self._partition_hint,
+        )

snowflake/snowpark_connect/date_time_format_mapping.py

@@ -24,6 +24,8 @@
 from pyspark.errors.exceptions.base import DateTimeException
 
 from snowflake.snowpark.types import DataType, StringType
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 
 
 # TODO: There are more patterns where spark may throw an error.
@@ -229,9 +231,11 @@ def convert_spark_format_to_snowflake(
     timestamp_input_type: DataType | None = None,
 ):
     if spark_format in {"Y", "w", "W"}:
-        raise DateTimeException(
+        exception = DateTimeException(
             f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter."
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT)
+        raise exception
     snowflake_format = ""
     i = 0
     n = len(spark_format)
@@ -299,28 +303,46 @@ def convert_spark_format_to_snowflake(
                     # Spark's 'a' would be at most 1 times
                     is_valid_a_pattern = spark_format[i : i + 2] != char * 2
                     if not is_valid_a_pattern:
-                        raise DateTimeException(
+                        exception = DateTimeException(
                             f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "h" | "K" | "k" | "H" | "m" | "s" | "d":
                     # Spark's characters would be at most 2 times
                     is_valid_2_patterns = spark_format[i : i + 3] != char * 3
                     if not is_valid_2_patterns:
-                        raise DateTimeException(
+                        exception = DateTimeException(
                             f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "D":
                     # Spark's 'D'' would be at most 3 times
                     is_valid_D_patterns = spark_format[i : i + 4] != char * 4
                     if not is_valid_D_patterns:
-                        raise DateTimeException(
+                        exception = DateTimeException(
                             f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "V":
                     # Spark's 'V' for Zone ID requires 'VV'. A single 'V' is invalid.
                     is_valid_vv_pattern = spark_format[i : i + 2] == "VV"
                     if not is_valid_vv_pattern:
-                        raise DateTimeException("Pattern letter count must be 2: V")
+                        exception = DateTimeException(
+                            "Pattern letter count must be 2: V"
+                        )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "O":
                     # Spark's 'O' would be either 1 or 4.
                     is_valid_o_or_oooo_pattern = spark_format[i : i + 2] != "OO" or (
@@ -328,28 +350,52 @@ def convert_spark_format_to_snowflake(
                         and spark_format[i : i + 5] != "OOOOO"
                     )
                     if not is_valid_o_or_oooo_pattern:
-                        raise DateTimeException(
+                        exception = DateTimeException(
                             "Pattern letter count must be 1 or 4: O"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "q" | "Q" | "z" | "E":
                     # Spark's characters would be at most 4 times
                     is_valid_4_patterns = spark_format[i : i + 5] != char * 5
                     if not is_valid_4_patterns:
-                        raise DateTimeException(f"Too many pattern letters: {char}")
+                        exception = DateTimeException(
+                            f"Too many pattern letters: {char}"
+                        )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "x" | "X" | "Z":
                     # Spark's 'x' or 'X' or 'z' or 'Z' would be at most 5 times
                     is_valid_xz_pattern = spark_format[i : i + 6] != char * 6
                     if not is_valid_xz_pattern:
-                        raise DateTimeException(f"Too many pattern letters: {char}")
+                        exception = DateTimeException(
+                            f"Too many pattern letters: {char}"
+                        )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "y":
                     # Spark's 'y' would be at most 6 times
                     is_valid_y_pattern = spark_format[i : i + 7] != char * 7
                     if not is_valid_y_pattern:
-                        raise DateTimeException(
+                        exception = DateTimeException(
                             f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "C" | "I":
-                    raise DateTimeException(f"Unknown pattern letter: {char}")
+                    exception = DateTimeException(f"Unknown pattern letter: {char}")
+                    attach_custom_error_code(
+                        exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                    )
+                    raise exception
 
             if (
                 spark_format[i : i + 5] in {"M" * 5, "L" * 5}
@@ -362,9 +408,13 @@ def convert_spark_format_to_snowflake(
                 or spark_format[i : i + 3] in {"kkk", "KKK"}
                 or spark_format[i : i + 10] == "SSSSSSSSSS"
             ):
-                raise DateTimeException(
+                exception = DateTimeException(
                     f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter."
                 )
+                attach_custom_error_code(
+                    exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                )
+                raise exception
 
             matched_pattern = False
 
@@ -375,7 +425,11 @@ def convert_spark_format_to_snowflake(
                         spark_key
                     ]
                     if isinstance(snowflake_equivalent, _UnsupportedSparkFormatPattern):
-                        raise DateTimeException(snowflake_equivalent.message)
+                        exception = DateTimeException(snowflake_equivalent.message)
+                        attach_custom_error_code(
+                            exception, ErrorCodes.UNSUPPORTED_OPERATION
+                        )
+                        raise exception
                     if snowflake_equivalent is not None:
                         snowflake_format += snowflake_equivalent
                     i += len(spark_key)
@@ -389,7 +443,11 @@ def convert_spark_format_to_snowflake(
                     isinstance(timestamp_input_type, StringType)
                     and char not in snowflake_time_format_separator
                 ):
-                    raise DateTimeException(f"Illegal pattern character: {char}")
+                    exception = DateTimeException(f"Illegal pattern character: {char}")
+                    attach_custom_error_code(
+                        exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                    )
+                    raise exception
 
                 snowflake_format += f'"{char}"'
                 i += 1

snowflake/snowpark_connect/error/error_codes.py

@@ -0,0 +1,50 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+"""
+Error code constants for Snowpark Connect.
+
+This module defines custom error codes that can be attached to exceptions
+and included in gRPC error responses.
+"""
+
+
+class ErrorCodes:
+    """Constants for Snowpark Connect custom error codes."""
+
+    # 1000-1999: Startup related errors
+    MISSING_DATABASE = 1001
+    MISSING_SCHEMA = 1002
+    RESOURCE_INITIALIZATION_FAILED = 1003
+    TCP_PORT_ALREADY_IN_USE = 1004
+    INVALID_SPARK_CONNECT_URL = 1005
+    INVALID_STARTUP_INPUT = 1006
+    INVALID_STARTUP_OPERATION = 1007
+    STARTUP_CONNECTION_FAILED = 1008
+
+    # 2000-2999: Configuration related errors
+    INVALID_CONFIG_VALUE = 2001
+    CONFIG_CHANGE_NOT_ALLOWED = 2002
+    CONFIG_NOT_ENABLED = 2003
+
+    # 3000-3999: User code errors
+    INVALID_SQL_SYNTAX = 3001
+    TYPE_MISMATCH = 3002
+    INVALID_CAST = 3003
+    INVALID_FUNCTION_ARGUMENT = 3004
+    ARRAY_INDEX_OUT_OF_BOUNDS = 3005
+    DIVISION_BY_ZERO = 3006
+    INVALID_INPUT = 3007
+    INVALID_OPERATION = 3008
+    INSUFFICIENT_INPUT = 3009
+
+    # 4000-4999: What we don't support
+    UNSUPPORTED_OPERATION = 4001
+    UNSUPPORTED_TYPE = 4002
+
+    # 5000-5999: Internal errors
+    INTERNAL_ERROR = 5001
+    TABLE_NOT_FOUND = 5002
+    COLUMN_NOT_FOUND = 5003
+    AMBIGUOUS_COLUMN_NAME = 5004