reifydb_value/util/
kernel.rs1use arrow_array::{Array, ArrayRef, BooleanArray, UInt64Array};
5use arrow_buffer::{BooleanBuffer, NullBuffer};
6use arrow_select::{
7 concat::concat,
8 filter::{FilterBuilder, FilterPredicate},
9 interleave::interleave,
10 take::take,
11 zip::zip,
12};
13
14use crate::{
15 Result,
16 error::{Diagnostic, Error},
17 util::bitmap,
18};
19
20pub fn predicate(mask: &BooleanBuffer, len: usize) -> FilterPredicate {
21 FilterBuilder::new(&BooleanArray::new(bitmap::resize(mask, len), None)).build()
22}
23
24pub fn shared_predicate(mask: &BooleanBuffer, len: usize) -> FilterPredicate {
25 FilterBuilder::new(&BooleanArray::new(bitmap::resize(mask, len), None)).optimize().build()
26}
27
28pub fn filtered<A>(array: &A, predicate: &FilterPredicate) -> A
29where
30 A: Array + Clone + 'static,
31{
32 let selected =
33 predicate.filter(array).expect("filter predicate is resized to the array length before it is built");
34 downcast(selected.as_ref())
35}
36
37pub fn kept_nulls<A>(source: &A, selected: &A) -> Option<NullBuffer>
38where
39 A: Array,
40{
41 match (source.logical_nulls().is_some(), selected.logical_nulls().is_some()) {
42 (true, false) => Some(NullBuffer::new_valid(selected.len())),
43 (true, true) => selected.logical_nulls(),
44 (false, _) => None,
45 }
46}
47
48pub fn rows_in_range(indices: &[usize], len: usize) -> Result<()> {
49 match indices.iter().find(|&&index| index >= len) {
50 Some(index) => Err(Error(Box::new(Diagnostic {
51 code: "INTERNAL_ERROR".to_string(),
52 message: format!("row index {index} out of range for a column of {len} rows"),
53 ..Diagnostic::default()
54 }))),
55 None => Ok(()),
56 }
57}
58
59pub fn taken<A>(array: &A, indices: &UInt64Array) -> A
60where
61 A: Array + Clone + 'static,
62{
63 let selected = take(array, indices, None).expect("every index is clamped into range before the take");
64 downcast(selected.as_ref())
65}
66
67pub fn interleaved<A>(source: &A, filler: &dyn Array, pairs: &[(usize, usize)]) -> A
68where
69 A: Array + Clone + 'static,
70{
71 downcast(picked(&[source as &dyn Array, filler], pairs).as_ref())
72}
73
74pub fn joined(arrays: &[&dyn Array]) -> ArrayRef {
75 concat(arrays).expect("all arrays share one type")
76}
77
78pub fn picked(arrays: &[&dyn Array], pairs: &[(usize, usize)]) -> ArrayRef {
79 interleave(arrays, pairs).expect("all arrays share one type and every pair is clamped into range")
80}
81
82pub fn merged(mask: &BooleanArray, truthy: &dyn Array, falsy: &dyn Array) -> ArrayRef {
83 zip(mask, &truthy, &falsy).expect("both sides share one type and the length of the mask")
84}
85
86pub fn downcast<A>(array: &dyn Array) -> A
87where
88 A: Array + Clone + 'static,
89{
90 array.as_any().downcast_ref::<A>().expect("an arrow kernel returns the array type it was given").clone()
91}