use crate::ArrayData;
use arrow_schema::DataType;
use super::{Extend, _MutableArrayData};
pub(super) fn build_extend(array: &ArrayData) -> Extend {
let size = match array.data_type() {
DataType::FixedSizeList(_, i) => *i as usize,
_ => unreachable!(),
};
if array.null_count() == 0 {
Box::new(
move |mutable: &mut _MutableArrayData,
index: usize,
start: usize,
len: usize| {
mutable.child_data.iter_mut().for_each(|child| {
child.extend(index, start * size, (start + len) * size)
})
},
)
} else {
Box::new(
move |mutable: &mut _MutableArrayData,
index: usize,
start: usize,
len: usize| {
(start..start + len).for_each(|i| {
if array.is_valid(i) {
mutable.child_data.iter_mut().for_each(|child| {
child.extend(index, i * size, (i + 1) * size)
})
} else {
mutable
.child_data
.iter_mut()
.for_each(|child| child.extend_nulls(size))
}
})
},
)
}
}
pub(super) fn extend_nulls(mutable: &mut _MutableArrayData, len: usize) {
let size = match mutable.data_type {
DataType::FixedSizeList(_, i) => i as usize,
_ => unreachable!(),
};
mutable
.child_data
.iter_mut()
.for_each(|child| child.extend_nulls(len * size))
}