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

snowflake/snowpark_connect/column_name_handler.py

@@ -27,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)
 
@@ -46,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]:
@@ -91,6 +93,14 @@ 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
@@ -137,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
@@ -170,7 +180,7 @@ class ColumnNameMap:
                 snowpark_name=snowpark_column_names[i],
                 qualifiers=column_qualifiers[i]
                 if column_qualifiers and column_qualifiers[i]
-                else {ColumnQualifier.no_qualifier()},
+                else set(),
                 catalog_info=catalog_info,
                 database_info=database_info,
             )
@@ -511,21 +521,6 @@ class ColumnNameMap:
             if self._quote_if_unquoted(c.snowpark_name) not in cols_to_drop
         ]
 
-    def get_qualifier_for_spark_column(
-        self,
-        spark_column_name: str,
-    ) -> ColumnQualifier:
-        """
-        Backward compatibility: returns the first qualifier for the given Spark column name.
-        Throws if more than one qualifier exists.
-        """
-        qualifiers = self.get_qualifiers_for_spark_column(spark_column_name)
-        if len(qualifiers) > 1:
-            raise ValueError(
-                "Shouldn't happen. Multiple qualifiers found; expected only one."
-            )
-        return next(iter(qualifiers))
-
     def get_qualifiers_for_spark_column(
         self,
         spark_column_name: str,
@@ -544,7 +539,7 @@ class ColumnNameMap:
         col = mapping.get(name)
 
         if col is None or len(col) == 0:
-            return {ColumnQualifier.no_qualifier()}
+            return set()
 
         return col[0].qualifiers
 
@@ -605,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({ColumnQualifier.no_qualifier()})
+                qualifiers.append(set())
             else:
                 spark_columns.append(c.spark_name)
                 snowpark_columns.append(c.snowpark_name)
@@ -615,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({ColumnQualifier.no_qualifier()})
+                qualifiers.append(set())
 
         return spark_columns, snowpark_columns, qualifiers
 
@@ -625,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__(
@@ -784,19 +840,28 @@ class JoinColumnNameMap(ColumnNameMap):
     def get_qualifiers_for_spark_column(
         self, spark_column_name: str
     ) -> set[ColumnQualifier]:
-        return {self.get_qualifier_for_spark_column(spark_column_name)}
-
-    def get_qualifier_for_spark_column(self, spark_column_name: str) -> ColumnQualifier:
-        qualifier_left = self.left_column_mapping.get_qualifier_for_spark_column(
+        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 (not qualifier_left.is_empty) and (not qualifier_right.is_empty):
+        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 qualifier_left.is_empty 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

snowflake/snowpark_connect/column_qualifier.py

@@ -23,10 +23,6 @@ class ColumnQualifier:
     def is_empty(self) -> bool:
         return len(self.parts) == 0
 
-    @classmethod
-    def no_qualifier(cls) -> ColumnQualifier:
-        return cls(())
-
     def all_qualified_names(self, name: str) -> list[str]:
         qualifier_parts = self.parts
         qualifier_prefixes = [

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(

snowflake/snowpark_connect/expression/map_sql_expression.py

@@ -418,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
@@ -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_qualifiers_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.

snowflake/snowpark_connect/expression/map_unresolved_function.py

@@ -24,7 +24,6 @@ from typing import List, Optional
 from urllib.parse import quote, unquote
 
 import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
-import pyspark.sql.functions as pyspark_functions
 from google.protobuf.message import Message
 from pyspark.errors.exceptions.base import (
     AnalysisException,
@@ -101,6 +100,7 @@ from snowflake.snowpark_connect.expression.map_unresolved_star import (
 )
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.relation.catalogs.utils import CURRENT_CATALOG_NAME
+from snowflake.snowpark_connect.relation.utils import is_aggregate_function
 from snowflake.snowpark_connect.type_mapping import (
     map_json_schema_to_snowpark,
     map_pyspark_types_to_snowpark_types,
@@ -400,9 +400,8 @@ def map_unresolved_function(
     result_type: Optional[DataType | List[DateType]] = None
     qualifier_parts: List[str] = []
 
-    pyspark_func = getattr(pyspark_functions, function_name, None)
-    if pyspark_func and pyspark_func.__doc__.lstrip().startswith("Aggregate function:"):
-        # Used by the GROUP BY ALL implementation. Far from ideal, but it seems to work...
+    # Check if this is an aggregate function (used by GROUP BY ALL implementation)
+    if is_aggregate_function(function_name):
         add_sql_aggregate_function()
 
     def _type_with_typer(col: Column) -> TypedColumn:
@@ -912,15 +911,28 @@ def map_unresolved_function(
                 ):
                     # String + YearMonthInterval: Spark tries to cast string to double first, throws error if it fails
                     result_type = StringType()
+                    raise_error = _raise_error_helper(StringType(), AnalysisException)
                     if isinstance(snowpark_typed_args[0].typ, StringType):
-                        result_exp = (
-                            snowpark_fn.cast(snowpark_args[0], "double")
-                            + snowpark_args[1]
-                        )
+                        # Try to cast string to double, if it fails (returns null), raise exception
+                        cast_result = snowpark_fn.try_cast(snowpark_args[0], "double")
+                        result_exp = snowpark_fn.when(
+                            cast_result.is_null(),
+                            raise_error(
+                                snowpark_fn.lit(
+                                    f'The value \'{snowpark_args[0]}\' of the type {snowpark_typed_args[0].typ} cannot be cast to "DOUBLE" because it is malformed. Correct the value as per the syntax, or change its target type. Use `try_cast` to tolerate malformed input and return NULL instead. If necessary set "spark.sql.ansi.enabled" to "false" to bypass this error.'
+                                )
+                            ),
+                        ).otherwise(cast_result + snowpark_args[1])
                     else:
-                        result_exp = snowpark_args[0] + snowpark_fn.cast(
-                            snowpark_args[1], "double"
-                        )
+                        cast_result = snowpark_fn.try_cast(snowpark_args[1], "double")
+                        result_exp = snowpark_fn.when(
+                            cast_result.is_null(),
+                            raise_error(
+                                snowpark_fn.lit(
+                                    f'The value \'{snowpark_args[0]}\' of the type {snowpark_typed_args[0].typ} cannot be cast to "DOUBLE" because it is malformed. Correct the value as per the syntax, or change its target type. Use `try_cast` to tolerate malformed input and return NULL instead. If necessary set "spark.sql.ansi.enabled" to "false" to bypass this error.'
+                                )
+                            ),
+                        ).otherwise(snowpark_args[0] + cast_result)
                 case (StringType(), t) | (t, StringType()) if isinstance(
                     t, DayTimeIntervalType
                 ):
@@ -6184,15 +6196,19 @@ def map_unresolved_function(
                 or isinstance(arg_type, DayTimeIntervalType)
                 else DoubleType()
             )
-
-        case "pow":
-            result_exp = snowpark_fn.pow(snowpark_args[0], snowpark_args[1])
-            result_type = DoubleType()
-        case "power":
-            spark_function_name = (
-                f"POWER({snowpark_arg_names[0]}, {snowpark_arg_names[1]})"
-            )
-            result_exp = snowpark_fn.pow(snowpark_args[0], snowpark_args[1])
+        case "pow" | "power":
+            spark_function_name = f"{function_name if function_name == 'pow' else function_name.upper()}({snowpark_arg_names[0]}, {snowpark_arg_names[1]})"
+            if not spark_sql_ansi_enabled:
+                snowpark_args = _validate_numeric_args(
+                    function_name, snowpark_typed_args, snowpark_args
+                )
+            result_exp = snowpark_fn.when(
+                snowpark_fn.equal_nan(snowpark_fn.cast(snowpark_args[0], FloatType()))
+                | snowpark_fn.equal_nan(
+                    snowpark_fn.cast(snowpark_args[1], FloatType())
+                ),
+                NAN,
+            ).otherwise(snowpark_fn.pow(snowpark_args[0], snowpark_args[1]))
             result_type = DoubleType()
         case "product":
             col = snowpark_args[0]
@@ -9458,15 +9474,21 @@ def map_unresolved_function(
                 result_exp = snowpark_fn.year(snowpark_fn.to_date(snowpark_args[0]))
             result_type = LongType()
         case binary_method if binary_method in ("to_binary", "try_to_binary"):
-            binary_format = "hex"
+            binary_format = snowpark_fn.lit("hex")
+            arg_str = snowpark_fn.cast(snowpark_args[0], StringType())
             if len(snowpark_args) > 1:
                 binary_format = snowpark_args[1]
             result_exp = snowpark_fn.when(
                 snowpark_args[0].isNull(), snowpark_fn.lit(None)
             ).otherwise(
                 snowpark_fn.function(binary_method)(
-                    snowpark_fn.cast(snowpark_args[0], StringType()), binary_format
-                ),
+                    snowpark_fn.when(
+                        (snowpark_fn.length(arg_str) % 2 == 1)
+                        & (snowpark_fn.lower(binary_format) == snowpark_fn.lit("hex")),
+                        snowpark_fn.concat(snowpark_fn.lit("0"), arg_str),
+                    ).otherwise(arg_str),
+                    binary_format,
+                )
             )
             result_type = BinaryType()
         case udtf_name if udtf_name.lower() in session._udtfs:
@@ -10705,12 +10727,18 @@ def _try_sum_helper(
                     return snowpark_fn.lit(None), new_type
                 else:
                     non_null_rows = snowpark_fn.count(col_name)
-                    return aggregate_sum / non_null_rows, new_type
+                    # Use _divnull to handle case when non_null_rows is 0
+                    return _divnull(aggregate_sum, non_null_rows), new_type
             else:
                 new_type = DecimalType(
                     precision=min(38, arg_type.precision + 10), scale=arg_type.scale
                 )
-                return aggregate_sum, new_type
+                # Return NULL when there are no non-null values (i.e., all values are NULL); this is handled using case/when to check for non-null values for both SUM and the sum component of AVG calculations.
+                non_null_rows = snowpark_fn.count(col_name)
+                result = snowpark_fn.when(
+                    non_null_rows == 0, snowpark_fn.lit(None)
+                ).otherwise(aggregate_sum)
+                return result, new_type
 
         case _:
             # If the input column is floating point (double and float are synonymous in Snowflake per
@@ -10728,9 +10756,16 @@ def _try_sum_helper(
                     return snowpark_fn.lit(None), DoubleType()
                 else:
                     non_null_rows = snowpark_fn.count(col_name)
-                    return aggregate_sum / non_null_rows, DoubleType()
+                    # Use _divnull to handle case when non_null_rows is 0
+                    return _divnull(aggregate_sum, non_null_rows), DoubleType()
             else:
-                return aggregate_sum, DoubleType()
+                # When all values are NULL, SUM should return NULL (not 0)
+                # Use case/when to return NULL when there are no non-null values (i.e., all values are NULL)
+                non_null_rows = snowpark_fn.count(col_name)
+                result = snowpark_fn.when(
+                    non_null_rows == 0, snowpark_fn.lit(None)
+                ).otherwise(aggregate_sum)
+                return result, DoubleType()
 
 
 def _get_type_precision(typ: DataType) -> tuple[int, int]:

snowflake/snowpark_connect/includes/python/pyspark/pandas/spark/__init__.py

@@ -0,0 +1,16 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#