Skip to content

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

Open
yongster wants to merge 1 commit into
apache:mainfrom
yongster:fix/ree-union-nullability
Open

fix(arrow-select): preserve nullability for REE and Union take#10994
yongster wants to merge 1 commit into
apache:mainfrom
yongster:fix/ree-union-nullability

Conversation

@yongster

@yongster yongster commented Sep 4, 2026

Copy link
Copy Markdown
Contributor

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

yongster commented Sep 4, 2026

Copy link
Copy Markdown
Contributor Author

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.

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

@Rich-T-kid Rich-T-kid 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.

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

Comment thread arrow-select/src/take.rs
}
}
DataType::Union(fields, UnionMode::Sparse) => {
if indices.null_count() > 0 && fields.iter().any(|(_, field)| !field.is_nullable()) {

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.

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?

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

Labels

arrow Changes to the arrow crate arrow-select

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Define nullability semantics for kernels on REE and Union arrays

2 participants