Add native support for CollectMetricsExec (df.observe) to preserve stage fusion #5124

Activity

1. andygrove

   added

   enhancementNew feature or requestNew feature or request

   on Jul 29, 2026
2. github-actions

   added

   requires-triage

   on Jul 29, 2026
3. andygrove

   added a commit that references this issue on Jul 29, 2026

   docs: link tracking issues for operator acceleration gaps

   ...

   5d71736
4. andygrove

   removed

   requires-triage

   on Aug 3, 2026
5. kazantsev-maksim commented on Aug 23, 2026

   

   kazantsev-maksim

   on Aug 23, 2026

   Contributor

   More actions

   take
6. github-actions

   assigned

   kazantsev-maksim

   on Aug 23, 2026
7. kazantsev-maksim commented on Aug 29, 2026

   

   kazantsev-maksim

   on Aug 29, 2026

   Last edited by kazantsev-maksim

   Contributor

   More actions

   Implementation plan for native CollectMetricsExec (df.observe)

   Picking this up. Below is the proposed plan and the constraints I expect to hit, so we can align before implementation starts.

   Goal

   Add a native pass-through operator that forwards batches unchanged (preserving stage fusion - the primary win called out in the issue) while computing the observed aggregate expressions, and report the results back so QueryExecutionListener / Observation events fire with the same values as vanilla Spark.

   Proposed design

   • Native operator (Rust). A pass-through ExecutionPlan that:
     • forwards every input RecordBatch downstream unchanged;
     • feeds each batch into DataFusion Accumulators built from the observed aggregate expressions;
     • on stream exhaustion, finalizes the accumulators and evaluates the metric expressions to produce the final per-partition metric values.
   • JVM operator + serde. CometCollectMetricsExec (serde in QueryPlanSerde) that:
     • serializes the aggregate + metric expressions into the native plan;
     • holds a real Spark AggregatingAccumulator, constructed from the original CollectMetricsExec.metricExpressions / child.output / conf, registered on the driver.
   • Reporting path. Publish the native-computed values back to the JVM through the existing CometMetricNode plumbing (the same channel used by bytes_scanned / output_rows for scan input metrics), then read them in a TaskCompletionListener in CometExecRDD.compute(), mirroring reportScanInputMetrics / reportSpillMetrics. The listener merges the values into the AggregatingAccumulator so driver-side aggregation and Observation work unchanged.

   Constraints we expect to hit (need decisions up front)

   1. Metric discovery is strictly type-based. QueryExecution.observedMetrics → CollectMetricsExec.collect(executedPlan) matches only case collector: CollectMetricsExec. Our node is a CometUnaryExec, and Spark's CollectMetricsExec is a case class (can't subclass), so once we replace the node collect() finds nothing and Observation.get returns empty/default.

      Solution: Keep a thin real CollectMetricsExec node in the executed plan as the accumulator "carrier" (populated by our listener), with the native pass-through underneath. This works with unmodified Spark; the compute subtree still fuses.

   2. Metric value types. The CometMetricNode channel (set(name, v: Long) → SQLMetric) carries only Long. observe allows arbitrary result types (count → Long, but sum(double), min(string), structs do not fit).

      Proposal:

      • Phase 1: support long-typed aggregates only (covers count / null-count - the most common data-quality checks), fall back to Spark otherwise.
      • Phase 2: add a separate channel carrying an InternalRow of metric values instead of reusing SQLMetric.

   3. Cross-partition / retry merge. Each partition finalizes independently; the merge must stay on the driver via the Spark AggregatingAccumulator, not in native code. This keeps us aligned with SPARK-58183 (last-attempt semantics) and SPARK-50007 (default values when the node is pruned).

   4. Fusion vs. fallback behavior (to verify). We need to confirm that with the chosen approach the child subtree stays native (unsupported/carrier node = a boundary) rather than triggering a full fallback of the whole subtree. This depends on how CometExecRule handles the node and needs to be validated during implementation.

   Scope guardrails

   observe already forbids distinct and non-deterministic aggregates, keeping the surface small. Conversion in QueryPlanSerde will be gated on: all metric aggregates supported natively and long-typed (phase 1), else fall back.

   Milestones

   1. Proto + QueryPlanSerde serde + gating (long-only, supported aggregates).
   2. Native pass-through operator computing and finalizing aggregates.
   3. Publish path via CometMetricNode + TaskCompletionListener → AggregatingAccumulator.
   4. Carrier-node discovery approach wired up; verify child stays native.
   5. End-to-end test (df.observe(count(lit(1)))) asserting observation.get matches vanilla Spark, plus multi-partition and pruned-node cases.

   Questions for reviewers

   • Is the carrier-node approach (CollectMetricsExec above the native operator) acceptable?
   • Is a long-only phase 1 acceptable to land first, with full-type support tracked separately?
   • Any concerns with the TaskCompletionListener → accumulator merge under task retries / speculation?

   CC @andygrove