[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. A plain UDF that references the accumulator is detected
(its closure is inspected) and rewritten by the analyzer to observe the per-row delta, wherever it
appears as a projected column, in a filter condition, or nested in a larger expression. Works on
classic and Spark Connect, in Scala and PySpark, so code migrates unchanged. If an accumulator UDF
ends up in a position the rule cannot observe, it fails fast rather than silently returning zero
(see below).

How it works (numeric, scalar UDF; ASCII):

  spark.accumulator("bad")  --referenced in a plain @udf-->  parse()
  df.withColumn("v", parse("raw"))

  Analyzer rule InjectObservedAccumulators rewrites the plan to:

     Project [ v = s.value ]                          <- user column, back to scalar
       |
     CollectMetrics(sum(s.delta) AS __oa_metric)      <- delta aggregated (i.e. df.observe)
       |
     Project [ s = udf(raw) => struct(value, delta) ] <- UDF emits value + per-row delta
       |
     <child>

  Executor: the rewritten UDF brackets each call -- reset buffer, run the user body
  (acc.add(..)), read the delta -- so exactly one delta per row is emitted and summed.

  Driver (acc.value):
    classic : a QueryExecutionListener copies the summed metric into a JVM registry.
    connect : the server forwards the metric; the client caches it per session.

Because the value rides through the query plan (not the scheduler's side channel), it is
exactly-once: task retries, speculation, and stage recomputation do not double count. Operator
UDFs (mapInPandas/applyInPandas/mapInArrow/applyInArrow) work the same way, except the
hidden delta column is added to the operator's output and observed directly (no struct wrapper).

// 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("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("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("rows")

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

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

Arbitrary (non-numeric) type via a custom merge — each task folds its partial, the partial is
serialized, gathered with collect_list, and folded on the driver with merge. Works in scalar
and operator UDFs, in both Python and Scala:

# PySpark: accumulate the set of distinct keys seen (merge = set union)
keys = spark.accumulator("keys", set(), merge=lambda a, v: a | {v})

@udf("string")
def f(s):
    keys.add(s)
    return s

df.withColumn("x", f("k")).count()
keys.value          # e.g. {"a", "b", "c"} -- folded on the driver

// Scala: the typed analog of AccumulatorV2 -- spark.accumulator[T](name, zero, merge)
val keys = spark.accumulator[Set[String]]("keys", Set.empty[String], _ ++ _)
val f = udf { (s: String) => keys.add(Set(s)); s }
df.withColumn("x", f($"k")).collect()
keys.value          // Set(...), folded on the driver

A single UDF may reference several accumulators (numeric and/or typed); each is detected and
harvested independently:

total = spark.accumulator("total")
keys = spark.accumulator("keys", set(), merge=lambda a, v: a | {v})

@udf("string")
def f(s):
    total.add(1)
    keys.add(s[0])
    return s

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

acc = spark.accumulator("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(name)
  (numeric) and SparkSession.accumulator[T](name, zero, merge) (typed, the AccumulatorV2 analog);
  ObservedAccumulator (add/value/doubleValue) and TypedObservedAccumulator[T]
  (add/value); a resolution rule InjectObservedAccumulators that detects the accumulator(s)
  (Scala: closure reflection; Python: marker) and rewrites the plan (materialize a
  struct(value, delta*) once, hoist one CollectMetrics aggregate per delta, project the value
  back); and a QueryExecutionListener that harvests the values.
• 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
  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 (numeric sums, exact ints, and -- for typed accumulators -- the collected
  partials the driver folds). Harvesting through the Scala Connect client is a follow-up (see edge cases).

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 hidden struct fields; the rule turns
those into observed aggregates. 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). 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 (and the typed
accumulator[T](name, zero, merge)), with add(...) (inside a UDF) and value/doubleValue (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 (sql/core): plain udf via the built-in rule,
  exactly-once single evaluation, cross-query accumulation, struct-to-scalar rewrite, a
  double-valued accumulator, foreachBatch, a typed custom-merge accumulator, multiple
  accumulators (numeric + typed) in one UDF, detection of an accumulator held indirectly
  (nested collection / object field), evaluation under both codegen (CODEGEN_ONLY) and
  interpreted (NO_CODEGEN), cross-session rejection, an accumulator UDF in a filter
  condition and composed in a larger expression, a fail-fast unobservable position, and a
  cross-check that the observed value matches a classic SparkContext accumulator across
  several cases.
• 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: scalar (row-at-a-time and vectorized pandas/Arrow), operator
  UDFs, numeric and custom-merge, multiple accumulators, exact-Long, cross-session rejection, the
  same cross-check against a classic SparkContext accumulator, filter / composed positions and
  the fail-fast unobservable position, and closure-detection coverage for nested containers,
  captured object attributes, and helper functions (the pure detection logic is also unit-verified
  without Spark).
• Spark Connect (Scala) — an E2E test in UserDefinedFunctionE2ETestSuite covers a numeric and
  a typed accumulator harvested through the Scala Connect client.

Run and passing locally on both classic and Spark Connect.

Module layout

The client-facing API (ObservedAccumulator / TypedObservedAccumulator) lives in sql/api,
shared by the classic and Spark Connect Scala clients. The analyzer rule, harvest listener, struct
wrapper, and JVM registry (InjectObservedAccumulators / ObservedAccumulatorRegistry) stay in
sql/core (they need catalyst). value comes from private[sql] SparkSession hooks
overridden per runtime: classic reads the JVM registry; the Scala Connect client reads a
client-side registry harvested from server responses (Scala Connect client harvest is a follow-up).

Automatic UDF detection

A plain UDF that merely calls acc.add() is recognized and rewritten automatically. Detection is
a bounded (depth- and node-capped), cycle-safe, best-effort walk of the closure graph on both
sides, so an accumulator is found however it is held — a captured variable/global, nested inside a
collection or map at any depth, an attribute of a captured object, or referenced by a helper
function the UDF calls. Framework "gateway" objects (a captured SparkSession/DataFrame/etc.)
are not traversed, and every access is guarded so detection never throws. Missing an accumulator
would silently return the zero value, so the walk errs toward finding too much (over-detection only
injects an observe whose delta is zero).

• Scala: InjectObservedAccumulators reflects over the ScalaUDF's closure and wraps the
  scalar UDF in ObservedAccumulatorStructWrapper, which resets the buffers, evaluates the UDF, and
  emits struct(value, delta*) — one delta field per referenced accumulator (numeric → Double,
  typed → serialized Binary partial). The wrapper has a doGenCode (participates in whole-stage
  codegen) and is non-deterministic.
• PySpark: UserDefinedFunction inspects a scalar UDF's closure the same way and rewrites it to
  emit the equivalent multi-delta struct.

The rewrite fires wherever the UDF occurs in a Project (a column or nested in a larger expression,
e.g. select(udf($"x") + 1)) or a Filter condition (e.g. filter(udf($"x"))) -- the struct is
materialized once in a lower Project so the delta is observed exactly once per input row.

Fail fast, not silent

Detection is a heuristic and the rewrite only reaches certain positions, so the risk is a silent
zero. To avoid that, after rewriting, an accumulator UDF still left in a position the rule cannot
observe -- inside an Aggregate/Join/Window/Sort expression, a higher-order-function lambda,
or nested as an argument to another UDF -- raises a clear error instead of returning the zero value.
(A closure whose accumulator detection missed entirely still cannot be caught this way.)

Cross-session safety

State is scoped per session: the classic JVM registry is keyed by the creating session's UUID (not
the bare name), so two sessions that each create an accumulator named "bad" do not collide; the
Connect client already scopes state per client. The harvest listener captures its registering
session, so a Structured Streaming micro-batch clone (a fresh session that shares the listener)
still attributes to the creating session.

An accumulator is harvested by the session that created it, so reading it while a different
SparkSession is active would silently return the zero value. Both the Scala and Python value
raise instead (skipping the check when no session is active, so ordinary single-session use never
trips).

Types

Any value type is supported. Two paths:

• Numeric — a fast path that avoids serialization: the per-row delta is a SQL number aggregated
  with sum. Python integer accumulators carry an exact Long delta (matching a classic
  LongAccumulator, no precision loss past 2^53); floating accumulators carry a Double. Scala:
  add(Long)/add(Double), value: Long / doubleValue: Double; PySpark value follows the
  zero type.
• Custom / arbitrary types — spark.accumulator(zero, name, merge=fn) (PySpark) /
  spark.accumulator[T](name, zero, merge) (Scala), the analog of a classic AccumulatorV2: each
  task folds its partial (add → merge), the partial is serialized, gathered with collect_list,
  and folded on the driver with merge. Works in scalar UDFs and in all four operator UDFs
  (mapInPandas/applyInPandas/mapInArrow/applyInArrow), on both classic and Connect.

Edge cases / follow-ons

• Higher-order functions (transform/filter/… with an accumulator UDF in the lambda) are a
  follow-up; for now such usage fails fast with a clear error (accumulate in a top-level UDF or
  an operator UDF instead) rather than silently mis-counting.
• Scala Connect client: a follow-up. The server-side analyzer rule is verified for Connect
  (it is the same rule the green PySpark Connect parity suite exercises), and the Scala client-side
  harvest is implemented, but it is not yet verified end-to-end: the connect-client-jvm E2E
  harness spawns its server from assembled jars that do not carry this PR's freshly-built
  sql/core, so the rule is absent on that server and the Scala client harvest path is never
  exercised there. (A classic closure-serialization test confirms detection survives the exact
  serialize/deserialize Connect does to ship a UDF.)

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

Generated-by: Claude Code (Opus 4.8)