fix: align Spark collect aggregate element nullability - #24767
fix: align Spark collect aggregate element nullability#24767goutamadwant wants to merge 2 commits into
Conversation
Codecov Report❌ Patch coverage is
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. 🚀 New features to boost your workflow:
|
|
Thanks @goutamadwant I'll check it today |
sunchao
left a comment
There was a problem hiding this comment.
[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 The fix now handles Arrow's conservative encoded-array nullability in |
comphead
left a comment
There was a problem hiding this comment.
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
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
left a comment
There was a problem hiding this comment.
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.
| .add_child_data(arr.to_data()) | ||
| .build() | ||
| .expect("single-row list array should contain valid data"); | ||
| return GenericListArray::from(data); |
There was a problem hiding this comment.
[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.
| let value = self.normalize_input(value)?; | ||
| self.inner.update_batch(&[value]) |
There was a problem hiding this comment.
[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.
| if value.data_type() == field.data_type() { | ||
| Ok(Arc::clone(value)) | ||
| } else { | ||
| Ok(cast(value.as_ref(), field.data_type())?) |
There was a problem hiding this comment.
[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
left a comment
There was a problem hiding this comment.
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); |
There was a problem hiding this comment.
[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)?; |
There was a problem hiding this comment.
[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.
Which issue does this PR close?
SparkCollectList/SparkCollectSetdeclare nullable list elements, diverging from Spark'scontainsNull = false#24741.Rationale for this change
Spark's
collect_listandcollect_setdrop null inputs and declare arrays withcontainsNull = 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?
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_listandcollect_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_listandcollect_setnow expose Arrow list schemas with non-nullable elements, matching Spark. There are no Rust API changes.