HyukjinKwon opened a new pull request, #57990:
URL: https://github.com/apache/spark/pull/57990
### What changes were proposed in this pull request?
This is a **[WIP] draft / proof-of-concept** (no JIRA yet) exploring an
*observe-backed
accumulator*: an accumulator whose value is carried through the query plan
and aggregated by a
`CollectMetrics` (`df.observe`) node instead of the scheduler's accumulator
side channel.
The user-facing surface mirrors a classic accumulator:
```scala
// Scala
val acc = spark.accumulator("bad")
val parse = acc.udf { (s: String) => try s.toDouble catch { case _:
Throwable => acc.add(); null } }
df.withColumn("v", parse($"raw")) // seamless: analyzer rule auto-inserts
the observe node
acc.value // Long, cumulative across queries
```
```python
# PySpark (classic and Spark Connect)
acc = spark.accumulator(0, "bad")
@acc.udf("double")
def parse(s):
try: return float(s)
except ValueError: acc.add(1); return None
parse.apply(df, "raw", out="v").count()
acc.value
```
Included:
- **Scala** (`sql/core/.../ObservedAccumulator.scala`):
`SparkSession.accumulator` (as an
implicit extension for now), `ObservedAccumulator` / `ObservedUDF`, a
resolution rule
`InjectObservedAccumulators` that detects a marker on the accumulator UDF
and rewrites the plan
to materialize the struct once, hoist `CollectMetrics(sum(delta))`, and
project the value back
(so `df.withColumn` is seamless), and a `QueryExecutionListener` that
harvests the value.
- **PySpark**: `spark.accumulator` on classic and Connect `SparkSession`,
implemented over the
public `observe` API (works in both).
- Design note: `OBSERVED_ACCUMULATOR_DESIGN.md`.
Because Catalyst cannot see the `add()` calls inside an opaque UDF body, the
UDF leaves a
detectable marker (its name) and emits its per-row delta as a hidden struct
field; the rule
turns that into an observed aggregate. The UDF is marked non-deterministic
so the optimizer
cannot duplicate it (the value field and the aggregated delta must come from
a single
evaluation per row).
### Why are the changes needed?
Classic accumulators updated inside **transformations** have no exactly-once
guarantee: task
retries, speculation, and stage recomputation can apply `add()` more than
once (Spark only
guarantees accumulator values for updates inside actions). An observe-backed
value is derived
from the rows that actually survive at the observe point, so it is
**exactly-once by
construction**.
**Scope / limitation (by design):** this is DataFrame/UDF-scoped. It cannot
back `add()` inside
arbitrary RDD closures (`rdd.foreach`, `mapPartitions`) because there is no
plan node or column
for `observe` to attach to, nor merge semantics that are not expressible as
a SQL aggregate. It
is therefore an API-compatible, more-correct accumulator for the
DataFrame/UDF path, not a
universal replacement for `SparkContext` accumulators.
### Does this PR introduce _any_ user-facing change?
Yes — a new API: `SparkSession.accumulator(...)` returning an
`ObservedAccumulator` with
`add(...)` (inside a UDF) and `value` (on the driver), in both Scala and
PySpark. No existing
behavior changes. As a draft, the API shape (naming, always-on wiring,
typing) is still open for
discussion.
### How was this patch tested?
- Scala: `ObservedAccumulatorSuite` (explicit `.apply` path) and
`ObservedAccumulatorSeamlessSuite`
(seamless `withColumn` via the analyzer rule) — 4/4 pass locally. They
cover correctness,
exactly-once single evaluation, cross-query accumulation, and
struct-to-scalar rewrite.
- PySpark: `python/pyspark/sql/tests/test_observed_accumulator.py`,
registered in
`dev/sparktestsupport/modules.py`. The implementation was additionally
verified end-to-end,
including surviving a forced task retry.
Known follow-ons before this could be non-draft: make seamless `value` block
on query completion
(it currently reads an async-listener-fed registry), wire the rule on by
default instead of via
`spark.sql.extensions`, support types beyond `Long`, and wire the seamless
path for Python UDFs.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
