vortex_array/arrays/extension/compute/
compare.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::ArrayRef;
7use crate::ExecutionCtx;
8use crate::IntoArray;
9use crate::array::ArrayView;
10use crate::arrays::ConstantArray;
11use crate::arrays::Extension;
12use crate::arrays::extension::ExtensionArrayExt;
13use crate::builtins::ArrayBuiltins;
14use crate::scalar_fn::fns::binary::CompareKernel;
15use crate::scalar_fn::fns::operators::CompareOperator;
16use crate::scalar_fn::fns::operators::Operator;
17
18impl CompareKernel for Extension {
19    fn compare(
20        lhs: ArrayView<'_, Extension>,
21        rhs: &ArrayRef,
22        operator: CompareOperator,
23        _ctx: &mut ExecutionCtx,
24    ) -> VortexResult<Option<ArrayRef>> {
25        // Storage values are only comparable when both sides share the same extension dtype
26        // (e.g. timestamps in different units must not compare their raw storage).
27        if !lhs.dtype().eq_ignore_nullability(rhs.dtype()) {
28            return Ok(None);
29        }
30
31        // If the RHS is a constant, we can extract the storage scalar.
32        if let Some(const_ext) = rhs.as_constant() {
33            let storage_scalar = const_ext.as_extension().to_storage_scalar();
34            return lhs
35                .storage_array()
36                .clone()
37                .binary(
38                    ConstantArray::new(storage_scalar, lhs.len()).into_array(),
39                    Operator::from(operator),
40                )
41                .map(Some);
42        }
43
44        // If the RHS is an extension array matching ours, we can extract the storage.
45        if let Some(rhs_ext) = rhs.as_opt::<Extension>() {
46            return lhs
47                .storage_array()
48                .clone()
49                .binary(rhs_ext.storage_array().clone(), Operator::from(operator))
50                .map(Some);
51        }
52
53        // Otherwise, we need the RHS to handle this comparison.
54        Ok(None)
55    }
56}
vortex_array/arrays/extension/compute/compare.rs

vortex_array/arrays/extension/compute/
compare.rs
