vortex_array/arrays/chunked/compute/
compare.rs

1use vortex_error::VortexResult;
2
3use crate::arrays::{ChunkedArray, ChunkedVTable};
4use crate::builders::{ArrayBuilder, BoolBuilder};
5use crate::compute::{CompareKernel, CompareKernelAdapter, Operator, compare};
6use crate::{Array, ArrayRef, register_kernel};
7
8impl CompareKernel for ChunkedVTable {
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 = rhs.slice(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(ChunkedVTable).lift());
36
37#[cfg(test)]
38mod tests {
39    use super::*;
40    use crate::IntoArray;
41    use crate::arrays::PrimitiveArray;
42
43    #[test]
44    fn empty_compare() {
45        let base = PrimitiveArray::from_iter(Vec::<u32>::new()).into_array();
46        let chunked =
47            ChunkedArray::try_new(vec![base.clone(), base.clone()], base.dtype().clone()).unwrap();
48        let chunked_empty = ChunkedArray::try_new(vec![], base.dtype().clone()).unwrap();
49
50        let r = compare(chunked.as_ref(), chunked_empty.as_ref(), Operator::Eq).unwrap();
51
52        assert!(r.is_empty());
53    }
54}
vortex_array/arrays/chunked/compute/compare.rs

vortex_array/arrays/chunked/compute/
compare.rs
