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

snowflake/snowpark_connect/relation/map_aggregate.py

@@ -4,15 +4,19 @@
 
 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,
-    with_column_map,
 )
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.expression.literal import get_literal_field_and_name
 from snowflake.snowpark_connect.expression.map_expression import (
     map_single_column_expression,
@@ -20,111 +24,171 @@ from snowflake.snowpark_connect.expression.map_expression import (
 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 temporary_pivot_expression
+from snowflake.snowpark_connect.utils.context import (
+    get_is_evaluating_sql,
+    set_current_grouping_columns,
+    temporary_pivot_expression,
+)
 
 
-def map_group_by_aggregate(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_group_by_aggregate(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Groups the DataFrame using the specified columns.
 
     Aggregations come in as expressions, which are mapped to `snowpark.Column`
     objects.
     """
-    input_df, columns = map_aggregate_helper(rel)
+    input_df_container, columns = map_aggregate_helper(rel)
+    input_df_actual = input_df_container.dataframe
+
     if len(columns.grouping_expressions()) == 0:
-        result = input_df.agg(*columns.aggregation_expressions())
+        result = input_df_actual.agg(*columns.aggregation_expressions())
     else:
-        result = input_df.group_by(*columns.grouping_expressions()).agg(
+        result = input_df_actual.group_by(*columns.grouping_expressions()).agg(
             *columns.aggregation_expressions()
         )
-    return with_column_map(
-        result,
-        columns.spark_names(),
-        columns.snowpark_names(),
-        columns.data_types(),
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=columns.spark_names(),
+        snowpark_column_names=columns.snowpark_names(),
+        snowpark_column_types=columns.data_types(),
         column_qualifiers=columns.get_qualifiers(),
-        parent_column_name_map=input_df._column_map,
+        parent_column_name_map=input_df_container.column_map,
     )
 
 
-def map_rollup_aggregate(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_rollup_aggregate(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Create a multidimensional rollup for the current DataFrame using the specified columns.
 
     Aggregations come in as expressions, which are mapped to `snowpark.Column`
     objects.
     """
-    input_df, columns = map_aggregate_helper(rel)
+    input_container, columns = map_aggregate_helper(rel)
+    input_df_actual = input_container.dataframe
+
     if len(columns.grouping_expressions()) == 0:
-        result = input_df.agg(*columns.aggregation_expressions())
+        result = input_df_actual.agg(*columns.aggregation_expressions())
     else:
-        result = input_df.rollup(*columns.grouping_expressions()).agg(
+        result = input_df_actual.rollup(*columns.grouping_expressions()).agg(
             *columns.aggregation_expressions()
         )
-    return with_column_map(
-        result,
-        columns.spark_names(),
-        columns.snowpark_names(),
-        columns.data_types(),
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=columns.spark_names(),
+        snowpark_column_names=columns.snowpark_names(),
+        snowpark_column_types=columns.data_types(),
         column_qualifiers=columns.get_qualifiers(),
-        parent_column_name_map=input_df._column_map,
+        parent_column_name_map=input_container.column_map,
     )
 
 
-def map_cube_aggregate(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_cube_aggregate(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Create a multidimensional cube for the current DataFrame using the specified columns.
 
     Aggregations come in as expressions, which are mapped to `snowpark.Column`
     objects.
     """
-    input_df, columns = map_aggregate_helper(rel)
+    input_container, columns = map_aggregate_helper(rel)
+    input_df_actual = input_container.dataframe
+
     if len(columns.grouping_expressions()) == 0:
-        result = input_df.agg(*columns.aggregation_expressions())
+        result = input_df_actual.agg(*columns.aggregation_expressions())
     else:
-        result = input_df.cube(*columns.grouping_expressions()).agg(
+        result = input_df_actual.cube(*columns.grouping_expressions()).agg(
             *columns.aggregation_expressions()
         )
-    return with_column_map(
-        result,
-        columns.spark_names(),
-        columns.snowpark_names(),
-        columns.data_types(),
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=columns.spark_names(),
+        snowpark_column_names=columns.snowpark_names(),
+        snowpark_column_types=columns.data_types(),
         column_qualifiers=columns.get_qualifiers(),
-        parent_column_name_map=input_df._column_map,
+        parent_column_name_map=input_container.column_map,
     )
 
 
-def map_pivot_aggregate(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_pivot_aggregate(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Pivots a column of the current DataFrame and performs the specified aggregation.
 
     There are 2 versions of the pivot function: one that requires the caller to specify the list of the distinct values
     to pivot on and one that does not.
     """
-    input_df, columns = map_aggregate_helper(rel, pivot=True, skip_alias=True)
+    input_container, columns = map_aggregate_helper(rel, pivot=True, skip_alias=True)
+    input_df_actual = input_container.dataframe
+
     pivot_column = map_single_column_expression(
-        rel.aggregate.pivot.col, input_df._column_map, ExpressionTyper(input_df)
+        rel.aggregate.pivot.col,
+        input_container.column_map,
+        ExpressionTyper(input_df_actual),
     )
     pivot_values = [
         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.pivot(
-            pivot_column[1].col, pivot_values if pivot_values else None
-        ).agg(*columns.aggregation_expressions())
+        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.group_by(*columns.grouping_expressions())
+            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_df._column_map.get_spark_column_name_from_snowpark_column_name(
+            input_container.column_map.get_spark_column_name_from_snowpark_column_name(
                 col, allow_non_exists=True
             )
         )
@@ -132,22 +196,57 @@ def map_pivot_aggregate(rel: relation_proto.Relation) -> snowpark.DataFrame:
         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 with_column_map(
-        result,
-        agg_name_list + spark_columns[len(agg_name_list) :],
-        result.columns,
+    return DataFrameContainer.create_with_column_mapping(
+        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_df._column_map,
+        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
@@ -162,7 +261,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
@@ -174,10 +273,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:
@@ -189,7 +288,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)
@@ -210,8 +309,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()
@@ -246,7 +351,8 @@ class _Columns:
 def map_aggregate_helper(
     rel: relation_proto.Relation, pivot: bool = False, skip_alias: bool = False
 ):
-    input_df = map_relation(rel.aggregate.input)
+    input_container = map_relation(rel.aggregate.input)
+    input_df = input_container.dataframe
     grouping_expressions = rel.aggregate.grouping_expressions
     expressions = rel.aggregate.aggregate_expressions
     groupings: list[_ColumnMetadata] = []
@@ -258,7 +364,7 @@ def map_aggregate_helper(
     with temporary_pivot_expression(pivot):
         for exp in grouping_expressions:
             new_name, snowpark_column = map_single_column_expression(
-                exp, input_df._column_map, typer
+                exp, input_container.column_map, typer
             )
             alias = make_column_names_snowpark_compatible(
                 [new_name], rel.common.plan_id, len(groupings)
@@ -275,9 +381,12 @@ def map_aggregate_helper(
                 )
             )
 
+        grouping_cols = [g.spark_name for g in groupings]
+        set_current_grouping_columns(grouping_cols)
+
         for exp in expressions:
             new_name, snowpark_column = map_single_column_expression(
-                exp, input_df._column_map, typer
+                exp, input_container.column_map, typer
             )
             alias = make_column_names_snowpark_compatible(
                 [new_name], rel.common.plan_id, len(groupings) + len(aggregations)
@@ -313,7 +422,7 @@ def map_aggregate_helper(
             )
 
         return (
-            input_df,
+            input_container,
             _Columns(
                 grouping_columns=groupings,
                 aggregation_columns=aggregations,

snowflake/snowpark_connect/relation/map_catalog.py

@@ -7,7 +7,7 @@ import re
 import pandas
 import pyspark.sql.connect.proto.catalog_pb2 as catalog_proto
 
-from snowflake import snowpark
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.catalogs import CATALOGS
 from snowflake.snowpark_connect.relation.catalogs.utils import (
     CURRENT_CATALOG_NAME,
@@ -22,7 +22,7 @@ from snowflake.snowpark_connect.utils.telemetry import (
 
 def map_catalog(
     rel: catalog_proto.Catalog,
-) -> pandas.DataFrame | snowpark.DataFrame:
+) -> DataFrameContainer | pandas.DataFrame:
     match rel.WhichOneof("cat_type"):
         # Database related APIs
         case "current_database":