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

snowflake/snowpark_connect/relation/map_sample_by.py

@@ -6,7 +6,7 @@ import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
 from snowflake import snowpark
-from snowflake.snowpark_connect.column_name_handler import set_schema_getter
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.expression.literal import get_literal_field_and_name
 from snowflake.snowpark_connect.expression.map_expression import (
     map_single_column_expression,
@@ -15,21 +15,28 @@ from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 
 
-def map_sample_by(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_sample_by(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Sample by an expression on the input DataFrame.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.sample_by.input)
+    input_container = map_relation(rel.sample_by.input)
+    input_df = input_container.dataframe
+
     exp: expressions_proto.Expression = rel.sample_by.col
     _, col_expr = map_single_column_expression(
-        exp, input_df._column_map, ExpressionTyper(input_df)
+        exp, input_container.column_map, ExpressionTyper(input_df)
     )
     fractions = {
         get_literal_field_and_name(frac.stratum)[0]: frac.fraction
         for frac in rel.sample_by.fractions
     }
     result: snowpark.DataFrame = input_df.sampleBy(col_expr.col, fractions)
-    result._column_map = input_df._column_map
-    result._table_name = input_df._table_name
-    set_schema_getter(result, lambda: input_df.schema)
-    return result
+    return DataFrameContainer(
+        result,
+        column_map=input_container.column_map,
+        table_name=input_container.table_name,
+        alias=input_container.alias,
+        cached_schema_getter=lambda: input_df.schema,
+    )

snowflake/snowpark_connect/relation/map_show_string.py

@@ -2,11 +2,17 @@
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
 
+import copy
+
 import pandas
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
 from snowflake import snowpark
 from snowflake.snowpark._internal.analyzer import analyzer_utils
+from snowflake.snowpark.functions import col
+from snowflake.snowpark.types import DateType, StringType, StructField, StructType
+from snowflake.snowpark_connect.column_name_handler import set_schema_getter
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 
 
@@ -18,12 +24,15 @@ def map_show_string(rel: relation_proto.Relation) -> pandas.DataFrame:
     message creates a string. The client expects this string to be packed into an Arrow
     Buffer object as a single cell.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.show_string.input)
+    input_df_container: DataFrameContainer = map_relation(rel.show_string.input)
+    raw_input_df = input_df_container.dataframe
+    input_df = _handle_datetype_columns(raw_input_df)
+
     show_string = input_df._show_string_spark(
         num_rows=rel.show_string.num_rows,
         truncate=rel.show_string.truncate,
         vertical=rel.show_string.vertical,
-        _spark_column_names=input_df._column_map.get_spark_columns(),
+        _spark_column_names=input_df_container.column_map.get_spark_columns(),
     )
     return pandas.DataFrame({"show_string": [show_string]})
 
@@ -32,13 +41,15 @@ def map_repr_html(rel: relation_proto.Relation) -> pandas.DataFrame:
     """
     Generate the html string representation of the input dataframe.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.html_string.input)
+    input_df_container: DataFrameContainer = map_relation(rel.html_string.input)
+    input_df = input_df_container.dataframe
+
     input_panda = input_df.toPandas()
     input_panda.rename(
         columns={
             analyzer_utils.unquote_if_quoted(
-                input_df._column_map.get_snowpark_columns()[i]
-            ): input_df._column_map.get_spark_columns()[i]
+                input_df_container.column_map.get_snowpark_columns()[i]
+            ): input_df_container.column_map.get_spark_columns()[i]
             for i in range(len(input_panda.columns))
         },
         inplace=True,
@@ -48,3 +59,29 @@ def map_repr_html(rel: relation_proto.Relation) -> pandas.DataFrame:
         max_rows=rel.html_string.num_rows,
     )
     return pandas.DataFrame({"html_string": [html_string]})
+
+
+def _handle_datetype_columns(input_df: snowpark.DataFrame) -> snowpark.DataFrame:
+    """
+    Maps DateType columns to strings it aims to allow showing the dates which are out of range of datetime.datetime.
+    """
+    new_column_mapping = []
+    new_fields = []
+    transformation_required = False
+    for field in input_df.schema:
+        if isinstance(field.datatype, DateType):
+            transformation_required = True
+            new_column_mapping.append(col(field.name).cast(StringType()))
+            new_fields.append(StructField(field.name, StringType()))
+        else:
+            new_column_mapping.append(col(field.name))
+            new_fields.append(field)
+
+    if not transformation_required:
+        return input_df
+
+    transformed_df = input_df.select(new_column_mapping)
+    set_schema_getter(transformed_df, lambda: StructType(new_fields))
+    transformed_df._column_map = copy.deepcopy(input_df._column_map)
+
+    return transformed_df

snowflake/snowpark_connect/relation/map_sql.py

@@ -26,15 +26,16 @@ from snowflake.snowpark._internal.analyzer.analyzer_utils import (
     unquote_if_quoted,
 )
 from snowflake.snowpark._internal.type_utils import convert_sp_to_sf_type
-from snowflake.snowpark._internal.utils import is_sql_select_statement
+from snowflake.snowpark._internal.utils import is_sql_select_statement, quote_name
 from snowflake.snowpark.functions import when_matched, when_not_matched
 from snowflake.snowpark_connect.config import (
-    auto_uppercase_ddl,
+    auto_uppercase_non_column_identifiers,
     get_boolean_session_config_param,
     global_config,
     set_config_param,
     unset_config_param,
 )
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.expression.map_expression import (
     ColumnNameMap,
     map_single_column_expression,
@@ -72,6 +73,7 @@ from ..expression.map_sql_expression import (
     map_logical_plan_expression,
     sql_parser,
 )
+from ..utils.identifiers import spark_to_sf_single_id
 
 _ctes = ContextVar[dict[str, relation_proto.Relation]]("_ctes", default={})
 
@@ -159,38 +161,32 @@ def parse_pos_args(
     return dict(zip(sorted(positions), pos_args))
 
 
-def execute_logical_plan(logical_plan) -> snowpark.DataFrame:
+def execute_logical_plan(logical_plan) -> DataFrameContainer:
     proto = map_logical_plan_relation(logical_plan)
     with push_evaluating_sql_scope():
         return map_relation(proto)
 
 
-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: jpype.JObject) -> str:
     return ".".join(
-        _spark_to_snowflake_single_id(str(part)) for part in as_java_list(multipart_id)
+        spark_to_sf_single_id(str(part)) for part in as_java_list(multipart_id)
     )
 
 
 def _rename_columns(
-    df: snowpark.DataFrame, user_specified_columns
+    df: snowpark.DataFrame, user_specified_columns, column_map: ColumnNameMap
 ) -> snowpark.DataFrame:
     user_columns = [str(col._1()) for col in as_java_list(user_specified_columns)]
 
     if user_columns:
         columns = zip(df.columns, user_columns)
     else:
-        columns = df._column_map.snowpark_to_spark_map().items()
+        columns = column_map.snowpark_to_spark_map().items()
 
     for orig_column, user_column in columns:
         df = df.with_column_renamed(
-            orig_column, _spark_to_snowflake_single_id(user_column)
+            orig_column, spark_to_sf_single_id(user_column, is_column=True)
         )
-
     return df
 
 
@@ -199,11 +195,12 @@ def _create_table_as_select(logical_plan, mode: str) -> None:
     name = get_relation_identifier_name(logical_plan.name())
     comment = logical_plan.tableSpec().comment()
 
-    df = execute_logical_plan(logical_plan.query())
-    columns = df._column_map.snowpark_to_spark_map().items()
+    container = execute_logical_plan(logical_plan.query())
+    df = container.dataframe
+    columns = container.column_map.snowpark_to_spark_map().items()
     for orig_column, user_column in columns:
         df = df.with_column_renamed(
-            orig_column, _spark_to_snowflake_single_id(user_column)
+            orig_column, spark_to_sf_single_id(user_column, is_column=True)
         )
 
     # TODO escaping should be handled by snowpark. remove when SNOW-2210271 is done
@@ -218,11 +215,11 @@ def _create_table_as_select(logical_plan, mode: str) -> None:
 
 
 def _spark_field_to_sql(field: jpype.JObject, is_column: bool) -> str:
-    # Column names will be uppercased according to "snowpark.connect.auto-uppercase.ddl",
+    # Column names will be uppercased according to "snowpark.connect.sql.identifiers.auto-uppercase",
     # and struct fields will be left as is. This should allow users to use the same names
     # in spark and Snowflake in most cases.
     if is_column:
-        name = _spark_to_snowflake_single_id(str(field.name()))
+        name = spark_to_sf_single_id(str(field.name()), is_column=True)
     else:
         name = quote_name_without_upper_casing(str(field.name()))
     data_type_str = _spark_datatype_to_sql(field.dataType())
@@ -257,10 +254,12 @@ def _normalize_identifiers(node):
     The identifiers need to be uppercased to match Snowflake's behaviour. Users can disable this by setting
     the `snowpark.connect.auto_uppercase_ddl` config to False.
     """
-    if isinstance(node, Identifier):
-        fixed = node.this.upper() if auto_uppercase_ddl() else node.this
-        return Identifier(this=fixed, quoted=True)
-    return node
+    if not isinstance(node, Identifier):
+        return node
+    elif auto_uppercase_non_column_identifiers():
+        return Identifier(this=node.this.upper(), quoted=True)
+    else:
+        return Identifier(this=node.this, quoted=True)
 
 
 def _remove_file_format_property(node):
@@ -328,7 +327,10 @@ def map_sql_to_pandas_df(
                 # Build Snowflake SQL from logical plan attributes
                 for col in as_java_list(columns_to_add):
                     # Follow the same pattern as AlterColumn for column name extraction
-                    col_name = ".".join(str(part) for part in as_java_list(col.name()))
+                    col_name = ".".join(
+                        spark_to_sf_single_id(part, is_column=True)
+                        for part in as_java_list(col.name())
+                    )
                     col_type = _spark_datatype_to_sql(col.dataType())
                     snowflake_sql = (
                         f"ALTER TABLE {table_name} ADD COLUMN {col_name} {col_type}"
@@ -341,9 +343,22 @@ def map_sql_to_pandas_df(
 
                 # Extract actual column name
                 column_name = ".".join(
-                    str(part) for part in as_java_list(column_obj.name())
+                    spark_to_sf_single_id(part, is_column=True)
+                    for part in as_java_list(column_obj.name())
                 )
 
+                if not global_config.spark_sql_caseSensitive:
+                    case_insensitive_name = next(
+                        (
+                            f.name
+                            for f in session.table(table_name).schema.fields
+                            if f.name.lower() == column_name.lower()
+                        ),
+                        None,
+                    )
+                    if case_insensitive_name:
+                        column_name = case_insensitive_name
+
                 # Build ALTER COLUMN command from logical plan attributes
                 alter_parts = []
 
@@ -381,9 +396,7 @@ def map_sql_to_pandas_df(
                 if_not_exists = "IF NOT EXISTS " if logical_plan.ifNotExists() else ""
                 session.sql(f"CREATE SCHEMA {if_not_exists}{name}").collect()
                 if previous_name is not None:
-                    session.sql(
-                        f"USE SCHEMA {_spark_to_snowflake_single_id(previous_name)}"
-                    ).collect()
+                    session.sql(f"USE SCHEMA {quote_name(previous_name)}").collect()
                 else:
                     # TODO: Unset the schema
                     pass
@@ -443,7 +456,8 @@ def map_sql_to_pandas_df(
                 else:
                     object_name: str = as_java_list(logical_plan.child().nameParts())[0]
                 _accessing_temp_object.set(False)
-                df = execute_logical_plan(logical_plan.query())
+                df_container = execute_logical_plan(logical_plan.query())
+                df = df_container.dataframe
                 if _accessing_temp_object.get():
                     raise AnalysisException(
                         f"[INVALID_TEMP_OBJ_REFERENCE] Cannot create the persistent object `{CURRENT_CATALOG_NAME}`.`{current_schema}`.`{object_name}` "
@@ -454,7 +468,9 @@ def map_sql_to_pandas_df(
                 name = get_relation_identifier_name(logical_plan.child())
                 comment = logical_plan.comment()
 
-                df = _rename_columns(df, logical_plan.userSpecifiedColumns())
+                df = _rename_columns(
+                    df, logical_plan.userSpecifiedColumns(), df_container.column_map
+                )
 
                 # TODO: Support logical_plan.replace() == False
                 df.create_or_replace_view(
@@ -464,7 +480,8 @@ def map_sql_to_pandas_df(
                     else None,
                 )
             case "CreateViewCommand":
-                df = execute_logical_plan(logical_plan.plan())
+                df_container = execute_logical_plan(logical_plan.plan())
+                df = df_container.dataframe
                 tmp_views = _get_current_temp_objects()
                 tmp_views.add(
                     (
@@ -475,7 +492,7 @@ def map_sql_to_pandas_df(
                 )
 
                 name = str(logical_plan.name().identifier())
-                name = _spark_to_snowflake_single_id(name)
+                name = spark_to_sf_single_id(name)
                 if isinstance(
                     logical_plan.viewType(),
                     jpype.JClass(
@@ -490,7 +507,9 @@ def map_sql_to_pandas_df(
                     else None
                 )
 
-                df = _rename_columns(df, logical_plan.userSpecifiedColumns())
+                df = _rename_columns(
+                    df, logical_plan.userSpecifiedColumns(), df_container.column_map
+                )
 
                 if logical_plan.replace():
                     df.create_or_replace_temp_view(
@@ -504,6 +523,7 @@ def map_sql_to_pandas_df(
                     )
             case "DescribeColumn":
                 name = get_relation_identifier_name(logical_plan.column())
+                # todo double check if this is correct
                 rows = session.sql(f"DESCRIBE TABLE {name}").collect()
             case "DescribeNamespace":
                 name = get_relation_identifier_name(logical_plan.namespace(), True)
@@ -521,11 +541,12 @@ def map_sql_to_pandas_df(
                 # This gets the schema without executing the query (similar to Spark's DESCRIBE QUERY)
                 # Get the inner query plan and convert it to SQL
                 inner_query_plan = logical_plan.plan()
-                df = execute_logical_plan(inner_query_plan)
+                df_container = execute_logical_plan(inner_query_plan)
+                df = df_container.dataframe
                 schema = df.schema
 
                 # Get original Spark column names using the column map from the original DataFrame
-                spark_columns = df._column_map.get_spark_columns()
+                spark_columns = df_container.column_map.get_spark_columns()
                 data = []
                 for i, field in enumerate(schema.fields):
                     # Use original Spark column name from column map
@@ -615,9 +636,9 @@ def map_sql_to_pandas_df(
                     "UnresolvedHaving",
                     "Distinct",
                 ):
-                    expr = execute_logical_plan(logical_plan.logicalPlan()).queries[
-                        "queries"
-                    ][0]
+                    expr = execute_logical_plan(
+                        logical_plan.logicalPlan()
+                    ).dataframe.queries["queries"][0]
                     final_sql = f"EXPLAIN USING TEXT {expr}"
                     rows = session.sql(final_sql).collect()
                 elif (
@@ -626,7 +647,7 @@ def map_sql_to_pandas_df(
                 ):
                     expr = execute_logical_plan(
                         logical_plan.logicalPlan().query()
-                    ).queries["queries"][0]
+                    ).dataframe.queries["queries"][0]
                     final_sql = f"EXPLAIN USING TEXT {expr}"
                     rows = session.sql(final_sql).collect()
                 else:
@@ -635,7 +656,8 @@ def map_sql_to_pandas_df(
                         f"{logical_plan_name} is not supported yet with EXPLAIN."
                     )
             case "InsertIntoStatement":
-                df = execute_logical_plan(logical_plan.query())
+                df_container = execute_logical_plan(logical_plan.query())
+                df = df_container.dataframe
                 queries = df.queries["queries"]
                 if len(queries) != 1:
                     raise SnowparkConnectNotImplementedError(
@@ -645,7 +667,7 @@ def map_sql_to_pandas_df(
                 name = get_relation_identifier_name(logical_plan.table(), True)
 
                 user_columns = [
-                    _spark_to_snowflake_single_id(str(col))
+                    spark_to_sf_single_id(str(col), is_column=True)
                     for col in as_java_list(logical_plan.userSpecifiedCols())
                 ]
                 overwrite_str = "OVERWRITE" if logical_plan.overwrite() else ""
@@ -751,32 +773,36 @@ def map_sql_to_pandas_df(
                             )
                     return assignments
 
-                source_df = map_relation(
+                source_df_container = map_relation(
                     map_logical_plan_relation(logical_plan.sourceTable())
                 )
-
+                source_df = source_df_container.dataframe
                 plan_id = gen_sql_plan_id()
-
-                target_df = map_relation(
+                target_df_container = map_relation(
                     map_logical_plan_relation(logical_plan.targetTable(), plan_id)
                 )
-                for col in target_df._column_map.columns:
+                target_df = target_df_container.dataframe
+
+                for col in target_df_container.column_map.columns:
                     target_df = target_df.with_column_renamed(
-                        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),
                     )
-                target_df = column_name_handler.with_column_map(
-                    target_df, target_df.columns, target_df.columns
+                target_df_container = DataFrameContainer.create_with_column_mapping(
+                    dataframe=target_df,
+                    spark_column_names=target_df.columns,
+                    snowpark_column_names=target_df.columns,
                 )
 
-                set_plan_id_map(plan_id, target_df)
+                set_plan_id_map(plan_id, target_df_container)
 
                 joined_df_before_condition: snowpark.DataFrame = source_df.join(
                     target_df
                 )
 
                 column_mapping_for_conditions = column_name_handler.JoinColumnNameMap(
-                    source_df,
-                    target_df,
+                    source_df_container.column_map,
+                    target_df_container.column_map,
                 )
                 typer_for_expressions = ExpressionTyper(joined_df_before_condition)
 
@@ -803,8 +829,8 @@ def map_sql_to_pandas_df(
                     ):
                         assignments = _get_assignments_from_action(
                             matched_action,
-                            source_df._column_map,
-                            target_df._column_map,
+                            source_df_container.column_map,
+                            target_df_container.column_map,
                             ExpressionTyper(source_df),
                             ExpressionTyper(target_df),
                         )
@@ -825,8 +851,8 @@ def map_sql_to_pandas_df(
                     ):
                         assignments = _get_assignments_from_action(
                             not_matched_action,
-                            source_df._column_map,
-                            target_df._column_map,
+                            source_df_container.column_map,
+                            target_df_container.column_map,
                             ExpressionTyper(source_df),
                             ExpressionTyper(target_df),
                         )
@@ -852,20 +878,28 @@ def map_sql_to_pandas_df(
                     source_df, merge_condition_typed_col.col, clauses
                 )
             case "DeleteFromTable":
-                df = map_relation(map_logical_plan_relation(logical_plan.table()))
-                for col in df._column_map.columns:
+                df_container = map_relation(
+                    map_logical_plan_relation(logical_plan.table())
+                )
+                df = df_container.dataframe
+                for col in df_container.column_map.columns:
                     df = df.with_column_renamed(
-                        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),
                     )
-                df = column_name_handler.with_column_map(df, df.columns, df.columns)
-
+                df_container = column_name_handler.create_with_column_mapping(
+                    dataframe=df,
+                    spark_column_names=df.columns,
+                    snowpark_column_names=df.columns,
+                )
+                df = df_container.dataframe
                 name = get_relation_identifier_name(logical_plan.table(), True)
                 (
                     condition_column_name,
                     condition_typed_col,
                 ) = map_single_column_expression(
                     map_logical_plan_expression(logical_plan.condition()),
-                    df._column_map,
+                    df_container.column_map,
                     ExpressionTyper(df),
                 )
                 session.table(name).delete(condition_typed_col.col)
@@ -873,9 +907,23 @@ def map_sql_to_pandas_df(
                 table_name = get_relation_identifier_name(logical_plan.table(), True)
                 column_obj = logical_plan.column()
                 old_column_name = ".".join(
-                    str(part) for part in as_java_list(column_obj.name())
+                    spark_to_sf_single_id(str(part), is_column=True)
+                    for part in as_java_list(column_obj.name())
+                )
+                if not global_config.spark_sql_caseSensitive:
+                    case_insensitive_name = next(
+                        (
+                            f.name
+                            for f in session.table(table_name).schema.fields
+                            if f.name.lower() == old_column_name.lower()
+                        ),
+                        None,
+                    )
+                    if case_insensitive_name:
+                        old_column_name = case_insensitive_name
+                new_column_name = spark_to_sf_single_id(
+                    str(logical_plan.newName()), is_column=True
                 )
-                new_column_name = str(logical_plan.newName())
 
                 # Pass through to Snowflake
                 snowflake_sql = f"ALTER TABLE {table_name} RENAME COLUMN {old_column_name} TO {new_column_name}"
@@ -971,6 +1019,32 @@ def map_sql_to_pandas_df(
                 else:
                     rows = session.sql("SHOW TABLES").collect()
 
+                # Return empty DataFrame with proper schema if no results
+                if not rows:
+                    if class_name == "ShowTableExtended":
+                        return (
+                            pandas.DataFrame(
+                                {
+                                    "namespace": [""],
+                                    "tableName": [""],
+                                    "isTemporary": [""],
+                                    "information": [""],
+                                }
+                            ),
+                            "",
+                        )
+                    else:
+                        return (
+                            pandas.DataFrame(
+                                {
+                                    "namespace": [""],
+                                    "tableName": [""],
+                                    "isTemporary": [""],
+                                }
+                            ),
+                            "",
+                        )
+
                 # Apply pattern filtering if pattern is provided
                 # This is workaround to filter using Python regex.
                 if pattern and rows:
@@ -1020,9 +1094,7 @@ def map_sql_to_pandas_df(
                     if db_name and multi_part_len == 2:
                         # Check db_name is same as in the full table name
                         if (
-                            _spark_to_snowflake_single_id(
-                                str(db_and_table_name[0])
-                            ).casefold()
+                            spark_to_sf_single_id(str(db_and_table_name[0])).casefold()
                             != db_name.casefold()
                         ):
                             raise AnalysisException(
@@ -1075,11 +1147,7 @@ def map_sql_to_pandas_df(
         if _is_sql_select_statement_helper(sql_string):
             return None, None
         session = snowpark.Session.get_active_session()
-        sql_df = session.sql(sql_string)
-        columns = sql_df.columns
-        column_name_handler.with_column_map(sql_df, columns, columns)
-        rows = sql_df.collect()
-
+        rows = session.sql(sql_string).collect()
     if rows:
         return pandas.DataFrame(rows), ""
     return pandas.DataFrame({"": [""]}), ""
@@ -1089,7 +1157,9 @@ def get_sql_passthrough() -> bool:
     return get_boolean_session_config_param("snowpark.connect.sql.passthrough")
 
 
-def map_sql(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_sql(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Map a SQL string to a DataFrame.
 
@@ -1112,7 +1182,11 @@ def map_sql(rel: relation_proto.Relation) -> snowpark.DataFrame:
         session = snowpark.Session.get_active_session()
         sql_df = session.sql(rel.sql.query)
         columns = sql_df.columns
-        return column_name_handler.with_column_map(sql_df, columns, columns)
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=sql_df,
+            spark_column_names=columns,
+            snowpark_column_names=columns,
+        )
 
 
 def map_logical_plan_relation(
@@ -1453,8 +1527,9 @@ def map_logical_plan_relation(
 
             # Need to find ids which are not part of values and remaining cols of df
             input_rel = map_logical_plan_relation(rel.child())
-            input_df: snowpark.DataFrame = map_relation(input_rel)
-            column_map = input_df._column_map
+            result = map_relation(input_rel)
+            input_df: snowpark.DataFrame = result.dataframe
+            column_map = result.column_map
             typer = ExpressionTyper(input_df)
             unpivot_spark_names = []
             for v in values:
@@ -1744,8 +1819,8 @@ def map_logical_plan_relation(
             # )
 
             # This is a workaround to fix the bug in snowpark where if we select posexplode with *, it would return wrong columns
-            input_df = map_relation(input_relation)
-            spark_columns = input_df._column_map.get_spark_columns()
+            input_container = map_relation(input_relation)
+            spark_columns = input_container.column_map.get_spark_columns()
             column_expressions = [
                 expressions_proto.Expression(
                     unresolved_attribute=expressions_proto.Expression.UnresolvedAttribute(
@@ -1796,7 +1871,9 @@ def get_relation_identifier_name(name_obj, is_multi_part: bool = False) -> str:
         expr = map_single_column_expression(
             expr_proto, m, ExpressionTyper.dummy_typer(session)
         )
-        name = session.range(1).select(expr[1].col).collect()[0][0]
+        name = spark_to_sf_single_id(
+            session.range(1).select(expr[1].col).collect()[0][0]
+        )
     else:
         if is_multi_part:
             name = _spark_to_snowflake(name_obj.multipartIdentifier())