snowpark-connect 0.31.0__py3-none-any.whl → 0.33.0__py3-none-any.whl

snowflake/snowpark_connect/expression/map_sql_expression.py

@@ -260,12 +260,47 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
     class_name = str(exp.getClass().getSimpleName())
     match class_name:
         case "AggregateExpression":
-            func_name = as_java_list(exp.children())[0].nodeName()
+            aggregate_func = as_java_list(exp.children())[0]
+            func_name = aggregate_func.nodeName()
             args = [
                 map_logical_plan_expression(e)
-                for e in list(as_java_list(as_java_list(exp.children())[0].children()))
+                for e in list(as_java_list(aggregate_func.children()))
             ]
-            proto = apply_filter_clause(func_name, args, exp)
+
+            # Special handling for percentile_cont and percentile_disc
+            # These functions have a 'reverse' property that indicates sort order
+            # Pass it as a 3rd argument (sort_order expression) without modifying children
+            if func_name.lower() in ("percentile_cont", "percentiledisc"):
+                # percentile_cont/disc should always have exactly 2 children: unresolved attribute and percentile value
+                if len(args) != 2:
+                    exception = AssertionError(
+                        f"{func_name} expected 2 args but got {len(args)}"
+                    )
+                    attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+                    raise exception
+
+                reverse = bool(aggregate_func.reverse())
+
+                direction = (
+                    expressions_proto.Expression.SortOrder.SORT_DIRECTION_DESCENDING
+                    if reverse
+                    else expressions_proto.Expression.SortOrder.SORT_DIRECTION_ASCENDING
+                )
+
+                sort_order_expr = expressions_proto.Expression(
+                    sort_order=expressions_proto.Expression.SortOrder(
+                        child=args[0],
+                        direction=direction,
+                    )
+                )
+                args.append(sort_order_expr)
+                proto = apply_filter_clause(func_name, [args[0]], exp)
+                # second arg is a literal value and it doesn't make sense to apply filter on it.
+                # also skips filtering on sort_order.
+                proto.unresolved_function.arguments.append(args[1])
+                proto.unresolved_function.arguments.append(sort_order_expr)
+            else:
+                proto = apply_filter_clause(func_name, args, exp)
         case "Alias":
             proto = expressions_proto.Expression(
                 alias=expressions_proto.Expression.Alias(
@@ -383,13 +418,21 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
                 )
             )
         case "Like" | "ILike" | "RLike":
+            arguments = [
+                map_logical_plan_expression(e)
+                for e in list(as_java_list(exp.children()))
+            ]
+            # exp.escapeChar() returns a JPype JChar - convert to string and create a literal
+            if getattr(exp, "escapeChar", None) is not None:
+                escape_char_str = str(exp.escapeChar())
+                escape_literal = expressions_proto.Expression(
+                    literal=expressions_proto.Expression.Literal(string=escape_char_str)
+                )
+                arguments.append(escape_literal)
             proto = expressions_proto.Expression(
                 unresolved_function=expressions_proto.Expression.UnresolvedFunction(
                     function_name=class_name.lower(),
-                    arguments=[
-                        map_logical_plan_expression(e)
-                        for e in list(as_java_list(exp.children()))
-                    ],
+                    arguments=arguments,
                 )
             )
         case "LikeAny" | "NotLikeAny" | "LikeAll" | "NotLikeAll":

snowflake/snowpark_connect/expression/map_unresolved_attribute.py

@@ -3,6 +3,7 @@
 #
 
 import re
+from typing import Any
 
 import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 from pyspark.errors.exceptions.connect import AnalysisException
@@ -247,7 +248,7 @@ def map_unresolved_attribute(
                 )
             )
             col = get_col(snowpark_name)
-            qualifiers = column_mapping.get_qualifier_for_spark_column(quoted_col_name)
+            qualifiers = column_mapping.get_qualifiers_for_spark_column(quoted_col_name)
             typed_col = TypedColumn(col, lambda: typer.type(col))
             typed_col.set_qualifiers(qualifiers)
             # Store matched columns info for later use
@@ -262,7 +263,7 @@ def map_unresolved_attribute(
                 )
             )
             col = get_col(snowpark_name)
-            qualifiers = column_mapping.get_qualifier_for_spark_column(quoted_col_name)
+            qualifiers = column_mapping.get_qualifiers_for_spark_column(quoted_col_name)
             typed_col = TypedColumn(col, lambda: typer.type(col))
             typed_col.set_qualifiers(qualifiers)
             return (matched_columns[0], typed_col)
@@ -275,12 +276,33 @@ def map_unresolved_attribute(
     else:
         quoted_attr_name = name_parts[0]
 
-    snowpark_name = column_mapping.get_snowpark_column_name_from_spark_column_name(
-        quoted_attr_name, allow_non_exists=True
-    )
+    # Helper function to try finding a column in current and outer scopes
+    def try_resolve_column(column_name: str) -> tuple[str | None, Any]:
+        # Try current scope
+        snowpark_name = column_mapping.get_snowpark_column_name_from_spark_column_name(
+            column_name, allow_non_exists=True
+        )
+        if snowpark_name is not None:
+            return snowpark_name, column_mapping
+
+        # Try outer scopes
+        for outer_df in get_outer_dataframes():
+            snowpark_name = (
+                outer_df.column_map.get_snowpark_column_name_from_spark_column_name(
+                    column_name, allow_non_exists=True
+                )
+            )
+            if snowpark_name is not None:
+                return snowpark_name, outer_df.column_map
+
+        return None, None
+
+    # Try to resolve the full qualified name first
+    snowpark_name, found_column_map = try_resolve_column(quoted_attr_name)
+
     if snowpark_name is not None:
         col = get_col(snowpark_name)
-        qualifiers = column_mapping.get_qualifier_for_spark_column(quoted_attr_name)
+        qualifiers = found_column_map.get_qualifiers_for_spark_column(quoted_attr_name)
     else:
         # this means it has to be a struct column with a field name
         snowpark_name: str | None = None
@@ -295,28 +317,43 @@ def map_unresolved_attribute(
         # For qualified names like "table.column.field", we need to find the column part
         for i in range(len(name_parts)):
             candidate_column = name_parts[i]
-            snowpark_name = (
-                column_mapping.get_snowpark_column_name_from_spark_column_name(
-                    candidate_column, allow_non_exists=True
-                )
-            )
+            snowpark_name, found_column_map = try_resolve_column(candidate_column)
+
             if snowpark_name is not None:
                 column_part_index = i
                 break
 
-            # Also try in outer dataframes
-            for outer_df in get_outer_dataframes():
-                snowpark_name = (
-                    outer_df.column_map.get_snowpark_column_name_from_spark_column_name(
-                        candidate_column, allow_non_exists=True
-                    )
-                )
-                if snowpark_name is not None:
-                    column_part_index = i
-                    break
+        # Validate qualifier scope: if we found a column but skipped prefix parts,
+        # those prefix parts could be valid qualifiers for the column
+        # We have prefix parts like 'nt1' in 'nt1.k' that were skipped
+        maybe_qualified = column_part_index > 0
+        if (
+            snowpark_name is not None
+            and maybe_qualified
+            and found_column_map is not None
+        ):
+            prefix_parts = name_parts[:column_part_index]
+            found_col_qualifiers = found_column_map.get_qualifiers_for_spark_column(
+                candidate_column
+            )
 
-            if snowpark_name is not None:
-                break
+            # Check if any qualifier matches the prefix
+            has_matching_qualifier = False
+            for qual in found_col_qualifiers:
+                if not qual.is_empty and len(qual.parts) >= len(prefix_parts):
+                    if qual.parts[-len(prefix_parts) :] == tuple(prefix_parts):
+                        has_matching_qualifier = True
+                        break
+
+            # If no matching qualifier, it's a scope violation
+            if not has_matching_qualifier:
+                # The prefix is not a valid qualifier for this column - scope violation!
+                exception = AnalysisException(
+                    f'[UNRESOLVED_COLUMN] Column "{attr_name}" cannot be resolved. '
+                    f'The table or alias "{".".join(prefix_parts)}" is not in scope or does not exist.'
+                )
+                attach_custom_error_code(exception, ErrorCodes.COLUMN_NOT_FOUND)
+                raise exception
 
         if snowpark_name is None:
             # Attempt LCA fallback.
@@ -338,7 +375,7 @@ def map_unresolved_attribute(
                 )
                 if snowpark_name is not None:
                     col = get_col(snowpark_name)
-                    qualifiers = column_mapping.get_qualifier_for_spark_column(
+                    qualifiers = column_mapping.get_qualifiers_for_spark_column(
                         unqualified_name
                     )
                     typed_col = TypedColumn(col, lambda: typer.type(col))
@@ -405,7 +442,7 @@ def map_unresolved_attribute(
         for field_name in path:
             col = col.getItem(field_name)
 
-        qualifiers = []
+        qualifiers = set()
 
     typed_col = TypedColumn(col, lambda: typer.type(col))
     typed_col.set_qualifiers(qualifiers)