Skip to content

fix: align nested collection buffer nullability before spilling - #5903

Open
ErikBPF wants to merge 4 commits into
apache:mainfrom
ErikBPF:investigate/5239-struct-spill
Open

ErikBPF wants to merge 4 commits into
apache:mainfrom
ErikBPF:investigate/5239-struct-spill

Conversation

@ErikBPF

@ErikBPF ErikBPF commented Sep 13, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Closes #5239.

Rationale for this change

Grouped collection aggregation over structs with required fields can fail
when final aggregation spills. Native collection state normalizes nested
fields to nullable, but Comet's intermediate output schema retains the
original required fields. The exchange casts the state back to that strict
schema, and final aggregation subsequently rejects nullable state emitted
for spilling.

What changes are included in this PR?

Use Spark's existing asNullable conversion for the element type of both
CollectList and CollectSet intermediate buffers. Add one regression
covering both functions over structs and arrays of structs, with required
fields created after the Parquet scan, Spark result comparison, native
execution, and a positive aggregate spill assertion for every case.

This fixes Comet's intermediate schema declaration independently of
apache/datafusion#24767, which changes upstream Spark accumulator typing.
When upgrading DataFusion to include that change, re-check compatibility
between its non-nullable collection elements and Comet's intermediate
buffer declaration (containsNull=true), using the forced-spill regression.
That declared-type difference is not a demonstrated failure at this PR's
current dependency pin.

How are these changes tested?

The follow-up extends the same regression with array<struct<required fields>>, keeping native execution, Spark-result comparison, and actual
spilling assertions for each shape. An array-only probe with the original
strict buffer declaration restored failed with a nested list/struct schema
nullability mismatch, confirming that the new case exercises the bug.

With the fix restored, the expanded targeted regression passed on both
Spark 4.1.3 and Spark 3.5.9: one test passed per profile, with zero failures
or cancellations. Both shapes execute both collection functions. The
Spark 3.5 run used a clean full-reactor build, and formatting/style checks
passed in both profiles. These runs used the release native library built
from the current PR source. The full aggregate suite was not rerun for
this test-only extension; its earlier results are recorded below.

The original struct-only regression failed on base
f69c4c81b9429e327ea95658530ae4ed4ed19635 with a strict-versus-nullable
struct schema mismatch during final aggregation, then passed with this
change using the release native library and Spark 4.1.3.

./mvnw -B -Prelease test -Dtest=none \
  '-Dsuites=org.apache.comet.exec.CometAggregateSuite collect_list and collect_set over non-nullable nested fields survive spilling'

On the original PR revision, the complete CometAggregateSuite passed:
99 succeeded, zero failed,
and the two pre-existing metric tests remained ignored. The full Maven
reactor, including Scalastyle and Spotless checks, succeeded.

The original struct-only regression also passed with Spark 3.5.9 / Scala
2.12.18 after a clean
full-reactor build using -Prelease -Pspark-3.5, including style and
formatting checks. Both Spark versions loaded the same release native
library. The full aggregate suite was run on Spark 4.1.3 only.

The first full-suite attempt hit the test unit's 1,024-task limit while
Spark created an AQE thread. Raising that task limit to 4,096 allowed the
suite to finish; the CPU and 32 GiB memory caps were unchanged.

Force final aggregation to spill so nested nullability mismatches cannot hide behind shuffle-only spills.

Refs apache#5239
Match native collection state so final aggregation can spill structs with required fields.

Closes apache#5239
@github-actions github-actions Bot added the bug Something isn't working label Sep 13, 2026
@ErikBPF

ErikBPF commented Sep 13, 2026

Copy link
Copy Markdown
Contributor Author

Extended the forced-spill regression to cover arrays of structs as well as structs. Required fields are constructed after the Parquet scan, so the new case exercises recursive nullability beneath an array. Each shape runs both collect_list and collect_set, compares Spark results, requires native aggregation, and asserts a positive spill count.

The PR description also records the DataFusion #24767 integration boundary: when upgrading that dependency, re-check outer collection-element nullability against Comet's intermediate buffer schema and rerun this regression. This follow-up does not change the current production declaration.

The new array-only case failed with the original strict buffer declaration restored, reporting the expected nested list/struct nullability mismatch. Restoring the fix made the expanded targeted regression pass on both Spark 4.1.3 and Spark 3.5.9, with one test passed per profile and no failures or cancellations. Both profiles passed formatting/style checks; Spark 3.5 used a clean full-reactor build. The native release library was rebuilt from the current PR source.

@rich7420 rich7420 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

@ErikBPF thanks for the patch

@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.

Correctness

Reviewed 6aa8ca8563818400c285640168082e893916b20b against 4abfd95114d61ad454f9ee269be1615f469f24e2. I found no actionable correctness or compatibility issue.

The prior declaration kept required fields inside collection buffers even though the native planner and collection accumulators normalize those fields to nullable. The shuffle reader can cast the partial state back to that stricter declaration. With DataFusion 55.1.0, final aggregation uses its input schema when materializing spill state, so the nullable accumulator output then fails Arrow's schema check. Applying asNullable to both buffer element types closes that mismatch at the declaration.

I compared this with the maintained Spark 3.5 and 4.0 implementations. Their conversion recursively handles structs, array elements, and map values while preserving field metadata. The change is confined to intermediate attributes. Spark's final result type, null-input filtering, empty-array result, list duplicates, and set deduplication remain unchanged. The existing safeguards for mixed Spark/native buffers and collection aggregates with PartialMerge remain in place.

The regression creates required fields after the Parquet scan, exercises both functions over structs and arrays of structs, compares sorted results with Spark, requires native aggregation, and asserts positive aggregate spilling for each shape. The author reports that restoring the strict declaration makes the array case fail and that the expanded test passes on Spark 3.5.9 and 4.1.3. I verified the source and dependency contracts, but did not independently execute those runs. The original reported base failure used f69c4c81, not this review's base. Product CI is awaiting approval and has no jobs, so there is no current CI test result to credit.

Performance

The added work walks the type definition when constructing the intermediate output attributes. It does not add processing per input row, serialization, or a second aggregate implementation. Keeping the declaration aligned also avoids the needless narrowing and subsequent widening at the shuffle boundary. I found no added data-copying path in this change. No benchmark was run, and I am not assigning a measured speedup to it.

Design

The fix belongs in adjustOutputForNativeState, which already translates Spark's intermediate attributes into the native buffer representation. Matching that declaration to the existing producer fixes the contract before the spill boundary. The current DataFusion collection state and Comet buffer both retain nullable outer list elements, so the patch appropriately leaves that separate contract unchanged.

Abstraction & complexity

Reusing Spark's recursive asNullable operation keeps the two collection cases small and consistent with the native planner. The Scala regression is justified by its generated data, memory configuration, native-plan checks, and spill metric assertion. It reuses one fixture for both nested shapes without adding a helper framework. I found no simplification or additional change worth requesting before merge.

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

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Native collect_set(struct) aggregate crashes on spill emit

3 participants