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

snowflake/snowpark_connect/relation/map_row_ops.py

@@ -20,12 +20,9 @@ from snowflake.snowpark.types import (
     NullType,
     ShortType,
 )
-from snowflake.snowpark_connect.column_name_handler import (
-    schema_getter,
-    set_schema_getter,
-    with_column_map,
-)
+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.expression.literal import get_literal_field_and_name
 from snowflake.snowpark_connect.expression.map_expression import (
     map_single_column_expression,
@@ -39,13 +36,14 @@ from snowflake.snowpark_connect.utils.telemetry import (
 
 def map_deduplicate(
     rel: relation_proto.Relation,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Deduplicate a DataFrame based on a Relation's deduplicate.
 
     The deduplicate is a list of columns that is applied to the DataFrame.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.deduplicate.input)
+    input_container = map_relation(rel.deduplicate.input)
+    input_df = input_container.dataframe
 
     if (
         rel.deduplicate.HasField("within_watermark")
@@ -62,23 +60,29 @@ def map_deduplicate(
         result: snowpark.DataFrame = input_df.drop_duplicates()
     else:
         result: snowpark.DataFrame = input_df.drop_duplicates(
-            *input_df._column_map.get_snowpark_column_names_from_spark_column_names(
+            *input_container.column_map.get_snowpark_column_names_from_spark_column_names(
                 list(rel.deduplicate.column_names)
             )
         )
-    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,
+        input_container.alias,
+        cached_schema_getter=lambda: input_df.schema,
+    )
 
-def map_dropna(rel: relation_proto.Relation) -> snowpark.DataFrame:
+
+def map_dropna(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Drop NA values from the input DataFrame.
-
-
     """
-    input_df: snowpark.DataFrame = map_relation(rel.drop_na.input)
+    input_container = map_relation(rel.drop_na.input)
+    input_df = input_container.dataframe
+
     if rel.drop_na.HasField("min_non_nulls"):
         thresh = rel.drop_na.min_non_nulls
         how = "all"
@@ -89,7 +93,9 @@ def map_dropna(rel: relation_proto.Relation) -> snowpark.DataFrame:
         columns: list[str] = [
             # Use the mapping to get the Snowpark internal column name
             # TODO: Verify the behavior of duplicate column names with dropna
-            input_df._column_map.get_snowpark_column_name_from_spark_column_name(c)
+            input_container.column_map.get_snowpark_column_name_from_spark_column_name(
+                c
+            )
             for c in rel.drop_na.cols
         ]
         result: snowpark.DataFrame = input_df.dropna(
@@ -97,22 +103,32 @@ def map_dropna(rel: relation_proto.Relation) -> snowpark.DataFrame:
         )
     else:
         result: snowpark.DataFrame = input_df.dropna(how=how, thresh=thresh)
-    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,
+        input_container.alias,
+        cached_schema_getter=lambda: input_df.schema,
+    )
 
-def map_fillna(rel: relation_proto.Relation) -> snowpark.DataFrame:
+
+def map_fillna(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Fill NA values in the DataFrame.
 
     The `fill_value` is a scalar value that will be used to replace NaN values.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.fill_na.input)
+    input_container = map_relation(rel.fill_na.input)
+    input_df = input_container.dataframe
+
     if len(rel.fill_na.cols) > 0:
         columns: list[str] = [
-            input_df._column_map.get_snowpark_column_name_from_spark_column_name(c)
+            input_container.column_map.get_snowpark_column_name_from_spark_column_name(
+                c
+            )
             for c in rel.fill_na.cols
         ]
         values = [get_literal_field_and_name(v)[0] for v in rel.fill_na.values]
@@ -142,26 +158,50 @@ def map_fillna(rel: relation_proto.Relation) -> snowpark.DataFrame:
                 for field in input_df.schema.fields
             }
         result = input_df.fillna(fill_value, include_decimal=True)
-    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,
+        input_container.alias,
+        cached_schema_getter=lambda: input_df.schema,
+    )
 
-def map_union(rel: relation_proto.Relation) -> snowpark.DataFrame:
+
+def map_union(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Union two DataFrames together.
 
     The two DataFrames must have the same schema.
     """
-    left_df: snowpark.DataFrame = map_relation(rel.set_op.left_input)
-    right_df: snowpark.DataFrame = map_relation(rel.set_op.right_input)
+    left_result = map_relation(rel.set_op.left_input)
+    right_result = 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
-    left_dtypes = [field.datatype for field in left_df.schema.fields]
-    right_dtypes = [field.datatype for field in right_df.schema.fields]
+    # 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
 
-    allow_missing_columns = bool(rel.set_op.allow_missing_columns)
+    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_dtypes = [field.datatype for field in left_schema.fields]
+    right_dtypes = [field.datatype for field in right_schema.fields]
 
     spark_sql_ansi_enabled = global_config.spark_sql_ansi_enabled
     if left_dtypes != right_dtypes and not rel.set_op.by_name:
@@ -212,13 +252,22 @@ def map_union(rel: relation_proto.Relation) -> snowpark.DataFrame:
                     target_right_dtypes.append(right_type)
 
         def cast_columns(
-            df: snowpark.DataFrame,
+            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
-            df_schema = df.schema  # Get current schema
+                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):
@@ -232,38 +281,46 @@ def map_union(rel: relation_proto.Relation) -> snowpark.DataFrame:
                     new_columns.append(df[col_name])
 
             new_df = df.select(new_columns)
-            return with_column_map(
-                new_df,
-                df._column_map.get_spark_columns(),
-                df._column_map.get_snowpark_columns(),
-                target_dtypes,
-                df._column_map.column_metadata,
-                parent_column_name_map=df._column_map,
+            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_df = cast_columns(left_df, left_dtypes, target_left_dtypes)
-        right_df = cast_columns(right_df, right_dtypes, target_right_dtypes)
+        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
 
     # Save the column names so that we can restore them after the union.
-    left_df_columns = left_df.columns
-
-    result: snowpark.DataFrame = None
+    left_df_columns = left_result.dataframe.columns
 
     if rel.set_op.by_name:
         # To use unionByName, we need to have the same column names.
         # We rename the columns back to their originals using the map
-        left_column_map = left_df._column_map
-        left_table_name = left_df._table_name
+        left_column_map = left_result.column_map
+        left_table_name = left_result.table_name
         left_schema_getter = schema_getter(left_df)
-        right_column_map = right_df._column_map
-
+        right_column_map = right_result.column_map
         columns_to_restore: dict[str, tuple[str, str]] = {}
 
         for column in right_df.columns:
             spark_name = (
                 right_column_map.get_spark_column_name_from_snowpark_column_name(column)
             )
-
             right_df = right_df.withColumnRenamed(column, spark_name)
             columns_to_restore[spark_name.upper()] = (spark_name, column)
 
@@ -271,11 +328,10 @@ def map_union(rel: relation_proto.Relation) -> snowpark.DataFrame:
             spark_name = (
                 left_column_map.get_spark_column_name_from_snowpark_column_name(column)
             )
-
             left_df = left_df.withColumnRenamed(column, spark_name)
             columns_to_restore[spark_name.upper()] = (spark_name, column)
 
-        result = left_df.union_all_by_name(
+        result = left_df.unionAllByName(
             right_df, allow_missing_columns=allow_missing_columns
         )
 
@@ -296,41 +352,42 @@ def map_union(rel: relation_proto.Relation) -> snowpark.DataFrame:
             right_df_col_metadata = right_column_map.column_metadata or {}
             merged_column_metadata = left_df_col_metadata | right_df_col_metadata
 
-            return with_column_map(
+            return DataFrameContainer.create_with_column_mapping(
                 result,
-                spark_columns,
-                snowpark_columns,
+                spark_column_names=spark_columns,
+                snowpark_column_names=snowpark_columns,
                 column_metadata=merged_column_metadata,
             )
 
         for i in range(len(left_df_columns)):
             result = result.withColumnRenamed(result.columns[i], left_df_columns[i])
 
-        result._column_map = left_column_map
-        result._table_name = left_table_name
-        set_schema_getter(result, left_schema_getter)
+        return DataFrameContainer(
+            result,
+            column_map=left_column_map,
+            table_name=left_table_name,
+            cached_schema_getter=left_schema_getter,
+        )
     elif rel.set_op.is_all:
         result = left_df.unionAll(right_df)
-        result._column_map = left_df._column_map
-        result._table_name = left_df._table_name
-        set_schema_getter(result, lambda: left_df.schema)
+        return DataFrameContainer(
+            result,
+            column_map=left_result.column_map,
+            cached_schema_getter=lambda: left_df.schema,
+        )
     else:
         result = left_df.union(right_df)
-        result._column_map = left_df._column_map
-        result._table_name = left_df._table_name
-        set_schema_getter(result, lambda: left_df.schema)
-
-    # union operation does not preserve column qualifiers
-    return with_column_map(
-        result,
-        result._column_map.get_spark_columns(),
-        result._column_map.get_snowpark_columns(),
-        column_metadata=result._column_map.column_metadata,
-        parent_column_name_map=result._column_map,
-    )
+        # union operation does not preserve column qualifiers
+        return DataFrameContainer(
+            result,
+            column_map=left_result.column_map,
+            cached_schema_getter=lambda: left_df.schema,
+        )
 
 
-def map_intersect(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_intersect(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Return a new DataFrame containing rows in both DataFrames:
 
@@ -363,36 +420,36 @@ def map_intersect(rel: relation_proto.Relation) -> snowpark.DataFrame:
     |  b|  3|
     +---+---+
     """
-    left_df: snowpark.DataFrame = map_relation(rel.set_op.left_input)
-    right_df: snowpark.DataFrame = map_relation(rel.set_op.right_input)
+    left_result = map_relation(rel.set_op.left_input)
+    right_result = map_relation(rel.set_op.right_input)
+    left_df = left_result.dataframe
+    right_df = right_result.dataframe
 
     if rel.set_op.is_all:
         left_df_with_row_number = utils.get_df_with_partition_row_number(
-            left_df, rel.set_op.left_input.common.plan_id, "left_row_number"
+            left_result, rel.set_op.left_input.common.plan_id, "left_row_number"
         )
         right_df_with_row_number = utils.get_df_with_partition_row_number(
-            right_df, rel.set_op.right_input.common.plan_id, "right_row_number"
+            right_result, rel.set_op.right_input.common.plan_id, "right_row_number"
         )
 
         result: snowpark.DataFrame = left_df_with_row_number.intersect(
             right_df_with_row_number
-        ).select(*left_df._column_map.get_snowpark_columns())
+        ).select(*left_result.column_map.get_snowpark_columns())
     else:
         result: snowpark.DataFrame = left_df.intersect(right_df)
 
-    # the result df keeps the column map of the original left_df
-    result = with_column_map(
-        result,
-        left_df._column_map.get_spark_columns(),
-        left_df._column_map.get_snowpark_columns(),
-        column_metadata=left_df._column_map.column_metadata,
+    return DataFrameContainer(
+        dataframe=result,
+        column_map=left_result.column_map,
+        table_name=left_result.table_name,
+        cached_schema_getter=lambda: left_df.schema,
     )
-    result._table_name = left_df._table_name
-    set_schema_getter(result, lambda: left_df.schema)
-    return result
 
 
-def map_except(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_except(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Return a new DataFrame containing rows in the left DataFrame but not in the right DataFrame.
 
@@ -426,8 +483,10 @@ def map_except(rel: relation_proto.Relation) -> snowpark.DataFrame:
     |  c|  4|
     +---+---+
     """
-    left_df: snowpark.DataFrame = map_relation(rel.set_op.left_input)
-    right_df: snowpark.DataFrame = map_relation(rel.set_op.right_input)
+    left_result = map_relation(rel.set_op.left_input)
+    right_result = map_relation(rel.set_op.right_input)
+    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,
@@ -453,91 +512,107 @@ def map_except(rel: relation_proto.Relation) -> snowpark.DataFrame:
         # +---+---+------------+
         # at the end we will do a select to exclude the row number column
         left_df_with_row_number = utils.get_df_with_partition_row_number(
-            left_df, rel.set_op.left_input.common.plan_id, "left_row_number"
+            left_result, rel.set_op.left_input.common.plan_id, "left_row_number"
         )
         right_df_with_row_number = utils.get_df_with_partition_row_number(
-            right_df, rel.set_op.right_input.common.plan_id, "right_row_number"
+            right_result, rel.set_op.right_input.common.plan_id, "right_row_number"
         )
 
         # Perform except use left_df_with_row_number and right_df_with_row_number,
         # and drop the row number column after except.
         result_df = left_df_with_row_number.except_(right_df_with_row_number).select(
-            *left_df._column_map.get_snowpark_columns()
+            *left_result.column_map.get_snowpark_columns()
         )
     else:
         result_df = left_df.except_(right_df)
 
     # the result df keeps the column map of the original left_df
     # union operation does not preserve column qualifiers
-    result_df = with_column_map(
-        result_df,
-        left_df._column_map.get_spark_columns(),
-        left_df._column_map.get_snowpark_columns(),
-        column_metadata=left_df._column_map.column_metadata,
+    return DataFrameContainer(
+        dataframe=result_df,
+        column_map=left_result.column_map,
+        table_name=left_result.table_name,
+        cached_schema_getter=lambda: left_df.schema,
     )
-    result_df._table_name = left_df._table_name
-    set_schema_getter(result_df, lambda: left_df.schema)
-    return result_df
 
 
 def map_filter(
     rel: relation_proto.Relation,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Filter a DataFrame based on a Relation's filter.
 
     The filter is a SQL expression that is applied to the DataFrame.
     """
-    input_df = map_relation(rel.filter.input)
+    input_container = map_relation(rel.filter.input)
+    input_df = input_container.dataframe
+
     typer = ExpressionTyper(input_df)
     _, condition = map_single_column_expression(
-        rel.filter.condition, input_df._column_map, typer
+        rel.filter.condition, input_container.column_map, typer
     )
     result = input_df.filter(condition.col)
-    result._column_map = input_df._column_map
-    result._alias = input_df._alias
-    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,
+        input_container.alias,
+        cached_schema_getter=lambda: input_df.schema,
+    )
 
 
 def map_limit(
     rel: relation_proto.Relation,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Limit a DataFrame based on a Relation's limit.
 
     The limit is an integer that is applied to the DataFrame.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.limit.input)
+    input_container = map_relation(rel.limit.input)
+    input_df = input_container.dataframe
+
     result: snowpark.DataFrame = input_df.limit(rel.limit.limit)
-    result._column_map = input_df._column_map
-    result._table_name = input_df._table_name
-    set_schema_getter(result, lambda: input_df.schema)
-    return result
+
+    return DataFrameContainer(
+        result,
+        column_map=input_container.column_map,
+        table_name=input_container.table_name,
+        alias=input_container.alias,
+        cached_schema_getter=lambda: input_df.schema,
+    )
 
 
 def map_offset(
     rel: relation_proto.Relation,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Offset a DataFrame based on a Relation's offset.
 
     The offset is an integer that is applied to the DataFrame.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.offset.input)
+    input_container = 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
     # support offset without limit.
     result: snowpark.DataFrame = input_df.limit(
         input_df.count(), offset=rel.offset.offset
     )
-    result._column_map = input_df._column_map
-    result._table_name = input_df._table_name
-    set_schema_getter(result, lambda: input_df.schema)
-    return result
+
+    return DataFrameContainer(
+        result,
+        column_map=input_container.column_map,
+        table_name=input_container.table_name,
+        alias=input_container.alias,
+        cached_schema_getter=lambda: input_df.schema,
+    )
 
 
-def map_replace(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_replace(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Replace values in the DataFrame.
 
@@ -545,10 +620,11 @@ def map_replace(rel: relation_proto.Relation) -> snowpark.DataFrame:
     values to replace. The values in the dictionary are the values to replace
     and the keys are the values to replace them with.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.replace.input)
+    result = map_relation(rel.replace.input)
+    input_df = result.dataframe
     ordered_columns = input_df.columns
-    column_map = input_df._column_map
-    table_name = input_df._table_name
+    column_map = result.column_map
+    table_name = result.table_name
     # note that seems like spark connect always send number values as double in rel.replace.replacements.
     to_replace = [
         get_literal_field_and_name(i.old_value)[0] for i in rel.replace.replacements
@@ -647,7 +723,7 @@ def map_replace(rel: relation_proto.Relation) -> snowpark.DataFrame:
 
     if len(rel.replace.cols) > 0:
         columns: list[str] = [
-            input_df._column_map.get_snowpark_column_name_from_spark_column_name(c)
+            column_map.get_snowpark_column_name_from_spark_column_name(c)
             for c in rel.replace.cols
         ]
         for c in columns:
@@ -657,18 +733,19 @@ def map_replace(rel: relation_proto.Relation) -> snowpark.DataFrame:
             input_df = input_df.with_column(c, replace_case_expr(c, to_replace, values))
 
     result = input_df.select(*[col(c) for c in ordered_columns])
-    result._column_map = column_map
-    result._table_name = table_name
-    return result
+
+    return DataFrameContainer(result, column_map=column_map, table_name=table_name)
 
 
 def map_sample(
     rel: relation_proto.Relation,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Sample a DataFrame based on a Relation's sample.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.sample.input)
+    input_container = 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")
@@ -691,26 +768,35 @@ def map_sample(
         )
     else:
         result: snowpark.DataFrame = input_df.sample(frac=frac)
-        result._column_map = input_df._column_map
-        result._table_name = input_df._table_name
-        set_schema_getter(result, lambda: input_df.schema)
-        return result
+        return DataFrameContainer(
+            result,
+            column_map=input_container.column_map,
+            table_name=input_container.table_name,
+            alias=input_container.alias,
+            cached_schema_getter=lambda: input_df.schema,
+        )
 
 
 def map_tail(
     rel: relation_proto.Relation,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Tail a DataFrame based on a Relation's tail.
 
     The tail is an integer that is applied to the DataFrame.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.tail.input)
+    input_container = map_relation(rel.tail.input)
+    input_df = input_container.dataframe
+
     num_rows = input_df.count()
     result: snowpark.DataFrame = input_df.limit(
         num_rows, offset=max(0, num_rows - rel.tail.limit)
     )
-    result._column_map = input_df._column_map
-    result._table_name = input_df._table_name
-    set_schema_getter(result, lambda: input_df.schema)
-    return result
+
+    return DataFrameContainer(
+        result,
+        column_map=input_container.column_map,
+        table_name=input_container.table_name,
+        alias=input_container.alias,
+        cached_schema_getter=lambda: input_df.schema,
+    )