snowpark-connect 0.27.0__py3-none-any.whl → 1.7.0__py3-none-any.whl

snowflake/snowpark_connect/__init__.py

@@ -10,6 +10,7 @@ sys.path.append(str(pathlib.Path(__file__).parent / "includes/python"))
 
 from .server import get_session  # noqa: E402, F401
 from .server import start_session  # noqa: E402, F401
+from .utils.session import skip_session_configuration  # noqa: E402, F401
 
 # Turn off catalog warning for Snowpark
 sp_logger = logging.getLogger("snowflake.snowpark")

snowflake/snowpark_connect/analyze_plan/map_tree_string.py

@@ -5,6 +5,9 @@
 import pyspark.sql.connect.proto.base_pb2 as proto_base
 
 from snowflake.snowpark_connect.relation.map_relation import map_relation
+from snowflake.snowpark_connect.relation.read.metadata_utils import (
+    without_internal_columns,
+)
 from snowflake.snowpark_connect.type_mapping import (
     SNOWPARK_TYPE_NAME_TO_PYSPARK_TYPE_NAME,
 )
@@ -16,12 +19,13 @@ def map_tree_string(
     # TODO: tracking the difference with pyspark in SNOW-1853347
     tree_string = request.tree_string
     snowpark_df_container = map_relation(tree_string.plan.root)
-    snowpark_df = snowpark_df_container.dataframe
-    column_map = snowpark_df_container.column_map
+    filtered_container = without_internal_columns(snowpark_df_container)
+    display_df = filtered_container.dataframe
+    filtered_column_mapping = filtered_container.column_map.snowpark_to_spark_map()
 
-    snowpark_tree_string = snowpark_df._format_schema(
+    snowpark_tree_string = display_df._format_schema(
         level=tree_string.level if tree_string.HasField("level") else None,
-        translate_columns=column_map.snowpark_to_spark_map(),
+        translate_columns=filtered_column_mapping,
         translate_types=SNOWPARK_TYPE_NAME_TO_PYSPARK_TYPE_NAME,
     )
     # workaround for the capitalization of nullable boolean value.

snowflake/snowpark_connect/client/__init__.py

@@ -0,0 +1,15 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+"""
+Lightweight Snowpark Connect client.
+"""
+
+from snowflake.snowpark_connect.client.server import (  # noqa: F401
+    get_session,
+    init_spark_session,
+    start_session,
+)
+
+__all__ = ["get_session", "init_spark_session", "start_session"]

snowflake/snowpark_connect/client/error_utils.py

@@ -0,0 +1,30 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+"""
+Lightweight error utilities for the Snowpark Connect client.
+
+This is a minimal version of snowflake.snowpark_connect.error.error_utils
+that avoids heavy dependencies like jpype that are only needed server-side.
+"""
+
+import threading
+
+# Thread-local storage for custom error codes when we can't attach them directly to exceptions
+_thread_local = threading.local()
+
+
+def attach_custom_error_code(exception: Exception, custom_error_code: int) -> Exception:
+    """
+    Attach a custom error code to any exception instance.
+    This allows us to add custom error codes to existing PySpark exceptions.
+    """
+    if not hasattr(exception, "custom_error_code"):
+        try:
+            exception.custom_error_code = custom_error_code
+        except (AttributeError, TypeError):
+            # Some exception types don't allow setting custom attributes
+            # Store the error code in thread-local storage for later retrieval
+            _thread_local.pending_error_code = custom_error_code
+    return exception

snowflake/snowpark_connect/client/exceptions.py

@@ -0,0 +1,36 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+from google.rpc import status_pb2
+
+
+class SparkConnectServerException(Exception):
+    def __init__(self, message) -> None:
+        super().__init__(message)
+
+
+class InvalidHostException(SparkConnectServerException):
+    def __init__(self, message) -> None:
+        super().__init__(message)
+
+
+class MissingAuthException(SparkConnectServerException):
+    def __init__(self) -> None:
+        super().__init__("Missing authorization information")
+
+
+class UnexpectedResponseException(SparkConnectServerException):
+    def __init__(self, message) -> None:
+        super().__init__(message)
+
+
+class QueryTimeoutException(SparkConnectServerException):
+    def __init__(self, message) -> None:
+        super().__init__(message)
+
+
+class GrpcErrorStatusException(SparkConnectServerException):
+    def __init__(self, status: status_pb2.Status) -> None:
+        super().__init__(f"Error code: {status.code}, message: {status.message}")
+        self.status = status

snowflake/snowpark_connect/client/query_results.py

@@ -0,0 +1,90 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+"""Query result fetching utilities for remote client."""
+
+import base64
+import logging
+from typing import Iterator, Tuple
+
+import pyarrow as pa
+from google.protobuf.message import DecodeError
+from google.rpc import code_pb2, status_pb2
+
+from snowflake.snowpark import Session
+from spark.connect import envelope_pb2
+
+logger = logging.getLogger(__name__)
+
+
+def fetch_query_result_as_arrow_batches(
+    session: Session, query_id: str, arrow_schema: pa.Schema
+) -> Iterator[Tuple[int, bytes]]:
+    """
+    Fetch query results as Arrow batches.
+    Used for large result sets that are fetched asynchronously.
+
+    Yields:
+        Tuples of (row_count, arrow_batch_bytes)
+    """
+    conn = session.connection
+    cursor = conn.cursor()
+
+    try:
+        cursor.get_results_from_sfqid(query_id)
+
+        for arrow_table in cursor.fetch_arrow_batches():
+            arrow_table = arrow_table.rename_columns(
+                [str(i) for i in range(arrow_table.num_columns)]
+            )
+            # Cast to expected schema
+            arrow_table = arrow_table.cast(arrow_schema, safe=False)
+
+            # Convert table to batches and serialize
+            record_batches = arrow_table.to_batches()
+            for batch in record_batches:
+                sink = pa.BufferOutputStream()
+                with pa.ipc.new_stream(sink, batch.schema) as writer:
+                    writer.write_batch(batch)
+                arrow_batch_bytes = sink.getvalue().to_pybytes()
+                yield batch.num_rows, arrow_batch_bytes
+    finally:
+        cursor.close()
+
+
+def fetch_query_result_as_protobuf(
+    session: Session, query_id: str
+) -> envelope_pb2.ResponseEnvelope:
+    """
+    Fetch query results as a ResponseEnvelope protobuf.
+    Polls for async query completion and decodes the result envelope.
+
+    Returns:
+        ResponseEnvelope containing the results
+    """
+    conn = session.connection
+    cursor = conn.cursor()
+
+    try:
+        cursor.get_results_from_sfqid(query_id)
+        # Python connector does not know how to handle new "RAW" result type, so for now we must use
+        # internal, private methods directly to work around this limitation
+        # TODO: Add proper RAW result handling in Python connector
+        cursor._prefetch_hook()  # block while polling till query is complete
+        data = cursor._result_set.batches[0]._data
+        resp_bytes = base64.b64decode(data)
+
+        try:
+            response_envelope = envelope_pb2.ResponseEnvelope()
+            response_envelope.ParseFromString(resp_bytes)
+            return response_envelope
+        except DecodeError:
+            logger.error(f"Failed to decode ResponseEnvelope for query_id: {query_id}")
+            return envelope_pb2.ResponseEnvelope(
+                status=status_pb2.Status(
+                    code=code_pb2.INTERNAL, message="Invalid operation or SQL"
+                )
+            )
+    finally:
+        cursor.close()