kazantsev-maksim commented on issue #5124:
URL:
https://github.com/apache/datafusion-comet/issues/5124#issuecomment-5463380477
**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?
**One honest note for us (not for the comment):** constraints #1 (type-based
discovery) and #2 (SQLMetric is Long-only) are confirmed from the Spark source
and the CometMetricNode.scala you shared. Constraint #4 (child stays native
with the carrier-node approach) I've flagged as "to verify" rather than
asserting it — we should confirm it against CometExecRule / planner.rs before
committing to the approach.
--
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]
