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

snowflake/snowpark_connect/relation/map_extension.py

@@ -345,7 +345,7 @@ def map_aggregate(
         return new_names[0], snowpark_column
 
     raw_groupings: list[tuple[str, TypedColumn]] = []
-    raw_aggregations: list[tuple[str, TypedColumn]] = []
+    raw_aggregations: list[tuple[str, TypedColumn, list[str]]] = []
 
     if not is_group_by_all:
         raw_groupings = [_map_column(exp) for exp in aggregate.grouping_expressions]
@@ -375,10 +375,21 @@ def map_aggregate(
     # Note: We don't clear the map here to preserve any parent context aliases
     from snowflake.snowpark_connect.utils.context import register_lca_alias
 
+    # If it's an unresolved attribute when its in aggregate.aggregate_expressions, we know it came from the parent map straight away
+    # in this case, we should see if the parent map has a qualifier for it and propagate that here, in case the order by references it in
+    # a qualified way later.
     agg_count = get_sql_aggregate_function_count()
     for exp in aggregate.aggregate_expressions:
         col = _map_column(exp)
-        raw_aggregations.append(col)
+        if exp.WhichOneof("expr_type") == "unresolved_attribute":
+            spark_name = col[0]
+            qualifiers = input_container.column_map.get_qualifier_for_spark_column(
+                spark_name
+            )
+        else:
+            qualifiers = []
+
+        raw_aggregations.append((col[0], col[1], qualifiers))
 
         # If this is an alias, register it in the LCA map for subsequent expressions
         if (
@@ -409,18 +420,20 @@ def map_aggregate(
     spark_columns: list[str] = []
     snowpark_columns: list[str] = []
     snowpark_column_types: list[snowpark_types.DataType] = []
+    all_qualifiers: list[list[str]] = []
 
     # Use grouping columns directly without aliases
     groupings = [col.col for _, col in raw_groupings]
 
     # Create aliases only for aggregation columns
     aggregations = []
-    for i, (spark_name, snowpark_column) in enumerate(raw_aggregations):
+    for i, (spark_name, snowpark_column, qualifiers) in enumerate(raw_aggregations):
         alias = make_column_names_snowpark_compatible([spark_name], plan_id, i)[0]
 
         spark_columns.append(spark_name)
         snowpark_columns.append(alias)
         snowpark_column_types.append(snowpark_column.typ)
+        all_qualifiers.append(qualifiers)
 
         aggregations.append(snowpark_column.col.alias(alias))
 
@@ -483,6 +496,7 @@ def map_aggregate(
             spark_column_names=spark_columns,
             snowpark_column_names=snowpark_columns,
             snowpark_column_types=snowpark_column_types,
+            column_qualifiers=all_qualifiers,
         ).column_map
 
         # Create hybrid column map that can resolve both input and aggregate contexts
@@ -494,7 +508,9 @@ def map_aggregate(
             aggregate_expressions=list(aggregate.aggregate_expressions),
             grouping_expressions=list(aggregate.grouping_expressions),
             spark_columns=spark_columns,
-            raw_aggregations=raw_aggregations,
+            raw_aggregations=[
+                (spark_name, col) for spark_name, col, _ in raw_aggregations
+            ],
         )
 
         # Map the HAVING condition using hybrid resolution
@@ -515,4 +531,5 @@ def map_aggregate(
         snowpark_column_names=snowpark_columns,
         snowpark_column_types=snowpark_column_types,
         parent_column_name_map=input_df._column_map,
+        column_qualifiers=all_qualifiers,
     )

snowflake/snowpark_connect/relation/map_map_partitions.py

@@ -12,7 +12,6 @@ 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.context import map_partitions_depth
 from snowflake.snowpark_connect.utils.pandas_udtf_utils import (
     create_pandas_udtf,
     create_pandas_udtf_with_arrow,
@@ -53,18 +52,18 @@ def _call_udtf(
         ).cast("int"),
     )
 
-    udtf_columns = input_df.columns + [snowpark_fn.col("_DUMMY_PARTITION_KEY")]
+    udtf_columns = [f"snowflake_jtf_{column}" for column in input_df.columns] + [
+        "_DUMMY_PARTITION_KEY"
+    ]
 
     tfc = snowpark_fn.call_table_function(udtf_name, *udtf_columns).over(
         partition_by=[snowpark_fn.col("_DUMMY_PARTITION_KEY")]
     )
 
-    # Use map_partitions_depth only when mapping non nested map_partitions
-    # When mapping chained functions additional column casting is necessary
-    if map_partitions_depth() == 1:
-        result_df_with_dummy = input_df_with_dummy.join_table_function(tfc)
-    else:
-        result_df_with_dummy = input_df_with_dummy.select(tfc)
+    # Overwrite the input_df columns to prevent name conflicts with UDTF output columns
+    result_df_with_dummy = input_df_with_dummy.to_df(udtf_columns).join_table_function(
+        tfc
+    )
 
     output_cols = [field.name for field in return_type.fields]
 

snowflake/snowpark_connect/relation/map_relation.py

@@ -16,7 +16,6 @@ from snowflake.snowpark_connect.utils.context import (
     get_plan_id_map,
     get_session_id,
     not_resolving_fun_args,
-    push_map_partitions,
     push_operation_scope,
     set_is_aggregate_function,
     set_plan_id_map,
@@ -185,8 +184,7 @@ def map_relation(
                     )
                 return cached_df
             case "map_partitions":
-                with push_map_partitions():
-                    result = map_map_partitions.map_map_partitions(rel)
+                result = map_map_partitions.map_map_partitions(rel)
             case "offset":
                 result = map_row_ops.map_offset(rel)
             case "project":

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
 
@@ -16,6 +17,7 @@ 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 +44,34 @@ 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 = session.read.options(snowpark_read_options)
         else:
-            reader = session.read.options(snowpark_options).schema(schema)
+            reader = session.read.options(snowpark_read_options).schema(schema)
         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 +96,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,8 +168,10 @@ 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]
@@ -120,23 +187,35 @@ def read_data(
                     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 = df.schema.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
 
@@ -253,20 +256,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 +279,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_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"'