snowpark-connect 0.28.1__py3-none-any.whl → 0.30.0__py3-none-any.whl

snowflake/snowpark_connect/relation/read/map_read.py

@@ -46,6 +46,9 @@ def map_read(
 
     Currently, the supported read formats are `csv`, `json` and `parquet`.
     """
+
+    materialize_df = True
+
     match rel.read.WhichOneof("read_type"):
         case "named_table":
             return map_read_table_or_file(rel)
@@ -99,6 +102,10 @@ def map_read(
                     for path in rel.read.data_source.paths
                 ]
 
+                # JSON already materializes the table internally
+                if read_format == "json":
+                    materialize_df = False
+
                 result = _read_file(
                     clean_source_paths, options, read_format, rel, schema, session
                 )
@@ -159,7 +166,9 @@ def map_read(
             raise SnowparkConnectNotImplementedError(f"Unsupported read type: {other}")
 
     return df_cache_map_put_if_absent(
-        (get_session_id(), rel.common.plan_id), lambda: result, materialize=True
+        (get_session_id(), rel.common.plan_id),
+        lambda: result,
+        materialize=materialize_df,
     )
 
 
@@ -205,6 +214,15 @@ def _get_supported_read_file_format(unparsed_identifier: str) -> str | None:
     return None
 
 
+def _quote_stage_path(stage_path: str) -> str:
+    """
+    Quote stage paths to escape any special characters.
+    """
+    if stage_path.startswith("@"):
+        return f"'{stage_path}'"
+    return stage_path
+
+
 def _read_file(
     clean_source_paths: list[str],
     options: dict,
@@ -218,6 +236,7 @@ def _read_file(
         session,
     )
     upload_files_if_needed(paths, clean_source_paths, session, read_format)
+    paths = [_quote_stage_path(path) for path in paths]
     match read_format:
         case "csv":
             from snowflake.snowpark_connect.relation.read.map_read_csv import (
@@ -285,8 +304,8 @@ def upload_files_if_needed(
 
     def _upload_dir(target: str, source: str) -> None:
         # overwrite=True will not remove all stale files in the target prefix
-
-        remove_command = f"REMOVE {target}/"
+        # Quote the target path to allow special characters.
+        remove_command = f"REMOVE '{target}/'"
         assert (
             "//" not in remove_command
         ), f"Remove command {remove_command} contains double slash"

snowflake/snowpark_connect/relation/read/map_read_csv.py

@@ -3,6 +3,7 @@
 #
 
 import copy
+from typing import Any
 
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
@@ -12,10 +13,15 @@ from snowflake.snowpark.dataframe_reader import DataFrameReader
 from snowflake.snowpark.types import StringType, StructField, StructType
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.read.map_read import CsvReaderConfig
+from snowflake.snowpark_connect.relation.read.metadata_utils import (
+    add_filename_metadata_to_reader,
+    get_non_metadata_fields,
+)
 from snowflake.snowpark_connect.relation.read.utils import (
     get_spark_column_names_from_snowpark_columns,
     rename_columns_as_snowflake_standard,
 )
+from snowflake.snowpark_connect.utils.io_utils import cached_file_format
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
@@ -42,21 +48,39 @@ def map_read_csv(
         )
     else:
         snowpark_options = options.convert_to_snowpark_args()
+        parse_header = snowpark_options.get("PARSE_HEADER", False)
+        file_format_options = _parse_csv_snowpark_options(snowpark_options)
+        file_format = cached_file_format(session, "csv", file_format_options)
+
+        snowpark_read_options = dict()
+        snowpark_read_options["FORMAT_NAME"] = file_format
+        snowpark_read_options["ENFORCE_EXISTING_FILE_FORMAT"] = True
+        snowpark_read_options["INFER_SCHEMA"] = snowpark_options.get(
+            "INFER_SCHEMA", False
+        )
+        snowpark_read_options["PATTERN"] = snowpark_options.get("PATTERN", None)
+
         raw_options = rel.read.data_source.options
+
         if schema is None or (
-            snowpark_options.get("PARSE_HEADER", False)
-            and raw_options.get("enforceSchema", "True").lower() == "false"
+            parse_header and raw_options.get("enforceSchema", "True").lower() == "false"
         ):  # Schema has to equals to header's format
-            reader = session.read.options(snowpark_options)
+            reader = add_filename_metadata_to_reader(
+                session.read.options(snowpark_options), raw_options
+            )
         else:
-            reader = session.read.options(snowpark_options).schema(schema)
+            reader = add_filename_metadata_to_reader(
+                session.read.options(snowpark_options).schema(schema), raw_options
+            )
         df = read_data(
             reader,
             schema,
             session,
             paths[0],
-            snowpark_options,
+            file_format_options,
+            snowpark_read_options,
             raw_options,
+            parse_header,
         )
         if len(paths) > 1:
             # TODO: figure out if this is what Spark does.
@@ -81,15 +105,65 @@ def map_read_csv(
         )
 
 
+_csv_file_format_allowed_options = {
+    "COMPRESSION",
+    "RECORD_DELIMITER",
+    "FIELD_DELIMITER",
+    "MULTI_LINE",
+    "FILE_EXTENSION",
+    "PARSE_HEADER",
+    "SKIP_HEADER",
+    "SKIP_BLANK_LINES",
+    "DATE_FORMAT",
+    "TIME_FORMAT",
+    "TIMESTAMP_FORMAT",
+    "BINARY_FORMAT",
+    "ESCAPE",
+    "ESCAPE_UNENCLOSED_FIELD",
+    "TRIM_SPACE",
+    "FIELD_OPTIONALLY_ENCLOSED_BY",
+    "NULL_IF",
+    "ERROR_ON_COLUMN_COUNT_MISMATCH",
+    "REPLACE_INVALID_CHARACTERS",
+    "EMPTY_FIELD_AS_NULL",
+    "SKIP_BYTE_ORDER_MARK",
+    "ENCODING",
+}
+
+
+def _parse_csv_snowpark_options(snowpark_options: dict[str, Any]) -> dict[str, Any]:
+    file_format_options = dict()
+    for key, value in snowpark_options.items():
+        upper_key = key.upper()
+        if upper_key in _csv_file_format_allowed_options:
+            file_format_options[upper_key] = value
+
+    # This option has to be removed, because we cannot use at the same time predefined file format and parse_header option
+    # Such combination causes snowpark to raise SQL compilation error: Invalid file format "PARSE_HEADER" is only allowed for CSV INFER_SCHEMA and MATCH_BY_COLUMN_NAME
+    parse_header = file_format_options.get("PARSE_HEADER", False)
+    if parse_header:
+        file_format_options["SKIP_HEADER"] = 1
+        del file_format_options["PARSE_HEADER"]
+
+    return file_format_options
+
+
 def get_header_names(
     session: snowpark.Session,
     path: list[str],
-    snowpark_options: dict,
+    file_format_options: dict,
+    snowpark_read_options: dict,
 ) -> list[str]:
-    snowpark_options_no_header = copy.copy(snowpark_options)
-    snowpark_options_no_header["PARSE_HEADER"] = False
+    no_header_file_format_options = copy.copy(file_format_options)
+    no_header_file_format_options["PARSE_HEADER"] = False
+    no_header_file_format_options.pop("SKIP_HEADER", None)
+
+    file_format = cached_file_format(session, "csv", no_header_file_format_options)
+    no_header_snowpark_read_options = copy.copy(snowpark_read_options)
+    no_header_snowpark_read_options["FORMAT_NAME"] = file_format
+    no_header_snowpark_read_options.pop("INFER_SCHEMA", None)
 
-    header_df = session.read.options(snowpark_options_no_header).csv(path).limit(1)
+    header_df = session.read.options(no_header_snowpark_read_options).csv(path).limit(1)
     header_data = header_df.collect()[0]
     return [
         f'"{header_data[i]}"'
@@ -103,40 +177,56 @@ def read_data(
     schema: snowpark.types.StructType | None,
     session: snowpark.Session,
     path: list[str],
-    snowpark_options: dict,
+    file_format_options: dict,
+    snowpark_read_options: dict,
     raw_options: dict,
+    parse_header: bool,
 ) -> snowpark.DataFrame:
     df = reader.csv(path)
     filename = path.strip("/").split("/")[-1]
+    non_metadata_fields = get_non_metadata_fields(df.schema.fields)
+
     if schema is not None:
-        if len(schema.fields) != len(df.schema.fields):
+        if len(schema.fields) != len(non_metadata_fields):
             raise Exception(f"csv load from {filename} failed.")
         if raw_options.get("enforceSchema", "True").lower() == "false":
             for i in range(len(schema.fields)):
                 if (
-                    schema.fields[i].name != df.schema.fields[i].name
-                    and f'"{schema.fields[i].name}"' != df.schema.fields[i].name
+                    schema.fields[i].name != non_metadata_fields[i].name
+                    and f'"{schema.fields[i].name}"' != non_metadata_fields[i].name
                 ):
                     raise Exception("CSV header does not conform to the schema")
         return df
 
-    headers = get_header_names(session, path, snowpark_options)
-
+    headers = get_header_names(
+        session, path, file_format_options, snowpark_read_options
+    )
+
+    df_schema_fields = non_metadata_fields
+    if len(headers) == len(df_schema_fields) and parse_header:
+        return df.select(
+            [
+                snowpark_fn.col(df_schema_fields[i].name).alias(headers[i])
+                for i in range(len(headers))
+            ]
+        )
     # Handle mismatch in column count between header and data
-    if (
-        len(df.schema.fields) == 1
-        and df.schema.fields[0].name.upper() == "C1"
-        and snowpark_options.get("PARSE_HEADER") is True
-        and len(headers) != len(df.schema.fields)
+    elif (
+        len(df_schema_fields) == 1
+        and df_schema_fields[0].name.upper() == "C1"
+        and parse_header
+        and len(headers) != len(df_schema_fields)
     ):
-        df = (
-            session.read.options(snowpark_options)
-            .schema(StructType([StructField(h, StringType(), True) for h in headers]))
-            .csv(path)
+        df = reader.schema(
+            StructType([StructField(h, StringType(), True) for h in headers])
+        ).csv(path)
+    elif not parse_header and len(headers) != len(df_schema_fields):
+        return df.select([df_schema_fields[i].name for i in range(len(headers))])
+    elif parse_header and len(headers) != len(df_schema_fields):
+        return df.select(
+            [
+                snowpark_fn.col(df_schema_fields[i].name).alias(headers[i])
+                for i in range(len(headers))
+            ]
         )
-    elif snowpark_options.get("PARSE_HEADER") is False and len(headers) != len(
-        df.schema.fields
-    ):
-        return df.select([df.schema.fields[i].name for i in range(len(headers))])
-
     return df

snowflake/snowpark_connect/relation/read/map_read_json.py

@@ -2,9 +2,12 @@
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
 
+import concurrent.futures
 import copy
 import json
+import os
 import typing
+import uuid
 from contextlib import suppress
 from datetime import datetime
 
@@ -26,6 +29,9 @@ from snowflake.snowpark.types import (
 )
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.read.map_read import JsonReaderConfig
+from snowflake.snowpark_connect.relation.read.metadata_utils import (
+    add_filename_metadata_to_reader,
+)
 from snowflake.snowpark_connect.relation.read.utils import (
     get_spark_column_names_from_snowpark_columns,
     rename_columns_as_snowflake_standard,
@@ -63,19 +69,26 @@ def map_read_json(
         )
     else:
         snowpark_options = options.convert_to_snowpark_args()
+        raw_options = rel.read.data_source.options
         snowpark_options["infer_schema"] = True
 
         rows_to_infer_schema = snowpark_options.pop("rowstoinferschema", 1000)
         dropFieldIfAllNull = snowpark_options.pop("dropfieldifallnull", False)
         batch_size = snowpark_options.pop("batchsize", 1000)
 
-        reader = session.read.options(snowpark_options)
+        reader = add_filename_metadata_to_reader(
+            session.read.options(snowpark_options), raw_options
+        )
 
         df = reader.json(paths[0])
         if len(paths) > 1:
             # TODO: figure out if this is what Spark does.
             for p in paths[1:]:
-                df = df.union_all(session.read.options(snowpark_options).json(p))
+                df = df.union_all(
+                    add_filename_metadata_to_reader(
+                        session.read.options(snowpark_options), raw_options
+                    ).json(p)
+                )
 
         if schema is None:
             schema = copy.deepcopy(df.schema)
@@ -253,20 +266,20 @@ def merge_row_schema(
     return schema
 
 
-def union_data_into_df(
-    result_df: snowpark.DataFrame,
-    data: typing.List[Row],
-    schema: StructType,
+def insert_data_chunk(
     session: snowpark.Session,
-) -> snowpark.DataFrame:
-    current_df = session.create_dataframe(
+    data: list[Row],
+    schema: StructType,
+    table_name: str,
+) -> None:
+    df = session.create_dataframe(
         data=data,
         schema=schema,
     )
-    if result_df is None:
-        return current_df
 
-    return result_df.union(current_df)
+    df.write.mode("append").save_as_table(
+        table_name, table_type="temp", table_exists=True
+    )
 
 
 def construct_dataframe_by_schema(
@@ -276,39 +289,47 @@ def construct_dataframe_by_schema(
     snowpark_options: dict,
     batch_size: int = 1000,
 ) -> snowpark.DataFrame:
-    result = None
+    table_name = "__sas_json_read_temp_" + uuid.uuid4().hex
+
+    # We can have more workers than CPU count, this is an IO-intensive task
+    max_workers = min(16, os.cpu_count() * 2)
 
     current_data = []
     progress = 0
-    for row in rows:
-        current_data.append(construct_row_by_schema(row, schema, snowpark_options))
-        if len(current_data) >= batch_size:
+
+    # Initialize the temp table
+    session.create_dataframe([], schema=schema).write.mode("append").save_as_table(
+        table_name, table_type="temp", table_exists=False
+    )
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as exc:
+        for row in rows:
+            current_data.append(construct_row_by_schema(row, schema, snowpark_options))
+            if len(current_data) >= batch_size:
+                progress += len(current_data)
+                exc.submit(
+                    insert_data_chunk,
+                    session,
+                    copy.deepcopy(current_data),
+                    schema,
+                    table_name,
+                )
+
+                logger.info(f"JSON reader: finished processing {progress} rows")
+                current_data.clear()
+
+        if len(current_data) > 0:
             progress += len(current_data)
-            result = union_data_into_df(
-                result,
-                current_data,
-                schema,
+            exc.submit(
+                insert_data_chunk,
                 session,
+                copy.deepcopy(current_data),
+                schema,
+                table_name,
             )
-
             logger.info(f"JSON reader: finished processing {progress} rows")
-            current_data = []
-
-    if len(current_data) > 0:
-        progress += len(current_data)
-        result = union_data_into_df(
-            result,
-            current_data,
-            schema,
-            session,
-        )
-
-        logger.info(f"JSON reader: finished processing {progress} rows")
-        current_data = []
 
-    if result is None:
-        raise ValueError("Dataframe cannot be empty")
-    return result
+    return session.table(table_name)
 
 
 def construct_row_by_schema(

snowflake/snowpark_connect/relation/read/map_read_parquet.py

@@ -22,6 +22,9 @@ 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.dataframe_container import DataFrameContainer
+from snowflake.snowpark_connect.relation.read.metadata_utils import (
+    add_filename_metadata_to_reader,
+)
 from snowflake.snowpark_connect.relation.read.reader_config import ReaderWriterConfig
 from snowflake.snowpark_connect.relation.read.utils import (
     rename_columns_as_snowflake_standard,
@@ -46,10 +49,13 @@ def map_read_parquet(
         )
 
     snowpark_options = options.convert_to_snowpark_args()
+    raw_options = rel.read.data_source.options
     assert schema is None, "Read PARQUET does not support user schema"
     assert len(paths) > 0, "Read PARQUET expects at least one path"
 
-    reader = session.read.options(snowpark_options)
+    reader = add_filename_metadata_to_reader(
+        session.read.options(snowpark_options), raw_options
+    )
 
     if len(paths) == 1:
         df = _read_parquet_with_partitions(session, reader, paths[0])

snowflake/snowpark_connect/relation/read/map_read_text.py

@@ -43,7 +43,12 @@ def read_text(
 ) -> snowpark.DataFrame:
     # TODO: handle stage name with double quotes
     files_paths = get_file_paths_from_stage(path, session)
-    stage_name = path.split("/")[0]
+    # Remove matching quotes from both ends of the path to get the stage name, if present.
+    if path and len(path) > 1 and path[0] == path[-1] and path[0] in ('"', "'"):
+        unquoted_path = path[1:-1]
+    else:
+        unquoted_path = path
+    stage_name = unquoted_path.split("/")[0]
     line_sep = options.get("lineSep") or "\n"
     column_name = (
         schema[0].name if schema is not None and len(schema.fields) > 0 else '"value"'

snowflake/snowpark_connect/relation/read/metadata_utils.py

@@ -0,0 +1,159 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+"""
+Utilities for handling internal metadata columns in file-based DataFrames.
+"""
+
+import os
+
+import pandas
+from pyspark.errors.exceptions.base import AnalysisException
+
+from snowflake import snowpark
+from snowflake.snowpark.column import METADATA_FILENAME
+from snowflake.snowpark.functions import col
+from snowflake.snowpark.types import StructField
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
+
+# Constant for the metadata filename column name
+METADATA_FILENAME_COLUMN = "METADATA$FILENAME"
+
+
+def add_filename_metadata_to_reader(
+    reader: snowpark.DataFrameReader,
+    options: dict | None = None,
+) -> snowpark.DataFrameReader:
+    """
+    Add filename metadata to a DataFrameReader based on configuration.
+
+    Args:
+        reader: Snowpark DataFrameReader instance
+        options: Dictionary of options to check for metadata configuration
+
+    Returns:
+        DataFrameReader with filename metadata enabled if configured, otherwise unchanged
+    """
+    # NOTE: SNOWPARK_POPULATE_FILE_METADATA_DEFAULT is an internal environment variable
+    # used only for CI testing to verify no metadata columns leak in regular file operations.
+    # This environment variable should NOT be exposed to end users. Users should only use snowpark.populateFileMetadata
+    # to enable metadata population.
+    metadata_default = os.environ.get(
+        "SNOWPARK_POPULATE_FILE_METADATA_DEFAULT", "false"
+    )
+
+    populate_metadata = (
+        options.get("snowpark.populateFileMetadata", metadata_default)
+        if options
+        else metadata_default
+    ).lower() == "true"
+
+    if populate_metadata:
+        return reader.with_metadata(METADATA_FILENAME)
+    else:
+        return reader
+
+
+def get_non_metadata_fields(schema_fields: list[StructField]) -> list[StructField]:
+    """
+    Filter out METADATA$FILENAME fields from a list of schema fields.
+
+    Args:
+        schema_fields: List of StructField objects from a DataFrame schema
+
+    Returns:
+        List of StructField objects excluding METADATA$FILENAME
+    """
+    return [field for field in schema_fields if field.name != METADATA_FILENAME_COLUMN]
+
+
+def get_non_metadata_column_names(schema_fields: list[StructField]) -> list[str]:
+    """
+    Get column names from schema fields, excluding METADATA$FILENAME.
+
+    Args:
+        schema_fields: List of StructField objects from a DataFrame schema
+
+    Returns:
+        List of column names (strings) excluding METADATA$FILENAME
+    """
+    return [
+        field.name for field in schema_fields if field.name != METADATA_FILENAME_COLUMN
+    ]
+
+
+def filter_metadata_column_name(column_names: list[str]) -> list[str]:
+    """
+    Get column names from column_names, excluding METADATA$FILENAME.
+
+    Returns:
+        List of column names (strings) excluding METADATA$FILENAME
+    """
+    return [
+        col_name for col_name in column_names if col_name != METADATA_FILENAME_COLUMN
+    ]
+
+
+def filter_metadata_columns(
+    result_container: DataFrameContainer | pandas.DataFrame | None,
+) -> DataFrameContainer | pandas.DataFrame | None:
+    """
+    Filter METADATA$FILENAME from DataFrame container for execution and write operations.
+
+    Args:
+        result_container: DataFrameContainer or pandas DataFrame to filter
+
+    Returns:
+        Filtered container (callers can access dataframe via container.dataframe)
+    """
+    # Handle pandas DataFrame case - return as-is
+    if isinstance(result_container, pandas.DataFrame):
+        return result_container
+
+    if result_container is None:
+        return None
+
+    result_df = result_container.dataframe
+    if not isinstance(result_df, snowpark.DataFrame):
+        return result_container
+
+    df_columns = result_container.column_map.get_snowpark_columns()
+    has_metadata_filename = any(name == METADATA_FILENAME_COLUMN for name in df_columns)
+
+    if not has_metadata_filename:
+        return result_container
+
+    non_metadata_columns = filter_metadata_column_name(df_columns)
+
+    if len(non_metadata_columns) == 0:
+        # DataFrame contains only metadata columns (METADATA$FILENAME), no actual data columns remaining.
+        # We don't have a way to return an empty dataframe.
+        raise AnalysisException(
+            "[DATAFRAME_MISSING_DATA_COLUMNS] Cannot perform operation on DataFrame that contains no data columns."
+        )
+
+    filtered_df = result_df.select([col(name) for name in non_metadata_columns])
+
+    original_spark_columns = result_container.column_map.get_spark_columns()
+    original_snowpark_columns = result_container.column_map.get_snowpark_columns()
+
+    filtered_spark_columns = []
+    filtered_snowpark_columns = []
+
+    for i, colname in enumerate(df_columns):
+        if colname != METADATA_FILENAME_COLUMN:
+            filtered_spark_columns.append(original_spark_columns[i])
+            filtered_snowpark_columns.append(original_snowpark_columns[i])
+
+    new_container = DataFrameContainer.create_with_column_mapping(
+        dataframe=filtered_df,
+        spark_column_names=filtered_spark_columns,
+        snowpark_column_names=filtered_snowpark_columns,
+        column_metadata=result_container.column_map.column_metadata,
+        table_name=result_container.table_name,
+        alias=result_container.alias,
+        partition_hint=result_container.partition_hint,
+    )
+
+    return new_container