snowpark-connect 0.31.0__tar.gz → 0.33.0__tar.gz

{snowpark_connect-0.31.0/src/snowpark_connect.egg-info → snowpark_connect-0.33.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowpark-connect
-Version: 0.31.0
+Version: 0.33.0
 Summary: Snowpark Connect for Spark
 Author: Snowflake, Inc
 License: Apache License, Version 2.0
@@ -9,11 +9,13 @@ Description-Content-Type: text/markdown
 License-File: LICENSE.txt
 License-File: LICENSE-binary
 License-File: NOTICE-binary
+Requires-Dist: snowpark-connect-deps-1==3.56.2
+Requires-Dist: snowpark-connect-deps-2==3.56.2
 Requires-Dist: certifi>=2025.1.31
 Requires-Dist: cloudpickle
 Requires-Dist: fsspec[http]
 Requires-Dist: jpype1
-Requires-Dist: protobuf<5.0,>=4.25.3
+Requires-Dist: protobuf<6.32.0,>=4.25.3
 Requires-Dist: s3fs>=2025.3.0
 Requires-Dist: snowflake.core<2,>=1.0.5
 Requires-Dist: snowflake-snowpark-python[pandas]<1.41.0,==1.40.0
@@ -29,11 +31,14 @@ Requires-Dist: grpcio-status<1.63,>=1.56.0
 Requires-Dist: googleapis-common-protos>=1.56.4
 Requires-Dist: numpy<2,>=1.15
 Requires-Dist: gcsfs>=2025.9.0
+Provides-Extra: jdk
+Requires-Dist: jdk4py==17.0.9.2; extra == "jdk"
 Dynamic: author
 Dynamic: description
 Dynamic: description-content-type
 Dynamic: license
 Dynamic: license-file
+Dynamic: provides-extra
 Dynamic: requires-dist
 Dynamic: requires-python
 Dynamic: summary

{snowpark_connect-0.31.0 → snowpark_connect-0.33.0}/setup.py

@@ -29,7 +29,6 @@ setup(
     package_data={
         "": ["*.json"],
         "snowflake.snowpark_connect": ["resources/*.jar"],
-        "snowflake.snowpark_connect.includes": ["jars/*.jar"],
     },
     package_dir={"": "src"},
     scripts=[
@@ -39,11 +38,13 @@ setup(
     ],
     python_requires=">=3.10,<3.13",
     install_requires=[
+        "snowpark-connect-deps-1==3.56.2",  # Spark JAR dependencies (71MB)
+        "snowpark-connect-deps-2==3.56.2",  # Other JAR dependencies (53MB)
         "certifi>=2025.1.31",  # prod-297255-inc0132291
         "cloudpickle",
         "fsspec[http]",
         "jpype1",
-        "protobuf>=4.25.3,<5.0",
+        "protobuf>=4.25.3,<6.32.0",
         "s3fs>=2025.3.0",  # prod-297255-inc0132291
         "snowflake.core>=1.0.5,<2",
         "snowflake-snowpark-python[pandas]==1.40.0,<1.41.0",
@@ -61,4 +62,7 @@ setup(
         "numpy>=1.15,<2",
         "gcsfs>=2025.9.0",
     ],
+    extras_require={
+        "jdk": ["jdk4py==17.0.9.2"],
+    },
 )

{snowpark_connect-0.31.0 → snowpark_connect-0.33.0}/src/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")

{snowpark_connect-0.31.0 → snowpark_connect-0.33.0}/src/snowflake/snowpark_connect/column_name_handler.py

@@ -13,12 +13,10 @@ 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.error.error_codes import ErrorCodes
 from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
@@ -29,6 +27,7 @@ from snowflake.snowpark_connect.utils.context import (
 from snowflake.snowpark_connect.utils.identifiers import (
     split_fully_qualified_spark_name,
 )
+from snowflake.snowpark_connect.utils.sequence import next_unique_num
 
 ALREADY_QUOTED = re.compile('^(".+")$', re.DOTALL)
 
@@ -48,6 +47,7 @@ def set_schema_getter(df: DataFrame, get_schema: Callable[[], StructType]) -> No
     df.__class__ = PatchedDataFrame
 
 
+# TODO replace plan_id-offset with single unique value
 def make_column_names_snowpark_compatible(
     names: list[str], plan_id: int, offset: int = 0
 ) -> list[str]:
@@ -93,35 +93,27 @@ def make_column_names_snowpark_compatible(
     ]
 
 
+def make_unique_snowpark_name(spark_name: str) -> str:
+    """
+    Returns a snowpark column name that's guaranteed to be unique in this session,
+    by appending "#<unique number>" to the given spark name.
+    """
+    return quote_name(f"{spark_name}#{next_unique_num()}")
+
+
 @dataclass(frozen=True)
 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:
@@ -133,13 +125,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'})
@@ -155,7 +147,7 @@ class ColumnNameMap:
 
         # Rename chain dictionary to track column renaming history
         self.rename_chains: dict[str, str] = {}  # old_name -> new_name mapping
-        self.current_columns: set[str] = set()  # Current column names
+        self.current_columns: set[str] = set()  # current column names
 
         # Parent ColumnNameMap classes
         self._parent_column_name_map = parent_column_name_map
@@ -186,21 +178,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 set(),
                 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')) ->
@@ -405,7 +394,7 @@ class ColumnNameMap:
         if spark_column_name not in self.spark_to_col:
             return False
 
-        columns = self.spark_to_col[spark_column_name]
+        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:
@@ -416,30 +405,7 @@ class ColumnNameMap:
         first_column = columns[0]
 
         for column in columns[1:]:
-            # Check snowpark_type attribute
-            # If one has the attribute but the other doesn't, they're different
-            if hasattr(first_column, "snowpark_type") != hasattr(
-                column, "snowpark_type"
-            ):
-                return False
-            # If both have the attribute and values differ, they're different expressions
-            if (
-                hasattr(first_column, "snowpark_type")
-                and hasattr(column, "snowpark_type")
-                and first_column.snowpark_type != column.snowpark_type
-            ):
-                return False
-
-            # Check qualifiers attribute
-            # If one has the attribute but the other doesn't, they're different
-            if hasattr(first_column, "qualifiers") != hasattr(column, "qualifiers"):
-                return False
-            # If both have the attribute and values differ, they might be from different contexts
-            if (
-                hasattr(first_column, "qualifiers")
-                and hasattr(column, "qualifiers")
-                and first_column.qualifiers != column.qualifiers
-            ):
+            if first_column.qualifiers != column.qualifiers:
                 return False
 
         # Additional safety check: ensure all snowpark names are actually in our mapping
@@ -500,32 +466,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
 
@@ -539,19 +503,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
@@ -559,13 +521,13 @@ class ColumnNameMap:
             if self._quote_if_unquoted(c.snowpark_name) not in cols_to_drop
         ]
 
-    def get_qualifier_for_spark_column(
+    def get_qualifiers_for_spark_column(
         self,
         spark_column_name: str,
-    ) -> list[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()
@@ -577,7 +539,7 @@ class ColumnNameMap:
         col = mapping.get(name)
 
         if col is None or len(col) == 0:
-            return []
+            return set()
 
         return col[0].qualifiers
 
@@ -609,7 +571,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.
@@ -638,7 +600,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(set())
             else:
                 spark_columns.append(c.spark_name)
                 snowpark_columns.append(c.snowpark_name)
@@ -648,7 +610,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(set())
 
         return spark_columns, snowpark_columns, qualifiers
 
@@ -658,6 +620,67 @@ class ColumnNameMap:
         else:
             return spark_name.upper()
 
+    def get_columns_after_join(
+        self, other: ColumnNameMap, join_columns: list[str]
+    ) -> list[ColumnNames]:
+        """
+        Returns a list of columns (names and qualifiers) after a using_columns join with the given column map
+        """
+
+        join_column_names = {self._normalized_spark_name(c) for c in join_columns}
+        other_join_columns: dict[str, ColumnNames] = {}
+        other_remaining_columns: list[ColumnNames] = []
+        for oc in other.columns:
+            col_name = self._normalized_spark_name(oc.spark_name)
+            # only take the first matching column
+            if col_name in join_column_names and col_name not in other_join_columns:
+                other_join_columns[col_name] = oc
+            else:
+                other_remaining_columns.append(oc)
+
+        joined_columns: list[ColumnNames] = []
+        visited: set[str] = set()
+        # add local columns first, we're in the left side of the join
+        for c in self.columns:
+            col_name = self._normalized_spark_name(c.spark_name)
+            if col_name in join_column_names and col_name not in visited:
+                visited.add(col_name)
+                qualifiers = c.qualifiers | other_join_columns[col_name].qualifiers
+                joined_columns.append(
+                    ColumnNames(c.spark_name, c.snowpark_name, qualifiers)
+                )
+            else:
+                joined_columns.append(c)
+
+        # add other columns, excluding join columns
+        return joined_columns + other_remaining_columns
+
+    def get_column_indexes(self, spark_names: list[str]) -> list[int]:
+        """
+        Returns the first positions of the given spark_names in this column mapping.
+        Used to reorder columns after a using_columns join.
+        """
+        # mapping from normalized spark name ot the first index of the column in the mapping
+        column_indexes = {}
+
+        for i, c in enumerate(self.columns):
+            col_name = self._normalized_spark_name(c.spark_name)
+            if col_name not in column_indexes:
+                column_indexes[col_name] = i
+
+        # return indexes for given columns
+        return [column_indexes[self._normalized_spark_name(c)] for c in spark_names]
+
+    def get_conflicting_snowpark_columns(self, other: ColumnNameMap) -> set[str]:
+        conflicting_columns = set()
+        snowpark_names = {c.snowpark_name for c in self.columns}
+
+        for c in other.columns:
+            if c.snowpark_name in snowpark_names:
+                conflicting_columns.add(c.snowpark_name)
+
+        return conflicting_columns
+
 
 class JoinColumnNameMap(ColumnNameMap):
     def __init__(
@@ -745,7 +768,9 @@ class JoinColumnNameMap(ColumnNameMap):
         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:
         exception = NotImplementedError("Method not implemented!")
         attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
@@ -788,42 +813,55 @@ class JoinColumnNameMap(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]]]:
         exception = NotImplementedError("Method not implemented!")
         attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
         raise exception
 
-    def get_qualifiers(self) -> list[list[str]]:
+    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]]:
+    ) -> 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]]]:
+        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_qualifier_for_spark_column(self, spark_column_name: str) -> list[str]:
-
-        qualifier_left = self.left_column_mapping.get_qualifier_for_spark_column(
+    def get_qualifiers_for_spark_column(
+        self, spark_column_name: str
+    ) -> set[ColumnQualifier]:
+        qualifiers_left = self.left_column_mapping.get_qualifiers_for_spark_column(
             spark_column_name
         )
-        qualifier_right = self.right_column_mapping.get_qualifier_for_spark_column(
+        qualifiers_right = self.right_column_mapping.get_qualifiers_for_spark_column(
             spark_column_name
         )
 
-        if (len(qualifier_left) > 0) and (len(qualifier_right) > 0):
+        if (len(qualifiers_left) > 0) and (len(qualifiers_right) > 0):
             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 qualifiers_right if len(qualifiers_left) == 0 else qualifiers_left
+
+    def get_columns_after_join(
+        self, other: ColumnNameMap, join_columns: list[str]
+    ) -> list[ColumnNames]:
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
+
+    def get_column_indexes(self, spark_names: list[str]) -> list[int]:
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception

snowpark_connect-0.33.0/src/snowflake/snowpark_connect/column_qualifier.py

@@ -0,0 +1,43 @@
+#
+# 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
+
+    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

{snowpark_connect-0.31.0 → snowpark_connect-0.33.0}/src/snowflake/snowpark_connect/dataframe_container.py

@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Callable
 
 from snowflake import snowpark
 from snowflake.snowpark.types import StructField, StructType
+from snowflake.snowpark_connect.column_qualifier import ColumnQualifier
 
 if TYPE_CHECKING:
     from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
@@ -58,7 +59,7 @@ 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,
         table_name: str | None = None,
         alias: str | None = None,
@@ -220,7 +221,7 @@ 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,
     ) -> ColumnNameMap:
         """Create a ColumnNameMap with the provided configuration."""

{snowpark_connect-0.31.0 → snowpark_connect-0.33.0}/src/snowflake/snowpark_connect/execute_plan/map_execution_command.py

@@ -54,9 +54,11 @@ def _create_column_rename_map(
             new_column_name = (
                 f"{new_column_name}_DEDUP_{column_counts[normalized_name] - 1}"
             )
-            renamed_cols.append(ColumnNames(new_column_name, col.snowpark_name, []))
+            renamed_cols.append(ColumnNames(new_column_name, col.snowpark_name, set()))
         else:
-            not_renamed_cols.append(ColumnNames(new_column_name, col.snowpark_name, []))
+            not_renamed_cols.append(
+                ColumnNames(new_column_name, col.snowpark_name, set())
+            )
 
     if len(renamed_cols) == 0:
         return {

{snowpark_connect-0.31.0 → snowpark_connect-0.33.0}/src/snowflake/snowpark_connect/expression/hybrid_column_map.py

@@ -148,14 +148,15 @@ class HybridColumnMap:
                     exp, self.aggregated_column_map, self.aggregated_typer
                 )
 
-        # For other expression types, try aggregated context first (likely references to computed values)
         try:
+            # 1. Evaluate the expression using the input grouping columns. i.e input_df.
+            # If not found, use the aggregate alias.
+            return map_expression(exp, self.input_column_map, self.input_typer)
+        except Exception:
+            # Fall back to input context
             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(

{snowpark_connect-0.31.0 → snowpark_connect-0.33.0}/src/snowflake/snowpark_connect/expression/map_expression.py

@@ -230,18 +230,18 @@ def map_expression(
                     | exp.sort_order.SORT_DIRECTION_ASCENDING
                 ):
                     if exp.sort_order.null_ordering == exp.sort_order.SORT_NULLS_LAST:
-                        return [child_name], snowpark_fn.asc_nulls_last(child_column)
+                        col = snowpark_fn.asc_nulls_last(child_column.col)
                     else:
                         # If nulls are not specified or null_ordering is FIRST in the sort order, Spark defaults to nulls
                         # first in the case of ascending sort order.
-                        return [child_name], snowpark_fn.asc_nulls_first(child_column)
+                        col = snowpark_fn.asc_nulls_first(child_column.col)
                 case exp.sort_order.SORT_DIRECTION_DESCENDING:
                     if exp.sort_order.null_ordering == exp.sort_order.SORT_NULLS_FIRST:
-                        return [child_name], snowpark_fn.desc_nulls_first(child_column)
+                        col = snowpark_fn.desc_nulls_first(child_column.col)
                     else:
                         # If nulls are not specified or null_ordering is LAST in the sort order, Spark defaults to nulls
                         # last in the case of descending sort order.
-                        return [child_name], snowpark_fn.desc_nulls_last(child_column)
+                        col = snowpark_fn.desc_nulls_last(child_column.col)
                 case _:
                     exception = ValueError(
                         f"Invalid sort direction {exp.sort_order.direction}"
@@ -250,6 +250,7 @@ def map_expression(
                         exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
                     )
                     raise exception
+            return [child_name], TypedColumn(col, lambda: typer.type(col))
         case "unresolved_attribute":
             col_name, col = map_att.map_unresolved_attribute(exp, column_mapping, typer)
             # Check if this is a multi-column regex expansion