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

snowflake/snowpark_connect/relation/map_stats.py

@@ -7,30 +7,34 @@ import ast
 import numpy as np
 import pandas
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
+from pyspark.errors.exceptions.base import AnalysisException
 
 import snowflake.snowpark.functions as fn
 import snowflake.snowpark.types as snowpark_types
 from snowflake import snowpark
 from snowflake.snowpark.exceptions import SnowparkSQLException
-from snowflake.snowpark_connect.column_name_handler import with_column_map
+from snowflake.snowpark_connect.config import get_boolean_session_config_param
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
 
 
 def map_corr(
     rel: relation_proto.Relation,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Find the correlation of two columns in the input DataFrame.
 
     Returns a pandas DataFrame because the correlation of two columns produces
     a scalar value.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.corr.input)
-    col1 = input_df._column_map.get_snowpark_column_name_from_spark_column_name(
+    input_container = map_relation(rel.corr.input)
+    input_df = input_container.dataframe
+
+    col1 = input_container.column_map.get_snowpark_column_name_from_spark_column_name(
         rel.corr.col1
     )
-    col2 = input_df._column_map.get_snowpark_column_name_from_spark_column_name(
+    col2 = input_container.column_map.get_snowpark_column_name_from_spark_column_name(
         rel.corr.col2
     )
     # TODO: Handle method, Snowpark does not support this yet.
@@ -44,18 +48,20 @@ def map_corr(
 
 def map_cov(
     rel: relation_proto.Relation,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Find the covariance of two columns in the input DataFrame.
 
     Returns a pandas DataFrame because the corvariance of two columns produces
     a scalar value.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.cov.input)
-    col1 = input_df._column_map.get_snowpark_column_name_from_spark_column_name(
+    input_container = map_relation(rel.cov.input)
+    input_df = input_container.dataframe
+
+    col1 = input_container.column_map.get_snowpark_column_name_from_spark_column_name(
         rel.cov.col1
     )
-    col2 = input_df._column_map.get_snowpark_column_name_from_spark_column_name(
+    col2 = input_container.column_map.get_snowpark_column_name_from_spark_column_name(
         rel.cov.col2
     )
     result: float = input_df.cov(col1, col2)
@@ -64,15 +70,40 @@ def map_cov(
 
 def map_approx_quantile(
     rel: relation_proto.Relation,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Find one or more approximate quantiles in the input DataFrame.
 
     Returns a pandas DataFrame because the approximate quantile produces a
     list of scalar values.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.approx_quantile.input)
-    cols = input_df._column_map.get_snowpark_column_names_from_spark_column_names(
+    input_container = map_relation(rel.approx_quantile.input)
+    input_df = input_container.dataframe
+
+    snowflake_compatible = get_boolean_session_config_param(
+        "enable_snowflake_extension_behavior"
+    )
+
+    if not snowflake_compatible:
+        # When Snowflake extension behavior is disabled, validate that all requested columns exist
+        requested_spark_cols = list(rel.approx_quantile.cols)
+        available_spark_cols = input_container.column_map.get_spark_columns()
+
+        for col_name in requested_spark_cols:
+            if col_name not in available_spark_cols:
+                # Find suggestions for the unresolved column
+                suggestions = [c for c in available_spark_cols if c != col_name]
+                suggestion_text = (
+                    f" Did you mean one of the following? [`{'`, `'.join(suggestions)}`]."
+                    if suggestions
+                    else ""
+                )
+
+                raise AnalysisException(
+                    f"[UNRESOLVED_COLUMN.WITH_SUGGESTION] A column or function parameter with name `{col_name}` cannot be resolved.{suggestion_text}"
+                )
+
+    cols = input_container.column_map.get_snowpark_column_names_from_spark_column_names(
         list(rel.approx_quantile.cols)
     )
     quantile = list(rel.approx_quantile.probabilities)
@@ -84,7 +115,7 @@ def map_approx_quantile(
 
 def map_describe(
     rel: relation_proto.Relation,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Computes basic statistics for numeric columns, which includes count, mean, stddev, min, and max.
     If no columns are provided, this function computes statistics for all numerical or string columns.
@@ -92,15 +123,19 @@ def map_describe(
 
     Returns a new DataFrame that provides basic statistics for the given DataFrame
     """
+    input_container = map_relation(rel.describe.input)
+    input_df = input_container.dataframe
+
     session = get_or_create_snowpark_session()
-    input_df: snowpark.DataFrame = map_relation(rel.describe.input)
     spark_cols = (
         list(rel.describe.cols)
         if rel.describe.cols
-        else input_df._column_map.get_spark_columns()
+        else input_container.column_map.get_spark_columns()
     )
     cols = [
-        input_df._column_map.get_snowpark_column_name_from_spark_column_name(column)
+        input_container.column_map.get_snowpark_column_name_from_spark_column_name(
+            column
+        )
         for column in spark_cols
     ]
 
@@ -121,13 +156,13 @@ def map_describe(
             if stat == row.SUMMARY:
                 ordered_statistics.append(row)
     ordered_desc_df = session.create_dataframe(ordered_statistics)
-    return _build_column_map_helper(ordered_desc_df, input_df)
+    return _build_column_map_helper_container(ordered_desc_df, input_container)
 
 
 # TODO: track missing Snowpark feature
 def map_summary(
     rel: relation_proto.Relation,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Computes specified statistics for numeric or string columns. Available statistics are: count, mean, stddev, min,
     max, arbitrary approximate percentiles specified as a percentage (e.g., 75%), count_distinct, and
@@ -137,12 +172,14 @@ def map_summary(
     Returns a new DataFrame that provides specified statistics for the given DataFrame.
     """
     session = get_or_create_snowpark_session()
-    input_df: snowpark.DataFrame = map_relation(rel.summary.input)
+    result = map_relation(rel.summary.input)
+    input_container: DataFrameContainer = result
+    input_df = input_container.dataframe
 
     numeric_and_string_spark_cols = [
         column
         for field, column in zip(
-            input_df.schema.fields, input_df._column_map.get_spark_columns()
+            input_df.schema.fields, input_container.column_map.get_spark_columns()
         )
         if isinstance(
             field.datatype, (snowpark_types._NumericType, snowpark_types.StringType)
@@ -151,7 +188,9 @@ def map_summary(
 
     # this is intentional to trigger ambigous column name is two columns of same name are provided
     numeric_and_string_snowpark_cols = [
-        input_df._column_map.get_snowpark_column_name_from_spark_column_name(column)
+        input_container.column_map.get_snowpark_column_name_from_spark_column_name(
+            column
+        )
         for column in numeric_and_string_spark_cols
     ]
 
@@ -221,9 +260,11 @@ def map_summary(
         # Modified quantile results, inserting [None, None, None] for string columns
         numeric_index = iter(approx_quantile_values)
         approx_quantile_values_including_string_columns = [
-            [str(value) for value in next(numeric_index)]
-            if col in eligible_columns
-            else [None] * len(quantiles)
+            (
+                [str(value) for value in next(numeric_index)]
+                if col in eligible_columns
+                else [None] * len(quantiles)
+            )
             for col in input_df.columns
         ]
 
@@ -248,22 +289,24 @@ def map_summary(
     spark_col_names = ["summary"]
     spark_col_names.extend(numeric_and_string_spark_cols)
 
-    return with_column_map(
-        ordered_summary_df,
-        spark_col_names,
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=ordered_summary_df,
+        spark_column_names=spark_col_names,
         snowpark_column_names=ordered_summary_df.columns,
     )
 
 
-def map_freq_items(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_freq_items(rel: relation_proto.Relation) -> DataFrameContainer:
     """
     Returns an approximation of the most frequent values in the input, along with their approximate frequencies.
     """
+    input_container = map_relation(rel.freq_items.input)
+    input_df = input_container.dataframe
+
     session = get_or_create_snowpark_session()
-    input_df: snowpark.DataFrame = map_relation(rel.freq_items.input)
     support = rel.freq_items.support
     spark_col_names = []
-    cols = input_df._column_map.get_snowpark_column_names_from_spark_column_names(
+    cols = input_container.column_map.get_snowpark_column_names_from_spark_column_names(
         list(rel.freq_items.cols)
     )
     approx_top_k_df = input_df.select(
@@ -289,11 +332,14 @@ def map_freq_items(rel: relation_proto.Relation) -> snowpark.DataFrame:
 
     for sp_col_name in cols:
         spark_col_names.append(
-            f"{input_df._column_map.get_spark_column_name_from_snowpark_column_name(sp_col_name)}_freqItems"
+            f"{input_container.column_map.get_spark_column_name_from_snowpark_column_name(sp_col_name)}_freqItems"
         )
     approx_top_k_df = session.createDataFrame([filtered_values], spark_col_names)
-    return with_column_map(
-        approx_top_k_df, spark_col_names, snowpark_column_names=spark_col_names
+
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=approx_top_k_df,
+        spark_column_names=spark_col_names,
+        snowpark_column_names=spark_col_names,
     )
 
 
@@ -306,19 +352,22 @@ def add_stat_to_df(
     return summary_df.union(session.createDataFrame(df_data, summary_df.schema))
 
 
-def _build_column_map_helper(
+def _build_column_map_helper_container(
     desc_df: snowpark.DataFrame,
-    input_df: snowpark.DataFrame,
-) -> snowpark.DataFrame:
+    input_container: DataFrameContainer,
+) -> DataFrameContainer:
+    """Container version of _build_column_map_helper."""
     spark_col_names = ["summary"]
     for i, sp_col_name in enumerate(desc_df.columns):
         if i != 0:
             spark_col_names.append(
-                input_df._column_map.get_spark_column_name_from_snowpark_column_name(
+                input_container.column_map.get_spark_column_name_from_snowpark_column_name(
                     sp_col_name
                 )
             )
 
-    return with_column_map(
-        desc_df, spark_col_names, snowpark_column_names=desc_df.columns
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=desc_df,
+        spark_column_names=spark_col_names,
+        snowpark_column_names=desc_df.columns,
     )

snowflake/snowpark_connect/relation/map_subquery_alias.py

@@ -4,29 +4,28 @@
 
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
-from snowflake import snowpark
-from snowflake.snowpark_connect.column_name_handler import with_column_map
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 
 
-def map_alias(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_alias(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Returns an aliased dataframe in which the columns can now be referenced to using col(<df alias>, <column name>).
     """
     alias: str = rel.subquery_alias.alias
     # we set reuse_parsed_plan=False because we need new expr_id for the attributes (output columns) in aliased snowpark dataframe
     # reuse_parsed_plan will lead to ambiguous column name for operations like joining two dataframes that are aliased from the same dataframe
-    input_df: snowpark.DataFrame = map_relation(
-        rel.subquery_alias.input, reuse_parsed_plan=False
-    )
-    input_df._alias = alias
-    qualifiers = [[alias]] * len(input_df._column_map.columns)
+    input_container = map_relation(rel.subquery_alias.input, reuse_parsed_plan=False)
+    qualifiers = [[alias]] * len(input_container.column_map.columns)
 
-    return with_column_map(
-        input_df,
-        input_df._column_map.get_spark_columns(),
-        input_df._column_map.get_snowpark_columns(),
-        column_metadata=input_df._column_map.column_metadata,
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=input_container.dataframe,
+        spark_column_names=input_container.column_map.get_spark_columns(),
+        snowpark_column_names=input_container.column_map.get_snowpark_columns(),
+        column_metadata=input_container.column_map.column_metadata,
         column_qualifiers=qualifiers,
-        parent_column_name_map=input_df._column_map.get_parent_column_name_map(),
+        parent_column_name_map=input_container.column_map.get_parent_column_name_map(),
+        alias=alias,
     )

snowflake/snowpark_connect/relation/map_udtf.py

@@ -8,7 +8,6 @@ import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 import pyspark.sql.connect.proto.types_pb2 as types_proto
 from pyspark.errors.exceptions.base import PySparkTypeError, PythonException
 
-from snowflake import snowpark
 from snowflake.snowpark.functions import col, parse_json
 from snowflake.snowpark.types import (
     ArrayType,
@@ -17,14 +16,12 @@ from snowflake.snowpark.types import (
     StructType,
     VariantType,
 )
-from snowflake.snowpark_connect.column_name_handler import (
-    ColumnNameMap,
-    with_column_map,
-)
+from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
 from snowflake.snowpark_connect.config import (
     get_boolean_session_config_param,
     global_config,
 )
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.expression.map_expression import (
     map_single_column_expression,
 )
@@ -34,10 +31,7 @@ from snowflake.snowpark_connect.type_mapping import (
     proto_to_snowpark_type,
 )
 from snowflake.snowpark_connect.utils.context import push_udtf_context
-from snowflake.snowpark_connect.utils.session import (
-    get_or_create_snowpark_session,
-    get_python_udxf_import_files,
-)
+from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
 from snowflake.snowpark_connect.utils.udtf_helper import (
     SnowparkUDTF,
     create_udtf_in_sproc,
@@ -45,6 +39,9 @@ from snowflake.snowpark_connect.utils.udtf_helper import (
     udtf_check,
 )
 from snowflake.snowpark_connect.utils.udtf_utils import create_udtf
+from snowflake.snowpark_connect.utils.udxf_import_utils import (
+    get_python_udxf_import_files,
+)
 
 
 def build_expected_types_from_parsed(
@@ -203,7 +200,7 @@ def is_arrow_enabled_in_udtf() -> bool:  # REINSTATED
 
 def map_common_inline_user_defined_table_function(
     rel: relation_proto.CommonInlineUserDefinedTableFunction,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     udtf_check(rel)
     session = get_or_create_snowpark_session()
     python_udft = rel.python_udtf
@@ -280,9 +277,9 @@ def map_common_inline_user_defined_table_function(
 
     snowpark_columns = [f.name for f in output_schema.fields]
 
-    return with_column_map(
-        df,
-        spark_column_names,
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=df,
+        spark_column_names=spark_column_names,
         snowpark_column_names=snowpark_columns,
         snowpark_column_types=snowpark_column_types,
     )

snowflake/snowpark_connect/relation/read/map_read.py

@@ -7,12 +7,14 @@ import json
 import logging
 import os
 import re
+from pathlib import Path
 
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
 from snowflake import snowpark
 from snowflake.snowpark.types import StructType
 from snowflake.snowpark_connect.config import global_config
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.io_utils import (
     convert_file_prefix_path,
     is_cloud_path,
@@ -38,7 +40,7 @@ logger = logging.getLogger("snowflake_connect_server")
 
 def map_read(
     rel: relation_proto.Relation,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Read a file into a Snowpark DataFrame.
 
@@ -91,9 +93,9 @@ def map_read(
             telemetry.report_io_read(read_format, options)
             session: snowpark.Session = get_or_create_snowpark_session()
             if len(rel.read.data_source.paths) > 0:
-                # Clean up trailing slashes from source paths to ensure consistent behavior
+                # Normalize paths to ensure consistent behavior
                 clean_source_paths = [
-                    path.rstrip("/") for path in rel.read.data_source.paths
+                    str(Path(path)) for path in rel.read.data_source.paths
                 ]
 
                 result = _read_file(
@@ -284,6 +286,9 @@ def upload_files_if_needed(
         # overwrite=True will not remove all stale files in the target prefix
 
         remove_command = f"REMOVE {target}/"
+        assert (
+            "//" not in remove_command
+        ), f"Remove command {remove_command} contains double slash"
         session.sql(remove_command).collect()
 
         try:

snowflake/snowpark_connect/relation/read/map_read_csv.py

@@ -10,7 +10,7 @@ import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
 from snowflake.snowpark.dataframe_reader import DataFrameReader
 from snowflake.snowpark.types import StringType, StructField, StructType
-from snowflake.snowpark_connect.column_name_handler import with_column_map
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.read.map_read import CsvReaderConfig
 from snowflake.snowpark_connect.relation.read.utils import (
     get_spark_column_names_from_snowpark_columns,
@@ -27,7 +27,7 @@ def map_read_csv(
     session: snowpark.Session,
     paths: list[str],
     options: CsvReaderConfig,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Read a CSV file into a Snowpark DataFrame.
 
@@ -73,11 +73,11 @@ def map_read_csv(
         renamed_df, snowpark_column_names = rename_columns_as_snowflake_standard(
             df, rel.common.plan_id
         )
-        return with_column_map(
-            renamed_df,
-            spark_column_names,
-            snowpark_column_names,
-            [f.datatype for f in df.schema.fields],
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=renamed_df,
+            spark_column_names=spark_column_names,
+            snowpark_column_names=snowpark_column_names,
+            snowpark_column_types=[f.datatype for f in df.schema.fields],
         )
 
 

snowflake/snowpark_connect/relation/read/map_read_jdbc.py

@@ -8,7 +8,7 @@ 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_connect.column_name_handler import with_column_map
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.read.jdbc_read_dbapi import JdbcDataFrameReader
 from snowflake.snowpark_connect.relation.read.utils import (
     Connection,
@@ -46,7 +46,7 @@ def map_read_jdbc(
     rel: relation_proto.Relation,
     session: snowpark.Session,
     options: dict[str, str],
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Read a table data or query data from a JDBC external datasource into a Snowpark DataFrame.
     """
@@ -98,11 +98,11 @@ def map_read_jdbc(
         renamed_df, snowpark_cols = rename_columns_as_snowflake_standard(
             df, rel.common.plan_id
         )
-        return with_column_map(
-            renamed_df,
-            true_names,
-            snowpark_cols,
-            [f.datatype for f in df.schema.fields],
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=renamed_df,
+            spark_column_names=true_names,
+            snowpark_column_names=snowpark_cols,
+            snowpark_column_types=[f.datatype for f in df.schema.fields],
         )
     except Exception as e:
         raise Exception(f"Error accessing JDBC datasource for read: {e}")

snowflake/snowpark_connect/relation/read/map_read_json.py

@@ -6,6 +6,7 @@ import copy
 import json
 import typing
 from contextlib import suppress
+from datetime import datetime
 
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
@@ -21,8 +22,9 @@ from snowflake.snowpark.types import (
     StringType,
     StructField,
     StructType,
+    TimestampType,
 )
-from snowflake.snowpark_connect.column_name_handler import with_column_map
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.read.map_read import JsonReaderConfig
 from snowflake.snowpark_connect.relation.read.utils import (
     get_spark_column_names_from_snowpark_columns,
@@ -43,7 +45,7 @@ def map_read_json(
     session: snowpark.Session,
     paths: list[str],
     options: JsonReaderConfig,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Read a JSON file into a Snowpark DataFrame.
 
@@ -105,11 +107,11 @@ def map_read_json(
         renamed_df, snowpark_column_names = rename_columns_as_snowflake_standard(
             df, rel.common.plan_id
         )
-        return with_column_map(
-            renamed_df,
-            spark_column_names,
-            snowpark_column_names,
-            [f.datatype for f in df.schema.fields],
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=renamed_df,
+            spark_column_names=spark_column_names,
+            snowpark_column_names=snowpark_column_names,
+            snowpark_column_types=[f.datatype for f in df.schema.fields],
         )
 
 
@@ -204,6 +206,8 @@ def merge_row_schema(
             next_level_content = row[col_name]
             if next_level_content is not None:
                 with suppress(json.JSONDecodeError):
+                    if isinstance(next_level_content, datetime):
+                        next_level_content = str(next_level_content)
                     next_level_content = json.loads(next_level_content)
                 if isinstance(next_level_content, dict):
                     sf.datatype = merge_json_schema(
@@ -235,6 +239,9 @@ def merge_row_schema(
                             sf.datatype.element_type,
                             dropFieldIfAllNull,
                         )
+        elif isinstance(sf.datatype, TimestampType):
+            sf.datatype = StringType()
+            columns_with_valid_contents.add(col_name)
         elif row[col_name] is not None:
             columns_with_valid_contents.add(col_name)
 
@@ -265,7 +272,7 @@ def construct_dataframe_by_schema(
     rows: typing.Iterator[Row],
     session: snowpark.Session,
     snowpark_options: dict,
-    batch_size: int = 100,
+    batch_size: int = 1000,
 ) -> snowpark.DataFrame:
     result = None
 
@@ -280,6 +287,8 @@ def construct_dataframe_by_schema(
                 session,
             )
 
+            current_data = []
+
     if len(current_data) > 0:
         result = union_data_into_df(
             result,
@@ -288,6 +297,8 @@ def construct_dataframe_by_schema(
             session,
         )
 
+        current_data = []
+
     if result is None:
         raise ValueError("Dataframe cannot be empty")
     return result

snowflake/snowpark_connect/relation/read/map_read_parquet.py

@@ -21,7 +21,7 @@ from snowflake.snowpark._internal.analyzer.analyzer_utils import (
 )
 from snowflake.snowpark.column import METADATA_FILENAME
 from snowflake.snowpark.types import DataType, DoubleType, IntegerType, StringType
-from snowflake.snowpark_connect.column_name_handler import with_column_map
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.read.reader_config import ReaderWriterConfig
 from snowflake.snowpark_connect.relation.read.utils import (
     rename_columns_as_snowflake_standard,
@@ -37,7 +37,7 @@ def map_read_parquet(
     session: snowpark.Session,
     paths: list[str],
     options: ReaderWriterConfig,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """Read a Parquet file into a Snowpark DataFrame."""
 
     if rel.read.is_streaming is True:
@@ -62,11 +62,11 @@ def map_read_parquet(
     renamed_df, snowpark_column_names = rename_columns_as_snowflake_standard(
         df, rel.common.plan_id
     )
-    return with_column_map(
-        renamed_df,
-        [analyzer_utils.unquote_if_quoted(c) for c in df.columns],
-        snowpark_column_names,
-        [f.datatype for f in df.schema.fields],
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=renamed_df,
+        spark_column_names=[analyzer_utils.unquote_if_quoted(c) for c in df.columns],
+        snowpark_column_names=snowpark_column_names,
+        snowpark_column_types=[f.datatype for f in df.schema.fields],
     )
 
 

snowflake/snowpark_connect/relation/read/map_read_socket.py

@@ -8,7 +8,7 @@ import pandas
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
 from snowflake import snowpark
-from snowflake.snowpark_connect.column_name_handler import with_column_map
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
@@ -24,7 +24,7 @@ def map_read_socket(
     rel: relation_proto.Relation,
     session: snowpark.Session,
     options: dict[str, str],
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     if rel.read.is_streaming is True:
         global full_data
         host = options.get("host", None)
@@ -50,7 +50,11 @@ def map_read_socket(
                     pandas.DataFrame({snowpark_cname: dataframe_data.split("\n")})
                 )
                 spark_cname = "value"
-                return with_column_map(df, [spark_cname], [snowpark_cname])
+                return DataFrameContainer.create_with_column_mapping(
+                    dataframe=df,
+                    spark_column_names=[spark_cname],
+                    snowpark_column_names=[snowpark_cname],
+                )
             except OSError as e:
                 raise Exception(f"Error connecting to {host}:{port} - {e}")
     else: