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
128#[derive(Clone, Debug)]
129pub struct DynamicComparisonExpr {
130    pub(crate) operator: CompareOperator,
131    pub(crate) rhs: Arc<Rhs>,
132    // Default value for the dynamic comparison.
133    pub(crate) default: bool,
134}
135
136impl DynamicComparisonExpr {
137    pub fn scalar(&self) -> Option<Scalar> {
138        (self.rhs.value)().map(|v| {
139            Scalar::try_new(self.rhs.dtype.clone(), Some(v))
140                .vortex_expect("`DynamicComparisonExpr` was invalid")
141        })
142    }
143}
144
145impl Display for DynamicComparisonExpr {
146    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
147        write!(
148            f,
149            "{} {}",
150            self.operator,
151            self.scalar()
152                .map_or_else(|| "<none>".to_string(), |v| v.to_string())
153        )
154    }
155}
156
157impl PartialEq for DynamicComparisonExpr {
158    fn eq(&self, other: &Self) -> bool {
159        self.operator == other.operator
160            && Arc::ptr_eq(&self.rhs, &other.rhs)
161            && self.default == other.default
162    }
163}
164impl Eq for DynamicComparisonExpr {}
165
166impl Hash for DynamicComparisonExpr {
167    fn hash<H: Hasher>(&self, state: &mut H) {
168        self.operator.hash(state);
169        Arc::as_ptr(&self.rhs).hash(state);
170        self.default.hash(state);
171    }
172}
173
174/// Hash and PartialEq are implemented based on the ptr of the value function, such that the
175/// internal value doesn't impact the hash of an expression tree.
176pub(crate) struct Rhs {
177    // The right-hand side value is a function that returns an `Option<ScalarValue>`.
178    pub(crate) value: Arc<dyn Fn() -> Option<ScalarValue> + Send + Sync>,
179    // The data type of the right-hand side value.
180    pub(crate) dtype: DType,
181}
182
183impl Rhs {
184    pub fn scalar(&self) -> Option<Scalar> {
185        (self.value)().map(|v| {
186            Scalar::try_new(self.dtype.clone(), Some(v)).vortex_expect("`Rhs` was invalid")
187        })
188    }
189}
190
191impl Debug for Rhs {
192    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
193        f.debug_struct("Rhs")
194            .field("value", &"<dyn Fn() -> Option<ScalarValue> + Send + Sync>")
195            .field("dtype", &self.dtype)
196            .finish()
197    }
198}
199
200/// A utility for checking whether any dynamic expressions have been updated.
201pub struct DynamicExprUpdates {
202    exprs: Box<[DynamicComparisonExpr]>,
203    // Track the latest observed versions of each dynamic expression, along with a version counter.
204    prev_versions: Mutex<(u64, Vec<Option<Scalar>>)>,
205}
206
207impl DynamicExprUpdates {
208    /// Track dynamic scalar functions contained in a bound expression tree.
209    pub fn new(expr: &BoundExpression) -> Option<Self> {
210        #[derive(Default)]
211        struct Visitor(Vec<DynamicComparisonExpr>);
212
213        impl NodeVisitor<'_> for Visitor {
214            type NodeTy = BoundExpression;
215
216            fn visit_down(&mut self, node: &'_ Self::NodeTy) -> VortexResult<TraversalOrder> {
217                if let Some(dynamic) = node
218                    .as_scalar()
219                    .and_then(|scalar_fn| scalar_fn.as_opt::<DynamicComparison>())
220                {
221                    self.0.push(dynamic.clone());
222                }
223                Ok(TraversalOrder::Continue)
224            }
225        }
226
227        let mut visitor = Visitor::default();
228        expr.accept(&mut visitor).vortex_expect("Infallible");
229
230        if visitor.0.is_empty() {
231            return None;
232        }
233
234        let exprs = visitor.0.into_boxed_slice();
235        let prev_versions = exprs
236            .iter()
237            .map(|expr| {
238                (expr.rhs.value)().map(|v| {
239                    Scalar::try_new(expr.rhs.dtype.clone(), Some(v))
240                        .vortex_expect("`DynamicExprUpdates` was invalid")
241                })
242            })
243            .collect();
244
245        Some(Self {
246            exprs,
247            prev_versions: Mutex::new((0, prev_versions)),
248        })
249    }
250
251    pub fn version(&self) -> u64 {
252        let mut guard = self.prev_versions.lock();
253
254        let mut updated = false;
255        for (i, expr) in self.exprs.iter().enumerate() {
256            let current = expr.scalar();
257            if current != guard.1[i] {
258                // At least one expression has been updated.
259                // We don't bail out early in order to avoid false positives for future calls
260                // to `is_updated`.
261                updated = true;
262                guard.1[i] = current;
263            }
264        }
265
266        if updated {
267            guard.0 += 1;
268        }
269
270        guard.0
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use std::sync::atomic::AtomicI32;
277    use std::sync::atomic::Ordering;
278
279    use vortex_buffer::buffer;
280    use vortex_error::VortexResult;
281
282    use super::*;
283    use crate::IntoArray;
284    use crate::VortexSessionExecute;
285    use crate::array_session;
286    use crate::arrays::BoolArray;
287    use crate::assert_arrays_eq;
288    use crate::dtype::DType;
289    use crate::dtype::Nullability;
290    use crate::dtype::PType;
291    use crate::expr::dynamic;
292    use crate::expr::root;
293
294    #[test]
295    fn is_not_strict() {
296        let expr = dynamic(
297            CompareOperator::Lt,
298            || None,
299            DType::Primitive(PType::I32, Nullability::NonNullable),
300            true,
301            root(),
302        );
303
304        assert!(!expr.signature().is_strict());
305    }
306
307    #[test]
308    fn return_dtype_bool() -> VortexResult<()> {
309        let expr = dynamic(
310            CompareOperator::Lt,
311            || Some(5i32.into()),
312            DType::Primitive(PType::I32, Nullability::NonNullable),
313            true,
314            root(),
315        );
316        let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
317        assert_eq!(
318            expr.return_dtype(&input_dtype)?,
319            DType::Bool(Nullability::NonNullable)
320        );
321        Ok(())
322    }
323
324    #[test]
325    fn execute_with_value() -> VortexResult<()> {
326        let mut ctx = array_session().create_execution_ctx();
327        let input = buffer![1i32, 5, 10].into_array();
328        let expr = dynamic(
329            CompareOperator::Lt,
330            || Some(5i32.into()),
331            DType::Primitive(PType::I32, Nullability::NonNullable),
332            true,
333            root(),
334        );
335        let result = input.apply(&expr)?;
336        assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx);
337        Ok(())
338    }
339
340    #[test]
341    fn execute_without_value_default_true() -> VortexResult<()> {
342        let mut ctx = array_session().create_execution_ctx();
343        let input = buffer![1i32, 5, 10].into_array();
344        let expr = dynamic(
345            CompareOperator::Lt,
346            || None,
347            DType::Primitive(PType::I32, Nullability::NonNullable),
348            true,
349            root(),
350        );
351        let result = input.apply(&expr)?;
352        assert_arrays_eq!(result, BoolArray::from_iter([true, true, true]), &mut ctx);
353        Ok(())
354    }
355
356    #[test]
357    fn execute_without_value_default_false() -> VortexResult<()> {
358        let mut ctx = array_session().create_execution_ctx();
359        let input = buffer![1i32, 5, 10].into_array();
360        let expr = dynamic(
361            CompareOperator::Lt,
362            || None,
363            DType::Primitive(PType::I32, Nullability::NonNullable),
364            false,
365            root(),
366        );
367        let result = input.apply(&expr)?;
368        assert_arrays_eq!(
369            result,
370            BoolArray::from_iter([false, false, false]),
371            &mut ctx
372        );
373        Ok(())
374    }
375
376    #[test]
377    fn execute_value_flips() -> VortexResult<()> {
378        let mut ctx = array_session().create_execution_ctx();
379        let threshold = Arc::new(AtomicI32::new(5));
380        let threshold_clone = Arc::clone(&threshold);
381        let expr = dynamic(
382            CompareOperator::Lt,
383            move || Some(threshold_clone.load(Ordering::SeqCst).into()),
384            DType::Primitive(PType::I32, Nullability::NonNullable),
385            true,
386            root(),
387        );
388        let input = buffer![1i32, 5, 10].into_array();
389
390        let result = input.clone().apply(&expr)?;
391        assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx);
392
393        threshold.store(10, Ordering::SeqCst);
394        let result = input.apply(&expr)?;
395        assert_arrays_eq!(result, BoolArray::from_iter([true, true, false]), &mut ctx);
396
397        Ok(())
398    }
399}
vortex_array/scalar_fn/fns/dynamic.rs

vortex_array/scalar_fn/fns/
dynamic.rs
