Skip to main content

vortex_array/arrays/dict/compute/
compare.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use super::DictArray;
7use super::DictVTable;
8use crate::Array;
9use crate::ArrayRef;
10use crate::ExecutionCtx;
11use crate::IntoArray;
12use crate::arrays::ConstantArray;
13use crate::compute::Operator;
14use crate::compute::compare;
15use crate::expr::CompareKernel;
16
17impl CompareKernel for DictVTable {
18    fn compare(
19        lhs: &DictArray,
20        rhs: &dyn Array,
21        operator: Operator,
22        _ctx: &mut ExecutionCtx,
23    ) -> VortexResult<Option<ArrayRef>> {
24        // if we have more values than codes, it is faster to canonicalise first.
25        if lhs.values().len() > lhs.codes().len() {
26            return Ok(None);
27        }
28
29        // If the RHS is constant, then we just need to compare against our encoded values.
30        if let Some(rhs) = rhs.as_constant() {
31            let compare_result = compare(
32                lhs.values(),
33                ConstantArray::new(rhs, lhs.values().len()).as_ref(),
34                operator,
35            )?;
36
37            // SAFETY: values len preserved, codes all still point to valid values
38            let result = unsafe {
39                DictArray::new_unchecked(lhs.codes().clone(), compare_result)
40                    .set_all_values_referenced(lhs.has_all_values_referenced())
41                    .into_array()
42            };
43
44            // We canonicalize the result because dictionary-encoded bools is dumb.
45            return Ok(Some(result.to_canonical()?.into_array()));
46        }
47
48        // It's a little more complex, but we could perform a comparison against the dictionary
49        // values in the future.
50        Ok(None)
51    }
52}