snowpark-connect 0.28.0__py3-none-any.whl → 0.29.0__py3-none-any.whl

snowflake/snowpark_connect/config.py

@@ -264,16 +264,22 @@ SESSION_CONFIG_KEY_WHITELIST = {
     "snowpark.connect.udtf.compatibility_mode",
     "snowpark.connect.views.duplicate_column_names_handling_mode",
     "enable_snowflake_extension_behavior",
+    "spark.hadoop.fs.s3a.server-side-encryption.key",
+    "spark.hadoop.fs.s3a.assumed.role.arn",
 }
-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
     )
 
 
@@ -283,7 +289,7 @@ class SessionConfig:
     default_session_config = {
         "snowpark.connect.sql.identifiers.auto-uppercase": "all_except_columns",
         "snowpark.connect.sql.passthrough": "false",
-        "snowpark.connect.cte.optimization_enabled": "true",
+        "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",
@@ -578,7 +584,10 @@ def set_snowflake_parameters(
             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 _:
             pass
 

snowflake/snowpark_connect/execute_plan/map_execution_command.py

@@ -1,18 +1,16 @@
 #
 # 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.types import StructField, StructType
 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.dataframe_container import DataFrameContainer
 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
@@ -28,10 +26,7 @@ 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)
+from snowflake.snowpark_connect.utils.temporary_view_cache import register_temp_view
 
 
 def _create_column_rename_map(
@@ -98,32 +93,35 @@ def map_execution_command(
             input_df = input_df_container.dataframe
             column_map = input_df_container.column_map
 
+            # TODO: Remove code handling deduplication. When view are not materialized we don't have to care about it.
             session_config = sessions_config[get_session_id()]
             duplicate_column_names_handling_mode = session_config[
                 "snowpark.connect.views.duplicate_column_names_handling_mode"
             ]
 
+            spark_columns = input_df_container.column_map.get_spark_columns()
             # 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)
-                )
+                rename_map = _create_column_rename_map(column_map.columns, True)
+                snowpark_columns = list(rename_map.values())
+                input_df = input_df.rename(rename_map)
             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
                 )
+                rename_map = _create_column_rename_map(remaining_columns, False)
+                snowpark_columns = list(rename_map.values())
+                spark_columns = list(dict.fromkeys(spark_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)
-                )
+                input_df = input_df.rename(rename_map)
             else:
                 # rename columns without deduplication
-                input_df = input_df.rename(
-                    _create_column_rename_map(column_map.columns, False)
-                )
+                rename_map = _create_column_rename_map(column_map.columns, True)
+                snowpark_columns = list(rename_map.values())
+                input_df = input_df.rename(rename_map)
 
             if req.is_global:
                 view_name = [global_config.spark_sql_globalTempDatabase, req.name]
@@ -132,18 +130,23 @@ def map_execution_command(
             view_name = [
                 spark_to_sf_single_id_with_unquoting(part) for part in view_name
             ]
+            joined_view_name = ".".join(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)
+            schema = StructType(
+                [
+                    StructField(field.name, field.datatype)
+                    for field in input_df.schema.fields
+                ]
+            )
+            input_df_container = DataFrameContainer.create_with_column_mapping(
+                dataframe=input_df,
+                spark_column_names=spark_columns,
+                snowpark_column_names=snowpark_columns,
+                parent_column_name_map=input_df_container.column_map,
+                cached_schema_getter=lambda: schema,
+            )
+
+            register_temp_view(joined_view_name, input_df_container, req.replace)
         case "write_stream_operation_start":
             match request.plan.command.write_stream_operation_start.format:
                 case "console":
@@ -207,43 +210,3 @@ def map_execution_command(
             raise 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()