This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11831-df4dc3ede59bd43b5f71847231212381d068bb72 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit 1a9c8590c7d9998571c876ba2fb1bf5f982bb5bc Author: Matt Butrovich <[email protected]> AuthorDate: Fri Oct 23 13:57:47 2026 +0000 fix(arrow-select): reject nulls in non-nullable columns in take_record_batch (#11831) # Which issue does this PR close? - Closes #11824. # Rationale for this change #10945 is not in a release yet, so this fix should go in 60.1.0 (#10525). This PR changes the same lines of `take_record_batch` as #11829, so it should merge after #11829. I will rebase it then. # What changes are included in this PR? `take_record_batch` builds its result with `RecordBatch::try_new_with_options` instead of `RecordBatch::new_unchecked`, so it checks nullability as 60.0.0 did. It passes the row count so a batch with no columns still works. # Are these changes tested? `take_record_batch` had no tests of its own, only its doc example and one test through `BatchCoalescer::push_batch_with_indices`. The new `test_take_record_batch_*` tests in `arrow-select/src/take.rs` cover a null index into non-nullable primitive, string, dictionary, struct, and fixed-size list columns, indices whose validity buffer has no nulls, and a batch with no columns. The benchmark bot shows `take_record_batch` within 2% of `main` ([default](https://github.com/apache/arrow-rs/pull/11831#issuecomment-6350836350), [aligned](https://github.com/apache/arrow-rs/pull/11831#issuecomment-6350853613)). # Are there any significant user-facing changes? No. This restores the 60.0.0 behavior. --------- Co-authored-by: Jefffrey <[email protected]> --- arrow-select/src/take.rs | 96 +++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 90 insertions(+), 6 deletions(-) diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index 410fff1a48..2b01a4b1f5 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -1726,9 +1726,10 @@ pub fn take_record_batch( .iter() .map(|c| take_impl::<_, false>(c.as_ref(), &indices)) .collect::<Result<Vec<_>, _>>()?; - // Safety: take returns one column per input column, of the same type and with - // indices.len() rows - Ok(unsafe { RecordBatch::new_unchecked(record_batch.schema(), columns, indices.len()) }) + // Null indices can put nulls in non-nullable columns, which this rejects. The + // row count is needed for a batch with no columns. + let options = RecordBatchOptions::new().with_row_count(Some(indices.len())); + RecordBatch::try_new_with_options(record_batch.schema(), columns, &options) }, d => Err(ArrowError::InvalidArgumentError(format!("Take only supported for integers, got {d:?}"))) ) @@ -4208,6 +4209,83 @@ mod tests { assert_eq!(result.len(), 2); } + /// Non-nullable columns of 3 rows, one per layout that represents a null differently + fn non_nullable_columns() -> Vec<ArrayRef> { + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["a", "bb", "ccc"])), + Arc::new(DictionaryArray::<Int8Type>::from_iter(["x", "y", "x"])), + Arc::new(StructArray::from(vec![( + Arc::new(Field::new("c", DataType::Int32, false)), + Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef, + )])), + Arc::new(FixedSizeListArray::from_iter_primitive::<Int32Type, _, _>( + vec![ + Some(vec![Some(1)]), + Some(vec![Some(2)]), + Some(vec![Some(3)]), + ], + 1, + )), + ] + } + + /// A batch with a nullable column followed by `column` in the non-nullable field "a" + fn batch_with_non_nullable(column: ArrayRef) -> RecordBatch { + let schema = Schema::new(vec![ + Field::new("n", DataType::Int32, true), + Field::new("a", column.data_type().clone(), false), + ]); + let n = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])); + RecordBatch::try_new(Arc::new(schema), vec![n, column]).unwrap() + } + + #[test] + fn test_take_record_batch_null_index_non_nullable_column() { + let indices = UInt32Array::from(vec![Some(2), None]); + for column in non_nullable_columns() { + let batch = batch_with_non_nullable(column); + let err = take_record_batch(&batch, &indices).unwrap_err(); + assert_eq!( + err.to_string(), + "Invalid argument error: Column 'a' is declared as non-nullable but contains null values", + "{}", + batch.column(1).data_type() + ); + } + } + + #[test] + fn test_take_record_batch_non_nullable_column() { + let indices = [ + UInt32Array::from(vec![2, 0, 2]), + // A validity buffer with no nulls + UInt32Array::new(vec![2, 0].into(), Some(NullBuffer::new_valid(2))), + ]; + for column in non_nullable_columns() { + let batch = batch_with_non_nullable(column); + for indices in &indices { + let taken = take_record_batch(&batch, indices).unwrap(); + + assert_eq!(taken.schema(), batch.schema()); + for (taken, column) in taken.columns().iter().zip(batch.columns()) { + assert_eq!(taken, &take(column, indices, None).unwrap()); + } + } + } + } + + #[test] + fn test_take_record_batch_no_columns_row_count() { + let options = RecordBatchOptions::new().with_row_count(Some(3)); + let batch = + RecordBatch::try_new_with_options(Arc::new(Schema::empty()), vec![], &options).unwrap(); + + let indices = UInt32Array::from(vec![Some(2), None, Some(0), Some(2)]); + let taken = take_record_batch(&batch, &indices).unwrap(); + assert_eq!(taken.num_rows(), 4); + } + /// Columns whose `take` reads no index, so taking them alone does not check bounds fn zero_width_columns(len: usize) -> Vec<ArrayRef> { vec![ @@ -4237,13 +4315,18 @@ mod tests { /// A batch of 3 rows with `first` followed by columns whose `take` reads every index fn batch_with_first_column(first: ArrayRef) -> RecordBatch { - RecordBatch::try_from_iter([ - ("first", first), + RecordBatch::try_from_iter_with_nullable([ + ("first", first, true), ( "i", Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])) as ArrayRef, + true, + ), + ( + "s", + Arc::new(StringArray::from(vec!["a", "bb", "ccc"])), + true, ), - ("s", Arc::new(StringArray::from(vec!["a", "bb", "ccc"]))), ( "l", Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![ @@ -4251,6 +4334,7 @@ mod tests { None, Some(vec![Some(3), None]), ])), + true, ), ]) .unwrap()
