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

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: ...

snowflake/snowpark_connect/relation/catalogs/snowflake_catalog.py

@@ -8,7 +8,10 @@ import typing
 import pandas
 import pyspark.sql.connect.proto.common_pb2 as common_proto
 import pyspark.sql.connect.proto.types_pb2 as types_proto
-from snowflake.core.exceptions import NotFoundError
+from pyspark.sql.connect.client.core import Retrying
+from snowflake.core.exceptions import APIError, NotFoundError
+from snowflake.core.schema import Schema
+from snowflake.core.table import Table, TableColumn
 
 from snowflake.snowpark import functions
 from snowflake.snowpark._internal.analyzer.analyzer_utils import (
@@ -22,6 +25,7 @@ from snowflake.snowpark_connect.config import (
     global_config,
 )
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
+from snowflake.snowpark_connect.error.exceptions import MaxRetryExceeded
 from snowflake.snowpark_connect.relation.catalogs.abstract_spark_catalog import (
     AbstractSparkCatalog,
     _get_current_snowflake_schema,
@@ -39,6 +43,37 @@ from snowflake.snowpark_connect.utils.telemetry import (
 from snowflake.snowpark_connect.utils.udf_cache import cached_udf
 
 
+def _is_retryable_api_error(e: Exception) -> bool:
+    """
+    Determine if an APIError should be retried.
+
+    Only retry on server errors, rate limiting, and transient network issues.
+    Don't retry on client errors like authentication, authorization, or validation failures.
+    """
+    if not isinstance(e, APIError):
+        return False
+
+    # Check if the error has a status_code attribute
+    if hasattr(e, "status_code"):
+        # Retry on server errors (5xx), rate limiting (429), and some client errors (400)
+        # 400 can be transient in some cases (like the original error trace shows)
+        return e.status_code in [400, 429, 500, 502, 503, 504]
+
+    # For APIErrors without explicit status codes, check the message
+    error_msg = str(e).lower()
+    retryable_patterns = [
+        "timeout",
+        "connection",
+        "network",
+        "unavailable",
+        "temporary",
+        "rate limit",
+        "throttle",
+    ]
+
+    return any(pattern in error_msg for pattern in retryable_patterns)
+
+
 def _normalize_identifier(identifier: str | None) -> str | None:
     if identifier is None:
         return None
@@ -73,10 +108,25 @@ class SnowflakeCatalog(AbstractSparkCatalog):
             )
         sp_catalog = get_or_create_snowpark_session().catalog
 
-        dbs = sp_catalog.list_schemas(
-            database=sf_quote(sf_database),
-            pattern=_normalize_identifier(sf_schema),
-        )
+        dbs: list[Schema] | None = None
+        for attempt in Retrying(
+            max_retries=5,
+            initial_backoff=100,  # 100ms
+            max_backoff=5000,  # 5 s
+            backoff_multiplier=2.0,
+            jitter=100,
+            min_jitter_threshold=200,
+            can_retry=_is_retryable_api_error,
+        ):
+            with attempt:
+                dbs = sp_catalog.list_schemas(
+                    database=sf_quote(sf_database),
+                    pattern=_normalize_identifier(sf_schema),
+                )
+        if dbs is None:
+            raise MaxRetryExceeded(
+                f"Failed to fetch databases {f'with pattern {pattern} ' if pattern is not None else ''}after all retry attempts"
+            )
         names: list[str] = list()
         catalogs: list[str] = list()
         descriptions: list[str | None] = list()
@@ -112,9 +162,24 @@ class SnowflakeCatalog(AbstractSparkCatalog):
             )
         sp_catalog = get_or_create_snowpark_session().catalog
 
-        db = sp_catalog.get_schema(
-            schema=sf_quote(sf_schema), database=sf_quote(sf_database)
-        )
+        db: Schema | None = None
+        for attempt in Retrying(
+            max_retries=5,
+            initial_backoff=100,  # 100ms
+            max_backoff=5000,  # 5 s
+            backoff_multiplier=2.0,
+            jitter=100,
+            min_jitter_threshold=200,
+            can_retry=_is_retryable_api_error,
+        ):
+            with attempt:
+                db = sp_catalog.get_schema(
+                    schema=sf_quote(sf_schema), database=sf_quote(sf_database)
+                )
+        if db is None:
+            raise MaxRetryExceeded(
+                f"Failed to fetch database {spark_dbName} after all retry attempts"
+            )
 
         name = unquote_if_quoted(db.name)
         return pandas.DataFrame(
@@ -241,11 +306,27 @@ class SnowflakeCatalog(AbstractSparkCatalog):
                 "Calling into another catalog is not currently supported"
             )
 
-        table = sp_catalog.get_table(
-            database=sf_quote(sf_database),
-            schema=sf_quote(sf_schema),
-            table_name=sf_quote(table_name),
-        )
+        table: Table | None = None
+        for attempt in Retrying(
+            max_retries=5,
+            initial_backoff=100,  # 100ms
+            max_backoff=5000,  # 5 s
+            backoff_multiplier=2.0,
+            jitter=100,
+            min_jitter_threshold=200,
+            can_retry=_is_retryable_api_error,
+        ):
+            with attempt:
+                table = sp_catalog.get_table(
+                    database=sf_quote(sf_database),
+                    schema=sf_quote(sf_schema),
+                    table_name=sf_quote(table_name),
+                )
+
+        if table is None:
+            raise MaxRetryExceeded(
+                f"Failed to fetch table {spark_tableName} after all retry attempts"
+            )
 
         return pandas.DataFrame(
             {
@@ -286,6 +367,7 @@ class SnowflakeCatalog(AbstractSparkCatalog):
     ) -> pandas.DataFrame:
         """List all columns in a table/view, optionally database name filter can be provided."""
         sp_catalog = get_or_create_snowpark_session().catalog
+        columns: list[TableColumn] | None = None
         if spark_dbName is None:
             catalog, sf_database, sf_schema, sf_table = _process_multi_layer_identifier(
                 spark_tableName
@@ -294,15 +376,39 @@ class SnowflakeCatalog(AbstractSparkCatalog):
                 raise SnowparkConnectNotImplementedError(
                     "Calling into another catalog is not currently supported"
                 )
-            columns = sp_catalog.list_columns(
-                database=sf_quote(sf_database),
-                schema=sf_quote(sf_schema),
-                table_name=sf_quote(sf_table),
-            )
+            for attempt in Retrying(
+                max_retries=5,
+                initial_backoff=100,  # 100ms
+                max_backoff=5000,  # 5 s
+                backoff_multiplier=2.0,
+                jitter=100,
+                min_jitter_threshold=200,
+                can_retry=_is_retryable_api_error,
+            ):
+                with attempt:
+                    columns = sp_catalog.list_columns(
+                        database=sf_quote(sf_database),
+                        schema=sf_quote(sf_schema),
+                        table_name=sf_quote(sf_table),
+                    )
         else:
-            columns = sp_catalog.list_columns(
-                schema=sf_quote(spark_dbName),
-                table_name=sf_quote(spark_tableName),
+            for attempt in Retrying(
+                max_retries=5,
+                initial_backoff=100,  # 100ms
+                max_backoff=5000,  # 5 s
+                backoff_multiplier=2.0,
+                jitter=100,
+                min_jitter_threshold=200,
+                can_retry=_is_retryable_api_error,
+            ):
+                with attempt:
+                    columns = sp_catalog.list_columns(
+                        schema=sf_quote(spark_dbName),
+                        table_name=sf_quote(spark_tableName),
+                    )
+        if columns is None:
+            raise MaxRetryExceeded(
+                f"Failed to fetch columns of {spark_tableName} after all retry attempts"
             )
         names: list[str] = list()
         descriptions: list[str | None] = list()

snowflake/snowpark_connect/relation/map_aggregate.py

@@ -4,10 +4,14 @@
 
 import re
 from dataclasses import dataclass
+from typing import Optional
 
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
+import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
+from snowflake.snowpark import Column
+from snowflake.snowpark._internal.analyzer.unary_expression import Alias
 from snowflake.snowpark.types import DataType
 from snowflake.snowpark_connect.column_name_handler import (
     make_column_names_snowpark_compatible,
@@ -21,6 +25,7 @@ from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 from snowflake.snowpark_connect.typed_column import TypedColumn
 from snowflake.snowpark_connect.utils.context import (
+    get_is_evaluating_sql,
     set_current_grouping_columns,
     temporary_pivot_expression,
 )
@@ -131,19 +136,109 @@ def map_pivot_aggregate(
         get_literal_field_and_name(lit)[0] for lit in rel.aggregate.pivot.values
     ]
 
+    used_columns = {pivot_column[1].col._expression.name}
+    if get_is_evaluating_sql():
+        # When evaluating SQL spark doesn't trim columns from the result
+        used_columns = {"*"}
+    else:
+        for expression in rel.aggregate.aggregate_expressions:
+            matched_identifiers = re.findall(
+                r'unparsed_identifier: "(.*)"', expression.__str__()
+            )
+            for identifier in matched_identifiers:
+                mapped_col = input_container.column_map.spark_to_col.get(
+                    identifier, None
+                )
+                if mapped_col:
+                    used_columns.add(mapped_col[0].snowpark_name)
+
     if len(columns.grouping_expressions()) == 0:
-        result = input_df_actual.pivot(
-            pivot_column[1].col, pivot_values if pivot_values else None
-        ).agg(*columns.aggregation_expressions())
+        # Snowpark doesn't support multiple aggregations in pivot without groupBy
+        # So we need to perform each aggregation separately and then combine results
+        if len(columns.aggregation_expressions(unalias=True)) > 1:
+            agg_expressions = columns.aggregation_expressions(unalias=True)
+            agg_metadata = columns.aggregation_columns
+            num_agg_functions = len(agg_expressions)
+
+            spark_names = []
+            pivot_results = []
+            for i, agg_expr in enumerate(agg_expressions):
+                pivot_result = (
+                    input_df_actual.select(*used_columns)
+                    .pivot(pivot_column[1].col, pivot_values if pivot_values else None)
+                    .agg(agg_expr)
+                )
+                for col_name in pivot_result.columns:
+                    spark_names.append(
+                        f"{pivot_column_name(col_name)}_{agg_metadata[i].spark_name}"
+                    )
+                pivot_results.append(pivot_result)
+
+            result = pivot_results[0]
+            for pivot_result in pivot_results[1:]:
+                result = result.cross_join(pivot_result)
+
+            pivot_columns_per_agg = len(pivot_results[0].columns)
+            reordered_spark_names = []
+            reordered_snowpark_names = []
+            reordered_types = []
+            column_selectors = []
+
+            for pivot_idx in range(pivot_columns_per_agg):
+                for agg_idx in range(num_agg_functions):
+                    current_pos = agg_idx * pivot_columns_per_agg + pivot_idx
+                    if current_pos < len(spark_names):
+                        idx = current_pos + 1  # 1-based indexing for Snowpark
+                        reordered_spark_names.append(spark_names[current_pos])
+                        reordered_snowpark_names.append(f"${idx}")
+                        reordered_types.append(
+                            result.schema.fields[current_pos].datatype
+                        )
+                        column_selectors.append(snowpark_fn.col(f"${idx}"))
+
+            return DataFrameContainer.create_with_column_mapping(
+                dataframe=result.select(*column_selectors),
+                spark_column_names=reordered_spark_names,
+                snowpark_column_names=reordered_snowpark_names,
+                column_qualifiers=[[]] * len(reordered_spark_names),
+                parent_column_name_map=input_container.column_map,
+                snowpark_column_types=reordered_types,
+            )
+        else:
+            result = (
+                input_df_actual.select(*used_columns)
+                .pivot(pivot_column[1].col, pivot_values if pivot_values else None)
+                .agg(*columns.aggregation_expressions(unalias=True))
+            )
     else:
         result = (
             input_df_actual.group_by(*columns.grouping_expressions())
             .pivot(pivot_column[1].col, pivot_values if pivot_values else None)
-            .agg(*columns.aggregation_expressions())
+            .agg(*columns.aggregation_expressions(unalias=True))
         )
 
+    agg_name_list = [c.spark_name for c in columns.grouping_columns]
+
+    # Calculate number of pivot values for proper Spark-compatible indexing
+    total_pivot_columns = len(result.columns) - len(agg_name_list)
+    num_pivot_values = (
+        total_pivot_columns // len(columns.aggregation_columns)
+        if len(columns.aggregation_columns) > 0
+        else 1
+    )
+
+    def _get_agg_exp_alias_for_col(col_index: int) -> Optional[str]:
+        if col_index < len(agg_name_list) or len(columns.aggregation_columns) <= 1:
+            return None
+        else:
+            index = (col_index - len(agg_name_list)) // num_pivot_values
+            return columns.aggregation_columns[index].spark_name
+
     spark_columns = []
-    for col in [string_parser(s) for s in result.columns]:
+    for col in [
+        pivot_column_name(c, _get_agg_exp_alias_for_col(i))
+        for i, c in enumerate(result.columns)
+    ]:
         spark_col = (
             input_container.column_map.get_spark_column_name_from_snowpark_column_name(
                 col, allow_non_exists=True
@@ -153,22 +248,57 @@ def map_pivot_aggregate(
         if spark_col is not None:
             spark_columns.append(spark_col)
         else:
-            spark_columns.append(col)
+            # Handle NULL column names to match Spark behavior (lowercase 'null')
+            if col == "NULL":
+                spark_columns.append(col.lower())
+            else:
+                spark_columns.append(col)
+
+    grouping_cols_count = len(agg_name_list)
+    pivot_cols = result.columns[grouping_cols_count:]
+    spark_pivot_cols = spark_columns[grouping_cols_count:]
+
+    num_agg_functions = len(columns.aggregation_columns)
+    num_pivot_values = len(pivot_cols) // num_agg_functions
+
+    reordered_snowpark_cols = []
+    reordered_spark_cols = []
+    column_indices = []  # 1-based indexing
+
+    for i in range(grouping_cols_count):
+        reordered_snowpark_cols.append(result.columns[i])
+        reordered_spark_cols.append(spark_columns[i])
+        column_indices.append(i + 1)
+
+    for pivot_idx in range(num_pivot_values):
+        for agg_idx in range(num_agg_functions):
+            current_pos = agg_idx * num_pivot_values + pivot_idx
+            if current_pos < len(pivot_cols):
+                reordered_snowpark_cols.append(pivot_cols[current_pos])
+                reordered_spark_cols.append(spark_pivot_cols[current_pos])
+                original_index = grouping_cols_count + current_pos
+                column_indices.append(original_index + 1)
+
+    reordered_result = result.select(
+        *[snowpark_fn.col(f"${idx}") for idx in column_indices]
+    )
 
-    agg_name_list = [c.spark_name for c in columns.grouping_columns]
     return DataFrameContainer.create_with_column_mapping(
-        dataframe=result,
-        spark_column_names=agg_name_list + spark_columns[len(agg_name_list) :],
-        snowpark_column_names=result.columns,
+        dataframe=reordered_result,
+        spark_column_names=reordered_spark_cols,
+        snowpark_column_names=[f"${idx}" for idx in column_indices],
         column_qualifiers=(
             columns.get_qualifiers()[: len(agg_name_list)]
-            + [[]] * (len(spark_columns) - len(agg_name_list))
+            + [[]] * (len(reordered_spark_cols) - len(agg_name_list))
         ),
         parent_column_name_map=input_container.column_map,
+        snowpark_column_types=[
+            result.schema.fields[idx - 1].datatype for idx in column_indices
+        ],
     )
 
 
-def string_parser(s):
+def pivot_column_name(snowpark_cname, opt_alias: Optional[str] = None) -> Optional[str]:
     # For values that are used as pivoted columns, the input and output are in the following format (outermost double quotes are part of the input):
 
     # 1. "'Java'" -> Java
@@ -183,7 +313,7 @@ def string_parser(s):
 
     try:
         # handling values that are used as pivoted columns
-        match = re.match(r'^"\'(.*)\'"$', s)
+        match = re.match(r'^"\'(.*)\'"$', snowpark_cname)
         # extract the content between the outermost double quote followed by a single quote "'
         content = match.group(1)
         # convert the escaped double quote to the actual double quote
@@ -195,10 +325,10 @@ def string_parser(s):
         content = re.sub(r"'", "", content)
         # replace the placeholder with the single quote which we want to preserve
         result = content.replace(escape_single_quote_placeholder, "'")
-        return result
+        return f"{result}_{opt_alias}" if opt_alias else result
     except Exception:
         # fallback to the original logic, handling aliased column names
-        double_quote_list = re.findall(r'"(.*?)"', s)
+        double_quote_list = re.findall(r'"(.*?)"', snowpark_cname)
         spark_string = ""
         for entry in list(filter(None, double_quote_list)):
             if "'" in entry:
@@ -210,7 +340,7 @@ def string_parser(s):
                 spark_string += entry
             else:
                 spark_string += '"' + entry + '"'
-        return s if spark_string == "" else spark_string
+        return snowpark_cname if spark_string == "" else spark_string
 
 
 @dataclass(frozen=True)
@@ -231,8 +361,14 @@ class _Columns:
     def grouping_expressions(self) -> list[snowpark.Column]:
         return [col.expression for col in self.grouping_columns]
 
-    def aggregation_expressions(self) -> list[snowpark.Column]:
-        return [col.expression for col in self.aggregation_columns]
+    def aggregation_expressions(self, unalias: bool = False) -> list[snowpark.Column]:
+        def _unalias(col: snowpark.Column) -> snowpark.Column:
+            if unalias and hasattr(col, "_expr1") and isinstance(col._expr1, Alias):
+                return _unalias(Column(col._expr1.child))
+            else:
+                return col
+
+        return [_unalias(col.expression) for col in self.aggregation_columns]
 
     def expressions(self) -> list[snowpark.Column]:
         return self.grouping_expressions() + self.aggregation_expressions()

snowflake/snowpark_connect/relation/map_column_ops.py

@@ -6,10 +6,12 @@ import ast
 import json
 import sys
 from collections import defaultdict
+from copy import copy
 
 import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 import pyspark.sql.connect.proto.types_pb2 as types_proto
+from pyspark.errors import PySparkValueError
 from pyspark.errors.exceptions.base import AnalysisException
 from pyspark.serializers import CloudPickleSerializer
 
@@ -44,6 +46,7 @@ from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 from snowflake.snowpark_connect.relation.utils import (
     TYPE_MAP_FOR_TO_SCHEMA,
+    can_sort_be_flattened,
     snowpark_functions_col,
 )
 from snowflake.snowpark_connect.type_mapping import (
@@ -266,6 +269,7 @@ def map_project(
 
             aliased_col = mapper.col.alias(snowpark_column)
             select_list.append(aliased_col)
+
             new_snowpark_columns.append(snowpark_column)
             new_spark_columns.append(spark_name)
             column_types.extend(mapper.types)
@@ -342,6 +346,12 @@ def map_sort(
 
     sort_order = sort.order
 
+    if not sort_order:
+        raise PySparkValueError(
+            error_class="CANNOT_BE_EMPTY",
+            message="At least one column must be specified.",
+        )
+
     if len(sort_order) == 1:
         parsed_col_name = split_fully_qualified_spark_name(
             sort_order[0].child.unresolved_attribute.unparsed_identifier
@@ -422,7 +432,30 @@ def map_sort(
     # TODO: sort.isglobal.
     if not order_specified:
         ascending = None
-    result = input_df.sort(cols, ascending=ascending)
+
+    select_statement = getattr(input_df, "_select_statement", None)
+    sort_expressions = [c._expression for c in cols]
+    if (
+        can_sort_be_flattened(select_statement, *sort_expressions)
+        and input_df._ops_after_agg is None
+    ):
+        # "flattened" order by that will allow using dropped columns
+        new = copy(select_statement)
+        new.from_ = select_statement.from_.to_subqueryable()
+        new.pre_actions = new.from_.pre_actions
+        new.post_actions = new.from_.post_actions
+        new.order_by = sort_expressions + (select_statement.order_by or [])
+        new.column_states = select_statement.column_states
+        new._merge_projection_complexity_with_subquery = False
+        new.df_ast_ids = (
+            select_statement.df_ast_ids.copy()
+            if select_statement.df_ast_ids is not None
+            else None
+        )
+        new.attributes = select_statement.attributes
+        result = input_df._with_plan(new)
+    else:
+        result = input_df.sort(cols, ascending=ascending)
 
     return DataFrameContainer(
         result,
@@ -666,10 +699,29 @@ def map_with_columns_renamed(
         )
 
     # Validate for naming conflicts
-    new_names_list = list(dict(rel.with_columns_renamed.rename_columns_map).values())
+    rename_map = dict(rel.with_columns_renamed.rename_columns_map)
+    new_names_list = list(rename_map.values())
     seen = set()
     for new_name in new_names_list:
-        if column_map.has_spark_column(new_name):
+        # Check if this new name conflicts with existing columns
+        # But allow renaming a column to a different case version of itself
+        is_case_insensitive_self_rename = False
+        if not global_config.spark_sql_caseSensitive:
+            # Find the source column(s) that map to this new name
+            source_columns = [
+                old_name
+                for old_name, new_name_candidate in rename_map.items()
+                if new_name_candidate == new_name
+            ]
+            # Check if any source column is the same as new name when case-insensitive
+            is_case_insensitive_self_rename = any(
+                source_col.lower() == new_name.lower() for source_col in source_columns
+            )
+
+        if (
+            column_map.has_spark_column(new_name)
+            and not is_case_insensitive_self_rename
+        ):
             # Spark doesn't allow reusing existing names, even if the result df will not contain duplicate columns
             raise _column_exists_error(new_name)
         if (global_config.spark_sql_caseSensitive and new_name in seen) or (
@@ -1056,14 +1108,12 @@ def map_group_map(
     snowpark_grouping_expressions: list[snowpark.Column] = []
     typer = ExpressionTyper(input_df)
     group_name_list: list[str] = []
-    qualifiers = []
     for exp in grouping_expressions:
         new_name, snowpark_column = map_single_column_expression(
             exp, input_container.column_map, typer
         )
         snowpark_grouping_expressions.append(snowpark_column.col)
         group_name_list.append(new_name)
-        qualifiers.append(snowpark_column.get_qualifiers())
     if rel.group_map.func.python_udf is None:
         raise ValueError("group_map relation without python udf is not supported")
 
@@ -1105,13 +1155,14 @@ def map_group_map(
         result = input_df.group_by(*snowpark_grouping_expressions).apply_in_pandas(
             callable_func, output_type
         )
-
-    qualifiers.extend([[]] * (len(result.columns) - len(group_name_list)))
+    # The UDTF `apply_in_pandas` generates a new table whose output schema
+    # can be entirely different from that of the input Snowpark DataFrame.
+    # As a result, the output DataFrame should not use qualifiers based on the input group by columns.
     return DataFrameContainer.create_with_column_mapping(
         dataframe=result,
         spark_column_names=[field.name for field in output_type],
         snowpark_column_names=result.columns,
-        column_qualifiers=qualifiers,
+        column_qualifiers=None,
         parent_column_name_map=input_container.column_map,
     )