Skip to content

[arrow-select] minor: Replace ArrayData with direct Array construction in concat #10377

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
liamzwbao wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

[arrow-select] minor: Replace ArrayData with direct Array construction in concat #10377

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 4 additions & 13 deletions arrow-select/src/concat.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,6 @@ use arrow_array::*;
use arrow_buffer::{
ArrowNativeType, BooleanBufferBuilder, MutableBuffer, NullBuffer, OffsetBuffer, 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};
Expand Down Expand Up @@ -436,9 +435,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)| {
Expand All @@ -457,15 +453,10 @@ 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()]);

// `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()))
Ok(Arc::new(RunArray::<R>::try_new(
&run_ends_array,
all_values.as_ref(),
)?))
}

macro_rules! dict_helper {
Expand Down
22 changes: 2 additions & 20 deletions arrow-select/src/nullif.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -577,31 +577,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))
}
}
194 changes: 64 additions & 130 deletions arrow-select/src/take.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1135,7 +1135,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;

Expand Down Expand Up @@ -1635,18 +1634,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,
Expand All @@ -1658,18 +1651,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,
Expand Down Expand Up @@ -1824,20 +1811,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)]);
Expand All @@ -1847,7 +1825,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),
Expand All @@ -1856,21 +1834,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);
}};
Expand All @@ -1879,7 +1852,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),
Expand All @@ -1889,22 +1862,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)]);
Expand All @@ -1914,7 +1876,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),
Expand All @@ -1924,21 +1886,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);
}};
Expand All @@ -1947,7 +1904,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),
Expand All @@ -1956,22 +1913,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)]);
Expand All @@ -1981,7 +1930,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),
Expand All @@ -1990,25 +1939,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);
}};
Expand Down Expand Up @@ -2294,19 +2235,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]);

Expand Down
Loading