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

snowflake/snowpark_connect/relation/map_column_ops.py

@@ -30,10 +30,9 @@ from snowflake.snowpark.table_function import _ExplodeFunctionCall
 from snowflake.snowpark.types import DataType, StructField, StructType, _NumericType
 from snowflake.snowpark_connect.column_name_handler import (
     make_column_names_snowpark_compatible,
-    set_schema_getter,
-    with_column_map,
 )
 from snowflake.snowpark_connect.config import global_config
+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_alias,
@@ -53,13 +52,13 @@ from snowflake.snowpark_connect.type_mapping import (
 )
 from snowflake.snowpark_connect.typed_column import TypedColumn
 from snowflake.snowpark_connect.utils import context
-from snowflake.snowpark_connect.utils.attribute_handling import (
-    split_fully_qualified_spark_name,
-)
 from snowflake.snowpark_connect.utils.context import (
     clear_lca_alias_map,
     register_lca_alias,
 )
+from snowflake.snowpark_connect.utils.identifiers import (
+    split_fully_qualified_spark_name,
+)
 from snowflake.snowpark_connect.utils.udtf_helper import (
     TEST_FLAG_FORCE_CREATE_SPROC,
     create_apply_udtf_in_sproc,
@@ -68,20 +67,21 @@ from snowflake.snowpark_connect.utils.udtf_helper import (
 
 def map_drop(
     rel: relation_proto.Relation,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Drop columns from a DataFrame.
 
     The drop is a list of expressions that is applied to the DataFrame.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.drop.input)
+    input_container = map_relation(rel.drop.input)
+    input_df = input_container.dataframe
     typer = ExpressionTyper(input_df)
     columns_to_drop_with_names = []
     for exp in rel.drop.columns:
         if exp.WhichOneof("expr_type") == "unresolved_attribute":
             try:
                 columns_to_drop_with_names.append(
-                    map_single_column_expression(exp, input_df._column_map, typer)
+                    map_single_column_expression(exp, input_container.column_map, typer)
                 )
             except AnalysisException as e:
                 if "[COLUMN_NOT_FOUND]" in e.message:
@@ -91,8 +91,8 @@ def map_drop(
     columns_to_drop: list[Column] = [
         col[1].col for col in columns_to_drop_with_names
     ] + [
-        snowpark_functions_col(c, input_df._column_map)
-        for c in input_df._column_map.get_snowpark_column_names_from_spark_column_names(
+        snowpark_functions_col(c, input_container.column_map)
+        for c in input_container.column_map.get_snowpark_column_names_from_spark_column_names(
             list(rel.drop.column_names)
         )
         if c is not None
@@ -100,7 +100,7 @@ def map_drop(
     # Sometimes we get a drop query with only invalid names. In this case, we return
     # the input DataFrame.
     if len(columns_to_drop) == 0:
-        return input_df
+        return input_container
 
     def _get_column_names_to_drop() -> list[str]:
         # more or less copied from Snowpark's DataFrame::drop
@@ -128,47 +128,52 @@ def map_drop(
     # Snowpark doesn't allow dropping all columns, so we have an EmptyDataFrame
     # object to handle these cases.
     try:
-        new_columns_names = input_df._column_map.get_snowpark_columns_after_drop(
+        column_map = input_container.column_map
+        new_columns_names = column_map.get_snowpark_columns_after_drop(
             _get_column_names_to_drop()
         )
         result: snowpark.DataFrame = input_df.drop(*columns_to_drop)
-        return with_column_map(
-            result,
-            input_df._column_map.get_spark_column_names_from_snowpark_column_names(
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=result,
+            spark_column_names=column_map.get_spark_column_names_from_snowpark_column_names(
                 new_columns_names
             ),
             snowpark_column_names=new_columns_names,
-            column_qualifiers=input_df._column_map.get_qualifiers_for_columns_after_drop(
+            column_qualifiers=column_map.get_qualifiers_for_columns_after_drop(
                 _get_column_names_to_drop()
             ),
-            parent_column_name_map=input_df._column_map,
+            parent_column_name_map=column_map,
         )
     except snowpark.exceptions.SnowparkColumnException:
         from snowflake.snowpark_connect.empty_dataframe import EmptyDataFrame
 
-        return EmptyDataFrame()
+        return DataFrameContainer(EmptyDataFrame())
 
 
-def map_project(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_project(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
-    Project column(s).
+    Project column(s) and return a container.
 
-    Projections come in as expressions, which are mapped to `snowpark.Column`
-    objects.
+    Projections come in as expressions, which are mapped to `snowpark.Column` objects.
     """
     if rel.project.HasField("input"):
-        input_df = map_relation(rel.project.input)
+        input_container = map_relation(rel.project.input)
+        input_df = input_container.dataframe
     else:
         # Create a dataframe to represent a OneRowRelation AST node.
         # XXX: Snowflake does not support 0-column tables, so create a dummy column;
         # its name does not seem to show up anywhere.
         session = snowpark.Session.get_active_session()
-        input_df = with_column_map(
-            session.create_dataframe([None], ["__DUMMY"]),
-            ["__DUMMY"],
-            ["__DUMMY"],
+        input_container = DataFrameContainer.create_with_column_mapping(
+            dataframe=session.create_dataframe([None], ["__DUMMY"]),
+            spark_column_names=["__DUMMY"],
+            snowpark_column_names=["__DUMMY"],
         )
-    context.set_df_before_projection(input_df)
+
+    input_df = input_container.dataframe
+    context.set_df_before_projection(input_container)
     expressions: list[expressions_proto.Expression] = rel.project.expressions
     if not expressions:
         # XXX: Snowflake does not support 0-column tables, so create a dummy column;
@@ -217,7 +222,7 @@ def map_project(rel: relation_proto.Relation) -> snowpark.DataFrame:
     )
 
     for exp in expressions:
-        new_spark_names, mapper = map_expression(exp, input_df._column_map, typer)
+        new_spark_names, mapper = map_expression(exp, input_container.column_map, typer)
         if len(new_spark_names) == 1 and not isinstance(
             mapper.col, _ExplodeFunctionCall
         ):
@@ -238,7 +243,7 @@ def map_project(rel: relation_proto.Relation) -> snowpark.DataFrame:
                 and not has_unresolved_star
             ):
                 # Try to get the existing Snowpark column name for this Spark column
-                existing_snowpark_name = input_df._column_map.get_snowpark_column_name_from_spark_column_name(
+                existing_snowpark_name = input_container.column_map.get_snowpark_column_name_from_spark_column_name(
                     spark_name, allow_non_exists=True
                 )
 
@@ -308,22 +313,28 @@ def map_project(rel: relation_proto.Relation) -> snowpark.DataFrame:
         result = result.toDF(*final_snowpark_columns)
         new_snowpark_columns = final_snowpark_columns
 
-    return with_column_map(
-        result,
-        new_spark_columns,
-        new_snowpark_columns,
-        column_types,
-        column_metadata=input_df._column_map.column_metadata,
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=new_spark_columns,
+        snowpark_column_names=new_snowpark_columns,
+        snowpark_column_types=column_types,
+        column_metadata=input_container.column_map.column_metadata,
         column_qualifiers=qualifiers,
-        parent_column_name_map=input_df._column_map,
+        parent_column_name_map=input_container.column_map,
+        table_name=input_container.table_name,
+        alias=input_container.alias,
     )
 
 
-def map_sort(sort: relation_proto.Sort) -> snowpark.DataFrame:
+def map_sort(
+    sort: relation_proto.Sort,
+) -> DataFrameContainer:
     """
-    Implements DataFrame.sort().
+    Implements DataFrame.sort() and return a container.
+
     """
-    input_df = map_relation(sort.input)
+    input_container = map_relation(sort.input)
+    input_df = input_container.dataframe
     cols = []
     ascending = []  # Ignored if all order values are set to "unspecified".
     order_specified = False
@@ -338,7 +349,7 @@ def map_sort(sort: relation_proto.Sort) -> snowpark.DataFrame:
         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
@@ -354,7 +365,7 @@ def map_sort(sort: relation_proto.Sort) -> snowpark.DataFrame:
                     direction=sort_order[0].direction,
                     null_ordering=sort_order[0].null_ordering,
                 )
-                for col in input_df._column_map.get_spark_columns()
+                for col in input_container.column_map.get_spark_columns()
             ]
 
     for so in sort_order:
@@ -370,7 +381,7 @@ def map_sort(sort: relation_proto.Sort) -> snowpark.DataFrame:
                 )
         else:
             _, typed_column = map_single_column_expression(
-                so.child, input_df._column_map, typer
+                so.child, input_container.column_map, typer
             )
             col = typed_column.col
 
@@ -412,29 +423,35 @@ def map_sort(sort: relation_proto.Sort) -> snowpark.DataFrame:
     if not order_specified:
         ascending = None
     result = input_df.sort(cols, ascending=ascending)
-    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,
+        input_container.column_map,
+        input_container.table_name,
+        cached_schema_getter=lambda: input_df.schema,
+    )
 
 
-def map_to_df(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_to_df(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
-    Transform the column names of the input DataFrame.
+    Transform the column names of the input DataFrame and return a container.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.to_df.input)
+    input_container = map_relation(rel.to_df.input)
+    input_df = input_container.dataframe
+
     new_column_names = list(rel.to_df.column_names)
-    if len(new_column_names) != len(input_df._column_map.columns):
+    if len(new_column_names) != len(input_container.column_map.columns):
         # TODO: Check error type here
         raise ValueError(
             "Number of column names must match number of columns in DataFrame"
         )
-
     snowpark_new_column_names = make_column_names_snowpark_compatible(
         new_column_names, rel.common.plan_id
     )
-
     result = input_df.toDF(*snowpark_new_column_names)
+
     if result._select_statement is not None:
         # do not allow snowpark to flatten the to_df result
         # TODO: remove after SNOW-2203706 is fixed
@@ -448,27 +465,33 @@ def map_to_df(rel: relation_proto.Relation) -> snowpark.DataFrame:
             ]
         )
 
-    set_schema_getter(result, _get_schema)
-    result_with_column_map = with_column_map(
-        result,
-        new_column_names,
+    result_container = DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=new_column_names,
         snowpark_column_names=snowpark_new_column_names,
+        parent_column_name_map=input_container.column_map,
+        table_name=input_container.table_name,
+        alias=input_container.alias,
+        cached_schema_getter=_get_schema,
     )
-    context.set_df_before_projection(result_with_column_map)
-    return result_with_column_map
+    context.set_df_before_projection(result_container)
+    return result_container
 
 
-def map_to_schema(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_to_schema(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Transform the column names of the input DataFrame.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.to_schema.input)
+    input_container = map_relation(rel.to_schema.input)
+    input_df = input_container.dataframe
     new_column_names = [field.name for field in rel.to_schema.schema.struct.fields]
     snowpark_new_column_names = make_column_names_snowpark_compatible(
         new_column_names, rel.common.plan_id
     )
     count_case_insensitive_column_names = defaultdict()
-    for key, value in input_df._column_map.spark_to_col.items():
+    for key, value in input_container.column_map.spark_to_col.items():
         count_case_insensitive_column_names[
             key.lower()
         ] = count_case_insensitive_column_names.get(key.lower(), 0) + len(value)
@@ -483,12 +506,12 @@ def map_to_schema(rel: relation_proto.Relation) -> snowpark.DataFrame:
         if field.name in already_existing_columns:
             if count_case_insensitive_column_names[field.name.lower()] > 1:
                 raise AnalysisException(
-                    f"[AMBIGUOUS_COLUMN_OR_FIELD] Column or field `{field.name}` is ambiguous and has {len(input_df._column_map.spark_to_col[field.name])} matches."
+                    f"[AMBIGUOUS_COLUMN_OR_FIELD] Column or field `{field.name}` is ambiguous and has {len(input_container.column_map.spark_to_col[field.name])} matches."
                 )
             snowpark_name = None
-            for name in input_df._column_map.spark_to_col:
+            for name in input_container.column_map.spark_to_col:
                 if name.lower() == field.name.lower():
-                    snowpark_name = input_df._column_map.spark_to_col[name][
+                    snowpark_name = input_container.column_map.spark_to_col[name][
                         0
                     ].snowpark_name
                     break
@@ -516,10 +539,10 @@ def map_to_schema(rel: relation_proto.Relation) -> snowpark.DataFrame:
         # All columns already exist, we're doing a simple update.
         snowpark_new_column_names = []
         for column in new_column_names:
-            for name in input_df._column_map.spark_to_col:
+            for name in input_container.column_map.spark_to_col:
                 if name.lower() == column.lower():
                     snowpark_new_column_names.append(
-                        input_df._column_map.spark_to_col[name][0].snowpark_name
+                        input_container.column_map.spark_to_col[name][0].snowpark_name
                     )
         result = input_df
     elif len(already_existing_columns) == 0:
@@ -540,16 +563,18 @@ def map_to_schema(rel: relation_proto.Relation) -> snowpark.DataFrame:
             # If the column doesn't already exist, append the new Snowpark name to columns_to_add
             if all(
                 spark_column.lower() != name.lower()
-                for name in input_df._column_map.spark_to_col
+                for name in input_container.column_map.spark_to_col
             ):
                 columns_to_add.append(snowpark_column)
                 new_snowpark_new_column_names.append(snowpark_column)
             else:
-                for name in input_df._column_map.spark_to_col:
+                for name in input_container.column_map.spark_to_col:
                     # If the column does exist, append the original Snowpark name, We don't need to add this column.
                     if name.lower() == spark_column.lower():
                         new_snowpark_new_column_names.append(
-                            input_df._column_map.spark_to_col[name][0].snowpark_name
+                            input_container.column_map.spark_to_col[name][
+                                0
+                            ].snowpark_name
                         )
         # Add all columns introduced by the new schema.
         new_columns = [
@@ -581,21 +606,24 @@ def map_to_schema(rel: relation_proto.Relation) -> snowpark.DataFrame:
                 column_metadata[field.name] = None
         else:
             column_metadata[field.name] = None
-    return with_column_map(
-        result_with_casting,
-        new_column_names,
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result_with_casting,
+        spark_column_names=new_column_names,
         snowpark_column_names=snowpark_new_column_names,
         snowpark_column_types=[field.datatype for field in snowpark_schema.fields],
         column_metadata=column_metadata,
-        parent_column_name_map=input_df._column_map,
+        parent_column_name_map=input_container.column_map,
     )
 
 
-def map_with_columns_renamed(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_with_columns_renamed(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
-    Rename columns in a DataFrame.
+    Rename columns in a DataFrame and return a container.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.with_columns_renamed.input)
+    input_container = map_relation(rel.with_columns_renamed.input)
+    input_df = input_container.dataframe
     rename_columns_map = dict(rel.with_columns_renamed.rename_columns_map)
 
     if not global_config.spark_sql_caseSensitive:
@@ -608,9 +636,11 @@ def map_with_columns_renamed(rel: relation_proto.Relation) -> snowpark.DataFrame
             k.lower(): v.lower() for k, v in rename_columns_map.items()
         }
 
+    column_map = input_container.column_map
+
     # re-construct the rename chains based on the input dataframe.
-    if input_df._column_map.rename_chains:
-        for key, value in input_df._column_map.rename_chains.items():
+    if input_container.column_map.rename_chains:
+        for key, value in input_container.column_map.rename_chains.items():
             if key in rename_columns_map:
                 # This is to handle the case where the same column is renamed multiple times.
                 # df.withColumnRenamed("a", "b").withColumnRenamed("a", "c")
@@ -628,19 +658,25 @@ def map_with_columns_renamed(rel: relation_proto.Relation) -> snowpark.DataFrame
                 # This just copies the renames from previous computed dataframe
                 rename_columns_map[key] = value
 
-    existing_columns = input_df._column_map.get_spark_columns()
+    existing_columns = input_container.column_map.get_spark_columns()
+
+    def _column_exists_error(name: str) -> AnalysisException:
+        return AnalysisException(
+            f"[COLUMN_ALREADY_EXISTS] The column `{name}` already exists. Consider to choose another name or rename the existing column."
+        )
 
     # Validate for naming conflicts
     new_names_list = list(dict(rel.with_columns_renamed.rename_columns_map).values())
     seen = set()
     for new_name in new_names_list:
+        if column_map.has_spark_column(new_name):
+            # Spark doesn't allow reusing existing names, even if the result df will not contain duplicate columns
+            raise _column_exists_error(new_name)
         if (global_config.spark_sql_caseSensitive and new_name in seen) or (
             not global_config.spark_sql_caseSensitive
             and new_name.lower() in [s.lower() for s in seen]
         ):
-            raise AnalysisException(
-                f"[COLUMN_ALREADY_EXISTS] The column `{new_name}` already exists. Consider to choose another name or rename the existing column."
-            )
+            raise _column_exists_error(new_name)
         seen.add(new_name)
 
     new_columns = []
@@ -656,25 +692,30 @@ def map_with_columns_renamed(rel: relation_proto.Relation) -> snowpark.DataFrame
 
     # Creating a new df to avoid updating the state of cached dataframe.
     new_df = input_df.select("*")
-    result_df = with_column_map(
-        new_df,
-        new_columns,
-        input_df._column_map.get_snowpark_columns(),
-        column_qualifiers=input_df._column_map.get_qualifiers(),
-        parent_column_name_map=input_df._column_map.get_parent_column_name_map(),
+    result_container = DataFrameContainer.create_with_column_mapping(
+        dataframe=new_df,
+        spark_column_names=new_columns,
+        snowpark_column_names=input_container.column_map.get_snowpark_columns(),
+        column_qualifiers=input_container.column_map.get_qualifiers(),
+        parent_column_name_map=input_container.column_map.get_parent_column_name_map(),
+        table_name=input_container.table_name,
+        alias=input_container.alias,
     )
-    result_df._column_map.rename_chains = rename_columns_map
+    result_container.column_map.rename_chains = rename_columns_map
 
-    return result_df
+    return result_container
 
 
-def map_with_columns(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_with_columns(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
-    Add columns to a DataFrame.
+    Add columns to a DataFrame and return a container.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.with_columns.input)
+    input_container = map_relation(rel.with_columns.input)
+    input_df = input_container.dataframe
     with_columns = [
-        map_alias(alias, input_df._column_map, ExpressionTyper(input_df))
+        map_alias(alias, input_container.column_map, ExpressionTyper(input_df))
         for alias in rel.with_columns.aliases
     ]
     # TODO: This list needs to contain all unique column names, but the code below doesn't
@@ -682,7 +723,7 @@ def map_with_columns(rel: relation_proto.Relation) -> snowpark.DataFrame:
     with_columns_names = []
     with_columns_exprs = []
     with_columns_types = []
-    with_column_offset = len(input_df._column_map.get_spark_columns())
+    with_column_offset = len(input_container.column_map.get_spark_columns())
     new_spark_names = []
     seen_columns = set()
     for names_list, expr in with_columns:
@@ -690,7 +731,7 @@ def map_with_columns(rel: relation_proto.Relation) -> snowpark.DataFrame:
             len(names_list) == 1
         ), f"Expected single column name, got {len(names_list)}: {names_list}"
         name = names_list[0]
-        name_normalized = input_df._column_map._normalized_spark_name(name)
+        name_normalized = input_container.column_map._normalized_spark_name(name)
         if name_normalized in seen_columns:
             raise ValueError(
                 f"[COLUMN_ALREADY_EXISTS] The column `{name}` already exists."
@@ -698,11 +739,9 @@ def map_with_columns(rel: relation_proto.Relation) -> snowpark.DataFrame:
         seen_columns.add(name_normalized)
         # If the column name is already in the DataFrame, we replace it, so we use the
         # mapping to get the correct column name.
-        if input_df._column_map.has_spark_column(name):
-            all_instances_of_spark_column_name = (
-                input_df._column_map.get_snowpark_column_names_from_spark_column_names(
-                    [name]
-                )
+        if input_container.column_map.has_spark_column(name):
+            all_instances_of_spark_column_name = input_container.column_map.get_snowpark_column_names_from_spark_column_names(
+                [name]
             )
             if len(all_instances_of_spark_column_name) == 0:
                 raise KeyError(f"Spark column name {name} does not exist")
@@ -729,7 +768,7 @@ def map_with_columns(rel: relation_proto.Relation) -> snowpark.DataFrame:
         new_spark_columns,
         new_snowpark_columns,
         qualifiers,
-    ) = input_df._column_map.with_columns(new_spark_names, with_columns_names)
+    ) = input_container.column_map.with_columns(new_spark_names, with_columns_names)
 
     # dedup the change in columns at snowpark name level, this is required by the with columns functions
     with_columns_names_deduped = []
@@ -751,32 +790,39 @@ def map_with_columns(rel: relation_proto.Relation) -> snowpark.DataFrame:
         + list(zip(with_columns_names, with_columns_types))
     )
 
-    column_metadata = input_df._column_map.column_metadata or {}
+    column_metadata = input_container.column_map.column_metadata or {}
     for alias in rel.with_columns.aliases:
         # this logic is triggered for df.withMetadata function.
         if alias.HasField("metadata") and len(alias.metadata.strip()) > 0:
             # spark sends list of alias names with only one element in the list with alias name.
             column_metadata[alias.name[0]] = json.loads(alias.metadata)
 
-    return with_column_map(
-        result,
-        new_spark_columns,
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=new_spark_columns,
         snowpark_column_names=new_snowpark_columns,
         snowpark_column_types=[
             snowpark_name_to_type.get(n) for n in new_snowpark_columns
         ],
         column_metadata=column_metadata,
         column_qualifiers=qualifiers,
-        parent_column_name_map=input_df._column_map,
+        parent_column_name_map=input_container.column_map,
+        table_name=input_container.table_name,
+        alias=input_container.alias,
     )
 
 
-def map_unpivot(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_unpivot(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     # Spark API:    df.unpivot([id_columns], [unpivot_columns], var_column, val_column)
     # Snowpark API: df.unpivot(val_column, var_column, [unpivot_columns])
     if rel.unpivot.HasField("values") and len(rel.unpivot.values.values) == 0:
         raise SparkException.unpivot_requires_value_columns()
 
+    input_container = map_relation(rel.unpivot.input)
+    input_df = input_container.dataframe
+
     def get_lease_common_ancestor_classes(types: list[snowpark.types.DataType]) -> set:
         mro_lists = [set(type.__class__.mro()) for type in types]
         common_ancestors = set.intersection(*mro_lists)
@@ -795,12 +841,15 @@ def map_unpivot(rel: relation_proto.Relation) -> snowpark.DataFrame:
         type_column_list = [
             (
                 f.datatype,
-                df._column_map.get_spark_column_name_from_snowpark_column_name(
-                    snowpark_functions_col(f.name, df._column_map).get_name()
+                input_container.column_map.get_spark_column_name_from_snowpark_column_name(
+                    snowpark_functions_col(
+                        f.name, input_container.column_map
+                    ).get_name()
                 ),
             )
             for f in df.schema.fields
-            if snowpark_functions_col(f.name, df._column_map).get_name() in col_names
+            if snowpark_functions_col(f.name, input_container.column_map).get_name()
+            in col_names
         ]
         type_iter, _ = zip(*type_column_list)
         type_list = list(type_iter)
@@ -837,7 +886,7 @@ def map_unpivot(rel: relation_proto.Relation) -> snowpark.DataFrame:
         typer = ExpressionTyper(input_df)
         for id_col in relation.unpivot.ids:
             spark_name, typed_column = map_single_column_expression(
-                id_col, df._column_map, typer
+                id_col, input_container.column_map, typer
             )
             id_col_names.append(typed_column.col.get_name())
             spark_columns.append(spark_name)
@@ -848,7 +897,7 @@ def map_unpivot(rel: relation_proto.Relation) -> snowpark.DataFrame:
         unpivot_spark_names = []
         for v in relation.unpivot.values.values:
             spark_name, typed_column = map_single_column_expression(
-                v, df._column_map, typer
+                v, input_container.column_map, typer
             )
             unpivot_col_names.append(typed_column.col.get_name())
             unpivot_spark_names.append(spark_name)
@@ -856,15 +905,19 @@ def map_unpivot(rel: relation_proto.Relation) -> snowpark.DataFrame:
         if not rel.unpivot.HasField("values"):
             # When `values` is `None`, all non-id columns will be unpivoted.
             for snowpark_name, spark_name in zip(
-                df._column_map.get_snowpark_columns(),
-                df._column_map.get_spark_columns(),
+                input_container.column_map.get_snowpark_columns(),
+                input_container.column_map.get_spark_columns(),
             ):
                 if (
-                    snowpark_functions_col(snowpark_name, df._column_map).get_name()
+                    snowpark_functions_col(
+                        snowpark_name, input_container.column_map
+                    ).get_name()
                     not in id_col_names
                 ):
                     unpivot_col_names.append(
-                        snowpark_functions_col(snowpark_name, df._column_map).get_name()
+                        snowpark_functions_col(
+                            snowpark_name, input_container.column_map
+                        ).get_name()
                     )
                     unpivot_spark_names.append(spark_name)
 
@@ -872,7 +925,6 @@ def map_unpivot(rel: relation_proto.Relation) -> snowpark.DataFrame:
         spark_columns.append(relation.unpivot.value_column_name)
         return spark_columns, id_col_names, unpivot_col_names, unpivot_spark_names
 
-    input_df: snowpark.DataFrame = map_relation(rel.unpivot.input)
     (
         spark_columns,
         id_col_names,
@@ -899,27 +951,35 @@ def map_unpivot(rel: relation_proto.Relation) -> snowpark.DataFrame:
     column_reverse_project = []
     snowpark_columns = []
     qualifiers = []
-    for c in input_df._column_map.get_snowpark_columns():
-        c_name = snowpark_functions_col(c, input_df._column_map).get_name()
+    for c in input_container.column_map.get_snowpark_columns():
+        c_name = snowpark_functions_col(c, input_container.column_map).get_name()
         if c_name in unpivot_col_names:
             if cast_type:
                 column_project.append(
-                    snowpark_functions_col(c, input_df._column_map)
+                    snowpark_functions_col(c, input_container.column_map)
                     .cast("DOUBLE")
                     .alias(c_name)
                 )
             else:
-                column_project.append(snowpark_functions_col(c, input_df._column_map))
+                column_project.append(
+                    snowpark_functions_col(c, input_container.column_map)
+                )
         if c_name in id_col_names:
             id_col_alias = "SES" + generate_random_alphanumeric().upper()
             column_project.append(
-                snowpark_functions_col(c, input_df._column_map).alias(id_col_alias)
+                snowpark_functions_col(c, input_container.column_map).alias(
+                    id_col_alias
+                )
             )
             column_reverse_project.append(
-                snowpark_functions_col(id_col_alias, input_df._column_map).alias(c)
+                snowpark_functions_col(id_col_alias, input_container.column_map).alias(
+                    c
+                )
             )
             snowpark_columns.append(c)
-            qualifiers.append(input_df._column_map.get_qualifier_for_spark_column(c))
+            qualifiers.append(
+                input_container.column_map.get_qualifier_for_spark_column(c)
+            )
 
     # Without the case when postprocessing, the result Spark dataframe is:
     # +---+------------+------+
@@ -941,7 +1001,7 @@ def map_unpivot(rel: relation_proto.Relation) -> snowpark.DataFrame:
         if post_process_variable_column is None:
             post_process_variable_column = snowpark_fn.when(
                 snowpark_functions_col(
-                    snowpark_variable_column_name, input_df._column_map
+                    snowpark_variable_column_name, input_container.column_map
                 )
                 == unquote_if_quoted(snowpark_name),
                 spark_name,
@@ -949,7 +1009,7 @@ def map_unpivot(rel: relation_proto.Relation) -> snowpark.DataFrame:
         else:
             post_process_variable_column = post_process_variable_column.when(
                 snowpark_functions_col(
-                    snowpark_variable_column_name, input_df._column_map
+                    snowpark_variable_column_name, input_container.column_map
                 )
                 == unquote_if_quoted(snowpark_name),
                 spark_name,
@@ -960,7 +1020,7 @@ def map_unpivot(rel: relation_proto.Relation) -> snowpark.DataFrame:
     )
     snowpark_columns.append(snowpark_variable_column_name)
     column_reverse_project.append(
-        snowpark_functions_col(snowpark_value_column_name, input_df._column_map)
+        snowpark_functions_col(snowpark_value_column_name, input_container.column_map)
     )
     snowpark_columns.append(snowpark_value_column_name)
     qualifiers.extend([[]] * 2)
@@ -975,20 +1035,23 @@ def map_unpivot(rel: relation_proto.Relation) -> snowpark.DataFrame:
         )
         .select(*column_reverse_project)
     )
-    return with_column_map(
-        result,
-        spark_columns,
-        snowpark_columns,
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=spark_columns,
+        snowpark_column_names=snowpark_columns,
         column_qualifiers=qualifiers,
-        parent_column_name_map=input_df._column_map,
+        parent_column_name_map=input_container.column_map,
     )
 
 
-def map_group_map(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_group_map(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Add columns to a DataFrame.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.group_map.input)
+    input_container = map_relation(rel.group_map.input)
+    input_df = input_container.dataframe
     grouping_expressions = rel.group_map.grouping_expressions
     snowpark_grouping_expressions: list[snowpark.Column] = []
     typer = ExpressionTyper(input_df)
@@ -996,7 +1059,7 @@ def map_group_map(rel: relation_proto.Relation) -> snowpark.DataFrame:
     qualifiers = []
     for exp in grouping_expressions:
         new_name, snowpark_column = map_single_column_expression(
-            exp, input_df._column_map, typer
+            exp, input_container.column_map, typer
         )
         snowpark_grouping_expressions.append(snowpark_column.col)
         group_name_list.append(new_name)
@@ -1013,9 +1076,9 @@ def map_group_map(rel: relation_proto.Relation) -> snowpark.DataFrame:
 
     if not is_compatible_python or TEST_FLAG_FORCE_CREATE_SPROC:
         original_columns = None
-        if input_df._column_map is not None:
+        if input_container.column_map is not None:
             original_columns = [
-                column.spark_name for column in input_df._column_map.columns
+                column.spark_name for column in input_container.column_map.columns
             ]
 
         apply_udtf_temp_name = create_apply_udtf_in_sproc(
@@ -1044,12 +1107,12 @@ def map_group_map(rel: relation_proto.Relation) -> snowpark.DataFrame:
         )
 
     qualifiers.extend([[]] * (len(result.columns) - len(group_name_list)))
-    return with_column_map(
-        result,
-        [field.name for field in output_type],
-        result.columns,
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=[field.name for field in output_type],
+        snowpark_column_names=result.columns,
         column_qualifiers=qualifiers,
-        parent_column_name_map=input_df._column_map,
+        parent_column_name_map=input_container.column_map,
     )