use crate::array::growable::GrowableFixedSizeList;
use crate::array::FixedSizeListArray;
use crate::array::{growable::Growable, PrimitiveArray};
use super::Index;
pub fn take<O: Index>(
values: &FixedSizeListArray,
indices: &PrimitiveArray<O>,
) -> FixedSizeListArray {
let mut capacity = 0;
let arrays = indices
.values()
.iter()
.map(|index| {
let index = index.to_usize();
let slice = values.clone().sliced(index, 1);
capacity += slice.len();
slice
})
.collect::<Vec<FixedSizeListArray>>();
let arrays = arrays.iter().collect();
if let Some(validity) = indices.validity() {
let mut growable: GrowableFixedSizeList =
GrowableFixedSizeList::new(arrays, true, capacity);
for index in 0..indices.len() {
if validity.get_bit(index) {
growable.extend(index, 0, 1);
} else {
growable.extend_validity(1)
}
}
growable.into()
} else {
let mut growable: GrowableFixedSizeList =
GrowableFixedSizeList::new(arrays, false, capacity);
for index in 0..indices.len() {
growable.extend(index, 0, 1);
}
growable.into()
}
}