Skip to content

yggdrasil.spark.ops

ops

Set-style operations on Spark DataFrames — mirror of :mod:yggdrasil.arrow.ops.

Exposes :func:dedup_spark_dataframe, :func:resample_spark_dataframe, and :func:fill_spark_dataframe. The signatures match their arrow counterparts (same parameter names, same defaults, same vocabulary for fill_strategy) so :class:yggdrasil.data.options.CastOptions can route a read through whichever engine it ends up holding without branching at the call site.

Distribution model

Two execution paths, picked by what the caller passes:

  • partition-by present — the op routes through df.groupBy(*partition_by).applyInArrow(...), delegating each group's work to :mod:yggdrasil.arrow.ops running inside the executor. Spark owns the shuffle; pyarrow owns the per-group algorithm. This is the canonical path for the entity-keyed cases the resample / dedup contract is designed for (one timeline per symbol, one row per tenant, …).

  • partition-by absent — Spark SQL window functions over an empty Window.partitionBy() (one global partition). Cheaper to plan than the grouped-arrow path on small inputs but funnels everything to one executor — fine when the caller really wants a single global timeline and the frame is bounded.

PySpark 4.0+ is assumed (the runtime pin in pyproject.toml); applyInArrow was added in 4.0 and is the natural arrow-typed group map.

dedup_spark_dataframe

dedup_spark_dataframe(
    df: "SparkDataFrame", keys: Sequence[str]
) -> "SparkDataFrame"

Drop duplicate rows on keys, keeping the first occurrence per group.

Mirrors :func:yggdrasil.arrow.ops.dedup_arrow_table. When keys is empty the input is returned unchanged so callers can call this unconditionally on every read pass — matching the arrow op's "zero-cost when the target schema has no unique columns" short-circuit.

"First occurrence" is anchored on :func:pyspark.sql.functions.monotonically_increasing_id, which preserves Spark's internal per-partition ordering. That order isn't a contract across shuffles, but it's the strongest "first row in input order" a distributed engine can offer without a user-supplied sort key — and it's stable across plan evaluations of the same job, which is what callers actually rely on.

Spark's dropDuplicates would be a one-liner but it picks an arbitrary row per duplicate group; the row_number() OVER (...) = 1 form keeps the contract aligned with the arrow path's "first non-null, first occurrence" answer.

resample_spark_dataframe

resample_spark_dataframe(
    df: "SparkDataFrame",
    *,
    time_column: str,
    sampling_seconds: int,
    partition_by: "Sequence[str] | None" = None,
    fill_strategy: "str | None" = "ffill"
) -> "SparkDataFrame"

Align df to a fixed sampling grid on time_column.

Mirrors :func:yggdrasil.arrow.ops.resample_arrow_table — timestamps are floored to the largest multiple of sampling_seconds <= the original, rows sharing a bucket collapse via "first occurrence", and fill_strategy runs the same ffill / bfill pass on the resampled output (per partition when partition_by is set, globally otherwise).

Distribution:

  • partition_by is non-empty → groupBy(*partition_by) .applyInArrow(...) delegates the bucket collapse + ffill to :func:yggdrasil.arrow.ops.resample_arrow_table running on each group's :class:pa.Table inside the executor. Spark owns the shuffle; pyarrow owns the per-group algorithm. This is the canonical path.
  • partition_by is empty → SQL window functions on a single global partition. Funnels the whole frame to one executor; use only when "one global timeline" is genuinely what the caller wants.

Short-circuits the same way as the arrow op:

  • sampling_seconds <= 0 → input returned unchanged.
  • time_column missing from the schema → input returned unchanged.
  • the column isn't a timestamp → input returned unchanged.

fill_spark_dataframe

fill_spark_dataframe(
    df: "SparkDataFrame",
    *,
    sort_by: "str | None" = None,
    partition_by: "Sequence[str] | None" = None,
    fill_strategy: "str | None" = "ffill",
    fill_columns: "Sequence[str] | None" = None
) -> "SparkDataFrame"

Forward / backward fill nulls per partition on a Spark frame.

Mirrors :func:yggdrasil.arrow.ops.fill_arrow_table — same semantics, same vocabulary for fill_strategy, same "nested types are skipped, partition boundaries are honored" rules. Uses :func:pyspark.sql.functions.last / first with ignorenulls=True over a partition-and-time window so the fill is fully push-down inside Spark's Catalyst optimiser — no arrow round-trip needed for the flat null-propagation case.