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

snowflake/snowpark_connect/relation/read/map_read_text.py

@@ -7,7 +7,7 @@ import typing
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
 from snowflake import snowpark
-from snowflake.snowpark_connect.column_name_handler import with_column_map
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.read.utils import (
     get_spark_column_names_from_snowpark_columns,
     rename_columns_as_snowflake_standard,
@@ -71,7 +71,7 @@ def map_read_text(
     schema: snowpark.types.StructType | None,
     session: snowpark.Session,
     paths: list[str],
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Read a TEXT file into a Snowpark DataFrame.
     """
@@ -98,9 +98,9 @@ def map_read_text(
     renamed_df, snowpark_column_names = rename_columns_as_snowflake_standard(
         df, rel.common.plan_id
     )
-    return with_column_map(
-        renamed_df,
-        spark_column_names,
-        snowpark_column_names,
-        [f.datatype for f in df.schema.fields],
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=renamed_df,
+        spark_column_names=spark_column_names,
+        snowpark_column_names=snowpark_column_names,
+        snowpark_column_types=[f.datatype for f in df.schema.fields],
     )

snowflake/snowpark_connect/relation/utils.py

@@ -32,6 +32,7 @@ from snowflake.snowpark_connect.column_name_handler import (
     ColumnNameMap,
     make_column_names_snowpark_compatible,
 )
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 
 TYPE_MAP_FOR_TO_SCHEMA = {
@@ -91,7 +92,9 @@ TYPE_MAP_FOR_TO_SCHEMA = {
 
 
 def get_df_with_partition_row_number(
-    df: snowpark.DataFrame, plan_id: int | None, row_number_column_name: str
+    container: DataFrameContainer,
+    plan_id: int | None,
+    row_number_column_name: str,
 ) -> snowpark.DataFrame:
     """
     Add a row number for each row in each partition for the given df, where
@@ -106,21 +109,24 @@ def get_df_with_partition_row_number(
     |  c|  4|                               |  c|  4|  0         |
     +---+---+                               +---+---+------------+
     """
+    df = container.dataframe
+    column_map = container.column_map
+
     row_number_snowpark_column_name = make_column_names_snowpark_compatible(
-        [row_number_column_name], plan_id, len(df._column_map.get_spark_columns())
+        [row_number_column_name], plan_id, len(column_map.get_spark_columns())
     )[0]
     row_number_snowpark_column = (
         snowpark_fn.row_number()
         .over(
             snowpark.window.Window.partition_by(
-                *df._column_map.get_snowpark_columns()
+                *column_map.get_snowpark_columns()
             ).order_by(snowpark_fn.lit(1))
         )
         .alias(row_number_snowpark_column_name)
     )
 
     df_with_partition_number = df.select(
-        *df._column_map.get_snowpark_columns(), row_number_snowpark_column
+        *column_map.get_snowpark_columns(), row_number_snowpark_column
     )
     return df_with_partition_number
 
@@ -197,7 +203,7 @@ def get_semantic_string(rel: relation_proto.Relation) -> str:
     """
     queries = [
         query
-        for query_list in map_relation(rel)._plan.execution_queries.values()
+        for query_list in map_relation(rel).dataframe._plan.execution_queries.values()
         for query in query_list
     ]
 

snowflake/snowpark_connect/relation/write/jdbc_write_dbapi.py

@@ -10,6 +10,7 @@ import snowflake.snowpark
 from snowflake import snowpark
 from snowflake.snowpark import DataFrameWriter
 from snowflake.snowpark.dataframe import DataFrame
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.read import jdbc_read_dbapi
 from snowflake.snowpark_connect.relation.read.jdbc_read_dbapi import JdbcDialect
 from snowflake.snowpark_connect.relation.read.utils import Connection
@@ -36,7 +37,7 @@ class JdbcDataFrameWriter(DataFrameWriter):
 
     def jdbc_write_dbapi(
         self,
-        input_df: DataFrame,
+        container: DataFrameContainer,
         create_connection: Callable[[dict[str, str]], "Connection"],
         close_connection: Callable[[Connection], None],
         table: str,
@@ -46,6 +47,7 @@ class JdbcDataFrameWriter(DataFrameWriter):
         Write a Snowpark Dataframe data into table of a JDBC datasource.
         """
 
+        input_df = container.dataframe
         conn = create_connection(self.jdbc_options)
         try:
             url = self.jdbc_options.get("url", None)
@@ -53,32 +55,32 @@ class JdbcDataFrameWriter(DataFrameWriter):
 
             table_exist = self._does_table_exist(conn, table)
             insert_query = self._generate_insert_query(
-                input_df,
+                container,
                 table,
             )
 
             match write_mode:
                 case "append":
                     if not table_exist:
-                        self._create_table(conn, table, input_df, jdbc_dialect)
+                        self._create_table(conn, table, container, jdbc_dialect)
                 case "errorifexists":
                     if table_exist:
                         raise ValueError(
                             "table is already exist and write mode is ERROR_IF_EXISTS"
                         )
                     else:
-                        self._create_table(conn, table, input_df, jdbc_dialect)
+                        self._create_table(conn, table, container, jdbc_dialect)
                 case "overwrite":
                     if table_exist:
                         self._drop_table(conn, table)
-                    self._create_table(conn, table, input_df, jdbc_dialect)
+                    self._create_table(conn, table, container, jdbc_dialect)
                 case "ignore":
                     if table_exist:
                         # With Ignore write mode, if table already exists, the save operation is expected
                         # to not save the contents of the DataFrame and to not change the existing data.
                         return
                     else:
-                        self._create_table(conn, table, input_df, jdbc_dialect)
+                        self._create_table(conn, table, container, jdbc_dialect)
                 case _:
                     raise ValueError(f"Invalid write mode value{write_mode}")
 
@@ -92,14 +94,14 @@ class JdbcDataFrameWriter(DataFrameWriter):
         finally:
             close_connection(conn)
 
-    def _generate_insert_query(self, input_df: DataFrame, table: str) -> str:
+    def _generate_insert_query(self, container: DataFrameContainer, table: str) -> str:
         """
         Generates INSERT statement with placeholders.
-        :param input_df: Snowpark dataframe to save
+        :param container: Snowpark dataframe container
         :param table: JDBC datasource table name
         :return: INSERT SQL statement
         """
-        true_names = input_df._column_map.get_spark_columns()
+        true_names = container.column_map.get_spark_columns()
         # quote each column name to match PySpark's case-sensitive column naming behavior.
         quoted_column_names = ",".join([f'"{col}"' for col in true_names])
         place_holders = ",".join(["?"] * len(true_names))
@@ -145,7 +147,7 @@ class JdbcDataFrameWriter(DataFrameWriter):
         self,
         conn: Connection,
         table: str,
-        input_df: DataFrame,
+        container,
         jdbc_dialect: JdbcDialect,
     ) -> None:
         """
@@ -154,14 +156,15 @@ class JdbcDataFrameWriter(DataFrameWriter):
 
         :param conn: A Python DBAPI connection over JDBC connection
         :param table: DBC datasource table name
-        :param input_df: Snowpark dataframe to save
+        :param container: Snowpark dataframe container
         :param jdbc_dialect: JDBC specific dialect
         :return: None
         """
+        input_df = container.dataframe
         columns_str = ""
         fields = input_df.schema.fields
         total_columns = len(fields)
-        column_map = input_df._column_map
+        column_map = container.column_map
 
         column_index = 0
         for field in fields:

snowflake/snowpark_connect/relation/write/map_write.py

@@ -9,20 +9,30 @@ from pathlib import Path
 import pyspark.sql.connect.proto.base_pb2 as proto_base
 import pyspark.sql.connect.proto.commands_pb2 as commands_proto
 from pyspark.errors.exceptions.base import AnalysisException
-from pyspark.sql.connect.types import StructType
 
 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.functions import col, lit, object_construct
+from snowflake.snowpark.types import (
+    ArrayType,
+    DataType,
+    DateType,
+    MapType,
+    StringType,
+    StructType,
+    TimestampType,
+    _NumericType,
+)
 from snowflake.snowpark_connect.config import (
-    auto_uppercase_ddl,
     global_config,
     sessions_config,
     str_to_bool,
 )
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.io_utils import (
     convert_file_prefix_path,
     is_cloud_path,
@@ -32,16 +42,19 @@ from snowflake.snowpark_connect.relation.read.reader_config import CsvWriterConf
 from snowflake.snowpark_connect.relation.stage_locator import get_paths_from_stage
 from snowflake.snowpark_connect.relation.utils import random_string
 from snowflake.snowpark_connect.type_mapping import snowpark_to_iceberg_type
-from snowflake.snowpark_connect.utils.attribute_handling import (
+from snowflake.snowpark_connect.utils.context import get_session_id
+from snowflake.snowpark_connect.utils.identifiers import (
+    spark_to_sf_single_id,
     split_fully_qualified_spark_name,
 )
-from snowflake.snowpark_connect.utils.context import get_session_id
 from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
     telemetry,
 )
 
+_column_order_for_write = "name"
+
 
 # TODO: We will revise/refactor this after changes for all formats are finalized.
 def clean_params(params):
@@ -85,14 +98,9 @@ def get_param_from_options(params, options, source):
         params["format_type_options"]["NULL_IF"] = options["nullValue"]
 
 
-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 _spark_to_snowflake(multipart_id: str) -> str:
     return ".".join(
-        _spark_to_snowflake_single_id(part)
+        spark_to_sf_single_id(part)
         for part in split_fully_qualified_spark_name(multipart_id)
     )
 
@@ -115,9 +123,8 @@ def map_write(request: proto_base.ExecutePlanRequest):
         case commands_proto.WriteOperation.SaveMode.SAVE_MODE_IGNORE:
             write_mode = "ignore"
 
-    input_df: snowpark.DataFrame = handle_column_names(
-        map_relation(write_op.input), write_op.source
-    )
+    result = map_relation(write_op.input)
+    input_df: snowpark.DataFrame = handle_column_names(result, write_op.source)
     session: snowpark.Session = get_or_create_snowpark_session()
 
     # Snowflake saveAsTable doesn't support format
@@ -198,7 +205,7 @@ def map_write(request: proto_base.ExecutePlanRequest):
             options = dict(write_op.options)
             if write_mode is None:
                 write_mode = "errorifexists"
-            map_write_jdbc(input_df, session, options, write_mode)
+            map_write_jdbc(result, session, options, write_mode)
         case "iceberg":
             table_name = (
                 write_op.path
@@ -220,7 +227,14 @@ def map_write(request: proto_base.ExecutePlanRequest):
                     snowpark_session=session,
                 )
                 write_mode = "append"
-            input_df.write.saveAsTable(table_name=snowpark_table_name, mode=write_mode)
+
+            _validate_schema_and_get_writer(
+                input_df, write_mode, snowpark_table_name
+            ).saveAsTable(
+                table_name=snowpark_table_name,
+                mode=write_mode,
+                column_order=_column_order_for_write,
+            )
         case _:
             snowpark_table_name = _spark_to_snowflake(write_op.table.table_name)
 
@@ -228,17 +242,23 @@ def map_write(request: proto_base.ExecutePlanRequest):
                 write_op.table.save_method
                 == commands_proto.WriteOperation.SaveTable.TableSaveMethod.TABLE_SAVE_METHOD_SAVE_AS_TABLE
             ):
-                input_df.write.saveAsTable(
+                _validate_schema_and_get_writer(
+                    input_df, write_mode, snowpark_table_name
+                ).saveAsTable(
                     table_name=snowpark_table_name,
                     mode=write_mode,
+                    column_order=_column_order_for_write,
                 )
             elif (
                 write_op.table.save_method
                 == commands_proto.WriteOperation.SaveTable.TableSaveMethod.TABLE_SAVE_METHOD_INSERT_INTO
             ):
-                input_df.write.saveAsTable(
+                _validate_schema_and_get_writer(
+                    input_df, write_mode, snowpark_table_name
+                ).saveAsTable(
                     table_name=snowpark_table_name,
                     mode=write_mode or "append",
+                    column_order=_column_order_for_write,
                 )
             else:
                 raise SnowparkConnectNotImplementedError(
@@ -265,10 +285,8 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
             )
 
     snowpark_table_name = _spark_to_snowflake(write_op.table_name)
-
-    input_df: snowpark.DataFrame = handle_column_names(
-        map_relation(write_op.input), "table"
-    )
+    result = map_relation(write_op.input)
+    input_df: snowpark.DataFrame = handle_column_names(result, "table")
     session: snowpark.Session = get_or_create_snowpark_session()
 
     if write_op.table_name is None or write_op.table_name == "":
@@ -304,18 +322,163 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                 schema=input_df.schema,
                 snowpark_session=session,
             )
-
-        input_df.write.saveAsTable(
+        _validate_schema_and_get_writer(
+            input_df, write_mode, snowpark_table_name
+        ).saveAsTable(
             table_name=snowpark_table_name,
             mode="append",
+            column_order=_column_order_for_write,
         )
     else:
-        input_df.write.saveAsTable(
+        _validate_schema_and_get_writer(
+            input_df, write_mode, snowpark_table_name
+        ).saveAsTable(
             table_name=snowpark_table_name,
             mode=write_mode,
+            column_order=_column_order_for_write,
         )
 
 
+def _validate_schema_and_get_writer(
+    input_df: snowpark.DataFrame, write_mode: str, snowpark_table_name: str
+) -> snowpark.DataFrameWriter:
+    if write_mode == "overwrite":
+        return input_df.write
+
+    table_schema = None
+    try:
+        table_schema = (
+            get_or_create_snowpark_session().table(snowpark_table_name).schema
+        )
+    except SnowparkSQLException as e:
+        msg = e.message
+        if "SQL compilation error" in msg and "does not exist" in msg:
+            pass
+        else:
+            raise e
+
+    if table_schema is None:
+        # If table does not exist, we can skip the schema validation
+        return input_df.write
+
+    _validate_schema_for_append(table_schema, input_df.schema, snowpark_table_name)
+
+    # if table exists and case sensitivity is not enabled, we need to rename the columns to match existing table schema
+    if not global_config.spark_sql_caseSensitive:
+
+        for field in input_df.schema.fields:
+            # Find the matching field in the table schema (case-insensitive)
+            col_name = field.name
+            renamed = col_name
+            matching_field = next(
+                (f for f in table_schema.fields if f.name.lower() == col_name.lower()),
+                None,
+            )
+            if matching_field is not None and matching_field != col_name:
+                renamed = matching_field.name
+                input_df = input_df.withColumnRenamed(col_name, renamed)
+                # Cast column if type does not match
+
+            if field.datatype != matching_field.datatype:
+                if isinstance(matching_field.datatype, StructType):
+                    input_df = input_df.withColumn(
+                        renamed,
+                        col(renamed).cast(matching_field.datatype, rename_fields=True),
+                    )
+                else:
+                    input_df = input_df.withColumn(
+                        renamed, col(renamed).cast(matching_field.datatype)
+                    )
+    return input_df.write
+
+
+def _validate_schema_for_append(
+    table_schema: DataType, data_schema: DataType, snowpark_table_name: str
+):
+    match (table_schema, data_schema):
+        case (_, _) if table_schema == data_schema:
+            return
+
+        case (StructType() as table_struct, StructType() as data_struct):
+
+            def _comparable_col_name(col: str) -> str:
+                return col if global_config.spark_sql_caseSensitive else col.lower()
+
+            def invalid_struct_schema():
+                raise AnalysisException(
+                    f"Cannot resolve columns for the existing table {snowpark_table_name} ({table_schema.simple_string()}) with the data schema ({data_schema.simple_string()})."
+                )
+
+            if len(table_struct.fields) != len(data_struct.fields):
+                raise AnalysisException(
+                    f"The column number of the existing table {snowpark_table_name} ({table_schema.simple_string()}) doesn't match the data schema ({data_schema.simple_string()}).)"
+                )
+
+            table_field_names = {
+                _comparable_col_name(field.name) for field in table_struct.fields
+            }
+            data_field_names = {
+                _comparable_col_name(field.name) for field in data_struct.fields
+            }
+
+            if table_field_names != data_field_names:
+                invalid_struct_schema()
+
+            for data_field in data_struct.fields:
+                matching_table_field = next(
+                    (
+                        f
+                        for f in table_struct.fields
+                        if _comparable_col_name(f.name)
+                        == _comparable_col_name(data_field.name)
+                    ),
+                    None,
+                )
+
+                if matching_table_field is None:
+                    invalid_struct_schema()
+                else:
+                    _validate_schema_for_append(
+                        matching_table_field.datatype,
+                        data_field.datatype,
+                        snowpark_table_name,
+                    )
+
+            return
+
+        case (StringType(), _) if not isinstance(
+            data_schema, (StructType, ArrayType, MapType, TimestampType, DateType)
+        ):
+            return
+
+        case (_, _) if isinstance(table_schema, _NumericType) and isinstance(
+            data_schema, _NumericType
+        ):
+            return
+
+        case (ArrayType() as table_array, ArrayType() as data_array):
+            _validate_schema_for_append(
+                table_array.element_type, data_array.element_type, snowpark_table_name
+            )
+
+        case (MapType() as table_map, MapType() as data_map):
+            _validate_schema_for_append(
+                table_map.key_type, data_map.key_type, snowpark_table_name
+            )
+            _validate_schema_for_append(
+                table_map.value_type, data_map.value_type, snowpark_table_name
+            )
+
+        case (TimestampType(), _) if isinstance(data_schema, (DateType, TimestampType)):
+            return
+        case (DateType(), _) if isinstance(data_schema, (DateType, TimestampType)):
+            return
+        case (_, _):
+            raise AnalysisException(
+                f"[INCOMPATIBLE_DATA_FOR_TABLE.CANNOT_SAFELY_CAST] Cannot write incompatible data for the table {snowpark_table_name}: Cannot safely cast {data_schema.simple_string()} to {table_schema.simple_string()}"
+            )
+
+
 def create_iceberg_table(
     snowpark_table_name: str,
     location: str,
@@ -323,7 +486,7 @@ def create_iceberg_table(
     snowpark_session: snowpark.Session,
 ):
     table_schema = [
-        f"{_spark_to_snowflake_single_id(field.name)} {snowpark_to_iceberg_type(field.datatype)}"
+        f"{spark_to_sf_single_id(unquote_if_quoted(field.name), is_column = True)} {snowpark_to_iceberg_type(field.datatype)}"
         for field in schema.fields
     ]
 
@@ -374,26 +537,22 @@ def rewrite_df(input_df: snowpark.DataFrame, source: str) -> snowpark.DataFrame:
     return rewritten_df.select(object_construct(*construct_key_values))
 
 
-def handle_column_names(df: snowpark.DataFrame, source: str) -> snowpark.DataFrame:
+def handle_column_names(
+    container: DataFrameContainer, source: str
+) -> snowpark.DataFrame:
     """
-    Handle column names.
-
-    Quote column name in these scenarios:
-        0. Not write to table
-        1. Customer enabled case sensitivity in config
+    Handle column names before write so they match spark schema.
     """
-    if not hasattr(df, "_column_map") or source == "jdbc":
+    df = container.dataframe
+    if source == "jdbc":
         # don't change column names for jdbc sources as we directly use spark column names for writing to the destination tables.
         return df
-    column_map = df._column_map
-    case_sensitive = global_config.spark_sql_caseSensitive
-    for column in df.columns:
-        spark_column_name = unquote_if_quoted(
-            column_map.get_spark_column_name_from_snowpark_column_name(column)
+    column_map = container.column_map
+
+    for column in column_map.columns:
+        df = df.withColumnRenamed(
+            column.snowpark_name, quote_name_without_upper_casing(column.spark_name)
         )
-        if source in ("csv", "parquet", "json") or case_sensitive:
-            spark_column_name = f'"{spark_column_name}"'
-        df = df.withColumnRenamed(column, spark_column_name)
     return df
 
 

snowflake/snowpark_connect/relation/write/map_write_jdbc.py

@@ -3,6 +3,7 @@
 #
 
 from snowflake import snowpark
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.read.map_read_jdbc import (
     close_connection,
     create_connection,
@@ -14,7 +15,7 @@ from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 
 
 def map_write_jdbc(
-    input_df: snowpark.DataFrame,
+    container: DataFrameContainer,
     session: snowpark.Session,
     options: dict[str, str],
     write_mode: str,
@@ -38,7 +39,7 @@ def map_write_jdbc(
 
     try:
         JdbcDataFrameWriter(session, jdbc_options).jdbc_write_dbapi(
-            input_df,
+            container,
             create_connection,
             close_connection,
             table=dbtable,

snowflake/snowpark_connect/server.py

@@ -112,10 +112,38 @@ _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE = 128 * 1024 * 1024
 _SPARK_CONNECT_GRPC_MAX_METADATA_SIZE = 64 * 1024  # 64kb
 
 
+def _sanitize_file_paths(text: str) -> str:
+    """
+    Sanitize file paths in error messages by replacing them with placeholders.
+    Only matches actual file paths, not module names or class names.
+    """
+    import re
+
+    # Pattern to match file paths in traceback "File" lines only
+    # This targets the specific format: File "/path/to/file.py", line XX
+    file_line_pattern = r'(File\s+["\'])([^"\']+)(["\'],\s+line\s+\d+)'
+
+    def replace_file_path(match):
+        return f"{match.group(1)}<redacted_file_path>{match.group(3)}"
+
+    return re.sub(file_line_pattern, replace_file_path, text)
+
+
 def _handle_exception(context, e: Exception):
     import traceback
 
-    traceback.print_exc()
+    # traceback.print_exc()
+    # SNOWFLAKE_SHOW_ERROR_TRACE controls sanitized traceback printing (default: false)
+    show_traceback = os.getenv("SNOWFLAKE_SHOW_ERROR_TRACE", "false").lower() == "true"
+
+    if show_traceback:
+        # Show detailed traceback (includes error info naturally)
+        error_traceback = traceback.format_exc()
+        sanitized_traceback = _sanitize_file_paths(error_traceback)
+        logger.error(sanitized_traceback)
+    else:
+        # Show only basic error information, no traceback
+        logger.error("Error: %s - %s", type(e).__name__, str(e))
 
     telemetry.report_request_failure(e)
 
@@ -195,12 +223,13 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
             telemetry.initialize_request_summary(request)
             match request.WhichOneof("analyze"):
                 case "schema":
-                    snowpark_df = map_relation(request.schema.plan.root)
+                    result = map_relation(request.schema.plan.root)
+                    snowpark_df = result.dataframe
                     snowpark_schema: snowpark.types.StructType = snowpark_df.schema
                     schema = proto_base.AnalyzePlanResponse.Schema(
                         schema=types_proto.DataType(
                             **snowpark_to_proto_type(
-                                snowpark_schema, snowpark_df._column_map, snowpark_df
+                                snowpark_schema, result.column_map, snowpark_df
                             )
                         )
                     )
@@ -262,7 +291,8 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
                     # Snowflake only exposes simplified execution plans, similar to Spark's optimized logical plans.
                     # Snowpark provides the execution plan IFF the dataframe maps to a single query.
                     # TODO: Do we need to return a Spark-like plan?
-                    snowpark_df = map_relation(request.explain.plan.root)
+                    result = map_relation(request.explain.plan.root)
+                    snowpark_df = result.dataframe
                     return proto_base.AnalyzePlanResponse(
                         session_id=request.session_id,
                         explain=proto_base.AnalyzePlanResponse.Explain(