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

snowflake/snowpark_connect/column_qualifier.py

@@ -0,0 +1,43 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from snowflake.snowpark._internal.analyzer.analyzer_utils import (
+    quote_name_without_upper_casing,
+)
+
+
+@dataclass(frozen=True)
+class ColumnQualifier:
+    parts: tuple[str, ...]
+
+    def __post_init__(self) -> None:
+        if not all(isinstance(x, str) for x in self.parts):
+            raise TypeError("ColumnQualifier.parts must be strings")
+
+    @property
+    def is_empty(self) -> bool:
+        return len(self.parts) == 0
+
+    def all_qualified_names(self, name: str) -> list[str]:
+        qualifier_parts = self.parts
+        qualifier_prefixes = [
+            ".".join(quote_name_without_upper_casing(x) for x in qualifier_parts[i:])
+            for i in range(len(qualifier_parts))
+        ]
+        return [f"{prefix}.{name}" for prefix in qualifier_prefixes]
+
+    def to_upper(self):
+        return ColumnQualifier(tuple(part.upper() for part in self.parts))
+
+    def matches(self, target: ColumnQualifier) -> bool:
+        if self.is_empty or target.is_empty:
+            return False
+        # If the column has fewer qualifiers than the target, it cannot match
+        if len(self.parts) < len(target.parts):
+            return False
+        return self.parts[-len(target.parts) :] == target.parts

snowflake/snowpark_connect/config.py

@@ -8,7 +8,7 @@ import re
 import sys
 from collections import defaultdict
 from copy import copy, deepcopy
-from typing import Any
+from typing import Any, Dict, Optional
 
 import jpype
 import pyspark.sql.connect.proto.base_pb2 as proto_base
@@ -17,11 +17,18 @@ from tzlocal import get_localzone_name
 from snowflake import snowpark
 from snowflake.snowpark._internal.analyzer.analyzer_utils import (
     quote_name_without_upper_casing,
+    unquote_if_quoted,
 )
 from snowflake.snowpark.exceptions import SnowparkSQLException
 from snowflake.snowpark.types import TimestampTimeZone, TimestampType
+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_support import set_integral_types_conversion
 from snowflake.snowpark_connect.utils.concurrent import SynchronizedDict
-from snowflake.snowpark_connect.utils.context import get_session_id
+from snowflake.snowpark_connect.utils.context import (
+    get_jpype_jclass_lock,
+    get_spark_session_id,
+)
 from snowflake.snowpark_connect.utils.external_udxf_cache import (
     clear_external_udxf_cache,
 )
@@ -139,9 +146,30 @@ class GlobalConfig:
         "spark.sql.parser.quotedRegexColumnNames": "false",
         # custom configs
         "snowpark.connect.version": ".".join(map(str, sas_version)),
+        "snowpark.connect.temporary.views.create_in_snowflake": "false",
         # Control whether repartition(n) on a DataFrame forces splitting into n files during writes
         # This matches spark behavior more closely, but introduces overhead.
         "snowflake.repartition.for.writes": "false",
+        "snowpark.connect.structured_types.fix": "true",
+        # Local relation optimization: Use List[Row] for small data, PyArrow for large data
+        # Enabled in production by default to improve performance for createDataFrame on small local relations.
+        # Disabled in tests by default unless explicitly enabled to stabilize flaky tests that are not applying row ordering.
+        # SNOW-2719980: Remove this flag after test fragility issues are resolved
+        "snowpark.connect.localRelation.optimizeSmallData": "true",
+        "spark.sql.execution.arrow.maxRecordsPerBatch": "10000",  # TODO: no-op
+        # USE_VECTORIZED_SCANNER will become the default in a future BCR; Snowflake recommends setting it to TRUE for new workloads.
+        # This significantly reduces latency for loading Parquet files by downloading only relevant columnar sections into memory.
+        "snowpark.connect.parquet.useVectorizedScanner": "true",
+        # USE_LOGICAL_TYPE enables proper handling of Parquet logical types (TIMESTAMP, DATE, DECIMAL).
+        # Without useLogicalType set to "true", Parquet TIMESTAMP (INT64 physical) is incorrectly read as NUMBER(38,0).
+        "snowpark.connect.parquet.useLogicalType": "false",
+        "spark.sql.legacy.dataset.nameNonStructGroupingKeyAsValue": "false",
+        "spark.sql.parquet.outputTimestampType": "TIMESTAMP_MILLIS",
+        "snowpark.connect.handleIntegralOverflow": "false",
+        "snowpark.connect.scala.version": "2.12",
+        # Control whether to convert decimal - to integral types and vice versa: DecimalType(p,0) <-> ByteType/ShortType/IntegerType/LongType
+        # Values: "client_default" (behavior based on client type), "enabled", "disabled"
+        "snowpark.connect.integralTypesEmulation": "client_default",
     }
 
     boolean_config_list = [
@@ -150,11 +178,16 @@ class GlobalConfig:
         "spark.sql.repl.eagerEval.enabled",
         "spark.sql.crossJoin.enabled",
         "spark.sql.caseSensitive",
+        "snowpark.connect.localRelation.optimizeSmallData",
+        "snowpark.connect.parquet.useVectorizedScanner",
+        "snowpark.connect.parquet.useLogicalType",
         "spark.sql.ansi.enabled",
         "spark.sql.legacy.allowHashOnMapType",
         "spark.Catalog.databaseFilterInformationSchema",
         "spark.sql.parser.quotedRegexColumnNames",
         "snowflake.repartition.for.writes",
+        "spark.sql.legacy.dataset.nameNonStructGroupingKeyAsValue",
+        "snowpark.connect.handleIntegralOverflow",
     ]
 
     int_config_list = [
@@ -171,8 +204,15 @@ class GlobalConfig:
         "spark.app.name": lambda session, name: setattr(
             session, "query_tag", f"Spark-Connect-App-Name={name}"
         ),
+        # TODO SNOW-2896871: Remove with version 1.10.0
         "snowpark.connect.udf.imports": lambda session, imports: parse_imports(
-            session, imports
+            session, imports, "python"
+        ),
+        "snowpark.connect.udf.python.imports": lambda session, imports: parse_imports(
+            session, imports, "python"
+        ),
+        "snowpark.connect.udf.java.imports": lambda session, imports: parse_imports(
+            session, imports, "java"
         ),
     }
 
@@ -257,21 +297,34 @@ SESSION_CONFIG_KEY_WHITELIST = {
     "spark.sql.execution.pythonUDTF.arrow.enabled",
     "spark.sql.tvf.allowMultipleTableArguments.enabled",
     "snowpark.connect.sql.passthrough",
+    "snowpark.connect.cte.optimization_enabled",
     "snowpark.connect.iceberg.external_volume",
     "snowpark.connect.sql.identifiers.auto-uppercase",
+    "snowpark.connect.sql.partition.external_table_location",
     "snowpark.connect.udtf.compatibility_mode",
     "snowpark.connect.views.duplicate_column_names_handling_mode",
-    "enable_snowflake_extension_behavior",
+    "snowpark.connect.temporary.views.create_in_snowflake",
+    "snowpark.connect.enable_snowflake_extension_behavior",
+    "spark.hadoop.fs.s3a.server-side-encryption.key",
+    "spark.hadoop.fs.s3a.assumed.role.arn",
+    "snowpark.connect.describe_cache_ttl_seconds",
+    "mapreduce.fileoutputcommitter.marksuccessfuljobs",
+    "spark.sql.parquet.enable.summary-metadata",
+    "parquet.enable.summary-metadata",
 }
-AZURE_SAS_KEY = re.compile(
+AZURE_ACCOUNT_KEY = re.compile(
     r"^fs\.azure\.sas\.[^\.]+\.[^\.]+\.blob\.core\.windows\.net$"
 )
+AZURE_SAS_KEY = re.compile(
+    r"^fs\.azure\.sas\.fixed\.token\.[^\.]+\.dfs\.core\.windows\.net$"
+)
 
 
 def valid_session_config_key(key: str):
     return (
         key in SESSION_CONFIG_KEY_WHITELIST  # AWS session keys
         or AZURE_SAS_KEY.match(key)  # Azure session keys
+        or AZURE_ACCOUNT_KEY.match(key)  # Azure account keys
     )
 
 
@@ -279,17 +332,23 @@ class SessionConfig:
     """This class contains the session configuration for the Spark Server."""
 
     default_session_config = {
-        "snowpark.connect.sql.identifiers.auto-uppercase": "all_except_columns",
         "snowpark.connect.sql.passthrough": "false",
+        "snowpark.connect.cte.optimization_enabled": "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",
+        "snowpark.connect.enable_snowflake_extension_behavior": "false",
+        "snowpark.connect.describe_cache_ttl_seconds": "300",
+        "snowpark.connect.sql.partition.external_table_location": None,
+        "mapreduce.fileoutputcommitter.marksuccessfuljobs": "false",
+        "spark.sql.parquet.enable.summary-metadata": "false",
+        "parquet.enable.summary-metadata": "false",
     }
 
     def __init__(self) -> None:
         self.config = deepcopy(self.default_session_config)
+        self.table_metadata: Dict[str, Dict[str, Any]] = {}
 
     def __getitem__(self, item: str) -> str:
         return self.get(item)
@@ -319,6 +378,11 @@ CONFIG_ALLOWED_VALUES: dict[str, tuple] = {
         "all",
         "none",
     ),
+    "snowpark.connect.integralTypesEmulation": (
+        "client_default",
+        "enabled",
+        "disabled",
+    ),
 }
 
 # Set some default configuration that are necessary for the driver.
@@ -344,9 +408,11 @@ def route_config_proto(
                 if not pair.HasField("value"):
                     from pyspark.errors import IllegalArgumentException
 
-                    raise IllegalArgumentException(
+                    exception = IllegalArgumentException(
                         f"Cannot set config '{pair.key}' to None"
                     )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_CONFIG_VALUE)
+                    raise exception
 
                 set_config_param(
                     config.session_id, pair.key, pair.value, snowpark_session
@@ -429,7 +495,11 @@ def route_config_proto(
                 pair.value = str(global_config.is_modifiable(key)).lower()
             return res
         case _:
-            raise SnowparkConnectNotImplementedError(f"Unexpected request {config}")
+            exception = SnowparkConnectNotImplementedError(
+                f"Unexpected request {config}"
+            )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def set_config_param(
@@ -469,19 +539,27 @@ def _verify_static_config_not_modified(key: str) -> None:
     # https://github.com/apache/spark/blob/v3.5.3/sql/core/src/main/scala/org/apache/spark/sql/RuntimeConfig.scala#L161
     # Spark does not allow to modify static configurations at runtime.
     if global_config.is_static_config(key) and global_config.is_set(key):
-        raise ValueError(f"Cannot modify the value of a static config: {key}")
+        exception = ValueError(f"Cannot modify the value of a static config: {key}")
+        attach_custom_error_code(exception, ErrorCodes.CONFIG_CHANGE_NOT_ALLOWED)
+        raise exception
 
 
 def _verify_is_valid_config_value(key: str, value: Any) -> None:
     if key in CONFIG_ALLOWED_VALUES and value not in CONFIG_ALLOWED_VALUES[key]:
-        raise ValueError(
+        exception = ValueError(
             f"Invalid value '{value}' for key '{key}'. Allowed values: {', '.join(CONFIG_ALLOWED_VALUES[key])}."
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_CONFIG_VALUE)
+        raise exception
 
 
 def _verify_is_not_readonly_config(key):
     if key in global_config.readonly_config_list:
-        raise ValueError(f"Config with key {key} is read-only and cannot be modified.")
+        exception = ValueError(
+            f"Config with key {key} is read-only and cannot be modified."
+        )
+        attach_custom_error_code(exception, ErrorCodes.CONFIG_CHANGE_NOT_ALLOWED)
+        raise exception
 
 
 def set_jvm_timezone(timezone_id: str):
@@ -498,10 +576,13 @@ def set_jvm_timezone(timezone_id: str):
         RuntimeError: If JVM is not started
     """
     if not jpype.isJVMStarted():
-        raise RuntimeError("JVM must be started before setting timezone")
+        exception = RuntimeError("JVM must be started before setting timezone")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     try:
-        TimeZone = jpype.JClass("java.util.TimeZone")
+        with get_jpype_jclass_lock():
+            TimeZone = jpype.JClass("java.util.TimeZone")
         new_timezone = TimeZone.getTimeZone(timezone_id)
         TimeZone.setDefault(new_timezone)
 
@@ -513,7 +594,9 @@ def set_jvm_timezone(timezone_id: str):
 def reset_jvm_timezone_to_system_default():
     """Reset JVM timezone to the system's default timezone"""
     if not jpype.isJVMStarted():
-        raise RuntimeError("JVM must be started first")
+        exception = RuntimeError("JVM must be started first")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     try:
         TimeZone = jpype.JClass("java.util.TimeZone")
@@ -522,9 +605,13 @@ def reset_jvm_timezone_to_system_default():
             f"Reset JVM timezone to system default: {TimeZone.getDefault().getID()}"
         )
     except jpype.JException as e:
-        raise RuntimeError(f"Java exception while resetting timezone: {e}")
+        exception = RuntimeError(f"Java exception while resetting timezone: {e}")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
     except Exception as e:
-        raise RuntimeError(f"Unexpected error resetting JVM timezone: {e}")
+        exception = RuntimeError(f"Unexpected error resetting JVM timezone: {e}")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
 
 def set_snowflake_parameters(
@@ -569,38 +656,137 @@ def set_snowflake_parameters(
                         snowpark_session.use_database(db)
                     case (prev, curr) if prev != curr:
                         snowpark_session.use_schema(prev)
+        case "snowpark.connect.cte.optimization_enabled":
+            # Set CTE optimization on the snowpark session
+            cte_enabled = str_to_bool(value)
+            snowpark_session.cte_optimization_enabled = cte_enabled
+            logger.info(f"Updated snowpark session CTE optimization: {cte_enabled}")
+        case "snowpark.connect.structured_types.fix":
+            # TODO: SNOW-2367714 Remove this once the fix is automatically enabled in Snowpark
+            snowpark.context._enable_fix_2360274 = str_to_bool(value)
+            logger.info(f"Updated snowpark session structured types fix: {value}")
+        case "spark.sql.parquet.outputTimestampType":
+            if value == "TIMESTAMP_MICROS":
+                snowpark_session.sql(
+                    "ALTER SESSION SET UNLOAD_PARQUET_TIME_TIMESTAMP_MILLIS = false"
+                ).collect()
+            else:
+                # Default: TIMESTAMP_MILLIS (or any other value)
+                snowpark_session.sql(
+                    "ALTER SESSION SET UNLOAD_PARQUET_TIME_TIMESTAMP_MILLIS = true"
+                ).collect()
+            logger.info(f"Updated parquet timestamp output type to: {value}")
+        case "snowpark.connect.scala.version":
+            # force java udf helper recreation
+            set_java_udf_creator_initialized_state(False)
+        case "snowpark.connect.integralTypesEmulation":
+            # "client_default" - don't change, let set_spark_version handle it
+            # "enabled" / "disabled" - explicitly set
+            if value.lower() == "enabled":
+                set_integral_types_conversion(True)
+            elif value.lower() == "disabled":
+                set_integral_types_conversion(False)
         case _:
             pass
 
 
 def get_boolean_session_config_param(name: str) -> bool:
-    session_config = sessions_config[get_session_id()]
+    session_config = sessions_config[get_spark_session_id()]
     return str_to_bool(session_config[name])
 
 
+def get_string_session_config_param(name: str) -> str:
+    session_config = sessions_config[get_spark_session_id()]
+    return str(session_config[name])
+
+
+def get_cte_optimization_enabled() -> bool:
+    """Get the CTE optimization configuration setting."""
+    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_parquet_metadata_generation_enabled() -> bool:
+    """
+    Get the Parquet metadata file generation configuration setting.
+    """
+    return get_boolean_session_config_param(
+        "spark.sql.parquet.enable.summary-metadata"
+    ) or get_boolean_session_config_param("parquet.enable.summary-metadata")
+
+
+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_spark_session_id()]
+    default_ttl: str = SessionConfig.default_session_config[
+        "snowpark.connect.describe_cache_ttl_seconds"
+    ]
+    try:
+        ttl_str = session_config.get(
+            "snowpark.connect.describe_cache_ttl_seconds", default_ttl
+        )
+        return int(ttl_str)
+    except ValueError:  # fallback to default ttl
+        return int(default_ttl)
+
+
+def should_create_temporary_view_in_snowflake() -> bool:
+    return str_to_bool(
+        global_config["snowpark.connect.temporary.views.create_in_snowflake"]
+    )
+
+
 def auto_uppercase_column_identifiers() -> bool:
-    session_config = sessions_config[get_session_id()]
-    return session_config[
+    session_config = sessions_config[get_spark_session_id()]
+    auto_upper_case_config = session_config[
         "snowpark.connect.sql.identifiers.auto-uppercase"
-    ].lower() in ("all", "only_columns")
+    ]
+    if auto_upper_case_config:
+        return auto_upper_case_config.lower() in ("all", "only_columns")
+
+    return not global_config.spark_sql_caseSensitive
 
 
 def auto_uppercase_non_column_identifiers() -> bool:
-    session_config = sessions_config[get_session_id()]
-    return session_config[
+    session_config = sessions_config[get_spark_session_id()]
+    auto_upper_case_config = session_config[
         "snowpark.connect.sql.identifiers.auto-uppercase"
-    ].lower() in ("all", "all_except_columns")
+    ]
+    if auto_upper_case_config:
+        return auto_upper_case_config.lower() in ("all", "all_except_columns")
+
+    return not global_config.spark_sql_caseSensitive
 
 
-def parse_imports(session: snowpark.Session, imports: str | None) -> None:
+def external_table_location() -> Optional[str]:
+    session_config = sessions_config[get_spark_session_id()]
+    return session_config.get(
+        "snowpark.connect.sql.partition.external_table_location", None
+    )
+
+
+def parse_imports(
+    session: snowpark.Session, imports: str | None, language: str
+) -> None:
     if not imports:
         return
 
     # UDF needs to be recreated to include new imports
     clear_external_udxf_cache(session)
+    if language == "java":
+
+        set_java_udf_creator_initialized_state(False)
 
     for udf_import in imports.strip("[] ").split(","):
-        session.add_import(udf_import)
+        udf_import = udf_import.strip()
+        if udf_import:
+            session.add_import(udf_import)
 
 
 def get_timestamp_type():
@@ -613,3 +799,100 @@ def get_timestamp_type():
             # shouldn't happen since `spark.sql.timestampType` is always defined, and `spark.conf.unset` sets it to default (TIMESTAMP_LTZ)
             timestamp_type = TimestampType(TimestampTimeZone.LTZ)
     return timestamp_type
+
+
+def record_table_metadata(
+    table_identifier: str,
+    table_type: str,
+    data_source: str,
+    supports_column_rename: bool = True,
+) -> None:
+    """
+    Record metadata about a table for Spark compatibility checks.
+
+    Args:
+        table_identifier: Full table identifier (catalog.database.table)
+        table_type: "v1" or "v2"
+        data_source: Source format (parquet, csv, iceberg, etc.)
+        supports_column_rename: Whether the table supports RENAME COLUMN
+    """
+    session_id = get_spark_session_id()
+    session_config = sessions_config[session_id]
+
+    # Normalize table identifier for consistent lookup
+    # Use the full catalog.database.table identifier to avoid conflicts
+    normalized_identifier = table_identifier.upper().strip('"')
+
+    session_config.table_metadata[normalized_identifier] = {
+        "table_type": table_type,
+        "data_source": data_source,
+        "supports_column_rename": supports_column_rename,
+    }
+
+
+def get_table_metadata(table_identifier: str) -> Dict[str, Any] | None:
+    """
+    Get stored metadata for a table.
+
+    Args:
+        table_identifier: Full table identifier (catalog.database.table)
+
+    Returns:
+        Table metadata dict or None if not found
+    """
+    session_id = get_spark_session_id()
+    session_config = sessions_config[session_id]
+
+    normalized_identifier = unquote_if_quoted(table_identifier).upper()
+
+    return session_config.table_metadata.get(normalized_identifier)
+
+
+def check_table_supports_operation(table_identifier: str, operation: str) -> bool:
+    """
+    Check if a table supports a given operation based on metadata and config.
+
+    Args:
+        table_identifier: Full table identifier (catalog.database.table)
+        operation: Operation to check (e.g., "rename_column")
+
+    Returns:
+        True if operation is supported, False if should be blocked
+    """
+    table_metadata = get_table_metadata(table_identifier)
+
+    if not table_metadata:
+        return True
+
+    session_id = get_spark_session_id()
+    session_config = sessions_config[session_id]
+    enable_extensions = str_to_bool(
+        session_config.get(
+            "snowpark.connect.enable_snowflake_extension_behavior", "false"
+        )
+    )
+
+    if enable_extensions:
+        return True
+
+    if operation == "rename_column":
+        return table_metadata.get("supports_column_rename", True)
+
+    return True
+
+
+def get_scala_version() -> str:
+    return global_config.get("snowpark.connect.scala.version")
+
+
+_java_udf_creator_initialized = False
+
+
+def is_java_udf_creator_initialized() -> bool:
+    global _java_udf_creator_initialized
+    return _java_udf_creator_initialized
+
+
+def set_java_udf_creator_initialized_state(value: bool) -> None:
+    global _java_udf_creator_initialized
+    _java_udf_creator_initialized = value

snowflake/snowpark_connect/constants.py

@@ -16,3 +16,5 @@ MAP_IN_ARROW_EVAL_TYPE = 207  # eval_type for mapInArrow operations
 COLUMN_METADATA_COLLISION_KEY = "{expr_id}_{key}"
 
 DUPLICATE_KEY_FOUND_ERROR_TEMPLATE = "Duplicate key found: {key}. You can set spark.sql.mapKeyDedupPolicy to LAST_WIN to deduplicate map keys with last wins policy."
+
+SPARK_VERSION = "3.5.3"