python/pyspark/__init__.py

#!pyspark-rs Drop-in replacement for pyspark client
#
# Exposes the Rust-backed Spark Connect client as a pure-Python package.

from pyspark.version import __version__  # noqa: F401  (tracks the Spark version we target)

# --- Client identity ---------------------------------------------------------
# This distribution is ``pyspark-client-rust``: the Spark Connect client backed by
# the native Rust engine (tonic), installed under the ``pyspark`` import name as a
# drop-in for the reference ``pyspark-client``. These markers let user code and bug
# reports tell, at runtime, which client is in use -- this one, not the reference
# Python client. See https://apache.github.io/spark-connect-rust/which-client/
__rust_client__: bool = True
__engine__: str = "rust"  # native Rust core via tonic -- not grpcio/py4j

from typing import Callable, TypeVar, Union

_F = TypeVar("_F", bound=Callable)


def since(version: Union[str, float]) -> Callable[[_F], _F]:
    """
    A decorator that annotates a function to append the version of Spark the function was added.
    For Connect-only clients, this is a no-op that just passes through the function.
    """
    def deco(f: _F) -> _F:
        # For Connect-only, we don't modify docstrings; this is just a pass-through decorator
        return f
    return deco


from pyspark.sql import (
    SparkSession,
    DataFrame,
    Column,
    Row,
)
from pyspark.storagelevel import StorageLevel

# For compatibility with testing harness, provide a stub SparkConf that works with
# Spark Connect. This is imported by the testing utils but we don't need a full
# implementation since we're connect-only.
class SparkConf:
    """Stub SparkConf for Spark Connect testing.

    For pure Connect mode, we don't need a full SparkConf with JVM support.
    This stub accepts configuration as upstream would.
    """
    def __init__(self, loadDefaults=True, _jvm=None, _jconf=None):
        self._jconf = None
        self.settings = {}

    def set(self, key, value):
        """Set a configuration value."""
        self.settings[key] = value
        return self

    def get(self, key, defaultValue=None):
        """Get a configuration value."""
        return self.settings.get(key, defaultValue)

    def remove(self, key):
        """Remove a configuration key."""
        if key in self.settings:
            del self.settings[key]
        return self

    def setAll(self, settings):
        """Set multiple configuration values."""
        for key, value in settings.items():
            self.set(key, value)
        return self

    def getAll(self):
        """Get all configuration values."""
        return list(self.settings.items())


__all__ = [
    "SparkSession",
    "DataFrame",
    "Column",
    "Row",
    "SparkConf",
    "StorageLevel",
]