datafusion_physical_expr/intervals/
cp_solver.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Constraint propagator/solver for custom [`PhysicalExpr`] graphs.
19//!
20//! The constraint propagator/solver in DataFusion uses interval arithmetic to
21//! perform mathematical operations on intervals, which represent a range of
22//! possible values rather than a single point value. This allows for the
23//! propagation of ranges through mathematical operations, and can be used to
24//! compute bounds for a complicated expression. The key idea is that by
25//! breaking down a complicated expression into simpler terms, and then
26//! combining the bounds for those simpler terms, one can obtain bounds for the
27//! overall expression.
28//!
29//! This way of using interval arithmetic to compute bounds for a complex
30//! expression by combining the bounds for the constituent terms within the
31//! original expression allows us to reason about the range of possible values
32//! of the expression. This information later can be used in range pruning of
33//! the provably unnecessary parts of `RecordBatch`es.
34//!
35//! # Example
36//!
37//! For example, consider a mathematical expression such as `x^2 + y = 4` \[1\].
38//! Since this expression would be a binary tree in [`PhysicalExpr`] notation,
39//! this type of an hierarchical computation is well-suited for a graph based
40//! implementation. In such an implementation, an equation system `f(x) = 0` is
41//! represented by a directed acyclic expression graph (DAEG).
42//!
43//! In order to use interval arithmetic to compute bounds for this expression,
44//! one would first determine intervals that represent the possible values of
45//! `x` and `y` Let's say that the interval for `x` is `[1, 2]` and the interval
46//! for `y` is `[-3, 1]`. In the chart below, you can see how the computation
47//! takes place.
48//!
49//! # References
50//!
51//! 1. Kabak, Mehmet Ozan. Analog Circuit Start-Up Behavior Analysis: An Interval
52//!    Arithmetic Based Approach, Chapter 4. Stanford University, 2015.
53//! 2. Moore, Ramon E. Interval analysis. Vol. 4. Englewood Cliffs: Prentice-Hall, 1966.
54//! 3. F. Messine, "Deterministic global optimization using interval constraint
55//!    propagation techniques," RAIRO-Operations Research, vol. 38, no. 04,
56//!    pp. 277-293, 2004.
57//!
58//! # Illustration
59//!
60//! ## Computing bounds for an expression using interval arithmetic
61//!
62//! ```text
63//!             +-----+                         +-----+
64//!        +----|  +  |----+               +----|  +  |----+
65//!        |    |     |    |               |    |     |    |
66//!        |    +-----+    |               |    +-----+    |
67//!        |               |               |               |
68//!    +-----+           +-----+       +-----+           +-----+
69//!    |   2 |           |  y  |       |   2 | [1, 4]    |  y  |
70//!    |[.]  |           |     |       |[.]  |           |     |
71//!    +-----+           +-----+       +-----+           +-----+
72//!       |                               |
73//!       |                               |
74//!     +---+                           +---+
75//!     | x | [1, 2]                    | x | [1, 2]
76//!     +---+                           +---+
77//!
78//!  (a) Bottom-up evaluation: Step 1 (b) Bottom up evaluation: Step 2
79//!
80//!                                      [1 - 3, 4 + 1] = [-2, 5]
81//!             +-----+                         +-----+
82//!        +----|  +  |----+               +----|  +  |----+
83//!        |    |     |    |               |    |     |    |
84//!        |    +-----+    |               |    +-----+    |
85//!        |               |               |               |
86//!    +-----+           +-----+       +-----+           +-----+
87//!    |   2 |[1, 4]     |  y  |       |   2 |[1, 4]     |  y  |
88//!    |[.]  |           |     |       |[.]  |           |     |
89//!    +-----+           +-----+       +-----+           +-----+
90//!       |              [-3, 1]          |              [-3, 1]
91//!       |                               |
92//!     +---+                           +---+
93//!     | x | [1, 2]                    | x | [1, 2]
94//!     +---+                           +---+
95//!
96//!  (c) Bottom-up evaluation: Step 3 (d) Bottom-up evaluation: Step 4
97//! ```
98//!
99//! ## Top-down constraint propagation using inverse semantics
100//!
101//! ```text
102//!    [-2, 5] ∩ [4, 4] = [4, 4]               [4, 4]
103//!            +-----+                         +-----+
104//!       +----|  +  |----+               +----|  +  |----+
105//!       |    |     |    |               |    |     |    |
106//!       |    +-----+    |               |    +-----+    |
107//!       |               |               |               |
108//!    +-----+           +-----+       +-----+           +-----+
109//!    |   2 | [1, 4]    |  y  |       |   2 | [1, 4]    |  y  | [0, 1]*
110//!    |[.]  |           |     |       |[.]  |           |     |
111//!    +-----+           +-----+       +-----+           +-----+
112//!      |              [-3, 1]          |
113//!      |                               |
114//!    +---+                           +---+
115//!    | x | [1, 2]                    | x | [1, 2]
116//!    +---+                           +---+
117//!
118//!  (a) Top-down propagation: Step 1 (b) Top-down propagation: Step 2
119//!
120//!                                     [1 - 3, 4 + 1] = [-2, 5]
121//!            +-----+                         +-----+
122//!       +----|  +  |----+               +----|  +  |----+
123//!       |    |     |    |               |    |     |    |
124//!       |    +-----+    |               |    +-----+    |
125//!       |               |               |               |
126//!    +-----+           +-----+       +-----+           +-----+
127//!    |   2 |[3, 4]**   |  y  |       |   2 |[3, 4]     |  y  |
128//!    |[.]  |           |     |       |[.]  |           |     |
129//!    +-----+           +-----+       +-----+           +-----+
130//!      |              [0, 1]           |              [-3, 1]
131//!      |                               |
132//!    +---+                           +---+
133//!    | x | [1, 2]                    | x | [sqrt(3), 2]***
134//!    +---+                           +---+
135//!
136//!  (c) Top-down propagation: Step 3  (d) Top-down propagation: Step 4
137//!
138//!    * [-3, 1] ∩ ([4, 4] - [1, 4]) = [0, 1]
139//!    ** [1, 4] ∩ ([4, 4] - [0, 1]) = [3, 4]
140//!    *** [1, 2] ∩ [sqrt(3), sqrt(4)] = [sqrt(3), 2]
141//! ```
142
143use std::collections::HashSet;
144use std::fmt::{Display, Formatter};
145use std::mem::{size_of, size_of_val};
146use std::sync::Arc;
147
148use super::utils::{
149    convert_duration_type_to_interval, convert_interval_type_to_duration, get_inverse_op,
150};
151use crate::expressions::{BinaryExpr, Literal};
152use crate::utils::{build_dag, ExprTreeNode};
153use crate::PhysicalExpr;
154
155use arrow::datatypes::{DataType, Schema};
156use datafusion_common::{internal_err, not_impl_err, Result};
157use datafusion_expr::interval_arithmetic::{apply_operator, satisfy_greater, Interval};
158use datafusion_expr::Operator;
159
160use petgraph::graph::NodeIndex;
161use petgraph::stable_graph::{DefaultIx, StableGraph};
162use petgraph::visit::{Bfs, Dfs, DfsPostOrder, EdgeRef};
163use petgraph::Outgoing;
164
165/// This object implements a directed acyclic expression graph (DAEG) that
166/// is used to compute ranges for expressions through interval arithmetic.
167#[derive(Clone, Debug)]
168pub struct ExprIntervalGraph {
169    graph: StableGraph<ExprIntervalGraphNode, usize>,
170    root: NodeIndex,
171}
172
173/// This object encapsulates all possible constraint propagation results.
174#[derive(PartialEq, Debug)]
175pub enum PropagationResult {
176    CannotPropagate,
177    Infeasible,
178    Success,
179}
180
181/// This is a node in the DAEG; it encapsulates a reference to the actual
182/// [`PhysicalExpr`] as well as an interval containing expression bounds.
183#[derive(Clone, Debug)]
184pub struct ExprIntervalGraphNode {
185    expr: Arc<dyn PhysicalExpr>,
186    interval: Interval,
187}
188
189impl PartialEq for ExprIntervalGraphNode {
190    fn eq(&self, other: &Self) -> bool {
191        self.expr.eq(&other.expr)
192    }
193}
194
195impl Display for ExprIntervalGraphNode {
196    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
197        write!(f, "{}", self.expr)
198    }
199}
200
201impl ExprIntervalGraphNode {
202    /// Constructs a new DAEG node with an `[-∞, ∞]` range.
203    pub fn new_unbounded(expr: Arc<dyn PhysicalExpr>, dt: &DataType) -> Result<Self> {
204        Interval::make_unbounded(dt)
205            .map(|interval| ExprIntervalGraphNode { expr, interval })
206    }
207
208    /// Constructs a new DAEG node with the given range.
209    pub fn new_with_interval(expr: Arc<dyn PhysicalExpr>, interval: Interval) -> Self {
210        ExprIntervalGraphNode { expr, interval }
211    }
212
213    /// Get the interval object representing the range of the expression.
214    pub fn interval(&self) -> &Interval {
215        &self.interval
216    }
217
218    /// This function creates a DAEG node from DataFusion's [`ExprTreeNode`]
219    /// object. Literals are created with definite, singleton intervals while
220    /// any other expression starts with an indefinite interval (`[-∞, ∞]`).
221    pub fn make_node(node: &ExprTreeNode<NodeIndex>, schema: &Schema) -> Result<Self> {
222        let expr = Arc::clone(&node.expr);
223        if let Some(literal) = expr.as_any().downcast_ref::<Literal>() {
224            let value = literal.value();
225            Interval::try_new(value.clone(), value.clone())
226                .map(|interval| Self::new_with_interval(expr, interval))
227        } else {
228            expr.data_type(schema)
229                .and_then(|dt| Self::new_unbounded(expr, &dt))
230        }
231    }
232}
233
234/// This function refines intervals `left_child` and `right_child` by applying
235/// constraint propagation through `parent` via operation. The main idea is
236/// that we can shrink ranges of variables x and y using parent interval p.
237///
238/// Assuming that x,y and p has ranges `[xL, xU]`, `[yL, yU]`, and `[pL, pU]`, we
239/// apply the following operations:
240/// - For plus operation, specifically, we would first do
241///     - `[xL, xU]` <- (`[pL, pU]` - `[yL, yU]`) ∩ `[xL, xU]`, and then
242///     - `[yL, yU]` <- (`[pL, pU]` - `[xL, xU]`) ∩ `[yL, yU]`.
243/// - For minus operation, specifically, we would first do
244///     - `[xL, xU]` <- (`[yL, yU]` + `[pL, pU]`) ∩ `[xL, xU]`, and then
245///     - `[yL, yU]` <- (`[xL, xU]` - `[pL, pU]`) ∩ `[yL, yU]`.
246/// - For multiplication operation, specifically, we would first do
247///     - `[xL, xU]` <- (`[pL, pU]` / `[yL, yU]`) ∩ `[xL, xU]`, and then
248///     - `[yL, yU]` <- (`[pL, pU]` / `[xL, xU]`) ∩ `[yL, yU]`.
249/// - For division operation, specifically, we would first do
250///     - `[xL, xU]` <- (`[yL, yU]` * `[pL, pU]`) ∩ `[xL, xU]`, and then
251///     - `[yL, yU]` <- (`[xL, xU]` / `[pL, pU]`) ∩ `[yL, yU]`.
252pub fn propagate_arithmetic(
253    op: &Operator,
254    parent: &Interval,
255    left_child: &Interval,
256    right_child: &Interval,
257) -> Result<Option<(Interval, Interval)>> {
258    let inverse_op = get_inverse_op(*op)?;
259    match (left_child.data_type(), right_child.data_type()) {
260        // If we have a child whose type is a time interval (i.e. DataType::Interval),
261        // we need special handling since timestamp differencing results in a
262        // Duration type.
263        (DataType::Timestamp(..), DataType::Interval(_)) => {
264            propagate_time_interval_at_right(
265                left_child,
266                right_child,
267                parent,
268                op,
269                &inverse_op,
270            )
271        }
272        (DataType::Interval(_), DataType::Timestamp(..)) => {
273            propagate_time_interval_at_left(
274                left_child,
275                right_child,
276                parent,
277                op,
278                &inverse_op,
279            )
280        }
281        _ => {
282            // First, propagate to the left:
283            match apply_operator(&inverse_op, parent, right_child)?
284                .intersect(left_child)?
285            {
286                // Left is feasible:
287                Some(value) => Ok(
288                    // Propagate to the right using the new left.
289                    propagate_right(&value, parent, right_child, op, &inverse_op)?
290                        .map(|right| (value, right)),
291                ),
292                // If the left child is infeasible, short-circuit.
293                None => Ok(None),
294            }
295        }
296    }
297}
298
299/// This function refines intervals `left_child` and `right_child` by applying
300/// comparison propagation through `parent` via operation. The main idea is
301/// that we can shrink ranges of variables x and y using parent interval p.
302/// Two intervals can be ordered in 6 ways for a Gt `>` operator:
303/// ```text
304///                           (1): Infeasible, short-circuit
305/// left:   |        ================                                               |
306/// right:  |                           ========================                    |
307///
308///                             (2): Update both interval
309/// left:   |              ======================                                   |
310/// right:  |                             ======================                    |
311///                                          |
312///                                          V
313/// left:   |                             =======                                   |
314/// right:  |                             =======                                   |
315///
316///                             (3): Update left interval
317/// left:   |                  ==============================                       |
318/// right:  |                           ==========                                  |
319///                                          |
320///                                          V
321/// left:   |                           =====================                       |
322/// right:  |                           ==========                                  |
323///
324///                             (4): Update right interval
325/// left:   |                           ==========                                  |
326/// right:  |                   ===========================                         |
327///                                          |
328///                                          V
329/// left:   |                           ==========                                  |
330/// right   |                   ==================                                  |
331///
332///                                   (5): No change
333/// left:   |                       ============================                    |
334/// right:  |               ===================                                     |
335///
336///                                   (6): No change
337/// left:   |                                    ====================               |
338/// right:  |                ===============                                        |
339///
340///         -inf --------------------------------------------------------------- +inf
341/// ```
342pub fn propagate_comparison(
343    op: &Operator,
344    parent: &Interval,
345    left_child: &Interval,
346    right_child: &Interval,
347) -> Result<Option<(Interval, Interval)>> {
348    if parent == &Interval::CERTAINLY_TRUE {
349        match op {
350            Operator::Eq => left_child.intersect(right_child).map(|result| {
351                result.map(|intersection| (intersection.clone(), intersection))
352            }),
353            Operator::Gt => satisfy_greater(left_child, right_child, true),
354            Operator::GtEq => satisfy_greater(left_child, right_child, false),
355            Operator::Lt => satisfy_greater(right_child, left_child, true)
356                .map(|t| t.map(reverse_tuple)),
357            Operator::LtEq => satisfy_greater(right_child, left_child, false)
358                .map(|t| t.map(reverse_tuple)),
359            _ => internal_err!(
360                "The operator must be a comparison operator to propagate intervals"
361            ),
362        }
363    } else if parent == &Interval::CERTAINLY_FALSE {
364        match op {
365            Operator::Eq => {
366                // TODO: Propagation is not possible until we support interval sets.
367                Ok(None)
368            }
369            Operator::Gt => satisfy_greater(right_child, left_child, false),
370            Operator::GtEq => satisfy_greater(right_child, left_child, true),
371            Operator::Lt => satisfy_greater(left_child, right_child, false)
372                .map(|t| t.map(reverse_tuple)),
373            Operator::LtEq => satisfy_greater(left_child, right_child, true)
374                .map(|t| t.map(reverse_tuple)),
375            _ => internal_err!(
376                "The operator must be a comparison operator to propagate intervals"
377            ),
378        }
379    } else {
380        // Uncertainty cannot change any end-point of the intervals.
381        Ok(None)
382    }
383}
384
385impl ExprIntervalGraph {
386    pub fn try_new(expr: Arc<dyn PhysicalExpr>, schema: &Schema) -> Result<Self> {
387        // Build the full graph:
388        let (root, graph) =
389            build_dag(expr, &|node| ExprIntervalGraphNode::make_node(node, schema))?;
390        Ok(Self { graph, root })
391    }
392
393    pub fn node_count(&self) -> usize {
394        self.graph.node_count()
395    }
396
397    /// Estimate size of bytes including `Self`.
398    pub fn size(&self) -> usize {
399        let node_memory_usage = self.graph.node_count()
400            * (size_of::<ExprIntervalGraphNode>() + size_of::<NodeIndex>());
401        let edge_memory_usage =
402            self.graph.edge_count() * (size_of::<usize>() + size_of::<NodeIndex>() * 2);
403
404        size_of_val(self) + node_memory_usage + edge_memory_usage
405    }
406
407    // Sometimes, we do not want to calculate and/or propagate intervals all
408    // way down to leaf expressions. For example, assume that we have a
409    // `SymmetricHashJoin` which has a child with an output ordering like:
410    //
411    // ```text
412    // PhysicalSortExpr {
413    //     expr: BinaryExpr('a', +, 'b'),
414    //     sort_option: ..
415    // }
416    // ```
417    //
418    // i.e. its output order comes from a clause like `ORDER BY a + b`. In such
419    // a case, we must calculate the interval for the `BinaryExpr(a, +, b)`
420    // instead of the columns inside this `BinaryExpr`, because this interval
421    // decides whether we prune or not. Therefore, children `PhysicalExpr`s of
422    // this `BinaryExpr` may be pruned for performance. The figure below
423    // explains this example visually.
424    //
425    // Note that we just remove the nodes from the DAEG, do not make any change
426    // to the plan itself.
427    //
428    // ```text
429    //
430    //                                  +-----+                                          +-----+
431    //                                  | GT  |                                          | GT  |
432    //                         +--------|     |-------+                         +--------|     |-------+
433    //                         |        +-----+       |                         |        +-----+       |
434    //                         |                      |                         |                      |
435    //                      +-----+                   |                      +-----+                   |
436    //                      |Cast |                   |                      |Cast |                   |
437    //                      |     |                   |             --\      |     |                   |
438    //                      +-----+                   |       ----------     +-----+                   |
439    //                         |                      |             --/         |                      |
440    //                         |                      |                         |                      |
441    //                      +-----+                +-----+                   +-----+                +-----+
442    //                   +--|Plus |--+          +--|Plus |--+                |Plus |             +--|Plus |--+
443    //                   |  |     |  |          |  |     |  |                |     |             |  |     |  |
444    //  Prune from here  |  +-----+  |          |  +-----+  |                +-----+             |  +-----+  |
445    //  ------------------------------------    |           |                                    |           |
446    //                   |           |          |           |                                    |           |
447    //                +-----+     +-----+    +-----+     +-----+                              +-----+     +-----+
448    //                | a   |     |  b  |    |  c  |     |  2  |                              |  c  |     |  2  |
449    //                |     |     |     |    |     |     |     |                              |     |     |     |
450    //                +-----+     +-----+    +-----+     +-----+                              +-----+     +-----+
451    //
452    // ```
453
454    /// This function associates stable node indices with [`PhysicalExpr`]s so
455    /// that we can match `Arc<dyn PhysicalExpr>` and NodeIndex objects during
456    /// membership tests.
457    pub fn gather_node_indices(
458        &mut self,
459        exprs: &[Arc<dyn PhysicalExpr>],
460    ) -> Vec<(Arc<dyn PhysicalExpr>, usize)> {
461        let graph = &self.graph;
462        let mut bfs = Bfs::new(graph, self.root);
463        // We collect the node indices (usize) of [PhysicalExpr]s in the order
464        // given by argument `exprs`. To preserve this order, we initialize each
465        // expression's node index with usize::MAX, and then find the corresponding
466        // node indices by traversing the graph.
467        let mut removals = vec![];
468        let mut expr_node_indices = exprs
469            .iter()
470            .map(|e| (Arc::clone(e), usize::MAX))
471            .collect::<Vec<_>>();
472        while let Some(node) = bfs.next(graph) {
473            // Get the plan corresponding to this node:
474            let expr = &graph[node].expr;
475            // If the current expression is among `exprs`, slate its children
476            // for removal:
477            if let Some(value) = exprs.iter().position(|e| expr.eq(e)) {
478                // Update the node index of the associated `PhysicalExpr`:
479                expr_node_indices[value].1 = node.index();
480                for edge in graph.edges_directed(node, Outgoing) {
481                    // Slate the child for removal, do not remove immediately.
482                    removals.push(edge.id());
483                }
484            }
485        }
486        for edge_idx in removals {
487            self.graph.remove_edge(edge_idx);
488        }
489        // Get the set of node indices reachable from the root node:
490        let connected_nodes = self.connected_nodes();
491        // Remove nodes not connected to the root node:
492        self.graph
493            .retain_nodes(|_, index| connected_nodes.contains(&index));
494        expr_node_indices
495    }
496
497    /// Returns the set of node indices reachable from the root node via a
498    /// simple depth-first search.
499    fn connected_nodes(&self) -> HashSet<NodeIndex> {
500        let mut nodes = HashSet::new();
501        let mut dfs = Dfs::new(&self.graph, self.root);
502        while let Some(node) = dfs.next(&self.graph) {
503            nodes.insert(node);
504        }
505        nodes
506    }
507
508    /// Updates intervals for all expressions in the DAEG by successive
509    /// bottom-up and top-down traversals.
510    pub fn update_ranges(
511        &mut self,
512        leaf_bounds: &mut [(usize, Interval)],
513        given_range: Interval,
514    ) -> Result<PropagationResult> {
515        self.assign_intervals(leaf_bounds);
516        let bounds = self.evaluate_bounds()?;
517        // There are three possible cases to consider:
518        // (1) given_range ⊇ bounds => Nothing to propagate
519        // (2) ∅ ⊂ (given_range ∩ bounds) ⊂ bounds => Can propagate
520        // (3) Disjoint sets => Infeasible
521        if given_range.contains(bounds)? == Interval::CERTAINLY_TRUE {
522            // First case:
523            Ok(PropagationResult::CannotPropagate)
524        } else if bounds.contains(&given_range)? != Interval::CERTAINLY_FALSE {
525            // Second case:
526            let result = self.propagate_constraints(given_range);
527            self.update_intervals(leaf_bounds);
528            result
529        } else {
530            // Third case:
531            Ok(PropagationResult::Infeasible)
532        }
533    }
534
535    /// This function assigns given ranges to expressions in the DAEG.
536    /// The argument `assignments` associates indices of sought expressions
537    /// with their corresponding new ranges.
538    pub fn assign_intervals(&mut self, assignments: &[(usize, Interval)]) {
539        for (index, interval) in assignments {
540            let node_index = NodeIndex::from(*index as DefaultIx);
541            self.graph[node_index].interval = interval.clone();
542        }
543    }
544
545    /// This function fetches ranges of expressions from the DAEG. The argument
546    /// `assignments` associates indices of sought expressions with their ranges,
547    /// which this function modifies to reflect the intervals in the DAEG.
548    pub fn update_intervals(&self, assignments: &mut [(usize, Interval)]) {
549        for (index, interval) in assignments.iter_mut() {
550            let node_index = NodeIndex::from(*index as DefaultIx);
551            *interval = self.graph[node_index].interval.clone();
552        }
553    }
554
555    /// Computes bounds for an expression using interval arithmetic via a
556    /// bottom-up traversal.
557    ///
558    /// # Examples
559    ///
560    /// ```
561    /// use arrow::datatypes::DataType;
562    /// use arrow::datatypes::Field;
563    /// use arrow::datatypes::Schema;
564    /// use datafusion_common::ScalarValue;
565    /// use datafusion_expr::interval_arithmetic::Interval;
566    /// use datafusion_expr::Operator;
567    /// use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal};
568    /// use datafusion_physical_expr::intervals::cp_solver::ExprIntervalGraph;
569    /// use datafusion_physical_expr::PhysicalExpr;
570    /// use std::sync::Arc;
571    ///
572    /// let expr = Arc::new(BinaryExpr::new(
573    ///     Arc::new(Column::new("gnz", 0)),
574    ///     Operator::Plus,
575    ///     Arc::new(Literal::new(ScalarValue::Int32(Some(10)))),
576    /// ));
577    ///
578    /// let schema = Schema::new(vec![Field::new("gnz".to_string(), DataType::Int32, true)]);
579    ///
580    /// let mut graph = ExprIntervalGraph::try_new(expr, &schema).unwrap();
581    /// // Do it once, while constructing.
582    /// let node_indices = graph.gather_node_indices(&[Arc::new(Column::new("gnz", 0))]);
583    /// let left_index = node_indices.get(0).unwrap().1;
584    ///
585    /// // Provide intervals for leaf variables (here, there is only one).
586    /// let intervals = vec![(left_index, Interval::make(Some(10), Some(20)).unwrap())];
587    ///
588    /// // Evaluate bounds for the composite expression:
589    /// graph.assign_intervals(&intervals);
590    /// assert_eq!(
591    ///     graph.evaluate_bounds().unwrap(),
592    ///     &Interval::make(Some(20), Some(30)).unwrap(),
593    /// )
594    /// ```
595    pub fn evaluate_bounds(&mut self) -> Result<&Interval> {
596        let mut dfs = DfsPostOrder::new(&self.graph, self.root);
597        while let Some(node) = dfs.next(&self.graph) {
598            let neighbors = self.graph.neighbors_directed(node, Outgoing);
599            let mut children_intervals = neighbors
600                .map(|child| self.graph[child].interval())
601                .collect::<Vec<_>>();
602            // If the current expression is a leaf, its interval should already
603            // be set externally, just continue with the evaluation procedure:
604            if !children_intervals.is_empty() {
605                // Reverse to align with `PhysicalExpr`'s children:
606                children_intervals.reverse();
607                self.graph[node].interval =
608                    self.graph[node].expr.evaluate_bounds(&children_intervals)?;
609            }
610        }
611        Ok(self.graph[self.root].interval())
612    }
613
614    /// Updates/shrinks bounds for leaf expressions using interval arithmetic
615    /// via a top-down traversal.
616    fn propagate_constraints(
617        &mut self,
618        given_range: Interval,
619    ) -> Result<PropagationResult> {
620        // Adjust the root node with the given range:
621        if let Some(interval) = self.graph[self.root].interval.intersect(given_range)? {
622            self.graph[self.root].interval = interval;
623        } else {
624            return Ok(PropagationResult::Infeasible);
625        }
626
627        let mut bfs = Bfs::new(&self.graph, self.root);
628
629        while let Some(node) = bfs.next(&self.graph) {
630            let neighbors = self.graph.neighbors_directed(node, Outgoing);
631            let mut children = neighbors.collect::<Vec<_>>();
632            // If the current expression is a leaf, its range is now final.
633            // So, just continue with the propagation procedure:
634            if children.is_empty() {
635                continue;
636            }
637            // Reverse to align with `PhysicalExpr`'s children:
638            children.reverse();
639            let children_intervals = children
640                .iter()
641                .map(|child| self.graph[*child].interval())
642                .collect::<Vec<_>>();
643            let node_interval = self.graph[node].interval();
644            // Special case: true OR could in principle be propagated by 3 interval sets,
645            // (i.e. left true, or right true, or both true) however we do not support this yet.
646            if node_interval == &Interval::CERTAINLY_TRUE
647                && self.graph[node]
648                    .expr
649                    .as_any()
650                    .downcast_ref::<BinaryExpr>()
651                    .is_some_and(|expr| expr.op() == &Operator::Or)
652            {
653                return not_impl_err!("OR operator cannot yet propagate true intervals");
654            }
655            let propagated_intervals = self.graph[node]
656                .expr
657                .propagate_constraints(node_interval, &children_intervals)?;
658            if let Some(propagated_intervals) = propagated_intervals {
659                for (child, interval) in children.into_iter().zip(propagated_intervals) {
660                    self.graph[child].interval = interval;
661                }
662            } else {
663                // The constraint is infeasible, report:
664                return Ok(PropagationResult::Infeasible);
665            }
666        }
667        Ok(PropagationResult::Success)
668    }
669
670    /// Returns the interval associated with the node at the given `index`.
671    pub fn get_interval(&self, index: usize) -> Interval {
672        self.graph[NodeIndex::new(index)].interval.clone()
673    }
674}
675
676/// This is a subfunction of the `propagate_arithmetic` function that propagates to the right child.
677fn propagate_right(
678    left: &Interval,
679    parent: &Interval,
680    right: &Interval,
681    op: &Operator,
682    inverse_op: &Operator,
683) -> Result<Option<Interval>> {
684    match op {
685        Operator::Minus => apply_operator(op, left, parent),
686        Operator::Plus => apply_operator(inverse_op, parent, left),
687        Operator::Divide => apply_operator(op, left, parent),
688        Operator::Multiply => apply_operator(inverse_op, parent, left),
689        _ => internal_err!("Interval arithmetic does not support the operator {}", op),
690    }?
691    .intersect(right)
692}
693
694/// During the propagation of [`Interval`] values on an [`ExprIntervalGraph`],
695/// if there exists a `timestamp - timestamp` operation, the result would be
696/// of type `Duration`. However, we may encounter a situation where a time interval
697/// is involved in an arithmetic operation with a `Duration` type. This function
698/// offers special handling for such cases, where the time interval resides on
699/// the left side of the operation.
700fn propagate_time_interval_at_left(
701    left_child: &Interval,
702    right_child: &Interval,
703    parent: &Interval,
704    op: &Operator,
705    inverse_op: &Operator,
706) -> Result<Option<(Interval, Interval)>> {
707    // We check if the child's time interval(s) has a non-zero month or day field(s).
708    // If so, we return it as is without propagating. Otherwise, we first convert
709    // the time intervals to the `Duration` type, then propagate, and then convert
710    // the bounds to time intervals again.
711    let result = if let Some(duration) = convert_interval_type_to_duration(left_child) {
712        match apply_operator(inverse_op, parent, right_child)?.intersect(duration)? {
713            Some(value) => {
714                let left = convert_duration_type_to_interval(&value);
715                let right = propagate_right(&value, parent, right_child, op, inverse_op)?;
716                match (left, right) {
717                    (Some(left), Some(right)) => Some((left, right)),
718                    _ => None,
719                }
720            }
721            None => None,
722        }
723    } else {
724        propagate_right(left_child, parent, right_child, op, inverse_op)?
725            .map(|right| (left_child.clone(), right))
726    };
727    Ok(result)
728}
729
730/// During the propagation of [`Interval`] values on an [`ExprIntervalGraph`],
731/// if there exists a `timestamp - timestamp` operation, the result would be
732/// of type `Duration`. However, we may encounter a situation where a time interval
733/// is involved in an arithmetic operation with a `Duration` type. This function
734/// offers special handling for such cases, where the time interval resides on
735/// the right side of the operation.
736fn propagate_time_interval_at_right(
737    left_child: &Interval,
738    right_child: &Interval,
739    parent: &Interval,
740    op: &Operator,
741    inverse_op: &Operator,
742) -> Result<Option<(Interval, Interval)>> {
743    // We check if the child's time interval(s) has a non-zero month or day field(s).
744    // If so, we return it as is without propagating. Otherwise, we first convert
745    // the time intervals to the `Duration` type, then propagate, and then convert
746    // the bounds to time intervals again.
747    let result = if let Some(duration) = convert_interval_type_to_duration(right_child) {
748        match apply_operator(inverse_op, parent, &duration)?.intersect(left_child)? {
749            Some(value) => {
750                propagate_right(left_child, parent, &duration, op, inverse_op)?
751                    .and_then(|right| convert_duration_type_to_interval(&right))
752                    .map(|right| (value, right))
753            }
754            None => None,
755        }
756    } else {
757        apply_operator(inverse_op, parent, right_child)?
758            .intersect(left_child)?
759            .map(|value| (value, right_child.clone()))
760    };
761    Ok(result)
762}
763
764fn reverse_tuple<T, U>((first, second): (T, U)) -> (U, T) {
765    (second, first)
766}
767
768#[cfg(test)]
769mod tests {
770    use super::*;
771    use crate::expressions::{BinaryExpr, Column};
772    use crate::intervals::test_utils::gen_conjunctive_numerical_expr;
773
774    use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano};
775    use arrow::datatypes::{Field, TimeUnit};
776    use datafusion_common::ScalarValue;
777
778    use itertools::Itertools;
779    use rand::rngs::StdRng;
780    use rand::{Rng, SeedableRng};
781    use rstest::*;
782
783    #[allow(clippy::too_many_arguments)]
784    fn experiment(
785        expr: Arc<dyn PhysicalExpr>,
786        exprs_with_interval: (Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>),
787        left_interval: Interval,
788        right_interval: Interval,
789        left_expected: Interval,
790        right_expected: Interval,
791        result: PropagationResult,
792        schema: &Schema,
793    ) -> Result<()> {
794        let col_stats = [
795            (Arc::clone(&exprs_with_interval.0), left_interval),
796            (Arc::clone(&exprs_with_interval.1), right_interval),
797        ];
798        let expected = [
799            (Arc::clone(&exprs_with_interval.0), left_expected),
800            (Arc::clone(&exprs_with_interval.1), right_expected),
801        ];
802        let mut graph = ExprIntervalGraph::try_new(expr, schema)?;
803        let expr_indexes = graph.gather_node_indices(
804            &col_stats.iter().map(|(e, _)| Arc::clone(e)).collect_vec(),
805        );
806
807        let mut col_stat_nodes = col_stats
808            .iter()
809            .zip(expr_indexes.iter())
810            .map(|((_, interval), (_, index))| (*index, interval.clone()))
811            .collect_vec();
812        let expected_nodes = expected
813            .iter()
814            .zip(expr_indexes.iter())
815            .map(|((_, interval), (_, index))| (*index, interval.clone()))
816            .collect_vec();
817
818        let exp_result =
819            graph.update_ranges(&mut col_stat_nodes[..], Interval::CERTAINLY_TRUE)?;
820        assert_eq!(exp_result, result);
821        col_stat_nodes.iter().zip(expected_nodes.iter()).for_each(
822            |((_, calculated_interval_node), (_, expected))| {
823                // NOTE: These randomized tests only check for conservative containment,
824                // not openness/closedness of endpoints.
825
826                // Calculated bounds are relaxed by 1 to cover all strict and
827                // and non-strict comparison cases since we have only closed bounds.
828                let one = ScalarValue::new_one(&expected.data_type()).unwrap();
829                assert!(
830                    calculated_interval_node.lower()
831                        <= &expected.lower().add(&one).unwrap(),
832                    "{}",
833                    format!(
834                        "Calculated {} must be less than or equal {}",
835                        calculated_interval_node.lower(),
836                        expected.lower()
837                    )
838                );
839                assert!(
840                    calculated_interval_node.upper()
841                        >= &expected.upper().sub(&one).unwrap(),
842                    "{}",
843                    format!(
844                        "Calculated {} must be greater than or equal {}",
845                        calculated_interval_node.upper(),
846                        expected.upper()
847                    )
848                );
849            },
850        );
851        Ok(())
852    }
853
854    macro_rules! generate_cases {
855        ($FUNC_NAME:ident, $TYPE:ty, $SCALAR:ident) => {
856            fn $FUNC_NAME<const ASC: bool>(
857                expr: Arc<dyn PhysicalExpr>,
858                left_col: Arc<dyn PhysicalExpr>,
859                right_col: Arc<dyn PhysicalExpr>,
860                seed: u64,
861                expr_left: $TYPE,
862                expr_right: $TYPE,
863            ) -> Result<()> {
864                let mut r = StdRng::seed_from_u64(seed);
865
866                let (left_given, right_given, left_expected, right_expected) = if ASC {
867                    let left = r.random_range((0 as $TYPE)..(1000 as $TYPE));
868                    let right = r.random_range((0 as $TYPE)..(1000 as $TYPE));
869                    (
870                        (Some(left), None),
871                        (Some(right), None),
872                        (Some(<$TYPE>::max(left, right + expr_left)), None),
873                        (Some(<$TYPE>::max(right, left + expr_right)), None),
874                    )
875                } else {
876                    let left = r.random_range((0 as $TYPE)..(1000 as $TYPE));
877                    let right = r.random_range((0 as $TYPE)..(1000 as $TYPE));
878                    (
879                        (None, Some(left)),
880                        (None, Some(right)),
881                        (None, Some(<$TYPE>::min(left, right + expr_left))),
882                        (None, Some(<$TYPE>::min(right, left + expr_right))),
883                    )
884                };
885
886                experiment(
887                    expr,
888                    (left_col.clone(), right_col.clone()),
889                    Interval::make(left_given.0, left_given.1).unwrap(),
890                    Interval::make(right_given.0, right_given.1).unwrap(),
891                    Interval::make(left_expected.0, left_expected.1).unwrap(),
892                    Interval::make(right_expected.0, right_expected.1).unwrap(),
893                    PropagationResult::Success,
894                    &Schema::new(vec![
895                        Field::new(
896                            left_col.as_any().downcast_ref::<Column>().unwrap().name(),
897                            DataType::$SCALAR,
898                            true,
899                        ),
900                        Field::new(
901                            right_col.as_any().downcast_ref::<Column>().unwrap().name(),
902                            DataType::$SCALAR,
903                            true,
904                        ),
905                    ]),
906                )
907            }
908        };
909    }
910    generate_cases!(generate_case_i32, i32, Int32);
911    generate_cases!(generate_case_i64, i64, Int64);
912    generate_cases!(generate_case_f32, f32, Float32);
913    generate_cases!(generate_case_f64, f64, Float64);
914
915    #[test]
916    fn testing_not_possible() -> Result<()> {
917        let left_col = Arc::new(Column::new("left_watermark", 0));
918        let right_col = Arc::new(Column::new("right_watermark", 0));
919
920        // left_watermark > right_watermark + 5
921        let left_and_1 = Arc::new(BinaryExpr::new(
922            Arc::clone(&left_col) as Arc<dyn PhysicalExpr>,
923            Operator::Plus,
924            Arc::new(Literal::new(ScalarValue::Int32(Some(5)))),
925        ));
926        let expr = Arc::new(BinaryExpr::new(
927            left_and_1,
928            Operator::Gt,
929            Arc::clone(&right_col) as Arc<dyn PhysicalExpr>,
930        ));
931        experiment(
932            expr,
933            (
934                Arc::clone(&left_col) as Arc<dyn PhysicalExpr>,
935                Arc::clone(&right_col) as Arc<dyn PhysicalExpr>,
936            ),
937            Interval::make(Some(10_i32), Some(20_i32))?,
938            Interval::make(Some(100), None)?,
939            Interval::make(Some(10), Some(20))?,
940            Interval::make(Some(100), None)?,
941            PropagationResult::Infeasible,
942            &Schema::new(vec![
943                Field::new(
944                    left_col.as_any().downcast_ref::<Column>().unwrap().name(),
945                    DataType::Int32,
946                    true,
947                ),
948                Field::new(
949                    right_col.as_any().downcast_ref::<Column>().unwrap().name(),
950                    DataType::Int32,
951                    true,
952                ),
953            ]),
954        )
955    }
956
957    macro_rules! integer_float_case_1 {
958        ($TEST_FUNC_NAME:ident, $GENERATE_CASE_FUNC_NAME:ident, $TYPE:ty, $SCALAR:ident) => {
959            #[rstest]
960            #[test]
961            fn $TEST_FUNC_NAME(
962                #[values(0, 1, 2, 3, 4, 12, 32, 314, 3124, 123, 125, 211, 215, 4123)]
963                seed: u64,
964                #[values(Operator::Gt, Operator::GtEq)] greater_op: Operator,
965                #[values(Operator::Lt, Operator::LtEq)] less_op: Operator,
966            ) -> Result<()> {
967                let left_col = Arc::new(Column::new("left_watermark", 0));
968                let right_col = Arc::new(Column::new("right_watermark", 0));
969
970                // left_watermark + 1 > right_watermark + 11 AND left_watermark + 3 < right_watermark + 33
971                let expr = gen_conjunctive_numerical_expr(
972                    left_col.clone(),
973                    right_col.clone(),
974                    (
975                        Operator::Plus,
976                        Operator::Plus,
977                        Operator::Plus,
978                        Operator::Plus,
979                    ),
980                    ScalarValue::$SCALAR(Some(1 as $TYPE)),
981                    ScalarValue::$SCALAR(Some(11 as $TYPE)),
982                    ScalarValue::$SCALAR(Some(3 as $TYPE)),
983                    ScalarValue::$SCALAR(Some(33 as $TYPE)),
984                    (greater_op, less_op),
985                );
986                // l > r + 10 AND r > l - 30
987                let l_gt_r = 10 as $TYPE;
988                let r_gt_l = -30 as $TYPE;
989                $GENERATE_CASE_FUNC_NAME::<true>(
990                    expr.clone(),
991                    left_col.clone(),
992                    right_col.clone(),
993                    seed,
994                    l_gt_r,
995                    r_gt_l,
996                )?;
997                // Descending tests
998                // r < l - 10 AND l < r + 30
999                let r_lt_l = -l_gt_r;
1000                let l_lt_r = -r_gt_l;
1001                $GENERATE_CASE_FUNC_NAME::<false>(
1002                    expr, left_col, right_col, seed, l_lt_r, r_lt_l,
1003                )
1004            }
1005        };
1006    }
1007
1008    integer_float_case_1!(case_1_i32, generate_case_i32, i32, Int32);
1009    integer_float_case_1!(case_1_i64, generate_case_i64, i64, Int64);
1010    integer_float_case_1!(case_1_f64, generate_case_f64, f64, Float64);
1011    integer_float_case_1!(case_1_f32, generate_case_f32, f32, Float32);
1012
1013    macro_rules! integer_float_case_2 {
1014        ($TEST_FUNC_NAME:ident, $GENERATE_CASE_FUNC_NAME:ident, $TYPE:ty, $SCALAR:ident) => {
1015            #[rstest]
1016            #[test]
1017            fn $TEST_FUNC_NAME(
1018                #[values(0, 1, 2, 3, 4, 12, 32, 314, 3124, 123, 125, 211, 215, 4123)]
1019                seed: u64,
1020                #[values(Operator::Gt, Operator::GtEq)] greater_op: Operator,
1021                #[values(Operator::Lt, Operator::LtEq)] less_op: Operator,
1022            ) -> Result<()> {
1023                let left_col = Arc::new(Column::new("left_watermark", 0));
1024                let right_col = Arc::new(Column::new("right_watermark", 0));
1025
1026                // left_watermark - 1 > right_watermark + 5 AND left_watermark + 3 < right_watermark + 10
1027                let expr = gen_conjunctive_numerical_expr(
1028                    left_col.clone(),
1029                    right_col.clone(),
1030                    (
1031                        Operator::Minus,
1032                        Operator::Plus,
1033                        Operator::Plus,
1034                        Operator::Plus,
1035                    ),
1036                    ScalarValue::$SCALAR(Some(1 as $TYPE)),
1037                    ScalarValue::$SCALAR(Some(5 as $TYPE)),
1038                    ScalarValue::$SCALAR(Some(3 as $TYPE)),
1039                    ScalarValue::$SCALAR(Some(10 as $TYPE)),
1040                    (greater_op, less_op),
1041                );
1042                // l > r + 6 AND r > l - 7
1043                let l_gt_r = 6 as $TYPE;
1044                let r_gt_l = -7 as $TYPE;
1045                $GENERATE_CASE_FUNC_NAME::<true>(
1046                    expr.clone(),
1047                    left_col.clone(),
1048                    right_col.clone(),
1049                    seed,
1050                    l_gt_r,
1051                    r_gt_l,
1052                )?;
1053                // Descending tests
1054                // r < l - 6 AND l < r + 7
1055                let r_lt_l = -l_gt_r;
1056                let l_lt_r = -r_gt_l;
1057                $GENERATE_CASE_FUNC_NAME::<false>(
1058                    expr, left_col, right_col, seed, l_lt_r, r_lt_l,
1059                )
1060            }
1061        };
1062    }
1063
1064    integer_float_case_2!(case_2_i32, generate_case_i32, i32, Int32);
1065    integer_float_case_2!(case_2_i64, generate_case_i64, i64, Int64);
1066    integer_float_case_2!(case_2_f64, generate_case_f64, f64, Float64);
1067    integer_float_case_2!(case_2_f32, generate_case_f32, f32, Float32);
1068
1069    macro_rules! integer_float_case_3 {
1070        ($TEST_FUNC_NAME:ident, $GENERATE_CASE_FUNC_NAME:ident, $TYPE:ty, $SCALAR:ident) => {
1071            #[rstest]
1072            #[test]
1073            fn $TEST_FUNC_NAME(
1074                #[values(0, 1, 2, 3, 4, 12, 32, 314, 3124, 123, 125, 211, 215, 4123)]
1075                seed: u64,
1076                #[values(Operator::Gt, Operator::GtEq)] greater_op: Operator,
1077                #[values(Operator::Lt, Operator::LtEq)] less_op: Operator,
1078            ) -> Result<()> {
1079                let left_col = Arc::new(Column::new("left_watermark", 0));
1080                let right_col = Arc::new(Column::new("right_watermark", 0));
1081
1082                // left_watermark - 1 > right_watermark + 5 AND left_watermark - 3 < right_watermark + 10
1083                let expr = gen_conjunctive_numerical_expr(
1084                    left_col.clone(),
1085                    right_col.clone(),
1086                    (
1087                        Operator::Minus,
1088                        Operator::Plus,
1089                        Operator::Minus,
1090                        Operator::Plus,
1091                    ),
1092                    ScalarValue::$SCALAR(Some(1 as $TYPE)),
1093                    ScalarValue::$SCALAR(Some(5 as $TYPE)),
1094                    ScalarValue::$SCALAR(Some(3 as $TYPE)),
1095                    ScalarValue::$SCALAR(Some(10 as $TYPE)),
1096                    (greater_op, less_op),
1097                );
1098                // l > r + 6 AND r > l - 13
1099                let l_gt_r = 6 as $TYPE;
1100                let r_gt_l = -13 as $TYPE;
1101                $GENERATE_CASE_FUNC_NAME::<true>(
1102                    expr.clone(),
1103                    left_col.clone(),
1104                    right_col.clone(),
1105                    seed,
1106                    l_gt_r,
1107                    r_gt_l,
1108                )?;
1109                // Descending tests
1110                // r < l - 6 AND l < r + 13
1111                let r_lt_l = -l_gt_r;
1112                let l_lt_r = -r_gt_l;
1113                $GENERATE_CASE_FUNC_NAME::<false>(
1114                    expr, left_col, right_col, seed, l_lt_r, r_lt_l,
1115                )
1116            }
1117        };
1118    }
1119
1120    integer_float_case_3!(case_3_i32, generate_case_i32, i32, Int32);
1121    integer_float_case_3!(case_3_i64, generate_case_i64, i64, Int64);
1122    integer_float_case_3!(case_3_f64, generate_case_f64, f64, Float64);
1123    integer_float_case_3!(case_3_f32, generate_case_f32, f32, Float32);
1124
1125    macro_rules! integer_float_case_4 {
1126        ($TEST_FUNC_NAME:ident, $GENERATE_CASE_FUNC_NAME:ident, $TYPE:ty, $SCALAR:ident) => {
1127            #[rstest]
1128            #[test]
1129            fn $TEST_FUNC_NAME(
1130                #[values(0, 1, 2, 3, 4, 12, 32, 314, 3124, 123, 125, 211, 215, 4123)]
1131                seed: u64,
1132                #[values(Operator::Gt, Operator::GtEq)] greater_op: Operator,
1133                #[values(Operator::Lt, Operator::LtEq)] less_op: Operator,
1134            ) -> Result<()> {
1135                let left_col = Arc::new(Column::new("left_watermark", 0));
1136                let right_col = Arc::new(Column::new("right_watermark", 0));
1137
1138                // left_watermark - 10 > right_watermark - 5 AND left_watermark - 30 < right_watermark - 3
1139                let expr = gen_conjunctive_numerical_expr(
1140                    left_col.clone(),
1141                    right_col.clone(),
1142                    (
1143                        Operator::Minus,
1144                        Operator::Minus,
1145                        Operator::Minus,
1146                        Operator::Plus,
1147                    ),
1148                    ScalarValue::$SCALAR(Some(10 as $TYPE)),
1149                    ScalarValue::$SCALAR(Some(5 as $TYPE)),
1150                    ScalarValue::$SCALAR(Some(3 as $TYPE)),
1151                    ScalarValue::$SCALAR(Some(10 as $TYPE)),
1152                    (greater_op, less_op),
1153                );
1154                // l > r + 5 AND r > l - 13
1155                let l_gt_r = 5 as $TYPE;
1156                let r_gt_l = -13 as $TYPE;
1157                $GENERATE_CASE_FUNC_NAME::<true>(
1158                    expr.clone(),
1159                    left_col.clone(),
1160                    right_col.clone(),
1161                    seed,
1162                    l_gt_r,
1163                    r_gt_l,
1164                )?;
1165                // Descending tests
1166                // r < l - 5 AND l < r + 13
1167                let r_lt_l = -l_gt_r;
1168                let l_lt_r = -r_gt_l;
1169                $GENERATE_CASE_FUNC_NAME::<false>(
1170                    expr, left_col, right_col, seed, l_lt_r, r_lt_l,
1171                )
1172            }
1173        };
1174    }
1175
1176    integer_float_case_4!(case_4_i32, generate_case_i32, i32, Int32);
1177    integer_float_case_4!(case_4_i64, generate_case_i64, i64, Int64);
1178    integer_float_case_4!(case_4_f64, generate_case_f64, f64, Float64);
1179    integer_float_case_4!(case_4_f32, generate_case_f32, f32, Float32);
1180
1181    macro_rules! integer_float_case_5 {
1182        ($TEST_FUNC_NAME:ident, $GENERATE_CASE_FUNC_NAME:ident, $TYPE:ty, $SCALAR:ident) => {
1183            #[rstest]
1184            #[test]
1185            fn $TEST_FUNC_NAME(
1186                #[values(0, 1, 2, 3, 4, 12, 32, 314, 3124, 123, 125, 211, 215, 4123)]
1187                seed: u64,
1188                #[values(Operator::Gt, Operator::GtEq)] greater_op: Operator,
1189                #[values(Operator::Lt, Operator::LtEq)] less_op: Operator,
1190            ) -> Result<()> {
1191                let left_col = Arc::new(Column::new("left_watermark", 0));
1192                let right_col = Arc::new(Column::new("right_watermark", 0));
1193
1194                // left_watermark - 10 > right_watermark - 5 AND left_watermark - 30 < right_watermark - 3
1195                let expr = gen_conjunctive_numerical_expr(
1196                    left_col.clone(),
1197                    right_col.clone(),
1198                    (
1199                        Operator::Minus,
1200                        Operator::Minus,
1201                        Operator::Minus,
1202                        Operator::Minus,
1203                    ),
1204                    ScalarValue::$SCALAR(Some(10 as $TYPE)),
1205                    ScalarValue::$SCALAR(Some(5 as $TYPE)),
1206                    ScalarValue::$SCALAR(Some(30 as $TYPE)),
1207                    ScalarValue::$SCALAR(Some(3 as $TYPE)),
1208                    (greater_op, less_op),
1209                );
1210                // l > r + 5 AND r > l - 27
1211                let l_gt_r = 5 as $TYPE;
1212                let r_gt_l = -27 as $TYPE;
1213                $GENERATE_CASE_FUNC_NAME::<true>(
1214                    expr.clone(),
1215                    left_col.clone(),
1216                    right_col.clone(),
1217                    seed,
1218                    l_gt_r,
1219                    r_gt_l,
1220                )?;
1221                // Descending tests
1222                // r < l - 5 AND l < r + 27
1223                let r_lt_l = -l_gt_r;
1224                let l_lt_r = -r_gt_l;
1225                $GENERATE_CASE_FUNC_NAME::<false>(
1226                    expr, left_col, right_col, seed, l_lt_r, r_lt_l,
1227                )
1228            }
1229        };
1230    }
1231
1232    integer_float_case_5!(case_5_i32, generate_case_i32, i32, Int32);
1233    integer_float_case_5!(case_5_i64, generate_case_i64, i64, Int64);
1234    integer_float_case_5!(case_5_f64, generate_case_f64, f64, Float64);
1235    integer_float_case_5!(case_5_f32, generate_case_f32, f32, Float32);
1236
1237    #[test]
1238    fn test_gather_node_indices_dont_remove() -> Result<()> {
1239        // Expression: a@0 + b@1 + 1 > a@0 - b@1, given a@0 + b@1.
1240        // Do not remove a@0 or b@1, only remove edges since a@0 - b@1 also
1241        // depends on leaf nodes a@0 and b@1.
1242        let left_expr = Arc::new(BinaryExpr::new(
1243            Arc::new(BinaryExpr::new(
1244                Arc::new(Column::new("a", 0)),
1245                Operator::Plus,
1246                Arc::new(Column::new("b", 1)),
1247            )),
1248            Operator::Plus,
1249            Arc::new(Literal::new(ScalarValue::Int32(Some(1)))),
1250        ));
1251
1252        let right_expr = Arc::new(BinaryExpr::new(
1253            Arc::new(Column::new("a", 0)),
1254            Operator::Minus,
1255            Arc::new(Column::new("b", 1)),
1256        ));
1257        let expr = Arc::new(BinaryExpr::new(left_expr, Operator::Gt, right_expr));
1258        let mut graph = ExprIntervalGraph::try_new(
1259            expr,
1260            &Schema::new(vec![
1261                Field::new("a", DataType::Int32, true),
1262                Field::new("b", DataType::Int32, true),
1263            ]),
1264        )
1265        .unwrap();
1266        // Define a test leaf node.
1267        let leaf_node = Arc::new(BinaryExpr::new(
1268            Arc::new(Column::new("a", 0)),
1269            Operator::Plus,
1270            Arc::new(Column::new("b", 1)),
1271        ));
1272        // Store the current node count.
1273        let prev_node_count = graph.node_count();
1274        // Gather the index of node in the expression graph that match the test leaf node.
1275        graph.gather_node_indices(&[leaf_node]);
1276        // Store the final node count.
1277        let final_node_count = graph.node_count();
1278        // Assert that the final node count is equal the previous node count.
1279        // This means we did not remove any node.
1280        assert_eq!(prev_node_count, final_node_count);
1281        Ok(())
1282    }
1283
1284    #[test]
1285    fn test_gather_node_indices_remove() -> Result<()> {
1286        // Expression: a@0 + b@1 + 1 > y@0 - z@1, given a@0 + b@1.
1287        // We expect to remove two nodes since we do not need a@ and b@.
1288        let left_expr = Arc::new(BinaryExpr::new(
1289            Arc::new(BinaryExpr::new(
1290                Arc::new(Column::new("a", 0)),
1291                Operator::Plus,
1292                Arc::new(Column::new("b", 1)),
1293            )),
1294            Operator::Plus,
1295            Arc::new(Literal::new(ScalarValue::Int32(Some(1)))),
1296        ));
1297
1298        let right_expr = Arc::new(BinaryExpr::new(
1299            Arc::new(Column::new("y", 0)),
1300            Operator::Minus,
1301            Arc::new(Column::new("z", 1)),
1302        ));
1303        let expr = Arc::new(BinaryExpr::new(left_expr, Operator::Gt, right_expr));
1304        let mut graph = ExprIntervalGraph::try_new(
1305            expr,
1306            &Schema::new(vec![
1307                Field::new("a", DataType::Int32, true),
1308                Field::new("b", DataType::Int32, true),
1309                Field::new("y", DataType::Int32, true),
1310                Field::new("z", DataType::Int32, true),
1311            ]),
1312        )
1313        .unwrap();
1314        // Define a test leaf node.
1315        let leaf_node = Arc::new(BinaryExpr::new(
1316            Arc::new(Column::new("a", 0)),
1317            Operator::Plus,
1318            Arc::new(Column::new("b", 1)),
1319        ));
1320        // Store the current node count.
1321        let prev_node_count = graph.node_count();
1322        // Gather the index of node in the expression graph that match the test leaf node.
1323        graph.gather_node_indices(&[leaf_node]);
1324        // Store the final node count.
1325        let final_node_count = graph.node_count();
1326        // Assert that the final node count is two less than the previous node
1327        // count; i.e. that we did remove two nodes.
1328        assert_eq!(prev_node_count, final_node_count + 2);
1329        Ok(())
1330    }
1331
1332    #[test]
1333    fn test_gather_node_indices_remove_one() -> Result<()> {
1334        // Expression: a@0 + b@1 + 1 > a@0 - z@1, given a@0 + b@1.
1335        // We expect to remove one nodesince we still need a@ but not b@.
1336        let left_expr = Arc::new(BinaryExpr::new(
1337            Arc::new(BinaryExpr::new(
1338                Arc::new(Column::new("a", 0)),
1339                Operator::Plus,
1340                Arc::new(Column::new("b", 1)),
1341            )),
1342            Operator::Plus,
1343            Arc::new(Literal::new(ScalarValue::Int32(Some(1)))),
1344        ));
1345
1346        let right_expr = Arc::new(BinaryExpr::new(
1347            Arc::new(Column::new("a", 0)),
1348            Operator::Minus,
1349            Arc::new(Column::new("z", 1)),
1350        ));
1351        let expr = Arc::new(BinaryExpr::new(left_expr, Operator::Gt, right_expr));
1352        let mut graph = ExprIntervalGraph::try_new(
1353            expr,
1354            &Schema::new(vec![
1355                Field::new("a", DataType::Int32, true),
1356                Field::new("b", DataType::Int32, true),
1357                Field::new("z", DataType::Int32, true),
1358            ]),
1359        )
1360        .unwrap();
1361        // Define a test leaf node.
1362        let leaf_node = Arc::new(BinaryExpr::new(
1363            Arc::new(Column::new("a", 0)),
1364            Operator::Plus,
1365            Arc::new(Column::new("b", 1)),
1366        ));
1367        // Store the current node count.
1368        let prev_node_count = graph.node_count();
1369        // Gather the index of node in the expression graph that match the test leaf node.
1370        graph.gather_node_indices(&[leaf_node]);
1371        // Store the final node count.
1372        let final_node_count = graph.node_count();
1373        // Assert that the final node count is one less than the previous node
1374        // count; i.e. that we did remove two nodes.
1375        assert_eq!(prev_node_count, final_node_count + 1);
1376        Ok(())
1377    }
1378
1379    #[test]
1380    fn test_gather_node_indices_cannot_provide() -> Result<()> {
1381        // Expression: a@0 + 1 + b@1 > y@0 - z@1 -> provide a@0 + b@1
1382        // TODO: We expect nodes a@0 and b@1 to be pruned, and intervals to be provided from the a@0 + b@1 node.
1383        // However, we do not have an exact node for a@0 + b@1 due to the binary tree structure of the expressions.
1384        // Pruning and interval providing for BinaryExpr expressions are more challenging without exact matches.
1385        // Currently, we only support exact matches for BinaryExprs, but we plan to extend support beyond exact matches in the future.
1386        let left_expr = Arc::new(BinaryExpr::new(
1387            Arc::new(BinaryExpr::new(
1388                Arc::new(Column::new("a", 0)),
1389                Operator::Plus,
1390                Arc::new(Literal::new(ScalarValue::Int32(Some(1)))),
1391            )),
1392            Operator::Plus,
1393            Arc::new(Column::new("b", 1)),
1394        ));
1395
1396        let right_expr = Arc::new(BinaryExpr::new(
1397            Arc::new(Column::new("y", 0)),
1398            Operator::Minus,
1399            Arc::new(Column::new("z", 1)),
1400        ));
1401        let expr = Arc::new(BinaryExpr::new(left_expr, Operator::Gt, right_expr));
1402        let mut graph = ExprIntervalGraph::try_new(
1403            expr,
1404            &Schema::new(vec![
1405                Field::new("a", DataType::Int32, true),
1406                Field::new("b", DataType::Int32, true),
1407                Field::new("y", DataType::Int32, true),
1408                Field::new("z", DataType::Int32, true),
1409            ]),
1410        )
1411        .unwrap();
1412        // Define a test leaf node.
1413        let leaf_node = Arc::new(BinaryExpr::new(
1414            Arc::new(Column::new("a", 0)),
1415            Operator::Plus,
1416            Arc::new(Column::new("b", 1)),
1417        ));
1418        // Store the current node count.
1419        let prev_node_count = graph.node_count();
1420        // Gather the index of node in the expression graph that match the test leaf node.
1421        graph.gather_node_indices(&[leaf_node]);
1422        // Store the final node count.
1423        let final_node_count = graph.node_count();
1424        // Assert that the final node count is equal the previous node count (i.e., no node was pruned).
1425        assert_eq!(prev_node_count, final_node_count);
1426        Ok(())
1427    }
1428
1429    #[test]
1430    fn test_propagate_constraints_singleton_interval_at_right() -> Result<()> {
1431        let expression = BinaryExpr::new(
1432            Arc::new(Column::new("ts_column", 0)),
1433            Operator::Plus,
1434            Arc::new(Literal::new(ScalarValue::new_interval_mdn(0, 1, 321))),
1435        );
1436        let parent = Interval::try_new(
1437            // 15.10.2020 - 10:11:12.000_000_321 AM
1438            ScalarValue::TimestampNanosecond(Some(1_602_756_672_000_000_321), None),
1439            // 16.10.2020 - 10:11:12.000_000_321 AM
1440            ScalarValue::TimestampNanosecond(Some(1_602_843_072_000_000_321), None),
1441        )?;
1442        let left_child = Interval::try_new(
1443            // 10.10.2020 - 10:11:12 AM
1444            ScalarValue::TimestampNanosecond(Some(1_602_324_672_000_000_000), None),
1445            // 20.10.2020 - 10:11:12 AM
1446            ScalarValue::TimestampNanosecond(Some(1_603_188_672_000_000_000), None),
1447        )?;
1448        let right_child = Interval::try_new(
1449            // 1 day 321 ns
1450            ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNano {
1451                months: 0,
1452                days: 1,
1453                nanoseconds: 321,
1454            })),
1455            // 1 day 321 ns
1456            ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNano {
1457                months: 0,
1458                days: 1,
1459                nanoseconds: 321,
1460            })),
1461        )?;
1462        let children = vec![&left_child, &right_child];
1463        let result = expression
1464            .propagate_constraints(&parent, &children)?
1465            .unwrap();
1466
1467        assert_eq!(
1468            vec![
1469                Interval::try_new(
1470                    // 14.10.2020 - 10:11:12 AM
1471                    ScalarValue::TimestampNanosecond(
1472                        Some(1_602_670_272_000_000_000),
1473                        None
1474                    ),
1475                    // 15.10.2020 - 10:11:12 AM
1476                    ScalarValue::TimestampNanosecond(
1477                        Some(1_602_756_672_000_000_000),
1478                        None
1479                    ),
1480                )?,
1481                Interval::try_new(
1482                    // 1 day 321 ns in Duration type
1483                    ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNano {
1484                        months: 0,
1485                        days: 1,
1486                        nanoseconds: 321,
1487                    })),
1488                    // 1 day 321 ns in Duration type
1489                    ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNano {
1490                        months: 0,
1491                        days: 1,
1492                        nanoseconds: 321,
1493                    })),
1494                )?
1495            ],
1496            result
1497        );
1498
1499        Ok(())
1500    }
1501
1502    #[test]
1503    fn test_propagate_constraints_column_interval_at_left() -> Result<()> {
1504        let expression = BinaryExpr::new(
1505            Arc::new(Column::new("interval_column", 1)),
1506            Operator::Plus,
1507            Arc::new(Column::new("ts_column", 0)),
1508        );
1509        let parent = Interval::try_new(
1510            // 15.10.2020 - 10:11:12 AM
1511            ScalarValue::TimestampMillisecond(Some(1_602_756_672_000), None),
1512            // 16.10.2020 - 10:11:12 AM
1513            ScalarValue::TimestampMillisecond(Some(1_602_843_072_000), None),
1514        )?;
1515        let right_child = Interval::try_new(
1516            // 10.10.2020 - 10:11:12 AM
1517            ScalarValue::TimestampMillisecond(Some(1_602_324_672_000), None),
1518            // 20.10.2020 - 10:11:12 AM
1519            ScalarValue::TimestampMillisecond(Some(1_603_188_672_000), None),
1520        )?;
1521        let left_child = Interval::try_new(
1522            // 2 days in millisecond
1523            ScalarValue::IntervalDayTime(Some(IntervalDayTime {
1524                days: 0,
1525                milliseconds: 172_800_000,
1526            })),
1527            // 10 days in millisecond
1528            ScalarValue::IntervalDayTime(Some(IntervalDayTime {
1529                days: 0,
1530                milliseconds: 864_000_000,
1531            })),
1532        )?;
1533        let children = vec![&left_child, &right_child];
1534        let result = expression
1535            .propagate_constraints(&parent, &children)?
1536            .unwrap();
1537
1538        assert_eq!(
1539            vec![
1540                Interval::try_new(
1541                    // 2 days in millisecond
1542                    ScalarValue::IntervalDayTime(Some(IntervalDayTime {
1543                        days: 0,
1544                        milliseconds: 172_800_000,
1545                    })),
1546                    // 6 days
1547                    ScalarValue::IntervalDayTime(Some(IntervalDayTime {
1548                        days: 0,
1549                        milliseconds: 518_400_000,
1550                    })),
1551                )?,
1552                Interval::try_new(
1553                    // 10.10.2020 - 10:11:12 AM
1554                    ScalarValue::TimestampMillisecond(Some(1_602_324_672_000), None),
1555                    // 14.10.2020 - 10:11:12 AM
1556                    ScalarValue::TimestampMillisecond(Some(1_602_670_272_000), None),
1557                )?
1558            ],
1559            result
1560        );
1561
1562        Ok(())
1563    }
1564
1565    #[test]
1566    fn test_propagate_comparison() -> Result<()> {
1567        // In the examples below:
1568        // `left` is unbounded: [?, ?],
1569        // `right` is known to be [1000,1000]
1570        // so `left` < `right` results in no new knowledge of `right` but knowing that `left` is now < 1000:` [?, 999]
1571        let left = Interval::make_unbounded(&DataType::Int64)?;
1572        let right = Interval::make(Some(1000_i64), Some(1000_i64))?;
1573        assert_eq!(
1574            (Some((
1575                Interval::make(None, Some(999_i64))?,
1576                Interval::make(Some(1000_i64), Some(1000_i64))?,
1577            ))),
1578            propagate_comparison(
1579                &Operator::Lt,
1580                &Interval::CERTAINLY_TRUE,
1581                &left,
1582                &right
1583            )?
1584        );
1585
1586        let left =
1587            Interval::make_unbounded(&DataType::Timestamp(TimeUnit::Nanosecond, None))?;
1588        let right = Interval::try_new(
1589            ScalarValue::TimestampNanosecond(Some(1000), None),
1590            ScalarValue::TimestampNanosecond(Some(1000), None),
1591        )?;
1592        assert_eq!(
1593            (Some((
1594                Interval::try_new(
1595                    ScalarValue::try_from(&DataType::Timestamp(
1596                        TimeUnit::Nanosecond,
1597                        None
1598                    ))
1599                    .unwrap(),
1600                    ScalarValue::TimestampNanosecond(Some(999), None),
1601                )?,
1602                Interval::try_new(
1603                    ScalarValue::TimestampNanosecond(Some(1000), None),
1604                    ScalarValue::TimestampNanosecond(Some(1000), None),
1605                )?
1606            ))),
1607            propagate_comparison(
1608                &Operator::Lt,
1609                &Interval::CERTAINLY_TRUE,
1610                &left,
1611                &right
1612            )?
1613        );
1614
1615        let left = Interval::make_unbounded(&DataType::Timestamp(
1616            TimeUnit::Nanosecond,
1617            Some("+05:00".into()),
1618        ))?;
1619        let right = Interval::try_new(
1620            ScalarValue::TimestampNanosecond(Some(1000), Some("+05:00".into())),
1621            ScalarValue::TimestampNanosecond(Some(1000), Some("+05:00".into())),
1622        )?;
1623        assert_eq!(
1624            (Some((
1625                Interval::try_new(
1626                    ScalarValue::try_from(&DataType::Timestamp(
1627                        TimeUnit::Nanosecond,
1628                        Some("+05:00".into()),
1629                    ))
1630                    .unwrap(),
1631                    ScalarValue::TimestampNanosecond(Some(999), Some("+05:00".into())),
1632                )?,
1633                Interval::try_new(
1634                    ScalarValue::TimestampNanosecond(Some(1000), Some("+05:00".into())),
1635                    ScalarValue::TimestampNanosecond(Some(1000), Some("+05:00".into())),
1636                )?
1637            ))),
1638            propagate_comparison(
1639                &Operator::Lt,
1640                &Interval::CERTAINLY_TRUE,
1641                &left,
1642                &right
1643            )?
1644        );
1645
1646        Ok(())
1647    }
1648
1649    #[test]
1650    fn test_propagate_or() -> Result<()> {
1651        let expr = Arc::new(BinaryExpr::new(
1652            Arc::new(Column::new("a", 0)),
1653            Operator::Or,
1654            Arc::new(Column::new("b", 1)),
1655        ));
1656        let parent = Interval::CERTAINLY_FALSE;
1657        let children_set = vec![
1658            vec![&Interval::CERTAINLY_FALSE, &Interval::UNCERTAIN],
1659            vec![&Interval::UNCERTAIN, &Interval::CERTAINLY_FALSE],
1660            vec![&Interval::CERTAINLY_FALSE, &Interval::CERTAINLY_FALSE],
1661            vec![&Interval::UNCERTAIN, &Interval::UNCERTAIN],
1662        ];
1663        for children in children_set {
1664            assert_eq!(
1665                expr.propagate_constraints(&parent, &children)?.unwrap(),
1666                vec![Interval::CERTAINLY_FALSE, Interval::CERTAINLY_FALSE],
1667            );
1668        }
1669
1670        let parent = Interval::CERTAINLY_FALSE;
1671        let children_set = vec![
1672            vec![&Interval::CERTAINLY_TRUE, &Interval::UNCERTAIN],
1673            vec![&Interval::UNCERTAIN, &Interval::CERTAINLY_TRUE],
1674        ];
1675        for children in children_set {
1676            assert_eq!(expr.propagate_constraints(&parent, &children)?, None,);
1677        }
1678
1679        let parent = Interval::CERTAINLY_TRUE;
1680        let children = vec![&Interval::CERTAINLY_FALSE, &Interval::UNCERTAIN];
1681        assert_eq!(
1682            expr.propagate_constraints(&parent, &children)?.unwrap(),
1683            vec![Interval::CERTAINLY_FALSE, Interval::CERTAINLY_TRUE]
1684        );
1685
1686        let parent = Interval::CERTAINLY_TRUE;
1687        let children = vec![&Interval::UNCERTAIN, &Interval::UNCERTAIN];
1688        assert_eq!(
1689            expr.propagate_constraints(&parent, &children)?.unwrap(),
1690            // Empty means unchanged intervals.
1691            vec![]
1692        );
1693
1694        Ok(())
1695    }
1696
1697    #[test]
1698    fn test_propagate_certainly_false_and() -> Result<()> {
1699        let expr = Arc::new(BinaryExpr::new(
1700            Arc::new(Column::new("a", 0)),
1701            Operator::And,
1702            Arc::new(Column::new("b", 1)),
1703        ));
1704        let parent = Interval::CERTAINLY_FALSE;
1705        let children_and_results_set = vec![
1706            (
1707                vec![&Interval::CERTAINLY_TRUE, &Interval::UNCERTAIN],
1708                vec![Interval::CERTAINLY_TRUE, Interval::CERTAINLY_FALSE],
1709            ),
1710            (
1711                vec![&Interval::UNCERTAIN, &Interval::CERTAINLY_TRUE],
1712                vec![Interval::CERTAINLY_FALSE, Interval::CERTAINLY_TRUE],
1713            ),
1714            (
1715                vec![&Interval::UNCERTAIN, &Interval::UNCERTAIN],
1716                // Empty means unchanged intervals.
1717                vec![],
1718            ),
1719            (
1720                vec![&Interval::CERTAINLY_FALSE, &Interval::UNCERTAIN],
1721                vec![],
1722            ),
1723        ];
1724        for (children, result) in children_and_results_set {
1725            assert_eq!(
1726                expr.propagate_constraints(&parent, &children)?.unwrap(),
1727                result
1728            );
1729        }
1730
1731        Ok(())
1732    }
1733}