Skip to content

Preserve columnar batches across observations in exchange producer stages #7324

Closed

Description

@pingzh

What is the problem the feature request solves?

For example, a native Parquet scan followed by observe(count(lit(1)).as("rows")) and a grouped aggregation currently falls back to Spark at CollectMetricsExec. The observation introduces a row boundary between otherwise columnar operators even though it forwards its input unchanged.

Preserve Spark's observation results while allowing Arrow batches to pass through eligible shuffle and broadcast exchange producer stages.

Describe the potential solution

Add a columnar collector that evaluates observation expressions with Spark's existing accumulator over row views of each batch, then forwards the original batch to Comet consumers. Keep Spark's observation discovery, task-completion publication and reset lifecycle.

Enable it only below a materializing exchange and where no intervening consumer can stop reading early. Keep Spark's row collector below limits, joins, partial iterator consumers and nonpositive replacement sampling, and for direct actions or writes without a materializing exchange. Observation counts must match the rows Spark would actually consume.

Cover AQE on/off, repeated execution, empty input, aggregate-state publication, partial consumers and stage fallback. Provide spark.comet.exec.collectMetrics.enabled to disable the columnar collector.

Additional context

Related to #5124. This is a bounded implementation that keeps metric evaluation in Spark and avoids row conversion in eligible producer stages. Native metric-expression evaluation and fusion into a single native execution block remain outside this change.

Willingness to contribute

I can contribute a fix for this feature independently.

Activity

  1. pingzh commented on Oct 29, 2026

    @pingzh
    ContributorAuthor

    take

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions