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

snowflake/snowpark_connect/expression/map_unresolved_star.py

@@ -34,6 +34,7 @@ def map_unresolved_star(
     column_mapping: ColumnNameMap,
     typer: ExpressionTyper,
 ) -> tuple[list[str], TypedColumn]:
+
     if exp.unresolved_star.HasField("unparsed_target"):
         unparsed_target = exp.unresolved_star.unparsed_target
         name_parts = split_fully_qualified_spark_name(unparsed_target)
@@ -102,7 +103,7 @@ def map_unresolved_star(
                 prefix_candidate_str = f"{prefix_candidate_str}.{name_parts[i]}"
             prefix_candidate = (
                 column_mapping.get_snowpark_column_name_from_spark_column_name(
-                    prefix_candidate_str, allow_non_exists=True
+                    prefix_candidate_str, allow_non_exists=True, is_qualified=(i > 0)
                 )
             )
             if prefix_candidate is None:
@@ -180,7 +181,7 @@ def map_unresolved_star_struct(
             prefix_candidate_str = f"{prefix_candidate_str}.{name_parts[i]}"
         prefix_candidate = (
             column_mapping.get_snowpark_column_name_from_spark_column_name(
-                prefix_candidate_str, allow_non_exists=True
+                prefix_candidate_str, allow_non_exists=True, is_qualified=(i > 0)
             )
         )
         if prefix_candidate is None:

snowflake/snowpark_connect/hidden_column.py

@@ -0,0 +1,39 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+
+class HiddenColumn:
+    """
+    Represents a hidden column in a Snowflake table.
+
+    Hidden columns are not visible in standard queries but can be accessed
+    directly if needed. This class provides a way to reference such columns
+    in Snowpark operations
+    """
+
+    def __init__(
+        self,
+        hidden_snowpark_name: str,
+        spark_name: str,
+        visible_snowpark_name: str,
+        qualifiers: list[str] | None = None,
+        original_position: int | None = None,
+    ) -> None:
+        """
+        Initializes a HiddenColumn instance.
+
+        Args:
+            name (str): The name of the hidden column.
+        """
+
+        # The Snowpark internal name for the hidden column
+        self.hidden_snowpark_name = hidden_snowpark_name
+        # The Spark name for the hidden column
+        self.spark_name = spark_name
+        # The left side visible Snowpark name for the dropped right side column
+        self.visible_snowpark_name = visible_snowpark_name
+        # Qualifiers for the hidden column (e.g., table or schema names)
+        self.qualifiers = qualifiers if qualifiers is not None else []
+        # The position of the hidden column in the original schema
+        self.original_position = original_position

snowflake/snowpark_connect/relation/map_column_ops.py

@@ -6,12 +6,10 @@ 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
 
@@ -46,7 +44,6 @@ 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 (
@@ -346,12 +343,6 @@ 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
@@ -433,29 +424,7 @@ def map_sort(
     if not order_specified:
         ascending = None
 
-    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)
+    result = input_df.sort(cols, ascending=ascending)
 
     return DataFrameContainer(
         result,
@@ -766,10 +735,17 @@ def map_with_columns(
     """
     input_container = map_relation(rel.with_columns.input)
     input_df = input_container.dataframe
-    with_columns = [
-        map_alias(alias, input_container.column_map, ExpressionTyper(input_df))
-        for alias in rel.with_columns.aliases
-    ]
+    with_columns = []
+    for alias in rel.with_columns.aliases:
+        spark_names, typed_alias = map_alias(
+            alias, input_container.column_map, ExpressionTyper(input_df)
+        )
+        register_lca_alias(spark_names[0], typed_alias)
+        with_columns.append((spark_names, typed_alias))
+
+    # we don't need lateral aliases anymore
+    clear_lca_alias_map()
+
     # TODO: This list needs to contain all unique column names, but the code below doesn't
     # guarantee that.
     with_columns_names = []
@@ -837,6 +813,12 @@ def map_with_columns(
         with_columns_names_deduped, with_columns_exprs_deduped
     ).select(*new_snowpark_columns)
 
+    # SNOW-2306644: the next projection after a withColumn call can completely remove the added column
+    # df.withColumn("new").select("foo").filter("new") will fail with a missing column error
+    # the column will be preserved if flattening is disabled
+    if hasattr(result, "_select_statement"):
+        result._select_statement.flatten_disabled = True
+
     snowpark_name_to_type = dict(
         [(f.name, f.datatype) for f in input_df.schema.fields]
         + list(zip(with_columns_names, with_columns_types))

snowflake/snowpark_connect/relation/map_extension.py

@@ -347,11 +347,54 @@ def map_aggregate(
     raw_groupings: list[tuple[str, TypedColumn]] = []
     raw_aggregations: list[tuple[str, TypedColumn]] = []
 
+    if not is_group_by_all:
+        raw_groupings = [_map_column(exp) for exp in aggregate.grouping_expressions]
+
+    # Determine grouping columns for context
+    # For GROUPING SETS, we need to extract the columns from the sets
+    grouping_columns_for_context = []
+    if aggregate.group_type == snowflake_proto.Aggregate.GROUP_TYPE_GROUPING_SETS:
+        # Use a list to preserve order, avoiding duplicates
+        for grouping_set in aggregate.grouping_sets:
+            for exp in grouping_set.grouping_set:
+                spark_names, _ = map_expression(exp, input_container.column_map, typer)
+                # map_expression always returns a list, get the first element
+                col_name = spark_names[0]
+                if col_name not in grouping_columns_for_context:
+                    grouping_columns_for_context.append(col_name)
+    else:
+        grouping_columns_for_context = [spark_name for spark_name, _ in raw_groupings]
+
+    # Set grouping columns context for processing aggregate expressions
+    # This context is needed for resolving grouping__id references
+    # TODO: This should properly handle nested queries with GROUP BY using push/pop
+    # Currently, nested queries may interfere with parent queries
+    set_current_grouping_columns(grouping_columns_for_context)
+
+    # LCA Support for aggregate expressions: Use the LCA alias map
+    # Note: We don't clear the map here to preserve any parent context aliases
+    from snowflake.snowpark_connect.utils.context import register_lca_alias
+
     agg_count = get_sql_aggregate_function_count()
     for exp in aggregate.aggregate_expressions:
         col = _map_column(exp)
         raw_aggregations.append(col)
 
+        # If this is an alias, register it in the LCA map for subsequent expressions
+        if (
+            exp.WhichOneof("expr_type") == "alias"
+            and exp.alias.name
+            and len(exp.alias.name) > 0
+        ):
+            alias_name = exp.alias.name[0]
+            spark_name, snowpark_column = col
+
+            # Register the alias pointing to the result of its expression
+            # This handles both simple aliases (k as lca) and complex ones (lca + 1 as col)
+            # The snowpark_column already contains the computed expression with its alias wrapper,
+            # which is fine - when referenced later, the column's value is what gets used
+            register_lca_alias(alias_name, snowpark_column)
+
         if is_group_by_all:
             new_agg_count = get_sql_aggregate_function_count()
             if new_agg_count == agg_count:
@@ -359,13 +402,6 @@ def map_aggregate(
             else:
                 agg_count = new_agg_count
 
-    if not is_group_by_all:
-        raw_groupings = [_map_column(exp) for exp in aggregate.grouping_expressions]
-
-    # Set the current grouping columns in context for grouping_id() function
-    grouping_spark_columns = [spark_name for spark_name, _ in raw_groupings]
-    set_current_grouping_columns(grouping_spark_columns)
-
     # Now create column name lists and assign aliases.
     # In case of GROUP BY ALL, even though groupings are a subset of aggregations,
     # they will have their own aliases so we can drop them later.
@@ -404,15 +440,20 @@ def map_aggregate(
         case snowflake_proto.Aggregate.GROUP_TYPE_CUBE:
             result = input_df.cube(groupings)
         case snowflake_proto.Aggregate.GROUP_TYPE_GROUPING_SETS:
-            # TODO: What do we do about groupings?
-            sets = (
-                [
-                    map_expression(exp, input_container.column_map, typer)[1].col
-                    for exp in grouping_sets.grouping_set
-                ]
-                for grouping_sets in aggregate.grouping_sets
+            # Map each grouping set to columns
+            sets_mapped = []
+            for grouping_set in aggregate.grouping_sets:
+                set_cols = []
+                for exp in grouping_set.grouping_set:
+                    _, typed_col = map_expression(
+                        exp, input_container.column_map, typer
+                    )
+                    set_cols.append(typed_col.col)
+                sets_mapped.append(set_cols)
+
+            result = input_df.group_by_grouping_sets(
+                snowpark.GroupingSets(*sets_mapped)
             )
-            result = input_df.group_by_grouping_sets(snowpark.GroupingSets(*sets))
         case other:
             raise SnowparkConnectNotImplementedError(
                 f"Unsupported GROUP BY type: {other}"