fix(arrow-select): preserve nullability for REE and Union take by yongster · Pull Request #10994 · apache/arrow-rs

yongster · 2026-09-04T17:46:12Z

take can introduce output nulls when its indices contain nulls.

Most Arrow arrays represent these with a top-level validity bitmap. However,
RunEndEncoded and Union arrays derive logical nullability from their child
arrays:

A RunEndEncoded array derives nullability from its values field.
A Dense Union represents a null through a nullable selected child.
A Sparse Union stores every child at every output position.

Previously, take could write nulls into children whose corresponding field
metadata was marked as non-nullable. This made the physical output inconsistent
with its declared schema.

Closes Define nullability semantics for kernels on REE and Union arrays #10992.
Return a compute error when null take indices would introduce nulls into a
RunEndEncoded array with a non-nullable values field.
For Dense Union arrays, select a nullable child to represent null take
indices.
Return a compute error for Sparse Union arrays with non-nullable fields when
take indices contain nulls, since every child would otherwise receive an
introduced null.
Add regression tests for each behavior.

This complements #10909 and ensures that arrays marked as non-nullable do not
produce output containing newly introduced nulls.

yongster · 2026-09-04T17:48:14Z

I also reviewed filter, interleave, concat, sort, and zip.

I did not identify the same metadata-consistency issue in these kernels: they
do not introduce nulls through null take indices in the way take does for
RunEndEncoded and Union arrays. This PR therefore intentionally limits its
scope to take.

Rich-T-kid

Looking at the arrow-spec, this looks mostly correct.

Rich-T-kid · 2026-09-04T18:49:12Z

            }
        }
        DataType::Union(fields, UnionMode::Sparse) => {
+            if indices.null_count() > 0 && fields.iter().any(|(_, field)| !field.is_nullable()) {


this seems to introduce a bug

#[test] fn test_take_union_builder_null_index_regression() { let mut builder = UnionBuilder::new_dense(); builder.append::<Int32Type>("a", 10).unwrap(); builder.append_null::<Int32Type>("a").unwrap(); let union = builder.build().unwrap(); // UnionBuilder currently declares union fields as non-nullable. let field = union.fields().iter().next().unwrap().1; assert!(!field.is_nullable()); // But it still represents logical nulls through the selected child. assert!(union.logical_nulls().unwrap().is_null(1)); let indices = UInt32Array::from(vec![Some(0), None, Some(1)]); // This should preserve take's normal null-index contract: // a null index produces a logical null in the output. let taken = take(&union, &indices, None).unwrap(); let taken = taken.as_union(); let logical_nulls = taken.logical_nulls().unwrap(); assert!(logical_nulls.is_valid(0)); assert!(logical_nulls.is_null(1)); assert!(logical_nulls.is_null(2)); }

with this PR at the take() call would cause an error. is this intended?

i think this is an existing issue with union builders?

UnionBuilder produces incorrect Union DataType #1637

ah, your right. interesting bug

Yes, intended. take follows the declared field nullability. UnionBuilder marking children as non-nullable after append_null is the existing #1637 issue and is out of scope here.

Jefffrey · 2026-09-05T00:59:12Z

            }
        }
        DataType::Union(fields, UnionMode::Sparse) => {
+            if indices.null_count() > 0 && fields.iter().any(|(_, field)| !field.is_nullable()) {


technically this might be too strict? as long as we have one child that is nullable we can select that type id, then for the rest we fill in nulls or some default; it shouldnt matter since they arent the selected child, even if their values arent null

Good call — that was stricter than dense. 4827226 now matches dense: a single nullable child is enough to represent a null index. Unselected children are taken with dummy indices so non-nullable fields do not get introduced nulls.

A sparse union with no nullable child still errors, same as dense.

Jefffrey · 2026-09-05T01:31:26Z

from a codex review it identified 2 edge cases:

for structs with non-nullable ree/union, currently we reject these because we try to take on them, even though we can mask them with the structs own null buffer (overly strict)
if we have nested unions where inner union has no nullable children, but as a field in the outer union it is marked as nullable, we can incorrect select it to be the null child (e.g. for dense unions)

(these are pretty extreme edge cases honestly, i wouldnt mind them being a separate issue; just wanted to point out for completeness)

Null take indices select a nullable child, matching dense unions. Other children are taken with dummy indices so non-nullable fields do not receive introduced nulls.

yongster · 2026-09-07T13:51:22Z

@Jefffrey I relaxed the sparse path as you suggested: mixed-nullability sparse unions now succeed when at least one child is nullable.

Agreed on the two Codex edge cases (struct-masked non-nullable REE/Union, nested union chosen as the null child). I'll leave those as a follow-up rather than expanding this PR.

Jefffrey · 2026-09-08T12:29:24Z

-        .next()
+        .find(|(_, field)| field.is_nullable())
        .map(|(type_id, _)| type_id)
        .ok_or_else(|| {


could use find_map here

Jefffrey · 2026-09-08T12:31:53Z

+    if values.is_empty() {
+        // Null indices would otherwise take dummy index 0, which is OOB.
+        return non_null_unspecified_values(values.data_type(), indices.len());
+    }


what if indices has a non-null index beyond bounds? or is that not a concern since it should already error for that case earlier?

Jefffrey · 2026-09-08T12:34:21Z

+    let dummy = IndexType::Native::from_usize(0).unwrap();
+    let mut values = indices.values().to_vec();
+    if let Some(nulls) = indices.nulls() {
+        for (idx, value) in values.iter_mut().enumerate() {
+            if nulls.is_null(idx) {
+                *value = dummy;
+            }
+        }
+    }
+    PrimitiveArray::new(ScalarBuffer::from(values), None)


Suggested change

let dummy = IndexType::Native::from_usize(0).unwrap();

let mut values = indices.values().to_vec();

if let Some(nulls) = indices.nulls() {

for (idx, value) in values.iter_mut().enumerate() {

if nulls.is_null(idx) {

*value = dummy;

}

}

}

PrimitiveArray::new(ScalarBuffer::from(values), None)

let dummy = IndexType::Native::ZERO;

let normalized = indices.iter().map(|idx| idx.unwrap_or(dummy));

PrimitiveArray::from_iter_values(normalized)

Jefffrey · 2026-09-08T12:35:05Z

+/// Builds an all-valid array of `len` whose values are unspecified.
+///
+/// Used for unused sparse-union child slots when the source child is empty.
+fn non_null_unspecified_values(data_type: &DataType, len: usize) -> Result<ArrayRef, ArrowError> {


personally id just inline this; that way we dont need a doc comment to explain for what use case this is for

fix(arrow-select): preserve nullability for REE and Union take

e7412ed

github-actions Bot added arrow Changes to the arrow crate arrow-select labels Sep 4, 2026

Rich-T-kid reviewed Sep 4, 2026

View reviewed changes

Jefffrey added the bug label Sep 5, 2026

Jefffrey reviewed Sep 5, 2026

View reviewed changes

Jefffrey mentioned this pull request Sep 5, 2026

fix(arrow-select): preserve nullif nulls for REE and Union #10991

Open

yongster and others added 2 commits September 6, 2026 07:04

Merge branch 'main' into fix/ree-union-nullability

0bc7ac3

fix(arrow-select): allow sparse union take with mixed child nullability

4827226

Null take indices select a nullable child, matching dense unions. Other children are taken with dummy indices so non-nullable fields do not receive introduced nulls.

yang3.xie and others added 3 commits September 7, 2026 22:22

chore: retrigger CI after flaky Archery JS build

a95a6e0

Merge branch 'main' into fix/ree-union-nullability

8f15ad4

Merge branch 'main' into fix/ree-union-nullability

b3a9d09

Jefffrey reviewed Sep 8, 2026

View reviewed changes

Uh oh!

Conversation

yongster commented Sep 4, 2026

Uh oh!

yongster commented Sep 4, 2026

Uh oh!

Rich-T-kid left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Jefffrey Sep 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Jefffrey commented Sep 5, 2026

Uh oh!

yongster commented Sep 7, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Jefffrey Sep 5, 2026 •

edited

Loading