`yggdrasil.spark.tests`¶

tests ¶

Unittest base class for PySpark tests.

Provides a single global SparkSession shared across all SparkTestCase subclasses in the process. Spark startup is expensive (~5-10s) — creating a session per test class would make the suite painfully slow.

Quick start¶

::

from yggdrasil.spark.tests import SparkTestCase

class TestMyStuff(SparkTestCase):
    def test_basic(self):
        df = self.spark.createDataFrame([(1, "a")], ["id", "val"])
        self.assertEqual(df.count(), 1)

    def test_arrow_roundtrip(self):
        import pyarrow as pa
        tbl = pa.table({"id": [1, 2], "val": ["a", "b"]})
        df = self.arrow_to_spark(tbl)
        self.assertSparkEqual(df, tbl)

    def test_with_scratch_path(self):
        # self.tmp_path is a fresh per-test directory, cleaned up after
        out = self.tmp_path / "data.parquet"
        self.spark.range(10).write.parquet(str(out))

pytest users¶

The module also exposes a spark fixture (session-scoped) so you can skip the class hierarchy if you prefer::

def test_something(spark):
    assert spark.range(5).count() == 5

Design notes¶

Session creation goes through :meth:PyEnv.spark_session with connect=False (force local-only — tests must not accidentally hit a Databricks workspace just because DATABRICKS_HOST happens to be set in the environment). tearDownClass never stops the session.
Because the session is shared, spark_extra_config only takes effect for the first class that triggers creation.
Arrow interop uses spark.sql.execution.arrow.pyspark.enabled=true by default — zero-copy-ish transfer for supported types.

SparkTestCase ¶

Bases: TestCase

Base class for Spark integration tests.

A single global SparkSession is created on first use and shared across every subclass in the process. Each test method also gets a fresh self.tmp_path (pathlib.Path) that is cleaned up in tearDown.

Attributes¶

spark : SparkSession The shared global session. Populated by setUpClass. tmp_path : pathlib.Path Per-test scratch directory. Populated by setUp.

Class attributes¶

spark_app_name : str Spark application name. Only effective for the first class that triggers session creation. spark_extra_config : dict[str, str] Extra Spark config entries. Same caveat — only the first class wins.

is_spark_connect `classmethod` ¶

is_spark_connect() -> bool

True when the shared session is a Spark Connect (Databricks Connect) session rather than a local-JVM one.

udf_supported `classmethod` ¶

udf_supported() -> bool

Whether Python UDFs can execute against the shared session.

Always true on a local-JVM session. On Spark Connect a Python UDF only runs when the client and server share the same minor Python version — otherwise the server rejects it ("client and server should have the same minor Python version"). Probed once (a tiny UDF round-trip) and cached globally, since the session is shared across the whole process.

skip_if_no_udf ¶

skip_if_no_udf(
    reason: str = "Python UDFs need a matching client/server Python minor version on Spark Connect",
) -> None

Skip the current test when Python UDFs can't run on the session (Spark Connect with a client/server Python-version mismatch).

skip_if_spark_connect ¶

skip_if_spark_connect(
    reason: str = "JVM-only API (sparkContext / rdd) is unavailable on Spark Connect",
) -> None

Skip the current test on a Spark Connect session — for tests that reach JVM-only APIs (sparkContext, rdd) absent in Connect.

df ¶

df(data: Iterable[Any], schema: Any = None) -> 'DataFrame'

Shorthand for self.spark.createDataFrame(data, schema).

arrow_to_spark ¶

arrow_to_spark(table: 'pa.Table') -> 'DataFrame'

Convert a pyarrow.Table to a Spark DataFrame via pandas.

Uses Arrow-backed pandas for zero-copy-ish transfer. Good enough for tests; don't use this on gigabyte-scale data.

spark_to_arrow ¶

spark_to_arrow(df: 'DataFrame') -> 'pa.Table'

Materialise a Spark DataFrame as a pyarrow.Table.

assertDataFrameEqual ¶

assertDataFrameEqual(
    actual: "DataFrame",
    expected: "DataFrame | pa.Table | list[dict[str, Any]]",
    *,
    ordered: bool = False,
    check_schema: bool = True
) -> None

Assert two DataFrames are equal.

Parameters¶

actual : DataFrame The DataFrame produced by the code under test. expected : DataFrame | pa.Table | list[dict] The reference value. Accepts a Spark DataFrame, a pyarrow Table, or a list of row dicts (useful for inline literals). ordered : bool, default False If False, both sides are sorted by all columns before compare. check_schema : bool, default True If True, schemas (names + types) must match exactly.

assertSchemaEqual ¶

assertSchemaEqual(
    actual: "DataFrame", expected_fields: list[tuple[str, Any]]
) -> None

Assert a DataFrame has exactly the given (name, dtype) fields.

dtype may be a pyspark.sql.types.DataType instance or its simpleString form ("int", "string", "array<long>", ...).

spark ¶

spark() -> 'SparkSession'

Session-scoped pytest fixture exposing the shared SparkSession.

spark_tmp_path ¶

spark_tmp_path(tmp_path: Path) -> Path

Per-test scratch path — just an alias for clarity in Spark tests.

yggdrasil.spark.tests¶