Skip to content

feat: preserve columnar batches across observations - #7325

Open
pingzh wants to merge 5 commits into
apache:mainfrom
pingzh:pingzh-columnar-observations
Open

pingzh wants to merge 5 commits into
apache:mainfrom
pingzh:pingzh-columnar-observations

Conversation

@pingzh

@pingzh pingzh commented Oct 29, 2026 •

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Closes #7324. Related to #5124; metric evaluation remains in Spark, so this does not implement the fully native collector proposed there.

Rationale for this change

A Parquet scan followed by observe(count(lit(1)).as("rows")) and a grouped aggregation currently introduces a row boundary at the observation. Preserve the observation results while passing Arrow batches between eligible native operators.

What changes are included in this PR?

  • Add CometCollectMetricsExec, enabled by spark.comet.exec.collectMetrics.enabled. Spark evaluates the metric expressions over batch row views and retains ownership of accumulator publication, observation discovery and resets. The original Arrow batch then feeds the downstream native block.
  • Require a materializing shuffle or broadcast exchange above the collector. Retain Spark collection below limits, joins, partial iterator consumers and nonpositive replacement sampling, and for direct actions or writes without such an exchange, so counts match Spark's row consumption.
  • Integrate native input boundaries, AQE revisits and stage reversion. Adopt the first completed accumulator state without copying large aggregate buffers, matching Spark.
  • Keep Arrow reader vectors independent from batches forwarded through observations, retaining buffers without copying data. Convert rows to UnsafeRows when metrics reference complex values, strings or calendar intervals, giving aggregate keys a consistent representation and keeping serialized expression caches independent of Arrow batches. Numeric-only metrics retain the direct row-view path. Preserve widened nested nullability in the observation scan schema during initial planning and AQE replanning, retaining existing accumulators.
  • Document the supported plan shapes and register CometCollectMetricsSuite in Linux and macOS CI, with Spark-version-specific aggregate test helpers.

How are these changes tested?

The focused suite compares output rows and observations with Spark, covering AQE on/off, empty inputs, repeated execution, aggregate-state publication, limits, joins, iterator consumers, sampling, stage reversion and disabled configuration. It also covers mixed struct representations in mode, JSON expression caches at task serialization, multi-batch structs, direct column closure, collect_list/collect_set over structs containing calendar intervals, and null padding in array_insert across the observation boundary, with AQE off/on and explicit replanning that checks collector identity and native projection. The null-padding regression reproduced the Arrow failure before both the initial-conversion fix and the replanning fix. The partial iterator tests include a trailing shuffle to exercise the row-consumption guard.

  • Spark 4.1.3: CometCollectMetricsSuite passed all 23 tests, including the two new complex-value regressions.
  • Spark 3.5.9 with -Pstrict-warnings: CometCollectMetricsSuite passed all 23 tests; main and test sources compiled successfully.
  • Tests used a previously built Apache Comet 1.3 native library with unchanged execution JNI signatures. CI will build the native library from this PR.
  • Spotless, Scala style checks, the full-tree syntactic Scalafix 0.14.6 CI check, Prettier for the changed Markdown, CI configuration validation and git diff --check passed. Two independent reviews approved the fixes.

Spark 4.1's optional strict-warning profile reports four preexisting errors in CometShuffleExchangeExec and CometArrowPythonRunnerBase; its normal profile passes. A maintainer needs to apply run-all-spark-profiles and run-spark-4.1-tests for broader version and Spark SQL coverage; the author account cannot add these labels.

@github-actions github-actions Bot added the enhancement New feature or request label Oct 29, 2026

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I reproduced two query failures on Spark 4.1.3 with the columnar collector enabled. One loses a struct’s child vectors between batches, and the other fails when collect_list copies a struct containing a calendar interval. Both queries succeed with only the collector disabled. I also left a test change that makes the partial-consumption guard observable.

The new suite passed all 17 tests locally using this commit’s JVM code and a previously built Comet 1.3.0 native library. The suite-registration and CI-configuration checks passed as well.

@ajsquared ajsquared left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 review finding.

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I reproduced two observation regressions on Spark 4.1.3 at fd08eebc. Both queries succeed on the pre-PR base and with spark.comet.exec.collectMetrics.enabled=false on this head.

@ajsquared ajsquared left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed the current changes and feedback; no independently confirmed unresolved P1 remains.

@pingzh
pingzh force-pushed the pingzh-columnar-observations branch from b750c0f to 96c5371 Compare October 30, 2026 14:22

@ajsquared ajsquared left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed the current changes and feedback; no independently confirmed unresolved P1 remains.

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed the current head and prior discussion. Focused Spark 4.1 collector CI and an Arrow ownership probe passed; local full Comet execution was unavailable, and Spark SQL, Iceberg and macOS CI jobs were skipped.

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

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Preserve columnar batches across observations in exchange producer stages

3 participants