snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/execute_plan/map_execution_root.py

@@ -0,0 +1,173 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+from collections import namedtuple
+from typing import Iterator
+
+import pandas
+import pyarrow as pa
+import pyspark.sql.connect.proto.base_pb2 as proto_base
+import pyspark.sql.connect.proto.types_pb2 as proto_types
+from pyarrow import Table
+
+import snowflake.snowpark_connect.tcm as tcm
+from snowflake import snowpark
+from snowflake.connector.cursor import SnowflakeCursor
+from snowflake.connector.errors import NotSupportedError
+from snowflake.snowpark._internal.analyzer.snowflake_plan import PlanQueryType
+from snowflake.snowpark._internal.utils import (
+    create_or_update_statement_params_with_query_tag,
+)
+from snowflake.snowpark_connect.constants import SERVER_SIDE_SESSION_ID
+from snowflake.snowpark_connect.execute_plan.utils import (
+    arrow_table_to_arrow_bytes,
+    pandas_to_arrow_batches_bytes,
+)
+from snowflake.snowpark_connect.relation.map_relation import map_relation
+from snowflake.snowpark_connect.type_mapping import (
+    map_snowpark_types_to_pyarrow_types,
+    snowpark_to_proto_type,
+)
+
+QueryResult = namedtuple("QueryResult", ["query_id", "arrow_schema", "spark_schema"])
+
+
+def _build_execute_plan_response(
+    row_count: int, data_bytes: bytes, schema, request: proto_base.ExecutePlanRequest
+):
+    return proto_base.ExecutePlanResponse(
+        session_id=request.session_id,
+        operation_id=SERVER_SIDE_SESSION_ID,
+        arrow_batch=proto_base.ExecutePlanResponse.ArrowBatch(
+            row_count=row_count,
+            data=data_bytes,
+        ),
+        schema=schema,
+    )
+
+
+def sproc_connector_fetch_arrow_batches_fix(self) -> Iterator[Table]:
+    self.check_can_use_arrow_resultset()
+    if self._prefetch_hook is not None:
+        self._prefetch_hook()
+    if self._query_result_format != "arrow":
+        raise NotSupportedError
+    return self._result_set._fetch_arrow_batches()
+
+
+# TODO: SNOW-2039432 use to_arrow_batches once it is fixed in sproc-python-connector
+# We need to do this monkey patching because sproc python connector has a bug where it does not call _prefetch_hook()
+# inside fetch_arrow_batches() which causes TCM, notebook and stored procs to throw error on df.collect()
+SnowflakeCursor.fetch_arrow_batches = sproc_connector_fetch_arrow_batches_fix
+
+SKIP_LEVELS_TWO = (
+    2  # limit traceback to return up to 2 stack trace entries from traceback object tb
+)
+
+
+# TODO: SNOW-2039432 use to_arrow_batches once it is fixed in sproc-python-connector
+# TODO: SNOW-2057291 remove once df.to_arrow() starts accepting to_iter parameter
+def to_arrow_batch_iter(result_df: snowpark.DataFrame) -> Iterator[Table]:
+    return result_df.session._conn.execute(
+        result_df._plan,
+        to_pandas=False,
+        to_iter=True,
+        to_arrow=True,
+        block=True,
+        _statement_params=create_or_update_statement_params_with_query_tag(
+            result_df._statement_params,
+            result_df.session.query_tag,
+            SKIP_LEVELS_TWO,
+            collect_stacktrace=result_df.session.conf.get(
+                "collect_stacktrace_in_query_tag"
+            ),
+        ),
+    )
+
+
+def map_execution_root(
+    request: proto_base.ExecutePlanRequest,
+) -> Iterator[proto_base.ExecutePlanResponse | QueryResult]:
+    result_df: snowpark.DataFrame | pandas.DataFrame = map_relation(request.plan.root)
+    if isinstance(result_df, snowpark.DataFrame):
+        snowpark_schema = result_df.schema
+        schema = snowpark_to_proto_type(
+            snowpark_schema, result_df._column_map, result_df
+        )
+        spark_columns = result_df._column_map.get_spark_columns()
+        if tcm.TCM_MODE:
+            # TCM result handling:
+            # - small result (only one batch): just return the executePlanResponse
+            # - large result (more than one batch): return a tuple with query UUID, arrow schema, and spark schema.
+            # If TCM_RETURN_QUERY_ID_FOR_SMALL_RESULT is true, all results will be treated as large result.
+            is_large_result = False
+            second_batch = False
+            first_arrow_table = None
+            with result_df.session.query_history() as qh:
+                for arrow_table in to_arrow_batch_iter(result_df):
+                    if second_batch:
+                        is_large_result = True
+                        break
+                    first_arrow_table = arrow_table
+                    second_batch = True
+                queries_cnt = len(
+                    result_df._plan.execution_queries[PlanQueryType.QUERIES]
+                )
+                # get query uuid from the last query; this may not be the last queries in query history because snowpark
+                # may run some post action queries, e.g., drop temp table.
+                query_id = qh.queries[queries_cnt - 1].query_id
+            if first_arrow_table is None:
+                # empty arrow batch iterator
+                pandas_df = result_df.to_pandas()
+                data_bytes = pandas_to_arrow_batches_bytes(pandas_df)
+                yield _build_execute_plan_response(0, data_bytes, schema, request)
+            elif not tcm.TCM_RETURN_QUERY_ID_FOR_SMALL_RESULT and not is_large_result:
+                data_bytes = arrow_table_to_arrow_bytes(
+                    first_arrow_table, snowpark_schema, spark_columns
+                )
+                yield _build_execute_plan_response(
+                    first_arrow_table.num_rows, data_bytes, schema, request
+                )
+            else:
+                # return query id and serialized schemas
+                arrow_schema = pa.schema(
+                    map_snowpark_types_to_pyarrow_types(
+                        snowpark_schema,
+                        pa.struct(first_arrow_table.schema),
+                        rename_struct_columns=True,
+                    )
+                )
+                serialized_arrow_schema = arrow_schema.serialize().to_pybytes()
+                spark_schema = proto_types.DataType(**schema)
+                yield QueryResult(
+                    query_id,
+                    serialized_arrow_schema,
+                    spark_schema.SerializeToString(),
+                )
+        else:
+            arrow_table_iter = to_arrow_batch_iter(result_df)
+            batch_count = 0
+            for arrow_table in arrow_table_iter:
+                if arrow_table.num_rows > 0:
+                    batch_count += 1
+                    data_bytes = arrow_table_to_arrow_bytes(
+                        arrow_table, snowpark_schema, spark_columns
+                    )
+                    yield _build_execute_plan_response(
+                        arrow_table.num_rows, data_bytes, schema, request
+                    )
+                else:
+                    break
+
+            # Empty result needs special processing
+            if batch_count == 0:
+                pandas_df = result_df.to_pandas()
+                data_bytes = pandas_to_arrow_batches_bytes(pandas_df)
+                yield _build_execute_plan_response(0, data_bytes, schema, request)
+    else:
+        pandas_df = result_df
+        data_bytes = pandas_to_arrow_batches_bytes(pandas_df)
+        row_count = len(pandas_df)
+        schema = None
+        yield _build_execute_plan_response(row_count, data_bytes, schema, request)

snowflake/snowpark_connect/execute_plan/utils.py

@@ -0,0 +1,183 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import pandas
+import pyarrow as pa
+import pyspark.sql.connect.proto.relations_pb2 as relation_proto
+from pyspark.sql.pandas.types import _dedup_names
+
+from snowflake.snowpark import types as sf_types
+from snowflake.snowpark_connect.type_mapping import map_snowpark_types_to_pyarrow_types
+from snowflake.snowpark_connect.utils.telemetry import (
+    SnowparkConnectNotImplementedError,
+)
+
+
+def is_streaming(rel: relation_proto.Relation) -> bool:
+    """
+    Check if the relation is a streaming relation.
+
+    A streaming relation is a relation that is the result of a streaming
+    operation. This is used to determine if the relation should be shown
+    immediately or if it should be stored in the session state for later use.
+    """
+    """Check if the relation is a streaming relation."""
+    try:
+        match rel.WhichOneof("rel_type"):
+            case "read":
+                return rel.read.is_streaming is True
+            case "project":
+                return is_streaming(rel.project.input)
+            case "filter":
+                return is_streaming(rel.filter.input)
+            case "join":
+                return is_streaming(rel.join.left) or is_streaming(rel.join.right)
+            case "set_op":
+                return is_streaming(rel.set_op.input)
+            case "sort":
+                return is_streaming(rel.sort.input)
+            case "limit":
+                return is_streaming(rel.limit.input)
+            case "aggregate":
+                return is_streaming(rel.aggregate.input)
+            case "sample":
+                return is_streaming(rel.sample.input)
+            case "offset":
+                return is_streaming(rel.offset.input)
+            case "deduplicate":
+                return is_streaming(rel.deduplicate.input)
+            case "subquery_alias":
+                return is_streaming(rel.subquery_alias.input)
+            case "repartition":
+                return is_streaming(rel.repartition.input)
+            case "to_df":
+                return is_streaming(rel.to_df.input)
+            case "with_columns_renamed":
+                return is_streaming(rel.with_columns_renamed.input)
+            case "show_string":
+                return is_streaming(rel.show_string.input)
+            case "drop":
+                return is_streaming(rel.drop.input)
+            case "tail":
+                return is_streaming(rel.tail.input)
+            case "with_columns":
+                return is_streaming(rel.with_columns.input)
+            case "hint":
+                return is_streaming(rel.hint.input)
+            case "unpivot":
+                return is_streaming(rel.unpivot.input)
+            case "to_schema":
+                return is_streaming(rel.to_schema.input)
+            case "repartition_by_expression":
+                return is_streaming(rel.repartition_by_expression.input)
+            case "map_partitions":
+                return is_streaming(rel.map_partitions.input)
+            case "collect_metrics":
+                return is_streaming(rel.collect_metrics.input)
+            case "parse":
+                return is_streaming(rel.parse.input)
+            case "group_map":
+                return is_streaming(rel.group_map.input)
+            case "co_group_map":
+                return is_streaming(rel.co_group_map.input)
+            case "with_watermark":
+                return is_streaming(rel.with_watermark.input)
+            case "apply_in_pandas_with_state":
+                return is_streaming(rel.apply_in_pandas.input)
+            case "html_string":
+                return is_streaming(rel.html_string.input)
+            case "cached_remote_relation":
+                raise SnowparkConnectNotImplementedError(
+                    "Cached remote relation not implemented"
+                )
+            case "common_inline_user_defined_table_function":
+                return is_streaming(rel.common_inline_user_defined_table_function.input)
+            case "fill_na":
+                return is_streaming(rel.fill_na.input)
+            case "drop_na":
+                return is_streaming(rel.drop_na.input)
+            case "replace":
+                return is_streaming(rel.replace.input)
+            case "stat":
+                return is_streaming(rel.stat.input)
+            case "summary":
+                return is_streaming(rel.summary.input)
+            case "crosstab":
+                return is_streaming(rel.crosstab.input)
+            case "describe":
+                return is_streaming(rel.describe.input)
+            case "cov":
+                return is_streaming(rel.cov.input)
+            case "corr":
+                return is_streaming(rel.corr.input)
+            case "approx_quantile":
+                return is_streaming(rel.approx_quantile.input)
+            case "freq_items":
+                return is_streaming(rel.freq_items.input)
+            case "sample_by":
+                return is_streaming(rel.sample_by.input)
+            case _:
+                return False
+    except AttributeError:
+        # This is a leaf node with no `input`.
+        return False
+
+
+def pandas_to_arrow_batches_bytes(
+    pandas_df: pandas.DataFrame, pyarrow_schema: pa.Schema = None
+) -> bytes:
+    """
+    Serialize a pandas DataFrame as Pyarrow encoded bytes.
+    """
+    # Pyarrow doesn't support duplicate column names, so we need to deduplicate them.
+    # It is important that the schema is passed in whatever message we send back to the
+    # client, otherwise the names will not be correct.
+    pandas_df.columns = _dedup_names(pandas_df.columns)
+    sink = pa.BufferOutputStream()
+    batch = pa.RecordBatch.from_pandas(pandas_df, schema=pyarrow_schema)
+    with pa.ipc.new_stream(sink, batch.schema) as writer:
+        writer.write_batch(batch)
+    return sink.getvalue().to_pybytes()
+
+
+def arrow_table_to_arrow_bytes(
+    table: pa.Table, snowpark_schema: sf_types.StructType, spark_columns: list
+) -> bytes:
+    """
+    Serialize a pyarrow.table as Pyarrow encoded bytes according to provided snowpark schema.
+    """
+    assert table.num_rows > 0, "Table must have at least one row"
+
+    # 1. rename column names to 0,1,2, etc. to avoid unmatching names due to undesired factors like quotes.
+    # 2. casting is required here because sometimes arrow table does use expected data type. E.g., for LongType,
+    #       pyarrow table uses decimal128(38,0), which converts to Decimal instead of Long on client side.
+    #
+    # note that we don't need to track the original column name, since this helper function only needs to generate arrow
+    # data bytes. When the arrow bytes are returned to spark connect client, an explicit schema would be passed along,
+    # which contains expected column name. E.g.,
+    #   return [
+    #         proto_base.ExecutePlanResponse(
+    #             session_id=request.session_id,
+    #             operation_id=SERVER_SIDE_SESSION_ID,
+    #             arrow_batch=proto_base.ExecutePlanResponse.ArrowBatch(
+    #                 row_count=row_count,
+    #                 data=arrow_bytes, # arrow bytes generated by this helper function
+    #             ),
+    #             schema=schema,    # schema containing correct column name
+    #         ),
+    #   ]
+    schema = pa.schema(
+        map_snowpark_types_to_pyarrow_types(
+            snowpark_schema, pa.struct(table.schema), rename_struct_columns=True
+        )
+    )
+    table = table.rename_columns([str(i) for i in range(table.num_columns)])
+    table = table.cast(schema, safe=False)
+    table = table.rename_columns(spark_columns)
+
+    sink = pa.BufferOutputStream()
+    with pa.ipc.new_stream(sink, table.schema) as writer:
+        writer.write_table(table)
+    arrow_bytes = sink.getvalue().to_pybytes()
+    return arrow_bytes

snowflake/snowpark_connect/expression/__init__.py

@@ -0,0 +1,3 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#

snowflake/snowpark_connect/expression/literal.py

@@ -0,0 +1,90 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import datetime
+import decimal
+from zoneinfo import ZoneInfo
+
+import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
+from tzlocal import get_localzone
+
+from snowflake.snowpark_connect.config import global_config
+from snowflake.snowpark_connect.utils.telemetry import (
+    SnowparkConnectNotImplementedError,
+)
+
+
+def get_literal_field_and_name(literal: expressions_proto.Expression.Literal):
+    match literal.WhichOneof("literal_type"):
+        case "byte":
+            return literal.byte, str(literal.byte)
+        case "short":
+            return literal.short, str(literal.short)
+        case "integer":
+            return literal.integer, str(literal.integer)
+        case "long":
+            return literal.long, str(literal.long)
+        case "float":
+            return (
+                literal.float,
+                str(literal.float) if literal.float == literal.float else "NaN",
+            )
+        case "double":
+            return (
+                literal.double,
+                str(literal.double) if literal.double == literal.double else "NaN",
+            )
+        case "string":
+            return literal.string, str(literal.string)
+        case "boolean":
+            return literal.boolean, str(literal.boolean)
+        case "date":
+            # Both snowflake and spark Date type don't consider time zones.
+            # Don't use datetime.date.fromtimestamp, which depends on local timezone
+            date = datetime.datetime.fromtimestamp(
+                literal.date * 86400, tz=datetime.timezone.utc
+            ).date()
+            return date, f"DATE '{date}'"
+        case "timestamp" | "timestamp_ntz" as t:
+            # Note - Clients need to ensure local_timezone is the same as spark_sql_session_timeZone config.
+            # No need to apply timezone for lit datetime, because we set the TIMEZONE parameter in snowpark session,
+            # the snowflake backend would convert the lit datetime correctly. However, for returned column name, the
+            # timezone needs to be added. Pyspark has a weird behavior that datetime.datetime always gets converted
+            # to local timezone before printing according to spark_sql_session_timeZone setting. Haven't found
+            # official doc about it, but this behavior is based on my testings.
+            tz = (
+                ZoneInfo(global_config.spark_sql_session_timeZone)
+                if hasattr(global_config, "spark_sql_session_timeZone")
+                else get_localzone()
+            )
+            if t == "timestamp":
+                microseconds = literal.timestamp
+            else:
+                microseconds = literal.timestamp_ntz
+            lit_dt = datetime.datetime.fromtimestamp(
+                microseconds // 1_000_000
+            ) + datetime.timedelta(microseconds=microseconds % 1_000_000)
+            tz_dt = datetime.datetime.fromtimestamp(
+                microseconds // 1_000_000, tz=tz
+            ) + datetime.timedelta(microseconds=microseconds % 1_000_000)
+            if t == "timestamp_ntz":
+                lit_dt = lit_dt.astimezone(datetime.timezone.utc)
+                tz_dt = tz_dt.astimezone(datetime.timezone.utc)
+            return lit_dt, f"{t.upper()} '{tz_dt.strftime('%Y-%m-%d %H:%M:%S')}'"
+        case "day_time_interval":
+            # TODO(SNOW-1920942): Snowflake SQL is missing an "interval" type.
+            timedelta = datetime.timedelta(
+                seconds=literal.day_time_interval / 1_000_000
+            )
+            str_value = f"INTERVAL '{literal.day_time_interval / 1_000_000} SECONDS'"
+            return timedelta, str_value
+        case "binary":
+            return literal.binary, str(literal.binary)
+        case "decimal":
+            # literal.decimal.precision & scale are ignored, as decimal.Decimal doesn't accept them
+            return decimal.Decimal(literal.decimal.value), literal.decimal.value
+        case "null" | None:
+            return None, "NULL"
+        case other:
+            raise SnowparkConnectNotImplementedError(f"Other Literal Type {other}")