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-11845-8953c9014c487676561b80a1e8a8c86ea44066bb in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit 03d40357b73c02e6454c7d352d6552f01bf500a8 Author: Prince Jain <[email protected]> AuthorDate: Mon Oct 26 04:14:21 2026 +0000 arrow-select: error on empty input and keep row count for column-less batches in interleave_record_batch (#11845) # Which issue does this PR close? - Closes #11833 # Rationale for this change `interleave_record_batch` read the schema from `record_batches[0]`, so calling it with no batches panicked with an index-out-of-bounds, and it built the result with `RecordBatch::try_new`, which rejects a batch with no columns and no explicit row count even though `indices.len()` provides one. # What changes are included in this PR? - Empty `record_batches` now returns `InvalidArgumentError`, matching the array-level `interleave`. - The result is built with `try_new_with_options` and `row_count = indices.len()`, so a column-less batch produces a batch with that many rows, matching `take_record_batch`. # Are these changes tested? Yes. Two tests in `interleave.rs`: empty input returns an `InvalidArgumentError`, and a 3-row schema-empty batch interleaved with `[(0, 0), (0, 2)]` yields a 2-row batch. Both fail on unmodified code (index panic and `InvalidArgumentError` respectively). Co-authored-by: RIchard Baah <[email protected]> --- arrow-select/src/interleave.rs | 36 +++++++++++++++++++++++++++++++++--- 1 file changed, 33 insertions(+), 3 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 00abaf638e..1ccee00b99 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -969,7 +969,12 @@ pub fn interleave_record_batch( record_batches: &[&RecordBatch], indices: &[(usize, usize)], ) -> Result<RecordBatch, ArrowError> { - let schema = record_batches[0].schema(); + let Some(first) = record_batches.first() else { + return Err(ArrowError::InvalidArgumentError( + "interleave_record_batch requires input of at least one record batch".to_string(), + )); + }; + let schema = first.schema(); let columns = (0..schema.fields().len()) .map(|i| { let column_values: Vec<&dyn Array> = record_batches @@ -979,7 +984,13 @@ pub fn interleave_record_batch( interleave(&column_values, indices) }) .collect::<Result<Vec<_>, _>>()?; - RecordBatch::try_new(schema, columns) + // The row count is set explicitly so that a batch with no columns still + // produces a batch with `indices.len()` rows, like `take_record_batch`. + RecordBatch::try_new_with_options( + schema, + columns, + &RecordBatchOptions::new().with_row_count(Some(indices.len())), + ) } #[cfg(test)] @@ -991,7 +1002,7 @@ mod tests { }; use arrow_array::types::{Decimal128Type, Int8Type, TimestampMicrosecondType}; use arrow_buffer::ScalarBuffer; - use arrow_schema::{Field, TimeUnit}; + use arrow_schema::{Field, Schema, TimeUnit}; #[test] fn test_primitive() { @@ -2533,4 +2544,23 @@ mod tests { assert_eq!(result.value(0).len(), 1); assert_eq!(result.value(3).len(), 1); } + + #[test] + fn test_interleave_record_batch_empty_input() { + let err = interleave_record_batch(&[], &[]).unwrap_err(); + assert!(matches!(err, ArrowError::InvalidArgumentError(_))); + } + + #[test] + fn test_interleave_record_batch_no_columns() { + // A batch with no columns but a known row count interleaves to a + // batch with `indices.len()` rows, matching `take_record_batch`. + let options = RecordBatchOptions::new().with_row_count(Some(3)); + let batch = + RecordBatch::try_new_with_options(Arc::new(Schema::empty()), vec![], &options).unwrap(); + + let result = interleave_record_batch(&[&batch], &[(0, 0), (0, 2)]).unwrap(); + assert_eq!(result.num_columns(), 0); + assert_eq!(result.num_rows(), 2); + } }
