vortex_array/scalar_fn/fns/
dynamic.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5use std::fmt::Display;
6use std::fmt::Formatter;
7use std::hash::Hash;
8use std::hash::Hasher;
9use std::sync::Arc;
10
11use parking_lot::Mutex;
12use vortex_error::VortexExpect;
13use vortex_error::VortexResult;
14use vortex_error::vortex_bail;
15use vortex_session::registry::CachedId;
16
17use crate::ArrayRef;
18use crate::ExecutionCtx;
19use crate::IntoArray;
20use crate::arrays::ConstantArray;
21use crate::dtype::DType;
22use crate::expr::BoundExpression;
23use crate::expr::display::ExprDisplay;
24use crate::expr::traversal::NodeExt;
25use crate::expr::traversal::NodeVisitor;
26use crate::expr::traversal::TraversalOrder;
27use crate::scalar::Scalar;
28use crate::scalar::ScalarValue;
29use crate::scalar_fn::Arity;
30use crate::scalar_fn::ChildName;
31use crate::scalar_fn::ExecutionArgs;
32use crate::scalar_fn::ScalarFnId;
33use crate::scalar_fn::ScalarFnVTable;
34use crate::scalar_fn::ScalarFnVTableExt;
35use crate::scalar_fn::VecExecutionArgs;
36use crate::scalar_fn::fns::binary::Binary;
37use crate::scalar_fn::fns::operators::CompareOperator;
38use crate::scalar_fn::fns::operators::Operator;
39
40/// A dynamic comparison expression can be used to capture a comparison to a value that can change
41/// during the execution of a query, such as when a compute engine pushes down an ORDER BY + LIMIT
42/// operation and is able to progressively tighten the bounds of the filter.
43#[derive(Clone)]
44pub struct DynamicComparison;
45
46impl ScalarFnVTable for DynamicComparison {
47    type Options = DynamicComparisonExpr;
48
49    fn id(&self) -> ScalarFnId {
50        static ID: CachedId = CachedId::new("vortex.dynamic");
51        *ID
52    }
53
54    fn arity(&self, _options: &Self::Options) -> Arity {
55        Arity::Exact(1)
56    }
57
58    fn child_name(&self, _instance: &Self::Options, child_idx: usize) -> ChildName {
59        match child_idx {
60            0 => ChildName::from("lhs"),
61            _ => unreachable!(),
62        }
63    }
64
65    fn fmt_sql(
66        &self,
67        dynamic: &DynamicComparisonExpr,
68        expr: &dyn ExprDisplay,
69        f: &mut Formatter<'_>,
70    ) -> std::fmt::Result {
71        Display::fmt(expr.display_child(0), f)?;
72        write!(f, " {} dynamic(", dynamic.operator)?;
73        match dynamic.scalar() {
74            None => write!(f, "scalar=<none>")?,
75            Some(scalar) => write!(f, "scalar={scalar}")?,
76        }
77        write!(f, ")")
78    }
79
80    fn return_dtype(
81        &self,
82        dynamic: &DynamicComparisonExpr,
83        arg_dtypes: &[DType],
84    ) -> VortexResult<DType> {
85        let lhs = &arg_dtypes[0];
86        if !dynamic.rhs.dtype.eq_ignore_nullability(lhs) {
87            vortex_bail!(
88                "Incompatible dtypes for dynamic comparison: expected {} (ignore nullability) but got {}",
89                &dynamic.rhs.dtype,
90                lhs
91            );
92        }
93        Ok(DType::Bool(
94            lhs.nullability() | dynamic.rhs.dtype.nullability(),
95        ))
96    }
97
98    fn execute(
99        &self,
100        data: &Self::Options,
101        args: &dyn ExecutionArgs,
102        ctx: &mut ExecutionCtx,
103    ) -> VortexResult<ArrayRef> {
104        if let Some(scalar) = data.rhs.scalar() {
105            let lhs = args.get(0)?;
106            let rhs = ConstantArray::new(scalar, args.row_count()).into_array();
107
108            let delegate_args = VecExecutionArgs::new(vec![lhs, rhs], args.row_count());
109            return Binary
110                .bind(Operator::from(data.operator))
111                .execute(&delegate_args, ctx);
112        }
113        let ret_dtype =
114            DType::Bool(args.get(0)?.dtype().nullability() | data.rhs.dtype.nullability());
115
116        Ok(ConstantArray::new(
117            Scalar::try_new(ret_dtype, Some(data.default.into()))?,
118            args.row_count(),
119        )
120        .into_array())
121    }
122
123    fn is_strict(&self, _options: &Self::Options) -> bool {
124        false
125    }
126
127    /// A dynamic comparison is a comparison, or a constant while it has no value to compare
128    /// against, and neither can fail.
129    fn is_infallible(&self, _options: &Self::Options) -> bool {
130        true
131    }
132}
133
134#[derive(Clone, Debug)]
135pub struct DynamicComparisonExpr {
136    pub(crate) operator: CompareOperator,
137    pub(crate) rhs: Arc<Rhs>,
138    // Default value for the dynamic comparison.
139    pub(crate) default: bool,
140}
141
142impl DynamicComparisonExpr {
143    pub fn scalar(&self) -> Option<Scalar> {
144        (self.rhs.value)().map(|v| {
145            Scalar::try_new(self.rhs.dtype.clone(), Some(v))
146                .vortex_expect("`DynamicComparisonExpr` was invalid")
147        })
148    }
149}
150
151impl Display for DynamicComparisonExpr {
152    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
153        write!(
154            f,
155            "{} {}",
156            self.operator,
157            self.scalar()
158                .map_or_else(|| "<none>".to_string(), |v| v.to_string())
159        )
160    }
161}
162
163impl PartialEq for DynamicComparisonExpr {
164    fn eq(&self, other: &Self) -> bool {
165        self.operator == other.operator
166            && Arc::ptr_eq(&self.rhs, &other.rhs)
167            && self.default == other.default
168    }
169}
170impl Eq for DynamicComparisonExpr {}
171
172impl Hash for DynamicComparisonExpr {
173    fn hash<H: Hasher>(&self, state: &mut H) {
174        self.operator.hash(state);
175        Arc::as_ptr(&self.rhs).hash(state);
176        self.default.hash(state);
177    }
178}
179
180/// Hash and PartialEq are implemented based on the ptr of the value function, such that the
181/// internal value doesn't impact the hash of an expression tree.
182pub(crate) struct Rhs {
183    // The right-hand side value is a function that returns an `Option<ScalarValue>`.
184    pub(crate) value: Arc<dyn Fn() -> Option<ScalarValue> + Send + Sync>,
185    // The data type of the right-hand side value.
186    pub(crate) dtype: DType,
187}
188
189impl Rhs {
190    pub fn scalar(&self) -> Option<Scalar> {
191        (self.value)().map(|v| {
192            Scalar::try_new(self.dtype.clone(), Some(v)).vortex_expect("`Rhs` was invalid")
193        })
194    }
195}
196
197impl Debug for Rhs {
198    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
199        f.debug_struct("Rhs")
200            .field("value", &"<dyn Fn() -> Option<ScalarValue> + Send + Sync>")
201            .field("dtype", &self.dtype)
202            .finish()
203    }
204}
205
206/// A utility for checking whether any dynamic expressions have been updated.
207pub struct DynamicExprUpdates {
208    exprs: Box<[DynamicComparisonExpr]>,
209    // Track the latest observed versions of each dynamic expression, along with a version counter.
210    prev_versions: Mutex<(u64, Vec<Option<Scalar>>)>,
211}
212
213impl DynamicExprUpdates {
214    /// Track dynamic scalar functions contained in a bound expression tree.
215    pub fn new(expr: &BoundExpression) -> Option<Self> {
216        #[derive(Default)]
217        struct Visitor(Vec<DynamicComparisonExpr>);
218
219        impl NodeVisitor<'_> for Visitor {
220            type NodeTy = BoundExpression;
221
222            fn visit_down(&mut self, node: &'_ Self::NodeTy) -> VortexResult<TraversalOrder> {
223                if let Some(dynamic) = node
224                    .as_scalar()
225                    .and_then(|scalar_fn| scalar_fn.as_opt::<DynamicComparison>())
226                {
227                    self.0.push(dynamic.clone());
228                }
229                Ok(TraversalOrder::Continue)
230            }
231        }
232
233        let mut visitor = Visitor::default();
234        expr.accept(&mut visitor).vortex_expect("Infallible");
235
236        if visitor.0.is_empty() {
237            return None;
238        }
239
240        let exprs = visitor.0.into_boxed_slice();
241        let prev_versions = exprs
242            .iter()
243            .map(|expr| {
244                (expr.rhs.value)().map(|v| {
245                    Scalar::try_new(expr.rhs.dtype.clone(), Some(v))
246                        .vortex_expect("`DynamicExprUpdates` was invalid")
247                })
248            })
249            .collect();
250
251        Some(Self {
252            exprs,
253            prev_versions: Mutex::new((0, prev_versions)),
254        })
255    }
256
257    pub fn version(&self) -> u64 {
258        let mut guard = self.prev_versions.lock();
259
260        let mut updated = false;
261        for (i, expr) in self.exprs.iter().enumerate() {
262            let current = expr.scalar();
263            if current != guard.1[i] {
264                // At least one expression has been updated.
265                // We don't bail out early in order to avoid false positives for future calls
266                // to `is_updated`.
267                updated = true;
268                guard.1[i] = current;
269            }
270        }
271
272        if updated {
273            guard.0 += 1;
274        }
275
276        guard.0
277    }
278}
279
280#[cfg(test)]
281mod tests {
282    use std::sync::atomic::AtomicI32;
283    use std::sync::atomic::Ordering;
284
285    use vortex_buffer::buffer;
286    use vortex_error::VortexResult;
287
288    use super::*;
289    use crate::IntoArray;
290    use crate::VortexSessionExecute;
291    use crate::array_session;
292    use crate::arrays::BoolArray;
293    use crate::assert_arrays_eq;
294    use crate::dtype::DType;
295    use crate::dtype::Nullability;
296    use crate::dtype::PType;
297    use crate::expr::dynamic;
298    use crate::expr::root;
299
300    #[test]
301    fn is_not_strict() {
302        let expr = dynamic(
303            CompareOperator::Lt,
304            || None,
305            DType::Primitive(PType::I32, Nullability::NonNullable),
306            true,
307            root(),
308        );
309
310        assert!(!expr.as_scalar().is_some_and(|f| f.signature().is_strict()));
311    }
312
313    #[test]
314    fn return_dtype_bool() -> VortexResult<()> {
315        let expr = dynamic(
316            CompareOperator::Lt,
317            || Some(5i32.into()),
318            DType::Primitive(PType::I32, Nullability::NonNullable),
319            true,
320            root(),
321        );
322        let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
323        assert_eq!(
324            expr.return_dtype(&input_dtype)?,
325            DType::Bool(Nullability::NonNullable)
326        );
327        Ok(())
328    }
329
330    #[test]
331    fn execute_with_value() -> VortexResult<()> {
332        let mut ctx = array_session().create_execution_ctx();
333        let input = buffer![1i32, 5, 10].into_array();
334        let expr = dynamic(
335            CompareOperator::Lt,
336            || Some(5i32.into()),
337            DType::Primitive(PType::I32, Nullability::NonNullable),
338            true,
339            root(),
340        );
341        let result = input.apply(&expr)?;
342        assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx);
343        Ok(())
344    }
345
346    #[test]
347    fn execute_without_value_default_true() -> VortexResult<()> {
348        let mut ctx = array_session().create_execution_ctx();
349        let input = buffer![1i32, 5, 10].into_array();
350        let expr = dynamic(
351            CompareOperator::Lt,
352            || None,
353            DType::Primitive(PType::I32, Nullability::NonNullable),
354            true,
355            root(),
356        );
357        let result = input.apply(&expr)?;
358        assert_arrays_eq!(result, BoolArray::from_iter([true, true, true]), &mut ctx);
359        Ok(())
360    }
361
362    #[test]
363    fn execute_without_value_default_false() -> VortexResult<()> {
364        let mut ctx = array_session().create_execution_ctx();
365        let input = buffer![1i32, 5, 10].into_array();
366        let expr = dynamic(
367            CompareOperator::Lt,
368            || None,
369            DType::Primitive(PType::I32, Nullability::NonNullable),
370            false,
371            root(),
372        );
373        let result = input.apply(&expr)?;
374        assert_arrays_eq!(
375            result,
376            BoolArray::from_iter([false, false, false]),
377            &mut ctx
378        );
379        Ok(())
380    }
381
382    #[test]
383    fn execute_value_flips() -> VortexResult<()> {
384        let mut ctx = array_session().create_execution_ctx();
385        let threshold = Arc::new(AtomicI32::new(5));
386        let threshold_clone = Arc::clone(&threshold);
387        let expr = dynamic(
388            CompareOperator::Lt,
389            move || Some(threshold_clone.load(Ordering::SeqCst).into()),
390            DType::Primitive(PType::I32, Nullability::NonNullable),
391            true,
392            root(),
393        );
394        let input = buffer![1i32, 5, 10].into_array();
395
396        let result = input.clone().apply(&expr)?;
397        assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx);
398
399        threshold.store(10, Ordering::SeqCst);
400        let result = input.apply(&expr)?;
401        assert_arrays_eq!(result, BoolArray::from_iter([true, true, false]), &mut ctx);
402
403        Ok(())
404    }
405}
vortex_array/scalar_fn/fns/dynamic.rs

vortex_array/scalar_fn/fns/
dynamic.rs
