pyspark-client 4.1.0.dev2__tar.gz → 4.1.0.dev4__tar.gz

{pyspark_client-4.1.0.dev2/pyspark_client.egg-info → pyspark_client-4.1.0.dev4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyspark-client
-Version: 4.1.0.dev2
+Version: 4.1.0.dev4
 Summary: Python Spark Connect client for Apache Spark
 Home-page: https://github.com/apache/spark/tree/master/python
 Author: Spark Developers
@@ -11,6 +11,7 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
 Classifier: Typing :: Typed
@@ -18,9 +19,10 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 Requires-Dist: pandas>=2.2.0
 Requires-Dist: pyarrow>=15.0.0
-Requires-Dist: grpcio>=1.67.0
-Requires-Dist: grpcio-status>=1.67.0
-Requires-Dist: googleapis-common-protos>=1.65.0
+Requires-Dist: grpcio>=1.76.0
+Requires-Dist: grpcio-status>=1.76.0
+Requires-Dist: googleapis-common-protos>=1.71.0
+Requires-Dist: zstandard>=0.25.0
 Requires-Dist: numpy>=1.21
 Requires-Dist: pyyaml>=3.11
 Dynamic: author

{pyspark_client-4.1.0.dev2 → pyspark_client-4.1.0.dev4}/pyspark/cloudpickle/__init__.py

@@ -3,7 +3,7 @@ from pyspark.cloudpickle.cloudpickle import *  # noqa
 
 __doc__ = cloudpickle.__doc__
 
-__version__ = "3.1.1"
+__version__ = "3.1.2"
 
 __all__ = [  # noqa
     "__version__",

{pyspark_client-4.1.0.dev2 → pyspark_client-4.1.0.dev4}/pyspark/cloudpickle/cloudpickle.py

@@ -783,6 +783,12 @@ def _class_getstate(obj):
 
     clsdict.pop("__dict__", None)  # unpicklable property object
 
+    if sys.version_info >= (3, 14):
+        # PEP-649/749: __annotate_func__ contains a closure that references the class
+        # dict. We need to exclude it from pickling. Python will recreate it when
+        # __annotations__ is accessed at unpickling time.
+        clsdict.pop("__annotate_func__", None)
+
     return (clsdict, {})
 
 
@@ -1190,6 +1196,10 @@ def _class_setstate(obj, state):
         for subclass in registry:
             obj.register(subclass)
 
+    # PEP-649/749: During pickling, we excluded the __annotate_func__ attribute but it
+    # will be created by Python. Subsequently, annotations will be recreated when
+    # __annotations__ is accessed.
+
     return obj
 
 
@@ -1301,12 +1311,9 @@ class Pickler(pickle.Pickler):
     def dump(self, obj):
         try:
             return super().dump(obj)
-        except RuntimeError as e:
-            if len(e.args) > 0 and "recursion" in e.args[0]:
-                msg = "Could not pickle object as excessively deep recursion required."
-                raise pickle.PicklingError(msg) from e
-            else:
-                raise
+        except RecursionError as e:
+            msg = "Could not pickle object as excessively deep recursion required."
+            raise pickle.PicklingError(msg) from e
 
     def __init__(self, file, protocol=None, buffer_callback=None):
         if protocol is None:

{pyspark_client-4.1.0.dev2 → pyspark_client-4.1.0.dev4}/pyspark/errors/error-conditions.json

@@ -549,6 +549,16 @@
       "<arg1> and <arg2> should be of the same length, got <arg1_length> and <arg2_length>."
     ]
   },
+  "MALFORMED_GEOGRAPHY": {
+    "message": [
+      "Geography binary is malformed. Please check the data source is valid."
+    ]
+  },
+  "MALFORMED_GEOMETRY": {
+    "message": [
+      "Geometry binary is malformed. Please check the data source is valid."
+    ]
+  },
   "MALFORMED_VARIANT": {
     "message": [
       "Variant binary is malformed. Please check the data source is valid."
@@ -1134,6 +1144,24 @@
       "Cannot serialize the function `<name>`. If you accessed the Spark session, or a DataFrame defined outside of the function, or any object that contains a Spark session, please be aware that they are not allowed in Spark Connect. For `foreachBatch`, please access the Spark session using `df.sparkSession`, where `df` is the first parameter in your `foreachBatch` function. For `StreamingQueryListener`, please access the Spark session using `self.spark`. For details please check out the PySpark doc for `foreachBatch` and `StreamingQueryListener`."
     ]
   },
+  "ST_INVALID_ALGORITHM_VALUE" : {
+    "message" : [
+      "Invalid or unsupported edge interpolation algorithm value: '<alg>'."
+    ],
+    "sqlState" : "22023"
+  },
+  "ST_INVALID_CRS_VALUE" : {
+    "message" : [
+      "Invalid or unsupported CRS (coordinate reference system) value: '<crs>'."
+    ],
+    "sqlState" : "22023"
+  },
+  "ST_INVALID_SRID_VALUE" : {
+    "message" : [
+      "Invalid or unsupported SRID (spatial reference identifier) value: <srid>."
+    ],
+    "sqlState" : "22023"
+  },
   "TEST_CLASS_NOT_COMPILED": {
     "message": [
       "<test_class_path> doesn't exist. Spark sql test classes are not compiled."

{pyspark_client-4.1.0.dev2 → pyspark_client-4.1.0.dev4}/pyspark/errors/exceptions/captured.py

@@ -107,7 +107,8 @@ class CapturedException(PySparkException):
         if self._origin is not None and is_instance_of(
             gw, self._origin, "org.apache.spark.SparkThrowable"
         ):
-            return self._origin.getCondition()
+            utils = SparkContext._jvm.PythonErrorUtils  # type: ignore[union-attr]
+            return utils.getCondition(self._origin)
         else:
             return None
 
@@ -118,7 +119,6 @@ class CapturedException(PySparkException):
     def getMessageParameters(self) -> Optional[Dict[str, str]]:
         from pyspark import SparkContext
         from py4j.java_gateway import is_instance_of
-        from py4j.protocol import Py4JError
 
         assert SparkContext._gateway is not None
 
@@ -126,38 +126,28 @@ class CapturedException(PySparkException):
         if self._origin is not None and is_instance_of(
             gw, self._origin, "org.apache.spark.SparkThrowable"
         ):
-            try:
-                return dict(self._origin.getMessageParameters())
-            except Py4JError as e:
-                if "py4j.Py4JException" in str(e) and "Method getMessageParameters" in str(e):
-                    return None
-                raise e
+            utils = SparkContext._jvm.PythonErrorUtils  # type: ignore[union-attr]
+            return dict(utils.getMessageParameters(self._origin))
         else:
             return None
 
     def getSqlState(self) -> Optional[str]:
         from pyspark import SparkContext
         from py4j.java_gateway import is_instance_of
-        from py4j.protocol import Py4JError
 
         assert SparkContext._gateway is not None
         gw = SparkContext._gateway
         if self._origin is not None and is_instance_of(
             gw, self._origin, "org.apache.spark.SparkThrowable"
         ):
-            try:
-                return self._origin.getSqlState()
-            except Py4JError as e:
-                if "py4j.Py4JException" in str(e) and "Method getSqlState" in str(e):
-                    return None
-                raise e
+            utils = SparkContext._jvm.PythonErrorUtils  # type: ignore[union-attr]
+            return utils.getSqlState(self._origin)
         else:
             return None
 
     def getMessage(self) -> str:
         from pyspark import SparkContext
         from py4j.java_gateway import is_instance_of
-        from py4j.protocol import Py4JError
 
         assert SparkContext._gateway is not None
         gw = SparkContext._gateway
@@ -165,21 +155,12 @@ class CapturedException(PySparkException):
         if self._origin is not None and is_instance_of(
             gw, self._origin, "org.apache.spark.SparkThrowable"
         ):
-            try:
-                error_class = self._origin.getCondition()
-            except Py4JError as e:
-                if "py4j.Py4JException" in str(e) and "Method getCondition" in str(e):
-                    return ""
-                raise e
-            try:
-                message_parameters = self._origin.getMessageParameters()
-            except Py4JError as e:
-                if "py4j.Py4JException" in str(e) and "Method getMessageParameters" in str(e):
-                    return ""
-                raise e
+            utils = SparkContext._jvm.PythonErrorUtils  # type: ignore[union-attr]
+            errorClass = utils.getCondition(self._origin)
+            messageParameters = utils.getMessageParameters(self._origin)
 
             error_message = getattr(gw.jvm, "org.apache.spark.SparkThrowableHelper").getMessage(
-                error_class, message_parameters
+                errorClass, messageParameters
             )
 
             return error_message
@@ -189,7 +170,6 @@ class CapturedException(PySparkException):
     def getQueryContext(self) -> List[BaseQueryContext]:
         from pyspark import SparkContext
         from py4j.java_gateway import is_instance_of
-        from py4j.protocol import Py4JError
 
         assert SparkContext._gateway is not None
 
@@ -198,13 +178,8 @@ class CapturedException(PySparkException):
             gw, self._origin, "org.apache.spark.SparkThrowable"
         ):
             contexts: List[BaseQueryContext] = []
-            try:
-                context = self._origin.getQueryContext()
-            except Py4JError as e:
-                if "py4j.Py4JException" in str(e) and "Method getQueryContext" in str(e):
-                    return []
-                raise e
-            for q in context:
+            utils = SparkContext._jvm.PythonErrorUtils  # type: ignore[union-attr]
+            for q in utils.getQueryContext(self._origin):
                 if q.contextType().toString() == "SQL":
                     contexts.append(SQLQueryContext(q))
                 else:

{pyspark_client-4.1.0.dev2 → pyspark_client-4.1.0.dev4}/pyspark/logger/logger.py

@@ -50,6 +50,10 @@ class JSONFormatter(logging.Formatter):
 
     default_msec_format = "%s.%03d"
 
+    def __init__(self, ensure_ascii: bool = False):
+        super().__init__()
+        self._ensure_ascii = ensure_ascii
+
     def format(self, record: logging.LogRecord) -> str:
         """
         Format the specified record as a JSON string.
@@ -69,7 +73,7 @@ class JSONFormatter(logging.Formatter):
             "level": record.levelname,
             "logger": record.name,
             "msg": record.getMessage(),
-            "context": record.__dict__.get("kwargs", {}),
+            "context": record.__dict__.get("context", {}),
         }
         if record.exc_info:
             exc_type, exc_value, exc_tb = record.exc_info
@@ -89,7 +93,7 @@ class JSONFormatter(logging.Formatter):
                 "msg": str(exc_value),
                 "stacktrace": structured_stacktrace,
             }
-        return json.dumps(log_entry, ensure_ascii=False)
+        return json.dumps(log_entry, ensure_ascii=self._ensure_ascii)
 
 
 class PySparkLogger(logging.Logger):
@@ -136,7 +140,19 @@ class PySparkLogger(logging.Logger):
     """
 
     def __init__(self, name: str = "PySparkLogger"):
+        from pyspark.logger.worker_io import JSONFormatterWithMarker
+
         super().__init__(name, level=logging.WARN)
+
+        root_logger = logging.getLogger()
+        if any(
+            isinstance(h, logging.StreamHandler)
+            and isinstance(h.formatter, JSONFormatterWithMarker)
+            for h in root_logger.handlers
+        ):
+            # Likely in the `capture_outputs` context, so don't add a handler
+            return
+
         _handler = logging.StreamHandler()
         self.addHandler(_handler)
 
@@ -291,7 +307,7 @@ class PySparkLogger(logging.Logger):
             msg=msg,
             args=args,
             exc_info=exc_info,
-            extra={"kwargs": kwargs},
+            extra={"context": kwargs},
             stack_info=stack_info,
             stacklevel=stacklevel,
         )

pyspark_client-4.1.0.dev4/pyspark/logger/worker_io.py

@@ -0,0 +1,297 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from contextlib import contextmanager
+import inspect
+import io
+import logging
+import os
+import sys
+import time
+from typing import BinaryIO, Callable, Generator, Iterable, Iterator, Optional, TextIO, Union
+from types import FrameType, TracebackType
+
+from pyspark.logger.logger import JSONFormatter
+
+
+class DelegatingTextIOWrapper(TextIO):
+    """A TextIO that delegates all operations to another TextIO object."""
+
+    def __init__(self, delegate: TextIO):
+        self._delegate = delegate
+
+    # Required TextIO properties
+    @property
+    def encoding(self) -> str:
+        return self._delegate.encoding
+
+    @property
+    def errors(self) -> Optional[str]:
+        return self._delegate.errors
+
+    @property
+    def newlines(self) -> Optional[Union[str, tuple[str, ...]]]:
+        return self._delegate.newlines
+
+    @property
+    def buffer(self) -> BinaryIO:
+        return self._delegate.buffer
+
+    @property
+    def mode(self) -> str:
+        return self._delegate.mode
+
+    @property
+    def name(self) -> str:
+        return self._delegate.name
+
+    @property
+    def line_buffering(self) -> int:
+        return self._delegate.line_buffering
+
+    @property
+    def closed(self) -> bool:
+        return self._delegate.closed
+
+    # Iterator protocol
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._delegate)
+
+    def __next__(self) -> str:
+        return next(self._delegate)
+
+    # Context manager protocol
+    def __enter__(self) -> TextIO:
+        return self._delegate.__enter__()
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        return self._delegate.__exit__(exc_type, exc_val, exc_tb)
+
+    # Core I/O methods
+    def write(self, s: str) -> int:
+        return self._delegate.write(s)
+
+    def writelines(self, lines: Iterable[str]) -> None:
+        return self._delegate.writelines(lines)
+
+    def read(self, size: int = -1) -> str:
+        return self._delegate.read(size)
+
+    def readline(self, size: int = -1) -> str:
+        return self._delegate.readline(size)
+
+    def readlines(self, hint: int = -1) -> list[str]:
+        return self._delegate.readlines(hint)
+
+    # Stream control methods
+    def close(self) -> None:
+        return self._delegate.close()
+
+    def flush(self) -> None:
+        return self._delegate.flush()
+
+    def seek(self, offset: int, whence: int = io.SEEK_SET) -> int:
+        return self._delegate.seek(offset, whence)
+
+    def tell(self) -> int:
+        return self._delegate.tell()
+
+    def truncate(self, size: Optional[int] = None) -> int:
+        return self._delegate.truncate(size)
+
+    # Stream capability methods
+    def fileno(self) -> int:
+        return self._delegate.fileno()
+
+    def isatty(self) -> bool:
+        return self._delegate.isatty()
+
+    def readable(self) -> bool:
+        return self._delegate.readable()
+
+    def seekable(self) -> bool:
+        return self._delegate.seekable()
+
+    def writable(self) -> bool:
+        return self._delegate.writable()
+
+
+class JSONFormatterWithMarker(JSONFormatter):
+    default_microsec_format = "%s.%06d"
+
+    def __init__(self, marker: str, worker_id: str, context_provider: Callable[[], dict[str, str]]):
+        super().__init__(ensure_ascii=True)
+        self._marker = marker
+        self._worker_id = worker_id
+        self._context_provider = context_provider
+
+    def format(self, record: logging.LogRecord) -> str:
+        context = self._context_provider()
+        if context:
+            context.update(record.__dict__.get("context", {}))
+            record.__dict__["context"] = context
+        return f"{self._marker}:{self._worker_id}:{super().format(record)}"
+
+    def formatTime(self, record: logging.LogRecord, datefmt: Optional[str] = None) -> str:
+        ct = self.converter(record.created)
+        if datefmt:
+            s = time.strftime(datefmt, ct)
+        else:
+            s = time.strftime(self.default_time_format, ct)
+            if self.default_microsec_format:
+                s = self.default_microsec_format % (
+                    s,
+                    int((record.created - int(record.created)) * 1000000),
+                )
+            elif self.default_msec_format:
+                s = self.default_msec_format % (s, record.msecs)
+            s = f"{s}{time.strftime('%z', ct)}"
+        return s
+
+
+class JsonOutput(DelegatingTextIOWrapper):
+    def __init__(
+        self,
+        delegate: TextIO,
+        json_out: TextIO,
+        logger_name: str,
+        log_level: int,
+        marker: str,
+        worker_id: str,
+        context_provider: Callable[[], dict[str, str]],
+    ):
+        super().__init__(delegate)
+        self._json_out = json_out
+        self._logger_name = logger_name
+        self._log_level = log_level
+        self._formatter = JSONFormatterWithMarker(marker, worker_id, context_provider)
+
+    def write(self, s: str) -> int:
+        if s.strip():
+            log_record = logging.LogRecord(
+                name=self._logger_name,
+                level=self._log_level,
+                pathname=None,  # type: ignore[arg-type]
+                lineno=None,  # type: ignore[arg-type]
+                msg=s.strip(),
+                args=None,
+                exc_info=None,
+                func=None,
+                sinfo=None,
+            )
+            self._json_out.write(f"{self._formatter.format(log_record)}\n")
+            self._json_out.flush()
+        return self._delegate.write(s)
+
+    def writelines(self, lines: Iterable[str]) -> None:
+        # Process each line through our JSON logging logic
+        for line in lines:
+            self.write(line)
+
+    def close(self) -> None:
+        pass
+
+
+def context_provider() -> dict[str, str]:
+    """
+    Provides context information for logging, including caller function name.
+    Finds the function name from the bottom of the stack, ignoring Python builtin
+    libraries and PySpark modules. Test packages are included.
+
+    Returns:
+        dict[str, str]: A dictionary containing context information including:
+            - func_name: Name of the function that initiated the logging
+            - class_name: Name of the class that initiated the logging if available
+    """
+
+    def is_pyspark_module(module_name: str) -> bool:
+        return module_name.startswith("pyspark.") and ".tests." not in module_name
+
+    bottom: Optional[FrameType] = None
+
+    # Get caller function information using inspect
+    try:
+        frame = inspect.currentframe()
+        is_in_pyspark_module = False
+
+        if frame:
+            while frame.f_back:
+                f_back = frame.f_back
+                module_name = f_back.f_globals.get("__name__", "")
+
+                if is_pyspark_module(module_name):
+                    if not is_in_pyspark_module:
+                        bottom = frame
+                        is_in_pyspark_module = True
+                else:
+                    is_in_pyspark_module = False
+
+                frame = f_back
+    except Exception:
+        # If anything goes wrong with introspection, don't fail the logging
+        # Just continue without caller information
+        pass
+
+    context = {}
+    if bottom:
+        context["func_name"] = bottom.f_code.co_name
+        if "self" in bottom.f_locals:
+            context["class_name"] = bottom.f_locals["self"].__class__.__name__
+        elif "cls" in bottom.f_locals:
+            context["class_name"] = bottom.f_locals["cls"].__name__
+    return context
+
+
+@contextmanager
+def capture_outputs(
+    context_provider: Callable[[], dict[str, str]] = context_provider
+) -> Generator[None, None, None]:
+    if "PYSPARK_SPARK_SESSION_UUID" in os.environ:
+        marker: str = "PYTHON_WORKER_LOGGING"
+        worker_id: str = str(os.getpid())
+        json_out = original_stdout = sys.stdout
+        delegate = original_stderr = sys.stderr
+
+        handler = logging.StreamHandler(json_out)
+        handler.setFormatter(JSONFormatterWithMarker(marker, worker_id, context_provider))
+        logger = logging.getLogger()
+        try:
+            sys.stdout = JsonOutput(
+                delegate, json_out, "stdout", logging.INFO, marker, worker_id, context_provider
+            )
+            sys.stderr = JsonOutput(
+                delegate, json_out, "stderr", logging.ERROR, marker, worker_id, context_provider
+            )
+            logger.addHandler(handler)
+            try:
+                yield
+            finally:
+                # Send an empty line to indicate the end of the outputs.
+                json_out.write(f"{marker}:{worker_id}:\n")
+                json_out.flush()
+        finally:
+            sys.stdout = original_stdout
+            sys.stderr = original_stderr
+            logger.removeHandler(handler)
+            handler.close()
+    else:
+        yield

{pyspark_client-4.1.0.dev2 → pyspark_client-4.1.0.dev4}/pyspark/pandas/supported_api_gen.py

@@ -38,7 +38,7 @@ from pyspark.pandas.exceptions import PandasNotImplementedError
 MAX_MISSING_PARAMS_SIZE = 5
 COMMON_PARAMETER_SET = {"kwargs", "args", "cls"}
 MODULE_GROUP_MATCH = [(pd, ps), (pdw, psw), (pdg, psg)]
-PANDAS_LATEST_VERSION = "2.3.2"
+PANDAS_LATEST_VERSION = "2.3.3"
 
 RST_HEADER = """
 =====================

{pyspark_client-4.1.0.dev2 → pyspark_client-4.1.0.dev4}/pyspark/pipelines/__init__.py

@@ -20,6 +20,7 @@ from pyspark.pipelines.api import (
     materialized_view,
     table,
     temporary_view,
+    create_sink,
 )
 
 __all__ = [
@@ -28,4 +29,5 @@ __all__ = [
     "materialized_view",
     "table",
     "temporary_view",
+    "create_sink",
 ]

pyspark_client-4.1.0.dev4/pyspark/pipelines/add_pipeline_analysis_context.py

@@ -0,0 +1,48 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from contextlib import contextmanager
+from typing import Generator, Optional
+from pyspark.sql import SparkSession
+
+from typing import Any, cast
+
+
+@contextmanager
+def add_pipeline_analysis_context(
+    spark: SparkSession, dataflow_graph_id: str, flow_name: Optional[str]
+) -> Generator[None, None, None]:
+    """
+    Context manager that add PipelineAnalysisContext extension to the user context
+    used for pipeline specific analysis.
+    """
+    extension_id = None
+    # Cast because mypy seems to think `spark` is a function, not an object.
+    # Likely related to SPARK-47544.
+    client = cast(Any, spark).client
+    try:
+        import pyspark.sql.connect.proto as pb2
+        from google.protobuf import any_pb2
+
+        analysis_context = pb2.PipelineAnalysisContext(
+            dataflow_graph_id=dataflow_graph_id, flow_name=flow_name
+        )
+        extension = any_pb2.Any()
+        extension.Pack(analysis_context)
+        extension_id = client.add_threadlocal_user_context_extension(extension)
+        yield
+    finally:
+        client.remove_user_context_extension(extension_id)