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::{CompareFn, Operator, compare, slice};
6use crate::{Array, ArrayRef};
7
8impl CompareFn<&ChunkedArray> 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
35#[cfg(test)]
36mod tests {
37    use super::*;
38    use crate::arrays::PrimitiveArray;
39
40    #[test]
41    fn empty_compare() {
42        let base = PrimitiveArray::from_iter(Vec::<u32>::new()).into_array();
43        let chunked =
44            ChunkedArray::try_new(vec![base.clone(), base.clone()], base.dtype().clone()).unwrap();
45        let chunked_empty = ChunkedArray::try_new(vec![], base.dtype().clone()).unwrap();
46
47        let r = compare(&chunked, &chunked_empty, Operator::Eq).unwrap();
48
49        assert!(r.is_empty());
50    }
51}