vortex_array/arrays/chunked/compute/
compare.rs

1use vortex_error::VortexResult;
2
3use crate::arrays::{ChunkedArray, ChunkedEncoding};
4use crate::builders::{ArrayBuilder, BoolBuilder};
5use crate::compute::{CompareKernel, CompareKernelAdapter, Operator, compare, slice};
6use crate::{Array, ArrayRef, register_kernel};
7
8impl CompareKernel for ChunkedEncoding {
9    fn compare(
10        &self,
11        lhs: &ChunkedArray,
12        rhs: &dyn Array,
13        operator: Operator,
14    ) -> VortexResult<Option<ArrayRef>> {
15        let mut idx = 0;
16
17        let mut bool_builder = BoolBuilder::with_capacity(
18            // nullable <= non-nullable
19            (lhs.dtype().is_nullable() || rhs.dtype().is_nullable()).into(),
20            lhs.len(),
21        );
22
23        for chunk in lhs.non_empty_chunks() {
24            let sliced = slice(rhs, idx, idx + chunk.len())?;
25            let cmp_result = compare(chunk, &sliced, operator)?;
26
27            bool_builder.extend_from_array(&cmp_result)?;
28            idx += chunk.len();
29        }
30
31        Ok(Some(bool_builder.finish()))
32    }
33}
34
35register_kernel!(CompareKernelAdapter(ChunkedEncoding).lift());
36
37#[cfg(test)]
38mod tests {
39    use super::*;
40    use crate::arrays::PrimitiveArray;
41
42    #[test]
43    fn empty_compare() {
44        let base = PrimitiveArray::from_iter(Vec::<u32>::new()).into_array();
45        let chunked =
46            ChunkedArray::try_new(vec![base.clone(), base.clone()], base.dtype().clone()).unwrap();
47        let chunked_empty = ChunkedArray::try_new(vec![], base.dtype().clone()).unwrap();
48
49        let r = compare(&chunked, &chunked_empty, Operator::Eq).unwrap();
50
51        assert!(r.is_empty());
52    }
53}
vortex_array/arrays/chunked/compute/compare.rs

vortex_array/arrays/chunked/compute/
compare.rs
