snowpark-connect 0.27.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

snowflake/snowpark_connect/execute_plan/map_execution_command.py

@@ -1,90 +1,29 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
-import re
-import uuid
-from collections import Counter
 
 import pyspark.sql.connect.proto.base_pb2 as proto_base
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
-from snowflake.snowpark import DataFrame, Session
-from snowflake.snowpark.exceptions import SnowparkSQLException
-from snowflake.snowpark_connect.column_name_handler import ColumnNames
-from snowflake.snowpark_connect.config import global_config, sessions_config
 from snowflake.snowpark_connect.constants import SERVER_SIDE_SESSION_ID
+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.execute_plan.utils import pandas_to_arrow_batches_bytes
 from snowflake.snowpark_connect.expression import map_udf
 from snowflake.snowpark_connect.relation import map_udtf
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 from snowflake.snowpark_connect.relation.map_sql import map_sql_to_pandas_df
-from snowflake.snowpark_connect.relation.write.map_write import map_write, map_write_v2
-from snowflake.snowpark_connect.utils.context import get_session_id
-from snowflake.snowpark_connect.utils.identifiers import (
-    spark_to_sf_single_id,
-    spark_to_sf_single_id_with_unquoting,
+from snowflake.snowpark_connect.relation.read.metadata_utils import (
+    without_internal_columns,
 )
+from snowflake.snowpark_connect.relation.write.map_write import map_write, map_write_v2
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
-
-_INTERNAL_VIEW_PREFIX = "__SC_RENAMED_V_"
-
-_CREATE_VIEW_PATTERN = re.compile(r"create\s+or\s+replace\s+view", re.IGNORECASE)
-
-
-def _create_column_rename_map(
-    columns: list[ColumnNames], rename_duplicated: bool
-) -> dict:
-    if rename_duplicated is False:
-        # if we are not renaming duplicated columns, we can just return the original names
-        return {
-            col.snowpark_name: spark_to_sf_single_id(col.spark_name, is_column=True)
-            for col in columns
-        }
-
-    column_counts = Counter()
-    not_renamed_cols = []
-    renamed_cols = []
-
-    for col in columns:
-        new_column_name = col.spark_name
-        normalized_name = new_column_name.lower()
-        column_counts[normalized_name] += 1
-
-        if column_counts[normalized_name] > 1:
-            new_column_name = (
-                f"{new_column_name}_DEDUP_{column_counts[normalized_name] - 1}"
-            )
-            renamed_cols.append(ColumnNames(new_column_name, col.snowpark_name, []))
-        else:
-            not_renamed_cols.append(ColumnNames(new_column_name, col.snowpark_name, []))
-
-    if len(renamed_cols) == 0:
-        return {
-            col.snowpark_name: spark_to_sf_single_id(col.spark_name, is_column=True)
-            for col in not_renamed_cols
-        }
-
-    # we need to make sure that we don't have duplicated names after renaming
-    # columns that were not renamed in this iteration should have priority over renamed duplicates
-    return _create_column_rename_map(not_renamed_cols + renamed_cols, True)
-
-
-def _find_duplicated_columns(
-    columns: list[ColumnNames],
-) -> (list[str], list[ColumnNames]):
-    duplicates = []
-    remaining_columns = []
-    seen = set()
-    for col in columns:
-        if col.spark_name in seen:
-            duplicates.append(col.snowpark_name)
-        else:
-            seen.add(col.spark_name)
-            remaining_columns.append(col)
-    return duplicates, remaining_columns
+from snowflake.snowpark_connect.utils.temporary_view_helper import (
+    create_temporary_view_from_dataframe,
+)
 
 
 def map_execution_command(
@@ -94,56 +33,10 @@ def map_execution_command(
     match request.plan.command.WhichOneof("command_type"):
         case "create_dataframe_view":
             req = request.plan.command.create_dataframe_view
-            input_df_container = map_relation(req.input)
-            input_df = input_df_container.dataframe
-            column_map = input_df_container.column_map
-
-            session_config = sessions_config[get_session_id()]
-            duplicate_column_names_handling_mode = session_config[
-                "snowpark.connect.views.duplicate_column_names_handling_mode"
-            ]
-
-            # rename columns to match spark names
-            if duplicate_column_names_handling_mode == "rename":
-                # deduplicate column names by appending _DEDUP_1, _DEDUP_2, etc.
-                input_df = input_df.rename(
-                    _create_column_rename_map(column_map.columns, True)
-                )
-            elif duplicate_column_names_handling_mode == "drop":
-                # Drop duplicate column names by removing all but the first occurrence.
-                duplicated_columns, remaining_columns = _find_duplicated_columns(
-                    column_map.columns
-                )
-                if len(duplicated_columns) > 0:
-                    input_df = input_df.drop(*duplicated_columns)
-                input_df = input_df.rename(
-                    _create_column_rename_map(remaining_columns, False)
-                )
-            else:
-                # rename columns without deduplication
-                input_df = input_df.rename(
-                    _create_column_rename_map(column_map.columns, False)
-                )
-
-            if req.is_global:
-                view_name = [global_config.spark_sql_globalTempDatabase, req.name]
-            else:
-                view_name = [req.name]
-            view_name = [
-                spark_to_sf_single_id_with_unquoting(part) for part in view_name
-            ]
-
-            if req.replace:
-                try:
-                    input_df.create_or_replace_temp_view(view_name)
-                except SnowparkSQLException as exc:
-                    if _is_error_caused_by_view_referencing_itself(exc):
-                        # This error is caused by statement with self reference like `CREATE VIEW A AS SELECT X FROM A`.
-                        _create_chained_view(input_df, view_name)
-                    else:
-                        raise
-            else:
-                input_df.create_temp_view(view_name)
+            input_df_container = without_internal_columns(map_relation(req.input))
+            create_temporary_view_from_dataframe(
+                input_df_container, req.name, req.is_global, req.replace
+            )
         case "write_stream_operation_start":
             match request.plan.command.write_stream_operation_start.format:
                 case "console":
@@ -204,46 +97,8 @@ def map_execution_command(
             map_udtf.register_udtf(request.plan.command.register_table_function)
 
         case other:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Command type {other} not implemented"
             )
-
-
-def _generate_random_builtin_view_name() -> str:
-    return _INTERNAL_VIEW_PREFIX + str(uuid.uuid4()).replace("-", "")
-
-
-def _is_error_caused_by_view_referencing_itself(exc: Exception) -> bool:
-    return "view definition refers to view being defined" in str(exc).lower()
-
-
-def _create_chained_view(input_df: DataFrame, view_name: str) -> None:
-    """
-    In order to create a view, which references itself, Spark would here take the previous
-    definition of A and paste it in place of `FROM A`. Snowflake would fail in such case, so
-    as a workaround, we create a chain of internal views instead. This function:
-    1. Renames previous definition of A to some internal name (instead of deleting).
-    2. Adjusts the DDL of a new statement to reference the name of a renmaed internal view, instead of itself.
-    """
-
-    session = Session.get_active_session()
-
-    view_name = ".".join(view_name)
-
-    tmp_name = _generate_random_builtin_view_name()
-    old_name_replacement = _generate_random_builtin_view_name()
-
-    input_df.create_or_replace_temp_view(tmp_name)
-
-    session.sql(f"ALTER VIEW {view_name} RENAME TO {old_name_replacement}").collect()
-
-    ddl: str = session.sql(f"SELECT GET_DDL('VIEW', '{tmp_name}')").collect()[0][0]
-
-    ddl = ddl.replace(view_name, old_name_replacement)
-
-    # GET_DDL result doesn't contain `TEMPORARY`, it's likely a bug.
-    ddl = _CREATE_VIEW_PATTERN.sub("create or replace temp view", ddl)
-
-    session.sql(ddl).collect()
-
-    session.sql(f"ALTER VIEW {tmp_name} RENAME TO {view_name}").collect()
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception

snowflake/snowpark_connect/execute_plan/map_execution_root.py

@@ -21,11 +21,16 @@ from snowflake.snowpark._internal.utils import (
 )
 from snowflake.snowpark_connect.constants import SERVER_SIDE_SESSION_ID
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
+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.execute_plan.utils import (
     arrow_table_to_arrow_bytes,
     pandas_to_arrow_batches_bytes,
 )
 from snowflake.snowpark_connect.relation.map_relation import map_relation
+from snowflake.snowpark_connect.relation.read.metadata_utils import (
+    without_internal_columns,
+)
 from snowflake.snowpark_connect.type_mapping import (
     map_snowpark_types_to_pyarrow_types,
     snowpark_to_proto_type,
@@ -53,7 +58,9 @@ def sproc_connector_fetch_arrow_batches_fix(self) -> Iterator[Table]:
     if self._prefetch_hook is not None:
         self._prefetch_hook()
     if self._query_result_format != "arrow":
-        raise NotSupportedError
+        exception = NotSupportedError()
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
     return self._result_set._fetch_arrow_batches()
 
 
@@ -92,14 +99,19 @@ def map_execution_root(
 ) -> Iterator[proto_base.ExecutePlanResponse | QueryResult]:
     result: DataFrameContainer | pandas.DataFrame = map_relation(request.plan.root)
     if isinstance(result, pandas.DataFrame):
-        result_df = result
+        pandas_df = result
+        data_bytes = pandas_to_arrow_batches_bytes(pandas_df)
+        row_count = len(pandas_df)
+        schema = None
+        yield _build_execute_plan_response(row_count, data_bytes, schema, request)
     else:
-        result_df = result.dataframe
-
-    if isinstance(result_df, snowpark.DataFrame):
-        snowpark_schema = result_df.schema
-        schema = snowpark_to_proto_type(snowpark_schema, result.column_map, result_df)
-        spark_columns = result.column_map.get_spark_columns()
+        filtered_result = without_internal_columns(result)
+        filtered_result_df = filtered_result.dataframe
+        snowpark_schema = filtered_result_df.schema
+        schema = snowpark_to_proto_type(
+            snowpark_schema, filtered_result.column_map, filtered_result_df
+        )
+        spark_columns = filtered_result.column_map.get_spark_columns()
         if tcm.TCM_MODE:
             # TCM result handling:
             # - small result (only one batch): just return the executePlanResponse
@@ -108,22 +120,22 @@ def map_execution_root(
             is_large_result = False
             second_batch = False
             first_arrow_table = None
-            with result_df.session.query_history() as qh:
-                for arrow_table in to_arrow_batch_iter(result_df):
+            with filtered_result_df.session.query_history() as qh:
+                for arrow_table in to_arrow_batch_iter(filtered_result_df):
                     if second_batch:
                         is_large_result = True
                         break
                     first_arrow_table = arrow_table
                     second_batch = True
                 queries_cnt = len(
-                    result_df._plan.execution_queries[PlanQueryType.QUERIES]
+                    filtered_result_df._plan.execution_queries[PlanQueryType.QUERIES]
                 )
                 # get query uuid from the last query; this may not be the last queries in query history because snowpark
                 # may run some post action queries, e.g., drop temp table.
                 query_id = qh.queries[queries_cnt - 1].query_id
             if first_arrow_table is None:
                 # empty arrow batch iterator
-                pandas_df = result_df.to_pandas()
+                pandas_df = filtered_result_df.to_pandas()
                 data_bytes = pandas_to_arrow_batches_bytes(pandas_df)
                 yield _build_execute_plan_response(0, data_bytes, schema, request)
             elif not tcm.TCM_RETURN_QUERY_ID_FOR_SMALL_RESULT and not is_large_result:
@@ -150,7 +162,7 @@ def map_execution_root(
                     spark_schema.SerializeToString(),
                 )
         else:
-            arrow_table_iter = to_arrow_batch_iter(result_df)
+            arrow_table_iter = to_arrow_batch_iter(filtered_result_df)
             batch_count = 0
             for arrow_table in arrow_table_iter:
                 if arrow_table.num_rows > 0:
@@ -166,12 +178,6 @@ def map_execution_root(
 
             # Empty result needs special processing
             if batch_count == 0:
-                pandas_df = result_df.to_pandas()
+                pandas_df = filtered_result_df.to_pandas()
                 data_bytes = pandas_to_arrow_batches_bytes(pandas_df)
                 yield _build_execute_plan_response(0, data_bytes, schema, request)
-    else:
-        pandas_df = result_df
-        data_bytes = pandas_to_arrow_batches_bytes(pandas_df)
-        row_count = len(pandas_df)
-        schema = None
-        yield _build_execute_plan_response(row_count, data_bytes, schema, request)

snowflake/snowpark_connect/execute_plan/utils.py

@@ -8,6 +8,8 @@ import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 from pyspark.sql.pandas.types import _dedup_names
 
 from snowflake.snowpark import types as sf_types
+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.type_mapping import map_snowpark_types_to_pyarrow_types
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
@@ -88,9 +90,11 @@ def is_streaming(rel: relation_proto.Relation) -> bool:
             case "html_string":
                 return is_streaming(rel.html_string.input)
             case "cached_remote_relation":
-                raise SnowparkConnectNotImplementedError(
+                exception = SnowparkConnectNotImplementedError(
                     "Cached remote relation not implemented"
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+                raise exception
             case "common_inline_user_defined_table_function":
                 return is_streaming(rel.common_inline_user_defined_table_function.input)
             case "fill_na":

snowflake/snowpark_connect/expression/function_defaults.py

@@ -7,6 +7,9 @@ from typing import Any
 import pyspark.sql.connect.proto.expressions_pb2 as expressions_pb2
 import pyspark.sql.connect.proto.types_pb2 as types_pb2
 
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
+
 
 @dataclass(frozen=True)
 class DefaultParameter:
@@ -154,7 +157,9 @@ def _create_literal_expression(value: Any) -> expressions_pb2.Expression:
         null_type.null.SetInParent()
         expr.literal.null.CopyFrom(null_type)
     else:
-        raise ValueError(f"Unsupported literal type: {value}")
+        exception = ValueError(f"Unsupported literal type: {value}")
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+        raise exception
 
     return expr
 
@@ -189,11 +194,13 @@ def inject_function_defaults(
 
     # Check if any required params are missing.
     if missing_arg_count > len(defaults):
-        raise ValueError(
+        exception = ValueError(
             f"Function '{function_name}' is missing required arguments. "
             f"Expected {total_args} args, got {current_arg_count}, "
             f"but only {len(defaults)} defaults are defined."
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT)
+        raise exception
 
     defaults_to_append = defaults[-missing_arg_count:]
     injected = False

snowflake/snowpark_connect/expression/hybrid_column_map.py

@@ -18,6 +18,9 @@ from snowflake import snowpark
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.typed_column import TypedColumn
+from snowflake.snowpark_connect.utils.identifiers import (
+    split_fully_qualified_spark_name,
+)
 
 
 class HybridColumnMap:
@@ -119,9 +122,11 @@ class HybridColumnMap:
         # Handle column references
         if expr_type == "unresolved_attribute":
             column_name = exp.unresolved_attribute.unparsed_identifier
+            name_parts = split_fully_qualified_spark_name(column_name)
+            alias_column_name = name_parts[0]
 
             # Check if it's an alias to an existing aggregate expression
-            if column_name in self.aggregate_aliases:
+            if alias_column_name in self.aggregate_aliases:
                 # Use the aggregated context to get the alias
                 return map_expression(
                     exp, self.aggregated_column_map, self.aggregated_typer
@@ -148,14 +153,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(
@@ -190,3 +196,45 @@ def create_hybrid_column_map_for_having(
         grouping_expressions=grouping_expressions,
         aggregate_aliases=aggregate_aliases,
     )
+
+
+def create_hybrid_column_map_for_order_by(
+    aggregate_metadata,  # AggregateMetadata type
+    aggregated_df: snowpark.DataFrame,
+    aggregated_column_map: ColumnNameMap,
+) -> HybridColumnMap:
+    """
+    Create a HybridColumnMap instance for ORDER BY clause resolution after aggregation.
+
+    This is similar to HAVING clause resolution - ORDER BY can reference:
+    1. Grouping columns (e.g., year, a)
+    2. Aggregate aliases (e.g., cnt)
+    3. Expressions on grouping columns (e.g., year(date) where date is pre-aggregation)
+
+    Args:
+        aggregate_metadata: Metadata from the aggregate operation
+        aggregated_df: The DataFrame after aggregation
+        aggregated_column_map: Column mapping for the aggregated DataFrame
+
+    Returns:
+        HybridColumnMap for resolving ORDER BY expressions
+    """
+    # Create typers for both contexts
+    input_typer = ExpressionTyper(aggregate_metadata.input_dataframe)
+    aggregated_typer = ExpressionTyper(aggregated_df)
+
+    # Build alias mapping from spark column names to aggregate expressions
+    aggregate_aliases = {}
+    for i, (spark_name, _) in enumerate(aggregate_metadata.raw_aggregations):
+        if i < len(aggregate_metadata.aggregate_expressions):
+            aggregate_aliases[spark_name] = aggregate_metadata.aggregate_expressions[i]
+
+    return HybridColumnMap(
+        input_column_map=aggregate_metadata.input_column_map,
+        input_typer=input_typer,
+        aggregated_column_map=aggregated_column_map,
+        aggregated_typer=aggregated_typer,
+        aggregate_expressions=aggregate_metadata.aggregate_expressions,
+        grouping_expressions=aggregate_metadata.grouping_expressions,
+        aggregate_aliases=aggregate_aliases,
+    )

snowflake/snowpark_connect/expression/literal.py

@@ -10,7 +10,8 @@ import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 from tzlocal import get_localzone
 
 from snowflake.snowpark_connect.config import global_config
-from snowflake.snowpark_connect.utils.context import get_is_evaluating_sql
+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.telemetry import (
     SnowparkConnectNotImplementedError,
 )
@@ -53,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}"
@@ -97,7 +99,29 @@ def get_literal_field_and_name(literal: expressions_proto.Expression.Literal):
                 *(get_literal_field_and_name(e) for e in literal.array.elements)
             )
             return array_values, f"ARRAY({', '.join(element_names)})"
+        case "struct":
+            struct_key_names = [
+                field.name for field in literal.struct.struct_type.struct.fields
+            ]
+            struct_values = [
+                get_literal_field_and_name(el)[0] for el in literal.struct.elements
+            ]
+
+            struct_dict = dict(zip(struct_key_names, struct_values))
+
+            struct_elements = [
+                item for pair in zip(struct_key_names, struct_values) for item in pair
+            ]
+
+            return (
+                struct_dict,
+                f"OBJECT_CONSTRUCT_KEEP_NULL({', '.join(str(x) for x in struct_elements)})",
+            )
         case "null" | None:
             return None, "NULL"
         case other:
-            raise SnowparkConnectNotImplementedError(f"Other Literal Type {other}")
+            exception = SnowparkConnectNotImplementedError(
+                f"Other Literal Type {other}"
+            )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception