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::Dict;
7use super::DictArray;
8use crate::ArrayRef;
9use crate::Canonical;
10use crate::ExecutionCtx;
11use crate::IntoArray;
12use crate::array::ArrayView;
13use crate::arrays::ConstantArray;
14use crate::arrays::dict::DictArrayExt;
15use crate::arrays::dict::DictArraySlotsExt;
16use crate::builtins::ArrayBuiltins;
17use crate::scalar_fn::fns::binary::CompareKernel;
18use crate::scalar_fn::fns::operators::CompareOperator;
19use crate::scalar_fn::fns::operators::Operator;
20
21impl CompareKernel for Dict {
22    fn compare(
23        lhs: ArrayView<'_, Dict>,
24        rhs: &ArrayRef,
25        operator: CompareOperator,
26        ctx: &mut ExecutionCtx,
27    ) -> VortexResult<Option<ArrayRef>> {
28        // if we have more values than codes, it is faster to canonicalise first.
29        if lhs.values().len() > lhs.codes().len() {
30            return Ok(None);
31        }
32
33        // If the RHS is constant, then we just need to compare against our encoded values.
34        if let Some(rhs) = rhs.as_constant() {
35            let compare_result = lhs.values().clone().binary(
36                ConstantArray::new(rhs, lhs.values().len()).into_array(),
37                Operator::from(operator),
38            )?;
39
40            // SAFETY: values len preserved, codes all still point to valid values
41            let result = unsafe {
42                DictArray::new_unchecked(lhs.codes().clone(), compare_result)
43                    .set_all_values_referenced(lhs.has_all_values_referenced())
44                    .into_array()
45            };
46
47            // We canonicalize the result because dictionary-encoded bools is dumb.
48            return Ok(Some(result.execute::<Canonical>(ctx)?.into_array()));
49        }
50
51        // It's a little more complex, but we could perform a comparison against the dictionary
52        // values in the future.
53        Ok(None)
54    }
55}