Skip to content

fix: align Spark collect aggregate element nullability - #24767

Open
goutamadwant wants to merge 2 commits into
apache:mainfrom
goutamadwant:fix-spark-collect-nullability
Open

fix: align Spark collect aggregate element nullability#24767
goutamadwant wants to merge 2 commits into
apache:mainfrom
goutamadwant:fix-spark-collect-nullability

Conversation

@goutamadwant

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

Spark's collect_list and collect_set drop null inputs and declare arrays with containsNull = false. The DataFusion Spark implementations declared nullable list elements, and their accumulator results could derive nested fields from runtime arrays. The resulting schema mismatch requires downstream casts; when runtime nested fields drift from the declared type, aggregate output validation can fail.

What changes are included in this PR?

  • Declare non-nullable list elements in the return and partial-state types for both aggregates.
  • Normalize partial state and final accumulator output to the declared list field, including nested child types.
  • Reuse matching child arrays without casting, while reconciling a differing nested runtime type only when needed.
  • Preserve the same exact type for empty and all-null results.

Are these changes tested?

Yes. Regression tests cover return and state fields, empty results, primitive values, nested struct nullability, partial-state merging, and both collect_list and collect_set.

The existing Spark aggregate and window SQLLogicTest suites also pass. The full required workspace test suite and workspace-wide clippy with all targets and features pass.

Are there any user-facing changes?

collect_list and collect_set now expose Arrow list schemas with non-nullable elements, matching Spark. There are no Rust API changes.

@github-actions github-actions Bot added the spark label Aug 29, 2026
@codecov-commenter

codecov-commenter commented Aug 29, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 84.15842% with 64 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.52%. Comparing base (c4910e0) to head (ddbadd3).
⚠️ Report is 23 commits behind head on main.

Files with missing lines Patch % Lines
datafusion/spark/src/function/aggregate/collect.rs 82.51% 16 Missing and 48 partials ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main   #24767      +/-   ##
==========================================
+ Coverage   81.47%   81.52%   +0.05%     
==========================================
  Files        1122     1123       +1     
  Lines      403629   406534    +2905     
  Branches   403629   406534    +2905     
==========================================
+ Hits       328866   331440    +2574     
- Misses      55510    55676     +166     
- Partials    19253    19418     +165     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@comphead

Copy link
Copy Markdown
Contributor

Thanks @goutamadwant I'll check it today

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

[P2] Handle encoded-array nullability before rebuilding the list

@goutamadwant The new reconstruction at collect.rs:96–98 panics for valid encoded input even when every collected element is logically non-null.

For example, construct a Dictionary<Int8, Utf8> with keys [0, 0] and dictionary values [Some("a"), None], using its actual datatype as the aggregate input type. The array passes ArrayData::validate_full() and both input rows are "a". Its logical_null_count() is zero, but Arrow 59.2's conservative DictionaryArray::is_nullable() is true because the dictionary retains an unused null entry.

The existing logical-null filtering and single-array concatenation retain that entry. The matching-type fast path then reuses the child array, while .with_field(field) makes the list item non-nullable. ListArray::new rejects the child's is_nullable() and panics with:

InvalidArgumentError("Non-nullable field of ListArray \"item\" cannot contain nulls")

I reproduced this at e3f28e6 through evaluate(), partial state(), and public SQL (SELECT collect_list(x) FROM dictionary_input after registering the batch and UDAF). The identical three probes pass with merge-base c4910e0 production code and return ["a", "a"].

Additional head/base controls confirm the same failure for collect_list over sparse unions and sliced run-end-encoded arrays, and for collect_set over sparse unions. The expanded matrix has 7 head panics versus all 11 cases passing on the merge base; dictionary/run-end-encoded collect_set controls pass, so those are not claimed as affected.

Please reconcile encoded-child nullability with Arrow's construction requirements while preserving Spark's non-nullable element contract, and add regression coverage for both partial state and final output. Merely making the constructor fallible would avoid the panic but still reject these valid inputs.

@goutamadwant

Copy link
Copy Markdown
Contributor Author

[P2] Handle encoded-array nullability before rebuilding the list

@goutamadwant The new reconstruction at collect.rs:96–98 panics for valid encoded input even when every collected element is logically non-null.

For example, construct a Dictionary<Int8, Utf8> with keys [0, 0] and dictionary values [Some("a"), None], using its actual datatype as the aggregate input type. The array passes ArrayData::validate_full() and both input rows are "a". Its logical_null_count() is zero, but Arrow 59.2's conservative DictionaryArray::is_nullable() is true because the dictionary retains an unused null entry.

The existing logical-null filtering and single-array concatenation retain that entry. The matching-type fast path then reuses the child array, while .with_field(field) makes the list item non-nullable. ListArray::new rejects the child's is_nullable() and panics with:

InvalidArgumentError("Non-nullable field of ListArray \"item\" cannot contain nulls")

I reproduced this at e3f28e6 through evaluate(), partial state(), and public SQL (SELECT collect_list(x) FROM dictionary_input after registering the batch and UDAF). The identical three probes pass with merge-base c4910e0 production code and return ["a", "a"].

Additional head/base controls confirm the same failure for collect_list over sparse unions and sliced run-end-encoded arrays, and for collect_set over sparse unions. The expanded matrix has 7 head panics versus all 11 cases passing on the merge base; dictionary/run-end-encoded collect_set controls pass, so those are not claimed as affected.

Please reconcile encoded-child nullability with Arrow's construction requirements while preserving Spark's non-nullable element contract, and add regression coverage for both partial state and final output. Merely making the constructor fallible would avoid the panic but still reject these valid inputs.

@sunchao addressed this in b8932a5d8.

The fix now handles Arrow's conservative encoded-array nullability in SingleRowListArrayBuilder: for a non-nullable list field, it checks the exact logical null count and preserves the encoded child buffers and datatype when there are no logical nulls. This keeps Spark's containsNull = false contract and avoids the panic without decoding the values or making the field nullable.
Let me know if you have any other comments or suggestions!

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

P1

P1-1 - Output array column stays nullable, Spark declares it non-nullable

Where: SparkCollectList / SparkCollectSet in collect.rs do not override
is_nullable().

Problem: AggregateUDFImpl::is_nullable() defaults to true (datafusion/expr/src/udaf.rs:547),
and udaf_default_return_field (udaf.rs:1201) stamps the output field's nullability from
it. So both aggregates report an output field nullable: true. Spark's Collect.nullable
is false (collect.scala:49); the array is never null. The PR fixed the element axis
(containsNull=false) but left the array axis, so the Spark schema fix is half-applied.

Why it matters: Divergence from the Spark schema this PR exists to match. It is a
schema/nullability contract mismatch for the Spark-compat layer, not a wrong-value bug
(the value is never null), so down-rank to P2 if the consumer (Comet) does not validate
field nullability. verdict: CONFIRMED on the divergence; PLAUSIBLE on a downstream
schema-check failure.

Fix: add fn is_nullable(&self) -> bool { false } to both impls. Safe and precedented:
default_value is already a non-null empty list (empty_list_scalar), which satisfies the
is_nullable=false contract noted at udaf.rs:546. Precedent: count.rs:296,
approx_distinct.rs:723, regr.rs:467. Assert it in
collect_types_have_non_nullable_elements (aggregate.is_nullable() is not currently checked).


P2

P2-1 - No .slt for the schema change; existing collect.slt is value-only

Where: datafusion/sqllogictest/test_files/spark/aggregate/collect.slt exists but every
query is query ? / query I? (values only). The PR adds Rust unit tests, no .slt.

Problem: The PR's entire behavioral change (element containsNull=false) is SQL-visible
via arrow_typeof, but no SQL-level test asserts it, and the existing SLT cannot catch it
(unchanged printed values). Checklist item 8 and Testing line 100: SQL-visible changes need
an .slt; Rust unit tests bypass the planner and type plumbing.

Fix: extend collect.slt with element-type assertions, e.g.
SELECT arrow_typeof(collect_list(a)) FROM (VALUES (1)) t(a); and the collect_set variant.
arrow_typeof renders the element field (List(Field { ... nullable: false ... })), so it
pins the change. The outer nullable=false (P1-1) is not visible via arrow_typeof; cover
that in the Rust test instead.

P2-2 - normalize_list_scalar null-list branch is untested

Where: collect.rs, the array.is_null(0) branch (new lines ~55-61, new_null_list).

Problem: This branch is only reachable through state() on an all-null / empty group
(ArrayAggAccumulator::state() -> evaluate() -> new_null_list(..., true, 1), then
normalized). It is real partial-aggregation behavior that flows into shuffle and
merge_batch. empty_results_have_non_nullable_elements calls only evaluate() (which
bypasses normalize via empty_list_scalar), and accumulator_state_and_output_preserve_nested_type
calls state() only after a non-null update_batch. The branch has no coverage.

Fix: add a test calling state() on a fresh (or all-null-updated) accumulator; assert one
null list whose element field is non-nullable and named item, for both distinct and
non-distinct.

P2-3 - New cast can turn a previously-working query into a runtime error (regression risk)

Where: collect.rs, normalize_list_scalar cast (new lines ~64-68):
cast(values.as_ref(), field.data_type()).

Problem: For nested elements (e.g. structs), when the runtime array's nested nullability
is looser than the declared type (arg_types[0]) and the runtime data actually contains
nulls under a field the declared type marks non-nullable, cast -> cast_struct_to_struct
-> StructArray::try_new errors ("Found unmasked nulls for non-nullable field",
../arrow-rs/arrow-array/src/array/struct_array.rs:169). Before this PR (nullable element,
no normalize) the value passed through. DataFusion nullability inference is imprecise across
joins, so this can regress a real query. It fails safe (errors, no corruption) and only when
values.data_type() != field.data_type(), so primitives and matching structs are unaffected.
verdict: PLAUSIBLE.

Fix / coverage: add a collect_list test over a struct whose nested field is declared
non-nullable but whose runtime input carries a null there, and decide the intended behavior
(error vs. widen). accumulator_state_and_output_preserve_nested_type differs in nullability
but uses no actual nulls, so it does not exercise this path.


@github-actions github-actions Bot added sqllogictest SQL Logic Tests (.slt) common Related to common crate labels Aug 30, 2026
@goutamadwant

Copy link
Copy Markdown
Contributor Author

P1

P1-1 - Output array column stays nullable, Spark declares it non-nullable

Where: SparkCollectList / SparkCollectSet in collect.rs do not override is_nullable().

Problem: AggregateUDFImpl::is_nullable() defaults to true (datafusion/expr/src/udaf.rs:547), and udaf_default_return_field (udaf.rs:1201) stamps the output field's nullability from it. So both aggregates report an output field nullable: true. Spark's Collect.nullable is false (collect.scala:49); the array is never null. The PR fixed the element axis (containsNull=false) but left the array axis, so the Spark schema fix is half-applied.

Why it matters: Divergence from the Spark schema this PR exists to match. It is a schema/nullability contract mismatch for the Spark-compat layer, not a wrong-value bug (the value is never null), so down-rank to P2 if the consumer (Comet) does not validate field nullability. verdict: CONFIRMED on the divergence; PLAUSIBLE on a downstream schema-check failure.

Fix: add fn is_nullable(&self) -> bool { false } to both impls. Safe and precedented: default_value is already a non-null empty list (empty_list_scalar), which satisfies the is_nullable=false contract noted at udaf.rs:546. Precedent: count.rs:296, approx_distinct.rs:723, regr.rs:467. Assert it in collect_types_have_non_nullable_elements (aggregate.is_nullable() is not currently checked).

P2

P2-1 - No .slt for the schema change; existing collect.slt is value-only

Where: datafusion/sqllogictest/test_files/spark/aggregate/collect.slt exists but every query is query ? / query I? (values only). The PR adds Rust unit tests, no .slt.

Problem: The PR's entire behavioral change (element containsNull=false) is SQL-visible via arrow_typeof, but no SQL-level test asserts it, and the existing SLT cannot catch it (unchanged printed values). Checklist item 8 and Testing line 100: SQL-visible changes need an .slt; Rust unit tests bypass the planner and type plumbing.

Fix: extend collect.slt with element-type assertions, e.g. SELECT arrow_typeof(collect_list(a)) FROM (VALUES (1)) t(a); and the collect_set variant. arrow_typeof renders the element field (List(Field { ... nullable: false ... })), so it pins the change. The outer nullable=false (P1-1) is not visible via arrow_typeof; cover that in the Rust test instead.

P2-2 - normalize_list_scalar null-list branch is untested

Where: collect.rs, the array.is_null(0) branch (new lines ~55-61, new_null_list).

Problem: This branch is only reachable through state() on an all-null / empty group (ArrayAggAccumulator::state() -> evaluate() -> new_null_list(..., true, 1), then normalized). It is real partial-aggregation behavior that flows into shuffle and merge_batch. empty_results_have_non_nullable_elements calls only evaluate() (which bypasses normalize via empty_list_scalar), and accumulator_state_and_output_preserve_nested_type calls state() only after a non-null update_batch. The branch has no coverage.

Fix: add a test calling state() on a fresh (or all-null-updated) accumulator; assert one null list whose element field is non-nullable and named item, for both distinct and non-distinct.

P2-3 - New cast can turn a previously-working query into a runtime error (regression risk)

Where: collect.rs, normalize_list_scalar cast (new lines ~64-68): cast(values.as_ref(), field.data_type()).

Problem: For nested elements (e.g. structs), when the runtime array's nested nullability is looser than the declared type (arg_types[0]) and the runtime data actually contains nulls under a field the declared type marks non-nullable, cast -> cast_struct_to_struct -> StructArray::try_new errors ("Found unmasked nulls for non-nullable field", ../arrow-rs/arrow-array/src/array/struct_array.rs:169). Before this PR (nullable element, no normalize) the value passed through. DataFusion nullability inference is imprecise across joins, so this can regress a real query. It fails safe (errors, no corruption) and only when values.data_type() != field.data_type(), so primitives and matching structs are unaffected. verdict: PLAUSIBLE.

Fix / coverage: add a collect_list test over a struct whose nested field is declared non-nullable but whose runtime input carries a null there, and decide the intended behavior (error vs. widen). accumulator_state_and_output_preserve_nested_type differs in nullability but uses no actual nulls, so it does not exercise this path.

Addressed in ddbadd3.

  • Both collect_list and collect_set now report a non-nullable outer result field, with return_field coverage.
  • collect.slt now pins List(non-null Int64) through arrow_typeof for both aggregates.
  • Empty partial state is covered for distinct and non-distinct paths, including state merge to the final empty non-null result.
  • Runtime input whose metadata is wider but whose data is valid is normalized to the declared type. An actual unmasked null under a declared non-nullable nested field is rejected fallibly at the accumulator boundary for both aggregates rather than widening the output schema or allowing collect_set to panic. Widening would recreate the AggregateExec schema mismatch this PR fixes.

Mutation checks confirmed each new test fails when its corresponding behavior is disabled. Please let me know if you ahve any other suggestions. thanks!

@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 ddbadd3 against merge base c4910e0. I found three reproducible P2 regressions, detailed inline.

Validation with Rust 1.97 and Arrow 59.2: the 10 existing collect tests pass. All seven additional probes pass with the merge-base production code; five fail on the PR head. The encoded-input failures reproduce through grouped and sliding-window SQL. The nested-schema failures reproduce through accumulator tests using the declared/runtime nullability mismatch this PR explicitly supports. I did not rerun the full workspace suite.

Comment on lines +622 to +625
.add_child_data(arr.to_data())
.build()
.expect("single-row list array should contain valid data");
return GenericListArray::from(data);

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.

[P2] Handle encoded nullability in downstream list operations

This preserves a dictionary child containing unused null entries. The resulting list passes ArrayData::validate_full(), but Arrow's concat and take constructors still reject its non-nullable element field.

I reproduced this with Dictionary<Int8, Utf8> keys [0, 0, 0, 0], dictionary values [Some("a"), None], group keys [0, 0, 1, 1], and row IDs [0, 1, 2, 3]. Both queries fail on this head:

SELECT collect_list(x) FROM t GROUP BY g ORDER BY g;
SELECT collect_list(x) OVER (
  ORDER BY id ROWS BETWEEN 1 PRECEDING AND CURRENT ROW
) FROM t ORDER BY id;

The error is Non-nullable field of ListArray "item" cannot contain nulls. Grouped/window output concatenates list scalars through ScalarValue::iter_to_array, reaching the conservative constructor again. Both queries pass with merge-base production code; the corresponding collect_set controls pass on both versions.

Please reconcile encoded-child nullability beyond this constructor and add grouped/window regression coverage. The existing global-aggregation test produces only one scalar and misses this path.

Comment on lines +258 to +259
let value = self.normalize_input(value)?;
self.inner.update_batch(&[value])

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.

[P2] Normalize retractions consistently with updates

For the declared-nonnullable/runtime-nullable nested Struct scenario covered by this PR, this update path initializes collect_set's RowConverter with the normalized schema. retract_batch still forwards the original runtime array.

Using declared Struct(required: Int32 non-null) and runtime Struct(required: Int32 nullable) containing [1, 2], update_batch succeeds but retracting the first original row now fails:

RowConverter column schema mismatch, expected Struct("required": non-null Int32) got Struct("required": Int32)

The identical update/retract probe passes on the merge base. Sliding windows pass slices of the original input to both operations, so this breaks retraction when the downstream runtime/declared schema mismatch occurs. This reproduction is at the accumulator API; I have not established a native SQL producer of that mismatch.

Please apply equivalent normalization during retraction and add an update/retract round-trip to the nested-schema regression coverage.

Comment on lines +245 to +248
if value.data_type() == field.data_type() {
Ok(Arc::clone(value))
} else {
Ok(cast(value.as_ref(), field.data_type())?)

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.

[P2] Drop ignored null rows before narrowing nested types

This cast processes rows before the inner accumulator drops null inputs. In the runtime/declared nested-nullability mismatch scenario, it can reject the backing payload of a row that both collect aggregates must ignore.

A concrete input has declared element type List<non-null Int32> and runtime type List<nullable Int32>, with offsets [0, 1, 2], child values [NULL, 1], and validity [false, true]. Its logical rows are [NULL, [1]]; the null list legally retains a null backing child, and the input passes ArrayData::validate_full().

Both collect_list and collect_set now fail during update_batch with Non-nullable field of ListArray "item" cannot contain nulls. The identical accumulator probes return [[1]] with merge-base production code. There is no null in the retained nested value.

Please remove ignored null rows before narrowing nested nullability, and cover this masked-payload case for both aggregates.

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

Thanks for the updates. I reviewed ddbadd312cee1dd117c4324d6e10305993fbd0dd and found two remaining issues in downstream compaction and sliding-window retraction, detailed inline.

All 38 current CI checks passed. The CI logs also show all 291 Spark unit tests passing, including the 10 collect regression tests. The two new findings are based on tracing the current DataFusion and Arrow 59.2.0 sources; I have not run their proposed reproducers locally.

.add_child_data(arr.to_data())
.build()
.expect("single-row list array should contain valid data");
return GenericListArray::from(data);

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.

[P2] Handle encoded children during scalar compaction too

The constructor workaround preserves a dictionary child's unused null entry, so the resulting non-nullable list still fails in an ordinary consumer. With the input from collect_list_handles_dictionary_with_unused_null (keys [0, 0], dictionary [Some("a"), None]), the source path for acc.evaluate()?.compacted() on the collect_list wrapper accumulator reaches compact_view_buffers's ListArray::new. Arrow's dictionary copy retains the unused null, so that constructor sees a non-nullable field and a conservatively nullable child and panics.

This also affects composition: ordered array_agg compacts each retained input scalar, and therefore fails if it receives this inner collect_list result. The pre-PR nullable result passed this constructor's nullability check.

Please preserve support for these valid encoded children through compaction as well, and extend the dictionary regression to compact the result or feed it into an ordered aggregate. This finding is source-traced against Arrow 59.2.0; I have not executed the proposed reproducer locally.

let [value] = values else {
return self.inner.update_batch(values);
};
let value = self.normalize_input(value)?;

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.

[P2] Normalize retracted inputs alongside updated inputs

This changes the type received by the inner accumulator, while retract_batch still forwards the original runtime array at lines 283–284. Reuse the valid input from accumulator_state_and_output_preserve_nested_type: declared Struct(required: Int32, nullable=false), runtime child nullable=true, and actual values [1, 2]. After updating a collect_set accumulator, retracting the first row of that same runtime array reaches RowConverter::append with a different nested type. Arrow's schema check compares nested nullability and returns RowConverter column schema mismatch.

Sliding windows pass slices of the original input to both operations, so the valid schema mismatch handled by the update path still causes a query error when rows leave the frame. Please apply the same normalization in retract_batch, and extend the nested-type test with update → retract → evaluate.

This is source-traced; I have not executed the reproducer locally or established that the complete SQL query succeeded on the pre-PR revision.

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

Labels

common Related to common crate spark sqllogictest SQL Logic Tests (.slt)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

bug: SparkCollectList/SparkCollectSet declare nullable list elements, diverging from Spark's containsNull = false

4 participants