snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/relation/read/map_read_json.py

@@ -0,0 +1,344 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import copy
+import json
+import typing
+from contextlib import suppress
+
+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.row import Row
+from snowflake.snowpark.types import (
+    ArrayType,
+    DataType,
+    DateType,
+    MapType,
+    NullType,
+    StringType,
+    StructField,
+    StructType,
+)
+from snowflake.snowpark_connect.column_name_handler import with_column_map
+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,
+    rename_columns_as_snowflake_standard,
+)
+from snowflake.snowpark_connect.type_mapping import (
+    cast_to_match_snowpark_type,
+    map_simple_types,
+)
+from snowflake.snowpark_connect.utils.telemetry import (
+    SnowparkConnectNotImplementedError,
+)
+
+
+def map_read_json(
+    rel: relation_proto.Relation,
+    schema: StructType | None,
+    session: snowpark.Session,
+    paths: list[str],
+    options: JsonReaderConfig,
+) -> snowpark.DataFrame:
+    """
+    Read a JSON file into a Snowpark DataFrame.
+
+    [JSON lines](http://jsonlines.org/) file format is supported.
+
+    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 JSON files."
+        )
+    else:
+        snowpark_options = options.convert_to_snowpark_args()
+        snowpark_options["infer_schema"] = True
+
+        rows_to_infer_schema = snowpark_options.pop("rowstoinferschema", 1000)
+        dropFieldIfAllNull = snowpark_options.pop("dropfieldifallnull", False)
+
+        reader = session.read.options(snowpark_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))
+
+        if schema is None:
+            schema = copy.deepcopy(df.schema)
+            infer_row_counts = 0
+
+            columns_with_valid_contents = set()
+            for row in df.to_local_iterator():
+                infer_row_counts += 1
+                if (
+                    rows_to_infer_schema != -1
+                    and infer_row_counts > rows_to_infer_schema
+                ):
+                    break
+                schema = merge_row_schema(
+                    schema, row, columns_with_valid_contents, dropFieldIfAllNull
+                )
+
+            if dropFieldIfAllNull:
+                schema.fields = [
+                    sf
+                    for sf in schema.fields
+                    if unquote_if_quoted(sf.name) in columns_with_valid_contents
+                ]
+
+        df = construct_dataframe_by_schema(
+            schema, df.to_local_iterator(), session, snowpark_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],
+        )
+
+
+def merge_json_schema(
+    content: typing.Any,
+    schema: StructType | None,
+    dropFieldIfAllNull: bool = False,
+) -> DataType:
+    if content is None:
+        if schema is not None:
+            return schema
+        return NullType()
+
+    if isinstance(content, str):
+        with suppress(json.JSONDecodeError):
+            json_content = json.loads(content)
+            if not isinstance(json_content, str):
+                content = json_content
+
+    if isinstance(content, dict):
+        current_schema = StructType()
+
+        existed_schema = {}
+        if schema is not None and schema.type_name() == "struct":
+            for sf in schema.fields:
+                existed_schema[sf.name] = sf.datatype
+
+        for k, v in content.items():
+            col_name = f'"{unquote_if_quoted(k)}"'
+            existed_data_type = existed_schema.get(col_name, None)
+            next_level_schema = merge_json_schema(
+                v, existed_data_type, dropFieldIfAllNull
+            )
+
+            if (
+                existed_data_type is not None
+                or not dropFieldIfAllNull
+                or not isinstance(next_level_schema, NullType)
+            ):
+                # Drop field if it's always null
+                current_schema.add(StructField(col_name, next_level_schema))
+            if existed_data_type is not None:
+                del existed_schema[col_name]
+
+        for k, v in existed_schema.items():
+            col_name = f'"{unquote_if_quoted(k)}"'
+            current_schema.add(StructField(col_name, v))
+
+    elif isinstance(content, list):
+        # ArrayType(*) need to have element schema inside, it would be NullType() as placeholder and keep updating while enumerating
+        inner_schema = NullType()
+        if schema is not None and schema.type_name() == "list":
+            inner_schema = schema.element_type
+        if len(content) > 0:
+            for v in content:
+                inner_schema = merge_json_schema(v, inner_schema, dropFieldIfAllNull)
+        if isinstance(inner_schema, NullType) and dropFieldIfAllNull:
+            return NullType()
+        current_schema = ArrayType(inner_schema)
+    else:
+        current_schema = map_simple_types(type(content).__name__)
+
+    # If there's conflict , use StringType
+    if (
+        schema is not None
+        and schema != NullType()
+        and current_schema is not None
+        and current_schema != NullType()
+        and schema.type_name() != current_schema.type_name()
+    ):
+        return StringType()
+
+    current_schema.structured = True
+    return current_schema
+
+
+def merge_row_schema(
+    schema: StructType | None,
+    row: Row,
+    columns_with_valid_contents: set,
+    dropFieldIfAllNull: bool = False,
+) -> StructType | NullType:
+    if row is None:
+        if schema is not None:
+            return schema
+        return NullType()
+    new_schema = StructType()
+
+    for sf in schema.fields:
+        col_name = unquote_if_quoted(sf.name)
+        if isinstance(sf.datatype, (StructType, MapType, StringType)):
+            next_level_content = row[col_name]
+            if next_level_content is not None:
+                with suppress(json.JSONDecodeError):
+                    next_level_content = json.loads(next_level_content)
+                if isinstance(next_level_content, dict):
+                    sf.datatype = merge_json_schema(
+                        next_level_content,
+                        None
+                        if not isinstance(sf.datatype, StructType)
+                        else sf.datatype,
+                        dropFieldIfAllNull,
+                    )
+                else:
+                    sf.datatype = StringType()
+                columns_with_valid_contents.add(col_name)
+
+        elif isinstance(sf.datatype, ArrayType):
+            content = row[col_name]
+            if content is not None:
+                with suppress(Exception):
+                    decoded_content = json.loads(content)
+                    if isinstance(decoded_content, list):
+                        content = decoded_content
+                if not isinstance(content, list):
+                    sf.datatype = StringType()
+                else:
+                    for v in content:
+                        if v is not None:
+                            columns_with_valid_contents.add(col_name)
+                        sf.datatype.element_type = merge_json_schema(
+                            v,
+                            sf.datatype.element_type,
+                            dropFieldIfAllNull,
+                        )
+        elif row[col_name] is not None:
+            columns_with_valid_contents.add(col_name)
+
+        sf.datatype.structured = True
+        new_schema.add(sf)
+
+    return schema
+
+
+def union_data_into_df(
+    result_df: snowpark.DataFrame,
+    data: typing.List[Row],
+    schema: StructType,
+    session: snowpark.Session,
+) -> snowpark.DataFrame:
+    current_df = session.create_dataframe(
+        data=data,
+        schema=schema,
+    )
+    if result_df is None:
+        return current_df
+
+    return result_df.union(current_df)
+
+
+def construct_dataframe_by_schema(
+    schema: StructType,
+    rows: typing.Iterator[Row],
+    session: snowpark.Session,
+    snowpark_options: dict,
+    batch_size: int = 100,
+) -> snowpark.DataFrame:
+    result = None
+
+    current_data = []
+    for row in rows:
+        current_data.append(construct_row_by_schema(row, schema, snowpark_options))
+        if len(current_data) >= batch_size:
+            result = union_data_into_df(
+                result,
+                current_data,
+                schema,
+                session,
+            )
+
+    if len(current_data) > 0:
+        result = union_data_into_df(
+            result,
+            current_data,
+            schema,
+            session,
+        )
+
+    if result is None:
+        raise ValueError("Dataframe cannot be empty")
+    return result
+
+
+def construct_row_by_schema(
+    content: typing.Any, schema: DataType, snowpark_options: dict
+) -> None | DataType:
+    if content is None:
+        return None
+    elif isinstance(schema, StructType):
+        result = {}
+        if isinstance(content, (dict, Row)):
+            for sf in schema.fields:
+                col_name = unquote_if_quoted(sf.name)
+                quoted_col_name = (
+                    f'"{col_name}"' if isinstance(content, Row) else col_name
+                )
+                result[quoted_col_name] = construct_row_by_schema(
+                    (content.as_dict() if isinstance(content, Row) else content).get(
+                        col_name, None
+                    ),
+                    sf.datatype,
+                    snowpark_options,
+                )
+        elif isinstance(content, str):
+            with suppress(json.JSONDecodeError):
+                decoded_content = json.loads(content)
+                if isinstance(decoded_content, dict):
+                    content = decoded_content
+            for sf in schema.fields:
+                col_name = unquote_if_quoted(sf.name)
+                result[col_name] = construct_row_by_schema(
+                    content.get(col_name, None), sf.datatype, snowpark_options
+                )
+        else:
+            raise SnowparkConnectNotImplementedError(
+                f"JSON construct {str(content)} to StructType failed"
+            )
+        return result
+    elif isinstance(schema, ArrayType):
+        result = []
+        inner_schema = schema.element_type
+        if isinstance(content, str):
+            content = json.loads(content)
+        for ele in content:
+            result.append(construct_row_by_schema(ele, inner_schema, snowpark_options))
+        return result
+    elif isinstance(schema, DateType):
+        return cast_to_match_snowpark_type(
+            schema, content, snowpark_options.get("DATE_FORMAT")
+        )
+
+    return cast_to_match_snowpark_type(schema, content)

snowflake/snowpark_connect/relation/read/map_read_parquet.py

@@ -0,0 +1,194 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import collections
+import re
+from collections.abc import Callable
+
+import pyspark.sql.connect.proto.relations_pb2 as relation_proto
+
+from snowflake import snowpark
+from snowflake.snowpark import (
+    DataFrame,
+    DataFrameReader,
+    Session,
+    functions as snowpark_fn,
+)
+from snowflake.snowpark._internal.analyzer import analyzer_utils
+from snowflake.snowpark._internal.analyzer.analyzer_utils import (
+    quote_name_without_upper_casing,
+)
+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.relation.read.reader_config import ReaderWriterConfig
+from snowflake.snowpark_connect.relation.read.utils import (
+    rename_columns_as_snowflake_standard,
+)
+from snowflake.snowpark_connect.utils.telemetry import (
+    SnowparkConnectNotImplementedError,
+)
+
+
+def map_read_parquet(
+    rel: relation_proto.Relation,
+    schema: str | None,
+    session: snowpark.Session,
+    paths: list[str],
+    options: ReaderWriterConfig,
+) -> snowpark.DataFrame:
+    """Read a Parquet file into a Snowpark DataFrame."""
+
+    if rel.read.is_streaming is True:
+        raise SnowparkConnectNotImplementedError(
+            "Streaming is not supported for Parquet files."
+        )
+
+    snowpark_options = options.convert_to_snowpark_args()
+    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)
+
+    if len(paths) == 1:
+        df = _read_parquet_with_partitions(session, reader, paths[0])
+    else:
+        df = _read_parquet_with_partitions(session, reader, paths[0])
+        for p in paths[1:]:
+            reader._user_schema = None
+            df = df.union_all(_read_parquet_with_partitions(session, reader, p))
+
+    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],
+    )
+
+
+def _read_parquet_with_partitions(
+    session: Session, reader: DataFrameReader, path: str
+) -> DataFrame:
+    """Reads parquet files and adds partition columns from subdirectories."""
+
+    partition_columns, inferred_types = _discover_partition_columns(session, path)
+    df = reader.with_metadata(METADATA_FILENAME).parquet(path)
+
+    if not partition_columns:
+        return df.drop(METADATA_FILENAME)
+
+    for col_name in partition_columns:
+        quoted_col_name = quote_name_without_upper_casing(col_name)
+        escaped_col_name = re.escape(col_name)
+        regex_pattern = rf"{escaped_col_name}=([^/]+)"
+
+        raw_value = snowpark_fn.regexp_extract(METADATA_FILENAME, regex_pattern, 1)
+        value_or_null = snowpark_fn.when(raw_value == "", None).otherwise(raw_value)
+
+        df = df.with_column(
+            quoted_col_name, snowpark_fn.cast(value_or_null, inferred_types[col_name])
+        )
+
+    return df.drop(METADATA_FILENAME)
+
+
+def _extract_partitions_from_path(path: str) -> dict[str, str]:
+    """Extracts partition key-value pairs from a path."""
+    partitions = {}
+    for segment in path.split("/"):
+        if "=" in segment:
+            col_name, value = _parse_partition_column(segment)
+            if col_name and value:
+                partitions[col_name] = value
+    return partitions
+
+
+def _discover_partition_columns(
+    session: Session, stage_path: str
+) -> tuple[list[str], dict[str, DataType]]:
+    """Discovers partition columns by analyzing subdirectory structure."""
+
+    partition_columns_values = collections.defaultdict(set)
+    dir_level_to_column_name = {}
+    base_partitions = _extract_partitions_from_path(stage_path)
+
+    path_segments_to_skip = len(stage_path.strip("/").split("/"))
+    if stage_path.startswith("@"):
+        path_segments_to_skip = 1
+        stage_parts = stage_path.split("/", 2)
+        if len(stage_parts) > 2:
+            additional_segments = len(stage_parts[2].strip("/").split("/"))
+            path_segments_to_skip += additional_segments
+
+    ls_result = session.sql(f"LS {stage_path}").collect()
+    if ls_result:
+        file_names = [row[0] for row in ls_result]
+
+        for file_path in file_names:
+            path_parts = file_path.strip("/").split("/")
+            path_segments_to_analyze = path_parts[path_segments_to_skip:]
+
+            for i, part in enumerate(path_segments_to_analyze):
+                if "=" in part and not part.endswith(".parquet"):
+                    key, value = part.split("=", 1)
+
+                    if key in base_partitions:
+                        continue
+
+                    if i not in dir_level_to_column_name:
+                        dir_level_to_column_name[i] = key
+                    elif dir_level_to_column_name[i] != key:
+                        raise ValueError(
+                            f"Conflicting partition column names detected: '{dir_level_to_column_name[i]}' and '{key}' "
+                            f"at the same directory level"
+                        )
+
+                    partition_columns_values[key].add(value)
+
+    seen_columns = set()
+    for level in sorted(dir_level_to_column_name.keys()):
+        col_name = dir_level_to_column_name[level]
+        if col_name in seen_columns:
+            raise ValueError(
+                f"Found partition column '{col_name}' at multiple directory levels. "
+                f"A partition column can only appear at a single level."
+            )
+        seen_columns.add(col_name)
+
+    ordered_columns = [
+        dir_level_to_column_name[level]
+        for level in sorted(dir_level_to_column_name.keys())
+    ]
+
+    inferred_types = {
+        col_name: _infer_partition_column_type(partition_columns_values[col_name])
+        for col_name in ordered_columns
+    }
+
+    return ordered_columns, inferred_types
+
+
+def _infer_partition_column_type(values: set[str]) -> DataType:
+    def _is_castable(value: str, type_: Callable) -> bool:
+        try:
+            type_(value)
+            return True
+        except ValueError:
+            return False
+
+    if all(_is_castable(value, int) for value in values):
+        return IntegerType()
+    if all(_is_castable(value, float) for value in values):
+        return DoubleType()
+    return StringType()
+
+
+def _parse_partition_column(name: str) -> tuple[str, str]:
+    """Extracts column name and partition value from a path segment."""
+    col_name, partition_value = name.split("=", maxsplit=1)
+
+    return col_name, partition_value

snowflake/snowpark_connect/relation/read/map_read_socket.py

@@ -0,0 +1,59 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import socket
+
+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.utils.telemetry import (
+    SnowparkConnectNotImplementedError,
+)
+
+# Full data is used at the moment to store all the data that has been read
+# from the socket. This is a temporary solution for demonstration purposes
+# and not scalable.
+# TODO: Use a stage in Snowflake to store the data.
+full_data = b""
+
+
+def map_read_socket(
+    rel: relation_proto.Relation,
+    session: snowpark.Session,
+    options: dict[str, str],
+) -> snowpark.DataFrame:
+    if rel.read.is_streaming is True:
+        global full_data
+        host = options.get("host", None)
+        port = options.get("port", None)
+        if not host or not port:
+            raise ValueError("Host and port must be provided in options.")
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            try:
+                s.connect((host, int(port)))
+                data = s.recv(1024)
+                full_data += data
+                while data:
+                    s.settimeout(1.25)
+                    try:
+                        data = s.recv(1024)
+                        full_data += data
+                    except socket.timeout:
+                        break
+                s.settimeout(None)
+                dataframe_data = full_data.decode("utf-8")
+                snowpark_cname = "VALUE"
+                df = session.create_dataframe(
+                    pandas.DataFrame({snowpark_cname: dataframe_data.split("\n")})
+                )
+                spark_cname = "value"
+                return with_column_map(df, [spark_cname], [snowpark_cname])
+            except OSError as e:
+                raise Exception(f"Error connecting to {host}:{port} - {e}")
+    else:
+        raise SnowparkConnectNotImplementedError(
+            "Socket reads are only supported in streaming mode."
+        )

snowflake/snowpark_connect/relation/read/map_read_table.py

@@ -0,0 +1,109 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import pyspark.sql.connect.proto.relations_pb2 as relation_proto
+from pyspark.errors.exceptions.base import AnalysisException
+
+from snowflake import snowpark
+from snowflake.snowpark._internal.analyzer.analyzer_utils import (
+    quote_name_without_upper_casing,
+    unquote_if_quoted,
+)
+from snowflake.snowpark.exceptions import SnowparkSQLException
+from snowflake.snowpark_connect.column_name_handler import with_column_map
+from snowflake.snowpark_connect.config import auto_uppercase_dml
+from snowflake.snowpark_connect.relation.read.utils import (
+    rename_columns_as_snowflake_standard,
+)
+from snowflake.snowpark_connect.utils.attribute_handling import (
+    split_fully_qualified_spark_name,
+)
+from snowflake.snowpark_connect.utils.session import _get_current_snowpark_session
+from snowflake.snowpark_connect.utils.telemetry import (
+    SnowparkConnectNotImplementedError,
+)
+
+
+def post_process_df(
+    df: snowpark.DataFrame, plan_id: int, source_table_name: str = None
+) -> snowpark.DataFrame:
+    try:
+        true_names = list(map(lambda x: unquote_if_quoted(x).lower(), df.columns))
+        renamed_df, snowpark_column_names = rename_columns_as_snowflake_standard(
+            df, plan_id
+        )
+        name_parts = split_fully_qualified_spark_name(source_table_name)
+
+        # If table name is not fully qualified (only has table name, no database),
+        # add current schema name to qualifiers so columns can be referenced with database prefix
+        # Note: In Spark, "database" corresponds to Snowflake "schema"
+        if source_table_name and len(name_parts) == 1:
+            session = _get_current_snowpark_session()
+            current_schema = session.get_current_schema()
+            if current_schema:
+                name_parts = [unquote_if_quoted(current_schema)] + name_parts
+
+        return with_column_map(
+            renamed_df,
+            true_names,
+            snowpark_column_names,
+            [f.datatype for f in df.schema.fields],
+            column_qualifiers=[name_parts] * len(true_names)
+            if source_table_name
+            else None,
+        )
+    except SnowparkSQLException as e:
+        # Check if this is a table/view not found error
+        # Snowflake error codes: 002003 (42S02) - Object does not exist or not authorized
+        if hasattr(e, "sql_error_code") and e.sql_error_code == 2003:
+            raise AnalysisException(
+                f"[TABLE_OR_VIEW_NOT_FOUND] The table or view cannot be found. {source_table_name}"
+            ) from None  # Suppress original exception to reduce message size
+        # Re-raise if it's not a table not found error
+        raise
+
+
+def get_table_from_name(
+    table_name: str, session: snowpark.Session, plan_id: int
+) -> snowpark.DataFrame:
+    snowpark_name = ".".join(
+        quote_name_without_upper_casing(part)
+        for part in split_fully_qualified_spark_name(table_name)
+    )
+
+    if auto_uppercase_dml():
+        snowpark_name = snowpark_name.upper()
+
+    df = session.read.table(snowpark_name)
+    post_processed_df = post_process_df(df, plan_id, table_name)
+    post_processed_df._table_name = table_name
+    return post_processed_df
+
+
+def get_table_from_query(
+    query: str, session: snowpark.Session, plan_id: int
+) -> snowpark.DataFrame:
+    df = session.sql(query)
+    return post_process_df(df, plan_id)
+
+
+def map_read_table(rel: relation_proto.Relation) -> snowpark.DataFrame:
+    """
+    Read a table into a Snowpark DataFrame.
+    """
+    session: snowpark.Session = _get_current_snowpark_session()
+    if rel.read.HasField("named_table"):
+        table_identifier = rel.read.named_table.unparsed_identifier
+    elif (
+        rel.read.data_source.HasField("format")
+        and rel.read.data_source.format.lower() == "iceberg"
+    ):
+        if len(rel.read.data_source.paths) != 1:
+            raise SnowparkConnectNotImplementedError(
+                f"Unexpected paths: {rel.read.data_source.paths}"
+            )
+        table_identifier = rel.read.data_source.paths[0]
+    else:
+        raise ValueError("The relation must have a table identifier.")
+    return get_table_from_name(table_identifier, session, rel.common.plan_id)