snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/relation/read/map_read.py

@@ -0,0 +1,367 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import concurrent.futures
+import json
+import logging
+import os
+import re
+
+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.relation.io_utils import (
+    convert_file_prefix_path,
+    is_cloud_path,
+)
+from snowflake.snowpark_connect.relation.read.map_read_table import map_read_table
+from snowflake.snowpark_connect.relation.read.reader_config import (
+    CsvReaderConfig,
+    JsonReaderConfig,
+    ParquetReaderConfig,
+)
+from snowflake.snowpark_connect.relation.stage_locator import get_paths_from_stage
+from snowflake.snowpark_connect.type_mapping import map_json_schema_to_snowpark
+from snowflake.snowpark_connect.utils.cache import df_cache_map_put_if_absent
+from snowflake.snowpark_connect.utils.context import get_session_id
+from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
+from snowflake.snowpark_connect.utils.telemetry import (
+    SnowparkConnectNotImplementedError,
+    telemetry,
+)
+
+logger = logging.getLogger("snowflake_connect_server")
+
+
+def map_read(
+    rel: relation_proto.Relation,
+) -> snowpark.DataFrame:
+    """
+    Read a file into a Snowpark DataFrame.
+
+    Currently, the supported read formats are `csv`, `json` and `parquet`.
+    """
+    match rel.read.WhichOneof("read_type"):
+        case "named_table":
+            return map_read_table_or_file(rel)
+
+        case "data_source":
+            read_format = None
+            if rel.read.data_source.HasField("format"):
+                read_format = rel.read.data_source.format
+            else:
+                read_format = global_config.get("spark.sql.sources.default")
+                if read_format is not None:
+                    read_format = read_format.split(".")[-1]
+
+            if not read_format:
+                # TODO: This should come from the config `spark.sql.sources.default`
+                # The default format is parquet, but users can override it.
+                read_format = "parquet"
+
+            if read_format.lower() == "iceberg":
+                telemetry.report_io_read("iceberg", None)
+                return map_read_table(rel)
+
+            if rel.read.data_source.schema == "":
+                schema = None
+            else:
+                try:
+                    parsed_schema = json.loads(rel.read.data_source.schema)
+                except json.JSONDecodeError:
+                    # It's a DDL-formatted string. e.g. "a INT, b DOUBLE"
+                    parsed_schema = {"fields": [], "type": "struct"}
+                    for field in rel.read.data_source.schema.split(","):
+                        name_and_type = field.split()
+                        assert (
+                            len(name_and_type) == 2
+                        ), f"Schema's definition {name_and_type} is invalid"
+                        parsed_schema["fields"].append(
+                            {
+                                "name": name_and_type[0],
+                                "nullable": True,
+                                "type": name_and_type[1],
+                            }
+                        )
+                schema = map_json_schema_to_snowpark(parsed_schema)
+            options = dict(rel.read.data_source.options)
+            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
+                clean_source_paths = [
+                    path.rstrip("/") for path in rel.read.data_source.paths
+                ]
+
+                result = _read_file(
+                    clean_source_paths, options, read_format, rel, schema, session
+                )
+            else:
+                match read_format:
+                    case "socket":
+                        from snowflake.snowpark_connect.relation.read.map_read_socket import (
+                            map_read_socket,
+                        )
+
+                        return map_read_socket(rel, session, options)
+
+                    case "jdbc":
+                        from snowflake.snowpark_connect.relation.read.map_read_jdbc import (
+                            map_read_jdbc,
+                        )
+
+                        return map_read_jdbc(rel, session, options)
+                    case "net.snowflake.spark.snowflake":
+                        options = {k.lower(): v for k, v in options.items()}
+                        QUERY_OPTION = "query"
+                        DBTABLE_OPTION = "dbtable"
+                        logger.warning(
+                            "Changing the Role, Warehouse, Database, or Schema via "
+                            "options is supported, but these will change for the "
+                            "entire session."
+                        )
+                        session.use_role(
+                            options.get("sfRole", session.get_current_role())
+                        )
+                        session.use_warehouse(
+                            options.get("sfWarehouse", session.get_current_warehouse())
+                        )
+                        session.use_database(
+                            options.get("sfDatabase", session.get_current_database())
+                        )
+                        session.use_schema(
+                            options.get("sfSchema", session.get_current_schema())
+                        )
+                        if QUERY_OPTION in options.keys():
+                            from .map_read_table import get_table_from_query
+
+                            return get_table_from_query(
+                                options[QUERY_OPTION], session, rel.common.plan_id
+                            )
+                        elif DBTABLE_OPTION in options.keys():
+                            from .map_read_table import get_table_from_name
+
+                            return get_table_from_name(
+                                options[DBTABLE_OPTION], session, rel.common.plan_id
+                            )
+                    case other:
+                        raise SnowparkConnectNotImplementedError(
+                            f"UNSUPPORTED FORMAT {other} WITH NO PATH"
+                        )
+        case other:
+            # TODO: Empty data source
+            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
+    )
+
+
+def map_read_table_or_file(rel):
+    read_named_table_from_file = (
+        rel.read.named_table.unparsed_identifier
+        and _get_supported_read_file_format(rel.read.named_table.unparsed_identifier)
+    )
+    if read_named_table_from_file:
+        # This case handles when user reads the file using the raw SQL
+        schema = None
+        read_format = _get_supported_read_file_format(
+            rel.read.named_table.unparsed_identifier
+        )
+        options = {}
+        telemetry.report_io_read(read_format, options)
+        session: snowpark.Session = get_or_create_snowpark_session()
+
+        clean_source_paths = [
+            re.sub(
+                rf"^{read_format}\.`([^`]+)`$",
+                r"\1",
+                rel.read.named_table.unparsed_identifier,
+            ).rstrip("/")
+        ]
+
+        return _read_file(
+            clean_source_paths, options, read_format, rel, schema, session
+        )
+    else:
+        return map_read_table(rel)
+
+
+def _get_supported_read_file_format(unparsed_identifier: str) -> str | None:
+    if unparsed_identifier.startswith("csv.`"):
+        return "csv"
+    elif unparsed_identifier.startswith("json.`"):
+        return "json"
+    elif unparsed_identifier.startswith("parquet.`"):
+        return "parquet"
+    elif unparsed_identifier.startswith("text.`"):
+        return "text"
+    return None
+
+
+def _read_file(
+    clean_source_paths: list[str],
+    options: dict,
+    read_format: str,
+    rel: relation_proto.Relation,
+    schema: StructType | None,
+    session: snowpark.Session,
+) -> snowpark.DataFrame:
+    paths = get_paths_from_stage(
+        clean_source_paths,
+        session,
+    )
+    upload_files_if_needed(paths, clean_source_paths, session, read_format)
+    match read_format:
+        case "csv":
+            from snowflake.snowpark_connect.relation.read.map_read_csv import (
+                map_read_csv,
+            )
+
+            return map_read_csv(rel, schema, session, paths, CsvReaderConfig(options))
+        case "json":
+            from snowflake.snowpark_connect.relation.read.map_read_json import (
+                map_read_json,
+            )
+
+            return map_read_json(rel, schema, session, paths, JsonReaderConfig(options))
+        case "parquet":
+            from snowflake.snowpark_connect.relation.read.map_read_parquet import (
+                map_read_parquet,
+            )
+
+            return map_read_parquet(
+                rel, schema, session, paths, ParquetReaderConfig(options)
+            )
+        case "text":
+            from snowflake.snowpark_connect.relation.read.map_read_text import (
+                map_read_text,
+            )
+
+            return map_read_text(rel, schema, session, paths)
+        case _:
+            raise SnowparkConnectNotImplementedError(
+                f"Unsupported format: {read_format}"
+            )
+
+
+def _skip_upload(path: str, read_format: str):
+    """
+    Determines whether to skip the upload of a file based on its format.
+    :param path: The path to the file.
+    :param read_format: The format for reading. Parquet formatting implies additional file skipping logic.
+    :return: True if the upload should be skipped, False otherwise.
+    """
+    if read_format == "parquet":
+        # Skip uploading files that are not parquet
+        return not path.endswith(".parquet")
+    return False
+
+
+def upload_files_if_needed(
+    stage_target_paths: list[str],
+    source_paths: list[str],
+    session: snowpark.Session,
+    read_format: str,
+) -> None:
+    """
+    Uploads file to stage if needed, preserving the underlying directory structure.
+    For parquet, the most common issue is a _SUCCESS.gz that causes reading to fail.
+    :param stage_target_paths: The paths to the staged files. They should be equal to the source_paths but with the stage name prefixed.
+    :param source_paths: The paths to the source files.
+    :param session: The Snowpark session.
+    :param read_format: The format for reading. Parquet formatting implies additional file skipping logic.
+    """
+
+    assert len(source_paths) == len(
+        stage_target_paths
+    ), "Source and target paths must have same length"
+
+    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}/"
+        session.sql(remove_command).collect()
+
+        try:
+            # Walk through all subdirectories
+            for root, dirs, files in os.walk(source):
+                if not files:
+                    continue
+
+                rel_path = os.path.relpath(root, source)
+
+                curr_target = target
+                if rel_path != ".":
+                    curr_target = f"{curr_target}/{rel_path}"
+
+                # If there are no directories, and this is not parquet where we need to filter files,
+                # we can use * to upload all files
+                if not dirs and read_format != "parquet":
+                    file_pattern = os.path.join(root, "*")
+                    # Ensure target ends with single slash
+                    pattern_target = f"{curr_target}/"
+                    try:
+                        session.file.put(
+                            file_pattern,
+                            pattern_target,
+                            auto_compress=False,
+                            overwrite=True,
+                        )
+                    except Exception as e:
+                        logger.error(
+                            f"Error uploading files {file_pattern} to {target}: {e}"
+                        )
+                        raise
+                # Otherwise, we need to upload files individually. Uploading with * pattern fails if the pattern
+                # matches a directory.from
+                else:
+                    for file in files:
+                        file_path = os.path.join(root, file)
+                        if _skip_upload(file_path, read_format):
+                            continue
+                        # Avoid double slashes in target path
+                        file_target = f"{curr_target}/{file}"
+                        try:
+                            session.file.put(
+                                file_path,
+                                file_target,
+                                auto_compress=False,
+                                overwrite=True,
+                            )
+                        except Exception as e:
+                            logger.error(
+                                f"Error uploading file {file_path} to {target}: {e}"
+                            )
+                            raise
+        except Exception as e:
+            logger.error(f"Error uploading directory {source} to {target}: {e}")
+            raise
+
+    def _upload_file(target: str, source: str) -> None:
+        session.file.put(source, target, auto_compress=False, overwrite=True)
+
+    with concurrent.futures.ThreadPoolExecutor(
+        max_workers=4, thread_name_prefix="LocalFileUploader_"
+    ) as exc:
+        for source_path, target_path in zip(source_paths, stage_target_paths):
+            if is_cloud_path(source_path):
+                continue
+
+            futures = []
+
+            source_path = convert_file_prefix_path(source_path)
+            # Upload local files
+            # since local files are the source of truth
+            # any existing files should be overwritten in the target prefix
+            if os.path.isdir(source_path):
+                futures.append(exc.submit(_upload_dir, target_path, source_path))
+            else:
+                futures.append(exc.submit(_upload_file, target_path, source_path))
+
+            # Check for exceptions - if we don't do this, they will be lost in the thread.
+            for future in concurrent.futures.as_completed(futures):
+                future.result()

snowflake/snowpark_connect/relation/read/map_read_csv.py

@@ -0,0 +1,142 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import copy
+
+import pyspark.sql.connect.proto.relations_pb2 as relation_proto
+
+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.relation.read.map_read import CsvReaderConfig
+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.telemetry import (
+    SnowparkConnectNotImplementedError,
+)
+
+
+def map_read_csv(
+    rel: relation_proto.Relation,
+    schema: snowpark.types.StructType | None,
+    session: snowpark.Session,
+    paths: list[str],
+    options: CsvReaderConfig,
+) -> snowpark.DataFrame:
+    """
+    Read a CSV file into a Snowpark DataFrame.
+
+    We leverage the stage that is already created in the map_read function that
+    calls this.
+    """
+
+    if rel.read.is_streaming is True:
+        # TODO: Structured streaming implementation.
+        raise SnowparkConnectNotImplementedError(
+            "Streaming is not supported for CSV files."
+        )
+    else:
+        snowpark_options = options.convert_to_snowpark_args()
+        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"
+        ):  # Schema has to equals to header's format
+            reader = session.read.options(snowpark_options)
+        else:
+            reader = session.read.options(snowpark_options).schema(schema)
+        df = read_data(
+            reader,
+            schema,
+            session,
+            paths[0],
+            snowpark_options,
+            raw_options,
+        )
+        if len(paths) > 1:
+            # TODO: figure out if this is what Spark does.
+            for p in paths[1:]:
+                df = df.union_all(reader.csv(p))
+
+        if schema is None:
+            df = df.select(
+                [snowpark_fn.col(c).cast("STRING").alias(c) for c in df.schema.names]
+            )
+
+        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],
+        )
+
+
+def get_header_names(
+    session: snowpark.Session,
+    path: list[str],
+    snowpark_options: dict,
+) -> list[str]:
+    snowpark_options_no_header = copy.copy(snowpark_options)
+    snowpark_options_no_header["PARSE_HEADER"] = False
+
+    header_df = session.read.options(snowpark_options_no_header).csv(path).limit(1)
+    header_data = header_df.collect()[0]
+    return [
+        f'"{header_data[i]}"'
+        for i in range(len(header_df.schema.fields))
+        if header_data[i] is not None
+    ]
+
+
+def read_data(
+    reader: DataFrameReader,
+    schema: snowpark.types.StructType | None,
+    session: snowpark.Session,
+    path: list[str],
+    snowpark_options: dict,
+    raw_options: dict,
+) -> snowpark.DataFrame:
+    df = reader.csv(path)
+    filename = path.strip("/").split("/")[-1]
+    if schema is not None:
+        if len(schema.fields) != len(df.schema.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
+                ):
+                    raise Exception("CSV header does not conform to the schema")
+        return df
+
+    headers = get_header_names(session, path, snowpark_options)
+
+    # 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)
+    ):
+        df = (
+            session.read.options(snowpark_options)
+            .schema(StructType([StructField(h, StringType(), True) for h in headers]))
+            .csv(path)
+        )
+    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_jdbc.py

@@ -0,0 +1,108 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import jaydebeapi
+import jpype
+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.relation.read.jdbc_read_dbapi import JdbcDataFrameReader
+from snowflake.snowpark_connect.relation.read.utils import (
+    Connection,
+    rename_columns_as_snowflake_standard,
+)
+from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
+
+
+def create_connection(jdbc_options: dict[str, str]) -> Connection:
+    url = jdbc_options.get("url", None)
+    driver = jdbc_options.get("driver", None)
+    if driver is None:
+        driver = (
+            jpype.java.sql.DriverManager.getDriver(url).getClass().getCanonicalName()
+        )
+    try:
+        return jaydebeapi.connect(driver, url, jdbc_options)
+    except Exception as e:
+        jpype.detachThreadFromJVM()
+        raise Exception(f"Error connecting JDBC datasource: {e}")
+
+
+def close_connection(conn: Connection) -> None:
+    if conn is not None:
+        conn.close()
+    # JVM main thread and SAS Server main thread creates deadlock.
+    # jaydebeapi has connect method which calls jpype internally.
+    # jpype start JVM if it's not started.
+    # Then jpype attach the current thread to the JVM.
+    # But jaydebeapi doesn't have any method to detach or shutdown the JVM.
+    jpype.detachThreadFromJVM()
+
+
+def map_read_jdbc(
+    rel: relation_proto.Relation,
+    session: snowpark.Session,
+    options: dict[str, str],
+) -> snowpark.DataFrame:
+    """
+    Read a table data or query data from a JDBC external datasource into a Snowpark DataFrame.
+    """
+
+    jdbc_options = options.copy()
+    dbtable = options.get("dbtable", None)
+    query = options.get("query", None)
+    partition_column = options.get("partitionColumn", None)
+    lower_bound = options.get("lowerBound", None)
+    upper_bound = options.get("upperBound", None)
+    num_partitions = options.get("numPartitions", None)
+    predicates = rel.read.data_source.predicates
+
+    logger.info(
+        f"dbtable={dbtable},query={query},partition_column={partition_column},lower_bound={lower_bound},upper_bound={upper_bound},num_partitions={num_partitions}"
+    )
+    if num_partitions is not None:
+        num_partitions = int(num_partitions)
+
+    if len(dbtable) == 0:
+        dbtable = None
+
+    if not dbtable and not query:
+        raise ValueError("Include dbtable or query is required option")
+
+    if query is not None and dbtable is not None:
+        raise ValueError(
+            "Not allowed to specify dbtable and query options at the same time"
+        )
+
+    if query is not None and partition_column is not None:
+        raise ValueError(
+            "Not allowed to specify partitionColumn and query options at the same time"
+        )
+
+    try:
+        df = JdbcDataFrameReader(session, jdbc_options).jdbc_read_dbapi(
+            create_connection,
+            close_connection,
+            table=dbtable,
+            query=query,
+            column=partition_column,
+            lower_bound=lower_bound,
+            upper_bound=upper_bound,
+            num_partitions=num_partitions,
+            predicates=predicates,
+        )
+        true_names = list(map(lambda x: unquote_if_quoted(x).lower(), df.columns))
+        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],
+        )
+    except Exception as e:
+        raise Exception(f"Error accessing JDBC datasource for read: {e}")