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

snowflake/snowpark_connect/execute_plan/map_execution_command.py

@@ -1,21 +1,13 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
-
 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._internal.analyzer.analyzer_utils import (
-    quote_name_without_upper_casing,
-)
 from snowflake.snowpark_connect.column_name_handler import ColumnNames
-from snowflake.snowpark_connect.config import (
-    auto_uppercase_ddl,
-    global_config,
-    sessions_config,
-)
+from snowflake.snowpark_connect.config import global_config, sessions_config
 from snowflake.snowpark_connect.constants import SERVER_SIDE_SESSION_ID
 from snowflake.snowpark_connect.execute_plan.utils import pandas_to_arrow_batches_bytes
 from snowflake.snowpark_connect.expression import map_udf
@@ -24,24 +16,23 @@ 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.utils.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
 
 
-def _spark_to_snowflake_single_id(name: str) -> str:
-    name = quote_name_without_upper_casing(name)
-    return name.upper() if auto_uppercase_ddl() else name
-
-
 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_snowflake_single_id(col.spark_name)
+            col.snowpark_name: spark_to_sf_single_id(col.spark_name, is_column=True)
             for col in columns
         }
 
@@ -64,7 +55,7 @@ def _create_column_rename_map(
 
     if len(renamed_cols) == 0:
         return {
-            col.snowpark_name: _spark_to_snowflake_single_id(col.spark_name)
+            col.snowpark_name: spark_to_sf_single_id(col.spark_name, is_column=True)
             for col in not_renamed_cols
         }
 
@@ -95,12 +86,9 @@ def map_execution_command(
     match request.plan.command.WhichOneof("command_type"):
         case "create_dataframe_view":
             req = request.plan.command.create_dataframe_view
-            input_df = map_relation(req.input)
-            # Use real column names when writing to sf.
-            assert hasattr(
-                input_df, "_column_map"
-            ), "input_df does not have the _column_map attribute"
-            column_map = input_df._column_map
+            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[
@@ -133,7 +121,9 @@ def map_execution_command(
                 view_name = [global_config.spark_sql_globalTempDatabase, req.name]
             else:
                 view_name = [req.name]
-            view_name = [_spark_to_snowflake_single_id(part) for part in view_name]
+            view_name = [
+                spark_to_sf_single_id_with_unquoting(part) for part in view_name
+            ]
 
             if req.replace:
                 input_df.create_or_replace_temp_view(view_name)

snowflake/snowpark_connect/execute_plan/map_execution_root.py

@@ -20,6 +20,7 @@ from snowflake.snowpark._internal.utils import (
     create_or_update_statement_params_with_query_tag,
 )
 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 (
     arrow_table_to_arrow_bytes,
     pandas_to_arrow_batches_bytes,
@@ -89,13 +90,16 @@ def to_arrow_batch_iter(result_df: snowpark.DataFrame) -> Iterator[Table]:
 def map_execution_root(
     request: proto_base.ExecutePlanRequest,
 ) -> Iterator[proto_base.ExecutePlanResponse | QueryResult]:
-    result_df: snowpark.DataFrame | pandas.DataFrame = map_relation(request.plan.root)
+    result: DataFrameContainer | pandas.DataFrame = map_relation(request.plan.root)
+    if isinstance(result, pandas.DataFrame):
+        result_df = result
+    else:
+        result_df = result.dataframe
+
     if isinstance(result_df, snowpark.DataFrame):
         snowpark_schema = result_df.schema
-        schema = snowpark_to_proto_type(
-            snowpark_schema, result_df._column_map, result_df
-        )
-        spark_columns = result_df._column_map.get_spark_columns()
+        schema = snowpark_to_proto_type(snowpark_schema, result.column_map, result_df)
+        spark_columns = result.column_map.get_spark_columns()
         if tcm.TCM_MODE:
             # TCM result handling:
             # - small result (only one batch): just return the executePlanResponse

snowflake/snowpark_connect/expression/map_extension.py

@@ -58,7 +58,8 @@ def map_extension(
             from snowflake.snowpark_connect.relation.map_relation import map_relation
 
             with push_evaluating_sql_scope():
-                df = map_relation(extension.subquery_expression.input)
+                df_container = map_relation(extension.subquery_expression.input)
+                df = df_container.dataframe
 
             queries = df.queries["queries"]
             if len(queries) != 1:

snowflake/snowpark_connect/expression/map_udf.py

@@ -13,10 +13,7 @@ from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.type_mapping import proto_to_snowpark_type
 from snowflake.snowpark_connect.typed_column import TypedColumn
-from snowflake.snowpark_connect.utils.session import (
-    get_or_create_snowpark_session,
-    get_python_udxf_import_files,
-)
+from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
 from snowflake.snowpark_connect.utils.udf_helper import (
     SnowparkUDF,
     gen_input_types,
@@ -28,6 +25,9 @@ from snowflake.snowpark_connect.utils.udf_helper import (
 from snowflake.snowpark_connect.utils.udf_utils import (
     ProcessCommonInlineUserDefinedFunction,
 )
+from snowflake.snowpark_connect.utils.udxf_import_utils import (
+    get_python_udxf_import_files,
+)
 
 
 def process_udf_return_type(

snowflake/snowpark_connect/expression/map_unresolved_attribute.py

@@ -17,15 +17,15 @@ from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
 from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.typed_column import TypedColumn
-from snowflake.snowpark_connect.utils.attribute_handling import (
-    split_fully_qualified_spark_name,
-)
 from snowflake.snowpark_connect.utils.context import (
     get_is_evaluating_sql,
     get_outer_dataframes,
     get_plan_id_map,
     resolve_lca_alias,
 )
+from snowflake.snowpark_connect.utils.identifiers import (
+    split_fully_qualified_spark_name,
+)
 
 SPARK_QUOTED = re.compile("^(`.*`)$", re.DOTALL)
 
@@ -46,11 +46,12 @@ def map_unresolved_attribute(
 
     if has_plan_id:
         plan_id = exp.unresolved_attribute.plan_id
-        target_df = get_plan_id_map(plan_id)
+        target_df_container = get_plan_id_map(plan_id)
+        target_df = target_df_container.dataframe
         assert (
             target_df is not None
         ), f"resolving an attribute of a unresolved dataframe {plan_id}"
-        column_mapping = target_df._column_map
+        column_mapping = target_df_container.column_map
         typer = ExpressionTyper(target_df)
 
     def get_col(snowpark_name):
@@ -146,8 +147,8 @@ def map_unresolved_attribute(
             name_parts[0], allow_non_exists=True
         )
         if snowpark_name is None:
-            for outer_df in get_outer_dataframes():
-                snowpark_name = outer_df._column_map.get_snowpark_column_name_from_spark_column_name(
+            for outer_df_container in get_outer_dataframes():
+                snowpark_name = outer_df_container.column_map.get_snowpark_column_name_from_spark_column_name(
                     name_parts[0], allow_non_exists=True
                 )
                 if snowpark_name is not None: