This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 314c52999e [arrow-select] perf: Replace `ArrayData` with direct
`Array` construction in concat (#10377)
314c52999e is described below
commit 314c52999e8308cabd914e26e7f46340ec145660
Author: Liam Bao <[email protected]>
AuthorDate: Mon Sep 28 07:18:49 2026 -0700
[arrow-select] perf: Replace `ArrayData` with direct `Array` construction
in concat (#10377)
# Which issue does this PR close?
- Part of #9298.
# Rationale for this change
`concat_run_arrays` built its result via `RunArray::try_new`, which
re-validates the run ends on every call even though they're already
known-valid here: inputs are valid `RunArray`s, so the adjusted run ends
are strictly increasing by construction.
# What changes are included in this PR?
- `concat.rs`: `concat_run_arrays` now builds the result directly via
`RunArray::new_unchecked(data_type, RunEndBuffer::new_unchecked(...),
values)`, skipping `try_new`'s revalidation.
- `concat.rs`: adds an explicit overflow guard and rejects up front if
the total doesn't fit in the run-end type `R`, before any `R::Native`
arithmetic runs.
- `nullif.rs`/`take.rs` changes are test-only cleanup.
# Are these changes tested?
Yes — existing test suite, plus a new regression test for the overflow
guard.
# Benchmark results
```
cargo bench -p arrow --bench concatenate_kernel --features test_utils --
"concat run"
```
this branch vs `main` (`RunArray<Int32Type>`, `n` arrays concatenated):
| Case | main | this PR | Change |
|---|---|---|---|
| logical=32, physical=4, n=4 | 773.6 ns | 563.4 ns | **−27%** |
| logical=1024, physical=128, n=8 | 3.33 µs | 2.68 µs | **−19%** |
| logical=8192, physical=1024, n=10 | 25.2 µs | 20.4 µs | **−19%** |
# Are there any user-facing changes?
No.
---
arrow-schema/src/error.rs | 4 +-
arrow-select/src/concat.rs | 47 ++++++---
arrow-select/src/nullif.rs | 22 +---
arrow-select/src/take.rs | 194 ++++++++++++------------------------
arrow/benches/concatenate_kernel.rs | 15 +++
5 files changed, 117 insertions(+), 165 deletions(-)
diff --git a/arrow-schema/src/error.rs b/arrow-schema/src/error.rs
index b9818b1b92..35d480d549 100644
--- a/arrow-schema/src/error.rs
+++ b/arrow-schema/src/error.rs
@@ -60,7 +60,7 @@ pub enum ArrowError {
CDataInterface(String),
/// Error when a dictionary key is bigger than the key type
DictionaryKeyOverflowError,
- /// Error when the run end index in a REE array is bigger than the array
length
+ /// Error when a run-end index is bigger than the run-end type
RunEndIndexOverflowError,
/// Error when the offset overflows.
OffsetOverflowError(usize),
@@ -129,7 +129,7 @@ impl Display for ArrowError {
write!(f, "Dictionary key bigger than the key type")
}
ArrowError::RunEndIndexOverflowError => {
- write!(f, "Run end encoded array index overflow error")
+ write!(f, "Run end index bigger than the run-end type")
}
ArrowError::OffsetOverflowError(offset) => {
write!(f, "Offset overflow error: {offset}")
diff --git a/arrow-select/src/concat.rs b/arrow-select/src/concat.rs
index 72dd239bbb..51b8082853 100644
--- a/arrow-select/src/concat.rs
+++ b/arrow-select/src/concat.rs
@@ -38,9 +38,9 @@ use arrow_array::cast::AsArray;
use arrow_array::types::*;
use arrow_array::*;
use arrow_buffer::{
- ArrowNativeType, BooleanBufferBuilder, MutableBuffer, NullBuffer,
OffsetBuffer, ScalarBuffer,
+ ArrowNativeType, BooleanBufferBuilder, MutableBuffer, NullBuffer,
OffsetBuffer, RunEndBuffer,
+ ScalarBuffer,
};
-use arrow_data::ArrayDataBuilder;
use arrow_data::transform::{Capacities, MutableArrayData};
use arrow_schema::{ArrowError, DataType, FieldRef, Fields, SchemaRef};
use std::{collections::HashSet, ops::Add, sync::Arc};
@@ -430,6 +430,12 @@ where
return Ok(new_empty_array(arrays[0].data_type()));
}
+ // Reject lengths that do not fit in `R` before any `R::Native` arithmetic
can wrap.
+ let total_len: usize = run_arrays.iter().map(|r| r.len()).sum();
+ if R::Native::from_usize(total_len).is_none() {
+ return Err(ArrowError::RunEndIndexOverflowError);
+ }
+
// The run ends need to be adjusted by the sum of the lengths of the
previous arrays.
let needed_run_end_adjustments = std::iter::once(R::default_value())
.chain(
@@ -442,9 +448,6 @@ where
)
.collect::<Vec<_>>();
- // This works out nicely to be the total (logical) length of the resulting
array.
- let total_len = needed_run_end_adjustments.last().unwrap().as_usize();
-
let run_ends_array =
PrimitiveArray::<R>::from_iter_values(run_arrays.iter().enumerate().flat_map(
move |(i, run_array)| {
@@ -463,15 +466,15 @@ where
let all_values = concat(&values_slices.iter().map(|x|
x.as_ref()).collect::<Vec<_>>())?;
- let builder = ArrayDataBuilder::new(run_arrays[0].data_type().clone())
- .len(total_len)
- .child_data(vec![run_ends_array.into_data(), all_values.into_data()]);
+ let data_type = run_arrays[0].data_type().clone();
+ let (_, run_ends_values, _) = run_ends_array.into_parts();
- // `build_unchecked` is used to avoid recursive validation of child arrays.
- let array_data = unsafe { builder.build_unchecked() };
- array_data.validate_data()?;
-
- Ok(Arc::<RunArray<R>>::new(array_data.into()))
+ // Safety: inputs are valid RunArrays; adjusted run ends are strictly
increasing
+ // and end at `total_len`. Physical length matches `all_values`.
+ let run_ends = unsafe { RunEndBuffer::new_unchecked(run_ends_values, 0,
total_len) };
+ Ok(Arc::new(unsafe {
+ RunArray::<R>::new_unchecked(data_type, run_ends, all_values)
+ }))
}
macro_rules! dict_helper {
@@ -1889,6 +1892,24 @@ mod tests {
assert_eq!(&[10, 20, 30, 40], values.values());
}
+ #[test]
+ fn test_concat_run_array_length_overflows_run_end_type() {
+ // 20_000 + 20_000 exceeds i16::MAX.
+ let array1 = RunArray::<Int16Type>::try_new(
+ &Int16Array::from(vec![20_000]),
+ &Int16Array::from(vec![1]),
+ )
+ .unwrap();
+ let array2 = RunArray::<Int16Type>::try_new(
+ &Int16Array::from(vec![20_000]),
+ &Int16Array::from(vec![2]),
+ )
+ .unwrap();
+
+ let err = concat(&[&array1, &array2]).unwrap_err();
+ assert!(matches!(err, ArrowError::RunEndIndexOverflowError));
+ }
+
#[test]
fn test_concat_sliced_run_array() {
// Slicing away first run in both arrays
diff --git a/arrow-select/src/nullif.rs b/arrow-select/src/nullif.rs
index aa09a28631..cd6d0ec4a2 100644
--- a/arrow-select/src/nullif.rs
+++ b/arrow-select/src/nullif.rs
@@ -578,31 +578,13 @@ mod tests {
/// Returns a new BooleanArray with no null buffer
fn remove_null_buffer(array: &BooleanArray) -> BooleanArray {
- make_array(
- array
- .into_data()
- .into_builder()
- .nulls(None)
- .build()
- .unwrap(),
- )
- .as_boolean()
- .clone()
+ BooleanArray::new(array.values().clone(), None)
}
/// Returns a new BooleanArray with a null buffer where all values are
valid
fn remove_null_values(array: &BooleanArray) -> BooleanArray {
let len = array.len();
let new_nulls = NullBuffer::from_iter(std::iter::repeat_n(true, len));
- make_array(
- array
- .into_data()
- .into_builder()
- .nulls(Some(new_nulls))
- .build()
- .unwrap(),
- )
- .as_boolean()
- .clone()
+ BooleanArray::new(array.values().clone(), Some(new_nulls))
}
}
diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs
index 8c9c4ba3da..4ed8d856d4 100644
--- a/arrow-select/src/take.rs
+++ b/arrow-select/src/take.rs
@@ -1674,7 +1674,6 @@ mod tests {
use super::*;
use arrow_array::builder::*;
use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano};
- use arrow_data::ArrayData;
use arrow_schema::{Field, Fields, TimeUnit, UnionFields};
use num_traits::ToPrimitive;
@@ -2174,18 +2173,12 @@ mod tests {
#[test]
fn test_take_bool_nullable_index() {
// indices where the masked invalid elements would be out of bounds
- let index_data = ArrayData::try_new(
- DataType::UInt32,
- 6,
- Some(Buffer::from_iter(vec![
+ let index = UInt32Array::new(
+ ScalarBuffer::from(vec![99, 0, 999, 1, 9999, 2]),
+ Some(NullBuffer::from(vec![
false, true, false, true, false, true,
])),
- 0,
- vec![Buffer::from_iter(vec![99, 0, 999, 1, 9999, 2])],
- vec![],
- )
- .unwrap();
- let index = UInt32Array::from(index_data);
+ );
test_take_boolean_arrays(
vec![Some(true), None, Some(false)],
&index,
@@ -2197,18 +2190,12 @@ mod tests {
#[test]
fn test_take_bool_nullable_index_nonnull_values() {
// indices where the masked invalid elements would be out of bounds
- let index_data = ArrayData::try_new(
- DataType::UInt32,
- 6,
- Some(Buffer::from_iter(vec![
+ let index = UInt32Array::new(
+ ScalarBuffer::from(vec![99, 0, 999, 1, 9999, 2]),
+ Some(NullBuffer::from(vec![
false, true, false, true, false, true,
])),
- 0,
- vec![Buffer::from_iter(vec![99, 0, 999, 1, 9999, 2])],
- vec![],
- )
- .unwrap();
- let index = UInt32Array::from(index_data);
+ );
test_take_boolean_arrays(
vec![Some(true), Some(true), Some(false)],
&index,
@@ -2464,20 +2451,11 @@ mod tests {
macro_rules! test_take_list {
($offset_type:ty, $list_data_type:ident, $list_array_type:ident) => {{
// Construct a value array, [[0,0,0], [-1,-2,-1], [], [2,3]]
- let value_data = Int32Array::from(vec![0, 0, 0, -1, -2, -1, 2,
3]).into_data();
- // Construct offsets
- let value_offsets: [$offset_type; 5] = [0, 3, 6, 6, 8];
- let value_offsets = Buffer::from_slice_ref(&value_offsets);
- // Construct a list array from the above two
- let list_data_type =
-
DataType::$list_data_type(Arc::new(Field::new_list_field(DataType::Int32,
false)));
- let list_data = ArrayData::builder(list_data_type.clone())
- .len(4)
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .build()
- .unwrap();
- let list_array = $list_array_type::from(list_data);
+ let values = Arc::new(Int32Array::from(vec![0, 0, 0, -1, -2, -1,
2, 3]));
+ let value_offsets =
+ OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0,
3, 6, 6, 8]));
+ let field = Arc::new(Field::new_list_field(DataType::Int32,
false));
+ let list_array = $list_array_type::new(Arc::clone(&field),
value_offsets, values, None);
// index returns: [[2,3], null, [-1,-2,-1], [], [0,0,0]]
let index = UInt32Array::from(vec![Some(3), None, Some(1),
Some(2), Some(0)]);
@@ -2487,7 +2465,7 @@ mod tests {
// construct a value array with expected results:
// [[2,3], null, [-1,-2,-1], [], [0,0,0]]
- let expected_data = Int32Array::from(vec![
+ let expected_values = Arc::new(Int32Array::from(vec![
Some(2),
Some(3),
Some(-1),
@@ -2496,21 +2474,16 @@ mod tests {
Some(0),
Some(0),
Some(0),
- ])
- .into_data();
- // construct offsets
- let expected_offsets: [$offset_type; 6] = [0, 2, 2, 5, 5, 8];
- let expected_offsets = Buffer::from_slice_ref(&expected_offsets);
- // construct list array from the two
- let expected_list_data = ArrayData::builder(list_data_type)
- .len(5)
+ ]));
+ let expected_offsets =
+ OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0,
2, 2, 5, 5, 8]));
+ let expected_list_array = $list_array_type::new(
+ field,
+ expected_offsets,
+ expected_values,
// null buffer remains the same as only the indices have nulls
- .nulls(index.nulls().cloned())
- .add_buffer(expected_offsets)
- .add_child_data(expected_data)
- .build()
- .unwrap();
- let expected_list_array =
$list_array_type::from(expected_list_data);
+ index.nulls().cloned(),
+ );
assert_eq!(a, &expected_list_array);
}};
@@ -2519,7 +2492,7 @@ mod tests {
macro_rules! test_take_list_with_value_nulls {
($offset_type:ty, $list_data_type:ident, $list_array_type:ident) => {{
// Construct a value array, [[0,null,0], [-1,-2,3], [null],
[5,null]]
- let value_data = Int32Array::from(vec![
+ let values = Arc::new(Int32Array::from(vec![
Some(0),
None,
Some(0),
@@ -2529,22 +2502,11 @@ mod tests {
None,
Some(5),
None,
- ])
- .into_data();
- // Construct offsets
- let value_offsets: [$offset_type; 5] = [0, 3, 6, 7, 9];
- let value_offsets = Buffer::from_slice_ref(&value_offsets);
- // Construct a list array from the above two
- let list_data_type =
-
DataType::$list_data_type(Arc::new(Field::new_list_field(DataType::Int32,
true)));
- let list_data = ArrayData::builder(list_data_type.clone())
- .len(4)
- .add_buffer(value_offsets)
- .null_bit_buffer(Some(Buffer::from([0b11111111])))
- .add_child_data(value_data)
- .build()
- .unwrap();
- let list_array = $list_array_type::from(list_data);
+ ]));
+ let value_offsets =
+ OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0,
3, 6, 7, 9]));
+ let field = Arc::new(Field::new_list_field(DataType::Int32, true));
+ let list_array = $list_array_type::new(Arc::clone(&field),
value_offsets, values, None);
// index returns: [[null], null, [-1,-2,3], [2,null], [0,null,0]]
let index = UInt32Array::from(vec![Some(2), None, Some(1),
Some(3), Some(0)]);
@@ -2554,7 +2516,7 @@ mod tests {
// construct a value array with expected results:
// [[null], null, [-1,-2,3], [5,null], [0,null,0]]
- let expected_data = Int32Array::from(vec![
+ let expected_values = Arc::new(Int32Array::from(vec![
None,
Some(-1),
Some(-2),
@@ -2564,21 +2526,16 @@ mod tests {
Some(0),
None,
Some(0),
- ])
- .into_data();
- // construct offsets
- let expected_offsets: [$offset_type; 6] = [0, 1, 1, 4, 6, 9];
- let expected_offsets = Buffer::from_slice_ref(&expected_offsets);
- // construct list array from the two
- let expected_list_data = ArrayData::builder(list_data_type)
- .len(5)
+ ]));
+ let expected_offsets =
+ OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0,
1, 1, 4, 6, 9]));
+ let expected_list_array = $list_array_type::new(
+ field,
+ expected_offsets,
+ expected_values,
// null buffer remains the same as only the indices have nulls
- .nulls(index.nulls().cloned())
- .add_buffer(expected_offsets)
- .add_child_data(expected_data)
- .build()
- .unwrap();
- let expected_list_array =
$list_array_type::from(expected_list_data);
+ index.nulls().cloned(),
+ );
assert_eq!(a, &expected_list_array);
}};
@@ -2587,7 +2544,7 @@ mod tests {
macro_rules! test_take_list_with_nulls {
($offset_type:ty, $list_data_type:ident, $list_array_type:ident) => {{
// Construct a value array, [[0,null,0], [-1,-2,3], null, [5,null]]
- let value_data = Int32Array::from(vec![
+ let values = Arc::new(Int32Array::from(vec![
Some(0),
None,
Some(0),
@@ -2596,22 +2553,14 @@ mod tests {
Some(3),
Some(5),
None,
- ])
- .into_data();
- // Construct offsets
- let value_offsets: [$offset_type; 5] = [0, 3, 6, 6, 8];
- let value_offsets = Buffer::from_slice_ref(&value_offsets);
- // Construct a list array from the above two
- let list_data_type =
-
DataType::$list_data_type(Arc::new(Field::new_list_field(DataType::Int32,
true)));
- let list_data = ArrayData::builder(list_data_type.clone())
- .len(4)
- .add_buffer(value_offsets)
- .null_bit_buffer(Some(Buffer::from([0b11111011])))
- .add_child_data(value_data)
- .build()
- .unwrap();
- let list_array = $list_array_type::from(list_data);
+ ]));
+ let value_offsets =
+ OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0,
3, 6, 6, 8]));
+ let field = Arc::new(Field::new_list_field(DataType::Int32, true));
+ // the entry at index 2 is null
+ let list_nulls = NullBuffer::from(vec![true, true, false, true]);
+ let list_array =
+ $list_array_type::new(Arc::clone(&field), value_offsets,
values, Some(list_nulls));
// index returns: [null, null, [-1,-2,3], [5,null], [0,null,0]]
let index = UInt32Array::from(vec![Some(2), None, Some(1),
Some(3), Some(0)]);
@@ -2621,7 +2570,7 @@ mod tests {
// construct a value array with expected results:
// [null, null, [-1,-2,3], [5,null], [0,null,0]]
- let expected_data = Int32Array::from(vec![
+ let expected_values = Arc::new(Int32Array::from(vec![
Some(-1),
Some(-2),
Some(3),
@@ -2630,25 +2579,17 @@ mod tests {
Some(0),
None,
Some(0),
- ])
- .into_data();
- // construct offsets
- let expected_offsets: [$offset_type; 6] = [0, 0, 0, 3, 5, 8];
- let expected_offsets = Buffer::from_slice_ref(&expected_offsets);
- // construct list array from the two
- let mut null_bits: [u8; 1] = [0; 1];
- bit_util::set_bit(&mut null_bits, 2);
- bit_util::set_bit(&mut null_bits, 3);
- bit_util::set_bit(&mut null_bits, 4);
- let expected_list_data = ArrayData::builder(list_data_type)
- .len(5)
- // null buffer must be recalculated as both values and indices
have nulls
- .null_bit_buffer(Some(Buffer::from(null_bits)))
- .add_buffer(expected_offsets)
- .add_child_data(expected_data)
- .build()
- .unwrap();
- let expected_list_array =
$list_array_type::from(expected_list_data);
+ ]));
+ let expected_offsets =
+ OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0,
0, 0, 3, 5, 8]));
+ // null buffer must be recalculated as both values and indices
have nulls
+ let expected_nulls = NullBuffer::from(vec![false, false, true,
true, true]);
+ let expected_list_array = $list_array_type::new(
+ field,
+ expected_offsets,
+ expected_values,
+ Some(expected_nulls),
+ );
assert_eq!(a, &expected_list_array);
}};
@@ -2958,19 +2899,12 @@ mod tests {
#[should_panic(expected = "index out of bounds: the len is 4 but the index
is 1000")]
fn test_take_list_out_of_bounds() {
// Construct a value array, [[0,0,0], [-1,-2,-1], [2,3]]
- let value_data = Int32Array::from(vec![0, 0, 0, -1, -2, -1, 2,
3]).into_data();
+ let values = Arc::new(Int32Array::from(vec![0, 0, 0, -1, -2, -1, 2,
3]));
// Construct offsets
- let value_offsets = Buffer::from_slice_ref([0, 3, 6, 8]);
+ let value_offsets =
OffsetBuffer::<i32>::new(ScalarBuffer::from(vec![0, 3, 6, 8]));
// Construct a list array from the above two
- let list_data_type =
- DataType::List(Arc::new(Field::new_list_field(DataType::Int32,
false)));
- let list_data = ArrayData::builder(list_data_type)
- .len(3)
- .add_buffer(value_offsets)
- .add_child_data(value_data)
- .build()
- .unwrap();
- let list_array = ListArray::from(list_data);
+ let field = Arc::new(Field::new_list_field(DataType::Int32, false));
+ let list_array = ListArray::new(field, value_offsets, values, None);
let index = UInt32Array::from(vec![1000]);
diff --git a/arrow/benches/concatenate_kernel.rs
b/arrow/benches/concatenate_kernel.rs
index 497c2e8f1e..6ae97c90a9 100644
--- a/arrow/benches/concatenate_kernel.rs
+++ b/arrow/benches/concatenate_kernel.rs
@@ -235,6 +235,21 @@ fn add_benchmark(c: &mut Criterion) {
|b| b.iter(|| bench_concat_arrays(&array_refs)),
);
}
+
+ // (name, logical_len, physical_len, num_arrays)
+ for (name, logical, physical, n) in [
+ ("small logical=32 physical=4 x4", 32usize, 4usize, 4usize),
+ ("logical=1024 physical=128 x8", 1024, 128, 8),
+ ("logical=8192 physical=1024 x10", 8192, 1024, 10),
+ ] {
+ let arrays: Vec<RunArray<Int32Type>> = (0..n)
+ .map(|_| create_primitive_run_array::<Int32Type,
Int32Type>(logical, physical))
+ .collect();
+ let array_refs: Vec<&dyn Array> = arrays.iter().map(|a| a as &dyn
Array).collect();
+ c.bench_function(&format!("concat run_end i32 {name}"), |b| {
+ b.iter(|| bench_concat_arrays(&array_refs))
+ });
+ }
}
criterion_group!(benches, add_benchmark);
