snowpark-connect 0.27.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

snowflake/snowpark_connect/relation/map_row_ops.py

@@ -1,15 +1,15 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
-
-
+import pandas
 import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 from pyspark.errors.exceptions.base import AnalysisException, IllegalArgumentException
 
 import snowflake.snowpark_connect.relation.utils as utils
 from snowflake import snowpark
-from snowflake.snowpark.functions import col, expr as snowpark_expr
+from snowflake.snowpark._internal.error_message import SnowparkClientExceptionMessages
+from snowflake.snowpark.functions import col, expr as snowpark_expr, lit
 from snowflake.snowpark.types import (
     BooleanType,
     ByteType,
@@ -20,21 +20,90 @@ from snowflake.snowpark.types import (
     LongType,
     NullType,
     ShortType,
+    StructField,
+    StructType,
+)
+from snowflake.snowpark_connect.column_name_handler import (
+    ColumnNameMap,
+    schema_getter,
+    set_schema_getter,
 )
-from snowflake.snowpark_connect.column_name_handler import ColumnNameMap, schema_getter
 from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 from snowflake.snowpark_connect.expression.literal import get_literal_field_and_name
 from snowflake.snowpark_connect.expression.map_expression import (
     map_single_column_expression,
 )
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.relation.map_relation import map_relation
+from snowflake.snowpark_connect.relation.read.metadata_utils import (
+    without_internal_columns,
+)
+from snowflake.snowpark_connect.utils.identifiers import (
+    split_fully_qualified_spark_name,
+)
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
 
 
+def cast_columns(
+    df_container: DataFrameContainer,
+    df_dtypes: list[snowpark.types.DataType],
+    target_dtypes: list[snowpark.types.DataType],
+    column_map: ColumnNameMap,
+):
+    df: snowpark.DataFrame = df_container.dataframe
+    if df_dtypes == target_dtypes:
+        return df_container
+    # Use cached schema if available to avoid triggering extra queries
+    if (
+        hasattr(df_container, "cached_schema_getter")
+        and df_container.cached_schema_getter is not None
+    ):
+        df_schema = df_container.cached_schema_getter()
+    else:
+        df_schema = df.schema  # Get current schema
+    new_columns = []
+
+    for i, field in enumerate(df_schema.fields):
+        col_name = field.name
+        current_type = field.datatype
+        target_type = target_dtypes[i]
+
+        if current_type != target_type:
+            new_columns.append(df[col_name].cast(target_type).alias(col_name))
+        else:
+            new_columns.append(df[col_name])
+
+    new_df = df.select(new_columns)
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=new_df,
+        spark_column_names=column_map.get_spark_columns(),
+        snowpark_column_names=column_map.get_snowpark_columns(),
+        snowpark_column_types=target_dtypes,
+        column_metadata=column_map.column_metadata,
+        parent_column_name_map=column_map,
+    )
+
+
+def get_schema_from_result(
+    result: DataFrameContainer,
+) -> StructType:
+    """
+    Get schema from a DataFrameContainer, using cached schema if available to avoid extra queries.
+    """
+    if (
+        hasattr(result, "cached_schema_getter")
+        and result.cached_schema_getter is not None
+    ):
+        return result.cached_schema_getter()
+    else:
+        return result.dataframe.schema
+
+
 def map_deduplicate(
     rel: relation_proto.Relation,
 ) -> DataFrameContainer:
@@ -43,16 +112,18 @@ def map_deduplicate(
 
     The deduplicate is a list of columns that is applied to the DataFrame.
     """
-    input_container = map_relation(rel.deduplicate.input)
+    input_container = without_internal_columns(map_relation(rel.deduplicate.input))
     input_df = input_container.dataframe
 
     if (
         rel.deduplicate.HasField("within_watermark")
         and rel.deduplicate.within_watermark
     ):
-        raise AnalysisException(
+        exception = AnalysisException(
             "dropDuplicatesWithinWatermark is not supported with batch DataFrames/DataSets"
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     if (
         rel.deduplicate.HasField("all_columns_as_keys")
@@ -81,7 +152,7 @@ def map_dropna(
     """
     Drop NA values from the input DataFrame.
     """
-    input_container = map_relation(rel.drop_na.input)
+    input_container = without_internal_columns(map_relation(rel.drop_na.input))
     input_df = input_container.dataframe
 
     if rel.drop_na.HasField("min_non_nulls"):
@@ -122,15 +193,23 @@ def map_fillna(
 
     The `fill_value` is a scalar value that will be used to replace NaN values.
     """
-    input_container = map_relation(rel.fill_na.input)
+    input_container = without_internal_columns(map_relation(rel.fill_na.input))
     input_df = input_container.dataframe
 
     if len(rel.fill_na.cols) > 0:
+        if rel.fill_na.cols == ["*"]:
+            # Expand "*" to all columns
+            spark_col_names = input_container.column_map.get_spark_columns()
+        else:
+            spark_col_names = list(rel.fill_na.cols)
+
+        # We don't validate the fully qualified spark name here as fillNa is no-op for structured type colums.
+        # It only works for scalar type columns like float, int, string or bool.
         columns: list[str] = [
             input_container.column_map.get_snowpark_column_name_from_spark_column_name(
-                c
+                split_fully_qualified_spark_name(c)[0]
             )
-            for c in rel.fill_na.cols
+            for c in spark_col_names
         ]
         values = [get_literal_field_and_name(v)[0] for v in rel.fill_na.values]
         if len(values) == 1:
@@ -177,29 +256,16 @@ def map_union(
 
     The two DataFrames must have the same schema.
     """
-    left_result = map_relation(rel.set_op.left_input)
-    right_result = map_relation(rel.set_op.right_input)
+    left_result = without_internal_columns(map_relation(rel.set_op.left_input))
+    right_result = without_internal_columns(map_relation(rel.set_op.right_input))
     left_df = left_result.dataframe
     right_df = right_result.dataframe
     allow_missing_columns = bool(rel.set_op.allow_missing_columns)
 
     # workaround for unstructured type vs structured type
     # Use cached schema if available to avoid triggering extra queries
-    if (
-        hasattr(left_result, "cached_schema_getter")
-        and left_result.cached_schema_getter is not None
-    ):
-        left_schema = left_result.cached_schema_getter()
-    else:
-        left_schema = left_df.schema
-
-    if (
-        hasattr(right_result, "cached_schema_getter")
-        and right_result.cached_schema_getter is not None
-    ):
-        right_schema = right_result.cached_schema_getter()
-    else:
-        right_schema = right_df.schema
+    left_schema = get_schema_from_result(left_result)
+    right_schema = get_schema_from_result(right_result)
 
     left_dtypes = [field.datatype for field in left_schema.fields]
     right_dtypes = [field.datatype for field in right_schema.fields]
@@ -207,7 +273,9 @@ def map_union(
     spark_sql_ansi_enabled = global_config.spark_sql_ansi_enabled
     if left_dtypes != right_dtypes and not rel.set_op.by_name:
         if len(left_dtypes) != len(right_dtypes):
-            raise AnalysisException("UNION: the number of columns must match")
+            exception = AnalysisException("UNION: the number of columns must match")
+            attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+            raise exception
         target_left_dtypes, target_right_dtypes = [], []
         for left_type, right_type in zip(left_dtypes, right_dtypes):
             match (left_type, right_type):
@@ -235,6 +303,29 @@ def map_union(
                     # Union of any type with null type is of the other type
                     target_left_dtypes.append(other_t)
                     target_right_dtypes.append(other_t)
+                case (snowpark.types.DecimalType(), snowpark.types.DecimalType()):
+                    # Widen decimal types to accommodate both sides
+                    # Calculate the maximum scale and maximum integer digits
+                    left_integer_digits = left_type.precision - left_type.scale
+                    right_integer_digits = right_type.precision - right_type.scale
+
+                    # The common type needs to accommodate:
+                    # - The maximum number of digits after the decimal point (scale)
+                    # - The maximum number of digits before the decimal point (integer digits)
+                    common_scale = max(left_type.scale, right_type.scale)
+                    common_integer_digits = max(
+                        left_integer_digits, right_integer_digits
+                    )
+                    common_precision = min(38, common_scale + common_integer_digits)
+
+                    # Ensure scale doesn't exceed precision
+                    common_scale = min(common_scale, common_precision)
+
+                    common_type = snowpark.types.DecimalType(
+                        common_precision, common_scale
+                    )
+                    target_left_dtypes.append(common_type)
+                    target_right_dtypes.append(common_type)
                 case (snowpark.types.BooleanType(), _) | (
                     _,
                     snowpark.types.BooleanType(),
@@ -243,54 +334,31 @@ def map_union(
                         not spark_sql_ansi_enabled
                         or snowpark.types.StringType() not in [left_type, right_type]
                     ):  # In ansi mode , string type union boolean type is acceptable
-                        raise AnalysisException(
+                        exception = AnalysisException(
                             f"""[INCOMPATIBLE_COLUMN_TYPE] UNION can only be performed on tables with compatible column types. "{str(left_type)}" type which is not compatible with "{str(right_type)}". """
                         )
+                        attach_custom_error_code(exception, ErrorCodes.TYPE_MISMATCH)
+                        raise exception
                     target_left_dtypes.append(left_type)
                     target_right_dtypes.append(right_type)
+                case (
+                    snowpark.types.TimestampType()
+                    | snowpark.types.DateType()
+                    | snowpark.types._NumericType(),
+                    snowpark.types.StringType(),
+                ) | (
+                    snowpark.types.StringType(),
+                    snowpark.types.TimestampType()
+                    | snowpark.types.DateType()
+                    | snowpark.types._NumericType(),
+                ) if not spark_sql_ansi_enabled:
+                    common_type = snowpark.types.StringType()
+                    target_left_dtypes.append(common_type)
+                    target_right_dtypes.append(common_type)
                 case _:
                     target_left_dtypes.append(left_type)
                     target_right_dtypes.append(right_type)
 
-        def cast_columns(
-            df_container: DataFrameContainer,
-            df_dtypes: list[snowpark.types.DataType],
-            target_dtypes: list[snowpark.types.DataType],
-            column_map: ColumnNameMap,
-        ):
-            df: snowpark.DataFrame = df_container.dataframe
-            if df_dtypes == target_dtypes:
-                return df_container
-            # Use cached schema if available to avoid triggering extra queries
-            if (
-                hasattr(df_container, "cached_schema_getter")
-                and df_container.cached_schema_getter is not None
-            ):
-                df_schema = df_container.cached_schema_getter()
-            else:
-                df_schema = df.schema  # Get current schema
-            new_columns = []
-
-            for i, field in enumerate(df_schema.fields):
-                col_name = field.name
-                current_type = field.datatype
-                target_type = target_dtypes[i]
-
-                if current_type != target_type:
-                    new_columns.append(df[col_name].cast(target_type).alias(col_name))
-                else:
-                    new_columns.append(df[col_name])
-
-            new_df = df.select(new_columns)
-            return DataFrameContainer.create_with_column_mapping(
-                dataframe=new_df,
-                spark_column_names=column_map.get_spark_columns(),
-                snowpark_column_names=column_map.get_snowpark_columns(),
-                snowpark_column_types=target_dtypes,
-                column_metadata=column_map.column_metadata,
-                parent_column_name_map=column_map,
-            )
-
         left_result = cast_columns(
             left_result,
             left_dtypes,
@@ -318,23 +386,37 @@ def map_union(
         right_column_map = right_result.column_map
         columns_to_restore: dict[str, tuple[str, str]] = {}
 
-        for column in right_df.columns:
+        original_right_schema = right_df.schema
+        right_renamed_fields = []
+        for field in original_right_schema.fields:
             spark_name = (
-                right_column_map.get_spark_column_name_from_snowpark_column_name(column)
+                right_column_map.get_spark_column_name_from_snowpark_column_name(
+                    field.name
+                )
+            )
+            right_df = right_df.withColumnRenamed(field.name, spark_name)
+            columns_to_restore[spark_name.upper()] = (spark_name, field.name)
+            right_renamed_fields.append(
+                StructField(spark_name, field.datatype, field.nullable)
             )
-            right_df = right_df.withColumnRenamed(column, spark_name)
-            columns_to_restore[spark_name.upper()] = (spark_name, column)
+        set_schema_getter(right_df, lambda: StructType(right_renamed_fields))
 
-        for column in left_df.columns:
+        original_left_schema = left_df.schema
+        left_renamed_fields = []
+        for field in original_left_schema.fields:
             spark_name = (
-                left_column_map.get_spark_column_name_from_snowpark_column_name(column)
+                left_column_map.get_spark_column_name_from_snowpark_column_name(
+                    field.name
+                )
+            )
+            left_df = left_df.withColumnRenamed(field.name, spark_name)
+            columns_to_restore[spark_name.upper()] = (spark_name, field.name)
+            left_renamed_fields.append(
+                StructField(spark_name, field.datatype, field.nullable)
             )
-            left_df = left_df.withColumnRenamed(column, spark_name)
-            columns_to_restore[spark_name.upper()] = (spark_name, column)
+        set_schema_getter(left_df, lambda: StructType(left_renamed_fields))
 
-        result = left_df.unionAllByName(
-            right_df, allow_missing_columns=allow_missing_columns
-        )
+        result = _union_by_name_optimized(left_df, right_df, allow_missing_columns)
 
         if allow_missing_columns:
             spark_columns = []
@@ -421,8 +503,8 @@ def map_intersect(
     |  b|  3|
     +---+---+
     """
-    left_result = map_relation(rel.set_op.left_input)
-    right_result = map_relation(rel.set_op.right_input)
+    left_result = without_internal_columns(map_relation(rel.set_op.left_input))
+    right_result = without_internal_columns(map_relation(rel.set_op.right_input))
     left_df = left_result.dataframe
     right_df = right_result.dataframe
 
@@ -484,11 +566,53 @@ def map_except(
     |  c|  4|
     +---+---+
     """
-    left_result = map_relation(rel.set_op.left_input)
-    right_result = map_relation(rel.set_op.right_input)
+    left_result = without_internal_columns(map_relation(rel.set_op.left_input))
+    right_result = without_internal_columns(map_relation(rel.set_op.right_input))
     left_df = left_result.dataframe
     right_df = right_result.dataframe
 
+    # workaround for unstructured type vs structured type
+    # Use cached schema if available to avoid triggering extra queries
+    left_schema = get_schema_from_result(left_result)
+    right_schema = get_schema_from_result(right_result)
+
+    left_dtypes = [field.datatype for field in left_schema.fields]
+    right_dtypes = [field.datatype for field in right_schema.fields]
+
+    if left_dtypes != right_dtypes and not rel.set_op.by_name:
+        if len(left_dtypes) != len(right_dtypes):
+            exception = AnalysisException("UNION: the number of columns must match")
+            attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+            raise exception
+        target_left_dtypes, target_right_dtypes = [], []
+        for left_type, right_type in zip(left_dtypes, right_dtypes):
+            match (left_type, right_type):
+                case (snowpark.types._NumericType(), snowpark.types.StringType()) | (
+                    snowpark.types.StringType(),
+                    snowpark.types._NumericType(),
+                ):
+                    common_type = snowpark.types.StringType()
+                    target_left_dtypes.append(common_type)
+                    target_right_dtypes.append(common_type)
+                case _:
+                    target_left_dtypes.append(left_type)
+                    target_right_dtypes.append(right_type)
+
+        left_result = cast_columns(
+            left_result,
+            left_dtypes,
+            target_left_dtypes,
+            left_result.column_map,
+        )
+        right_result = cast_columns(
+            right_result,
+            right_dtypes,
+            target_right_dtypes,
+            right_result.column_map,
+        )
+        left_df = left_result.dataframe
+        right_df = right_result.dataframe
+
     if rel.set_op.is_all:
         # Snowflake except removes all duplicated rows. In order to handle the case,
         # we add a partition row number column to the df to make duplicated rows unique to
@@ -573,13 +697,18 @@ def map_filter(
 
 def map_limit(
     rel: relation_proto.Relation,
-) -> DataFrameContainer:
+) -> DataFrameContainer | pandas.DataFrame:
     """
     Limit a DataFrame based on a Relation's limit.
 
     The limit is an integer that is applied to the DataFrame.
     """
-    input_container = map_relation(rel.limit.input)
+
+    input_container = without_internal_columns(map_relation(rel.limit.input))
+
+    if isinstance(input_container, pandas.DataFrame):
+        return input_container.head(rel.limit.limit)
+
     input_df = input_container.dataframe
 
     result: snowpark.DataFrame = input_df.limit(rel.limit.limit)
@@ -601,7 +730,7 @@ def map_offset(
 
     The offset is an integer that is applied to the DataFrame.
     """
-    input_container = map_relation(rel.offset.input)
+    input_container = without_internal_columns(map_relation(rel.offset.input))
     input_df = input_container.dataframe
 
     # TODO: This is a terrible way to have to do this, but Snowpark does not
@@ -629,7 +758,7 @@ def map_replace(
     values to replace. The values in the dictionary are the values to replace
     and the keys are the values to replace them with.
     """
-    result = map_relation(rel.replace.input)
+    result = without_internal_columns(map_relation(rel.replace.input))
     input_df = result.dataframe
     ordered_columns = input_df.columns
     column_map = result.column_map
@@ -752,12 +881,14 @@ def map_sample(
     """
     Sample a DataFrame based on a Relation's sample.
     """
-    input_container = map_relation(rel.sample.input)
+    input_container = without_internal_columns(map_relation(rel.sample.input))
     input_df = input_container.dataframe
 
     frac = rel.sample.upper_bound - rel.sample.lower_bound
     if frac < 0 or frac > 1:
-        raise IllegalArgumentException("Sample fraction must be between 0 and 1")
+        exception = IllegalArgumentException("Sample fraction must be between 0 and 1")
+        attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+        raise exception
     # The seed argument is not supported here. There are a number of reasons that implementing
     # this will be complicated in Snowflake. Here is a list of complications:
     #
@@ -772,9 +903,11 @@ def map_sample(
     # these issues.
     if rel.sample.with_replacement:
         # TODO: Use a random number generator with ROW_NUMBER and SELECT.
-        raise SnowparkConnectNotImplementedError(
+        exception = SnowparkConnectNotImplementedError(
             "Sample with replacement is not supported"
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
     else:
         result: snowpark.DataFrame = input_df.sample(frac=frac)
         return DataFrameContainer(
@@ -794,7 +927,7 @@ def map_tail(
 
     The tail is an integer that is applied to the DataFrame.
     """
-    input_container = map_relation(rel.tail.input)
+    input_container = without_internal_columns(map_relation(rel.tail.input))
     input_df = input_container.dataframe
 
     num_rows = input_df.count()
@@ -809,3 +942,89 @@ def map_tail(
         alias=input_container.alias,
         cached_schema_getter=lambda: input_df.schema,
     )
+
+
+def _union_by_name_optimized(
+    left_df: snowpark.DataFrame,
+    right_df: snowpark.DataFrame,
+    allow_missing_columns: bool = False,
+) -> snowpark.DataFrame:
+    """
+    This implementation is an optimized version of Snowpark's Dataframe::_union_by_name_internal.
+    The only change is, that it avoids redundant schema queries that occur in the standard Snowpark,
+    by reusing already-fetched/calculated schemas.
+    """
+
+    left_schema = left_df.schema
+    right_schema = right_df.schema
+
+    left_cols = {field.name for field in left_schema.fields}
+    right_cols = {field.name for field in right_schema.fields}
+    right_field_map = {field.name: field for field in right_schema.fields}
+
+    missing_left = right_cols - left_cols
+    missing_right = left_cols - right_cols
+
+    def add_nulls(
+        missing_cols: set[str], to_df: snowpark.DataFrame, from_df: snowpark.DataFrame
+    ) -> snowpark.DataFrame:
+        dt_map = {field.name: field.datatype for field in from_df.schema.fields}
+        result = to_df.select(
+            "*",
+            *[lit(None).cast(dt_map[col]).alias(col) for col in missing_cols],
+        )
+
+        result_fields = []
+        for field in to_df.schema.fields:
+            result_fields.append(
+                StructField(field.name, field.datatype, field.nullable)
+            )
+        for col_name in missing_cols:
+            from_field = next(
+                field for field in from_df.schema.fields if field.name == col_name
+            )
+            result_fields.append(
+                StructField(col_name, from_field.datatype, from_field.nullable)
+            )
+
+        set_schema_getter(result, lambda: StructType(result_fields))
+
+        return result
+
+    if missing_left or missing_right:
+        if allow_missing_columns:
+            left = left_df
+            right = right_df
+            if missing_left:
+                left = add_nulls(missing_left, left, right)
+            if missing_right:
+                right = add_nulls(missing_right, right, left)
+            result = left._union_by_name_internal(right, is_all=True)
+
+            result_fields = []
+            for field in left_schema.fields:
+                result_fields.append(
+                    StructField(field.name, field.datatype, field.nullable)
+                )
+            for col_name in missing_left:
+                right_field = right_field_map[col_name]
+                result_fields.append(
+                    StructField(col_name, right_field.datatype, right_field.nullable)
+                )
+
+            set_schema_getter(result, lambda: StructType(result_fields))
+            return result
+        else:
+            exception = (
+                SnowparkClientExceptionMessages.DF_CANNOT_RESOLVE_COLUMN_NAME_AMONG(
+                    missing_left, missing_right
+                )
+            )
+            attach_custom_error_code(exception, ErrorCodes.COLUMN_NOT_FOUND)
+            raise exception
+
+    result = left_df.unionAllByName(
+        right_df, allow_missing_columns=allow_missing_columns
+    )
+    set_schema_getter(result, lambda: left_df.schema)
+    return result

snowflake/snowpark_connect/relation/map_show_string.py

@@ -15,6 +15,9 @@ from snowflake.snowpark_connect.column_name_handler import set_schema_getter
 from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.map_relation import map_relation
+from snowflake.snowpark_connect.relation.read.metadata_utils import (
+    without_internal_columns,
+)
 
 
 def map_show_string(rel: relation_proto.Relation) -> pandas.DataFrame:
@@ -26,14 +29,17 @@ def map_show_string(rel: relation_proto.Relation) -> pandas.DataFrame:
     Buffer object as a single cell.
     """
     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)
+    filtered_container = without_internal_columns(input_df_container)
+    display_df = filtered_container.dataframe
+    display_spark_columns = filtered_container.column_map.get_spark_columns()
+
+    input_df = _handle_datetype_columns(display_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_container.column_map.get_spark_columns(),
+        _spark_column_names=display_spark_columns,
         _spark_session_tz=global_config.spark_sql_session_timeZone,
     )
     return pandas.DataFrame({"show_string": [show_string]})
@@ -44,14 +50,15 @@ def map_repr_html(rel: relation_proto.Relation) -> pandas.DataFrame:
     Generate the html string representation of the input dataframe.
     """
     input_df_container: DataFrameContainer = map_relation(rel.html_string.input)
-    input_df = input_df_container.dataframe
 
+    filtered_container = without_internal_columns(input_df_container)
+    input_df = filtered_container.dataframe
     input_panda = input_df.toPandas()
     input_panda.rename(
         columns={
             analyzer_utils.unquote_if_quoted(
-                input_df_container.column_map.get_snowpark_columns()[i]
-            ): input_df_container.column_map.get_spark_columns()[i]
+                filtered_container.column_map.get_snowpark_columns()[i]
+            ): filtered_container.column_map.get_spark_columns()[i]
             for i in range(len(input_panda.columns))
         },
         inplace=True,