vortex_onpair/compute/
compare.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::ArrayRef;
5use vortex_array::ArrayView;
6use vortex_array::ExecutionCtx;
7use vortex_array::IntoArray;
8use vortex_array::arrays::BoolArray;
9use vortex_array::arrays::ConstantArray;
10use vortex_array::builtins::ArrayBuiltins;
11use vortex_array::dtype::DType;
12use vortex_array::scalar::Scalar;
13use vortex_array::scalar_fn::fns::binary::CompareKernel;
14use vortex_array::scalar_fn::fns::operators::CompareOperator;
15use vortex_buffer::BitBuffer;
16use vortex_error::VortexResult;
17
18use crate::OnPair;
19use crate::OnPairArraySlotsExt;
20
21impl CompareKernel for OnPair {
22    fn compare(
23        lhs: ArrayView<'_, Self>,
24        rhs: &ArrayRef,
25        operator: CompareOperator,
26        ctx: &mut ExecutionCtx,
27    ) -> VortexResult<Option<ArrayRef>> {
28        let Some(constant) = rhs.as_constant() else {
29            return Ok(None);
30        };
31        let is_empty = match constant.dtype() {
32            DType::Utf8(_) => constant.as_utf8().is_empty(),
33            DType::Binary(_) => constant.as_binary().is_empty(),
34            _ => return Ok(None),
35        };
36        if is_empty != Some(true) {
37            return Ok(None);
38        }
39
40        let lengths = lhs.uncompressed_lengths();
41        let buffer = match operator {
42            // every value is greater than an empty string
43            CompareOperator::Gte => BitBuffer::new_set(lhs.len()),
44            // no value is less than an empty string
45            CompareOperator::Lt => BitBuffer::new_unset(lhs.len()),
46            _ => lengths
47                .binary(
48                    ConstantArray::new(Scalar::zero_value(lengths.dtype()), lengths.len())
49                        .into_array(),
50                    operator.into(),
51                )?
52                .execute(ctx)?,
53        };
54        Ok(Some(
55            BoolArray::new(
56                buffer,
57                lhs.validity()?
58                    .union_nullability(constant.dtype().nullability()),
59            )
60            .into_array(),
61        ))
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use std::sync::LazyLock;
68
69    use rstest::rstest;
70    use vortex_array::IntoArray;
71    use vortex_array::VortexSessionExecute;
72    use vortex_array::array_session;
73    use vortex_array::arrays::BoolArray;
74    use vortex_array::arrays::ConstantArray;
75    use vortex_array::arrays::VarBinArray;
76    use vortex_array::assert_arrays_eq;
77    use vortex_array::builtins::ArrayBuiltins;
78    use vortex_array::dtype::DType;
79    use vortex_array::dtype::Nullability;
80    use vortex_array::scalar::Scalar;
81    use vortex_array::scalar_fn::fns::operators::Operator;
82    use vortex_error::VortexResult;
83    use vortex_session::VortexSession;
84
85    use crate::compress::DEFAULT_DICT12_CONFIG;
86    use crate::compress::onpair_compress;
87
88    static SESSION: LazyLock<VortexSession> = LazyLock::new(array_session);
89
90    #[cfg_attr(miri, ignore)]
91    #[rstest]
92    #[case(Operator::Eq, [true, false, true, false])]
93    #[case(Operator::NotEq, [false, true, false, true])]
94    #[case(Operator::Gt, [false, true, false, true])]
95    #[case(Operator::Gte, [true, true, true, true])]
96    #[case(Operator::Lt, [false, false, false, false])]
97    #[case(Operator::Lte, [true, false, true, false])]
98    fn compare_empty_string(#[case] op: Operator, #[case] expected: [bool; 4]) -> VortexResult<()> {
99        let input = VarBinArray::from_iter(
100            [Some(""), Some("a"), Some(""), Some("bbb")],
101            DType::Utf8(Nullability::NonNullable),
102        );
103        let mut ctx = SESSION.create_execution_ctx();
104        let arr = onpair_compress(input.as_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?.into_array();
105
106        let result = arr
107            .binary(ConstantArray::new("", input.len()).into_array(), op)?
108            .execute::<BoolArray>(&mut ctx)?;
109        assert_arrays_eq!(&result, &BoolArray::from_iter(expected), &mut ctx);
110        Ok(())
111    }
112
113    #[cfg_attr(miri, ignore)]
114    #[test]
115    fn compare_empty_string_nullable() -> VortexResult<()> {
116        let input = VarBinArray::from_iter(
117            [Some(""), None, Some("x")],
118            DType::Utf8(Nullability::Nullable),
119        );
120        let mut ctx = SESSION.create_execution_ctx();
121        let arr = onpair_compress(input.as_array(), DEFAULT_DICT12_CONFIG, &mut ctx)?.into_array();
122
123        let eq_empty = arr
124            .clone()
125            .binary(ConstantArray::new("", arr.len()).into_array(), Operator::Eq)?
126            .execute::<BoolArray>(&mut ctx)?;
127        assert_arrays_eq!(
128            &eq_empty,
129            &BoolArray::from_iter([Some(true), None, Some(false)]),
130            &mut ctx
131        );
132
133        let null_rhs =
134            ConstantArray::new(Scalar::null(DType::Utf8(Nullability::Nullable)), arr.len());
135        let eq_null = arr
136            .binary(null_rhs.into_array(), Operator::Eq)?
137            .execute::<BoolArray>(&mut ctx)?;
138        assert_arrays_eq!(
139            &eq_null,
140            &BoolArray::from_iter([None::<bool>, None, None]),
141            &mut ctx
142        );
143        Ok(())
144    }
145}
vortex_onpair/compute/compare.rs

vortex_onpair/compute/
compare.rs
