`yggdrasil.spark.executor`¶

executor ¶

:class:SparkStatementExecutor — minimal :class:StatementExecutor for running SQL through a :class:pyspark.sql.SparkSession.

Spark is synchronous from the SQL submission perspective, so this executor is essentially a thin shim around session.sql(text) that returns a terminal :class:SparkStatementResult. It plugs into :class:StatementBatch like any other executor; the wait phase is a no-op because each result is already terminal.

This module is kept independent of any Databricks-specific code so it can be used standalone — open-source Spark, local PySpark, or composed into :class:SQLEngine alongside a Databricks warehouse executor.

SparkStatementExecutor ¶

SparkStatementExecutor(
    spark_session: Optional["SparkSession"] = None, *args, **kwargs
)

Bases: StatementExecutor[SparkPreparedStatement, SparkStatementResult, SparkStatementBatch]

Run statements through a SparkSession.

The session is resolved in priority order:

The session attached to the incoming :class:SparkPreparedStatement.
The session pinned on this executor (spark_session field).
:meth:PyEnv.spark_session — creates one if necessary.

Lazy resolution means the executor is cheap to construct even in environments where pyspark isn't installed; the import only fires when a statement actually runs.

Singleton-cached for the process lifetime — Spark is a per-JVM singleton already, so two callers asking for a Spark executor share one instance. The pinned spark_session is rebindable in place; it doesn't participate in singleton identity (raw SparkSession objects aren't reliably hashable across processes anyway).

opened `property` ¶

opened: bool

True iff :meth:_acquire has run and :meth:_release hasn't.

closed `property` ¶

closed: bool

Inverse of :attr:opened.

default `classmethod` ¶

default() -> 'SparkStatementExecutor'

Return a default executor — no shared registry, just a fresh handle.

resolve_session ¶

resolve_session(
    statement: Optional[SparkPreparedStatement] = None, *, create: bool = True
) -> Optional["SparkSession"]

Resolve the SparkSession used to execute statement.

Precedence: per-statement → executor-pinned → environment. When create=True (the default) a missing session is materialized via :meth:PyEnv.spark_session; with create=False and no session reachable, None is returned.

has_session ¶

has_session() -> bool

Whether a SparkSession is reachable without creating a new one.

Useful for engines that compose this executor and want to fall back to a different backend when Spark isn't available.

sql ¶

sql(text: str, *, row_limit: Optional[int] = None) -> SparkStatementResult

Shortcut: run a raw SQL string and return the terminal result.

Equivalent to execute(SparkPreparedStatement(text, row_limit=row_limit)) but skips the coercion round-trip.

parallelize ¶

parallelize(
    function: Callable[[IN], OUT],
    inputs: Any,
    *,
    schema: "Any | None" = None,
    byte_size: int = 128 * 1024 * 1024
) -> "SparkDataset"

Distribute function over inputs via Spark executors.

Resolves a :class:SparkSession through :meth:resolve_session, then delegates to :meth:Dataset.parallelize — the session resolution is the only value-add over calling Dataset.parallelize directly (it picks the right Databricks Connect / classic / local session for the caller's context).

open ¶

open() -> 'Disposable'

Acquire the resource and cascade into owned children.

Order:

Run our own :meth:_acquire (subclass body).
Flip :attr:opened to True and mark _self_opened.
For each owned child, in registration order:
If the child is already opened, just :meth:_claim it. It stays self-opened — the existing self-open is what keeps it alive after we let go.
Otherwise, call :meth:open on the child (which recursively cascades into ITS owned children), then clear the child's _self_opened flag so the child knows its open is parent-driven, then :meth:_claim it. Without that flag clear, the eventual :meth:_unclaim would refuse to close — it would see "I'm self-opened, someone explicitly opened me, leave me alone."

Both branches record the child in our per-frame scratch list so :meth:_release knows what to unclaim.

Transactional rollback: if any child's open or claim raises, we walk back through the children we already touched (in reverse), unclaim each, then call our own :meth:_release with committed=False and re-raise the original exception. From the caller's view, the open atomically either succeeded with the whole graph live, or failed with nothing changed.

Not reentrant: raises :class:RuntimeError if already opened. Nesting is expressed via with self: blocks, not via paired :meth:open calls.

commit ¶

commit()

Commit current state

rollback ¶

rollback()

Rollback current state

close ¶

close(force: bool = False) -> None

Release the resource and cascade into owned children.

Order:

If currently held open by an outside parent claim (_claim_count > 0) AND we are not in self-opened state, this is a no-op — the parents that opened us still need us live. (Handled inside :meth:_do_close.)
Walk our scratch list of acquired children in REVERSE registration order; :meth:_unclaim each. A child whose claim count hits zero and isn't otherwise self-opened closes itself.
Run :meth:_before_release, then :meth:_release — with committed reflecting the dirty bit (cleared on exception by __exit__).

Idempotent: no-op when already closed, unless force.

force=True runs teardown even when :attr:closed. Intended for error-recovery paths where subclass state might be inconsistent.

Does NOT touch :attr:depth — the with-stack counter belongs to :meth:__enter__/:meth:__exit__ exclusively. If a caller calls :meth:close inside an active with block, the outer :meth:__exit__ will harmlessly skip the now-no-op close on unwind.

mark_dirty ¶

mark_dirty() -> None

Signal pending mutations — commit on next clean :meth:close.

to_singleton ¶

to_singleton(ttl: Any = ...) -> 'Singleton'

Promote this instance into the per-class _INSTANCES cache.

Hot listing paths (iterdir / _ls / glob) build children with singleton_ttl=False so the bounded cache doesn't fill up with thousands of short-lived entries. When a caller decides one of those children is worth keeping around (handing it to a long-running worker, returning it from an API), :meth:to_singleton registers self into the cache so the next constructor call with the same key collapses to the same instance.

ttl defaults to the subclass's _SINGLETON_TTL (... = no caching, None = process lifetime, or a seconds count). When a different instance is already cached under this key, that pre-existing one wins and is returned unchanged — the cache is the source of truth.

invalidate_singleton ¶

invalidate_singleton(remove_global: bool = True) -> None

Pop self from the per-class _INSTANCES cache.

Mutating ops on a Singleton-cached object (writes, deletes, schema invalidations on a Databricks table, put_object on an :class:S3Path) want to make sure the next caller asking for the same key gets a fresh build rather than collapsing onto this stale handle — that's what remove_global=True (the default) does. The pop is :meth:identity-guarded: only an entry that still points at self is removed, so a concurrent re-construction that already raced past this thread is left alone.

remove_global=False is a no-op. The keyword exists so subclass invalidators (invalidate_singleton, _invalidate_entity_tag_cache, …) can offer the same switch without branching at the call site.

prepare ¶

prepare(statement: 'PS | PreparedStatement | str') -> PS

Coerce statement into this executor's prepared-statement type.

Mirrors :meth:Session.prepare_request_before_send: takes whatever the caller passed (raw string, cross-backend :class:PreparedStatement, already-typed instance) and returns the concrete :attr:_PREPARED_CLASS every downstream hook expects. Subclasses that need to inject per-statement defaults (warehouse routing, catalog binding, SELECT-rewrite for cluster execution) override this — same shape as Session's hook.

send ¶

send(statement: 'PS | PreparedStatement | str', *, start: bool = True) -> SR

Dispatch statement and return its tracking :class:StatementResult.

Mirrors :meth:Session.send. start=True (default) fires the backend submission eagerly — the result comes back in flight (or already terminal for synchronous backends). start=False returns the idled :class:StatementResult whose backend submission is deferred until :meth:StatementResult.start fires.

The returned result is always bound to this executor — every subclass _submit_statement is supposed to thread executor=self through the constructor, but that's easy to forget and downstream code (StatementResult.wait, retry, raise_for_status) needs the back-reference. Setting it here when it's missing makes the contract enforceable from one place instead of audited per backend.

execute ¶

execute(
    statement: "PS | PreparedStatement | str",
    *,
    options: Optional[ExecutionOptions] = None,
    wait: WaitingConfigArg = True,
    raise_error: bool = True,
    start: bool = True
) -> SR

Submit a single statement and optionally wait for completion.

Two ways to pass execution policy:

Per-call kwargs wait / raise_error (ergonomic, matches the previous public API).
An :class:ExecutionOptions via options= (when you want to reuse the same policy across many calls or compose from layered defaults).

The two are merged: options provides the base, kwargs override any field they explicitly set. Unknown kwargs go nowhere — they are not forwarded to the backend. Use a typed :class:PreparedStatement subclass for backend-specific configuration (parameters, byte limits, routing, etc.).

execute_many ¶

execute_many(
    statements: Iterable["PS | PreparedStatement | str"],
    *,
    options: Optional[ExecutionOptions] = None,
    wait: WaitingConfigArg = True,
    raise_error: bool = True,
    parallel: Optional[int] = None,
    **batch_kwargs: Any
) -> SB

Run several statements as a batch and return the populated batch.

Convenience wrapper around :meth:batch: enqueues every statement, submits, and (by default) waits. parallel controls the wait phase only — submission itself is sequential, since most backends either accept fast or reject fast.

**batch_kwargs are forwarded to the batch constructor (e.g. external_paths= for :class:WarehouseStatementBatch).

batch ¶

batch(
    statements: Optional[Iterable["PS | PreparedStatement | str"]] = None,
    *,
    executor: "StatementExecutor | None" = None,
    parallel: Optional[int] = None,
    **kwargs: Any
) -> SB

Construct a batch bound to this executor.

cancel_all ¶

cancel_all() -> None

Best-effort cancel every live result this executor has produced.

yggdrasil.spark.executor¶