Skip to main content

reifydb_value/util/
kernel.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use 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}