snowpark-connect 0.20.2__py3-none-any.whl → 0.22.1__py3-none-any.whl

snowflake/snowpark_connect/analyze_plan/map_tree_string.py

@@ -15,8 +15,9 @@ def map_tree_string(
 ) -> proto_base.AnalyzePlanResponse:
     # TODO: tracking the difference with pyspark in SNOW-1853347
     tree_string = request.tree_string
-    snowpark_df = map_relation(tree_string.plan.root)
-    column_map = snowpark_df._column_map
+    snowpark_df_container = map_relation(tree_string.plan.root)
+    snowpark_df = snowpark_df_container.dataframe
+    column_map = snowpark_df_container.column_map
 
     snowpark_tree_string = snowpark_df._format_schema(
         level=tree_string.level if tree_string.HasField("level") else None,

snowflake/snowpark_connect/column_name_handler.py

@@ -12,14 +12,13 @@ from functools import cached_property
 
 from pyspark.errors.exceptions.base import AnalysisException
 
-from snowflake import snowpark
 from snowflake.snowpark import DataFrame
 from snowflake.snowpark._internal.analyzer.analyzer_utils import (
     quote_name_without_upper_casing,
     unquote_if_quoted,
 )
 from snowflake.snowpark._internal.utils import quote_name
-from snowflake.snowpark.types import DataType, StructField, StructType
+from snowflake.snowpark.types import StructType
 from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.utils.context import get_current_operation_scope
 
@@ -41,64 +40,6 @@ def set_schema_getter(df: DataFrame, get_schema: Callable[[], StructType]) -> No
     df.__class__ = PatchedDataFrame
 
 
-def with_column_map(
-    result_df: snowpark.DataFrame,
-    spark_column_names: list[str],
-    snowpark_column_names: list[str],
-    snowpark_column_types: list[DataType] = None,
-    column_metadata: dict | None = None,
-    column_qualifiers: list[list[str]] | None = None,
-    parent_column_name_map: ColumnNameMap | None = None,
-) -> snowpark.DataFrame:
-    """
-    Build a mapping from the DataFrame's column names to the Spark column names.
-
-    This is used to track the original column names and handle column naming differences
-    between Spark and Snowpark.
-
-    The elements in result_df.columns and the elements in spark_column_names must be a one-to-one mapping.
-
-    Args:
-        result_df (snowpark.DataFrame): The DataFrame to map.
-        spark_column_names (list[str]): The Spark column names.
-        snowpark_column_names (list[str]): The Snowpark column names.
-        snowpark_column_types (list[DataType], optional): The Snowpark column types. **if provided df.schema will be overridden with inferred schema**
-        column_metadata (dict, optional): Metadata for the columns.
-        column_qualifiers (list[list[str]], optional): Qualifiers for the columns, used to handle table aliases or DataFrame aliases.
-        parent_column_name_map (ColumnNameMap, optional): A ColumnNameMap, that came from the dataframe used to create result_df (parent df)
-
-    Returns:
-        snowpark.DataFrame: The mapped DataFrame.
-    """
-    assert len(snowpark_column_names) == len(
-        spark_column_names
-    ), "Number of Spark column names must match number of columns in DataFrame"
-    result_df._column_map = ColumnNameMap(
-        spark_column_names,
-        snowpark_column_names,
-        column_metadata=column_metadata,
-        column_qualifiers=column_qualifiers,
-        parent_column_name_map=parent_column_name_map,
-    )
-    result_df._table_name = None
-
-    if snowpark_column_types is not None:
-        assert len(snowpark_column_names) == len(
-            snowpark_column_types
-        ), "Number of Snowpark column names and types must match"
-
-        set_schema_getter(
-            result_df,
-            lambda: StructType(
-                [
-                    StructField(n, t, _is_column=False)
-                    for n, t in zip(snowpark_column_names, snowpark_column_types)
-                ]
-            ),
-        )
-    return result_df
-
-
 def make_column_names_snowpark_compatible(
     names: list[str], plan_id: int, offset: int = 0
 ) -> list[str]:
@@ -189,7 +130,7 @@ class ColumnNameMap:
         column_qualifiers: Optional qualifiers for the columns, used to handle table aliases or DataFrame aliases.
         parent_column_name_map: parent ColumnNameMap
         """
-        self.columns = []
+        self.columns: list[ColumnNames] = []
         self.spark_to_col = defaultdict(list)
         self.uppercase_spark_to_col = defaultdict(list)
         self.snowpark_to_col = defaultdict(list)
@@ -602,11 +543,11 @@ class ColumnNameMap:
 class JoinColumnNameMap(ColumnNameMap):
     def __init__(
         self,
-        left_input: snowpark.DataFrame,
-        right_input: snowpark.DataFrame,
+        left_colmap: ColumnNameMap,
+        right_colmap: ColumnNameMap,
     ) -> None:
-        self.left_column_mapping: ColumnNameMap = left_input._column_map
-        self.right_column_mapping: ColumnNameMap = right_input._column_map
+        self.left_column_mapping: ColumnNameMap = left_colmap
+        self.right_column_mapping: ColumnNameMap = right_colmap
 
     def get_snowpark_column_name_from_spark_column_name(
         self,

snowflake/snowpark_connect/config.py

@@ -9,7 +9,7 @@ import re
 import sys
 import time
 from collections import defaultdict
-from copy import copy
+from copy import copy, deepcopy
 from typing import Any
 
 import jpype
@@ -33,7 +33,7 @@ from snowflake.snowpark_connect.version import VERSION as sas_version
 
 
 def str_to_bool(boolean_str: str) -> bool:
-    assert boolean_str in [
+    assert boolean_str in (
         "True",
         "true",
         "False",
@@ -41,7 +41,7 @@ def str_to_bool(boolean_str: str) -> bool:
         "1",
         "0",
         "",  # This is the default value, equivalent to False.
-    ], f"Invalid boolean value: {boolean_str}"
+    ), f"Invalid boolean value: {boolean_str}"
     return boolean_str in ["True", "true", "1"]
 
 
@@ -131,6 +131,7 @@ class GlobalConfig:
         "spark.sql.caseSensitive": "false",
         "spark.sql.mapKeyDedupPolicy": "EXCEPTION",
         "spark.sql.ansi.enabled": "false",
+        "spark.sql.legacy.allowHashOnMapType": "false",
         "spark.sql.sources.default": "parquet",
         "spark.Catalog.databaseFilterInformationSchema": "false",
         "spark.sql.parser.quotedRegexColumnNames": "false",
@@ -145,6 +146,7 @@ class GlobalConfig:
         "spark.sql.crossJoin.enabled",
         "spark.sql.caseSensitive",
         "spark.sql.ansi.enabled",
+        "spark.sql.legacy.allowHashOnMapType",
         "spark.Catalog.databaseFilterInformationSchema",
         "spark.sql.parser.quotedRegexColumnNames",
     ]
@@ -166,6 +168,9 @@ class GlobalConfig:
         "snowpark.connect.udf.packages": lambda session, packages: session.add_packages(
             *packages.strip("[] ").split(",")
         ),
+        "snowpark.connect.udf.imports": lambda session, imports: parse_imports(
+            session, imports
+        ),
     }
 
     float_config_list = []
@@ -250,10 +255,10 @@ SESSION_CONFIG_KEY_WHITELIST = {
     "spark.sql.tvf.allowMultipleTableArguments.enabled",
     "snowpark.connect.sql.passthrough",
     "snowpark.connect.iceberg.external_volume",
-    "snowpark.connect.auto-uppercase.ddl",
-    "snowpark.connect.auto-uppercase.dml",
+    "snowpark.connect.sql.identifiers.auto-uppercase",
     "snowpark.connect.udtf.compatibility_mode",
     "snowpark.connect.views.duplicate_column_names_handling_mode",
+    "enable_snowflake_extension_behavior",
 }
 AZURE_SAS_KEY = re.compile(
     r"^fs\.azure\.sas\.[^\.]+\.[^\.]+\.blob\.core\.windows\.net$"
@@ -271,17 +276,17 @@ class SessionConfig:
     """This class contains the session configuration for the Spark Server."""
 
     default_session_config = {
-        "snowpark.connect.auto-uppercase.ddl": "true",
-        "snowpark.connect.auto-uppercase.dml": "true",
+        "snowpark.connect.sql.identifiers.auto-uppercase": "all_except_columns",
         "snowpark.connect.sql.passthrough": "false",
         "snowpark.connect.udtf.compatibility_mode": "false",
         "snowpark.connect.views.duplicate_column_names_handling_mode": "rename",
         "spark.sql.execution.pythonUDTF.arrow.enabled": "false",
         "spark.sql.tvf.allowMultipleTableArguments.enabled": "true",
+        "enable_snowflake_extension_behavior": "false",
     }
 
     def __init__(self) -> None:
-        self.config = copy(self.default_session_config)
+        self.config = deepcopy(self.default_session_config)
 
     def __getitem__(self, item: str) -> str:
         return self.get(item)
@@ -304,7 +309,13 @@ CONFIG_ALLOWED_VALUES: dict[str, tuple] = {
         "rename",
         "fail",
         "drop",
-    )
+    ),
+    "snowpark.connect.sql.identifiers.auto-uppercase": (
+        "all_except_columns",
+        "only_columns",
+        "all",
+        "none",
+    ),
 }
 
 # Set some default configuration that are necessary for the driver.
@@ -324,7 +335,7 @@ def route_config_proto(
     match op_type:
         case "set":
             logger.info("SET")
-
+            telemetry.report_config_set(config.operation.set.pairs)
             for pair in config.operation.set.pairs:
                 # Check if the value field is present, not present when invalid fields are set in conf.
                 if not pair.HasField("value"):
@@ -334,7 +345,6 @@ def route_config_proto(
                         f"Cannot set config '{pair.key}' to None"
                     )
 
-                telemetry.report_config_set(pair.key, pair.value)
                 set_config_param(
                     config.session_id, pair.key, pair.value, snowpark_session
                 )
@@ -342,14 +352,15 @@ def route_config_proto(
             return proto_base.ConfigResponse(session_id=config.session_id)
         case "unset":
             logger.info("UNSET")
+            telemetry.report_config_unset(config.operation.unset.keys)
             for key in config.operation.unset.keys:
-                telemetry.report_config_unset(key)
                 unset_config_param(config.session_id, key, snowpark_session)
 
             return proto_base.ConfigResponse(session_id=config.session_id)
         case "get":
             logger.info("GET")
             res = proto_base.ConfigResponse(session_id=config.session_id)
+            telemetry.report_config_get(config.operation.get.keys)
             for key in config.operation.get.keys:
                 pair = res.pairs.add()
                 pair.key = key
@@ -359,6 +370,9 @@ def route_config_proto(
             return res
         case "get_with_default":
             logger.info("GET_WITH_DEFAULT")
+            telemetry.report_config_get(
+                [pair.key for pair in config.operation.get_with_default.pairs]
+            )
             result_pairs = [
                 proto_base.KeyValue(
                     key=pair.key,
@@ -375,6 +389,7 @@ def route_config_proto(
         case "get_option":
             logger.info("GET_OPTION")
             res = proto_base.ConfigResponse(session_id=config.session_id)
+            telemetry.report_config_get(config.operation.get_option.keys)
             for key in config.operation.get_option.keys:
                 pair = res.pairs.add()
                 pair.key = key
@@ -403,6 +418,7 @@ def route_config_proto(
         case "is_modifiable":
             logger.info("IS_MODIFIABLE")
             res = proto_base.ConfigResponse(session_id=config.session_id)
+            telemetry.report_config_get(config.operation.is_modifiable.keys)
             for key in config.operation.is_modifiable.keys:
                 pair = res.pairs.add()
                 pair.key = key
@@ -533,7 +549,7 @@ def set_snowflake_parameters(
                 value = global_config.default_static_global_config.get(key)
 
             snowpark_name = quote_name_without_upper_casing(value)
-            if auto_uppercase_ddl():
+            if auto_uppercase_non_column_identifiers():
                 snowpark_name = snowpark_name.upper()
 
             # Create the schema on demand. Before creating it, however,
@@ -568,9 +584,23 @@ def get_boolean_session_config_param(name: str) -> bool:
     return str_to_bool(session_config[name])
 
 
-def auto_uppercase_dml() -> bool:
-    return get_boolean_session_config_param("snowpark.connect.auto-uppercase.dml")
+def auto_uppercase_column_identifiers() -> bool:
+    session_config = sessions_config[get_session_id()]
+    return session_config[
+        "snowpark.connect.sql.identifiers.auto-uppercase"
+    ].lower() in ("all", "only_columns")
+
+
+def auto_uppercase_non_column_identifiers() -> bool:
+    session_config = sessions_config[get_session_id()]
+    return session_config[
+        "snowpark.connect.sql.identifiers.auto-uppercase"
+    ].lower() in ("all", "all_except_columns")
+
 
+def parse_imports(session: snowpark.Session, imports: str | None) -> None:
+    if not imports:
+        return
 
-def auto_uppercase_ddl() -> bool:
-    return get_boolean_session_config_param("snowpark.connect.auto-uppercase.ddl")
+    for udf_import in imports.strip("[] ").split(","):
+        session.add_import(udf_import)

snowflake/snowpark_connect/dataframe_container.py

@@ -0,0 +1,242 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable
+
+from snowflake import snowpark
+from snowflake.snowpark.types import StructField, StructType
+
+if TYPE_CHECKING:
+    from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
+
+
+class DataFrameContainer:
+    """
+    A container class that wraps a Snowpark DataFrame along with additional metadata.
+
+    This class provides a unified interface for managing Snowpark DataFrames along with
+    their column mappings, schema information, and metadata.
+    """
+
+    def __init__(
+        self,
+        dataframe: snowpark.DataFrame,
+        column_map: ColumnNameMap | None = None,
+        table_name: str | None = None,
+        alias: str | None = None,
+        cached_schema_getter: Callable[[], StructType] | None = None,
+    ) -> None:
+        """
+        Initialize a new DataFrameContainer.
+
+        Args:
+            dataframe: The underlying Snowpark DataFrame
+            column_map: Optional column name mapping
+            table_name: Optional table name for the DataFrame
+            alias: Optional alias for the DataFrame
+            cached_schema_getter: Optional function to get cached schema
+        """
+        self._dataframe = dataframe
+        self._column_map = self._create_default_column_map(column_map)
+        self._table_name = table_name
+        self._alias = alias
+
+        if cached_schema_getter is not None:
+            self._apply_cached_schema_getter(cached_schema_getter)
+
+    @classmethod
+    def create_with_column_mapping(
+        cls,
+        dataframe: snowpark.DataFrame,
+        spark_column_names: list[str],
+        snowpark_column_names: list[str],
+        snowpark_column_types: list | None = None,
+        column_metadata: dict | None = None,
+        column_qualifiers: list[list[str]] | None = None,
+        parent_column_name_map: ColumnNameMap | None = None,
+        table_name: str | None = None,
+        alias: str | None = None,
+        cached_schema_getter: Callable[[], StructType] | None = None,
+    ) -> DataFrameContainer:
+        """
+        Create a new container with complete column mapping configuration.
+
+        Args:
+            dataframe: The underlying Snowpark DataFrame
+            spark_column_names: List of Spark column names
+            snowpark_column_names: List of corresponding Snowpark column names
+            snowpark_column_types: Optional list of column types
+            column_metadata: Optional metadata dictionary
+            column_qualifiers: Optional column qualifiers
+            parent_column_name_map: Optional parent column name map
+            table_name: Optional table name
+            alias: Optional alias
+            cached_schema_getter: Optional function to get cached schema
+
+        Returns:
+            A new DataFrameContainer instance
+
+        Raises:
+            AssertionError: If column names and types don't match expected lengths
+        """
+        # Validate inputs
+        cls._validate_column_mapping_inputs(
+            spark_column_names, snowpark_column_names, snowpark_column_types
+        )
+
+        column_map = cls._create_column_map(
+            spark_column_names,
+            snowpark_column_names,
+            column_metadata,
+            column_qualifiers,
+            parent_column_name_map,
+        )
+
+        # Determine the schema getter to use
+        final_schema_getter = None
+
+        if cached_schema_getter is not None:
+            # Use the provided schema getter
+            final_schema_getter = cached_schema_getter
+        elif snowpark_column_types is not None:
+            # Create schema from types and wrap in function
+            schema = cls._create_schema_from_types(
+                snowpark_column_names, snowpark_column_types
+            )
+            if schema is not None:
+
+                def get_schema():
+                    return schema
+
+                final_schema_getter = get_schema
+
+        return cls(
+            dataframe=dataframe,
+            column_map=column_map,
+            table_name=table_name,
+            alias=alias,
+            cached_schema_getter=final_schema_getter,
+        )
+
+    @property
+    def dataframe(self) -> snowpark.DataFrame:
+        """Get the underlying Snowpark DataFrame."""
+        # Ensure the DataFrame has the _column_map attribute for backward compatibility
+        # Some of the snowpark code needs references to _column_map
+        self._dataframe._column_map = self._column_map
+        return self._dataframe
+
+    @property
+    def column_map(self) -> ColumnNameMap:
+        """Get the column name mapping."""
+        return self._column_map
+
+    @column_map.setter
+    def column_map(self, value: ColumnNameMap) -> None:
+        """Set the column name mapping."""
+        self._column_map = value
+
+    @property
+    def table_name(self) -> str | None:
+        """Get the table name."""
+        return self._table_name
+
+    @table_name.setter
+    def table_name(self, value: str | None) -> None:
+        """Set the table name."""
+        self._table_name = value
+
+    @property
+    def alias(self) -> str | None:
+        """Get the alias name."""
+        return self._alias
+
+    @alias.setter
+    def alias(self, value: str | None) -> None:
+        """Set the alias name."""
+        self._alias = value
+
+    def _create_default_column_map(
+        self, column_map: ColumnNameMap | None
+    ) -> ColumnNameMap:
+        """Create a default column map if none provided."""
+        if column_map is not None:
+            return column_map
+
+        from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
+
+        return ColumnNameMap([], [])
+
+    def _apply_cached_schema_getter(
+        self, schema_getter: Callable[[], StructType]
+    ) -> None:
+        """Apply a cached schema getter to the dataframe."""
+        from snowflake.snowpark_connect.column_name_handler import set_schema_getter
+
+        set_schema_getter(self._dataframe, schema_getter)
+
+    @staticmethod
+    def _validate_column_mapping_inputs(
+        spark_column_names: list[str],
+        snowpark_column_names: list[str],
+        snowpark_column_types: list | None = None,
+    ) -> None:
+        """
+        Validate inputs for column mapping creation.
+
+        Raises:
+            AssertionError: If validation fails
+        """
+        assert len(snowpark_column_names) == len(
+            spark_column_names
+        ), "Number of Spark column names must match number of columns in DataFrame"
+
+        if snowpark_column_types is not None:
+            assert len(snowpark_column_names) == len(
+                snowpark_column_types
+            ), "Number of Snowpark column names and types must match"
+
+    @staticmethod
+    def _create_column_map(
+        spark_column_names: list[str],
+        snowpark_column_names: list[str],
+        column_metadata: dict | None = None,
+        column_qualifiers: list[list[str]] | None = None,
+        parent_column_name_map: ColumnNameMap | None = None,
+    ) -> ColumnNameMap:
+        """Create a ColumnNameMap with the provided configuration."""
+        from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
+
+        return ColumnNameMap(
+            spark_column_names,
+            snowpark_column_names,
+            column_metadata=column_metadata,
+            column_qualifiers=column_qualifiers,
+            parent_column_name_map=parent_column_name_map,
+        )
+
+    @staticmethod
+    def _create_schema_from_types(
+        snowpark_column_names: list[str],
+        snowpark_column_types: list | None,
+    ) -> StructType | None:
+        """
+        Create a StructType schema from column names and types.
+
+        Returns:
+            StructType if types are provided, None otherwise
+        """
+        if snowpark_column_types is None:
+            return None
+
+        return StructType(
+            [
+                StructField(name, column_type, _is_column=False)
+                for name, column_type in zip(
+                    snowpark_column_names, snowpark_column_types
+                )
+            ]
+        )

snowflake/snowpark_connect/error/error_utils.py

@@ -28,7 +28,9 @@ from pyspark.errors.exceptions.base import (
     PySparkException,
     PythonException,
     SparkRuntimeException,
+    UnsupportedOperationException,
 )
+from pyspark.errors.exceptions.connect import SparkConnectGrpcException
 from snowflake.core.exceptions import NotFoundError
 
 from snowflake.connector.errors import ProgrammingError
@@ -49,7 +51,9 @@ SPARK_PYTHON_TO_JAVA_EXCEPTION = {
     ArrayIndexOutOfBoundsException: "java.lang.ArrayIndexOutOfBoundsException",
     NumberFormatException: "java.lang.NumberFormatException",
     SparkRuntimeException: "org.apache.spark.SparkRuntimeException",
+    SparkConnectGrpcException: "pyspark.errors.exceptions.connect.SparkConnectGrpcException",
     PythonException: "org.apache.spark.api.python.PythonException",
+    UnsupportedOperationException: "java.lang.UnsupportedOperationException",
 }
 
 WINDOW_FUNCTION_ANALYSIS_EXCEPTION_SQL_ERROR_CODE = {1005, 2303}
@@ -68,6 +72,9 @@ init_multi_args_exception_pattern = (
 terminate_multi_args_exception_pattern = (
     r"terminate\(\) missing \d+ required positional argument"
 )
+snowpark_connect_exception_pattern = re.compile(
+    r"\[snowpark-connect-exception(?::(\w+))?\]\s*(.+?)'\s*is not recognized"
+)
 
 
 def contains_udtf_select(sql_string):
@@ -100,6 +107,19 @@ def _get_converted_known_sql_or_custom_exception(
         return SparkRuntimeException(
             message="Unexpected value for start in function slice: SQL array indices start at 1."
         )
+    match = snowpark_connect_exception_pattern.search(
+        ex.message if hasattr(ex, "message") else str(ex)
+    )
+    if match:
+        class_name = match.group(1)
+        message = match.group(2)
+        exception_class = (
+            globals().get(class_name, SparkConnectGrpcException)
+            if class_name
+            else SparkConnectGrpcException
+        )
+        return exception_class(message=message)
+
     if "select with no columns" in msg and contains_udtf_select(query):
         # We try our best to detect if the SQL string contains a UDTF call and the output schema is empty.
         return PythonException(message=f"[UDTF_RETURN_SCHEMA_MISMATCH] {ex.message}")
@@ -131,6 +151,11 @@ def _get_converted_known_sql_or_custom_exception(
                 message=f"[UDTF_EXEC_ERROR] User defined table function encountered an error in the terminate method: {ex.message}"
             )
 
+        if "failed to split string, provided pattern:" in msg:
+            return IllegalArgumentException(
+                message=f"Failed to split string using provided pattern. {ex.message}"
+            )
+
         if "100357" in msg and "wrong tuple size for returned value" in msg:
             return PythonException(
                 message=f"[UDTF_RETURN_SCHEMA_MISMATCH] The number of columns in the result does not match the specified schema. {ex.message}"