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

snowflake/snowpark_connect/relation/map_crosstab.py

@@ -6,26 +6,41 @@ import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
 import snowflake.snowpark.functions as fn
 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.map_relation import map_relation
 
 
 def map_crosstab(
     rel: relation_proto.Relation,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Perform a crosstab on the input DataFrame.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.crosstab.input)
-    col1 = input_df._column_map.get_snowpark_column_name_from_spark_column_name(
+    input_container = map_relation(rel.crosstab.input)
+    input_df = input_container.dataframe
+
+    col1 = input_container.column_map.get_snowpark_column_name_from_spark_column_name(
         rel.crosstab.col1
     )
-    col2 = input_df._column_map.get_snowpark_column_name_from_spark_column_name(
+    col2 = input_container.column_map.get_snowpark_column_name_from_spark_column_name(
         rel.crosstab.col2
     )
     input_df = input_df.select(
         fn.col(col1).cast("string").alias(col1), fn.col(col2).cast("string").alias(col2)
     )
+
+    # Handle empty DataFrame case
+    if input_df.count() == 0:
+        # For empty DataFrame, return a DataFrame with just the first column name
+        result = input_df.select(
+            fn.lit(f"{rel.crosstab.col1}_{rel.crosstab.col2}").alias("c0")
+        )
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=result,
+            spark_column_names=[f"{rel.crosstab.col1}_{rel.crosstab.col2}"],
+            snowpark_column_names=["c0"],
+        )
+
     result: snowpark.DataFrame = input_df.crosstab(col1, col2)
     new_columns = [f"{rel.crosstab.col1}_{rel.crosstab.col2}"] + [
         (
@@ -45,4 +60,8 @@ def map_crosstab(
     result = result.rename(
         dict(zip(result.columns, [f"c{i}" for i in range(len(result.columns))]))
     )
-    return with_column_map(result, new_columns, result.columns)
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=new_columns,
+        snowpark_column_names=result.columns,
+    )

snowflake/snowpark_connect/relation/map_extension.py

@@ -14,26 +14,29 @@ from snowflake import snowpark
 from snowflake.snowpark_connect.column_name_handler import (
     ColumnNameMap,
     make_column_names_snowpark_compatible,
-    with_column_map,
 )
 from snowflake.snowpark_connect.config import get_boolean_session_config_param
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.expression.map_expression import map_expression
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 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_sql_aggregate_function_count,
     push_outer_dataframe,
+    set_current_grouping_columns,
+)
+from snowflake.snowpark_connect.utils.identifiers import (
+    split_fully_qualified_spark_name,
 )
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
 
 
-def map_extension(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_extension(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     The Extension relation type contains any extensions we use for adding new
     functionality to Spark Connect.
@@ -46,7 +49,8 @@ def map_extension(rel: relation_proto.Relation) -> snowpark.DataFrame:
     match extension.WhichOneof("op"):
         case "rdd_map":
             rdd_map = extension.rdd_map
-            input_df: snowpark.DataFrame = map_relation(rdd_map.input)
+            result = map_relation(rdd_map.input)
+            input_df = result.dataframe
 
             column_name = "_RDD_"
             if len(input_df.columns) > 1:
@@ -67,32 +71,39 @@ def map_extension(rel: relation_proto.Relation) -> snowpark.DataFrame:
                 replace=True,
             )
             result = input_df.select(func(column_name).as_(column_name))
-            return with_column_map(result, [column_name], [column_name], [return_type])
+            return DataFrameContainer.create_with_column_mapping(
+                dataframe=result,
+                spark_column_names=[column_name],
+                snowpark_column_names=[column_name],
+                snowpark_column_types=[return_type],
+            )
         case "subquery_column_aliases":
             subquery_aliases = extension.subquery_column_aliases
             rel.extension.Unpack(subquery_aliases)
-            input_df: snowpark.DataFrame = map_relation(subquery_aliases.input)
-            snowpark_col_names = input_df._column_map.get_snowpark_columns()
+            result = map_relation(subquery_aliases.input)
+            input_df = result.dataframe
+            snowpark_col_names = result.column_map.get_snowpark_columns()
             if len(subquery_aliases.aliases) != len(snowpark_col_names):
                 raise AnalysisException(
                     "Number of column aliases does not match number of columns. "
                     f"Number of column aliases: {len(subquery_aliases.aliases)}; "
                     f"number of columns: {len(snowpark_col_names)}."
                 )
-            return with_column_map(
-                input_df,
-                subquery_aliases.aliases,
-                snowpark_col_names,
-                column_qualifiers=input_df._column_map.get_qualifiers(),
+            return DataFrameContainer.create_with_column_mapping(
+                dataframe=input_df,
+                spark_column_names=subquery_aliases.aliases,
+                snowpark_column_names=snowpark_col_names,
+                column_qualifiers=result.column_map.get_qualifiers(),
             )
         case "lateral_join":
             lateral_join = extension.lateral_join
-            left_df: snowpark.DataFrame = map_relation(lateral_join.left)
+            left_result = map_relation(lateral_join.left)
+            left_df = left_result.dataframe
 
             udtf_info = get_udtf_project(lateral_join.right)
             if udtf_info:
                 return handle_lateral_join_with_udtf(
-                    left_df, lateral_join.right, udtf_info
+                    left_result, lateral_join.right, udtf_info
                 )
 
             left_queries = left_df.queries["queries"]
@@ -101,8 +112,9 @@ def map_extension(rel: relation_proto.Relation) -> snowpark.DataFrame:
                     f"Unexpected number of queries: {len(left_queries)}"
                 )
             left_query = left_queries[0]
-            with push_outer_dataframe(left_df):
-                right_df: snowpark.DataFrame = map_relation(lateral_join.right)
+            with push_outer_dataframe(left_result):
+                right_result = map_relation(lateral_join.right)
+                right_df = right_result.dataframe
             right_queries = right_df.queries["queries"]
             if len(right_queries) != 1:
                 raise SnowparkConnectNotImplementedError(
@@ -112,14 +124,14 @@ def map_extension(rel: relation_proto.Relation) -> snowpark.DataFrame:
             input_df_sql = f"WITH __left AS ({left_query}) SELECT * FROM __left INNER JOIN LATERAL ({right_query})"
             session = snowpark.Session.get_active_session()
             input_df = session.sql(input_df_sql)
-            return with_column_map(
-                input_df,
-                left_df._column_map.get_spark_columns()
-                + right_df._column_map.get_spark_columns(),
-                left_df._column_map.get_snowpark_columns()
-                + right_df._column_map.get_snowpark_columns(),
-                column_qualifiers=left_df._column_map.get_qualifiers()
-                + right_df._column_map.get_qualifiers(),
+            return DataFrameContainer.create_with_column_mapping(
+                dataframe=input_df,
+                spark_column_names=left_result.column_map.get_spark_columns()
+                + right_result.column_map.get_spark_columns(),
+                snowpark_column_names=left_result.column_map.get_snowpark_columns()
+                + right_result.column_map.get_snowpark_columns(),
+                column_qualifiers=left_result.column_map.get_qualifiers()
+                + right_result.column_map.get_qualifiers(),
             )
 
         case "udtf_with_table_arguments":
@@ -165,13 +177,13 @@ def handle_udtf_with_table_arguments(
         raise ValueError(f"UDTF '{udtf_info.function_name}' not found.")
     _udtf_obj, udtf_spark_output_names = session._udtfs[udtf_name_lower]
 
-    table_dfs = []
+    table_containers = []
     for table_arg_info in udtf_info.table_arguments:
-        table_df = map_relation(table_arg_info.table_argument)
-        table_dfs.append((table_df, table_arg_info.table_argument_idx))
+        result = map_relation(table_arg_info.table_argument)
+        table_containers.append((result, table_arg_info.table_argument_idx))
 
-    if len(table_dfs) == 1:
-        base_df = table_dfs[0][0]
+    if len(table_containers) == 1:
+        base_df = table_containers[0][0].dataframe
     else:
         if not get_boolean_session_config_param(
             "spark.sql.tvf.allowMultipleTableArguments.enabled"
@@ -181,11 +193,11 @@ def handle_udtf_with_table_arguments(
                 "Please set `spark.sql.tvf.allowMultipleTableArguments.enabled` to `true`"
             )
 
-        base_df = table_dfs[0][0]
+        base_df = table_containers[0][0].dataframe
         first_table_col_count = len(base_df.columns)
 
-        for table_df, _ in table_dfs[1:]:
-            base_df = base_df.cross_join(table_df)
+        for table_container, _ in table_containers[1:]:
+            base_df = base_df.cross_join(table_container.dataframe)
 
         # Ensure deterministic ordering to match Spark's Cartesian product behavior
         # For two tables A and B, Spark produces: for each B row, iterate through A rows
@@ -206,9 +218,9 @@ def handle_udtf_with_table_arguments(
         scalar_args.append(typed_column.col)
 
     table_arg_variants = []
-    for table_df, table_arg_idx in table_dfs:
-        table_columns = table_df._column_map.get_snowpark_columns()
-        spark_columns = table_df._column_map.get_spark_columns()
+    for table_container, table_arg_idx in table_containers:
+        table_columns = table_container.column_map.get_snowpark_columns()
+        spark_columns = table_container.column_map.get_spark_columns()
 
         # Create a structure that supports both positional and named access
         # Format: {"__fields__": ["col1", "col2"], "__values__": [val1, val2]}
@@ -247,15 +259,15 @@ def handle_udtf_with_table_arguments(
 
     final_df = result_df.select(*udtf_output_columns)
 
-    return with_column_map(
-        final_df,
-        udtf_spark_output_names,
-        udtf_output_columns,
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=final_df,
+        spark_column_names=udtf_spark_output_names,
+        snowpark_column_names=udtf_output_columns,
     )
 
 
 def handle_lateral_join_with_udtf(
-    left_df: snowpark.DataFrame,
+    left_result: DataFrameContainer,
     udtf_relation: relation_proto.Relation,
     udtf_info: tuple[snowpark.udtf.UserDefinedTableFunction, list],
 ) -> snowpark.DataFrame:
@@ -269,7 +281,8 @@ def handle_lateral_join_with_udtf(
     _udtf_obj, udtf_spark_output_names = udtf_info
 
     typer = ExpressionTyper.dummy_typer(session)
-    left_column_map = left_df._column_map
+    left_column_map = left_result.column_map
+    left_df = left_result.dataframe
     table_func = snowpark_fn.table_function(_udtf_obj.name)
     udtf_args = [
         map_expression(arg_proto, left_column_map, typer)[1].col
@@ -278,11 +291,12 @@ def handle_lateral_join_with_udtf(
     udtf_args_variant = [snowpark_fn.to_variant(arg) for arg in udtf_args]
     result_df = left_df.join_table_function(table_func(*udtf_args_variant))
 
-    return with_column_map(
-        result_df,
-        left_df._column_map.get_spark_columns() + udtf_spark_output_names,
-        result_df.columns,
-        column_qualifiers=left_df._column_map.get_qualifiers()
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result_df,
+        spark_column_names=left_result.column_map.get_spark_columns()
+        + udtf_spark_output_names,
+        snowpark_column_names=result_df.columns,
+        column_qualifiers=left_result.column_map.get_qualifiers()
         + [[]] * len(udtf_spark_output_names),
     )
 
@@ -290,7 +304,8 @@ def handle_lateral_join_with_udtf(
 def map_aggregate(
     aggregate: snowflake_proto.Aggregate, plan_id: int
 ) -> snowpark.DataFrame:
-    input_df: snowpark.DataFrame = map_relation(aggregate.input)
+    input_container = map_relation(aggregate.input)
+    input_df: snowpark.DataFrame = input_container.dataframe
 
     # Detect the "GROUP BY ALL" case:
     # - it's a plain GROUP BY (not ROLLUP, CUBE, etc.)
@@ -307,7 +322,7 @@ def map_aggregate(
         if (
             len(parsed_col_name) == 1
             and parsed_col_name[0].lower() == "all"
-            and input_df._column_map.get_snowpark_column_name_from_spark_column_name(
+            and input_container.column_map.get_snowpark_column_name_from_spark_column_name(
                 parsed_col_name[0], allow_non_exists=True
             )
             is None
@@ -320,7 +335,9 @@ def map_aggregate(
     typer = ExpressionTyper(input_df)
 
     def _map_column(exp: expression_proto.Expression) -> tuple[str, TypedColumn]:
-        new_names, snowpark_column = map_expression(exp, input_df._column_map, typer)
+        new_names, snowpark_column = map_expression(
+            exp, input_container.column_map, typer
+        )
         if len(new_names) != 1:
             raise SnowparkConnectNotImplementedError(
                 "Multi-column aggregate expressions are not supported"
@@ -345,6 +362,10 @@ def map_aggregate(
     if not is_group_by_all:
         raw_groupings = [_map_column(exp) for exp in aggregate.grouping_expressions]
 
+    # Set the current grouping columns in context for grouping_id() function
+    grouping_spark_columns = [spark_name for spark_name, _ in raw_groupings]
+    set_current_grouping_columns(grouping_spark_columns)
+
     # Now create column name lists and assign aliases.
     # In case of GROUP BY ALL, even though groupings are a subset of aggregations,
     # they will have their own aliases so we can drop them later.
@@ -378,7 +399,7 @@ def map_aggregate(
             # TODO: What do we do about groupings?
             sets = (
                 [
-                    map_expression(exp, input_df._column_map, typer)[1].col
+                    map_expression(exp, input_container.column_map, typer)[1].col
                     for exp in grouping_sets.grouping_set
                 ]
                 for grouping_sets in aggregate.grouping_sets
@@ -397,16 +418,20 @@ def map_aggregate(
         result = result.select(result.columns[-len(spark_columns) :])
 
     # Build a parent column map that includes groupings.
-    result = with_column_map(
-        result, spark_columns, snowpark_columns, snowpark_column_types
+    result_container = DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=spark_columns,
+        snowpark_column_names=snowpark_columns,
+        snowpark_column_types=snowpark_column_types,
     )
 
     # Drop the groupings.
     grouping_count = len(groupings)
-    return with_column_map(
+
+    return DataFrameContainer.create_with_column_mapping(
         result.drop(snowpark_columns[:grouping_count]),
         spark_columns[grouping_count:],
         snowpark_columns[grouping_count:],
         snowpark_column_types[grouping_count:],
-        parent_column_name_map=result._column_map,
+        parent_column_name_map=result_container.column_map,
     )

snowflake/snowpark_connect/relation/map_join.py

@@ -8,14 +8,10 @@ import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
 import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
-from snowflake.snowpark_connect.column_name_handler import (
-    ColumnNameMap,
-    JoinColumnNameMap,
-    set_schema_getter,
-    with_column_map,
-)
+from snowflake.snowpark_connect.column_name_handler import JoinColumnNameMap
 from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.constants import COLUMN_METADATA_COLLISION_KEY
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.error.error_utils import SparkException
 from snowflake.snowpark_connect.expression.map_expression import (
     map_single_column_expression,
@@ -38,15 +34,18 @@ from snowflake.snowpark_connect.utils.telemetry import (
 USING_COLUMN_NOT_FOUND_ERROR = "[UNRESOLVED_USING_COLUMN_FOR_JOIN] USING column `{0}` not found on the {1} side of the join. The {1}-side columns: {2}"
 
 
-def map_join(rel: relation_proto.Relation) -> snowpark.DataFrame:
-    left_input: snowpark.DataFrame = map_relation(rel.join.left)
-    right_input: snowpark.DataFrame = map_relation(rel.join.right)
+def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
+    left_container: DataFrameContainer = map_relation(rel.join.left)
+    right_container: DataFrameContainer = map_relation(rel.join.right)
+
+    left_input: snowpark.DataFrame = left_container.dataframe
+    right_input: snowpark.DataFrame = right_container.dataframe
     is_natural_join = rel.join.join_type >= NATURAL_JOIN_TYPE_BASE
     using_columns = rel.join.using_columns
     if is_natural_join:
         rel.join.join_type -= NATURAL_JOIN_TYPE_BASE
-        left_spark_columns = left_input._column_map.get_spark_columns()
-        right_spark_columns = right_input._column_map.get_spark_columns()
+        left_spark_columns = left_container.column_map.get_spark_columns()
+        right_spark_columns = right_container.column_map.get_spark_columns()
         common_spark_columns = [
             x for x in left_spark_columns if x in right_spark_columns
         ]
@@ -79,8 +78,8 @@ def map_join(rel: relation_proto.Relation) -> snowpark.DataFrame:
     if rel.join.HasField("join_condition"):
         assert not using_columns
 
-        left_columns = list(left_input._column_map.spark_to_col.keys())
-        right_columns = list(right_input._column_map.spark_to_col.keys())
+        left_columns = list(left_container.column_map.spark_to_col.keys())
+        right_columns = list(right_container.column_map.spark_to_col.keys())
 
         # All PySpark join types are in the format of JOIN_TYPE_XXX.
         # We remove the first 10 characters (JOIN_TYPE_) and replace all underscores with spaces to match the exception.
@@ -90,15 +89,15 @@ def map_join(rel: relation_proto.Relation) -> snowpark.DataFrame:
         with push_sql_scope(), push_evaluating_join_condition(
             pyspark_join_type, left_columns, right_columns
         ):
-            if left_input._alias is not None:
-                set_sql_plan_name(left_input._alias, rel.join.left.common.plan_id)
-            if right_input._alias is not None:
-                set_sql_plan_name(right_input._alias, rel.join.right.common.plan_id)
+            if left_container.alias is not None:
+                set_sql_plan_name(left_container.alias, rel.join.left.common.plan_id)
+            if right_container.alias is not None:
+                set_sql_plan_name(right_container.alias, rel.join.right.common.plan_id)
             _, join_expression = map_single_column_expression(
                 rel.join.join_condition,
                 column_mapping=JoinColumnNameMap(
-                    left_input,
-                    right_input,
+                    left_container.column_map,
+                    right_container.column_map,
                 ),
                 typer=JoinExpressionTyper(left_input, right_input),
             )
@@ -111,7 +110,7 @@ def map_join(rel: relation_proto.Relation) -> snowpark.DataFrame:
         )
     elif using_columns:
         if any(
-            left_input._column_map.get_snowpark_column_name_from_spark_column_name(
+            left_container.column_map.get_snowpark_column_name_from_spark_column_name(
                 c, allow_non_exists=True, return_first=True
             )
             is None
@@ -124,17 +123,17 @@ def map_join(rel: relation_proto.Relation) -> snowpark.DataFrame:
                     next(
                         c
                         for c in using_columns
-                        if left_input._column_map.get_snowpark_column_name_from_spark_column_name(
+                        if left_container.column_map.get_snowpark_column_name_from_spark_column_name(
                             c, allow_non_exists=True, return_first=True
                         )
                         is None
                     ),
                     "left",
-                    left_input._column_map.get_spark_columns(),
+                    left_container.column_map.get_spark_columns(),
                 )
             )
         if any(
-            right_input._column_map.get_snowpark_column_name_from_spark_column_name(
+            right_container.column_map.get_snowpark_column_name_from_spark_column_name(
                 c, allow_non_exists=True, return_first=True
             )
             is None
@@ -147,26 +146,26 @@ def map_join(rel: relation_proto.Relation) -> snowpark.DataFrame:
                     next(
                         c
                         for c in using_columns
-                        if right_input._column_map.get_snowpark_column_name_from_spark_column_name(
+                        if right_container.column_map.get_snowpark_column_name_from_spark_column_name(
                             c, allow_non_exists=True, return_first=True
                         )
                         is None
                     ),
                     "right",
-                    right_input._column_map.get_spark_columns(),
+                    right_container.column_map.get_spark_columns(),
                 )
             )
 
         # Round trip the using columns through the column map to get the correct names
         # in order to support case sensitivity.
         # TODO: case_corrected_left_columns / case_corrected_right_columns may no longer be required as Snowpark dataframe preserves the column casing now.
-        case_corrected_left_columns = left_input._column_map.get_spark_column_names_from_snowpark_column_names(
-            left_input._column_map.get_snowpark_column_names_from_spark_column_names(
+        case_corrected_left_columns = left_container.column_map.get_spark_column_names_from_snowpark_column_names(
+            left_container.column_map.get_snowpark_column_names_from_spark_column_names(
                 list(using_columns), return_first=True
             )
         )
-        case_corrected_right_columns = right_input._column_map.get_spark_column_names_from_snowpark_column_names(
-            right_input._column_map.get_snowpark_column_names_from_spark_column_names(
+        case_corrected_right_columns = right_container.column_map.get_spark_column_names_from_snowpark_column_names(
+            right_container.column_map.get_snowpark_column_names_from_spark_column_names(
                 list(using_columns), return_first=True
             )
         )
@@ -177,12 +176,12 @@ def map_join(rel: relation_proto.Relation) -> snowpark.DataFrame:
         snowpark_using_columns = [
             (
                 left_input[
-                    left_input._column_map.get_snowpark_column_name_from_spark_column_name(
+                    left_container.column_map.get_snowpark_column_name_from_spark_column_name(
                         lft, return_first=True
                     )
                 ],
                 right_input[
-                    right_input._column_map.get_snowpark_column_name_from_spark_column_name(
+                    right_container.column_map.get_snowpark_column_name_from_spark_column_name(
                         r, return_first=True
                     )
                 ],
@@ -231,45 +230,49 @@ def map_join(rel: relation_proto.Relation) -> snowpark.DataFrame:
         # - LEFT SEMI JOIN: Returns left rows that have matches in right table (no right columns)
         # - LEFT ANTI JOIN: Returns left rows that have NO matches in right table (no right columns)
         # Both preserve only the columns from the left DataFrame without adding any columns from the right.
-        spark_cols_after_join: list[str] = left_input._column_map.get_spark_columns()
-        qualifiers = left_input._column_map.get_qualifiers()
+        spark_cols_after_join: list[str] = left_container.column_map.get_spark_columns()
+        qualifiers = left_container.column_map.get_qualifiers()
     else:
         # Add Spark columns and plan_ids from left DF
         spark_cols_after_join: list[str] = list(
-            left_input._column_map.get_spark_columns()
+            left_container.column_map.get_spark_columns()
         ) + [
             spark_col
-            for i, spark_col in enumerate(right_input._column_map.get_spark_columns())
+            for i, spark_col in enumerate(
+                right_container.column_map.get_spark_columns()
+            )
             if spark_col not in case_corrected_right_columns
             or spark_col
-            in right_input._column_map.get_spark_columns()[
+            in right_container.column_map.get_spark_columns()[
                 :i
             ]  # this is to make sure we only remove the column once
         ]
 
-        qualifiers = list(left_input._column_map.get_qualifiers()) + [
-            right_input._column_map.get_qualifier_for_spark_column(spark_col)
-            for i, spark_col in enumerate(right_input._column_map.get_spark_columns())
+        qualifiers = list(left_container.column_map.get_qualifiers()) + [
+            right_container.column_map.get_qualifier_for_spark_column(spark_col)
+            for i, spark_col in enumerate(
+                right_container.column_map.get_spark_columns()
+            )
             if spark_col not in case_corrected_right_columns
             or spark_col
-            in right_input._column_map.get_spark_columns()[
+            in right_container.column_map.get_spark_columns()[
                 :i
             ]  # this is to make sure we only remove the column once]
         ]
 
     column_metadata = {}
-    if left_input._column_map.column_metadata:
-        column_metadata.update(left_input._column_map.column_metadata)
+    if left_container.column_map.column_metadata:
+        column_metadata.update(left_container.column_map.column_metadata)
 
-    if right_input._column_map.column_metadata:
-        for key, value in right_input._column_map.column_metadata.items():
+    if right_container.column_map.column_metadata:
+        for key, value in right_container.column_map.column_metadata.items():
             if key not in column_metadata:
                 column_metadata[key] = value
             else:
                 # In case of collision, use snowpark's column's expr_id as prefix.
                 # this is a temporary solution until SNOW-1926440 is resolved.
                 try:
-                    snowpark_name = right_input._column_map.get_snowpark_column_name_from_spark_column_name(
+                    snowpark_name = right_container.column_map.get_snowpark_column_name_from_spark_column_name(
                         key
                     )
                     expr_id = right_input[snowpark_name]._expression.expr_id
@@ -281,10 +284,10 @@ def map_join(rel: relation_proto.Relation) -> snowpark.DataFrame:
                     # ignore any errors that happens while fetching the metadata
                     pass
 
-    result_df = with_column_map(
-        result,
-        spark_cols_after_join,
-        result.columns,
+    result_container = DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=spark_cols_after_join,
+        snowpark_column_names=result.columns,
         column_metadata=column_metadata,
         column_qualifiers=qualifiers,
     )
@@ -298,7 +301,7 @@ def map_join(rel: relation_proto.Relation) -> snowpark.DataFrame:
         and rel.join.right.common.HasField("plan_id")
     ):
         right_plan_id = rel.join.right.common.plan_id
-        set_plan_id_map(right_plan_id, result_df)
+        set_plan_id_map(right_plan_id, result_container)
 
     # For FULL OUTER joins, we also need to map the left dataframe's plan_id
     # since both columns are replaced with a coalesced column
@@ -309,7 +312,7 @@ def map_join(rel: relation_proto.Relation) -> snowpark.DataFrame:
         and rel.join.left.common.HasField("plan_id")
     ):
         left_plan_id = rel.join.left.common.plan_id
-        set_plan_id_map(left_plan_id, result_df)
+        set_plan_id_map(left_plan_id, result_container)
 
     if rel.join.using_columns:
         # When join 'using_columns', the 'join columns' should go first in result DF.
@@ -323,19 +326,25 @@ def map_join(rel: relation_proto.Relation) -> snowpark.DataFrame:
             remaining = [el for i, el in enumerate(lst) if i not in idxs_to_shift]
             return to_move + remaining
 
-        reordered_df = result_df.select(
-            [snowpark_fn.col(c) for c in reorder(result_df.columns)]
+        # Create reordered DataFrame
+        reordered_df = result_container.dataframe.select(
+            [snowpark_fn.col(c) for c in reorder(result_container.dataframe.columns)]
         )
-        reordered_df._column_map = ColumnNameMap(
-            spark_column_names=reorder(result_df._column_map.get_spark_columns()),
-            snowpark_column_names=reorder(result_df._column_map.get_snowpark_columns()),
+
+        # Create new container with reordered metadata
+        original_df = result_container.dataframe
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=reordered_df,
+            spark_column_names=reorder(result_container.column_map.get_spark_columns()),
+            snowpark_column_names=reorder(
+                result_container.column_map.get_snowpark_columns()
+            ),
             column_metadata=column_metadata,
             column_qualifiers=reorder(qualifiers),
+            table_name=result_container.table_name,
+            cached_schema_getter=lambda: snowpark.types.StructType(
+                reorder(original_df.schema.fields)
+            ),
         )
-        reordered_df._table_name = result_df._table_name
-        set_schema_getter(
-            reordered_df,
-            lambda: snowpark.types.StructType(reorder(result_df.schema.fields)),
-        )
-        return reordered_df
-    return result_df
+
+    return result_container