[WIP][POC][SQL][PYTHON] Observe-backed accumulator (SparkSession.accumulator)

[HyukjinKwon]

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 matches a classic accumulator — spark.accumulator(...), acc.add(...)
inside a UDF, acc.value on the driver — with no wrapper API. A plain UDF that references the
accumulator is detected (its closure is inspected) and rewritten by the analyzer to observe the
per-row delta. Works on classic and Spark Connect, so code migrates unchanged.

// Scala: a plain UDF that just calls acc.add()
val acc = spark.accumulator("bad")
val parse = udf { (s: String) => try s.toDouble catch { case _: Throwable => acc.add(); null } }
df.withColumn("v", parse($"raw")).write.parquet(path)  // records `acc` automatically
acc.value        // Long, cumulative across queries

# PySpark (classic and Spark Connect): a plain @udf
acc = spark.accumulator(0, "bad")

@udf("double")
def parse(s):
    try:
        return float(s)
    except ValueError:
        acc.add(1)
        return None

df.withColumn("v", parse("raw")).count()
acc.value

More examples

Double-valued accumulator (add(Double) → doubleValue):

val total = spark.accumulator("total")
val f = udf { (x: Long) => total.add(x.toDouble * 0.5); x }
df.withColumn("y", f($"x")).collect()
total.doubleValue     // Double; total.value gives the rounded Long

Vectorized pandas_udf (call acc.add once per batch):

acc = spark.accumulator(0, "bad")

@pandas_udf("double")
def parse(s: pd.Series) -> pd.Series:
    parsed = pd.to_numeric(s, errors="coerce")
    acc.add(int(parsed.isna().sum()))   # bad rows in this batch
    return parsed

df.withColumn("v", parse("raw")).count()
acc.value

mapInPandas and applyInPandas (operator UDFs):

rows = spark.accumulator(0, "rows")

def count_rows(it):
    for pdf in it:
        rows.add(len(pdf))
        yield pdf

df.mapInPandas(count_rows, df.schema).count()

def head1(pdf):
    rows.add(len(pdf))
    return pdf.head(1)

df.groupBy("k").applyInPandas(head1, df.schema).count()
rows.value

Structured Streaming via foreachBatch (value accumulates across micro-batches):

acc = spark.accumulator(0, "bad")   # parse() defined as above

stream.writeStream.foreachBatch(
    lambda batch, _id: batch.withColumn("v", parse("raw")).write.mode("append").save(path)
).start()
# acc.value grows per micro-batch, like a classic accumulator

Included:

• Scala (sql/core/.../ObservedAccumulator.scala + sql/api): SparkSession.accumulator
  (implicit extension); ObservedAccumulator with add/value/doubleValue; a resolution rule
  InjectObservedAccumulators that detects the accumulator (Scala: closure reflection; Python:
  marker) and rewrites the plan (materialize the struct once, hoist CollectMetrics(sum(delta)),
  project the value back); and a QueryExecutionListener that harvests the value.
• The rule is a built-in session resolution rule (BaseSessionStateBuilder, post-hoc), so no
  spark.sql.extensions wiring. It matches both ScalaUDF and PythonUDF (it runs on the JVM
  logical plan, which PySpark builds too) and reads value/delta by struct ordinal, so it is
  language-agnostic. It must run during analysis (not the optimizer) because it changes the
  column's output type from a struct to the scalar value, and a DataFrame's schema comes from the
  analyzed plan.
• PySpark: spark.accumulator on classic and Connect sessions. The UDF is tagged so the JVM
  rule fires for df.withColumn; acc.value reads the harvested value over py4j (classic) or
  from the Connect client registry (is_remote() selects the source).
• Spark Connect (PySpark): the rule runs server-side, so it fires for Connect plans too. The
  Connect server forwards observed metrics from server-injected CollectMetrics nodes (keyed by
  __oa_metric_<name>, no client Observation), and the Python Connect client captures them into a
  per-client registry (mirroring the existing __python_accumulator__ path).

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 (value and 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. 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): DataFrame/UDF-scoped. It cannot back add() inside arbitrary
RDD closures (no plan node for observe to attach to), nor merge semantics not expressible as a SQL
aggregate. 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 Scala and PySpark. No existing behavior
changes. As a draft, the API shape is open for discussion.

How was this patch tested?

• Scala: ObservedAccumulatorSuite (explicit observe path) and ObservedAccumulatorSeamlessSuite
  (all plain udf, via the built-in rule) — pass locally. Cover correctness, exactly-once single
  evaluation, cross-query accumulation, struct-to-scalar rewrite, a double-valued accumulator,
  and foreachBatch.
• PySpark: python/pyspark/sql/tests/test_observed_accumulator.py plus a Connect parity suite
  (tests/connect/test_parity_observed_accumulator.py), registered in
  dev/sparktestsupport/modules.py — plain @udf seamless, cumulative, double-valued, and closure
  detection. The Python closure detection logic is also unit-verified without Spark.

Module layout

The client-facing API (ObservedAccumulator: add/value/doubleValue) lives in sql/api,
shared by the classic and Spark Connect Scala clients. The analyzer rule, harvest listener, and
JVM registry (InjectObservedAccumulators / ObservedAccumulatorRegistry) stay in sql/core
(they need catalyst). value's seamless component comes from a private[sql] SparkSession hook
overridden per runtime: classic reads the JVM registry; the Scala Connect client reads a
client-side registry harvested from server responses. The harvest listener registers lazily and
idempotently on first spark.accumulator(...).

100% seamless (no wrapper)

A plain UDF that merely calls acc.add() is recognized and rewritten automatically — no wrapper
or special decorator:

• Scala: InjectObservedAccumulators reflects over the ScalaUDF's closure
  (ClosureCleaner-style, following $outer/nested lambdas) for a captured ObservedAccumulator,
  and wraps the scalar UDF in ObservedAccumulatorStructWrapper (resets the buffer, evaluates the
  UDF, reads the delta → struct(value, delta); CodegenFallback, non-deterministic).
• PySpark: UserDefinedFunction.__init__ inspects a scalar UDF's __closure__/__globals__
  for a captured ObservedAccumulator and rewrites it the same way.

Streaming

Works inside foreachBatch, matching classic accumulators — the per-micro-batch batchDF is an
ordinary batch DataFrame, so the built-in rule fires and the listener harvests per batch
(cumulative across batches). Runtime-verified on classic.

UDF coverage

Seamless support (a plain UDF that references the accumulator is detected via closure inspection
and rewritten — no wrapper) spans:

• Row-at-a-time scalar udf, plain and Arrow-optimized — classic runtime-verified.
• Vectorized scalar pandas_udf / Arrow scalar udf — the transform emits a struct batch with
  the batch-total delta on its first row (observe sum totals it); call acc.add(count) once per
  batch.
• Operator UDFs mapInPandas / mapInArrow / applyInPandas / applyInArrow — these are
  separate operators the projection rule can't touch, so a hidden __oa_delta column is added to
  the operator's output and an observe node is injected in Python and then dropped (harvested by
  name/key exactly like the scalar path — no extra JVM change). Hooked on both classic and
  Connect (classic mixins + Connect _map_partitions/GroupedData). applyIn* handles both the
  single-DataFrame/Table and the iterator forms. (See edge cases below.)

Classic ⇄ Connect parity: the seamless hooks are wired in both the classic and Connect Python
UDF classes and operator methods, and acc.value reads the classic JVM registry or the Connect
client registry via is_remote() — so code migrates classic ⇄ Connect unchanged. A Connect parity
test runs the shared suite on a Connect session.

Everything except the row-at-a-time scalar classic path is CI/real-build-verified only (needs
pandas/pyarrow, a full build, and a Connect harness); the Python closure detection + schema helpers
are unit-verified.

Types

• Numeric Long/Double — the SQL sum fast path: add(Long)/add(Double); value: Long
  and doubleValue: Double; PySpark value follows the zero type.
• Custom / arbitrary merges — spark.accumulator(zero, name, merge=fn): each task folds its
  partial (add → merge), the partial is serialized, gathered with collect_list, and folded on
  the driver with merge (bounded by #partitions, not #rows — the general AccumulatorV2 pattern).
  Implemented for PySpark pandas mapInPandas/applyInPandas. Arrow-custom and scalar-custom
  (no partition boundary in a scalar UDF) are follow-ups.

Edge cases / follow-ons

• mapIn*: numeric delta added strictly after the final output batch, or in a partition that emits
  zero rows, is not captured (documented; rare).
• Custom (arbitrary-merge) accumulators: pandas operators only so far (Arrow + scalar to follow).
• value blocks on query completion on classic (drains the listener bus); Connect harvests
  synchronously.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Opus 4.8)