snowpark-connect 0.32.0__py3-none-any.whl → 1.0.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():x}")
+
+
 @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,71 @@ class ColumnNameMap:
         else:
             return spark_name.upper()
 
+    def get_columns_after_join(
+        self, right: ColumnNameMap, join_columns: list[str], join_type: str
+    ) -> list[ColumnNames]:
+        """
+        Returns a list of columns (names and qualifiers) after a using_columns join with the given column map
+        """
+
+        # first, let's gather right-side join columns for qualifier lookup
+        # and the remaining columns to append them to the result
+        join_column_names = [self._normalized_spark_name(c) for c in join_columns]
+        right_join_columns: dict[str, ColumnNames] = {}
+        right_remaining_columns: list[ColumnNames] = []
+        for oc in right.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 right_join_columns:
+                right_join_columns[col_name] = oc
+            else:
+                right_remaining_columns.append(oc)
+
+        # now gather left-side columns
+        left_join_columns: dict[str, ColumnNames] = {}
+        left_remaining_columns: list[ColumnNames] = []
+        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 left_join_columns:
+                # only assign join-side qualifier for outer joins
+                match join_type:
+                    case "left":
+                        qualifiers = c.qualifiers
+                    case "right":
+                        qualifiers = right_join_columns[col_name].qualifiers
+                    case _:
+                        qualifiers = (
+                            c.qualifiers | right_join_columns[col_name].qualifiers
+                        )
+                left_join_columns[col_name] = ColumnNames(
+                    c.spark_name, c.snowpark_name, qualifiers
+                )
+            else:
+                left_remaining_columns.append(c)
+
+        # join columns go first in the user-given order,
+        # then the remaining left-side columns, then remaining right-side columns
+        match join_type:
+            case "right":
+                ordered_join_columns = [
+                    right_join_columns[name] for name in join_column_names
+                ]
+            case _:
+                ordered_join_columns = [
+                    left_join_columns[name] for name in join_column_names
+                ]
+        return ordered_join_columns + left_remaining_columns + right_remaining_columns
+
+    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__(
@@ -681,19 +741,6 @@ class JoinColumnNameMap(ColumnNameMap):
             else snowpark_column_name_in_left
         )
 
-        # this means that the reference is for the column in right dataframe but same snowpark name exist in left dataframe as well
-        # or vice versa, so we need to append _left or _right to the snowpark name
-        if (
-            snowpark_name in self.left_column_mapping.get_snowpark_columns()
-            and snowpark_column_name_in_right is not None
-        ):
-            snowpark_name = quote_name(f"{unquote_if_quoted(snowpark_name)}_right")
-        elif (
-            snowpark_name in self.right_column_mapping.get_snowpark_columns()
-            and snowpark_column_name_in_left is not None
-        ):
-            snowpark_name = quote_name(f"{unquote_if_quoted(snowpark_name)}_left")
-
         return snowpark_name
 
     def get_snowpark_column_names_from_spark_column_names(
@@ -784,19 +831,23 @@ 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, right: ColumnNameMap, join_columns: list[str], join_type: str
+    ) -> list[ColumnNames]:
+        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/config.py

@@ -275,6 +275,7 @@ SESSION_CONFIG_KEY_WHITELIST = {
     "spark.hadoop.fs.s3a.server-side-encryption.key",
     "spark.hadoop.fs.s3a.assumed.role.arn",
     "snowpark.connect.describe_cache_ttl_seconds",
+    "mapreduce.fileoutputcommitter.marksuccessfuljobs",
 }
 AZURE_ACCOUNT_KEY = re.compile(
     r"^fs\.azure\.sas\.[^\.]+\.[^\.]+\.blob\.core\.windows\.net$"
@@ -304,6 +305,7 @@ class SessionConfig:
         "spark.sql.tvf.allowMultipleTableArguments.enabled": "true",
         "snowpark.connect.enable_snowflake_extension_behavior": "false",
         "snowpark.connect.describe_cache_ttl_seconds": "300",
+        "mapreduce.fileoutputcommitter.marksuccessfuljobs": "false",
     }
 
     def __init__(self) -> None:
@@ -639,6 +641,13 @@ def get_cte_optimization_enabled() -> bool:
     return get_boolean_session_config_param("snowpark.connect.cte.optimization_enabled")
 
 
+def get_success_file_generation_enabled() -> bool:
+    """Get the _SUCCESS file generation configuration setting."""
+    return get_boolean_session_config_param(
+        "mapreduce.fileoutputcommitter.marksuccessfuljobs"
+    )
+
+
 def get_describe_cache_ttl_seconds() -> int:
     """Get the describe query cache TTL from session config, with a default fallback."""
     session_config: SessionConfig = sessions_config[get_session_id()]

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/literal.py

@@ -12,7 +12,6 @@ from tzlocal import get_localzone
 from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.error.error_codes import ErrorCodes
 from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
-from snowflake.snowpark_connect.utils.context import get_is_evaluating_sql
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
@@ -55,20 +54,21 @@ def get_literal_field_and_name(literal: expressions_proto.Expression.Literal):
                 microseconds = literal.timestamp
             else:
                 microseconds = literal.timestamp_ntz
-            lit_dt = datetime.datetime.fromtimestamp(
-                microseconds // 1_000_000
-            ) + datetime.timedelta(microseconds=microseconds % 1_000_000)
-            tz_dt = datetime.datetime.fromtimestamp(
-                microseconds // 1_000_000, tz=local_tz
+
+            dt_utc = datetime.datetime.fromtimestamp(
+                microseconds // 1_000_000, tz=datetime.timezone.utc
             ) + datetime.timedelta(microseconds=microseconds % 1_000_000)
+
             if t == "timestamp_ntz":
-                lit_dt = lit_dt.astimezone(datetime.timezone.utc)
-                tz_dt = tz_dt.astimezone(datetime.timezone.utc)
-            elif not get_is_evaluating_sql():
+                # For timestamp_ntz, display in UTC
+                lit_dt = dt_utc.replace(tzinfo=None)
+                tz_dt = dt_utc
+            else:
+                # For timestamp_ltz, always display in session timezone
                 config_tz = global_config.spark_sql_session_timeZone
-                config_tz = ZoneInfo(config_tz) if config_tz else local_tz
-                tz_dt = tz_dt.astimezone(config_tz)
-                lit_dt = lit_dt.astimezone(local_tz)
+                display_tz = ZoneInfo(config_tz) if config_tz else local_tz
+                tz_dt = dt_utc.astimezone(display_tz)
+                lit_dt = tz_dt.replace(tzinfo=None)
 
             def _format_timestamp(dt) -> str:
                 without_micros = f"{dt.year:04d}-{dt.month:02d}-{dt.day:02d} {dt.hour:02d}:{dt.minute:02d}:{dt.second:02d}"

snowflake/snowpark_connect/expression/map_sql_expression.py

@@ -67,10 +67,16 @@ def sql_parser():
     """
 
     ts_type = global_config.spark_sql_timestampType
+    session_tz = global_config.spark_sql_session_timeZone
 
     if ts_type is not None:
         _get_sql_conf().get().setConfString("spark.sql.timestampType", str(ts_type))
 
+    if session_tz is not None:
+        _get_sql_conf().get().setConfString(
+            "spark.sql.session.timeZone", str(session_tz)
+        )
+
     return _get_sql_parser()
 
 
@@ -418,13 +424,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
@@ -69,6 +70,143 @@ def _get_catalog_database_from_column_map(
     return catalog_database_info
 
 
+def _try_resolve_column_in_scopes(
+    column_name: str, column_mapping: ColumnNameMap
+) -> tuple[str | None, Any]:
+    """
+    Try to resolve a column name in current and outer scopes.
+
+    Args:
+        column_name: The column name to resolve
+        column_mapping: The column mapping for the current scope
+
+    Returns:
+        Tuple of (snowpark_name, found_column_map) or (None, None) if not found
+    """
+    # 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
+
+
+def _find_column_with_qualifier_match(
+    name_parts: list[str],
+    column_mapping: ColumnNameMap,
+) -> tuple[int, str | None, Any]:
+    """
+    Find the column position in name_parts where the prefix matches a qualifier.
+
+    In Spark, table qualifiers have at most 3 parts:
+    - 1 part: table only (e.g., 't1') → ColumnQualifier(('t1',))
+    - 2 parts: database.table (e.g., 'mydb.t5') → ColumnQualifier(('mydb', 't5'))
+    - 3 parts: catalog.database.table (e.g., 'cat.mydb.t5') → ColumnQualifier(('cat', 'mydb', 't5'))
+
+    Examples of how this works (suffix matching):
+    1) Input: "mydb1.t5.t5.i1" with qualifier ('mydb1', 't5')
+       - At i=2: prefix=['mydb1','t5'], matches qualifier suffix ('mydb1', 't5') → Column found!
+       - Remaining ['i1'] is treated as field access
+
+    2) Input: "t5.t5.i1" with qualifier ('mydb1', 't5')
+       - At i=1: prefix=['t5'], matches qualifier suffix ('t5',) → Column found!
+       - Remaining ['i1'] is treated as field access
+
+    3) Input: "cat.mydb.t5.t5.i1" with qualifier ('cat', 'mydb', 't5')
+       - At i=3: prefix=['cat','mydb','t5'], matches qualifier suffix → Column found!
+       - Remaining ['i1'] is treated as field access
+
+    The key insight: if the prefix before a candidate matches the END (suffix) of a qualifier,
+    then that position is the column reference. This allows partial qualification (e.g., just table
+    name instead of full database.table)
+
+    Args:
+        name_parts: The parts of the qualified name (e.g., ['mydb1', 't5', 't5', 'i1'])
+        column_mapping: The column mapping to resolve columns against
+
+    Returns:
+        Tuple of (column_part_index, snowpark_name, found_column_map)
+        Returns (0, None, None) if no valid column found
+
+    Raises:
+        AnalysisException: If a column is found but with invalid qualifier (scope violation)
+    """
+    # Track if we found a column but with wrong qualifier (scope violation)
+    scope_violation = None
+
+    for i in range(len(name_parts)):
+        candidate_column = name_parts[i]
+        snowpark_name, found_column_map = _try_resolve_column_in_scopes(
+            candidate_column, column_mapping
+        )
+
+        if snowpark_name is not None:
+            candidate_qualifiers = found_column_map.get_qualifiers_for_spark_column(
+                candidate_column
+            )
+            prefix_parts = name_parts[:i]
+
+            # Check if this is a valid column reference position
+            # A valid position is where the prefix exactly matches one of the qualifiers
+            is_valid_reference = False
+
+            if i == 0:
+                # No prefix (unqualified access)
+                # Always valid - Spark allows unqualified access to any column
+                # The remaining parts (name_parts[1:]) will be treated as
+                # struct/map/array field access (e.g., "person.address.city" where
+                # person is the column and address.city is the field path)
+                is_valid_reference = True
+            else:
+                # Has prefix - check if it matches the end (suffix) of any qualifier
+                # Spark allows partial qualification, so for qualifier ('mydb1', 't5'):
+                # - Can access as mydb1.t5.t5.i1 (full qualifier match)
+                # - Can access as t5.t5.i1 (suffix match - just table part)
+                # e.g., for "t5.t5.i1", when i=1, prefix=['t5'] matches suffix of ('mydb1', 't5')
+                # If valid, the remaining parts (name_parts[i+1:]) will be treated as
+                # struct/map/array field access (e.g., ['i1'] is a field in column t5)
+                for qual in candidate_qualifiers:
+                    if len(qual.parts) >= len(prefix_parts) and qual.parts[
+                        -len(prefix_parts) :
+                    ] == tuple(prefix_parts):
+                        is_valid_reference = True
+                        break
+
+            if is_valid_reference:
+                # This is the actual column reference
+                return (i, snowpark_name, found_column_map)
+            elif i > 0:
+                # Found column but qualifier doesn't match - this is a scope violation
+                # e.g., SELECT nt1.k where k exists but nt1 is not its qualifier
+                attr_name = ".".join(name_parts)
+                scope_violation = (attr_name, ".".join(prefix_parts))
+
+    # If we detected a scope violation, throw error
+    if scope_violation:
+        attr_name, invalid_qualifier = scope_violation
+        exception = AnalysisException(
+            f'[UNRESOLVED_COLUMN] Column "{attr_name}" cannot be resolved. '
+            f'The table or alias "{invalid_qualifier}" is not in scope or does not exist.'
+        )
+        attach_custom_error_code(exception, ErrorCodes.COLUMN_NOT_FOUND)
+        raise exception
+
+    # No valid column found
+    return (0, None, None)
+
+
 def map_unresolved_attribute(
     exp: expressions_proto.Expression,
     column_mapping: ColumnNameMap,
@@ -275,12 +413,14 @@ 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
+    # Try to resolve the full qualified name first
+    snowpark_name, found_column_map = _try_resolve_column_in_scopes(
+        quoted_attr_name, column_mapping
     )
+
     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
@@ -291,32 +431,13 @@ def map_unresolved_attribute(
             original_attr_name, column_mapping
         )
 
-        # Try to find the column name in different parts of the name_parts array
-        # 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
-                )
-            )
-            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
-
-            if snowpark_name is not None:
-                break
+        # Find the column by matching qualifiers with the prefix parts
+        # Note: This may raise AnalysisException if a scope violation is detected
+        (
+            column_part_index,
+            snowpark_name,
+            found_column_map,
+        ) = _find_column_with_qualifier_match(name_parts, column_mapping)
 
         if snowpark_name is None:
             # Attempt LCA fallback.