vortex_expr/exprs/
between.rs

1use std::any::Any;
2use std::fmt::{Debug, Display};
3use std::sync::Arc;
4
5use vortex_array::ArrayRef;
6use vortex_array::compute::{BetweenOptions, between};
7use vortex_dtype::DType;
8use vortex_dtype::DType::Bool;
9use vortex_error::VortexResult;
10
11use crate::{AnalysisExpr, BinaryExpr, ExprRef, Scope, ScopeDType, VortexExpr};
12
13#[derive(Debug, Eq, Hash)]
14#[allow(clippy::derived_hash_with_manual_eq)]
15pub struct Between {
16    arr: ExprRef,
17    lower: ExprRef,
18    upper: ExprRef,
19    options: BetweenOptions,
20}
21
22impl Between {
23    pub fn between(
24        arr: ExprRef,
25        lower: ExprRef,
26        upper: ExprRef,
27        options: BetweenOptions,
28    ) -> ExprRef {
29        Arc::new(Self {
30            arr,
31            lower,
32            upper,
33            options,
34        })
35    }
36
37    pub fn to_binary_expr(&self) -> ExprRef {
38        let lhs = BinaryExpr::new_expr(
39            self.lower.clone(),
40            self.options.lower_strict.to_operator().into(),
41            self.arr.clone(),
42        );
43        let rhs = BinaryExpr::new_expr(
44            self.arr.clone(),
45            self.options.upper_strict.to_operator().into(),
46            self.upper.clone(),
47        );
48        BinaryExpr::new_expr(lhs, crate::Operator::And, rhs)
49    }
50}
51
52impl Display for Between {
53    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54        write!(
55            f,
56            "({} {} {} {} {})",
57            self.lower,
58            self.options.lower_strict.to_operator(),
59            self.arr,
60            self.options.upper_strict.to_operator(),
61            self.upper
62        )
63    }
64}
65
66impl PartialEq for Between {
67    fn eq(&self, other: &Between) -> bool {
68        self.arr.eq(&other.arr)
69            && other.lower.eq(&self.lower)
70            && other.upper.eq(&self.upper)
71            && self.options == other.options
72    }
73}
74
75#[cfg(feature = "proto")]
76pub(crate) mod proto {
77    use vortex_array::compute::{BetweenOptions, StrictComparison};
78    use vortex_error::{VortexResult, vortex_bail};
79    use vortex_proto::expr::kind;
80    use vortex_proto::expr::kind::Kind;
81
82    use crate::between::Between;
83    use crate::{ExprDeserialize, ExprRef, ExprSerializable, Id};
84
85    pub(crate) struct BetweenSerde;
86
87    impl Id for BetweenSerde {
88        fn id(&self) -> &'static str {
89            "between"
90        }
91    }
92
93    impl ExprSerializable for Between {
94        fn id(&self) -> &'static str {
95            BetweenSerde.id()
96        }
97
98        fn serialize_kind(&self) -> VortexResult<Kind> {
99            Ok(Kind::Between(kind::Between {
100                lower_strict: self.options.lower_strict == StrictComparison::Strict,
101                upper_strict: self.options.upper_strict == StrictComparison::Strict,
102            }))
103        }
104    }
105
106    impl ExprDeserialize for BetweenSerde {
107        fn deserialize(&self, kind: &Kind, children: Vec<ExprRef>) -> VortexResult<ExprRef> {
108            let Kind::Between(between) = kind else {
109                vortex_bail!("wrong kind {:?}, want between", kind)
110            };
111
112            Ok(Between::between(
113                children[0].clone(),
114                children[1].clone(),
115                children[2].clone(),
116                BetweenOptions {
117                    lower_strict: if between.lower_strict {
118                        StrictComparison::Strict
119                    } else {
120                        StrictComparison::NonStrict
121                    },
122                    upper_strict: if between.upper_strict {
123                        StrictComparison::Strict
124                    } else {
125                        StrictComparison::NonStrict
126                    },
127                },
128            ))
129        }
130    }
131}
132
133impl AnalysisExpr for Between {}
134
135impl VortexExpr for Between {
136    fn as_any(&self) -> &dyn Any {
137        self
138    }
139
140    fn unchecked_evaluate(&self, scope: &Scope) -> VortexResult<ArrayRef> {
141        let arr_val = self.arr.unchecked_evaluate(scope)?;
142        let lower_arr_val = self.lower.unchecked_evaluate(scope)?;
143        let upper_arr_val = self.upper.unchecked_evaluate(scope)?;
144
145        between(&arr_val, &lower_arr_val, &upper_arr_val, &self.options)
146    }
147
148    fn children(&self) -> Vec<&ExprRef> {
149        vec![&self.arr, &self.lower, &self.upper]
150    }
151
152    fn replacing_children(self: Arc<Self>, children: Vec<ExprRef>) -> ExprRef {
153        Arc::new(Self {
154            arr: children[0].clone(),
155            lower: children[1].clone(),
156            upper: children[2].clone(),
157            options: self.options.clone(),
158        })
159    }
160
161    fn return_dtype(&self, ctx: &ScopeDType) -> VortexResult<DType> {
162        let arr_dt = self.arr.return_dtype(ctx)?;
163        let lower_dt = self.lower.return_dtype(ctx)?;
164        let upper_dt = self.upper.return_dtype(ctx)?;
165
166        assert!(
167            arr_dt.eq_ignore_nullability(&lower_dt),
168            "array dtype {}, lower dtype {}",
169            arr_dt,
170            lower_dt
171        );
172        assert!(
173            arr_dt.eq_ignore_nullability(&upper_dt),
174            "array dtype {}, upper dtype {}",
175            arr_dt,
176            upper_dt
177        );
178
179        Ok(Bool(
180            arr_dt.nullability() | lower_dt.nullability() | upper_dt.nullability(),
181        ))
182    }
183}
vortex_expr/exprs/between.rs

vortex_expr/exprs/
between.rs
