Skip to main content

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::DynArray;
8use crate::ExecutionCtx;
9use crate::IntoArray;
10use crate::arrays::ConstantArray;
11use crate::arrays::ExtensionArray;
12use crate::arrays::ExtensionVTable;
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 ExtensionVTable {
19    fn compare(
20        lhs: &ExtensionArray,
21        rhs: &ArrayRef,
22        operator: CompareOperator,
23        _ctx: &mut ExecutionCtx,
24    ) -> VortexResult<Option<ArrayRef>> {
25        // If the RHS is a constant, we can extract the storage scalar.
26        if let Some(const_ext) = rhs.as_constant() {
27            let storage_scalar = const_ext.as_extension().to_storage_scalar();
28            return lhs
29                .storage()
30                .to_array()
31                .binary(
32                    ConstantArray::new(storage_scalar, lhs.len()).into_array(),
33                    Operator::from(operator),
34                )
35                .map(Some);
36        }
37
38        // If the RHS is an extension array matching ours, we can extract the storage.
39        if let Some(rhs_ext) = rhs.as_opt::<ExtensionVTable>() {
40            return lhs
41                .storage()
42                .to_array()
43                .binary(rhs_ext.storage().to_array(), Operator::from(operator))
44                .map(Some);
45        }
46
47        // Otherwise, we need the RHS to handle this comparison.
48        Ok(None)
49    }
50}