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

snowflake/snowpark_connect/relation/map_local_relation.py

@@ -11,12 +11,13 @@ import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
 from snowflake import snowpark
 from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted
+from snowflake.snowpark._internal.utils import is_in_stored_procedure
 from snowflake.snowpark.types import LongType, StructField, StructType
 from snowflake.snowpark_connect import tcm
 from snowflake.snowpark_connect.column_name_handler import (
     make_column_names_snowpark_compatible,
-    with_column_map,
 )
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.type_mapping import (
     map_json_schema_to_snowpark,
     map_pyarrow_to_snowpark_types,
@@ -126,10 +127,8 @@ def map_pylist_cell_to_python_object(cell, type: pa.lib.DataType):
                 map_pylist_cell_to_python_object(obj, list_type.value_type)
                 for obj in cell
             ]
-        case map_type if (
-            cell is not None
-            and isinstance(type, pa.lib.MapType)
-            and all(isinstance(obj, tuple) and len(obj) == 2 for obj in cell)
+        case map_type if cell is not None and isinstance(type, pa.lib.MapType) and all(
+            isinstance(obj, tuple) and len(obj) == 2 for obj in cell
         ):
             # the MapType in arrow becomes list in pylist_df,
             # e.g. {"Car": "Honda", "Bike": "Yamaha"} --> [("Car", "Honda"), ("Bike", "Yamaha")] , and causes some
@@ -171,7 +170,9 @@ def map_pandas_cell_to_python_object(cell):
     return res if res == res else None
 
 
-def map_local_relation(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_local_relation(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     if rel.local_relation.HasField("data"):
         data = pa.BufferReader(rel.local_relation.data)
         with pa.ipc.open_stream(data) as reader:
@@ -204,9 +205,9 @@ def map_local_relation(rel: relation_proto.Relation) -> snowpark.DataFrame:
             # Only create the pandas dataframe for empty dataframe cases.
             pandas_df = table.to_pandas()
             snowpark_df: snowpark.DataFrame = session.create_dataframe(pandas_df)
-            return with_column_map(
-                snowpark_df,
-                spark_column_names,
+            return DataFrameContainer.create_with_column_mapping(
+                dataframe=snowpark_df,
+                spark_column_names=spark_column_names,
                 snowpark_column_names=new_columns,
                 column_metadata=column_metadata,
             )
@@ -230,16 +231,22 @@ def map_local_relation(rel: relation_proto.Relation) -> snowpark.DataFrame:
         # Special characters in the schema currently break create_dataframe with arrow
         # https://snowflakecomputing.atlassian.net/browse/SNOW-2199291
         current_schema = session.get_current_schema()
+
+        # _create_temp_stage() changes were not ported to the internal connector, leading to this
+        # error on TCM and in notebooks (sproc):
+        # TypeError: _create_temp_stage() takes 7 positional arguments but 8 were given
         use_pyarrow = (
-            re.match(
+            not is_in_stored_procedure()
+            # TODO: SNOW-2220726 investigate why use_pyarrow failed in TCM:
+            and not tcm.TCM_MODE
+            and re.match(
                 # See https://docs.snowflake.com/en/sql-reference/identifiers-syntax
                 r"[A-Za-z_][A-Za-z0-9_\$]*",
                 # Schema may be double-quoted.
                 current_schema.strip('"') if current_schema is not None else "",
             )
             is not None
-        ) and not tcm.TCM_MODE  # TODO: SNOW-2220726 investigate why use_pyarrow failed in TCM:
-        # TypeError: _create_temp_stage() takes 7 positional arguments but 8 were given
+        )
 
         if use_pyarrow:
             snowpark_df: snowpark.DataFrame = session.create_dataframe(
@@ -257,6 +264,7 @@ def map_local_relation(rel: relation_proto.Relation) -> snowpark.DataFrame:
             ]
 
             snowpark_df = snowpark_df.select(*casted_columns)
+
         else:
             pylist_df = [
                 list(row)
@@ -285,9 +293,9 @@ def map_local_relation(rel: relation_proto.Relation) -> snowpark.DataFrame:
                 snowpark_schema,
             )
 
-        return with_column_map(
-            snowpark_df,
-            spark_column_names,
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=snowpark_df,
+            spark_column_names=spark_column_names,
             snowpark_column_names=new_columns,
             column_metadata=column_metadata,
             snowpark_column_types=[f.datatype for f in snowpark_schema.fields],
@@ -305,9 +313,9 @@ def map_local_relation(rel: relation_proto.Relation) -> snowpark.DataFrame:
             [],
             snowpark_schema,
         )
-        return with_column_map(
-            snowpark_df,
-            spark_column_names,
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=snowpark_df,
+            spark_column_names=spark_column_names,
             snowpark_column_names=new_columns,
             column_metadata=column_metadata,
         )
@@ -317,10 +325,17 @@ def map_local_relation(rel: relation_proto.Relation) -> snowpark.DataFrame:
         )
 
 
-def map_range(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_range(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     session = get_or_create_snowpark_session()
     new_columns = make_column_names_snowpark_compatible(["id"], rel.common.plan_id)
     result = session.range(
         rel.range.start, rel.range.end, rel.range.step
     ).with_column_renamed("ID", new_columns[0])
-    return with_column_map(result, ["id"], new_columns, [LongType()])
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=["id"],
+        snowpark_column_names=new_columns,
+        snowpark_column_types=[LongType()],
+    )

snowflake/snowpark_connect/relation/map_map_partitions.py

@@ -8,9 +8,9 @@ from pyspark.sql.connect.proto.expressions_pb2 import CommonInlineUserDefinedFun
 import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
 from snowflake.snowpark.types import StructType
-from snowflake.snowpark_connect.column_name_handler import with_column_map
 from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.constants import MAP_IN_ARROW_EVAL_TYPE
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 from snowflake.snowpark_connect.type_mapping import proto_to_snowpark_type
 from snowflake.snowpark_connect.utils.pandas_udtf_utils import create_pandas_udtf
@@ -30,13 +30,16 @@ from snowflake.snowpark_connect.utils.udtf_helper import (
 )
 
 
-def map_map_partitions(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_map_partitions(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Map a function over the partitions of the input DataFrame.
 
     This is a simple wrapper around the `mapInPandas` method in Snowpark.
     """
-    input_df = map_relation(rel.map_partitions.input)
+    input_container = map_relation(rel.map_partitions.input)
+    input_df = input_container.dataframe
     udf_proto = rel.map_partitions.func
     udf_check(udf_proto)
 
@@ -44,7 +47,7 @@ def map_map_partitions(rel: relation_proto.Relation) -> snowpark.DataFrame:
     eval_type = udf_proto.python_udf.eval_type
 
     if eval_type == MAP_IN_ARROW_EVAL_TYPE:
-        return _map_in_arrow_with_pandas_udtf(input_df, udf_proto)
+        return _map_in_arrow_with_pandas_udtf(input_container, udf_proto)
     else:
         return _map_partitions_with_udf(input_df, udf_proto)
 
@@ -76,22 +79,24 @@ def _call_udtf(
     # Only return the output columns.
     result_df = result_df_with_dummy.select(*output_cols)
 
-    return with_column_map(
-        result_df,
-        output_cols,
-        output_cols,
-        [field.datatype for field in return_type.fields],
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result_df,
+        spark_column_names=output_cols,
+        snowpark_column_names=output_cols,
+        snowpark_column_types=[field.datatype for field in return_type.fields],
     )
 
 
 def _map_in_arrow_with_pandas_udtf(
-    input_df: snowpark.DataFrame, udf_proto: CommonInlineUserDefinedFunction
+    input_df_container: DataFrameContainer,
+    udf_proto: CommonInlineUserDefinedFunction,
 ) -> snowpark.DataFrame:
     """
     Handle mapInArrow using pandas_udtf for partition-level Arrow processing.
     """
+    input_df = input_df_container.dataframe
     input_schema = input_df.schema
-    spark_column_names = input_df._column_map.get_spark_columns()
+    spark_column_names = input_df_container.column_map.get_spark_columns()
     return_type = proto_to_snowpark_type(udf_proto.python_udf.output_type)
     if require_creating_udtf_in_sproc(udf_proto):
         udtf_name = create_pandas_udtf_in_sproc(
@@ -138,9 +143,9 @@ def _map_partitions_with_udf(
     udf_column_name = "UDF_OUTPUT"
     snowpark_columns = [snowpark_fn.col(name) for name in input_df.columns]
     result = input_df.select(snowpark_fn.call_udf(snowpark_udf.name, *snowpark_columns))
-    return with_column_map(
-        result,
-        [udf_column_name],
-        [udf_column_name],
-        [snowpark_udf.return_type],
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result,
+        spark_column_names=[udf_column_name],
+        snowpark_column_names=[udf_column_name],
+        snowpark_column_types=[snowpark_udf.return_type],
     )

snowflake/snowpark_connect/relation/map_relation.py

@@ -7,8 +7,7 @@ import copy
 import pandas
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
-from snowflake import snowpark
-from snowflake.snowpark_connect.column_name_handler import set_schema_getter
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.utils.cache import (
     df_cache_map_get,
     df_cache_map_put_if_absent,
@@ -32,8 +31,8 @@ NATURAL_JOIN_TYPE_BASE = 22
 
 def map_relation(
     rel: relation_proto.Relation, reuse_parsed_plan: bool = True
-) -> snowpark.DataFrame | pandas.DataFrame:
-    """Map a Spark Protobuf Relation message to a Snowpark DataFrame or pandas DataFrame.
+) -> DataFrameContainer | pandas.DataFrame:
+    """Map a Spark Protobuf Relation message to a DataFrameContainer or pandas DataFrame.
 
     NOTE: We return a pandas DataFrame object when the return value of the operation is a
     scalar value. The client expects these as an Arrow buffer with this return value packed
@@ -43,11 +42,11 @@ def map_relation(
 
     Args:
         rel (relation_proto.Relation): The Spark Protobuf Relation message to map.
-        reuse_parsed_plan (bool, optional): If True, reuses previously parsed df from cache
+        reuse_parsed_plan (bool, optional): If True, reuses previously parsed container from cache
             to avoid redundant operations.
 
     Returns:
-        snowpark.DataFrame | pandas.DataFrame: The Snowpark DataFrame or pandas DataFrame
+        DataFrameContainer | pandas.DataFrame: The DataFrameContainer or pandas DataFrame
         that corresponds to the input Spark Protobuf Relation message.
     """
     # TODO: from snowflake_connect_server.relation import map_extension
@@ -75,22 +74,27 @@ def map_relation(
         # Check for cached relation
         cache_entry = df_cache_map_get((get_session_id(), rel.common.plan_id))
         if cache_entry is not None:
-            if isinstance(cache_entry, snowpark.DataFrame):
+            if isinstance(cache_entry, DataFrameContainer):
                 set_plan_id_map(rel.common.plan_id, cache_entry)
             return cache_entry
 
         # If df is not cached, check if we have parsed the plan
-        cached_df = get_plan_id_map(rel.common.plan_id)
-        if cached_df is not None:
+        cached_container = get_plan_id_map(rel.common.plan_id)
+        if cached_container is not None:
+            cached_df = cached_container.dataframe
             result = copy.copy(cached_df)
-            result._column_map = copy.deepcopy(cached_df._column_map)
-            result._table_name = copy.deepcopy(cached_df._table_name)
-            set_schema_getter(result, lambda: cached_df.schema)
+            # Create new container without triggering schema access
+            result_container = DataFrameContainer(
+                result,
+                column_map=copy.deepcopy(cached_container.column_map),
+                table_name=copy.deepcopy(cached_container.table_name),
+                alias=cached_container.alias,
+                cached_schema_getter=lambda: cached_df.schema,
+            )
             # If we don't make a copy of the df._output, the expression IDs for attributes in Snowpark DataFrames will differ from those stored in the cache,
             # leading to errors during query execution.
             result._output = cached_df._output
-
-            return result
+            return result_container
 
     if rel.WhichOneof("rel_type") is not None:
         logger.info(rel.WhichOneof("rel_type").upper())
@@ -99,7 +103,7 @@ def map_relation(
         # type was incorrectly routed here.
         raise SnowparkConnectNotImplementedError("No Relation Type")
 
-    result: snowpark.DataFrame | pandas.DataFrame
+    result: DataFrameContainer | pandas.DataFrame
     operation = rel.WhichOneof("rel_type")
     with push_operation_scope(operation):
         match operation:
@@ -248,6 +252,8 @@ def map_relation(
             case other:
                 raise SnowparkConnectNotImplementedError(f"Other Relation {other}")
 
-        if isinstance(result, snowpark.DataFrame):
+        # Store container in plan cache
+        if isinstance(result, DataFrameContainer):
             set_plan_id_map(rel.common.plan_id, result)
+
         return result