snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/relation/read/map_read_text.py

@@ -0,0 +1,106 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import typing
+
+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.relation.read.utils import (
+    get_spark_column_names_from_snowpark_columns,
+    rename_columns_as_snowflake_standard,
+)
+from snowflake.snowpark_connect.utils.io_utils import file_format
+from snowflake.snowpark_connect.utils.telemetry import (
+    SnowparkConnectNotImplementedError,
+)
+
+
+def get_file_paths_from_stage(
+    path: str,
+    session: snowpark.Session,
+) -> typing.List[str]:
+    files_paths = []
+    for listed_path_row in session.sql(f"LIST {path}").collect():
+        listed_path = listed_path_row[0].split("/")
+        if listed_path_row[0].startswith("s3://") or listed_path_row[0].startswith(
+            "s3a://"
+        ):
+            listed_path = listed_path[3:]
+        else:
+            listed_path = listed_path[1:]
+        files_paths.append("/".join(listed_path))
+    return files_paths
+
+
+def read_text(
+    path: str,
+    schema: snowpark.types.StructType | None,
+    session: snowpark.Session,
+    options: typing.MutableMapping[str, str],
+) -> snowpark.DataFrame:
+    # TODO: handle stage name with double quotes
+    files_paths = get_file_paths_from_stage(path, session)
+    stage_name = 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"'
+    )
+    default_column_name = "TEXT"
+
+    result = []
+    separator = (
+        None if options.get("wholetext", "False").lower() == "true" else line_sep
+    )
+    text_file_format = file_format(
+        session, options.get("compression", "auto"), separator
+    )
+    for fp in files_paths:
+        content = session.sql(
+            f"SELECT T.$1 AS {default_column_name} FROM {stage_name}/{fp} (FILE_FORMAT => {text_file_format}) AS T"
+        ).collect()
+        for row in content:
+            result.append(row[0])
+    return session.createDataFrame(result, [column_name])
+
+
+def map_read_text(
+    rel: relation_proto.Relation,
+    schema: snowpark.types.StructType | None,
+    session: snowpark.Session,
+    paths: list[str],
+) -> snowpark.DataFrame:
+    """
+    Read a TEXT file into a Snowpark DataFrame.
+    """
+    if rel.read.is_streaming is True:
+        # TODO: Structured streaming implementation.
+        raise SnowparkConnectNotImplementedError(
+            "Streaming is not supported for CSV files."
+        )
+
+    df = read_text(paths[0], schema, session, rel.read.data_source.options)
+    if len(paths) > 1:
+        for p in paths[1:]:
+            df = df.union_all(
+                read_text(
+                    p,
+                    schema,
+                    session,
+                    rel.read.data_source.options,
+                )
+            )
+
+    spark_column_names = get_spark_column_names_from_snowpark_columns(df.columns)
+
+    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],
+    )

snowflake/snowpark_connect/relation/read/reader_config.py

@@ -0,0 +1,399 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+from dataclasses import dataclass
+from typing import Any
+
+from snowflake.snowpark_connect.config import str_to_bool
+from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
+
+
+@dataclass
+class _Config:
+    default_config: dict[str, str]
+    supported_options: set[str]
+    boolean_config_list: list[str]
+    int_config_list: list[str]
+    float_config_list: list[str]
+
+
+# TODO: There is a issue where we don't differentiate between our defaults, and what user has explicitly provided.
+# For explicit options that are provided, a visible warning will be better.
+class ReaderWriterConfig:
+
+    # Default global configuration for Snowpark Connect.
+    # All keys must be lowercase in order to preserve case configuration insensitivity.
+    # This mimics how the Spark works.
+    default_global_config = {
+        # TODO: Snowpark does not support mode argument
+        # "mode": "PERMISSIVE",
+        # TODO: Snowpark does not support locale argument
+        # "locale": "en-US",
+        "compression": "auto",
+    }
+
+    def __init__(self, config: _Config, options: dict[str, str]) -> None:
+        self.config = lowercase_dict_keys(
+            self.default_global_config
+        ) | lowercase_dict_keys(config.default_config)
+        self.supported_options = {option.lower() for option in config.supported_options}
+        self.boolean_config_list = [
+            option.lower() for option in config.boolean_config_list
+        ]
+        self.int_config_list = [option.lower() for option in config.int_config_list]
+        self.float_config_list = [option.lower() for option in config.float_config_list]
+
+        for key, value in options.items():
+            self.config[key.lower()] = value
+
+    def _get_config_setting(self, key: str) -> bool | int | float | str | None:
+        """Get the configuration setting for the key based on the setting type."""
+        if key in self.boolean_config_list:
+            return str_to_bool(self.config[key])
+        elif key in self.int_config_list:
+            return int(self.config[key])
+        elif key in self.float_config_list:
+            return float(self.config[key])
+        else:
+            return self.config[key]
+
+    # TODO: When we convert into args, we cannot only convert the key, we need to adjust the value also.
+    # For example, for differences in timestamp format.
+    def convert_to_snowpark_args(self) -> dict[str, Any]:
+        snowpark_config = {}
+
+        for key, value in self.config.items():
+            if key in self.supported_options:
+                snowpark_config[key] = value
+            else:
+                logger.debug(f"unsupported reader option: {key}")
+
+        for key in snowpark_config.keys():
+            snowpark_config[key] = self._get_config_setting(key)
+        return snowpark_config
+
+    def get(self, key: str) -> str | None:
+        return self.config.get(key.lower(), None)
+
+
+def lowercase_dict_keys(d: dict[str, Any]) -> dict[str, Any]:
+    """Convert all keys in the dictionary to lowercase."""
+    return {key.lower(): value for key, value in d.items()}
+
+
+def lowercase_set(s: set[str]) -> set[str]:
+    return {value.lower() for value in s}
+
+
+# Config has to be lowercased, because it is publicly available
+CSV_READ_SUPPORTED_OPTIONS = lowercase_set(
+    {
+        "schema",
+        "sep",
+        "encoding",
+        "quote",
+        # escape has different semantics in snowpark, but should work for standard use-cases
+        "escape",
+        # "comment", # Comment is not supported
+        "header",
+        "inferSchema",
+        # "ignoreLeadingWhiteSpace",
+        # "ignoreTrailingWhiteSpace",
+        "nullValue",
+        # "nanValue",
+        # "positiveInf",
+        # "negativeInf",
+        "dateFormat",
+        "timestampFormat",
+        # "maxColumns",
+        # "maxCharsPerColumn",
+        # "maxMalformedLogPerPartition",
+        # "mode",
+        # "columnNameOfCorruptRecord",
+        "multiLine",
+        # "charToEscapeQuoteEscaping",
+        # "samplingRatio",
+        # "enforceSchema",
+        # "emptyValue",
+        # "locale",
+        "lineSep",
+        "pathGlobFilter",
+        # "recursiveFileLookup",
+        # "modifiedBefore",
+        # "modifiedAfter",
+        # "unescapedQuoteHandling",
+        "compression",
+        # "escapeQuotes",
+        # "quoteAll",
+    }
+)
+
+# Config has to be lowercased, because it is publicly available
+CSV_READ_DEFAULT_CONFIG = lowercase_dict_keys(
+    {
+        "header": "false",
+        "inferSchema": "true",
+        # TODO: This default is ok for reads, but it should be removed for writes because it will lead to
+        # quoting even when it is not necessary.
+        "quote": '"',
+        "comment": "#",
+        # TODO nullValue of "" is correct for Spark, Snowflake's default is \N.
+        # However, Snowflake will refuse to write when nullValue is empty and fields aren't
+        # optionally enclosed. Hence, we are not changing the default nullValue.
+        # We need to look more to see if this is a good default for both write+read, or if we need to adjust.
+        # "nullValue": "",
+        # TODO: Snowpark does not support NaN value argument.
+        # "nanValue": "NaN",
+        "dateFormat": "yyyy-MM-dd",
+        "timestampFormat": "YYYY-MM-DD HH24:MI:SS.FF6",
+        # TODO: Snowpark does not support maxColumns argument
+        # "maxColumns": "20480",
+        # TODO: Snowpark does not support maxCharsPerColumn argument
+        # "maxCharsPerColumn": "1000000",
+        # TODO: Snowpark does not support maxMalformedLogPerPartition argument
+        # "maxMalformedLogPerPartition": "10",
+        "charset": "UTF-8",
+        # TODO: Snowpark does not support multiLine argument
+        "multiLine": "false",
+        # TODO: Snowpark does not support ignoreLeadingWhiteSpace argument
+        # "ignoreLeadingWhiteSpace": "false",
+        # TODO: Snowpark does not support ignoreTrailingWhiteSpace argument
+        # "ignoreTrailingWhiteSpace": "false",
+        # TODO: Snowpark does not support samplingRatio argument
+        # "samplingRatio": "1.0",
+        # TODO: Snowpark does not support emptyValue argument
+        # "emptyValue": "",
+        "lineSep": "\n",
+        "sep": ",",
+        # TODO: Snowpark does not support escapeQuotes argument
+        # "escapeQuotes": "true",
+        # TODO: Snowpark does not support quoteAll argument
+        # "quoteAll": "false",
+        "escape": "\\\\",
+    }
+)
+
+
+def csv_convert_to_snowpark_args(snowpark_config: dict[str, Any]) -> dict[str, Any]:
+    renamed_args = {
+        "inferSchema": "INFER_SCHEMA",
+        # TODO: quote in Spark and FIELD_OPTIONALLY_ENCLOSED_BY in Snowflake are not actually the same.
+        "quote": "FIELD_OPTIONALLY_ENCLOSED_BY",
+        "nullValue": "NULL_IF",
+        "dateFormat": "DATE_FORMAT",
+        "timestampFormat": "TIMESTAMP_FORMAT",
+        "lineSep": "RECORD_DELIMITER",
+        "sep": "FIELD_DELIMITER",
+        "header": "PARSE_HEADER",
+        "pathGlobFilter": "PATTERN",
+        "multiLine": "MULTI_LINE",
+    }
+    renamed_args = lowercase_dict_keys(renamed_args)
+
+    # spark does not escape unenclosed fields
+    snowpark_config["ESCAPE_UNENCLOSED_FIELD"] = "NONE"
+    snowpark_config["ERROR_ON_COLUMN_COUNT_MISMATCH"] = False
+    # snowpark_config["EMPTY_FIELD_AS_NULL"] = True
+
+    # Fix the escape character if it is provided.
+    # TODO SNOW-2081726: This seems to be a Snowpark bug
+    if snowpark_config["escape"] and snowpark_config["escape"] == "\\":
+        snowpark_config["escape"] = "\\\\"
+
+    # Rename the keys to match the Snowpark configuration.
+    for spark_arg, snowpark_arg in renamed_args.items():
+        if spark_arg not in snowpark_config:
+            continue
+        snowpark_config[snowpark_arg] = snowpark_config[spark_arg]
+        del snowpark_config[spark_arg]
+
+    return snowpark_config
+
+
+class CsvReaderConfig(ReaderWriterConfig):
+    # spark reader options that snowpark is able to handle.
+    # Spark options are here: https://spark.apache.org/docs/latest/sql-data-sources-csv.html
+    # Snowpark options are here: https://docs.snowflake.com/en/sql-reference/sql/create-file-format
+    def __init__(self, options: dict[str, str]) -> None:
+        super().__init__(
+            _Config(
+                default_config=CSV_READ_DEFAULT_CONFIG,
+                supported_options=CSV_READ_SUPPORTED_OPTIONS,
+                boolean_config_list=[
+                    "header",
+                    "inferSchema",
+                    "multiLine",
+                    "ignoreLeadingWhiteSpace",
+                    "ignoreTrailingWhiteSpace",
+                    "escapeQuotes",
+                    "quoteAll",
+                ],
+                int_config_list=[
+                    "maxColumns",
+                    "maxCharsPerColumn",
+                    "maxMalformedLogPerPartition",
+                ],
+                float_config_list=["samplingRatio"],
+            ),
+            options,
+        )
+
+    def convert_to_snowpark_args(self) -> dict[str, Any]:
+        snowpark_config = super().convert_to_snowpark_args()
+        return csv_convert_to_snowpark_args(snowpark_config)
+
+
+# TODO: This is just a first pass, we need to differentiate more clearly between read and write configs.
+class CsvWriterConfig(ReaderWriterConfig):
+    # spark reader options that snowpark is able to handle.
+    # Spark options are here: https://spark.apache.org/docs/latest/sql-data-sources-csv.html
+    # Snowpark options are here: https://docs.snowflake.com/en/sql-reference/sql/create-file-format
+
+    def __init__(self, options: dict[str, str]) -> None:
+        super().__init__(
+            _Config(
+                default_config=dict(
+                    {
+                        key: value
+                        for key, value in CSV_READ_DEFAULT_CONFIG.items()
+                        # Quote is removed here because it behaves very differently in Snowpark compared to Spark.
+                        if key not in ["quote", "inferSchema", "compression"]
+                    },
+                    **(
+                        {
+                            "compression": "none"  # When writing files compression should be provided by the user
+                        }
+                    ),
+                ),
+                supported_options=CSV_READ_SUPPORTED_OPTIONS - {"inferSchema"},
+                boolean_config_list=[
+                    "header",
+                    "multiLine",
+                    "ignoreLeadingWhiteSpace",
+                    "ignoreTrailingWhiteSpace",
+                    "escapeQuotes",
+                    "quoteAll",
+                ],
+                int_config_list=[
+                    "maxColumns",
+                    "maxCharsPerColumn",
+                    "maxMalformedLogPerPartition",
+                ],
+                float_config_list=["samplingRatio"],
+            ),
+            options,
+        )
+
+    def convert_to_snowpark_args(self) -> dict[str, Any]:
+        snowpark_config = super().convert_to_snowpark_args()
+        return csv_convert_to_snowpark_args(snowpark_config)
+
+
+class JsonReaderConfig(ReaderWriterConfig):
+    def __init__(self, options: dict[str, str]) -> None:
+        super().__init__(
+            _Config(
+                default_config={
+                    # TODO: primitivesAsString: Union[bool, str, None] = None,
+                    # TODO: prefersDecimal: Union[bool, str, None] = None,
+                    # TODO: allowComments: Union[bool, str, None] = None,
+                    # TODO: allowUnquotedFieldNames: Union[bool, str, None] = None,
+                    # TODO: allowSingleQuotes: Union[bool, str, None] = None,
+                    # TODO: allowNumericLeadingZero: Union[bool, str, None] = None,
+                    # TODO: allowBackslashEscapingAnyCharacter: Union[bool, str, None] = None,
+                    # TODO: columnNameOfCorruptRecord: Optional[str] = None,
+                    "dateFormat": "auto",
+                    "timestampFormat": "auto",
+                    # TODO: multiLine: Union[bool, str, None] = None,
+                    # TODO: allowUnquotedControlChars: Union[bool, str, None] = None,
+                    # TODO: lineSep: Optional[str] = None,
+                    # TODO: samplingRatio: Union[str, float, None] = None,
+                    # TODO: dropFieldIfAllNull: Union[bool, str, None] = None,
+                    # TODO: encoding: Optional[str] = None,
+                    # TODO: pathGlobFilter: Union[bool, str, None] = None,
+                    # TODO: recursiveFileLookup: Union[bool, str, None] = None,
+                    # TODO: modifiedBefore: Union[bool, str, None] = None,
+                    # TODO: modifiedAfter: Union[bool, str, None] = None,
+                    # TODO: allowNonNumericNumbers: Union[bool, str, None] = None,
+                    "rowsToInferSchema": 1000,
+                },
+                supported_options={
+                    "schema",
+                    # "primitivesAsString",
+                    # "prefersDecimal",
+                    # "allowComments",
+                    # "allowUnquotedFieldNames",
+                    # "allowSingleQuotes",
+                    # "allowNumericLeadingZero",
+                    # "allowBackslashEscapingAnyCharacter",
+                    # "mode",
+                    # "columnNameOfCorruptRecord",
+                    "dateFormat",
+                    "timestampFormat",
+                    "multiLine",
+                    # "allowUnquotedControlChars",
+                    # "lineSep",
+                    # "samplingRatio",
+                    "dropFieldIfAllNull",
+                    "encoding",
+                    # "locale",
+                    # "pathGlobFilter",
+                    # "recursiveFileLookup",
+                    # "modifiedBefore",
+                    # "modifiedAfter",
+                    # "allowNonNumericNumbers",
+                    "compression",
+                    # "ignoreNullFields",
+                    "rowsToInferSchema",
+                },
+                boolean_config_list=[
+                    "multiLine",
+                    "dropFieldIfAllNull",
+                ],
+                int_config_list=["rowsToInferSchema"],
+                float_config_list=["samplingRatio"],
+            ),
+            options,
+        )
+
+    def convert_to_snowpark_args(self) -> dict[str, Any]:
+        renamed_args = {
+            "inferSchema": "INFER_SCHEMA",
+            "dateFormat": "DATE_FORMAT",
+            "timestampFormat": "TIMESTAMP_FORMAT",
+            "multiLine": "STRIP_OUTER_ARRAY",
+        }
+        renamed_args = lowercase_dict_keys(renamed_args)
+        snowpark_config = super().convert_to_snowpark_args()
+        # Rename the keys to match the Snowpark configuration.
+        for spark_arg, snowpark_arg in renamed_args.items():
+            if spark_arg not in snowpark_config:
+                continue
+            snowpark_config[snowpark_arg] = snowpark_config[spark_arg]
+            del snowpark_config[spark_arg]
+        return snowpark_config
+
+
+class ParquetReaderConfig(ReaderWriterConfig):
+    def __init__(self, options: dict[str, str]) -> None:
+        super().__init__(
+            _Config(
+                default_config={},
+                supported_options={
+                    # "mergeSchema",
+                    # "pathGlobFilter",
+                    # "recursiveFileLookup",
+                    # "modifiedBefore",
+                    # "modifiedAfter",
+                    # "datetimeRebaseMode",
+                    # "int96RebaseMode",
+                    # "mode",
+                    "compression",
+                },
+                boolean_config_list=[],
+                int_config_list=[],
+                float_config_list=[],
+            ),
+            options,
+        )

snowflake/snowpark_connect/relation/read/utils.py

@@ -0,0 +1,155 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import functools
+import re
+import time
+from collections.abc import Callable
+from typing import (  # noqa: F401
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Generator,
+    Iterator,
+    List,
+    NewType,
+    Optional,
+    Protocol,
+    Tuple,
+    Type,
+    Union,
+    get_args,
+    get_origin,
+)
+
+from snowflake import snowpark
+from snowflake.snowpark._internal.analyzer import analyzer_utils
+from snowflake.snowpark.exceptions import SnowparkClientException
+from snowflake.snowpark_connect.column_name_handler import (
+    make_column_names_snowpark_compatible,
+)
+from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
+
+STATEMENT_PARAMS_DATA_SOURCE = "SNOWPARK_PYTHON_DATASOURCE"
+DATA_SOURCE_DBAPI_SIGNATURE = "DataFrameReader.dbapi"
+DATA_SOURCE_SQL_COMMENT = (
+    f"/* Python:snowflake.snowpark.{DATA_SOURCE_DBAPI_SIGNATURE} */"
+)
+
+INDEXED_COLUMN_NAME_PATTERN = re.compile(r"(^\"c)(\d+)(\"$)")
+
+
+def subtract_one(match: re.Match[str]) -> str:
+    """Spark column names are 0 indexed, Snowpark is 1 indexed."""
+    return f"_c{str(int(match.group(2)) - 1)}"
+
+
+def get_spark_column_names_from_snowpark_columns(
+    snowpark_column_names: List[str],
+) -> List[str]:
+    return [
+        analyzer_utils.unquote_if_quoted(
+            INDEXED_COLUMN_NAME_PATTERN.sub(subtract_one, c)
+        )
+        for c in snowpark_column_names
+    ]
+
+
+def rename_columns_as_snowflake_standard(
+    df: snowpark.DataFrame, plan_id: int
+) -> tuple[snowpark.DataFrame, list[str]]:
+    """
+    Renames the columns of a Snowflake DataFrame to follow a standard format.
+    Args:
+        df (snowpark.DataFrame): The input Snowflake DataFrame.
+
+    Returns:
+        tuple[snowpark.DataFrame, list[str]]: A tuple containing the modified DataFrame
+        with renamed columns and a list of the new column names.
+    """
+
+    if df.columns is None or len(df.columns) == 0:
+        return df, []
+
+    new_columns = make_column_names_snowpark_compatible(df.columns, plan_id)
+    return (
+        df.select(
+            *(df.col(orig).alias(alias) for orig, alias in zip(df.columns, new_columns))
+        ),
+        new_columns,
+    )
+
+
+class Connection(Protocol):
+    """External datasource connection created from user-input create_connection function."""
+
+    def cursor(self) -> "Cursor":
+        pass
+
+    def close(self):
+        pass
+
+    def commit(self):
+        pass
+
+    def rollback(self):
+        pass
+
+
+class Cursor(Protocol):
+    """Cursor created from external datasource connection"""
+
+    def execute(self, sql: str, *params: Any) -> "Cursor":
+        pass
+
+    def fetchall(self):
+        pass
+
+    def fetchone(self):
+        pass
+
+    def close(self):
+        pass
+
+
+def exponential_backoff(
+    func: Callable | None = None,
+    max_retry_count: int = 3,
+    initial_retry_delay_ms: int = 50,
+    exponential_backoff_base: int = 2,
+) -> Callable:
+    if func is None:
+        return functools.partial(
+            exponential_backoff,
+            max_retry_count=max_retry_count,
+            initial_retry_delay_ms=initial_retry_delay_ms,
+            exponential_backoff_base=exponential_backoff_base,
+        )
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        error = None
+        for retry_count in range(max_retry_count):
+            try:
+                return func(*args, **kwargs)
+            except Exception as e:
+                error = e
+
+                delay_ms = (
+                    exponential_backoff_base**retry_count * initial_retry_delay_ms
+                )
+                time.sleep(delay_ms / 1000.0)
+                logger.debug(
+                    f"Function '{func.__name__}' failed with {error.__repr__()}, retry count: {retry_count}, retrying ..."
+                )
+        error = SnowparkClientException(
+            message=f"failed to run '{func.__name__}', got {error.__repr__()}"
+        )
+        logger.debug(
+            f"Function '{func.__name__}' failed with {error.__repr__()}, exceed max retry time"
+        )
+
+        return error
+
+    return wrapper