snowpark-connect 0.30.1__py3-none-any.whl → 0.32.0__py3-none-any.whl

snowflake/snowpark_connect/__init__.py

@@ -10,6 +10,7 @@ sys.path.append(str(pathlib.Path(__file__).parent / "includes/python"))
 
 from .server import get_session  # noqa: E402, F401
 from .server import start_session  # noqa: E402, F401
+from .utils.session import skip_session_configuration  # noqa: E402, F401
 
 # Turn off catalog warning for Snowpark
 sp_logger = logging.getLogger("snowflake.snowpark")

snowflake/snowpark_connect/column_name_handler.py

@@ -13,14 +13,17 @@ from functools import cached_property
 from pyspark.errors.exceptions.base import AnalysisException
 
 from snowflake.snowpark import DataFrame
-from snowflake.snowpark._internal.analyzer.analyzer_utils import (
-    quote_name_without_upper_casing,
-    unquote_if_quoted,
-)
+from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted
 from snowflake.snowpark._internal.utils import quote_name
 from snowflake.snowpark.types import StructType
+from snowflake.snowpark_connect.column_qualifier import ColumnQualifier
 from snowflake.snowpark_connect.config import global_config
-from snowflake.snowpark_connect.utils.context import get_current_operation_scope
+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.utils.context import (
+    get_current_operation_scope,
+    get_is_processing_order_by,
+)
 from snowflake.snowpark_connect.utils.identifiers import (
     split_fully_qualified_spark_name,
 )
@@ -92,31 +95,15 @@ def make_column_names_snowpark_compatible(
 class ColumnNames:
     spark_name: str
     snowpark_name: str
-    qualifiers: list[str]
+    qualifiers: set[ColumnQualifier]
     catalog_info: str | None = None  # Catalog from fully qualified name
     database_info: str | None = None  # Database from fully qualified name
 
-
-def get_list_of_spark_names_for_column(column_names: ColumnNames) -> list[str]:
-    """
-    Returns a list of Spark names for a given ColumnNames object.
-    This is useful when a single Spark name maps to multiple names due to table alias.
-
-    For example, if the column name is 'id' and the qualifiers are ['db', 'table'],
-    then the possible Spark names are:
-      ['id', 'db.table.id', 'table.id']
-    """
-    spark_name = column_names.spark_name
-    qualifiers = column_names.qualifiers
-
-    qualifier_suffixes_list = [
-        ".".join(quote_name_without_upper_casing(x) for x in qualifiers[i:])
-        for i in range(len(qualifiers))
-    ]
-    return [spark_name] + [
-        f"{qualifier_suffix}.{spark_name}"
-        for qualifier_suffix in qualifier_suffixes_list
-    ]
+    def all_spark_names_including_qualified_names(self):
+        all_names = [self.spark_name]
+        for qualifier in self.qualifiers:
+            all_names.extend(qualifier.all_qualified_names(self.spark_name))
+        return all_names
 
 
 class ColumnNameMap:
@@ -128,13 +115,13 @@ class ColumnNameMap:
             [], bool
         ] = lambda: global_config.spark_sql_caseSensitive,
         column_metadata: dict | None = None,
-        column_qualifiers: list[list[str]] | None = None,
+        column_qualifiers: list[set[ColumnQualifier]] = None,
         parent_column_name_map: ColumnNameMap | None = None,
     ) -> None:
         """
         spark_column_names: Original spark column names
         snowpark_column_names: Snowpark column names
-        column_metadata: This field is used to store metadata related to columns. Since Snowpark’s Struct type does not support metadata,
+        column_metadata: This field is used to store metadata related to columns. Since Snowpark's Struct type does not support metadata,
         we use this attribute to store any metadata related to the columns.
         The key is the original Spark column name, and the value is the metadata.
         example: Dict('age', {'foo': 'bar'})
@@ -142,7 +129,7 @@ class ColumnNameMap:
         parent_column_name_map: parent ColumnNameMap
         """
         self.columns: list[ColumnNames] = []
-        self.spark_to_col = defaultdict(list)
+        self.spark_to_col: defaultdict[str, list[ColumnNames]] = defaultdict(list)
         self.uppercase_spark_to_col = defaultdict(list)
         self.snowpark_to_col = defaultdict(list)
         self.is_case_sensitive = is_case_sensitive
@@ -181,21 +168,18 @@ class ColumnNameMap:
             c = ColumnNames(
                 spark_name=spark_name,
                 snowpark_name=snowpark_column_names[i],
-                qualifiers=column_qualifiers[i] if column_qualifiers else [],
+                qualifiers=column_qualifiers[i]
+                if column_qualifiers and column_qualifiers[i]
+                else {ColumnQualifier.no_qualifier()},
                 catalog_info=catalog_info,
                 database_info=database_info,
             )
             self.columns.append(c)
 
-            # we want to store all the spark names including qualifiers (these are generated from table alias or dataframe alias)
-            spark_names_including_qualifier = get_list_of_spark_names_for_column(c)
-
-            for spark_name_including_qualifier in spark_names_including_qualifier:
+            for spark_name in c.all_spark_names_including_qualified_names():
                 # the same spark name can map to multiple snowpark names
-                self.spark_to_col[spark_name_including_qualifier].append(c)
-                self.uppercase_spark_to_col[
-                    spark_name_including_qualifier.upper()
-                ].append(c)
+                self.spark_to_col[spark_name].append(c)
+                self.uppercase_spark_to_col[spark_name.upper()].append(c)
 
             # the same snowpark name can map to multiple spark column
             # e.g. df.select(date_format('dt', 'yyy'), date_format('dt', 'yyyy')) ->
@@ -353,18 +337,77 @@ class ColumnNameMap:
 
         snowpark_names_len = len(snowpark_names)
         if snowpark_names_len > 1:
-            raise AnalysisException(
-                f"Ambiguous spark column name {spark_column_name}, potential snowpark column names {snowpark_names}"
-            )
+            # Check if this is a case where we have identical expressions that can be safely resolved to the first one
+            # This commonly happens with GROUP BY expressions that also appear in SELECT clauses
+            if (
+                get_is_processing_order_by()
+                and self._can_resolve_ambiguous_identical_expressions(
+                    resolved_name, snowpark_names
+                )
+            ):
+                # All the ambiguous columns represent the same expression, so we can safely use the first one
+                return snowpark_names[0]
+            else:
+                exception = AnalysisException(
+                    f"Ambiguous spark column name {spark_column_name}, potential snowpark column names {snowpark_names}"
+                )
+                attach_custom_error_code(exception, ErrorCodes.AMBIGUOUS_COLUMN_NAME)
+                raise exception
         elif snowpark_names_len == 0:
             if allow_non_exists:
                 return None
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Spark column name {spark_column_name} does not exist"
                 )
+                attach_custom_error_code(exception, ErrorCodes.COLUMN_NOT_FOUND)
+                raise exception
         return snowpark_names[0]
 
+    def _can_resolve_ambiguous_identical_expressions(
+        self, spark_column_name: str, snowpark_names: list[str]
+    ) -> bool:
+        """
+        Determine if ambiguous columns represent identical expressions that can be safely resolved to the first one.
+
+        This handles the common case where the same expression (like a UDF call) appears multiple times
+        in a SELECT clause within a GROUP BY query. Since they're the same expression operating on the
+        same grouped data, they will have identical values, so we can safely resolve to any of them.
+
+        Args:
+            spark_column_name: The Spark column name that has multiple mappings, make sure resolve this reforehand
+            snowpark_names: List of Snowpark column names that map to this Spark column name
+
+        Returns:
+            True if we can safely resolve to the first snowpark column, False otherwise
+        """
+        if spark_column_name not in self.spark_to_col:
+            return False
+
+        columns: list[ColumnNames] = self.spark_to_col[spark_column_name]
+
+        # If we don't have multiple columns, there's no ambiguity to resolve
+        if len(columns) <= 1:
+            return False
+
+        # Check if all the snowpark names correspond to columns that have identical underlying expressions
+        # We'll compare the actual column objects to see if they represent the same computation
+        first_column = columns[0]
+
+        for column in columns[1:]:
+            if first_column.qualifiers != column.qualifiers:
+                return False
+
+        # Additional safety check: ensure all snowpark names are actually in our mapping
+        for snowpark_name in snowpark_names:
+            if snowpark_name not in self.snowpark_to_col:
+                return False
+
+        # If we reach here, the columns appear to be identical expressions from the same context
+        # This commonly happens in GROUP BY scenarios where the same expression appears in both
+        # the grouping clause and the select clause
+        return True
+
     def get_spark_column_names_from_snowpark_column_names(
         self,
         snowpark_column_names: list[str],
@@ -390,16 +433,20 @@ class ColumnNameMap:
         )
         spark_names_len = len(spark_names)
         if spark_names_len > 1:
-            raise AnalysisException(
+            exception = AnalysisException(
                 f"Ambiguous snowpark column name {snowpark_column_name}, potential spark column names {spark_names}"
             )
+            attach_custom_error_code(exception, ErrorCodes.AMBIGUOUS_COLUMN_NAME)
+            raise exception
         elif spark_names_len == 0:
             if allow_non_exists:
                 return None
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Snowpark column name {snowpark_column_name} does not exist"
                 )
+                attach_custom_error_code(exception, ErrorCodes.COLUMN_NOT_FOUND)
+                raise exception
         return spark_names[0]
 
     def get_spark_column_name(self, idx: int) -> str:
@@ -409,32 +456,30 @@ class ColumnNameMap:
         return [c.spark_name for c in self.columns]
 
     def get_spark_and_snowpark_columns_with_qualifier_for_qualifier(
-        self, qualifiers_input: list[str]
-    ) -> tuple[list[str], list[str], list[list[str]]]:
+        self, target_qualifier: ColumnQualifier
+    ) -> tuple[list[str], list[str], list[set[ColumnQualifier]]]:
         """
-        Returns the Spark and Snowpark column names along with their qualifiers for the specified qualifiers.
-        If a column does not have a qualifier, it will be None.
+        Returns the Spark and Snowpark column names along with their qualifiers for the specified qualifier.
         """
-        spark_columns = []
-        snowpark_columns = []
-        qualifiers = []
+        spark_columns: list[str] = []
+        snowpark_columns: list[str] = []
+        qualifiers: list[set[ColumnQualifier]] = []
 
+        normalized_qualifier = target_qualifier
         if not self.is_case_sensitive():
-            qualifiers_input = [q.upper() for q in qualifiers_input]
+            normalized_qualifier = target_qualifier.to_upper()
 
-        for c in self.columns:
-            col_qualifiers = (
-                [q.upper() for q in c.qualifiers]
+        for column in self.columns:
+            # Normalize all qualifiers for comparison
+            column_qualifiers: set[ColumnQualifier] = (
+                {q.to_upper() for q in iter(column.qualifiers)}
                 if not self.is_case_sensitive()
-                else c.qualifiers
+                else column.qualifiers
             )
-            if len(col_qualifiers) < len(qualifiers_input):
-                # If the column has fewer qualifiers than the input, it cannot match
-                continue
-            if col_qualifiers[-len(qualifiers_input) :] == qualifiers_input:
-                spark_columns.append(c.spark_name)
-                snowpark_columns.append(c.snowpark_name)
-                qualifiers.append(c.qualifiers)
+            if any([q.matches(normalized_qualifier) for q in column_qualifiers]):
+                spark_columns.append(column.spark_name)
+                snowpark_columns.append(column.snowpark_name)
+                qualifiers.append(column.qualifiers)
 
         return spark_columns, snowpark_columns, qualifiers
 
@@ -448,19 +493,17 @@ class ColumnNameMap:
             if self._quote_if_unquoted(c) not in cols_to_drop
         ]
 
-    def get_qualifiers(self) -> list[list[str]]:
+    def get_qualifiers(self) -> list[set[ColumnQualifier]]:
         """
         Returns the qualifiers for the columns.
-        If a column does not have a qualifier, it will be None.
         """
         return [c.qualifiers for c in self.columns]
 
     def get_qualifiers_for_columns_after_drop(
         self, cols_to_drop: list[str]
-    ) -> list[list[str]]:
+    ) -> list[set[ColumnQualifier]]:
         """
         Returns the qualifiers for the columns after dropping the specified columns.
-        If a column is dropped, its qualifier will be None.
         """
         return [
             c.qualifiers
@@ -471,10 +514,25 @@ class ColumnNameMap:
     def get_qualifier_for_spark_column(
         self,
         spark_column_name: str,
-    ) -> list[str]:
+    ) -> ColumnQualifier:
+        """
+        Backward compatibility: returns the first qualifier for the given Spark column name.
+        Throws if more than one qualifier exists.
+        """
+        qualifiers = self.get_qualifiers_for_spark_column(spark_column_name)
+        if len(qualifiers) > 1:
+            raise ValueError(
+                "Shouldn't happen. Multiple qualifiers found; expected only one."
+            )
+        return next(iter(qualifiers))
+
+    def get_qualifiers_for_spark_column(
+        self,
+        spark_column_name: str,
+    ) -> set[ColumnQualifier]:
         """
         Returns the qualifier for the specified Spark column name.
-        If the column does not exist, returns None.
+        If the column does not exist, returns empty ColumnQualifier.
         """
         if not self.is_case_sensitive():
             name = spark_column_name.upper()
@@ -486,7 +544,7 @@ class ColumnNameMap:
         col = mapping.get(name)
 
         if col is None or len(col) == 0:
-            return []
+            return {ColumnQualifier.no_qualifier()}
 
         return col[0].qualifiers
 
@@ -518,7 +576,7 @@ class ColumnNameMap:
 
     def with_columns(
         self, new_spark_columns: list[str], new_snowpark_columns: list[str]
-    ) -> tuple[list[str], list[str], list[list[str]]]:
+    ) -> tuple[list[str], list[str], list[set[ColumnQualifier]]]:
         """
         Returns an ordered list of spark and snowpark column names after adding the new columns through a withColumns call.
         All replaced columns retain their ordering in the dataframe. The new columns are added to the end of the list.
@@ -547,7 +605,7 @@ class ColumnNameMap:
                 removed_index.add(index)
                 spark_columns.append(new_spark_columns[index])
                 snowpark_columns.append(new_snowpark_columns[index])
-                qualifiers.append([])
+                qualifiers.append({ColumnQualifier.no_qualifier()})
             else:
                 spark_columns.append(c.spark_name)
                 snowpark_columns.append(c.snowpark_name)
@@ -557,7 +615,7 @@ class ColumnNameMap:
             if i not in removed_index:
                 spark_columns.append(new_spark_columns[i])
                 snowpark_columns.append(new_snowpark_columns[i])
-                qualifiers.append([])
+                qualifiers.append({ColumnQualifier.no_qualifier()})
 
         return spark_columns, snowpark_columns, qualifiers
 
@@ -604,14 +662,18 @@ class JoinColumnNameMap(ColumnNameMap):
             if allow_non_exists:
                 return None
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Spark column name {spark_column_name} does not exist in either left or right DataFrame"
                 )
+                attach_custom_error_code(exception, ErrorCodes.COLUMN_NOT_FOUND)
+                raise exception
 
         if (snowpark_column_name_in_right is not None) and (
             snowpark_column_name_in_left is not None
         ):
-            raise AnalysisException(f"Ambiguous column name {spark_column_name}")
+            exception = AnalysisException(f"Ambiguous column name {spark_column_name}")
+            attach_custom_error_code(exception, ErrorCodes.AMBIGUOUS_COLUMN_NAME)
+            raise exception
 
         snowpark_name = (
             snowpark_column_name_in_right
@@ -637,60 +699,94 @@ class JoinColumnNameMap(ColumnNameMap):
     def get_snowpark_column_names_from_spark_column_names(
         self, spark_column_names: list[str], return_first: bool = False
     ) -> list[str]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_spark_column_names_from_snowpark_column_names(
         self,
         snowpark_column_names: list[str],
     ) -> list[str]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_spark_column_name_from_snowpark_column_name(
-        self, snowpark_column_name: str
+        self,
+        snowpark_column_name: str,
+        allow_non_exists: bool = False,
     ) -> str:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_spark_columns(self) -> list[str]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_snowpark_columns(self) -> list[str]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_snowpark_columns_after_drop(self, cols_to_drop: list[str]) -> list[str]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_renamed_nested_column_name(self, name) -> str | None:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def has_spark_column(self, spark_column_name: str) -> bool:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def snowpark_to_spark_map(self) -> dict[str, str]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def spark_to_snowpark_for_pattern(self, pattern: str) -> list[tuple[str, str]]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def with_columns(
         self, new_spark_columns: list[str], new_snowpark_columns: list[str]
-    ) -> tuple[list[str], list[str], list[list[str]]]:
-        raise NotImplementedError("Method not implemented!")
+    ) -> tuple[list[str], list[str], list[set[ColumnQualifier]]]:
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
-    def get_qualifiers(self) -> list[list[str]]:
-        raise NotImplementedError("Method not implemented!")
+    def get_qualifiers(self) -> list[set[ColumnQualifier]]:
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_qualifiers_for_columns_after_drop(
         self, cols_to_drop: list[str]
-    ) -> list[list[str]]:
-        raise NotImplementedError("Method not implemented!")
+    ) -> list[set[ColumnQualifier]]:
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_spark_and_snowpark_columns_with_qualifier_for_qualifier(
-        self, qualifiers_input: list[str]
-    ) -> tuple[list[str], list[str], list[list[str]]]:
-        raise NotImplementedError("Method not implemented!")
-
-    def get_qualifier_for_spark_column(self, spark_column_name: str) -> list[str]:
-
+        self, target_qualifier: list[str]
+    ) -> tuple[list[str], list[str], list[set[ColumnQualifier]]]:
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
+
+    def get_qualifiers_for_spark_column(
+        self, spark_column_name: str
+    ) -> set[ColumnQualifier]:
+        return {self.get_qualifier_for_spark_column(spark_column_name)}
+
+    def get_qualifier_for_spark_column(self, spark_column_name: str) -> ColumnQualifier:
         qualifier_left = self.left_column_mapping.get_qualifier_for_spark_column(
             spark_column_name
         )
@@ -698,7 +794,9 @@ class JoinColumnNameMap(ColumnNameMap):
             spark_column_name
         )
 
-        if (len(qualifier_left) > 0) and (len(qualifier_right) > 0):
-            raise AnalysisException(f"Ambiguous column name {spark_column_name}")
+        if (not qualifier_left.is_empty) and (not qualifier_right.is_empty):
+            exception = AnalysisException(f"Ambiguous column name {spark_column_name}")
+            attach_custom_error_code(exception, ErrorCodes.AMBIGUOUS_COLUMN_NAME)
+            raise exception
 
-        return qualifier_right if len(qualifier_left) == 0 else qualifier_left
+        return qualifier_right if qualifier_left.is_empty else qualifier_left

snowflake/snowpark_connect/column_qualifier.py

@@ -0,0 +1,47 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from snowflake.snowpark._internal.analyzer.analyzer_utils import (
+    quote_name_without_upper_casing,
+)
+
+
+@dataclass(frozen=True)
+class ColumnQualifier:
+    parts: tuple[str, ...]
+
+    def __post_init__(self) -> None:
+        if not all(isinstance(x, str) for x in self.parts):
+            raise TypeError("ColumnQualifier.parts must be strings")
+
+    @property
+    def is_empty(self) -> bool:
+        return len(self.parts) == 0
+
+    @classmethod
+    def no_qualifier(cls) -> ColumnQualifier:
+        return cls(())
+
+    def all_qualified_names(self, name: str) -> list[str]:
+        qualifier_parts = self.parts
+        qualifier_prefixes = [
+            ".".join(quote_name_without_upper_casing(x) for x in qualifier_parts[i:])
+            for i in range(len(qualifier_parts))
+        ]
+        return [f"{prefix}.{name}" for prefix in qualifier_prefixes]
+
+    def to_upper(self):
+        return ColumnQualifier(tuple(part.upper() for part in self.parts))
+
+    def matches(self, target: ColumnQualifier) -> bool:
+        if self.is_empty or target.is_empty:
+            return False
+        # If the column has fewer qualifiers than the target, it cannot match
+        if len(self.parts) < len(target.parts):
+            return False
+        return self.parts[-len(target.parts) :] == target.parts