use std::borrow::Cow;
use arrow_array::{
Array, GenericByteArray, LargeBinaryArray, LargeStringArray, builder::GenericByteBuilder, types::ByteArrayType,
};
use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer, OffsetBuffer};
use arrow_select::filter::FilterPredicate;
use crate::{
Result,
util::{bitmap, kernel},
value::{Value, blob::Blob, value_type::ValueType},
};
pub fn get<T>(array: &GenericByteArray<T>, index: usize) -> Option<&T::Native>
where
T: ByteArrayType<Offset = i64>,
{
if index < array.len() {
Some(array.value(index))
} else {
None
}
}
fn row_bytes<T>(array: &GenericByteArray<T>, index: usize) -> &[u8]
where
T: ByteArrayType<Offset = i64>,
{
<T::Native as AsRef<[u8]>>::as_ref(array.value(index))
}
pub fn equals<T>(left: &GenericByteArray<T>, right: &GenericByteArray<T>) -> bool
where
T: ByteArrayType<Offset = i64>,
{
left.len() == right.len() && (0..left.len()).all(|i| row_bytes(left, i) == row_bytes(right, i))
}
pub fn empty<T>() -> GenericByteArray<T>
where
T: ByteArrayType<Offset = i64>,
{
GenericByteArray::new(OffsetBuffer::new_empty(), Buffer::from_vec(Vec::<u8>::new()), None)
}
fn data_byte_len<T>(array: &GenericByteArray<T>) -> usize
where
T: ByteArrayType<Offset = i64>,
{
let offsets = array.value_offsets();
(offsets[offsets.len() - 1] - offsets[0]) as usize
}
pub fn compact_parts<T>(array: &GenericByteArray<T>) -> (&[u8], Cow<'_, [i64]>)
where
T: ByteArrayType<Offset = i64>,
{
let offsets = array.value_offsets();
let first = offsets[0];
let last = offsets[offsets.len() - 1];
let data = &array.value_data()[first as usize..last as usize];
let offsets = if first == 0 {
Cow::Borrowed(offsets)
} else {
Cow::Owned(offsets.iter().map(|offset| offset - first).collect())
};
(data, offsets)
}
pub fn slice<T>(array: &GenericByteArray<T>, start: usize, end: usize) -> GenericByteArray<T>
where
T: ByteArrayType<Offset = i64>,
{
let len = array.len();
let start = start.min(len);
let end = end.min(len);
if start >= end {
return attach_nulls(empty(), bitmap::slice_nulls(array.logical_nulls().as_ref(), start, end));
}
array.slice(start, end - start)
}
pub fn take<T>(array: &GenericByteArray<T>, num: usize) -> GenericByteArray<T>
where
T: ByteArrayType<Offset = i64>,
{
slice(array, 0, num)
}
pub fn filter<T>(array: &GenericByteArray<T>, mask: &BooleanBuffer) -> GenericByteArray<T>
where
T: ByteArrayType<Offset = i64>,
{
filter_with(array, &kernel::predicate(mask, array.len()))
}
pub fn filter_with<T>(array: &GenericByteArray<T>, predicate: &FilterPredicate) -> GenericByteArray<T>
where
T: ByteArrayType<Offset = i64>,
{
let selected = kernel::filtered(array, predicate);
let nulls = kernel::kept_nulls(array, &selected);
attach_nulls(selected, nulls)
}
pub fn reorder<T>(array: &GenericByteArray<T>, indices: &[usize]) -> Result<GenericByteArray<T>>
where
T: ByteArrayType<Offset = i64>,
{
kernel::rows_in_range(indices, array.len())?;
let mut builder = GenericByteBuilder::<T>::with_capacity(indices.len(), data_byte_len(array));
for &idx in indices {
builder.append_value(array.value(idx));
}
Ok(attach_nulls(builder.finish(), bitmap::reorder_nulls(array.logical_nulls().as_ref(), indices)))
}
pub fn attach_nulls<T>(array: GenericByteArray<T>, nulls: Option<NullBuffer>) -> GenericByteArray<T>
where
T: ByteArrayType<Offset = i64>,
{
bitmap::assert_nulls_len(nulls.as_ref(), array.len());
let (offsets, values, _) = array.into_parts();
unsafe { GenericByteArray::new_unchecked(offsets, values, nulls) }
}
pub fn capacity<T>(array: &GenericByteArray<T>) -> usize
where
T: ByteArrayType<Offset = i64>,
{
let offsets = array.offsets().inner().inner();
if offsets.strong_count() == 1 {
(offsets.capacity() / size_of::<i64>()).saturating_sub(1)
} else {
array.len()
}
}
pub fn heap_size<T>(array: &GenericByteArray<T>) -> usize
where
T: ByteArrayType<Offset = i64>,
{
let values = array.values();
let data = if values.strong_count() == 1 {
values.capacity()
} else {
data_byte_len(array)
};
data + (capacity(array) + 1) * size_of::<i64>()
}
pub fn blob_array(values: &[Blob]) -> LargeBinaryArray {
LargeBinaryArray::from_iter_values(values.iter().map(|blob| blob.as_bytes()))
}
pub fn utf8_get_value(array: &LargeStringArray, index: usize) -> Value {
match get(array, index) {
Some(s) => Value::Utf8(s.to_string()),
None => Value::none_of(ValueType::Utf8),
}
}
pub fn utf8_as_string(array: &LargeStringArray, index: usize) -> String {
match get(array, index) {
Some(s) => s.to_string(),
None => "none".to_string(),
}
}
pub fn blob_get_value(array: &LargeBinaryArray, index: usize) -> Value {
match get(array, index) {
Some(bytes) => Value::Blob(Blob::new(bytes.to_vec())),
None => Value::none_of(ValueType::Blob),
}
}
pub fn blob_as_string(array: &LargeBinaryArray, index: usize) -> String {
match get(array, index) {
Some(bytes) => Blob::new(bytes.to_vec()).to_string(),
None => "none".to_string(),
}
}
#[cfg(test)]
mod tests {
use arrow_array::builder::LargeStringBuilder;
use super::*;
mod utf8 {
use super::*;
#[test]
fn test_slice() {
let container = LargeStringArray::from(vec![
"one".to_string(),
"two".to_string(),
"three".to_string(),
"four".to_string(),
]);
let sliced = slice(&container, 1, 3);
assert_eq!(sliced.len(), 2);
assert_eq!(get(&sliced, 0), Some("two"));
assert_eq!(get(&sliced, 1), Some("three"));
}
#[test]
fn test_filter() {
let container = LargeStringArray::from(vec![
"keep".to_string(),
"drop".to_string(),
"keep".to_string(),
"drop".to_string(),
]);
let mask = BooleanBuffer::from(vec![true, false, true, false]);
let container = filter(&container, &mask);
assert_eq!(container.len(), 2);
assert_eq!(get(&container, 0), Some("keep"));
assert_eq!(get(&container, 1), Some("keep"));
}
#[test]
fn test_reorder() {
let container = LargeStringArray::from(vec![
"first".to_string(),
"second".to_string(),
"third".to_string(),
]);
let indices = [2, 0, 1];
let container = reorder(&container, &indices).unwrap();
assert_eq!(container.len(), 3);
assert_eq!(get(&container, 0), Some("third"));
assert_eq!(get(&container, 1), Some("first"));
assert_eq!(get(&container, 2), Some("second"));
}
#[test]
fn test_reorder_with_out_of_bounds() {
let container = LargeStringArray::from(vec!["a".to_string(), "b".to_string()]);
let error = reorder(&container, &[1, 5, 0]).unwrap_err();
assert_eq!(error.diagnostic().message, "row index 5 out of range for a column of 2 rows");
}
#[test]
fn test_empty_strings() {
let mut builder = LargeStringBuilder::with_capacity(2, 0);
builder.append_value("");
builder.append_value("");
let container = builder.finish();
assert_eq!(container.len(), 2);
assert_eq!(get(&container, 0), Some(""));
assert_eq!(get(&container, 1), Some(""));
assert!(0 < container.len());
assert!(1 < container.len());
}
#[test]
fn test_data_bytes_and_offsets_match_zero_copy_layout() {
let container = LargeStringArray::from(vec!["aa".to_string(), "bb".to_string()]);
assert_eq!(compact_parts(&container).0, b"aabb");
assert_eq!(&*compact_parts(&container).1, &[0i64, 2, 4]);
}
}
mod blob {
use super::*;
#[test]
fn test_data_bytes_and_offsets_match_zero_copy_layout() {
let container = blob_array(&[Blob::new(vec![0xAA, 0xBB]), Blob::new(vec![0xCC])]);
assert_eq!(compact_parts(&container).0, &[0xAAu8, 0xBB, 0xCC]);
assert_eq!(&*compact_parts(&container).1, &[0i64, 2, 3]);
}
}
mod varlen {
use super::*;
fn abcd() -> LargeBinaryArray {
LargeBinaryArray::from_iter_values([b"aa".as_slice(), b"bb", b"cc", b"dd"])
}
#[test]
fn filter_in_place_keeps_matching_elements() {
let c = LargeBinaryArray::from_iter_values([b"yes".as_slice(), b"no", b"yes", b"no"]);
let c = filter(&c, &BooleanBuffer::collect_bool(4, |i| i % 2 == 0));
assert_eq!(c.len(), 2);
assert_eq!(get(&c, 0), Some(b"yes".as_slice()));
assert_eq!(get(&c, 1), Some(b"yes".as_slice()));
}
#[test]
fn reorder_refuses_an_out_of_range_index() {
let c = LargeBinaryArray::from_iter_values([b"a".as_slice(), b"b"]);
let error = reorder(&c, &[1, 100, 0]).unwrap_err();
assert_eq!(error.diagnostic().message, "row index 100 out of range for a column of 2 rows");
}
#[test]
fn take_n_truncates() {
let c = LargeBinaryArray::from_iter_values([b"a".as_slice(), b"b", b"c"]);
let t = take(&c, 2);
assert_eq!(t.len(), 2);
assert_eq!(get(&t, 0), Some(b"a".as_slice()));
assert_eq!(get(&t, 1), Some(b"b".as_slice()));
}
#[test]
fn slice_extracts_subrange_with_rebased_offsets() {
let c = LargeBinaryArray::from_iter_values([b"aa".as_slice(), b"bb", b"cc", b"dd"]);
let s = slice(&c, 1, 3);
assert_eq!(s.len(), 2);
assert_eq!(get(&s, 0), Some(b"bb".as_slice()));
assert_eq!(get(&s, 1), Some(b"cc".as_slice()));
assert_eq!(&*compact_parts(&s).1, &[0i64, 2, 4]);
}
#[test]
fn equality_compares_logical_content() {
let a = LargeBinaryArray::from_iter_values([b"x".as_slice(), b"y"]);
let b = LargeBinaryArray::from_iter_values([b"x".as_slice(), b"y"]);
assert!(equals(&a, &b));
}
#[test]
fn frozen_slice_compact_parts_are_rebased() {
let c = abcd();
let s = slice(&c, 1, 3);
let (data, offsets) = compact_parts(&s);
assert_eq!(data, b"bbcc");
assert_eq!(&*offsets, &[0i64, 2, 4]);
assert!(matches!(offsets, Cow::Owned(_)));
}
#[test]
fn unsliced_compact_parts_borrow_offsets() {
let c = abcd();
let (data, offsets) = compact_parts(&c);
assert_eq!(data, b"aabbccdd");
assert_eq!(&*offsets, &[0i64, 2, 4, 6, 8]);
assert!(matches!(offsets, Cow::Borrowed(_)));
}
#[test]
fn empty_range_slice_of_frozen_container_is_empty() {
let c = abcd();
for (start, end) in [(2, 2), (3, 1), (10, 20)] {
let s = slice(&c, start, end);
assert_eq!(s.len(), 0);
assert_eq!(data_byte_len(&s), 0);
assert_eq!(compact_parts(&s).0, b"");
}
}
#[test]
fn filter_on_frozen_slice_keeps_right_rows() {
let s = slice(&abcd(), 1, 4);
let s = filter(&s, &BooleanBuffer::collect_bool(3, |i| i != 1));
assert_eq!(s.len(), 2);
assert_eq!(get(&s, 0), Some(b"bb".as_slice()));
assert_eq!(get(&s, 1), Some(b"dd".as_slice()));
assert_eq!(s.value_data(), b"bbdd");
assert_eq!(s.value_offsets(), &[0i64, 2, 4]);
}
}
}