Skip to main content

datafusion_pruning/
pruning_predicate.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//! [`PruningPredicate`] to apply filter [`Expr`] to prune "containers"
19//! based on statistics (e.g. Parquet Row Groups)
20//!
21//! [`Expr`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/enum.Expr.html
22use std::collections::HashSet;
23use std::sync::Arc;
24
25use arrow::array::AsArray;
26use arrow::{
27    array::{ArrayRef, BooleanArray, new_null_array},
28    datatypes::{DataType, Field, Schema, SchemaRef},
29    record_batch::{RecordBatch, RecordBatchOptions},
30};
31// pub use for backwards compatibility
32pub use datafusion_common::pruning::PruningStatistics;
33use datafusion_physical_expr::simplifier::PhysicalExprSimplifier;
34use datafusion_physical_plan::metrics::Count;
35use log::{debug, trace};
36
37use datafusion_common::error::Result;
38use datafusion_common::tree_node::{TransformedResult, TreeNodeRecursion};
39use datafusion_common::{
40    _internal_datafusion_err, Column, DFSchema, assert_eq_or_internal_err,
41};
42use datafusion_common::{
43    ScalarValue, internal_datafusion_err, plan_datafusion_err, plan_err,
44    tree_node::{Transformed, TreeNode},
45};
46use datafusion_expr_common::casts::try_cast_literal_to_type;
47use datafusion_expr_common::operator::Operator;
48use datafusion_physical_expr::utils::{Guarantee, LiteralGuarantee};
49use datafusion_physical_expr::{PhysicalExprRef, expressions as phys_expr};
50use datafusion_physical_expr_common::physical_expr::snapshot_physical_expr_opt;
51use datafusion_physical_plan::{ColumnarValue, PhysicalExpr};
52
53/// Used to prove that arbitrary predicates (boolean expression) can not
54/// possibly evaluate to `true` given information about a column provided by
55/// [`PruningStatistics`].
56///
57/// # Introduction
58///
59/// `PruningPredicate` analyzes filter expressions using statistics such as
60/// min/max values and null counts, attempting to prove a "container" (e.g.
61/// Parquet Row Group) can be skipped without reading the actual data,
62/// potentially leading to significant performance improvements.
63///
64/// For example, `PruningPredicate`s are used to prune Parquet Row Groups based
65/// on the min/max values found in the Parquet metadata. If the
66/// `PruningPredicate` can prove that the filter can never evaluate to `true`
67/// for any row in the Row Group, the entire Row Group is skipped during query
68/// execution.
69///
70/// The `PruningPredicate` API is general, and can be used for pruning other
71/// types of containers (e.g. files) based on statistics that may be known from
72/// external catalogs (e.g. Delta Lake) or other sources. How this works is a
73/// subtle topic.  See the Background and Implementation section for details.
74///
75/// `PruningPredicate` supports:
76///
77/// 1. Arbitrary expressions (including user defined functions)
78///
79/// 2. Vectorized evaluation (provide more than one set of statistics at a time)
80///    so it is suitable for pruning 1000s of containers.
81///
82/// 3. Any source of information that implements the [`PruningStatistics`] trait
83///    (not just Parquet metadata).
84///
85/// # Example
86///
87/// See the [`pruning.rs` example in the `datafusion-examples`] for a complete
88/// example of how to use `PruningPredicate` to prune files based on min/max
89/// values.
90///
91/// [`pruning.rs` example in the `datafusion-examples`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/query_planning/pruning.rs
92///
93/// Given an expression like `x = 5` and statistics for 3 containers (Row
94/// Groups, files, etc) `A`, `B`, and `C`:
95///
96/// ```text
97///   A: {x_min = 0, x_max = 4}
98///   B: {x_min = 2, x_max = 10}
99///   C: {x_min = 5, x_max = 8}
100/// ```
101///
102/// `PruningPredicate` will conclude that the rows in container `A` can never
103/// be true (as the maximum value is only `4`), so it can be pruned:
104///
105/// ```text
106/// A: false (no rows could possibly match x = 5)
107/// B: true  (rows might match x = 5)
108/// C: true  (rows might match x = 5)
109/// ```
110///
111/// See [`PruningPredicateBuilder`] and [`PruningPredicate::prune`] for more information.
112///
113/// # Background
114///
115/// ## Boolean Tri-state logic
116///
117/// To understand the details of the rest of this documentation, it is important
118/// to understand how the tri-state boolean logic in SQL works. As this is
119/// somewhat esoteric, we review it here.
120///
121/// SQL has a notion of `NULL` that represents the value is `“unknown”` and this
122/// uncertainty propagates through expressions. SQL `NULL` behaves very
123/// differently than the `NULL` in most other languages where it is a special,
124/// sentinel value (e.g. `0` in `C/C++`). While representing uncertainty with
125/// `NULL` is powerful and elegant, SQL `NULL`s are often deeply confusing when
126/// first encountered as they behave differently than most programmers may
127/// expect.
128///
129/// In most other programming languages,
130/// * `a == NULL` evaluates to `true` if `a` also had the value `NULL`
131/// * `a == NULL` evaluates to `false` if `a` has any other value
132///
133/// However, in SQL `a = NULL` **always** evaluates to `NULL` (never `true` or
134/// `false`):
135///
136/// Expression    | Result
137/// ------------- | ---------
138/// `1 = NULL`    | `NULL`
139/// `NULL = NULL` | `NULL`
140///
141/// Also important is how `AND` and `OR` works with tri-state boolean logic as
142/// (perhaps counterintuitively) the result is **not** always NULL. While
143/// consistent with the notion of `NULL` representing “unknown”, this is again,
144/// often deeply confusing 🤯 when first encountered.
145///
146/// Expression       | Result    | Intuition
147/// ---------------  | --------- | -----------
148/// `NULL AND true`  |   `NULL`  | The `NULL` stands for “unknown” and if it were `true` or `false` the overall expression value could change
149/// `NULL AND false` |  `false`  | If the `NULL` was either `true` or `false` the overall expression is still `false`
150/// `NULL AND NULL`  | `NULL`    |
151///
152/// Expression      | Result    | Intuition
153/// --------------- | --------- | ----------
154/// `NULL OR true`  | `true`    |  If the `NULL` was either `true` or `false` the overall expression is still `true`
155/// `NULL OR false` | `NULL`    |  The `NULL` stands for “unknown” and if it were `true` or `false` the overall expression value could change
156/// `NULL OR NULL`  |  `NULL`   |
157///
158/// ## SQL Filter Semantics
159///
160/// The SQL `WHERE` clause has a boolean expression, often called a filter or
161/// predicate. The semantics of this predicate are that the query evaluates the
162/// predicate for each row in the input tables and:
163///
164/// * Rows that evaluate to `true` are returned in the query results
165///
166/// * Rows that evaluate to `false` are not returned (“filtered out” or “pruned” or “skipped”).
167///
168/// * Rows that evaluate to `NULL` are **NOT** returned (also “filtered out”).
169///   Note: *this treatment of `NULL` is **DIFFERENT** than how `NULL` is treated
170///   in the rewritten predicate described below.*
171///
172/// # `PruningPredicate` Implementation
173///
174/// Armed with the information in the Background section, we can now understand
175/// how the `PruningPredicate` logic works.
176///
177/// ## Interface
178///
179/// **Inputs**
180/// 1. An input schema describing what columns exist
181///
182/// 2. A predicate (expression that evaluates to a boolean)
183///
184/// 3. [`PruningStatistics`] that provides information about columns in that
185///    schema, for multiple “containers”. For each column in each container, it
186///    provides optional information on contained values, min_values, max_values,
187///    null_counts counts, and row_counts counts.
188///
189/// **Outputs**:
190/// A (non null) boolean value for each container:
191/// * `true`: There MAY be rows that match the predicate
192///
193/// * `false`: There are no rows that could possibly match the predicate (the
194///   predicate can never possibly be true). The container can be pruned (skipped)
195///   entirely.
196///
197/// While `PruningPredicate` will never return a `NULL` value, the
198/// rewritten predicate (as returned by `build_predicate_expression` and used internally
199/// by `PruningPredicate`) may evaluate to `NULL` when some of the min/max values
200/// or null / row counts are not known.
201///
202/// In order to be correct, `PruningPredicate` must return false
203/// **only** if it can determine that for all rows in the container, the
204/// predicate could never evaluate to `true` (always evaluates to either `NULL`
205/// or `false`).
206///
207/// ## Contains Analysis and Min/Max Rewrite
208///
209/// `PruningPredicate` works by first analyzing the predicate to see what
210/// [`LiteralGuarantee`] must hold for the predicate to be true.
211///
212/// Then, the `PruningPredicate` rewrites the original predicate into an
213/// expression that references the min/max values of each column in the original
214/// predicate.
215///
216/// When the min/max values are actually substituted in to this expression and
217/// evaluated, the result means
218///
219/// * `true`: there MAY be rows that pass the predicate, **KEEPS** the container
220///
221/// * `NULL`: there MAY be rows that pass the predicate, **KEEPS** the container
222///   Note that rewritten predicate can evaluate to NULL when some of
223///   the min/max values are not known. *Note that this is different than
224///   the SQL filter semantics where `NULL` means the row is filtered
225///   out.*
226///
227/// * `false`: there are no rows that could possibly match the predicate,
228///   **PRUNES** the container
229///
230/// For example, given a column `x`, the `x_min`, `x_max`, `x_null_count`, and
231/// `x_row_count` represent the minimum and maximum values, the null count of
232/// column `x`, and the row count of column `x`, provided by the `PruningStatistics`.
233/// `x_null_count` and `x_row_count` are used to handle the case where the column `x`
234/// is known to be all `NULL`s. Note this is different from knowing nothing about
235/// the column `x`, which confusingly is encoded by returning `NULL` for the min/max
236/// values from [`PruningStatistics::max_values`] and [`PruningStatistics::min_values`].
237///
238/// Here are some examples of the rewritten predicates:
239///
240/// Original Predicate | Rewritten Predicate
241/// ------------------ | --------------------
242/// `x = 5` | `x_null_count != x_row_count AND (x_min <= 5 AND 5 <= x_max)`
243/// `x < 5` | `x_null_count != x_row_count AND (x_min < 5)`
244/// `x = 5 AND y = 10` | `x_null_count != x_row_count AND (x_min <= 5 AND 5 <= x_max) AND y_null_count != y_row_count (y_min <= 10 AND 10 <= y_max)`
245/// `x IS NULL`  | `x_null_count > 0`
246/// `x IS NOT NULL`  | `x_null_count != row_count`
247/// `CAST(x as int) = 5` | `x_null_count != x_row_count (CAST(x_min as int) <= 5 AND 5 <= CAST(x_max as int))`
248///
249/// ## Predicate Evaluation
250/// The PruningPredicate works in two passes
251///
252/// **First pass**:  For each `LiteralGuarantee` calls
253/// [`PruningStatistics::contained`] and rules out containers where the
254/// LiteralGuarantees are not satisfied
255///
256/// **Second Pass**: Evaluates the rewritten expression using the
257/// min/max/null_counts/row_counts values for each column for each container. For any
258/// container that this expression evaluates to `false`, it rules out those
259/// containers.
260///
261///
262/// ### Example 1
263///
264/// Given the predicate, `x = 5 AND y = 10`, the rewritten predicate would look like:
265///
266/// ```sql
267/// x_null_count != x_row_count AND (x_min <= 5 AND 5 <= x_max)
268/// AND
269/// y_null_count != y_row_count AND (y_min <= 10 AND 10 <= y_max)
270/// ```
271///
272/// If we know that for a given container, `x` is between `1 and 100` and we know that
273/// `y` is between `4` and `7`, we know nothing about the null count and row count of
274/// `x` and `y`, the input statistics might look like:
275///
276/// Column   | Value
277/// -------- | -----
278/// `x_min`  | `1`
279/// `x_max`  | `100`
280/// `x_null_count` | `null`
281/// `x_row_count`  | `null`
282/// `y_min`  | `4`
283/// `y_max`  | `7`
284/// `y_null_count` | `null`
285/// `y_row_count`  | `null`
286///
287/// When these statistics values are substituted in to the rewritten predicate and
288/// simplified, the result is `false`:
289///
290/// * `null != null AND (1 <= 5 AND 5 <= 100) AND null != null AND (4 <= 10 AND 10 <= 7)`
291/// * `null = null` is `null` which is not true, so the AND moves on to the next clause
292/// * `null and (1 <= 5 AND 5 <= 100) AND null AND (4 <= 10 AND 10 <= 7)`
293/// * evaluating the clauses further we get:
294/// * `null and true and null and false`
295/// * `null and false`
296/// * `false`
297///
298/// Returning `false` means the container can be pruned, which matches the
299/// intuition that  `x = 5 AND y = 10` can’t be true for any row if all values of `y`
300/// are `7` or less.
301///
302/// Note that if we had ended up with `null AND true AND null AND true` the result
303/// would have been `null`.
304/// `null` is treated the same as`true`, because we can't prove that the predicate is `false.`
305///
306/// If, for some other container, we knew `y` was between the values `4` and
307/// `15`, then the rewritten predicate evaluates to `true` (verifying this is
308/// left as an exercise to the reader -- are you still here?), and the container
309/// **could not** be pruned. The intuition is that there may be rows where the
310/// predicate *might* evaluate to `true`, and the only way to find out is to do
311/// more analysis, for example by actually reading the data and evaluating the
312/// predicate row by row.
313///
314/// ### Example 2
315///
316/// Given the same predicate, `x = 5 AND y = 10`, the rewritten predicate would
317/// look like the same as example 1:
318///
319/// ```sql
320/// x_null_count != x_row_count AND (x_min <= 5 AND 5 <= x_max)
321/// AND
322/// y_null_count != y_row_count AND (y_min <= 10 AND 10 <= y_max)
323/// ```
324///
325/// If we know that for another given container, `x_min` is NULL and `x_max` is
326/// NULL (the min/max values are unknown), `x_null_count` is `100` and `x_row_count`
327///  is `100`; we know that `y` is between `4` and `7`, but we know nothing about
328/// the null count and row count of `y`. The input statistics might look like:
329///
330/// Column   | Value
331/// -------- | -----
332/// `x_min`  | `null`
333/// `x_max`  | `null`
334/// `x_null_count` | `100`
335/// `x_row_count`  | `100`
336/// `y_min`  | `4`
337/// `y_max`  | `7`
338/// `y_null_count` | `null`
339/// `y_row_count`  | `null`
340///
341/// When these statistics values are substituted in to the rewritten predicate and
342/// simplified, the result is `false`:
343///
344/// * `100 != 100 AND (null <= 5 AND 5 <= null) AND null = null AND (4 <= 10 AND 10 <= 7)`
345/// * `false AND null AND null AND false`
346/// * `false AND false`
347/// * `false`
348///
349/// Returning `false` means the container can be pruned, which matches the
350/// intuition that  `x = 5 AND y = 10` can’t be true because all values in `x`
351/// are known to be NULL.
352///
353/// # Related Work
354///
355/// [`PruningPredicate`] implements the type of min/max pruning described in
356/// Section `3.3.3` of the [`Snowflake SIGMOD Paper`]. The technique is
357/// described by various research such as [small materialized aggregates], [zone
358/// maps], and [data skipping].
359///
360/// [`Snowflake SIGMOD Paper`]: https://dl.acm.org/doi/10.1145/2882903.2903741
361/// [small materialized aggregates]: https://www.vldb.org/conf/1998/p476.pdf
362/// [zone maps]: https://dl.acm.org/doi/10.1007/978-3-642-03730-6_10
363/// [data skipping]: https://dl.acm.org/doi/10.1145/2588555.2610515
364#[derive(Debug, Clone)]
365pub struct PruningPredicate {
366    /// The input schema against which the predicate will be evaluated
367    schema: SchemaRef,
368    /// A min/max pruning predicate (rewritten in terms of column min/max
369    /// values, which are supplied by statistics)
370    predicate_expr: Arc<dyn PhysicalExpr>,
371    /// Description of which statistics are required to evaluate `predicate_expr`
372    required_columns: RequiredColumns,
373    /// Original physical predicate from which this predicate expr is derived
374    /// (required for serialization)
375    orig_expr: Arc<dyn PhysicalExpr>,
376    /// [`LiteralGuarantee`]s used to try and prove a predicate can not possibly
377    /// evaluate to `true`.
378    ///
379    /// See [`PruningPredicate::literal_guarantees`] for more details.
380    literal_guarantees: Vec<LiteralGuarantee>,
381}
382
383/// Build a pruning predicate from an optional predicate expression.
384/// If the predicate is None or the predicate cannot be converted to a pruning
385/// predicate, return None.
386/// If there is an error creating the pruning predicate it is recorded by incrementing
387/// the `predicate_creation_errors` counter.
388pub fn build_pruning_predicate(
389    predicate: Arc<dyn PhysicalExpr>,
390    file_schema: &SchemaRef,
391    predicate_creation_errors: &Count,
392) -> Option<Arc<PruningPredicate>> {
393    PruningPredicateBuilder::new()
394        .with_file_schema(Arc::clone(file_schema))
395        .with_error_counter(predicate_creation_errors)
396        .build(predicate)
397}
398
399/// Builder for a [`PruningPredicate`]. Groups optional configuration —
400/// `IN (...)` rewrite cap, error counter — so future additions do not
401/// churn the top-level API.
402///
403/// The two entry points are:
404///  - [`Self::build`]: convenience for scan sites that already track a
405///    `predicate_creation_errors` counter. Returns `Some(Arc<..>)` when the
406///    resulting predicate can actually prune, `None` when it is trivially
407///    true or when construction failed (in which case the error counter is
408///    incremented if one was supplied).
409///  - [`Self::try_build`]: returns a raw `Result<PruningPredicate>` for
410///    callers that want to surface errors themselves.
411///
412#[derive(Default)]
413pub struct PruningPredicateBuilder<'a> {
414    file_schema: Option<SchemaRef>,
415    error_counter: Option<&'a Count>,
416    max_in_list_size: usize,
417}
418
419impl<'a> PruningPredicateBuilder<'a> {
420    /// Create a new builder with the default pruning predicate configuration.
421    pub fn new() -> Self {
422        Self {
423            file_schema: None,
424            error_counter: None,
425            max_in_list_size: MAX_IN_LIST_SIZE,
426        }
427    }
428
429    /// Set the schema of the container that will be pruned (typically the
430    /// parquet file schema).
431    pub fn with_file_schema(mut self, file_schema: SchemaRef) -> Self {
432        self.file_schema = Some(file_schema);
433        self
434    }
435
436    /// Metric counter incremented once per predicate that fails to build.
437    /// Only consulted by [`Self::build`]; [`Self::try_build`] surfaces the
438    /// error directly.
439    pub fn with_error_counter(mut self, error_counter: &'a Count) -> Self {
440        self.error_counter = Some(error_counter);
441        self
442    }
443
444    /// Cap on the size of `IN (...)` lists that will be rewritten into per-
445    /// value min/max statistics checks. Lists longer than this fall back to
446    /// the unhandled-predicate hook (typically "keep the container").
447    ///
448    /// Query engines typically pass
449    /// `datafusion.execution.parquet.max_in_list_size` here.
450    pub fn with_max_in_list_size(mut self, max_in_list_size: usize) -> Self {
451        self.max_in_list_size = max_in_list_size;
452        self
453    }
454
455    /// Build a [`PruningPredicate`] wrapped in `Some(Arc<..>)` when it can
456    /// prune, `None` when it is trivially true or when construction fails.
457    /// If [`Self::with_error_counter`] was set, construction failures are
458    /// recorded there.
459    pub fn build(
460        self,
461        predicate: Arc<dyn PhysicalExpr>,
462    ) -> Option<Arc<PruningPredicate>> {
463        let error_counter = self.error_counter;
464        match self.try_build(predicate) {
465            Ok(pruning_predicate) => {
466                if !pruning_predicate.always_true() {
467                    return Some(Arc::new(pruning_predicate));
468                }
469            }
470            Err(e) => {
471                debug!("Could not create pruning predicate for: {e}");
472                if let Some(counter) = error_counter {
473                    counter.add(1);
474                }
475            }
476        }
477        None
478    }
479
480    /// Build a [`PruningPredicate`], returning the construction error
481    /// directly. Callers that want the always-true predicate elided or
482    /// errors folded into a counter should use [`Self::build`] instead.
483    pub fn try_build(
484        self,
485        mut predicate: Arc<dyn PhysicalExpr>,
486    ) -> Result<PruningPredicate> {
487        let file_schema = self.file_schema.ok_or_else(|| {
488            _internal_datafusion_err!(
489                "PruningPredicateBuilder requires a file schema (call `with_file_schema`)"
490            )
491        })?;
492
493        // Get a (simpler) snapshot of the physical expr here to use with `PruningPredicate`.
494        // In particular this unravels any `DynamicFilterPhysicalExpr`s by snapshotting them
495        // so that PruningPredicate can work with a static expression.
496        let tf = snapshot_physical_expr_opt(predicate)?;
497        if tf.transformed {
498            // If we had an expression such as Dynamic(part_col < 5 and col < 10)
499            // (this could come from something like `select * from t order by part_col, col, limit 10`)
500            // after snapshotting and because `DynamicFilterPhysicalExpr` applies child replacements to its
501            // children after snapshotting and previously `replace_columns_with_literals` may have been called with partition values
502            // the expression we have now is `8 < 5 and col < 10`.
503            // Thus we need as simplifier pass to get `false and col < 10` => `false` here.
504            let simplifier = PhysicalExprSimplifier::new(&file_schema);
505            predicate = simplifier.simplify(tf.data)?;
506        } else {
507            predicate = tf.data;
508        }
509        let unhandled_hook = Arc::new(ConstantUnhandledPredicateHook::default()) as _;
510
511        // build predicate expression once
512        let mut required_columns = RequiredColumns::new();
513        let predicate_expr = build_predicate_expression(
514            &predicate,
515            &file_schema,
516            &mut required_columns,
517            &unhandled_hook,
518            self.max_in_list_size,
519        );
520        let predicate_schema = required_columns.schema();
521        // Simplify the newly created predicate to get rid of redundant casts, comparisons, etc.
522        let predicate_expr =
523            PhysicalExprSimplifier::new(&predicate_schema).simplify(predicate_expr)?;
524        let literal_guarantees = LiteralGuarantee::analyze(&predicate);
525
526        Ok(PruningPredicate {
527            schema: file_schema,
528            predicate_expr,
529            required_columns,
530            orig_expr: predicate,
531            literal_guarantees,
532        })
533    }
534}
535
536/// Rewrites predicates that [`PredicateRewriter`] can not handle, e.g. certain
537/// complex expressions or predicates that reference columns that are not in the
538/// schema.
539pub trait UnhandledPredicateHook {
540    /// Called when a predicate can not be rewritten in terms of statistics or
541    /// references a column that is not in the schema.
542    fn handle(&self, expr: &Arc<dyn PhysicalExpr>) -> Arc<dyn PhysicalExpr>;
543}
544
545/// The default handling for unhandled predicates is to return a constant `true`
546/// (meaning don't prune the container)
547#[derive(Debug, Clone)]
548struct ConstantUnhandledPredicateHook {
549    default: Arc<dyn PhysicalExpr>,
550}
551
552impl Default for ConstantUnhandledPredicateHook {
553    fn default() -> Self {
554        Self {
555            default: Arc::new(phys_expr::Literal::new(ScalarValue::from(true))),
556        }
557    }
558}
559
560impl UnhandledPredicateHook for ConstantUnhandledPredicateHook {
561    fn handle(&self, _expr: &Arc<dyn PhysicalExpr>) -> Arc<dyn PhysicalExpr> {
562        Arc::clone(&self.default)
563    }
564}
565
566impl PruningPredicate {
567    /// Try to create a new instance of [`PruningPredicate`]
568    ///
569    /// This will translate the provided `expr` filter expression into
570    /// a *pruning predicate*.
571    ///
572    /// A pruning predicate is one that has been rewritten in terms of
573    /// the min and max values of column references and that evaluates
574    /// to FALSE if the filter predicate would evaluate FALSE *for
575    /// every row* whose values fell within the min / max ranges (aka
576    /// could be pruned).
577    ///
578    /// The pruning predicate evaluates to TRUE or NULL
579    /// if the filter predicate *might* evaluate to TRUE for at least
580    /// one row whose values fell within the min/max ranges (in other
581    /// words they might pass the predicate)
582    ///
583    /// For example, the filter expression `(column / 2) = 4` becomes
584    /// the pruning predicate
585    /// `(column_min / 2) <= 4 && 4 <= (column_max / 2))`
586    ///
587    /// See the struct level documentation on [`PruningPredicate`] for more
588    /// details.
589    ///
590    /// Note that `PruningPredicate` does not attempt to normalize or simplify
591    /// the input expression unless calling [`snapshot_physical_expr_opt`]
592    /// returns a new expression.
593    /// It is recommended that you pass the expressions through [`PhysicalExprSimplifier`]
594    /// before calling this method to make sure the expressions can be used for pruning.
595    ///
596    /// Use [`PruningPredicateBuilder`] to construct new pruning predicates.
597    #[deprecated(since = "55.0.0", note = "Use PruningPredicateBuilder instead")]
598    pub fn try_new(expr: Arc<dyn PhysicalExpr>, schema: SchemaRef) -> Result<Self> {
599        PruningPredicateBuilder::new()
600            .with_file_schema(schema)
601            .try_build(expr)
602    }
603
604    /// For each set of statistics, evaluates the pruning predicate
605    /// and returns a `bool` with the following meaning for a
606    /// all rows whose values match the statistics:
607    ///
608    /// `true`: There MAY be rows that match the predicate
609    ///
610    /// `false`: There are no rows that could possibly match the predicate
611    ///
612    /// Note: the predicate passed to `prune` should already be simplified as
613    /// much as possible (e.g. this pass doesn't handle some
614    /// expressions like `b = false`, but it does handle the
615    /// simplified version `b`. See [`ExprSimplifier`] to simplify expressions.
616    ///
617    /// [`ExprSimplifier`]: https://docs.rs/datafusion/latest/datafusion/optimizer/simplify_expressions/struct.ExprSimplifier.html
618    pub fn prune<S: PruningStatistics + ?Sized>(
619        &self,
620        statistics: &S,
621    ) -> Result<Vec<bool>> {
622        let mut builder = BoolVecBuilder::new(statistics.num_containers());
623
624        // Try to prove the predicate can't be true for the containers based on
625        // literal guarantees
626        for literal_guarantee in &self.literal_guarantees {
627            let LiteralGuarantee {
628                column,
629                guarantee,
630                literals,
631            } = literal_guarantee;
632            if let Some(results) = statistics.contained(column, literals) {
633                match guarantee {
634                    // `In` means the values in the column must be one of the
635                    // values in the set for the predicate to evaluate to true.
636                    // If `contained` returns false, that means the column is
637                    // not any of the values so we can prune the container
638                    Guarantee::In => builder.combine_array(&results),
639                    // `NotIn` means the values in the column must not be
640                    // any of the values in the set for the predicate to
641                    // evaluate to true. If `contained` returns true, it means the
642                    // column is only in the set of values so we can prune the
643                    // container
644                    Guarantee::NotIn => {
645                        builder.combine_array(&arrow::compute::not(&results)?)
646                    }
647                }
648                // if all containers are pruned (has rows that DEFINITELY DO NOT pass the predicate)
649                // can return early without evaluating the rest of predicates.
650                if builder.check_all_pruned() {
651                    return Ok(builder.build());
652                }
653            }
654        }
655
656        // Next, try to prove the predicate can't be true for the containers based
657        // on min/max values
658
659        // build a RecordBatch that contains the min/max values in the
660        // appropriate statistics columns for the min/max predicate
661        let statistics_batch =
662            build_statistics_record_batch(statistics, &self.required_columns)?;
663
664        // Evaluate the pruning predicate on that record batch and append any results to the builder
665        builder.combine_value(self.predicate_expr.evaluate(&statistics_batch)?);
666
667        Ok(builder.build())
668    }
669
670    /// Return a reference to the input schema
671    pub fn schema(&self) -> &SchemaRef {
672        &self.schema
673    }
674
675    /// Returns a reference to the physical expr used to construct this pruning predicate
676    pub fn orig_expr(&self) -> &Arc<dyn PhysicalExpr> {
677        &self.orig_expr
678    }
679
680    /// Returns a reference to the predicate expr
681    pub fn predicate_expr(&self) -> &Arc<dyn PhysicalExpr> {
682        &self.predicate_expr
683    }
684
685    /// Returns a reference to the literal guarantees
686    ///
687    /// Note that **All** `LiteralGuarantee`s must be satisfied for the
688    /// expression to possibly be `true`. If any is not satisfied, the
689    /// expression is guaranteed to be `null` or `false`.
690    pub fn literal_guarantees(&self) -> &[LiteralGuarantee] {
691        &self.literal_guarantees
692    }
693
694    /// Returns true if this pruning predicate can not prune anything.
695    ///
696    /// This happens if the predicate is a literal `true`  and
697    /// literal_guarantees is empty.
698    ///
699    /// This can happen when a predicate is simplified to a constant `true`
700    pub fn always_true(&self) -> bool {
701        is_always_true(&self.predicate_expr) && self.literal_guarantees.is_empty()
702    }
703
704    pub fn required_columns(&self) -> &RequiredColumns {
705        &self.required_columns
706    }
707
708    /// Names of the columns that are known to be / not be in a set
709    /// of literals (constants). These are the columns the that may be passed to
710    /// [`PruningStatistics::contained`] during pruning.
711    ///
712    /// This is useful to avoid fetching statistics for columns that will not be
713    /// used in the predicate. For example, it can be used to avoid reading
714    /// unneeded bloom filters (a non trivial operation).
715    pub fn literal_columns(&self) -> Vec<String> {
716        let mut seen = HashSet::new();
717        self.literal_guarantees
718            .iter()
719            .map(|e| &e.column.name)
720            // avoid duplicates
721            .filter(|name| seen.insert(*name))
722            .map(|s| s.to_string())
723            .collect()
724    }
725}
726
727/// Builds the return `Vec` for [`PruningPredicate::prune`].
728#[derive(Debug)]
729struct BoolVecBuilder {
730    /// One element per container. Each element is
731    /// * `true`: if the container has row that may pass the predicate
732    /// * `false`: if the container has rows that DEFINITELY DO NOT pass the predicate
733    inner: Vec<bool>,
734}
735
736impl BoolVecBuilder {
737    /// Create a new `BoolVecBuilder` with `num_containers` elements
738    fn new(num_containers: usize) -> Self {
739        Self {
740            // assume by default all containers may pass the predicate
741            inner: vec![true; num_containers],
742        }
743    }
744
745    /// Combines result `array` for a conjunct (e.g. `AND` clause) of a
746    /// predicate into the currently in progress array.
747    ///
748    /// Each `array` element is:
749    /// * `true`: container has row that may pass the predicate
750    /// * `false`: all container rows DEFINITELY DO NOT pass the predicate
751    /// * `null`: container may or may not have rows that pass the predicate
752    fn combine_array(&mut self, array: &BooleanArray) {
753        assert_eq!(array.len(), self.inner.len());
754        for (cur, new) in self.inner.iter_mut().zip(array.iter()) {
755            // `false` for this conjunct means we know for sure no rows could
756            // pass the predicate and thus we set the corresponding container
757            // location to false.
758            if let Some(false) = new {
759                *cur = false;
760            }
761        }
762    }
763
764    /// Combines the results in the [`ColumnarValue`] to the currently in
765    /// progress array, following the same rules as [`Self::combine_array`].
766    ///
767    /// # Panics
768    /// If `value` is not boolean
769    fn combine_value(&mut self, value: ColumnarValue) {
770        match value {
771            ColumnarValue::Array(array) => {
772                self.combine_array(array.as_boolean());
773            }
774            ColumnarValue::Scalar(ScalarValue::Boolean(Some(false))) => {
775                // False means all containers can not pass the predicate
776                self.inner = vec![false; self.inner.len()];
777            }
778            _ => {
779                // Null or true means the rows in container may pass this
780                // conjunct so we can't prune any containers based on that
781            }
782        }
783    }
784
785    /// Convert this builder into a Vec of bools
786    fn build(self) -> Vec<bool> {
787        self.inner
788    }
789
790    /// Check all containers has rows that DEFINITELY DO NOT pass the predicate
791    fn check_all_pruned(&self) -> bool {
792        self.inner.iter().all(|&x| !x)
793    }
794}
795
796fn is_always_true(expr: &Arc<dyn PhysicalExpr>) -> bool {
797    expr.downcast_ref::<phys_expr::Literal>()
798        .map(|l| matches!(l.value(), ScalarValue::Boolean(Some(true))))
799        .unwrap_or_default()
800}
801
802fn is_always_false(expr: &Arc<dyn PhysicalExpr>) -> bool {
803    expr.downcast_ref::<phys_expr::Literal>()
804        .map(|l| matches!(l.value(), ScalarValue::Boolean(Some(false))))
805        .unwrap_or_default()
806}
807
808/// Describes which columns statistics are necessary to evaluate a
809/// [`PruningPredicate`].
810///
811/// This structure permits reading and creating the minimum number statistics,
812/// which is important since statistics may be non trivial to read (e.g. large
813/// strings or when there are 1000s of columns).
814///
815/// Handles creating references to the min/max statistics
816/// for columns as well as recording which statistics are needed
817#[derive(Debug, Default, Clone)]
818pub struct RequiredColumns {
819    /// The statistics required to evaluate this predicate:
820    /// * The unqualified column in the input schema
821    /// * Statistics type (e.g. Min or Max or Null_Count)
822    /// * The field the statistics value should be placed in for
823    ///   pruning predicate evaluation (e.g. `min_value` or `max_value`)
824    columns: Vec<(phys_expr::Column, StatisticsType, Field)>,
825}
826
827impl RequiredColumns {
828    fn new() -> Self {
829        Self::default()
830    }
831
832    /// Returns Some(column) if this is a single column predicate.
833    ///
834    /// Returns None if this is a multi-column predicate.
835    ///
836    /// Examples:
837    /// * `a > 5 OR a < 10` returns `Some(a)`
838    /// * `a > 5 OR b < 10` returns `None`
839    /// * `true` returns None
840    pub fn single_column(&self) -> Option<&phys_expr::Column> {
841        if self.columns.windows(2).all(|w| {
842            // check if all columns are the same (ignoring statistics and field)
843            let c1 = &w[0].0;
844            let c2 = &w[1].0;
845            c1 == c2
846        }) {
847            self.columns.first().map(|r| &r.0)
848        } else {
849            None
850        }
851    }
852
853    /// Returns a schema that describes the columns required to evaluate this
854    /// pruning predicate.
855    /// The schema contains the fields for each column in `self.columns` with
856    /// the appropriate data type for the statistics.
857    /// Order matters, this same order is used to evaluate the
858    /// pruning predicate.
859    fn schema(&self) -> Schema {
860        let fields = self
861            .columns
862            .iter()
863            .map(|(_c, _t, f)| f.clone())
864            .collect::<Vec<_>>();
865        Schema::new(fields)
866    }
867
868    /// Returns an iterator over items in columns (see doc on
869    /// `self.columns` for details)
870    pub(crate) fn iter(
871        &self,
872    ) -> impl Iterator<Item = &(phys_expr::Column, StatisticsType, Field)> {
873        self.columns.iter()
874    }
875
876    fn find_stat_column(
877        &self,
878        column: &phys_expr::Column,
879        statistics_type: StatisticsType,
880    ) -> Option<usize> {
881        match statistics_type {
882            StatisticsType::RowCount => {
883                // Use the first row count we find, if any
884                self.columns
885                    .iter()
886                    .enumerate()
887                    .find(|(_i, (_c, t, _f))| t == &statistics_type)
888                    .map(|(i, (_c, _t, _f))| i)
889            }
890            _ => self
891                .columns
892                .iter()
893                .enumerate()
894                .find(|(_i, (c, t, _f))| c == column && t == &statistics_type)
895                .map(|(i, (_c, _t, _f))| i),
896        }
897    }
898
899    /// Rewrites column_expr so that all appearances of column
900    /// are replaced with a reference to either the min or max
901    /// statistics column, while keeping track that a reference to the statistics
902    /// column is required
903    ///
904    /// for example, an expression like `col("foo") > 5`, when called
905    /// with Max would result in an expression like `col("foo_max") >
906    /// 5` with the appropriate entry noted in self.columns
907    fn stat_column_expr(
908        &mut self,
909        column: &phys_expr::Column,
910        column_expr: &Arc<dyn PhysicalExpr>,
911        field: &Field,
912        stat_type: StatisticsType,
913    ) -> Result<Arc<dyn PhysicalExpr>> {
914        let (idx, need_to_insert) = match self.find_stat_column(column, stat_type) {
915            Some(idx) => (idx, false),
916            None => (self.columns.len(), true),
917        };
918
919        let column_name = column.name();
920        let stat_column_name = match stat_type {
921            StatisticsType::Min => format!("{column_name}_min"),
922            StatisticsType::Max => format!("{column_name}_max"),
923            StatisticsType::NullCount => format!("{column_name}_null_count"),
924            StatisticsType::RowCount => "row_count".to_string(),
925        };
926
927        let stat_column = phys_expr::Column::new(&stat_column_name, idx);
928
929        // only add statistics column if not previously added
930        if need_to_insert {
931            // may be null if statistics are not present
932            let nullable = true;
933            let stat_field =
934                Field::new(stat_column.name(), field.data_type().clone(), nullable);
935            self.columns.push((column.clone(), stat_type, stat_field));
936        }
937        rewrite_column_expr(Arc::clone(column_expr), column, &stat_column)
938    }
939
940    /// rewrite col --> col_min
941    fn min_column_expr(
942        &mut self,
943        column: &phys_expr::Column,
944        column_expr: &Arc<dyn PhysicalExpr>,
945        field: &Field,
946    ) -> Result<Arc<dyn PhysicalExpr>> {
947        self.stat_column_expr(column, column_expr, field, StatisticsType::Min)
948    }
949
950    /// rewrite col --> col_max
951    fn max_column_expr(
952        &mut self,
953        column: &phys_expr::Column,
954        column_expr: &Arc<dyn PhysicalExpr>,
955        field: &Field,
956    ) -> Result<Arc<dyn PhysicalExpr>> {
957        self.stat_column_expr(column, column_expr, field, StatisticsType::Max)
958    }
959
960    /// rewrite col --> col_null_count
961    fn null_count_column_expr(
962        &mut self,
963        column: &phys_expr::Column,
964        column_expr: &Arc<dyn PhysicalExpr>,
965        field: &Field,
966    ) -> Result<Arc<dyn PhysicalExpr>> {
967        self.stat_column_expr(column, column_expr, field, StatisticsType::NullCount)
968    }
969
970    /// rewrite col --> col_row_count
971    fn row_count_column_expr(
972        &mut self,
973        column: &phys_expr::Column,
974        column_expr: &Arc<dyn PhysicalExpr>,
975        field: &Field,
976    ) -> Result<Arc<dyn PhysicalExpr>> {
977        self.stat_column_expr(column, column_expr, field, StatisticsType::RowCount)
978    }
979}
980
981impl From<Vec<(phys_expr::Column, StatisticsType, Field)>> for RequiredColumns {
982    fn from(columns: Vec<(phys_expr::Column, StatisticsType, Field)>) -> Self {
983        Self { columns }
984    }
985}
986
987/// Build a RecordBatch from a list of statistics, creating arrays,
988/// with one row for each PruningStatistics and columns specified in
989/// the required_columns parameter.
990///
991/// For example, if the requested columns are
992/// ```text
993/// ("s1", Min, Field:s1_min)
994/// ("s2", Max, field:s2_max)
995/// ```
996///
997/// And the input statistics had
998/// ```text
999/// S1(Min: 5, Max: 10)
1000/// S2(Min: 99, Max: 1000)
1001/// S3(Min: 1, Max: 2)
1002/// ```
1003///
1004/// Then this function would build a record batch with 2 columns and
1005/// one row s1_min and s2_max as follows (s3 is not requested):
1006///
1007/// ```text
1008/// s1_min | s2_max
1009/// -------+--------
1010///   5    | 1000
1011/// ```
1012fn build_statistics_record_batch<S: PruningStatistics + ?Sized>(
1013    statistics: &S,
1014    required_columns: &RequiredColumns,
1015) -> Result<RecordBatch> {
1016    let mut arrays = Vec::<ArrayRef>::new();
1017    // For each needed statistics column:
1018    for (column, statistics_type, stat_field) in required_columns.iter() {
1019        let column = Column::from_name(column.name());
1020        let data_type = stat_field.data_type();
1021
1022        let num_containers = statistics.num_containers();
1023
1024        let array = match statistics_type {
1025            StatisticsType::Min => statistics.min_values(&column),
1026            StatisticsType::Max => statistics.max_values(&column),
1027            StatisticsType::NullCount => statistics.null_counts(&column),
1028            StatisticsType::RowCount => statistics.row_counts(),
1029        };
1030        let array = array.unwrap_or_else(|| new_null_array(data_type, num_containers));
1031
1032        assert_eq_or_internal_err!(
1033            num_containers,
1034            array.len(),
1035            "mismatched statistics length. Expected {}, got {}",
1036            num_containers,
1037            array.len()
1038        );
1039
1040        // cast statistics array to required data type (e.g. parquet
1041        // provides timestamp statistics as "Int64")
1042        let array = arrow::compute::cast(&array, data_type)?;
1043
1044        arrays.push(array);
1045    }
1046
1047    let schema = Arc::new(required_columns.schema());
1048    // provide the count in case there were no needed statistics
1049    let mut options = RecordBatchOptions::default();
1050    options.row_count = Some(statistics.num_containers());
1051
1052    trace!("Creating statistics batch for {required_columns:#?} with {arrays:#?}");
1053
1054    RecordBatch::try_new_with_options(schema, arrays, &options).map_err(|err| {
1055        plan_datafusion_err!("Can not create statistics record batch: {err}")
1056    })
1057}
1058
1059struct PruningExpressionBuilder<'a> {
1060    column: phys_expr::Column,
1061    column_expr: Arc<dyn PhysicalExpr>,
1062    op: Operator,
1063    scalar_expr: Arc<dyn PhysicalExpr>,
1064    field: &'a Field,
1065    required_columns: &'a mut RequiredColumns,
1066}
1067
1068impl<'a> PruningExpressionBuilder<'a> {
1069    fn try_new(
1070        left: &'a Arc<dyn PhysicalExpr>,
1071        right: &'a Arc<dyn PhysicalExpr>,
1072        left_columns: ColumnReferenceCount,
1073        right_columns: ColumnReferenceCount,
1074        op: Operator,
1075        schema: &'a SchemaRef,
1076        required_columns: &'a mut RequiredColumns,
1077    ) -> Result<Self> {
1078        // find column name; input could be a more complicated expression
1079        let (column_expr, scalar_expr, column, correct_operator) = match (
1080            left_columns,
1081            right_columns,
1082        ) {
1083            (ColumnReferenceCount::One(column), ColumnReferenceCount::Zero) => {
1084                (left, right, column, op)
1085            }
1086            (ColumnReferenceCount::Zero, ColumnReferenceCount::One(column)) => {
1087                (right, left, column, reverse_operator(op)?)
1088            }
1089            (ColumnReferenceCount::One(_), ColumnReferenceCount::One(_)) => {
1090                // both sides have one column - not supported
1091                return plan_err!(
1092                    "Expression not supported for pruning: left has 1 column, right has 1 column"
1093                );
1094            }
1095            (ColumnReferenceCount::Zero, ColumnReferenceCount::Zero) => {
1096                // both sides are literals - should be handled before calling try_new
1097                return plan_err!(
1098                    "Pruning literal expressions is not supported, please call PhysicalExprSimplifier first"
1099                );
1100            }
1101            (ColumnReferenceCount::Many, _) | (_, ColumnReferenceCount::Many) => {
1102                return plan_err!(
1103                    "Expression not supported for pruning: left or right has multiple columns"
1104                );
1105            }
1106        };
1107
1108        let df_schema = DFSchema::try_from(Arc::clone(schema))?;
1109        let (column_expr, correct_operator, scalar_expr) = rewrite_expr_to_prunable(
1110            column_expr,
1111            correct_operator,
1112            scalar_expr,
1113            df_schema,
1114        )?;
1115        let field = match schema.column_with_name(column.name()) {
1116            Some((_, f)) => f,
1117            _ => {
1118                return plan_err!("Field not found in schema");
1119            }
1120        };
1121
1122        Ok(Self {
1123            column,
1124            column_expr,
1125            op: correct_operator,
1126            scalar_expr,
1127            field,
1128            required_columns,
1129        })
1130    }
1131
1132    fn op(&self) -> Operator {
1133        self.op
1134    }
1135
1136    fn scalar_expr(&self) -> &Arc<dyn PhysicalExpr> {
1137        &self.scalar_expr
1138    }
1139
1140    fn min_column_expr(&mut self) -> Result<Arc<dyn PhysicalExpr>> {
1141        self.required_columns
1142            .min_column_expr(&self.column, &self.column_expr, self.field)
1143    }
1144
1145    fn max_column_expr(&mut self) -> Result<Arc<dyn PhysicalExpr>> {
1146        self.required_columns
1147            .max_column_expr(&self.column, &self.column_expr, self.field)
1148    }
1149
1150    /// This function is to simply retune the `null_count` physical expression no matter what the
1151    /// predicate expression is
1152    ///
1153    /// i.e., x > 5 => x_null_count,
1154    ///       cast(x as int) < 10 => x_null_count,
1155    ///       try_cast(x as float) < 10.0 => x_null_count
1156    fn null_count_column_expr(&mut self) -> Result<Arc<dyn PhysicalExpr>> {
1157        // Retune to [`phys_expr::Column`]
1158        let column_expr = Arc::new(self.column.clone()) as _;
1159
1160        // null_count is DataType::UInt64, which is different from the column's data type (i.e. self.field)
1161        let null_count_field = &Field::new(self.field.name(), DataType::UInt64, true);
1162
1163        self.required_columns.null_count_column_expr(
1164            &self.column,
1165            &column_expr,
1166            null_count_field,
1167        )
1168    }
1169
1170    /// This function is to simply retune the `row_count` physical expression no matter what the
1171    /// predicate expression is
1172    ///
1173    /// i.e., x > 5 => x_row_count,
1174    ///       cast(x as int) < 10 => x_row_count,
1175    ///       try_cast(x as float) < 10.0 => x_row_count
1176    fn row_count_column_expr(&mut self) -> Result<Arc<dyn PhysicalExpr>> {
1177        // Retune to [`phys_expr::Column`]
1178        let column_expr = Arc::new(self.column.clone()) as _;
1179
1180        // row_count is DataType::UInt64, which is different from the column's data type (i.e. self.field)
1181        let row_count_field = &Field::new(self.field.name(), DataType::UInt64, true);
1182
1183        self.required_columns.row_count_column_expr(
1184            &self.column,
1185            &column_expr,
1186            row_count_field,
1187        )
1188    }
1189}
1190
1191/// This function is designed to rewrite the column_expr to
1192/// ensure the column_expr is monotonically increasing.
1193///
1194/// For example,
1195/// 1. `col > 10`
1196/// 2. `-col > 10` should be rewritten to `col < -10`
1197/// 3. `!col = true` would be rewritten to `col = !true`
1198/// 4. `abs(a - 10) > 0` not supported
1199/// 5. `cast(can_prunable_expr) > 10`
1200/// 6. `try_cast(can_prunable_expr) > 10`
1201///
1202/// More rewrite rules are still in progress.
1203fn rewrite_expr_to_prunable(
1204    column_expr: &PhysicalExprRef,
1205    op: Operator,
1206    scalar_expr: &PhysicalExprRef,
1207    schema: DFSchema,
1208) -> Result<(PhysicalExprRef, Operator, PhysicalExprRef)> {
1209    if !is_compare_op(op) {
1210        return plan_err!("rewrite_expr_to_prunable only support compare expression");
1211    }
1212
1213    if column_expr.downcast_ref::<phys_expr::Column>().is_some() {
1214        // `col op lit()`
1215        Ok((Arc::clone(column_expr), op, Arc::clone(scalar_expr)))
1216    } else if let Some(cast) = column_expr.downcast_ref::<phys_expr::CastExpr>() {
1217        // `cast(col) op lit()`
1218        let (left, op, right) = rewrite_cast_child_to_prunable(
1219            cast.expr(),
1220            cast.cast_type(),
1221            op,
1222            scalar_expr,
1223            schema,
1224        )?;
1225        let left = Arc::new(phys_expr::CastExpr::new_with_target_field(
1226            left,
1227            Arc::clone(cast.target_field()),
1228            None,
1229        ));
1230        // PruningPredicate does not support pruning on nested fields yet.
1231        // End-to-end nested-field pruning also requires Parquet statistics
1232        // extraction to agree with PruningPredicate on a stats representation
1233        // for nested field expressions.
1234        Ok((left, op, right))
1235    } else if let Some(try_cast) = column_expr.downcast_ref::<phys_expr::TryCastExpr>() {
1236        // `try_cast(col) op lit()`
1237        let (left, op, right) = rewrite_cast_child_to_prunable(
1238            try_cast.expr(),
1239            try_cast.cast_type(),
1240            op,
1241            scalar_expr,
1242            schema,
1243        )?;
1244        let left = Arc::new(phys_expr::TryCastExpr::new(
1245            left,
1246            try_cast.cast_type().clone(),
1247        ));
1248        Ok((left, op, right))
1249    } else if let Some(neg) = column_expr.downcast_ref::<phys_expr::NegativeExpr>() {
1250        // `-col > lit()`  --> `col < -lit()`
1251        let (left, op, right) =
1252            rewrite_expr_to_prunable(neg.arg(), op, scalar_expr, schema)?;
1253        let right = Arc::new(phys_expr::NegativeExpr::new(right));
1254        Ok((left, reverse_operator(op)?, right))
1255    } else if let Some(not) = column_expr.downcast_ref::<phys_expr::NotExpr>() {
1256        // `!col = true` --> `col = !true`
1257        if !matches!(
1258            op,
1259            Operator::Eq
1260                | Operator::NotEq
1261                | Operator::IsDistinctFrom
1262                | Operator::IsNotDistinctFrom
1263        ) {
1264            return plan_err!(
1265                "Not with operator other than Eq / NotEq / IsDistinctFrom / IsNotDistinctFrom is not supported"
1266            );
1267        }
1268        if not.arg().downcast_ref::<phys_expr::Column>().is_some() {
1269            let left = Arc::clone(not.arg());
1270            let right = Arc::new(phys_expr::NotExpr::new(Arc::clone(scalar_expr)));
1271            Ok((left, reverse_operator(op)?, right))
1272        } else {
1273            plan_err!("Not with complex expression {column_expr:?} is not supported")
1274        }
1275    } else {
1276        plan_err!("column expression {column_expr:?} is not supported")
1277    }
1278}
1279
1280fn rewrite_cast_child_to_prunable(
1281    cast_child_expr: &PhysicalExprRef,
1282    cast_type: &DataType,
1283    op: Operator,
1284    scalar_expr: &PhysicalExprRef,
1285    schema: DFSchema,
1286) -> Result<(PhysicalExprRef, Operator, PhysicalExprRef)> {
1287    verify_support_type_for_prune(
1288        &cast_child_expr.data_type(schema.as_arrow())?,
1289        cast_type,
1290    )?;
1291    rewrite_expr_to_prunable(cast_child_expr, op, scalar_expr, schema)
1292}
1293
1294fn is_compare_op(op: Operator) -> bool {
1295    matches!(
1296        op,
1297        Operator::Eq
1298            | Operator::NotEq
1299            | Operator::Lt
1300            | Operator::LtEq
1301            | Operator::Gt
1302            | Operator::GtEq
1303            | Operator::IsDistinctFrom
1304            | Operator::IsNotDistinctFrom
1305            | Operator::LikeMatch
1306            | Operator::NotLikeMatch
1307    )
1308}
1309
1310// The pruning logic is based on the comparing the min/max bounds.
1311// Must make sure the two type has order.
1312// For example, casts from string to numbers is not correct.
1313// Because the "13" is less than "3" with UTF8 comparison order.
1314fn verify_support_type_for_prune(from_type: &DataType, to_type: &DataType) -> Result<()> {
1315    // Dictionary casts are always supported as long as the value types are supported
1316    let from_type = match from_type {
1317        DataType::Dictionary(_, t) => {
1318            return verify_support_type_for_prune(t.as_ref(), to_type);
1319        }
1320        _ => from_type,
1321    };
1322    let to_type = match to_type {
1323        DataType::Dictionary(_, t) => {
1324            return verify_support_type_for_prune(from_type, t.as_ref());
1325        }
1326        _ => to_type,
1327    };
1328    // If both types are strings or both are not strings (number, timestamp, etc)
1329    // then we can compare them.
1330    // PruningPredicate does not support casting of strings to numbers and such.
1331    if from_type.is_string() == to_type.is_string() {
1332        Ok(())
1333    } else {
1334        plan_err!(
1335            "Try Cast/Cast with from type {from_type} to type {to_type} is not supported"
1336        )
1337    }
1338}
1339
1340/// replaces a column with an old name with a new name in an expression
1341fn rewrite_column_expr(
1342    e: Arc<dyn PhysicalExpr>,
1343    column_old: &phys_expr::Column,
1344    column_new: &phys_expr::Column,
1345) -> Result<Arc<dyn PhysicalExpr>> {
1346    e.transform(|expr| {
1347        if let Some(column) = expr.downcast_ref::<phys_expr::Column>()
1348            && column == column_old
1349        {
1350            return Ok(Transformed::yes(Arc::new(column_new.clone())));
1351        }
1352
1353        Ok(Transformed::no(expr))
1354    })
1355    .data()
1356}
1357
1358fn reverse_operator(op: Operator) -> Result<Operator> {
1359    op.swap().ok_or_else(|| {
1360        internal_datafusion_err!(
1361            "Could not reverse operator {op} while building pruning predicate"
1362        )
1363    })
1364}
1365
1366/// Given a column reference to `column`, returns a pruning
1367/// expression in terms of the min and max that will evaluate to true
1368/// if the column may contain values, and false if definitely does not
1369/// contain values
1370fn build_single_column_expr(
1371    column: &phys_expr::Column,
1372    schema: &Schema,
1373    required_columns: &mut RequiredColumns,
1374    is_not: bool, // if true, treat as !col
1375) -> Option<Arc<dyn PhysicalExpr>> {
1376    let field = schema.field_with_name(column.name()).ok()?;
1377
1378    if *field.data_type() == DataType::Boolean {
1379        let col_ref = Arc::new(column.clone()) as _;
1380
1381        let min = required_columns
1382            .min_column_expr(column, &col_ref, field)
1383            .ok()?;
1384        let max = required_columns
1385            .max_column_expr(column, &col_ref, field)
1386            .ok()?;
1387
1388        // remember -- we want an expression that is:
1389        // TRUE: if there may be rows that match
1390        // FALSE: if there are no rows that match
1391        if is_not {
1392            // The only way we know a column couldn't match is if both the min and max are true
1393            // !(min && max)
1394            Some(Arc::new(phys_expr::NotExpr::new(Arc::new(
1395                phys_expr::BinaryExpr::new(min, Operator::And, max),
1396            ))))
1397        } else {
1398            // the only way we know a column couldn't match is if both the min and max are false
1399            // !(!min && !max) --> min || max
1400            Some(Arc::new(phys_expr::BinaryExpr::new(min, Operator::Or, max)))
1401        }
1402    } else {
1403        None
1404    }
1405}
1406
1407/// Given an expression reference to `expr`, if `expr` is a column expression,
1408/// returns a pruning expression in terms of IsNull that will evaluate to true
1409/// if the column may contain null, and false if definitely does not
1410/// contain null.
1411/// If `with_not` is true, build a pruning expression for `col IS NOT NULL`: `col_count != col_null_count`
1412/// The pruning expression evaluates to true ONLY if the column definitely CONTAINS
1413/// at least one NULL value.  In this case we can know that `IS NOT NULL` can not be true and
1414/// thus can prune the row group / value
1415fn build_is_null_column_expr(
1416    expr: &Arc<dyn PhysicalExpr>,
1417    schema: &Schema,
1418    required_columns: &mut RequiredColumns,
1419    with_not: bool,
1420) -> Option<Arc<dyn PhysicalExpr>> {
1421    if let Some(col) = expr.downcast_ref::<phys_expr::Column>() {
1422        let field = schema.field_with_name(col.name()).ok()?;
1423
1424        let null_count_field = &Field::new(field.name(), DataType::UInt64, true);
1425        if with_not {
1426            if let Ok(row_count_expr) =
1427                required_columns.row_count_column_expr(col, expr, null_count_field)
1428            {
1429                required_columns
1430                    .null_count_column_expr(col, expr, null_count_field)
1431                    .map(|null_count_column_expr| {
1432                        // IsNotNull(column) => null_count != row_count
1433                        Arc::new(phys_expr::BinaryExpr::new(
1434                            null_count_column_expr,
1435                            Operator::NotEq,
1436                            row_count_expr,
1437                        )) as _
1438                    })
1439                    .ok()
1440            } else {
1441                None
1442            }
1443        } else {
1444            required_columns
1445                .null_count_column_expr(col, expr, null_count_field)
1446                .map(|null_count_column_expr| {
1447                    // IsNull(column) => null_count > 0
1448                    Arc::new(phys_expr::BinaryExpr::new(
1449                        null_count_column_expr,
1450                        Operator::Gt,
1451                        Arc::new(phys_expr::Literal::new(ScalarValue::UInt64(Some(0)))),
1452                    )) as _
1453                })
1454                .ok()
1455        }
1456    } else {
1457        None
1458    }
1459}
1460
1461/// Default maximum number of entries in an `IN (...)` list that will be
1462/// rewritten into a chain of per-value min/max checks by
1463/// `build_predicate_expression`. Callers threading a [`PredicateRewriter`]
1464/// can override this via [`PredicateRewriter::with_max_in_list_size`], and
1465/// query engines can wire it from the
1466/// `datafusion.execution.parquet.max_in_list_size` config option.
1467pub const MAX_IN_LIST_SIZE: usize = 20;
1468
1469/// Rewrite a predicate expression in terms of statistics (min/max/null_counts)
1470/// for use as a [`PruningPredicate`].
1471pub struct PredicateRewriter {
1472    unhandled_hook: Arc<dyn UnhandledPredicateHook>,
1473    max_in_list_size: usize,
1474}
1475
1476impl Default for PredicateRewriter {
1477    fn default() -> Self {
1478        Self {
1479            unhandled_hook: Arc::new(ConstantUnhandledPredicateHook::default()),
1480            max_in_list_size: MAX_IN_LIST_SIZE,
1481        }
1482    }
1483}
1484
1485impl PredicateRewriter {
1486    /// Create a new `PredicateRewriter`
1487    pub fn new() -> Self {
1488        Self::default()
1489    }
1490
1491    /// Set the unhandled hook to be used when a predicate can not be rewritten
1492    pub fn with_unhandled_hook(
1493        mut self,
1494        unhandled_hook: Arc<dyn UnhandledPredicateHook>,
1495    ) -> Self {
1496        self.unhandled_hook = unhandled_hook;
1497        self
1498    }
1499
1500    /// Set the maximum size of an `IN (...)` list that will be rewritten into a
1501    /// chain of per-value statistics checks. Lists longer than this fall back
1502    /// to the unhandled-predicate hook (typically "keep the container"),
1503    /// effectively skipping container-level pruning for large IN lists.
1504    ///
1505    /// The default (see [`MAX_IN_LIST_SIZE`]) preserves the
1506    /// historical behaviour. Callers wiring config through can override via
1507    /// `datafusion.execution.max_in_list_size`.
1508    pub fn with_max_in_list_size(mut self, max_in_list_size: usize) -> Self {
1509        self.max_in_list_size = max_in_list_size;
1510        self
1511    }
1512
1513    /// Translate logical filter expression into pruning predicate
1514    /// expression that will evaluate to FALSE if it can be determined no
1515    /// rows between the min/max values could pass the predicates.
1516    ///
1517    /// Any predicates that can not be translated will be passed to `unhandled_hook`.
1518    ///
1519    /// Returns the pruning predicate as an [`PhysicalExpr`]
1520    ///
1521    /// Notice: `IN (...)` lists longer than `max_in_list_size` (default
1522    /// [`MAX_IN_LIST_SIZE`]) fall back to calling `unhandled_hook`.
1523    pub fn rewrite_predicate_to_statistics_predicate(
1524        &self,
1525        expr: &Arc<dyn PhysicalExpr>,
1526        schema: &Schema,
1527    ) -> Arc<dyn PhysicalExpr> {
1528        let mut required_columns = RequiredColumns::new();
1529        build_predicate_expression(
1530            expr,
1531            &Arc::new(schema.clone()),
1532            &mut required_columns,
1533            &self.unhandled_hook,
1534            self.max_in_list_size,
1535        )
1536    }
1537}
1538
1539/// Translate logical filter expression into pruning predicate
1540/// expression that will evaluate to FALSE if it can be determined no
1541/// rows between the min/max values could pass the predicates.
1542///
1543/// Any predicates that can not be translated will be passed to `unhandled_hook`.
1544///
1545/// Returns the pruning predicate as an [`PhysicalExpr`]
1546///
1547/// `max_in_list_size` is the largest `IN (...)` list that will be rewritten
1548/// into a chain of per-value statistics checks; longer lists fall back to
1549/// `unhandled_hook`.
1550fn build_predicate_expression(
1551    expr: &Arc<dyn PhysicalExpr>,
1552    schema: &SchemaRef,
1553    required_columns: &mut RequiredColumns,
1554    unhandled_hook: &Arc<dyn UnhandledPredicateHook>,
1555    max_in_list_size: usize,
1556) -> Arc<dyn PhysicalExpr> {
1557    if is_always_false(expr) {
1558        // Shouldn't return `unhandled_hook.handle(expr)`
1559        // Because it will transfer false to true.
1560        return Arc::clone(expr);
1561    }
1562    // predicate expression can only be a binary expression
1563    if let Some(is_null) = expr.downcast_ref::<phys_expr::IsNullExpr>() {
1564        return build_is_null_column_expr(is_null.arg(), schema, required_columns, false)
1565            .unwrap_or_else(|| unhandled_hook.handle(expr));
1566    }
1567    if let Some(is_not_null) = expr.downcast_ref::<phys_expr::IsNotNullExpr>() {
1568        return build_is_null_column_expr(
1569            is_not_null.arg(),
1570            schema,
1571            required_columns,
1572            true,
1573        )
1574        .unwrap_or_else(|| unhandled_hook.handle(expr));
1575    }
1576    if let Some(col) = expr.downcast_ref::<phys_expr::Column>() {
1577        return build_single_column_expr(col, schema, required_columns, false)
1578            .unwrap_or_else(|| unhandled_hook.handle(expr));
1579    }
1580    if let Some(not) = expr.downcast_ref::<phys_expr::NotExpr>() {
1581        // match !col (don't do so recursively)
1582        if let Some(col) = not.arg().downcast_ref::<phys_expr::Column>() {
1583            return build_single_column_expr(col, schema, required_columns, true)
1584                .unwrap_or_else(|| unhandled_hook.handle(expr));
1585        } else {
1586            return unhandled_hook.handle(expr);
1587        }
1588    }
1589    if let Some(in_list) = expr.downcast_ref::<phys_expr::InListExpr>() {
1590        if !in_list.list().is_empty() && in_list.list().len() <= max_in_list_size {
1591            let eq_op = if in_list.negated() {
1592                Operator::NotEq
1593            } else {
1594                Operator::Eq
1595            };
1596            let re_op = if in_list.negated() {
1597                Operator::And
1598            } else {
1599                Operator::Or
1600            };
1601            let change_expr = in_list
1602                .list()
1603                .iter()
1604                .map(|e| {
1605                    Arc::new(phys_expr::BinaryExpr::new(
1606                        Arc::clone(in_list.expr()),
1607                        eq_op,
1608                        Arc::clone(e),
1609                    )) as _
1610                })
1611                .reduce(|a, b| Arc::new(phys_expr::BinaryExpr::new(a, re_op, b)) as _)
1612                .unwrap();
1613            return build_predicate_expression(
1614                &change_expr,
1615                schema,
1616                required_columns,
1617                unhandled_hook,
1618                max_in_list_size,
1619            );
1620        } else {
1621            return unhandled_hook.handle(expr);
1622        }
1623    }
1624
1625    let (left, op, right) = {
1626        if let Some(bin_expr) = expr.downcast_ref::<phys_expr::BinaryExpr>() {
1627            (
1628                Arc::clone(bin_expr.left()),
1629                *bin_expr.op(),
1630                Arc::clone(bin_expr.right()),
1631            )
1632        } else if let Some(like_expr) = expr.downcast_ref::<phys_expr::LikeExpr>() {
1633            if like_expr.case_insensitive() {
1634                return unhandled_hook.handle(expr);
1635            }
1636            let op = match (like_expr.negated(), like_expr.case_insensitive()) {
1637                (false, false) => Operator::LikeMatch,
1638                (true, false) => Operator::NotLikeMatch,
1639                (false, true) => Operator::ILikeMatch,
1640                (true, true) => Operator::NotILikeMatch,
1641            };
1642            (
1643                Arc::clone(like_expr.expr()),
1644                op,
1645                Arc::clone(like_expr.pattern()),
1646            )
1647        } else {
1648            return unhandled_hook.handle(expr);
1649        }
1650    };
1651
1652    if op == Operator::And || op == Operator::Or {
1653        let left_expr = build_predicate_expression(
1654            &left,
1655            schema,
1656            required_columns,
1657            unhandled_hook,
1658            max_in_list_size,
1659        );
1660        let right_expr = build_predicate_expression(
1661            &right,
1662            schema,
1663            required_columns,
1664            unhandled_hook,
1665            max_in_list_size,
1666        );
1667        // simplify boolean expression if applicable
1668        let expr = match (&left_expr, op, &right_expr) {
1669            (left, Operator::And, right)
1670                if is_always_false(left) || is_always_false(right) =>
1671            {
1672                Arc::new(phys_expr::Literal::new(ScalarValue::Boolean(Some(false))))
1673            }
1674            (left, Operator::And, _) if is_always_true(left) => right_expr,
1675            (_, Operator::And, right) if is_always_true(right) => left_expr,
1676            (left, Operator::Or, right)
1677                if is_always_true(left) || is_always_true(right) =>
1678            {
1679                Arc::new(phys_expr::Literal::new(ScalarValue::Boolean(Some(true))))
1680            }
1681            (left, Operator::Or, _) if is_always_false(left) => right_expr,
1682            (_, Operator::Or, right) if is_always_false(right) => left_expr,
1683
1684            _ => Arc::new(phys_expr::BinaryExpr::new(left_expr, op, right_expr)),
1685        };
1686        return expr;
1687    }
1688
1689    let left_columns = ColumnReferenceCount::from_expression(&left);
1690    let right_columns = ColumnReferenceCount::from_expression(&right);
1691    let expr_builder = PruningExpressionBuilder::try_new(
1692        &left,
1693        &right,
1694        left_columns,
1695        right_columns,
1696        op,
1697        schema,
1698        required_columns,
1699    );
1700    let mut expr_builder = match expr_builder {
1701        Ok(builder) => builder,
1702        // allow partial failure in predicate expression generation
1703        // this can still produce a useful predicate when multiple conditions are joined using AND
1704        Err(e) => {
1705            debug!("Error building pruning expression: {e}");
1706            return unhandled_hook.handle(expr);
1707        }
1708    };
1709
1710    build_statistics_expr(&mut expr_builder)
1711        .unwrap_or_else(|_| unhandled_hook.handle(expr))
1712}
1713
1714/// Count of distinct column references in an expression.
1715/// This is the same as [`collect_columns`] but optimized to stop counting
1716/// once more than one distinct column is found.
1717///
1718/// For example, in expression `col1 + col2`, the count is `Many`.
1719/// In expression `col1 + 5`, the count is `One`.
1720/// In expression `5 + 10`, the count is `Zero`.
1721///
1722/// [`collect_columns`]: datafusion_physical_expr::utils::collect_columns
1723#[derive(Debug, PartialEq, Eq)]
1724enum ColumnReferenceCount {
1725    /// no column references
1726    Zero,
1727    /// Only one column reference
1728    One(phys_expr::Column),
1729    /// More than one column reference
1730    Many,
1731}
1732
1733impl ColumnReferenceCount {
1734    /// Count the number of distinct column references in an expression
1735    fn from_expression(expr: &Arc<dyn PhysicalExpr>) -> Self {
1736        let mut seen = HashSet::<phys_expr::Column>::new();
1737        expr.apply(|expr| {
1738            if let Some(column) = expr.downcast_ref::<phys_expr::Column>() {
1739                seen.insert(column.clone());
1740                if seen.len() > 1 {
1741                    return Ok(TreeNodeRecursion::Stop);
1742                }
1743            }
1744            Ok(TreeNodeRecursion::Continue)
1745        })
1746        // pre_visit always returns OK, so this will always too
1747        .expect("no way to return error during recursion");
1748        match seen.len() {
1749            0 => ColumnReferenceCount::Zero,
1750            1 => ColumnReferenceCount::One(
1751                seen.into_iter().next().expect("just checked len==1"),
1752            ),
1753            _ => ColumnReferenceCount::Many,
1754        }
1755    }
1756}
1757
1758fn build_statistics_expr(
1759    expr_builder: &mut PruningExpressionBuilder,
1760) -> Result<Arc<dyn PhysicalExpr>> {
1761    let statistics_expr: Arc<dyn PhysicalExpr> = match expr_builder.op() {
1762        Operator::NotEq => build_ne_statistics_expr(expr_builder)?,
1763        Operator::Eq => {
1764            // column = literal => (min, max) = literal => min <= literal && literal <= max
1765            // (column / 2) = 4 => (column_min / 2) <= 4 && 4 <= (column_max / 2)
1766            build_eq_statistics_expr(expr_builder)?
1767        }
1768        Operator::IsDistinctFrom => return build_is_distinct_from(expr_builder),
1769        Operator::IsNotDistinctFrom => return build_is_not_distinct_from(expr_builder),
1770        Operator::NotLikeMatch => build_not_like_match(expr_builder)?,
1771        Operator::LikeMatch => build_like_match(expr_builder).ok_or_else(|| {
1772            plan_datafusion_err!(
1773                "LIKE expression with wildcard at the beginning is not supported"
1774            )
1775        })?,
1776        Operator::Gt => {
1777            // column > literal => (min, max) > literal => max > literal
1778            Arc::new(phys_expr::BinaryExpr::new(
1779                expr_builder.max_column_expr()?,
1780                Operator::Gt,
1781                Arc::clone(expr_builder.scalar_expr()),
1782            ))
1783        }
1784        Operator::GtEq => {
1785            // column >= literal => (min, max) >= literal => max >= literal
1786            Arc::new(phys_expr::BinaryExpr::new(
1787                expr_builder.max_column_expr()?,
1788                Operator::GtEq,
1789                Arc::clone(expr_builder.scalar_expr()),
1790            ))
1791        }
1792        Operator::Lt => {
1793            // column < literal => (min, max) < literal => min < literal
1794            Arc::new(phys_expr::BinaryExpr::new(
1795                expr_builder.min_column_expr()?,
1796                Operator::Lt,
1797                Arc::clone(expr_builder.scalar_expr()),
1798            ))
1799        }
1800        Operator::LtEq => {
1801            // column <= literal => (min, max) <= literal => min <= literal
1802            Arc::new(phys_expr::BinaryExpr::new(
1803                expr_builder.min_column_expr()?,
1804                Operator::LtEq,
1805                Arc::clone(expr_builder.scalar_expr()),
1806            ))
1807        }
1808        // other expressions are not supported
1809        _ => {
1810            return plan_err!(
1811                "expressions other than (neq, eq, gt, gteq, lt, lteq) are not supported"
1812            );
1813        }
1814    };
1815    let statistics_expr = wrap_null_count_check_expr(statistics_expr, expr_builder)?;
1816    Ok(statistics_expr)
1817}
1818
1819fn binary_expr(
1820    left: Arc<dyn PhysicalExpr>,
1821    op: Operator,
1822    right: Arc<dyn PhysicalExpr>,
1823) -> Arc<dyn PhysicalExpr> {
1824    Arc::new(phys_expr::BinaryExpr::new(left, op, right))
1825}
1826
1827fn and_expr(
1828    left: Arc<dyn PhysicalExpr>,
1829    right: Arc<dyn PhysicalExpr>,
1830) -> Arc<dyn PhysicalExpr> {
1831    binary_expr(left, Operator::And, right)
1832}
1833
1834fn or_expr(
1835    left: Arc<dyn PhysicalExpr>,
1836    right: Arc<dyn PhysicalExpr>,
1837) -> Arc<dyn PhysicalExpr> {
1838    binary_expr(left, Operator::Or, right)
1839}
1840
1841fn build_eq_statistics_expr(
1842    expr_builder: &mut PruningExpressionBuilder,
1843) -> Result<Arc<dyn PhysicalExpr>> {
1844    let min_column_expr = expr_builder.min_column_expr()?;
1845    let max_column_expr = expr_builder.max_column_expr()?;
1846    Ok(and_expr(
1847        binary_expr(
1848            min_column_expr,
1849            Operator::LtEq,
1850            Arc::clone(expr_builder.scalar_expr()),
1851        ),
1852        binary_expr(
1853            Arc::clone(expr_builder.scalar_expr()),
1854            Operator::LtEq,
1855            max_column_expr,
1856        ),
1857    ))
1858}
1859
1860fn build_ne_statistics_expr(
1861    expr_builder: &mut PruningExpressionBuilder,
1862) -> Result<Arc<dyn PhysicalExpr>> {
1863    let min_column_expr = expr_builder.min_column_expr()?;
1864    let max_column_expr = expr_builder.max_column_expr()?;
1865    Ok(or_expr(
1866        binary_expr(
1867            min_column_expr,
1868            Operator::NotEq,
1869            Arc::clone(expr_builder.scalar_expr()),
1870        ),
1871        binary_expr(
1872            Arc::clone(expr_builder.scalar_expr()),
1873            Operator::NotEq,
1874            max_column_expr,
1875        ),
1876    ))
1877}
1878
1879fn column_has_nulls_expr(
1880    expr_builder: &mut PruningExpressionBuilder,
1881) -> Result<Arc<dyn PhysicalExpr>> {
1882    Ok(binary_expr(
1883        expr_builder.null_count_column_expr()?,
1884        Operator::Gt,
1885        Arc::new(phys_expr::Literal::new(ScalarValue::UInt64(Some(0)))),
1886    ))
1887}
1888
1889fn column_has_non_nulls_expr(
1890    expr_builder: &mut PruningExpressionBuilder,
1891) -> Result<Arc<dyn PhysicalExpr>> {
1892    Ok(binary_expr(
1893        expr_builder.null_count_column_expr()?,
1894        Operator::NotEq,
1895        expr_builder.row_count_column_expr()?,
1896    ))
1897}
1898
1899fn build_is_distinct_from(
1900    expr_builder: &mut PruningExpressionBuilder,
1901) -> Result<Arc<dyn PhysicalExpr>> {
1902    let scalar_expr = Arc::clone(expr_builder.scalar_expr());
1903
1904    Ok(or_expr(
1905        and_expr(
1906            Arc::new(phys_expr::IsNullExpr::new(Arc::clone(&scalar_expr))),
1907            column_has_non_nulls_expr(expr_builder)?,
1908        ),
1909        and_expr(
1910            Arc::new(phys_expr::IsNotNullExpr::new(scalar_expr)),
1911            or_expr(
1912                column_has_nulls_expr(expr_builder)?,
1913                build_ne_statistics_expr(expr_builder)?,
1914            ),
1915        ),
1916    ))
1917}
1918
1919fn build_is_not_distinct_from(
1920    expr_builder: &mut PruningExpressionBuilder,
1921) -> Result<Arc<dyn PhysicalExpr>> {
1922    let scalar_expr = Arc::clone(expr_builder.scalar_expr());
1923
1924    Ok(or_expr(
1925        and_expr(
1926            Arc::new(phys_expr::IsNullExpr::new(Arc::clone(&scalar_expr))),
1927            column_has_nulls_expr(expr_builder)?,
1928        ),
1929        and_expr(
1930            Arc::new(phys_expr::IsNotNullExpr::new(scalar_expr)),
1931            and_expr(
1932                column_has_non_nulls_expr(expr_builder)?,
1933                build_eq_statistics_expr(expr_builder)?,
1934            ),
1935        ),
1936    ))
1937}
1938
1939/// returns the string literal of the scalar value if it is a string
1940fn unpack_string(s: &ScalarValue) -> Option<&str> {
1941    s.try_as_str().flatten()
1942}
1943
1944fn extract_string_literal(expr: &Arc<dyn PhysicalExpr>) -> Option<&str> {
1945    if let Some(lit) = expr.downcast_ref::<phys_expr::Literal>() {
1946        let s = unpack_string(lit.value())?;
1947        return Some(s);
1948    }
1949    None
1950}
1951
1952/// Wrap a string in a `Literal` whose `ScalarValue` matches `target_type`
1953fn string_literal_as(value: String, target_type: &DataType) -> Arc<dyn PhysicalExpr> {
1954    let utf8 = ScalarValue::Utf8(Some(value));
1955    let scalar = try_cast_literal_to_type(&utf8, target_type).unwrap_or(utf8);
1956    Arc::new(phys_expr::Literal::new(scalar))
1957}
1958
1959/// Convert `column LIKE literal` where P is a constant prefix of the literal
1960/// to a range check on the column: `P <= column && column < P'`, where P' is the
1961/// lowest string after all P* strings.
1962fn build_like_match(
1963    expr_builder: &mut PruningExpressionBuilder,
1964) -> Option<Arc<dyn PhysicalExpr>> {
1965    // column LIKE literal => (min, max) LIKE literal split at unescaped % => min <= split literal && split literal <= max
1966    // column LIKE 'foo%' => min <= 'foo' && 'foo' <= max
1967    // column LIKE 'foo\_%' => min <= 'foo_' && 'foo_' <= max (the _ is escaped)
1968    // column LIKE 'foo\%%' => min <= 'foo%' && 'foo%' <= max (the % is escaped)
1969    // column LIKE '%foo' => min <= '' && '' <= max => true
1970    // column LIKE '%foo%' => min <= '' && '' <= max => true
1971    // column LIKE 'foo' => min <= 'foo' && 'foo' <= max
1972
1973    // TODO Handle ILIKE perhaps by making the min lowercase and max uppercase
1974    //  this may involve building the physical expressions that call lower() and upper()
1975    let min_column_expr = expr_builder.min_column_expr().ok()?;
1976    let max_column_expr = expr_builder.max_column_expr().ok()?;
1977    let scalar_expr = expr_builder.scalar_expr();
1978    // Synthesized bounds must match the column type (e.g. `Utf8View`).
1979    let target_type = expr_builder.field.data_type();
1980    // check that the scalar is a string literal
1981    let s = extract_string_literal(scalar_expr)?;
1982    // ANSI SQL specifies two wildcards: % and _. % matches zero or more characters, _ matches exactly one character.
1983    let (decoded_prefix, rest) = split_constant_prefix(s);
1984    let has_wildcard = !rest.is_empty();
1985    if has_wildcard && decoded_prefix.is_empty() {
1986        // there's no filtering we could possibly do, return None and have this be handled by the unhandled hook
1987        return None;
1988    }
1989    let (lower_bound, upper_bound) = if has_wildcard {
1990        let incremented_prefix = increment_utf8(&decoded_prefix)?;
1991        let lower_bound_lit = string_literal_as(decoded_prefix, target_type);
1992        let upper_bound_lit = string_literal_as(incremented_prefix, target_type);
1993        (lower_bound_lit, upper_bound_lit)
1994    } else {
1995        // the like expression is a literal and can be converted into a comparison
1996        let bound = string_literal_as(decoded_prefix, target_type);
1997        (Arc::clone(&bound), bound)
1998    };
1999    let lower_bound_expr = Arc::new(phys_expr::BinaryExpr::new(
2000        lower_bound,
2001        Operator::LtEq,
2002        Arc::clone(&max_column_expr),
2003    ));
2004    let upper_bound_expr = Arc::new(phys_expr::BinaryExpr::new(
2005        Arc::clone(&min_column_expr),
2006        Operator::LtEq,
2007        upper_bound,
2008    ));
2009    let combined = Arc::new(phys_expr::BinaryExpr::new(
2010        upper_bound_expr,
2011        Operator::And,
2012        lower_bound_expr,
2013    ));
2014    Some(combined)
2015}
2016
2017// For predicate `col NOT LIKE 'const_prefix%'`, we rewrite it as `(col_min NOT LIKE 'const_prefix%' OR col_max NOT LIKE 'const_prefix%')`.
2018//
2019// The intuition is that if both `col_min` and `col_max` begin with `const_prefix` that means
2020// **all** data in this row group begins with `const_prefix` as well (and therefore the predicate
2021// looking for rows that don't begin with `const_prefix` can never be true)
2022fn build_not_like_match(
2023    expr_builder: &mut PruningExpressionBuilder<'_>,
2024) -> Result<Arc<dyn PhysicalExpr>> {
2025    // col NOT LIKE 'const_prefix%' -> !(col_min LIKE 'const_prefix%' && col_max LIKE 'const_prefix%') -> (col_min NOT LIKE 'const_prefix%' || col_max NOT LIKE 'const_prefix%')
2026
2027    let min_column_expr = expr_builder.min_column_expr()?;
2028    let max_column_expr = expr_builder.max_column_expr()?;
2029
2030    let scalar_expr = expr_builder.scalar_expr();
2031
2032    let pattern = extract_string_literal(scalar_expr).ok_or_else(|| {
2033        plan_datafusion_err!("cannot extract literal from NOT LIKE expression")
2034    })?;
2035
2036    let (const_prefix, remaining) = split_constant_prefix(pattern);
2037    if const_prefix.is_empty() || remaining != "%" {
2038        // we can not handle `%` at the beginning or in the middle of the pattern
2039        // Example: For pattern "foo%bar", the row group might include values like
2040        // ["foobar", "food", "foodbar"], making it unsafe to prune.
2041        // Even if the min/max values in the group (e.g., "foobar" and "foodbar")
2042        // match the pattern, intermediate values like "food" may not
2043        // match the full pattern "foo%bar", making pruning unsafe.
2044        // (truncate foo%bar to foo% have same problem)
2045
2046        // we can not handle pattern containing `_`
2047        // Example: For pattern "foo_", row groups might contain ["fooa", "fooaa", "foob"],
2048        // which means not every row is guaranteed to match the pattern.
2049        return Err(plan_datafusion_err!(
2050            "NOT LIKE expressions only support constant_prefix+wildcard`%`"
2051        ));
2052    }
2053
2054    let min_col_not_like_epxr = Arc::new(phys_expr::LikeExpr::new(
2055        true,
2056        false,
2057        Arc::clone(&min_column_expr),
2058        Arc::clone(scalar_expr),
2059    ));
2060
2061    let max_col_not_like_expr = Arc::new(phys_expr::LikeExpr::new(
2062        true,
2063        false,
2064        Arc::clone(&max_column_expr),
2065        Arc::clone(scalar_expr),
2066    ));
2067
2068    Ok(Arc::new(phys_expr::BinaryExpr::new(
2069        min_col_not_like_epxr,
2070        Operator::Or,
2071        max_col_not_like_expr,
2072    )))
2073}
2074
2075/// Returns unescaped constant prefix of a LIKE pattern (possibly empty) and the remaining pattern (possibly empty)
2076fn split_constant_prefix(pattern: &str) -> (String, &str) {
2077    let mut prefix = String::with_capacity(pattern.len());
2078    let mut iter = pattern.char_indices();
2079    while let Some((idx, c)) = iter.next() {
2080        match c {
2081            '%' | '_' => return (prefix, &pattern[idx..]),
2082            '\\' => match iter.next() {
2083                Some((_, escaped)) => prefix.push(escaped),
2084                None => prefix.push('\\'),
2085            },
2086            _ => prefix.push(c),
2087        }
2088    }
2089    (prefix, "")
2090}
2091
2092/// Increment a UTF8 string by one, returning `None` if it can't be incremented.
2093/// This makes it so that the returned string will always compare greater than the input string
2094/// or any other string with the same prefix.
2095/// This is necessary since the statistics may have been truncated: if we have a min statistic
2096/// of "fo" that may have originally been "foz" or anything else with the prefix "fo".
2097/// E.g. `increment_utf8("foo") >= "foo"` and `increment_utf8("foo") >= "fooz"`
2098/// In this example `increment_utf8("foo") == "fop"
2099fn increment_utf8(data: &str) -> Option<String> {
2100    // Helper function to check if a character is valid to use
2101    fn is_valid_unicode(c: char) -> bool {
2102        let cp = c as u32;
2103
2104        // Filter out non-characters (https://www.unicode.org/versions/corrigendum9.html)
2105        if [0xFFFE, 0xFFFF].contains(&cp) || (0xFDD0..=0xFDEF).contains(&cp) {
2106            return false;
2107        }
2108
2109        // Filter out private use area
2110        if cp >= 0x110000 {
2111            return false;
2112        }
2113
2114        true
2115    }
2116
2117    // Convert string to vector of code points
2118    let mut code_points: Vec<char> = data.chars().collect();
2119
2120    // Work backwards through code points
2121    for idx in (0..code_points.len()).rev() {
2122        let original = code_points[idx] as u32;
2123
2124        // Try incrementing the code point
2125        if let Some(next_char) = char::from_u32(original + 1)
2126            && is_valid_unicode(next_char)
2127        {
2128            code_points[idx] = next_char;
2129            // truncate the string to the current index
2130            code_points.truncate(idx + 1);
2131            return Some(code_points.into_iter().collect());
2132        }
2133    }
2134
2135    None
2136}
2137
2138/// Wrap the statistics expression in a check that skips the expression if the column is all nulls.
2139///
2140/// This is important not only as an optimization but also because statistics may not be
2141/// accurate for columns that are all nulls.
2142/// For example, for an `int` column `x` with all nulls, the min/max/null_count statistics
2143/// might be set to 0 and evaluating `x = 0` would incorrectly include the column.
2144///
2145/// For example:
2146///
2147/// `x_min <= 10 AND 10 <= x_max`
2148///
2149/// will become
2150///
2151/// ```sql
2152/// x_null_count != x_row_count AND (x_min <= 10 AND 10 <= x_max)
2153/// ````
2154///
2155/// If the column is known to be all nulls, then the expression
2156/// `x_null_count = x_row_count` will be true, which will cause the
2157/// boolean expression to return false. Therefore, prune out the container.
2158fn wrap_null_count_check_expr(
2159    statistics_expr: Arc<dyn PhysicalExpr>,
2160    expr_builder: &mut PruningExpressionBuilder,
2161) -> Result<Arc<dyn PhysicalExpr>> {
2162    // (x_null_count != x_row_count) AND (<statistics_expr>)
2163    Ok(and_expr(
2164        column_has_non_nulls_expr(expr_builder)?,
2165        statistics_expr,
2166    ))
2167}
2168
2169#[derive(Debug, Copy, Clone, PartialEq, Eq)]
2170pub(crate) enum StatisticsType {
2171    Min,
2172    Max,
2173    NullCount,
2174    RowCount,
2175}
2176
2177#[cfg(test)]
2178mod tests {
2179    use std::collections::HashMap;
2180    use std::ops::{Not, Rem};
2181
2182    use super::*;
2183    use datafusion_common::test_util::batches_to_string;
2184    use datafusion_expr::{and, col, lit, or};
2185    use datafusion_physical_expr::utils::collect_columns;
2186    use insta::assert_snapshot;
2187
2188    use arrow::array::Decimal128Array;
2189    use arrow::{
2190        array::{BinaryArray, Int32Array, Int64Array, StringArray, UInt64Array},
2191        datatypes::TimeUnit,
2192    };
2193    use datafusion_expr::expr::InList;
2194    use datafusion_expr::{BinaryExpr, Expr, cast, is_null, try_cast};
2195    use datafusion_functions_nested::expr_fn::{array_has, make_array};
2196    use datafusion_physical_expr::expressions::{
2197        self as phys_expr, DynamicFilterPhysicalExpr,
2198    };
2199    use datafusion_physical_expr::planner::logical2physical;
2200    use itertools::Itertools;
2201
2202    #[derive(Debug, Default)]
2203    /// Mock statistic provider for tests
2204    ///
2205    /// Each row represents the statistics for a "container" (which
2206    /// might represent an entire parquet file, or directory of files,
2207    /// or some other collection of data for which we had statistics)
2208    ///
2209    /// Note All `ArrayRefs` must be the same size.
2210    struct ContainerStats {
2211        min: Option<ArrayRef>,
2212        max: Option<ArrayRef>,
2213        /// Optional values
2214        null_counts: Option<ArrayRef>,
2215        row_counts: Option<ArrayRef>,
2216        /// Optional known values (e.g. mimic a bloom filter)
2217        /// (value, contained)
2218        /// If present, all BooleanArrays must be the same size as min/max
2219        contained: Vec<(HashSet<ScalarValue>, BooleanArray)>,
2220    }
2221
2222    impl ContainerStats {
2223        fn new() -> Self {
2224            Default::default()
2225        }
2226        fn new_decimal128(
2227            min: impl IntoIterator<Item = Option<i128>>,
2228            max: impl IntoIterator<Item = Option<i128>>,
2229            precision: u8,
2230            scale: i8,
2231        ) -> Self {
2232            Self::new()
2233                .with_min(Arc::new(
2234                    min.into_iter()
2235                        .collect::<Decimal128Array>()
2236                        .with_precision_and_scale(precision, scale)
2237                        .unwrap(),
2238                ))
2239                .with_max(Arc::new(
2240                    max.into_iter()
2241                        .collect::<Decimal128Array>()
2242                        .with_precision_and_scale(precision, scale)
2243                        .unwrap(),
2244                ))
2245        }
2246
2247        fn new_i64(
2248            min: impl IntoIterator<Item = Option<i64>>,
2249            max: impl IntoIterator<Item = Option<i64>>,
2250        ) -> Self {
2251            Self::new()
2252                .with_min(Arc::new(min.into_iter().collect::<Int64Array>()))
2253                .with_max(Arc::new(max.into_iter().collect::<Int64Array>()))
2254        }
2255
2256        fn new_i32(
2257            min: impl IntoIterator<Item = Option<i32>>,
2258            max: impl IntoIterator<Item = Option<i32>>,
2259        ) -> Self {
2260            Self::new()
2261                .with_min(Arc::new(min.into_iter().collect::<Int32Array>()))
2262                .with_max(Arc::new(max.into_iter().collect::<Int32Array>()))
2263        }
2264
2265        fn new_utf8<'a>(
2266            min: impl IntoIterator<Item = Option<&'a str>>,
2267            max: impl IntoIterator<Item = Option<&'a str>>,
2268        ) -> Self {
2269            Self::new()
2270                .with_min(Arc::new(min.into_iter().collect::<StringArray>()))
2271                .with_max(Arc::new(max.into_iter().collect::<StringArray>()))
2272        }
2273
2274        fn new_bool(
2275            min: impl IntoIterator<Item = Option<bool>>,
2276            max: impl IntoIterator<Item = Option<bool>>,
2277        ) -> Self {
2278            Self::new()
2279                .with_min(Arc::new(min.into_iter().collect::<BooleanArray>()))
2280                .with_max(Arc::new(max.into_iter().collect::<BooleanArray>()))
2281        }
2282
2283        fn min(&self) -> Option<ArrayRef> {
2284            self.min.clone()
2285        }
2286
2287        fn max(&self) -> Option<ArrayRef> {
2288            self.max.clone()
2289        }
2290
2291        fn null_counts(&self) -> Option<ArrayRef> {
2292            self.null_counts.clone()
2293        }
2294
2295        fn row_counts(&self) -> Option<ArrayRef> {
2296            self.row_counts.clone()
2297        }
2298
2299        /// return an iterator over all arrays in this statistics
2300        fn arrays(&self) -> Vec<ArrayRef> {
2301            let contained_arrays = self
2302                .contained
2303                .iter()
2304                .map(|(_values, contained)| Arc::new(contained.clone()) as ArrayRef);
2305
2306            [
2307                self.min.as_ref().cloned(),
2308                self.max.as_ref().cloned(),
2309                self.null_counts.as_ref().cloned(),
2310                self.row_counts.as_ref().cloned(),
2311            ]
2312            .into_iter()
2313            .flatten()
2314            .chain(contained_arrays)
2315            .collect()
2316        }
2317
2318        /// Returns the number of containers represented by this statistics This
2319        /// picks the length of the first array as all arrays must have the same
2320        /// length (which is verified by `assert_invariants`).
2321        fn len(&self) -> usize {
2322            // pick the first non zero length
2323            self.arrays().iter().map(|a| a.len()).next().unwrap_or(0)
2324        }
2325
2326        /// Ensure that the lengths of all arrays are consistent
2327        fn assert_invariants(&self) {
2328            let mut prev_len = None;
2329
2330            for len in self.arrays().iter().map(|a| a.len()) {
2331                // Get a length, if we don't already have one
2332                match prev_len {
2333                    None => {
2334                        prev_len = Some(len);
2335                    }
2336                    Some(prev_len) => {
2337                        assert_eq!(prev_len, len);
2338                    }
2339                }
2340            }
2341        }
2342
2343        /// Add min values
2344        fn with_min(mut self, min: ArrayRef) -> Self {
2345            self.min = Some(min);
2346            self
2347        }
2348
2349        /// Add max values
2350        fn with_max(mut self, max: ArrayRef) -> Self {
2351            self.max = Some(max);
2352            self
2353        }
2354
2355        /// Add null counts. There must be the same number of null counts as
2356        /// there are containers
2357        fn with_null_counts(
2358            mut self,
2359            counts: impl IntoIterator<Item = Option<u64>>,
2360        ) -> Self {
2361            let null_counts: ArrayRef =
2362                Arc::new(counts.into_iter().collect::<UInt64Array>());
2363
2364            self.assert_invariants();
2365            self.null_counts = Some(null_counts);
2366            self
2367        }
2368
2369        /// Add row counts. There must be the same number of row counts as
2370        /// there are containers
2371        fn with_row_counts(
2372            mut self,
2373            counts: impl IntoIterator<Item = Option<u64>>,
2374        ) -> Self {
2375            let row_counts: ArrayRef =
2376                Arc::new(counts.into_iter().collect::<UInt64Array>());
2377
2378            self.assert_invariants();
2379            self.row_counts = Some(row_counts);
2380            self
2381        }
2382
2383        /// Add contained information.
2384        #[allow(clippy::allow_attributes, clippy::mutable_key_type)] // ScalarValue has interior mutability but is intentionally used as hash key
2385        pub fn with_contained(
2386            mut self,
2387            values: impl IntoIterator<Item = ScalarValue>,
2388            contained: impl IntoIterator<Item = Option<bool>>,
2389        ) -> Self {
2390            let contained: BooleanArray = contained.into_iter().collect();
2391            let values: HashSet<_> = values.into_iter().collect();
2392
2393            self.contained.push((values, contained));
2394            self.assert_invariants();
2395            self
2396        }
2397
2398        /// get any contained information for the specified values
2399        #[allow(clippy::allow_attributes, clippy::mutable_key_type)] // ScalarValue has interior mutability but is intentionally used as hash key
2400        fn contained(&self, find_values: &HashSet<ScalarValue>) -> Option<BooleanArray> {
2401            // find the one with the matching values
2402            self.contained
2403                .iter()
2404                .find(|(values, _contained)| values == find_values)
2405                .map(|(_values, contained)| contained.clone())
2406        }
2407    }
2408
2409    #[derive(Debug, Default)]
2410    struct TestStatistics {
2411        // key: column name
2412        stats: HashMap<Column, ContainerStats>,
2413    }
2414
2415    impl TestStatistics {
2416        fn new() -> Self {
2417            Self::default()
2418        }
2419
2420        fn with(
2421            mut self,
2422            name: impl Into<String>,
2423            container_stats: ContainerStats,
2424        ) -> Self {
2425            let col = Column::from_name(name.into());
2426            self.stats.insert(col, container_stats);
2427            self
2428        }
2429
2430        /// Add null counts for the specified column.
2431        /// There must be the same number of null counts as
2432        /// there are containers
2433        fn with_null_counts(
2434            mut self,
2435            name: impl Into<String>,
2436            counts: impl IntoIterator<Item = Option<u64>>,
2437        ) -> Self {
2438            let col = Column::from_name(name.into());
2439
2440            // take stats out and update them
2441            let container_stats = self
2442                .stats
2443                .remove(&col)
2444                .unwrap_or_default()
2445                .with_null_counts(counts);
2446
2447            // put stats back in
2448            self.stats.insert(col, container_stats);
2449            self
2450        }
2451
2452        /// Add row counts for the specified column.
2453        /// There must be the same number of row counts as
2454        /// there are containers
2455        fn with_row_counts(
2456            mut self,
2457            name: impl Into<String>,
2458            counts: impl IntoIterator<Item = Option<u64>>,
2459        ) -> Self {
2460            let col = Column::from_name(name.into());
2461
2462            // take stats out and update them
2463            let container_stats = self
2464                .stats
2465                .remove(&col)
2466                .unwrap_or_default()
2467                .with_row_counts(counts);
2468
2469            // put stats back in
2470            self.stats.insert(col, container_stats);
2471            self
2472        }
2473
2474        /// Add contained information for the specified column.
2475        fn with_contained(
2476            mut self,
2477            name: impl Into<String>,
2478            values: impl IntoIterator<Item = ScalarValue>,
2479            contained: impl IntoIterator<Item = Option<bool>>,
2480        ) -> Self {
2481            let col = Column::from_name(name.into());
2482
2483            // take stats out and update them
2484            let container_stats = self
2485                .stats
2486                .remove(&col)
2487                .unwrap_or_default()
2488                .with_contained(values, contained);
2489
2490            // put stats back in
2491            self.stats.insert(col, container_stats);
2492            self
2493        }
2494    }
2495
2496    impl PruningStatistics for TestStatistics {
2497        fn min_values(&self, column: &Column) -> Option<ArrayRef> {
2498            self.stats
2499                .get(column)
2500                .map(|container_stats| container_stats.min())
2501                .unwrap_or(None)
2502        }
2503
2504        fn max_values(&self, column: &Column) -> Option<ArrayRef> {
2505            self.stats
2506                .get(column)
2507                .map(|container_stats| container_stats.max())
2508                .unwrap_or(None)
2509        }
2510
2511        fn num_containers(&self) -> usize {
2512            self.stats
2513                .values()
2514                .next()
2515                .map(|container_stats| container_stats.len())
2516                .unwrap_or(0)
2517        }
2518
2519        fn null_counts(&self, column: &Column) -> Option<ArrayRef> {
2520            self.stats
2521                .get(column)
2522                .map(|container_stats| container_stats.null_counts())
2523                .unwrap_or(None)
2524        }
2525
2526        fn row_counts(&self) -> Option<ArrayRef> {
2527            self.stats
2528                .values()
2529                .find_map(|container_stats| container_stats.row_counts())
2530        }
2531
2532        fn contained(
2533            &self,
2534            column: &Column,
2535            values: &HashSet<ScalarValue>,
2536        ) -> Option<BooleanArray> {
2537            self.stats
2538                .get(column)
2539                .and_then(|container_stats| container_stats.contained(values))
2540        }
2541    }
2542
2543    /// Returns the specified min/max container values
2544    struct OneContainerStats {
2545        min_values: Option<ArrayRef>,
2546        max_values: Option<ArrayRef>,
2547        num_containers: usize,
2548    }
2549
2550    impl PruningStatistics for OneContainerStats {
2551        fn min_values(&self, _column: &Column) -> Option<ArrayRef> {
2552            self.min_values.clone()
2553        }
2554
2555        fn max_values(&self, _column: &Column) -> Option<ArrayRef> {
2556            self.max_values.clone()
2557        }
2558
2559        fn num_containers(&self) -> usize {
2560            self.num_containers
2561        }
2562
2563        fn null_counts(&self, _column: &Column) -> Option<ArrayRef> {
2564            None
2565        }
2566
2567        fn row_counts(&self) -> Option<ArrayRef> {
2568            None
2569        }
2570
2571        fn contained(
2572            &self,
2573            _column: &Column,
2574            _values: &HashSet<ScalarValue>,
2575        ) -> Option<BooleanArray> {
2576            None
2577        }
2578    }
2579
2580    /// Row count should only be referenced once in the pruning expression, even if we need the row count
2581    /// for multiple columns.
2582    #[test]
2583    fn test_unique_row_count_field_and_column() {
2584        // c1 = 100 AND c2 = 200
2585        let schema: SchemaRef = Arc::new(Schema::new(vec![
2586            Field::new("c1", DataType::Int32, true),
2587            Field::new("c2", DataType::Int32, true),
2588        ]));
2589        let expr = col("c1").eq(lit(100)).and(col("c2").eq(lit(200)));
2590        let expr = logical2physical(&expr, &schema);
2591        let p = PruningPredicateBuilder::new()
2592            .with_file_schema(Arc::clone(&schema))
2593            .try_build(expr)
2594            .unwrap();
2595        // note pruning expression refers to row_count twice
2596        assert_eq!(
2597            "c1_null_count@2 != row_count@3 AND c1_min@0 <= 100 AND 100 <= c1_max@1 AND c2_null_count@6 != row_count@3 AND c2_min@4 <= 200 AND 200 <= c2_max@5",
2598            p.predicate_expr.to_string()
2599        );
2600
2601        // Fields in required schema should be unique, otherwise when creating batches
2602        // it will fail because of duplicate field names
2603        let mut fields = HashSet::new();
2604        for (_col, _ty, field) in p.required_columns().iter() {
2605            let was_new = fields.insert(field);
2606            if !was_new {
2607                panic!(
2608                    "Duplicate field in required schema: {field:?}. Previous fields:\n{fields:#?}"
2609                );
2610            }
2611        }
2612    }
2613
2614    #[test]
2615    fn prune_all_rows_null_counts() {
2616        // if null_count = row_count then we should prune the container for i = 0
2617        // regardless of the statistics
2618        let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
2619        let statistics = TestStatistics::new().with(
2620            "i",
2621            ContainerStats::new_i32(
2622                vec![Some(0)], // min
2623                vec![Some(0)], // max
2624            )
2625            .with_null_counts(vec![Some(1)])
2626            .with_row_counts(vec![Some(1)]),
2627        );
2628        let expected_ret = &[false];
2629        prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
2630
2631        // this should be true even if the container stats are missing
2632        let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
2633        let container_stats = ContainerStats {
2634            min: Some(Arc::new(Int32Array::from(vec![None]))),
2635            max: Some(Arc::new(Int32Array::from(vec![None]))),
2636            null_counts: Some(Arc::new(UInt64Array::from(vec![Some(1)]))),
2637            row_counts: Some(Arc::new(UInt64Array::from(vec![Some(1)]))),
2638            ..ContainerStats::default()
2639        };
2640        let statistics = TestStatistics::new().with("i", container_stats);
2641        let expected_ret = &[false];
2642        prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
2643
2644        // If the null counts themselves are missing we should be able to fall back to the stats
2645        let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
2646        let container_stats = ContainerStats {
2647            min: Some(Arc::new(Int32Array::from(vec![Some(0)]))),
2648            max: Some(Arc::new(Int32Array::from(vec![Some(0)]))),
2649            null_counts: Some(Arc::new(UInt64Array::from(vec![None]))),
2650            row_counts: Some(Arc::new(UInt64Array::from(vec![Some(1)]))),
2651            ..ContainerStats::default()
2652        };
2653        let statistics = TestStatistics::new().with("i", container_stats);
2654        let expected_ret = &[true];
2655        prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
2656        let expected_ret = &[false];
2657        prune_with_expr(col("i").gt(lit(0)), &schema, &statistics, expected_ret);
2658
2659        // Same for the row counts
2660        let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
2661        let container_stats = ContainerStats {
2662            min: Some(Arc::new(Int32Array::from(vec![Some(0)]))),
2663            max: Some(Arc::new(Int32Array::from(vec![Some(0)]))),
2664            null_counts: Some(Arc::new(UInt64Array::from(vec![Some(1)]))),
2665            row_counts: Some(Arc::new(UInt64Array::from(vec![None]))),
2666            ..ContainerStats::default()
2667        };
2668        let statistics = TestStatistics::new().with("i", container_stats);
2669        let expected_ret = &[true];
2670        prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
2671        let expected_ret = &[false];
2672        prune_with_expr(col("i").gt(lit(0)), &schema, &statistics, expected_ret);
2673    }
2674
2675    #[test]
2676    fn prune_missing_statistics() {
2677        // If the min or max stats are missing we should not prune
2678        // (unless we know all rows are null, see `prune_all_rows_null_counts`)
2679        let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
2680        let container_stats = ContainerStats {
2681            min: Some(Arc::new(Int32Array::from(vec![None, Some(0)]))),
2682            max: Some(Arc::new(Int32Array::from(vec![Some(0), None]))),
2683            null_counts: Some(Arc::new(UInt64Array::from(vec![Some(0), Some(0)]))),
2684            row_counts: Some(Arc::new(UInt64Array::from(vec![Some(1), Some(1)]))),
2685            ..ContainerStats::default()
2686        };
2687        let statistics = TestStatistics::new().with("i", container_stats);
2688        let expected_ret = &[true, true];
2689        prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
2690        let expected_ret = &[false, true];
2691        prune_with_expr(col("i").gt(lit(0)), &schema, &statistics, expected_ret);
2692        let expected_ret = &[true, false];
2693        prune_with_expr(col("i").lt(lit(0)), &schema, &statistics, expected_ret);
2694    }
2695
2696    #[test]
2697    fn prune_null_stats() {
2698        // if null_count = row_count then we should prune the container for i = 0
2699        // regardless of the statistics
2700        let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
2701
2702        let statistics = TestStatistics::new().with(
2703            "i",
2704            ContainerStats::new_i32(
2705                vec![Some(0)], // min
2706                vec![Some(0)], // max
2707            )
2708            .with_null_counts(vec![Some(1)])
2709            .with_row_counts(vec![Some(1)]),
2710        );
2711
2712        let expected_ret = &[false];
2713
2714        // i = 0
2715        prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
2716    }
2717
2718    #[test]
2719    fn test_build_statistics_record_batch() {
2720        // Request a record batch with of s1_min, s2_max, s3_max, s3_min
2721        let required_columns = RequiredColumns::from(vec![
2722            // min of original column s1, named s1_min
2723            (
2724                phys_expr::Column::new("s1", 1),
2725                StatisticsType::Min,
2726                Field::new("s1_min", DataType::Int32, true),
2727            ),
2728            // max of original column s2, named s2_max
2729            (
2730                phys_expr::Column::new("s2", 2),
2731                StatisticsType::Max,
2732                Field::new("s2_max", DataType::Int32, true),
2733            ),
2734            // max of original column s3, named s3_max
2735            (
2736                phys_expr::Column::new("s3", 3),
2737                StatisticsType::Max,
2738                Field::new("s3_max", DataType::Utf8, true),
2739            ),
2740            // min of original column s3, named s3_min
2741            (
2742                phys_expr::Column::new("s3", 3),
2743                StatisticsType::Min,
2744                Field::new("s3_min", DataType::Utf8, true),
2745            ),
2746        ]);
2747
2748        let statistics = TestStatistics::new()
2749            .with(
2750                "s1",
2751                ContainerStats::new_i32(
2752                    vec![None, None, Some(9), None],  // min
2753                    vec![Some(10), None, None, None], // max
2754                ),
2755            )
2756            .with(
2757                "s2",
2758                ContainerStats::new_i32(
2759                    vec![Some(2), None, None, None],  // min
2760                    vec![Some(20), None, None, None], // max
2761                ),
2762            )
2763            .with(
2764                "s3",
2765                ContainerStats::new_utf8(
2766                    vec![Some("a"), None, None, None],      // min
2767                    vec![Some("q"), None, Some("r"), None], // max
2768                ),
2769            );
2770
2771        let batch =
2772            build_statistics_record_batch(&statistics, &required_columns).unwrap();
2773        assert_snapshot!(batches_to_string(&[batch]), @r"
2774        +--------+--------+--------+--------+
2775        | s1_min | s2_max | s3_max | s3_min |
2776        +--------+--------+--------+--------+
2777        |        | 20     | q      | a      |
2778        |        |        |        |        |
2779        | 9      |        | r      |        |
2780        |        |        |        |        |
2781        +--------+--------+--------+--------+
2782        ");
2783    }
2784
2785    #[test]
2786    fn test_build_statistics_casting() {
2787        // Test requesting a Timestamp column, but getting statistics as Int64
2788        // which is what Parquet does
2789
2790        // Request a record batch with of s1_min as a timestamp
2791        let required_columns = RequiredColumns::from(vec![(
2792            phys_expr::Column::new("s3", 3),
2793            StatisticsType::Min,
2794            Field::new(
2795                "s1_min",
2796                DataType::Timestamp(TimeUnit::Nanosecond, None),
2797                true,
2798            ),
2799        )]);
2800
2801        // Note the statistics pass back i64 (not timestamp)
2802        let statistics = OneContainerStats {
2803            min_values: Some(Arc::new(Int64Array::from(vec![Some(10)]))),
2804            max_values: Some(Arc::new(Int64Array::from(vec![Some(20)]))),
2805            num_containers: 1,
2806        };
2807
2808        let batch =
2809            build_statistics_record_batch(&statistics, &required_columns).unwrap();
2810
2811        assert_snapshot!(batches_to_string(&[batch]), @r"
2812        +-------------------------------+
2813        | s1_min                        |
2814        +-------------------------------+
2815        | 1970-01-01T00:00:00.000000010 |
2816        +-------------------------------+
2817        ");
2818    }
2819
2820    #[test]
2821    fn test_build_statistics_no_required_stats() {
2822        let required_columns = RequiredColumns::new();
2823
2824        let statistics = OneContainerStats {
2825            min_values: Some(Arc::new(Int64Array::from(vec![Some(10)]))),
2826            max_values: Some(Arc::new(Int64Array::from(vec![Some(20)]))),
2827            num_containers: 1,
2828        };
2829
2830        let batch =
2831            build_statistics_record_batch(&statistics, &required_columns).unwrap();
2832        assert_eq!(batch.num_rows(), 1); // had 1 container
2833    }
2834
2835    #[test]
2836    fn test_build_statistics_inconsistent_types() {
2837        // Test requesting a Utf8 column when the stats return some other type
2838
2839        // Request a record batch with of s1_min as a timestamp
2840        let required_columns = RequiredColumns::from(vec![(
2841            phys_expr::Column::new("s3", 3),
2842            StatisticsType::Min,
2843            Field::new("s1_min", DataType::Utf8, true),
2844        )]);
2845
2846        // Note the statistics return an invalid UTF-8 sequence which will be converted to null
2847        let statistics = OneContainerStats {
2848            min_values: Some(Arc::new(BinaryArray::from(vec![&[255u8] as &[u8]]))),
2849            max_values: None,
2850            num_containers: 1,
2851        };
2852
2853        let batch =
2854            build_statistics_record_batch(&statistics, &required_columns).unwrap();
2855        assert_snapshot!(batches_to_string(&[batch]), @r"
2856        +--------+
2857        | s1_min |
2858        +--------+
2859        |        |
2860        +--------+
2861        ");
2862    }
2863
2864    #[test]
2865    fn test_build_statistics_inconsistent_length() {
2866        // return an inconsistent length to the actual statistics arrays
2867        let required_columns = RequiredColumns::from(vec![(
2868            phys_expr::Column::new("s1", 3),
2869            StatisticsType::Min,
2870            Field::new("s1_min", DataType::Int64, true),
2871        )]);
2872
2873        // Note the statistics pass back i64 (not timestamp)
2874        let statistics = OneContainerStats {
2875            min_values: Some(Arc::new(Int64Array::from(vec![Some(10)]))),
2876            max_values: Some(Arc::new(Int64Array::from(vec![Some(20)]))),
2877            num_containers: 3,
2878        };
2879
2880        let result =
2881            build_statistics_record_batch(&statistics, &required_columns).unwrap_err();
2882        assert!(
2883            result
2884                .to_string()
2885                .contains("mismatched statistics length. Expected 3, got 1"),
2886            "{}",
2887            result
2888        );
2889    }
2890
2891    #[test]
2892    fn row_group_predicate_eq() -> Result<()> {
2893        let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
2894        let expected_expr =
2895            "c1_null_count@2 != row_count@3 AND c1_min@0 <= 1 AND 1 <= c1_max@1";
2896
2897        // test column on the left
2898        let expr = col("c1").eq(lit(1));
2899        let predicate_expr =
2900            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2901        assert_eq!(predicate_expr.to_string(), expected_expr);
2902
2903        // test column on the right
2904        let expr = lit(1).eq(col("c1"));
2905        let predicate_expr =
2906            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2907        assert_eq!(predicate_expr.to_string(), expected_expr);
2908
2909        Ok(())
2910    }
2911
2912    #[test]
2913    fn row_group_predicate_not_eq() -> Result<()> {
2914        let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
2915        let expected_expr =
2916            "c1_null_count@2 != row_count@3 AND (c1_min@0 != 1 OR 1 != c1_max@1)";
2917
2918        // test column on the left
2919        let expr = col("c1").not_eq(lit(1));
2920        let predicate_expr =
2921            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2922        assert_eq!(predicate_expr.to_string(), expected_expr);
2923
2924        // test column on the right
2925        let expr = lit(1).not_eq(col("c1"));
2926        let predicate_expr =
2927            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2928        assert_eq!(predicate_expr.to_string(), expected_expr);
2929
2930        Ok(())
2931    }
2932
2933    #[test]
2934    fn row_group_predicate_gt() -> Result<()> {
2935        let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
2936        let expected_expr = "c1_null_count@1 != row_count@2 AND c1_max@0 > 1";
2937
2938        // test column on the left
2939        let expr = col("c1").gt(lit(1));
2940        let predicate_expr =
2941            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2942        assert_eq!(predicate_expr.to_string(), expected_expr);
2943
2944        // test column on the right
2945        let expr = lit(1).lt(col("c1"));
2946        let predicate_expr =
2947            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2948        assert_eq!(predicate_expr.to_string(), expected_expr);
2949
2950        Ok(())
2951    }
2952
2953    #[test]
2954    fn row_group_predicate_gt_eq() -> Result<()> {
2955        let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
2956        let expected_expr = "c1_null_count@1 != row_count@2 AND c1_max@0 >= 1";
2957
2958        // test column on the left
2959        let expr = col("c1").gt_eq(lit(1));
2960        let predicate_expr =
2961            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2962        assert_eq!(predicate_expr.to_string(), expected_expr);
2963        // test column on the right
2964        let expr = lit(1).lt_eq(col("c1"));
2965        let predicate_expr =
2966            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2967        assert_eq!(predicate_expr.to_string(), expected_expr);
2968
2969        Ok(())
2970    }
2971
2972    #[test]
2973    fn row_group_predicate_lt() -> Result<()> {
2974        let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
2975        let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 < 1";
2976
2977        // test column on the left
2978        let expr = col("c1").lt(lit(1));
2979        let predicate_expr =
2980            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2981        assert_eq!(predicate_expr.to_string(), expected_expr);
2982
2983        // test column on the right
2984        let expr = lit(1).gt(col("c1"));
2985        let predicate_expr =
2986            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2987        assert_eq!(predicate_expr.to_string(), expected_expr);
2988
2989        Ok(())
2990    }
2991
2992    #[test]
2993    fn row_group_predicate_lt_eq() -> Result<()> {
2994        let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
2995        let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 <= 1";
2996
2997        // test column on the left
2998        let expr = col("c1").lt_eq(lit(1));
2999        let predicate_expr =
3000            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3001        assert_eq!(predicate_expr.to_string(), expected_expr);
3002        // test column on the right
3003        let expr = lit(1).gt_eq(col("c1"));
3004        let predicate_expr =
3005            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3006        assert_eq!(predicate_expr.to_string(), expected_expr);
3007
3008        Ok(())
3009    }
3010
3011    #[test]
3012    fn row_group_predicate_and() -> Result<()> {
3013        let schema = Schema::new(vec![
3014            Field::new("c1", DataType::Int32, false),
3015            Field::new("c2", DataType::Int32, false),
3016            Field::new("c3", DataType::Int32, false),
3017        ]);
3018        // test AND operator joining supported c1 < 1 expression and unsupported c2 > c3 expression
3019        let expr = col("c1").lt(lit(1)).and(col("c2").lt(col("c3")));
3020        let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 < 1";
3021        let predicate_expr =
3022            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3023        assert_eq!(predicate_expr.to_string(), expected_expr);
3024
3025        Ok(())
3026    }
3027
3028    #[test]
3029    fn row_group_predicate_or() -> Result<()> {
3030        let schema = Schema::new(vec![
3031            Field::new("c1", DataType::Int32, false),
3032            Field::new("c2", DataType::Int32, false),
3033        ]);
3034        // test OR operator joining supported c1 < 1 expression and unsupported c2 % 2 = 0 expression
3035        let expr = col("c1").lt(lit(1)).or(col("c2").rem(lit(2)).eq(lit(0)));
3036        let expected_expr = "true";
3037        let predicate_expr =
3038            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3039        assert_eq!(predicate_expr.to_string(), expected_expr);
3040
3041        Ok(())
3042    }
3043
3044    #[test]
3045    fn row_group_predicate_not() -> Result<()> {
3046        let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
3047        let expected_expr = "true";
3048
3049        let expr = col("c1").not();
3050        let predicate_expr =
3051            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3052        assert_eq!(predicate_expr.to_string(), expected_expr);
3053
3054        Ok(())
3055    }
3056
3057    #[test]
3058    fn row_group_predicate_not_bool() -> Result<()> {
3059        let schema = Schema::new(vec![Field::new("c1", DataType::Boolean, false)]);
3060        let expected_expr = "NOT c1_min@0 AND c1_max@1";
3061
3062        let expr = col("c1").not();
3063        let predicate_expr =
3064            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3065        assert_eq!(predicate_expr.to_string(), expected_expr);
3066
3067        Ok(())
3068    }
3069
3070    #[test]
3071    fn row_group_predicate_bool() -> Result<()> {
3072        let schema = Schema::new(vec![Field::new("c1", DataType::Boolean, false)]);
3073        let expected_expr = "c1_min@0 OR c1_max@1";
3074
3075        let expr = col("c1");
3076        let predicate_expr =
3077            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3078        assert_eq!(predicate_expr.to_string(), expected_expr);
3079
3080        Ok(())
3081    }
3082
3083    /// Test that non-boolean literal expressions don't prune any containers and error gracefully by not pruning anything instead of e.g. panicking
3084    #[test]
3085    fn row_group_predicate_non_boolean() {
3086        let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)]));
3087        let statistics = TestStatistics::new()
3088            .with("c1", ContainerStats::new_i32(vec![Some(0)], vec![Some(10)]));
3089        let expected_ret = &[true];
3090        prune_with_expr(lit(1), &schema, &statistics, expected_ret);
3091    }
3092
3093    // Test that literal-to-literal comparisons are correctly evaluated.
3094    // When both sides are constants, the expression should be evaluated directly
3095    // and if it's false, all containers should be pruned.
3096    #[test]
3097    fn row_group_predicate_literal_false() {
3098        // lit(1) = lit(2) is always false, so all containers should be pruned
3099        let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)]));
3100        let statistics = TestStatistics::new()
3101            .with("c1", ContainerStats::new_i32(vec![Some(0)], vec![Some(10)]));
3102        let expected_ret = &[false];
3103        prune_with_simplified_expr(lit(1).eq(lit(2)), &schema, &statistics, expected_ret);
3104    }
3105
3106    /// Test nested/complex literal expression trees.
3107    /// This is an integration test that PhysicalExprSimplifier + PruningPredicate work together as expected.
3108    #[test]
3109    fn row_group_predicate_literal_true() {
3110        // lit(1) = lit(1) is always true, so no containers should be pruned
3111        let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)]));
3112        let statistics = TestStatistics::new()
3113            .with("c1", ContainerStats::new_i32(vec![Some(0)], vec![Some(10)]));
3114        let expected_ret = &[true];
3115        prune_with_simplified_expr(lit(1).eq(lit(1)), &schema, &statistics, expected_ret);
3116    }
3117
3118    /// Test nested/complex literal expression trees.
3119    /// This is an integration test that PhysicalExprSimplifier + PruningPredicate work together as expected.
3120    #[test]
3121    fn row_group_predicate_literal_null() {
3122        // lit(1) = null is always null, so no containers should be pruned
3123        let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)]));
3124        let statistics = TestStatistics::new()
3125            .with("c1", ContainerStats::new_i32(vec![Some(0)], vec![Some(10)]));
3126        let expected_ret = &[true];
3127        prune_with_simplified_expr(
3128            lit(1).eq(lit(ScalarValue::Null)),
3129            &schema,
3130            &statistics,
3131            expected_ret,
3132        );
3133    }
3134
3135    /// Test nested/complex literal expression trees.
3136    /// This is an integration test that PhysicalExprSimplifier + PruningPredicate work together as expected.
3137    #[test]
3138    fn row_group_predicate_complex_literals() {
3139        let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)]));
3140        let statistics = TestStatistics::new()
3141            .with("c1", ContainerStats::new_i32(vec![Some(0)], vec![Some(10)]));
3142
3143        // (1 + 2) > 0 is always true
3144        prune_with_simplified_expr(
3145            (lit(1) + lit(2)).gt(lit(0)),
3146            &schema,
3147            &statistics,
3148            &[true],
3149        );
3150
3151        // (1 + 2) < 0 is always false
3152        prune_with_simplified_expr(
3153            (lit(1) + lit(2)).lt(lit(0)),
3154            &schema,
3155            &statistics,
3156            &[false],
3157        );
3158
3159        // Nested AND of literals: true AND false = false
3160        prune_with_simplified_expr(
3161            lit(true).and(lit(false)),
3162            &schema,
3163            &statistics,
3164            &[false],
3165        );
3166
3167        // Nested OR of literals: true OR false = true
3168        prune_with_simplified_expr(
3169            lit(true).or(lit(false)),
3170            &schema,
3171            &statistics,
3172            &[true],
3173        );
3174
3175        // Complex nested: (1 < 2) AND (3 > 1) = true AND true = true
3176        prune_with_simplified_expr(
3177            lit(1).lt(lit(2)).and(lit(3).gt(lit(1))),
3178            &schema,
3179            &statistics,
3180            &[true],
3181        );
3182
3183        // Complex nested: (1 > 2) OR (3 < 1) = false OR false = false
3184        prune_with_simplified_expr(
3185            lit(1).gt(lit(2)).or(lit(3).lt(lit(1))),
3186            &schema,
3187            &statistics,
3188            &[false],
3189        );
3190    }
3191
3192    /// Integration test demonstrating that a dynamic filter with replaced children as literals will be snapshotted, simplified and then pruned correctly.
3193    #[test]
3194    fn row_group_predicate_dynamic_filter_with_literals() {
3195        let schema = Arc::new(Schema::new(vec![
3196            Field::new("c1", DataType::Int32, true),
3197            Field::new("part", DataType::Utf8, true),
3198        ]));
3199        let statistics = TestStatistics::new()
3200            // Note that we have no stats, pruning can only happen via partition value pruning from the dynamic filter
3201            .with_row_counts("c1", vec![Some(10)]);
3202        let dynamic_filter_expr = col("c1").gt(lit(5)).and(col("part").eq(lit("B")));
3203        let phys_expr = logical2physical(&dynamic_filter_expr, &schema);
3204        let children = collect_columns(&phys_expr)
3205            .iter()
3206            .map(|c| Arc::new(c.clone()) as Arc<dyn PhysicalExpr>)
3207            .collect_vec();
3208        let dynamic_phys_expr =
3209            Arc::new(DynamicFilterPhysicalExpr::new(children, phys_expr))
3210                as Arc<dyn PhysicalExpr>;
3211        // Simulate the partition value substitution that would happen in ParquetOpener
3212        let remapped_expr = dynamic_phys_expr
3213            .children()
3214            .into_iter()
3215            .map(|child_expr| {
3216                let Some(col_expr) = child_expr.downcast_ref::<phys_expr::Column>()
3217                else {
3218                    return Arc::clone(child_expr);
3219                };
3220                if col_expr.name() == "part" {
3221                    // simulate dynamic filter replacement with literal "A"
3222                    Arc::new(phys_expr::Literal::new(ScalarValue::Utf8(Some(
3223                        "A".to_string(),
3224                    )))) as Arc<dyn PhysicalExpr>
3225                } else {
3226                    Arc::clone(child_expr)
3227                }
3228            })
3229            .collect_vec();
3230        let dynamic_filter_expr =
3231            dynamic_phys_expr.with_new_children(remapped_expr).unwrap();
3232        // After substitution the expression is c1 > 5 AND part = "B" which should prune the file since the partition value is "A"
3233        let expected = &[false];
3234        let p = PruningPredicateBuilder::new()
3235            .with_file_schema(Arc::clone(&schema))
3236            .try_build(dynamic_filter_expr)
3237            .unwrap();
3238        let result = p.prune(&statistics).unwrap();
3239        assert_eq!(result, expected);
3240    }
3241
3242    #[test]
3243    fn row_group_predicate_lt_bool() -> Result<()> {
3244        let schema = Schema::new(vec![Field::new("c1", DataType::Boolean, false)]);
3245        let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 < true";
3246
3247        // DF doesn't support arithmetic on boolean columns so
3248        // this predicate will error when evaluated
3249        let expr = col("c1").lt(lit(true));
3250        let predicate_expr =
3251            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3252        assert_eq!(predicate_expr.to_string(), expected_expr);
3253
3254        Ok(())
3255    }
3256
3257    #[test]
3258    fn row_group_predicate_required_columns() -> Result<()> {
3259        let schema = Schema::new(vec![
3260            Field::new("c1", DataType::Int32, false),
3261            Field::new("c2", DataType::Int32, false),
3262        ]);
3263        let mut required_columns = RequiredColumns::new();
3264        // c1 < 1 and (c2 = 2 or c2 = 3)
3265        let expr = col("c1")
3266            .lt(lit(1))
3267            .and(col("c2").eq(lit(2)).or(col("c2").eq(lit(3))));
3268        let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 < 1 AND (c2_null_count@5 != row_count@2 AND c2_min@3 <= 2 AND 2 <= c2_max@4 OR c2_null_count@5 != row_count@2 AND c2_min@3 <= 3 AND 3 <= c2_max@4)";
3269        let predicate_expr =
3270            test_build_predicate_expression(&expr, &schema, &mut required_columns);
3271        assert_eq!(predicate_expr.to_string(), expected_expr);
3272        println!("required_columns: {required_columns:#?}"); // for debugging assertions below
3273        // c1 < 1 should add c1_min
3274        let c1_min_field = Field::new("c1_min", DataType::Int32, false);
3275        assert_eq!(
3276            required_columns.columns[0],
3277            (
3278                phys_expr::Column::new("c1", 0),
3279                StatisticsType::Min,
3280                c1_min_field.with_nullable(true) // could be nullable if stats are not present
3281            )
3282        );
3283        // c1 < 1 should add c1_null_count
3284        let c1_null_count_field = Field::new("c1_null_count", DataType::UInt64, false);
3285        assert_eq!(
3286            required_columns.columns[1],
3287            (
3288                phys_expr::Column::new("c1", 0),
3289                StatisticsType::NullCount,
3290                c1_null_count_field.with_nullable(true) // could be nullable if stats are not present
3291            )
3292        );
3293        // c1 < 1 should add row_count
3294        let row_count_field = Field::new("row_count", DataType::UInt64, false);
3295        assert_eq!(
3296            required_columns.columns[2],
3297            (
3298                phys_expr::Column::new("c1", 0),
3299                StatisticsType::RowCount,
3300                row_count_field.with_nullable(true) // could be nullable if stats are not present
3301            )
3302        );
3303        // c2 = 2 should add c2_min and c2_max
3304        let c2_min_field = Field::new("c2_min", DataType::Int32, false);
3305        assert_eq!(
3306            required_columns.columns[3],
3307            (
3308                phys_expr::Column::new("c2", 1),
3309                StatisticsType::Min,
3310                c2_min_field.with_nullable(true) // could be nullable if stats are not present
3311            )
3312        );
3313        let c2_max_field = Field::new("c2_max", DataType::Int32, false);
3314        assert_eq!(
3315            required_columns.columns[4],
3316            (
3317                phys_expr::Column::new("c2", 1),
3318                StatisticsType::Max,
3319                c2_max_field.with_nullable(true) // could be nullable if stats are not present
3320            )
3321        );
3322        // c2 = 2 should add c2_null_count
3323        let c2_null_count_field = Field::new("c2_null_count", DataType::UInt64, false);
3324        assert_eq!(
3325            required_columns.columns[5],
3326            (
3327                phys_expr::Column::new("c2", 1),
3328                StatisticsType::NullCount,
3329                c2_null_count_field.with_nullable(true) // could be nullable if stats are not present
3330            )
3331        );
3332        // c2 = 1 should add row_count
3333        let row_count_field = Field::new("row_count", DataType::UInt64, false);
3334        assert_eq!(
3335            required_columns.columns[2],
3336            (
3337                phys_expr::Column::new("c1", 0),
3338                StatisticsType::RowCount,
3339                row_count_field.with_nullable(true) // could be nullable if stats are not present
3340            )
3341        );
3342        // c2 = 3 shouldn't add any new statistics fields
3343        assert_eq!(required_columns.columns.len(), 6);
3344
3345        Ok(())
3346    }
3347
3348    #[test]
3349    fn row_group_predicate_in_list() -> Result<()> {
3350        let schema = Schema::new(vec![
3351            Field::new("c1", DataType::Int32, false),
3352            Field::new("c2", DataType::Int32, false),
3353        ]);
3354        // test c1 in(1, 2, 3)
3355        let expr = Expr::InList(InList::new(
3356            Box::new(col("c1")),
3357            vec![lit(1), lit(2), lit(3)],
3358            false,
3359        ));
3360        let expected_expr = "c1_null_count@2 != row_count@3 AND c1_min@0 <= 1 AND 1 <= c1_max@1 OR c1_null_count@2 != row_count@3 AND c1_min@0 <= 2 AND 2 <= c1_max@1 OR c1_null_count@2 != row_count@3 AND c1_min@0 <= 3 AND 3 <= c1_max@1";
3361        let predicate_expr =
3362            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3363        assert_eq!(predicate_expr.to_string(), expected_expr);
3364
3365        Ok(())
3366    }
3367
3368    #[test]
3369    fn row_group_predicate_in_list_empty() -> Result<()> {
3370        let schema = Schema::new(vec![
3371            Field::new("c1", DataType::Int32, false),
3372            Field::new("c2", DataType::Int32, false),
3373        ]);
3374        // test c1 in()
3375        let expr = Expr::InList(InList::new(Box::new(col("c1")), vec![], false));
3376        let expected_expr = "true";
3377        let predicate_expr =
3378            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3379        assert_eq!(predicate_expr.to_string(), expected_expr);
3380
3381        Ok(())
3382    }
3383
3384    #[test]
3385    fn row_group_predicate_in_list_negated() -> Result<()> {
3386        let schema = Schema::new(vec![
3387            Field::new("c1", DataType::Int32, false),
3388            Field::new("c2", DataType::Int32, false),
3389        ]);
3390        // test c1 not in(1, 2, 3)
3391        let expr = Expr::InList(InList::new(
3392            Box::new(col("c1")),
3393            vec![lit(1), lit(2), lit(3)],
3394            true,
3395        ));
3396        let expected_expr = "c1_null_count@2 != row_count@3 AND (c1_min@0 != 1 OR 1 != c1_max@1) AND c1_null_count@2 != row_count@3 AND (c1_min@0 != 2 OR 2 != c1_max@1) AND c1_null_count@2 != row_count@3 AND (c1_min@0 != 3 OR 3 != c1_max@1)";
3397        let predicate_expr =
3398            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3399        assert_eq!(predicate_expr.to_string(), expected_expr);
3400
3401        Ok(())
3402    }
3403
3404    #[test]
3405    fn row_group_predicate_between() -> Result<()> {
3406        let schema = Schema::new(vec![
3407            Field::new("c1", DataType::Int32, false),
3408            Field::new("c2", DataType::Int32, false),
3409        ]);
3410
3411        // test c1 BETWEEN 1 AND 5
3412        let expr1 = col("c1").between(lit(1), lit(5));
3413
3414        // test 1 <= c1 <= 5
3415        let expr2 = col("c1").gt_eq(lit(1)).and(col("c1").lt_eq(lit(5)));
3416
3417        let predicate_expr1 =
3418            test_build_predicate_expression(&expr1, &schema, &mut RequiredColumns::new());
3419
3420        let predicate_expr2 =
3421            test_build_predicate_expression(&expr2, &schema, &mut RequiredColumns::new());
3422        assert_eq!(predicate_expr1.to_string(), predicate_expr2.to_string());
3423
3424        Ok(())
3425    }
3426
3427    #[test]
3428    fn row_group_predicate_between_with_in_list() -> Result<()> {
3429        let schema = Schema::new(vec![
3430            Field::new("c1", DataType::Int32, false),
3431            Field::new("c2", DataType::Int32, false),
3432        ]);
3433        // test c1 in(1, 2)
3434        let expr1 = col("c1").in_list(vec![lit(1), lit(2)], false);
3435
3436        // test c2 BETWEEN 4 AND 5
3437        let expr2 = col("c2").between(lit(4), lit(5));
3438
3439        // test c1 in(1, 2) and c2 BETWEEN 4 AND 5
3440        let expr3 = expr1.and(expr2);
3441
3442        let expected_expr = "(c1_null_count@2 != row_count@3 AND c1_min@0 <= 1 AND 1 <= c1_max@1 OR c1_null_count@2 != row_count@3 AND c1_min@0 <= 2 AND 2 <= c1_max@1) AND c2_null_count@5 != row_count@3 AND c2_max@4 >= 4 AND c2_null_count@5 != row_count@3 AND c2_min@6 <= 5";
3443        let predicate_expr =
3444            test_build_predicate_expression(&expr3, &schema, &mut RequiredColumns::new());
3445        assert_eq!(predicate_expr.to_string(), expected_expr);
3446
3447        Ok(())
3448    }
3449
3450    #[test]
3451    fn row_group_predicate_in_list_to_many_values() -> Result<()> {
3452        let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
3453        // test c1 in(1..21)
3454        // in pruning.rs has MAX_IN_LIST_SIZE = 20, more than this value will be rewrite
3455        // always true
3456        let expr = col("c1").in_list((1..=21).map(lit).collect(), false);
3457
3458        let expected_expr = "true";
3459        let predicate_expr =
3460            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3461        assert_eq!(predicate_expr.to_string(), expected_expr);
3462
3463        Ok(())
3464    }
3465
3466    // With the configurable cap, a caller that raises
3467    // `max_in_list_size` above the default gets the IN list rewritten
3468    // into a per-value min/max chain instead of falling through to `true`.
3469    // This verifies both `PredicateRewriter::with_max_in_list_size` and the
3470    // recursive OR path inside `build_predicate_expression`.
3471    #[test]
3472    fn row_group_predicate_in_list_rewritten_at_raised_cap() -> Result<()> {
3473        let schema =
3474            Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)]));
3475        // 25 items — above the default 20, below a raised cap of 32.
3476        let expr = col("c1").in_list((1..=25).map(lit).collect(), false);
3477        let physical = logical2physical(&expr, &schema);
3478        let rewriter = PredicateRewriter::new().with_max_in_list_size(32);
3479        let predicate_expr =
3480            rewriter.rewrite_predicate_to_statistics_predicate(&physical, &schema);
3481        // At the raised cap, IN is rewritten into per-value min/max checks
3482        // OR'd together; the resulting predicate must not collapse to
3483        // `true` (which is what the default cap produces).
3484        assert_ne!(
3485            predicate_expr.to_string(),
3486            "true",
3487            "IN(25) with raised cap must rewrite into a statistics-based predicate, not fall through to `true`"
3488        );
3489        // Sanity: the rewritten predicate references per-value literals.
3490        assert!(
3491            predicate_expr.to_string().contains(" <= 1 ")
3492                && predicate_expr.to_string().contains(" <= 25 "),
3493            "rewritten predicate should include per-value bounds for each IN entry, got: {predicate_expr}"
3494        );
3495        Ok(())
3496    }
3497
3498    // Guard: when the cap is 0 (opt-out) the IN branch is skipped entirely
3499    // regardless of list length, so even a small IN falls through to the
3500    // unhandled hook.
3501    #[test]
3502    fn row_group_predicate_in_list_disabled_at_zero_cap() -> Result<()> {
3503        let schema =
3504            Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)]));
3505        let expr = col("c1").in_list(vec![lit(1), lit(2), lit(3)], false);
3506        let physical = logical2physical(&expr, &schema);
3507        let rewriter = PredicateRewriter::new().with_max_in_list_size(0);
3508        let predicate_expr =
3509            rewriter.rewrite_predicate_to_statistics_predicate(&physical, &schema);
3510        assert_eq!(
3511            predicate_expr.to_string(),
3512            "true",
3513            "cap=0 must skip IN rewrite even for small lists"
3514        );
3515        Ok(())
3516    }
3517
3518    // The high-level [`PruningPredicateBuilder`] should thread
3519    // `max_in_list_size` all the way through: a 25-item IN with the default
3520    // cap must fall through to the unhandled hook (`predicate_expr = true`),
3521    // while a raised cap produces a real per-value statistics predicate.
3522    #[test]
3523    fn pruning_predicate_builder_threads_max_in_list_size() -> Result<()> {
3524        let schema =
3525            Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)]));
3526        let expr = col("c1").in_list((1..=25).map(lit).collect(), false);
3527        let physical = logical2physical(&expr, &schema);
3528
3529        // With the default cap the IN branch bails out and the pruning
3530        // predicate expression collapses to `true` (i.e., no container
3531        // pruning based on stats).
3532        let default_pp = PruningPredicateBuilder::new()
3533            .with_file_schema(Arc::clone(&schema))
3534            .try_build(Arc::clone(&physical))?;
3535        assert_eq!(
3536            default_pp.predicate_expr().to_string(),
3537            "true",
3538            "default cap must fall through to `true` for 25-item IN"
3539        );
3540
3541        // Raising the cap produces a real statistics predicate with per-
3542        // value bounds.
3543        let raised_pp = PruningPredicateBuilder::new()
3544            .with_file_schema(Arc::clone(&schema))
3545            .with_max_in_list_size(32)
3546            .try_build(physical)?;
3547        let raised_expr = raised_pp.predicate_expr().to_string();
3548        assert_ne!(
3549            raised_expr, "true",
3550            "raised cap must produce a real statistics predicate for 25-item IN"
3551        );
3552        assert!(
3553            raised_expr.contains(" <= 1 ") && raised_expr.contains(" <= 25 "),
3554            "raised-cap predicate should include per-value bounds, got: {raised_expr}"
3555        );
3556        Ok(())
3557    }
3558
3559    #[test]
3560    #[expect(deprecated)]
3561    fn deprecated_try_new_delegates_to_builder() -> Result<()> {
3562        let schema =
3563            Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)]));
3564        let expr = logical2physical(&col("c1").eq(lit(1)), &schema);
3565
3566        let deprecated =
3567            PruningPredicate::try_new(Arc::clone(&expr), Arc::clone(&schema))?;
3568        let builder = PruningPredicateBuilder::new()
3569            .with_file_schema(schema)
3570            .try_build(expr)?;
3571
3572        assert_eq!(
3573            deprecated.predicate_expr().to_string(),
3574            builder.predicate_expr().to_string()
3575        );
3576        assert_eq!(
3577            deprecated.required_columns().schema(),
3578            builder.required_columns().schema()
3579        );
3580        Ok(())
3581    }
3582
3583    #[test]
3584    fn row_group_predicate_cast_int_int() -> Result<()> {
3585        let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
3586        let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Int64) <= 1 AND 1 <= CAST(c1_max@1 AS Int64)";
3587
3588        // test cast(c1 as int64) = 1
3589        // test column on the left
3590        let expr = cast(col("c1"), DataType::Int64).eq(lit(ScalarValue::Int64(Some(1))));
3591        let predicate_expr =
3592            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3593        assert_eq!(predicate_expr.to_string(), expected_expr);
3594
3595        // test column on the right
3596        let expr = lit(ScalarValue::Int64(Some(1))).eq(cast(col("c1"), DataType::Int64));
3597        let predicate_expr =
3598            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3599        assert_eq!(predicate_expr.to_string(), expected_expr);
3600
3601        let expected_expr =
3602            "c1_null_count@1 != row_count@2 AND TRY_CAST(c1_max@0 AS Int64) > 1";
3603
3604        // test column on the left
3605        let expr =
3606            try_cast(col("c1"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(1))));
3607        let predicate_expr =
3608            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3609        assert_eq!(predicate_expr.to_string(), expected_expr);
3610
3611        // test column on the right
3612        let expr =
3613            lit(ScalarValue::Int64(Some(1))).lt(try_cast(col("c1"), DataType::Int64));
3614        let predicate_expr =
3615            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3616        assert_eq!(predicate_expr.to_string(), expected_expr);
3617
3618        Ok(())
3619    }
3620
3621    #[test]
3622    fn row_group_predicate_cast_string_string() -> Result<()> {
3623        let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]);
3624        let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Utf8) <= 1 AND 1 <= CAST(c1_max@1 AS Utf8)";
3625
3626        // test column on the left
3627        let expr = cast(col("c1"), DataType::Utf8)
3628            .eq(lit(ScalarValue::Utf8(Some("1".to_string()))));
3629        let predicate_expr =
3630            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3631        assert_eq!(predicate_expr.to_string(), expected_expr);
3632
3633        // test column on the right
3634        let expr = lit(ScalarValue::Utf8(Some("1".to_string())))
3635            .eq(cast(col("c1"), DataType::Utf8));
3636        let predicate_expr =
3637            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3638        assert_eq!(predicate_expr.to_string(), expected_expr);
3639
3640        Ok(())
3641    }
3642
3643    #[test]
3644    fn row_group_predicate_cast_string_int() -> Result<()> {
3645        let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]);
3646        let expected_expr = "true";
3647
3648        // test column on the left
3649        let expr = cast(col("c1"), DataType::Int32).eq(lit(ScalarValue::Int32(Some(1))));
3650        let predicate_expr =
3651            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3652        assert_eq!(predicate_expr.to_string(), expected_expr);
3653
3654        // test column on the right
3655        let expr = lit(ScalarValue::Int32(Some(1))).eq(cast(col("c1"), DataType::Int32));
3656        let predicate_expr =
3657            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3658        assert_eq!(predicate_expr.to_string(), expected_expr);
3659
3660        Ok(())
3661    }
3662
3663    #[test]
3664    fn row_group_predicate_cast_int_string() -> Result<()> {
3665        let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
3666        let expected_expr = "true";
3667
3668        // test column on the left
3669        let expr = cast(col("c1"), DataType::Utf8)
3670            .eq(lit(ScalarValue::Utf8(Some("1".to_string()))));
3671        let predicate_expr =
3672            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3673        assert_eq!(predicate_expr.to_string(), expected_expr);
3674
3675        // test column on the right
3676        let expr = lit(ScalarValue::Utf8(Some("1".to_string())))
3677            .eq(cast(col("c1"), DataType::Utf8));
3678        let predicate_expr =
3679            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3680        assert_eq!(predicate_expr.to_string(), expected_expr);
3681
3682        Ok(())
3683    }
3684
3685    #[test]
3686    fn row_group_predicate_date_date() -> Result<()> {
3687        let schema = Schema::new(vec![Field::new("c1", DataType::Date32, false)]);
3688        let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Date64) <= 1970-01-01 AND 1970-01-01 <= CAST(c1_max@1 AS Date64)";
3689
3690        // test column on the left
3691        let expr =
3692            cast(col("c1"), DataType::Date64).eq(lit(ScalarValue::Date64(Some(123))));
3693        let predicate_expr =
3694            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3695        assert_eq!(predicate_expr.to_string(), expected_expr);
3696
3697        // test column on the right
3698        let expr =
3699            lit(ScalarValue::Date64(Some(123))).eq(cast(col("c1"), DataType::Date64));
3700        let predicate_expr =
3701            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3702        assert_eq!(predicate_expr.to_string(), expected_expr);
3703
3704        Ok(())
3705    }
3706
3707    #[test]
3708    fn row_group_predicate_dict_string_date() -> Result<()> {
3709        // Test with Dictionary<UInt8, Utf8> for the literal
3710        let schema = Schema::new(vec![Field::new("c1", DataType::Date32, false)]);
3711        let expected_expr = "true";
3712
3713        // test column on the left
3714        let expr = cast(
3715            col("c1"),
3716            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
3717        )
3718        .eq(lit(ScalarValue::Utf8(Some("2024-01-01".to_string()))));
3719        let predicate_expr =
3720            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3721        assert_eq!(predicate_expr.to_string(), expected_expr);
3722
3723        // test column on the right
3724        let expr = lit(ScalarValue::Utf8(Some("2024-01-01".to_string()))).eq(cast(
3725            col("c1"),
3726            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
3727        ));
3728        let predicate_expr =
3729            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3730        assert_eq!(predicate_expr.to_string(), expected_expr);
3731
3732        Ok(())
3733    }
3734
3735    #[test]
3736    fn row_group_predicate_date_dict_string() -> Result<()> {
3737        // Test with Dictionary<UInt8, Utf8> for the column
3738        let schema = Schema::new(vec![Field::new(
3739            "c1",
3740            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
3741            false,
3742        )]);
3743        let expected_expr = "true";
3744
3745        // test column on the left
3746        let expr =
3747            cast(col("c1"), DataType::Date32).eq(lit(ScalarValue::Date32(Some(123))));
3748        let predicate_expr =
3749            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3750        assert_eq!(predicate_expr.to_string(), expected_expr);
3751
3752        // test column on the right
3753        let expr =
3754            lit(ScalarValue::Date32(Some(123))).eq(cast(col("c1"), DataType::Date32));
3755        let predicate_expr =
3756            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3757        assert_eq!(predicate_expr.to_string(), expected_expr);
3758
3759        Ok(())
3760    }
3761
3762    #[test]
3763    fn row_group_predicate_dict_dict_same_value_type() -> Result<()> {
3764        // Test with Dictionary types that have the same value type but different key types
3765        let schema = Schema::new(vec![Field::new(
3766            "c1",
3767            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
3768            false,
3769        )]);
3770
3771        // Direct comparison with no cast
3772        let expr = col("c1").eq(lit(ScalarValue::Utf8(Some("test".to_string()))));
3773        let predicate_expr =
3774            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3775        let expected_expr =
3776            "c1_null_count@2 != row_count@3 AND c1_min@0 <= test AND test <= c1_max@1";
3777        assert_eq!(predicate_expr.to_string(), expected_expr);
3778
3779        // Test with column cast to a dictionary with different key type
3780        let expr = cast(
3781            col("c1"),
3782            DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3783        )
3784        .eq(lit(ScalarValue::Utf8(Some("test".to_string()))));
3785        let predicate_expr =
3786            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3787        let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Dictionary(UInt16, Utf8)) <= test AND test <= CAST(c1_max@1 AS Dictionary(UInt16, Utf8))";
3788        assert_eq!(predicate_expr.to_string(), expected_expr);
3789
3790        Ok(())
3791    }
3792
3793    #[test]
3794    fn row_group_predicate_dict_dict_different_value_type() -> Result<()> {
3795        // Test with Dictionary types that have different value types
3796        let schema = Schema::new(vec![Field::new(
3797            "c1",
3798            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Int32)),
3799            false,
3800        )]);
3801        let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Int64) <= 123 AND 123 <= CAST(c1_max@1 AS Int64)";
3802
3803        // Test with literal of a different type
3804        let expr =
3805            cast(col("c1"), DataType::Int64).eq(lit(ScalarValue::Int64(Some(123))));
3806        let predicate_expr =
3807            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3808        assert_eq!(predicate_expr.to_string(), expected_expr);
3809
3810        Ok(())
3811    }
3812
3813    #[test]
3814    fn row_group_predicate_nested_dict() -> Result<()> {
3815        // Test with nested Dictionary types
3816        let schema = Schema::new(vec![Field::new(
3817            "c1",
3818            DataType::Dictionary(
3819                Box::new(DataType::UInt8),
3820                Box::new(DataType::Dictionary(
3821                    Box::new(DataType::UInt16),
3822                    Box::new(DataType::Utf8),
3823                )),
3824            ),
3825            false,
3826        )]);
3827        let expected_expr =
3828            "c1_null_count@2 != row_count@3 AND c1_min@0 <= test AND test <= c1_max@1";
3829
3830        // Test with a simple literal
3831        let expr = col("c1").eq(lit(ScalarValue::Utf8(Some("test".to_string()))));
3832        let predicate_expr =
3833            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3834        assert_eq!(predicate_expr.to_string(), expected_expr);
3835
3836        Ok(())
3837    }
3838
3839    #[test]
3840    fn row_group_predicate_dict_date_dict_date() -> Result<()> {
3841        // Test with dictionary-wrapped date types for both sides
3842        let schema = Schema::new(vec![Field::new(
3843            "c1",
3844            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Date32)),
3845            false,
3846        )]);
3847        let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Dictionary(UInt16, Date64)) <= 1970-01-01 AND 1970-01-01 <= CAST(c1_max@1 AS Dictionary(UInt16, Date64))";
3848
3849        // Test with a cast to a different date type
3850        let expr = cast(
3851            col("c1"),
3852            DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Date64)),
3853        )
3854        .eq(lit(ScalarValue::Date64(Some(123))));
3855        let predicate_expr =
3856            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3857        assert_eq!(predicate_expr.to_string(), expected_expr);
3858
3859        Ok(())
3860    }
3861
3862    #[test]
3863    fn row_group_predicate_date_string() -> Result<()> {
3864        let schema = Schema::new(vec![Field::new("c1", DataType::Utf8, false)]);
3865        let expected_expr = "true";
3866
3867        // test column on the left
3868        let expr =
3869            cast(col("c1"), DataType::Date32).eq(lit(ScalarValue::Date32(Some(123))));
3870        let predicate_expr =
3871            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3872        assert_eq!(predicate_expr.to_string(), expected_expr);
3873
3874        // test column on the right
3875        let expr =
3876            lit(ScalarValue::Date32(Some(123))).eq(cast(col("c1"), DataType::Date32));
3877        let predicate_expr =
3878            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3879        assert_eq!(predicate_expr.to_string(), expected_expr);
3880
3881        Ok(())
3882    }
3883
3884    #[test]
3885    fn row_group_predicate_string_date() -> Result<()> {
3886        let schema = Schema::new(vec![Field::new("c1", DataType::Date32, false)]);
3887        let expected_expr = "true";
3888
3889        // test column on the left
3890        let expr = cast(col("c1"), DataType::Utf8)
3891            .eq(lit(ScalarValue::Utf8(Some("2024-01-01".to_string()))));
3892        let predicate_expr =
3893            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3894        assert_eq!(predicate_expr.to_string(), expected_expr);
3895
3896        // test column on the right
3897        let expr = lit(ScalarValue::Utf8(Some("2024-01-01".to_string())))
3898            .eq(cast(col("c1"), DataType::Utf8));
3899        let predicate_expr =
3900            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3901        assert_eq!(predicate_expr.to_string(), expected_expr);
3902
3903        Ok(())
3904    }
3905
3906    #[test]
3907    fn row_group_predicate_cast_list() -> Result<()> {
3908        let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
3909        // test cast(c1 as int64) in int64(1, 2, 3)
3910        let expr = Expr::InList(InList::new(
3911            Box::new(cast(col("c1"), DataType::Int64)),
3912            vec![
3913                lit(ScalarValue::Int64(Some(1))),
3914                lit(ScalarValue::Int64(Some(2))),
3915                lit(ScalarValue::Int64(Some(3))),
3916            ],
3917            false,
3918        ));
3919        let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Int64) <= 1 AND 1 <= CAST(c1_max@1 AS Int64) OR c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Int64) <= 2 AND 2 <= CAST(c1_max@1 AS Int64) OR c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Int64) <= 3 AND 3 <= CAST(c1_max@1 AS Int64)";
3920        let predicate_expr =
3921            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3922        assert_eq!(predicate_expr.to_string(), expected_expr);
3923
3924        let expr = Expr::InList(InList::new(
3925            Box::new(cast(col("c1"), DataType::Int64)),
3926            vec![
3927                lit(ScalarValue::Int64(Some(1))),
3928                lit(ScalarValue::Int64(Some(2))),
3929                lit(ScalarValue::Int64(Some(3))),
3930            ],
3931            true,
3932        ));
3933        let expected_expr = "c1_null_count@2 != row_count@3 AND (CAST(c1_min@0 AS Int64) != 1 OR 1 != CAST(c1_max@1 AS Int64)) AND c1_null_count@2 != row_count@3 AND (CAST(c1_min@0 AS Int64) != 2 OR 2 != CAST(c1_max@1 AS Int64)) AND c1_null_count@2 != row_count@3 AND (CAST(c1_min@0 AS Int64) != 3 OR 3 != CAST(c1_max@1 AS Int64))";
3934        let predicate_expr =
3935            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3936        assert_eq!(predicate_expr.to_string(), expected_expr);
3937
3938        Ok(())
3939    }
3940
3941    #[test]
3942    fn prune_decimal_data() {
3943        // decimal(9,2)
3944        let schema = Arc::new(Schema::new(vec![Field::new(
3945            "s1",
3946            DataType::Decimal128(9, 2),
3947            true,
3948        )]));
3949
3950        prune_with_expr(
3951            // s1 > 5
3952            col("s1").gt(lit(ScalarValue::Decimal128(Some(500), 9, 2))),
3953            &schema,
3954            // If the data is written by spark, the physical data type is INT32 in the parquet
3955            // So we use the INT32 type of statistic.
3956            &TestStatistics::new().with(
3957                "s1",
3958                ContainerStats::new_i32(
3959                    vec![Some(0), Some(4), None, Some(3)], // min
3960                    vec![Some(5), Some(6), Some(4), None], // max
3961                ),
3962            ),
3963            &[false, true, false, true],
3964        );
3965
3966        prune_with_expr(
3967            // with cast column to other type
3968            cast(col("s1"), DataType::Decimal128(14, 3))
3969                .gt(lit(ScalarValue::Decimal128(Some(5000), 14, 3))),
3970            &schema,
3971            &TestStatistics::new().with(
3972                "s1",
3973                ContainerStats::new_i32(
3974                    vec![Some(0), Some(4), None, Some(3)], // min
3975                    vec![Some(5), Some(6), Some(4), None], // max
3976                ),
3977            ),
3978            &[false, true, false, true],
3979        );
3980
3981        prune_with_expr(
3982            // with try cast column to other type
3983            try_cast(col("s1"), DataType::Decimal128(14, 3))
3984                .gt(lit(ScalarValue::Decimal128(Some(5000), 14, 3))),
3985            &schema,
3986            &TestStatistics::new().with(
3987                "s1",
3988                ContainerStats::new_i32(
3989                    vec![Some(0), Some(4), None, Some(3)], // min
3990                    vec![Some(5), Some(6), Some(4), None], // max
3991                ),
3992            ),
3993            &[false, true, false, true],
3994        );
3995
3996        // decimal(18,2)
3997        let schema = Arc::new(Schema::new(vec![Field::new(
3998            "s1",
3999            DataType::Decimal128(18, 2),
4000            true,
4001        )]));
4002        prune_with_expr(
4003            // s1 > 5
4004            col("s1").gt(lit(ScalarValue::Decimal128(Some(500), 18, 2))),
4005            &schema,
4006            // If the data is written by spark, the physical data type is INT64 in the parquet
4007            // So we use the INT32 type of statistic.
4008            &TestStatistics::new().with(
4009                "s1",
4010                ContainerStats::new_i64(
4011                    vec![Some(0), Some(4), None, Some(3)], // min
4012                    vec![Some(5), Some(6), Some(4), None], // max
4013                ),
4014            ),
4015            &[false, true, false, true],
4016        );
4017
4018        // decimal(23,2)
4019        let schema = Arc::new(Schema::new(vec![Field::new(
4020            "s1",
4021            DataType::Decimal128(23, 2),
4022            true,
4023        )]));
4024
4025        prune_with_expr(
4026            // s1 > 5
4027            col("s1").gt(lit(ScalarValue::Decimal128(Some(500), 23, 2))),
4028            &schema,
4029            &TestStatistics::new().with(
4030                "s1",
4031                ContainerStats::new_decimal128(
4032                    vec![Some(0), Some(400), None, Some(300)], // min
4033                    vec![Some(500), Some(600), Some(400), None], // max
4034                    23,
4035                    2,
4036                ),
4037            ),
4038            &[false, true, false, true],
4039        );
4040    }
4041
4042    #[test]
4043    fn prune_api() {
4044        let schema = Arc::new(Schema::new(vec![
4045            Field::new("s1", DataType::Utf8, true),
4046            Field::new("s2", DataType::Int32, true),
4047        ]));
4048
4049        let statistics = TestStatistics::new().with(
4050            "s2",
4051            ContainerStats::new_i32(
4052                vec![Some(0), Some(4), None, Some(3)], // min
4053                vec![Some(5), Some(6), None, None],    // max
4054            ),
4055        );
4056        prune_with_expr(
4057            // Prune using s2 > 5
4058            col("s2").gt(lit(5)),
4059            &schema,
4060            &statistics,
4061            // s2 [0, 5] ==> no rows should pass
4062            // s2 [4, 6] ==> some rows could pass
4063            // No stats for s2 ==> some rows could pass
4064            // s2 [3, None] (null max) ==> some rows could pass
4065            &[false, true, true, true],
4066        );
4067
4068        prune_with_expr(
4069            // filter with cast
4070            cast(col("s2"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(5)))),
4071            &schema,
4072            &statistics,
4073            &[false, true, true, true],
4074        );
4075    }
4076
4077    #[test]
4078    fn prune_not_eq_data() {
4079        let schema = Arc::new(Schema::new(vec![Field::new("s1", DataType::Utf8, true)]));
4080
4081        prune_with_expr(
4082            // Prune using s2 != 'M'
4083            col("s1").not_eq(lit("M")),
4084            &schema,
4085            &TestStatistics::new().with(
4086                "s1",
4087                ContainerStats::new_utf8(
4088                    vec![Some("A"), Some("A"), Some("N"), Some("M"), None, Some("A")], // min
4089                    vec![Some("Z"), Some("L"), Some("Z"), Some("M"), None, None], // max
4090                ),
4091            ),
4092            // s1 [A, Z] ==> might have values that pass predicate
4093            // s1 [A, L] ==> all rows pass the predicate
4094            // s1 [N, Z] ==> all rows pass the predicate
4095            // s1 [M, M] ==> all rows do not pass the predicate
4096            // No stats for s2 ==> some rows could pass
4097            // s2 [3, None] (null max) ==> some rows could pass
4098            &[true, true, true, false, true, true],
4099        );
4100    }
4101
4102    /// Creates setup for boolean chunk pruning
4103    ///
4104    /// For predicate "b1" (boolean expr)
4105    /// b1 [false, false] ==> no rows can pass (not keep)
4106    /// b1 [false, true] ==> some rows could pass (must keep)
4107    /// b1 [true, true] ==> all rows must pass (must keep)
4108    /// b1 [NULL, NULL]  ==> unknown (must keep)
4109    /// b1 [false, NULL]  ==> unknown (must keep)
4110    ///
4111    /// For predicate "!b1" (boolean expr)
4112    /// b1 [false, false] ==> all rows pass (must keep)
4113    /// b1 [false, true] ==> some rows could pass (must keep)
4114    /// b1 [true, true] ==> no rows can pass (not keep)
4115    /// b1 [NULL, NULL]  ==> unknown (must keep)
4116    /// b1 [false, NULL]  ==> unknown (must keep)
4117    fn bool_setup() -> (SchemaRef, TestStatistics, Vec<bool>, Vec<bool>) {
4118        let schema =
4119            Arc::new(Schema::new(vec![Field::new("b1", DataType::Boolean, true)]));
4120
4121        let statistics = TestStatistics::new().with(
4122            "b1",
4123            ContainerStats::new_bool(
4124                vec![Some(false), Some(false), Some(true), None, Some(false)], // min
4125                vec![Some(false), Some(true), Some(true), None, None],         // max
4126            ),
4127        );
4128        let expected_true = vec![false, true, true, true, true];
4129        let expected_false = vec![true, true, false, true, true];
4130
4131        (schema, statistics, expected_true, expected_false)
4132    }
4133
4134    #[test]
4135    fn prune_bool_const_expr() {
4136        let (schema, statistics, _, _) = bool_setup();
4137
4138        prune_with_expr(
4139            // true
4140            lit(true),
4141            &schema,
4142            &statistics,
4143            &[true, true, true, true, true],
4144        );
4145
4146        prune_with_expr(
4147            // false
4148            lit(false),
4149            &schema,
4150            &statistics,
4151            &[false, false, false, false, false],
4152        );
4153    }
4154
4155    #[test]
4156    fn prune_bool_column() {
4157        let (schema, statistics, expected_true, _) = bool_setup();
4158
4159        prune_with_expr(
4160            // b1
4161            col("b1"),
4162            &schema,
4163            &statistics,
4164            &expected_true,
4165        );
4166    }
4167
4168    #[test]
4169    fn prune_bool_not_column() {
4170        let (schema, statistics, _, expected_false) = bool_setup();
4171
4172        prune_with_expr(
4173            // !b1
4174            col("b1").not(),
4175            &schema,
4176            &statistics,
4177            &expected_false,
4178        );
4179    }
4180
4181    #[test]
4182    fn prune_bool_column_eq_true() {
4183        let (schema, statistics, expected_true, _) = bool_setup();
4184
4185        prune_with_expr(
4186            // b1 = true
4187            col("b1").eq(lit(true)),
4188            &schema,
4189            &statistics,
4190            &expected_true,
4191        );
4192    }
4193
4194    #[test]
4195    fn prune_bool_not_column_eq_true() {
4196        let (schema, statistics, _, expected_false) = bool_setup();
4197
4198        prune_with_expr(
4199            // !b1 = true
4200            col("b1").not().eq(lit(true)),
4201            &schema,
4202            &statistics,
4203            &expected_false,
4204        );
4205    }
4206
4207    /// Creates a setup for chunk pruning, modeling a int32 column "i"
4208    /// with 5 different containers (e.g. RowGroups). They have [min,
4209    /// max]:
4210    ///
4211    /// i [-5, 5]
4212    /// i [1, 11]
4213    /// i [-11, -1]
4214    /// i [NULL, NULL]
4215    /// i [1, NULL]
4216    fn int32_setup() -> (SchemaRef, TestStatistics) {
4217        let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
4218
4219        let statistics = TestStatistics::new().with(
4220            "i",
4221            ContainerStats::new_i32(
4222                vec![Some(-5), Some(1), Some(-11), None, Some(1)], // min
4223                vec![Some(5), Some(11), Some(-1), None, None],     // max
4224            ),
4225        );
4226        (schema, statistics)
4227    }
4228
4229    #[test]
4230    fn prune_int32_col_gt_zero() {
4231        let (schema, statistics) = int32_setup();
4232
4233        // Expression "i > 0" and "-i < 0"
4234        // i [-5, 5] ==> some rows could pass (must keep)
4235        // i [1, 11] ==> all rows must pass (must keep)
4236        // i [-11, -1] ==>  no rows can pass (not keep)
4237        // i [NULL, NULL]  ==> unknown (must keep)
4238        // i [1, NULL]  ==> unknown (must keep)
4239        let expected_ret = &[true, true, false, true, true];
4240
4241        // i > 0
4242        prune_with_expr(col("i").gt(lit(0)), &schema, &statistics, expected_ret);
4243
4244        // -i < 0
4245        prune_with_expr(
4246            Expr::Negative(Box::new(col("i"))).lt(lit(0)),
4247            &schema,
4248            &statistics,
4249            expected_ret,
4250        );
4251    }
4252
4253    #[test]
4254    fn prune_int32_col_lte_zero() {
4255        let (schema, statistics) = int32_setup();
4256
4257        // Expression "i <= 0" and "-i >= 0"
4258        // i [-5, 5] ==> some rows could pass (must keep)
4259        // i [1, 11] ==> no rows can pass (not keep)
4260        // i [-11, -1] ==>  all rows must pass (must keep)
4261        // i [NULL, NULL]  ==> unknown (must keep)
4262        // i [1, NULL]  ==> no rows can pass (not keep)
4263        let expected_ret = &[true, false, true, true, false];
4264
4265        prune_with_expr(
4266            // i <= 0
4267            col("i").lt_eq(lit(0)),
4268            &schema,
4269            &statistics,
4270            expected_ret,
4271        );
4272
4273        prune_with_expr(
4274            // -i >= 0
4275            Expr::Negative(Box::new(col("i"))).gt_eq(lit(0)),
4276            &schema,
4277            &statistics,
4278            expected_ret,
4279        );
4280    }
4281
4282    #[test]
4283    fn prune_int32_col_lte_zero_cast() {
4284        let (schema, statistics) = int32_setup();
4285
4286        // Expression "cast(i as utf8) <= '0'"
4287        // i [-5, 5] ==> some rows could pass (must keep)
4288        // i [1, 11] ==> no rows can pass in theory, -0.22 (conservatively keep)
4289        // i [-11, -1] ==>  no rows could pass in theory (conservatively keep)
4290        // i [NULL, NULL]  ==> unknown (must keep)
4291        // i [1, NULL]  ==> no rows can pass (conservatively keep)
4292        let expected_ret = &[true, true, true, true, true];
4293
4294        prune_with_expr(
4295            // cast(i as utf8) <= 0
4296            cast(col("i"), DataType::Utf8).lt_eq(lit("0")),
4297            &schema,
4298            &statistics,
4299            expected_ret,
4300        );
4301
4302        prune_with_expr(
4303            // try_cast(i as utf8) <= 0
4304            try_cast(col("i"), DataType::Utf8).lt_eq(lit("0")),
4305            &schema,
4306            &statistics,
4307            expected_ret,
4308        );
4309
4310        prune_with_expr(
4311            // cast(-i as utf8) >= 0
4312            cast(Expr::Negative(Box::new(col("i"))), DataType::Utf8).gt_eq(lit("0")),
4313            &schema,
4314            &statistics,
4315            expected_ret,
4316        );
4317
4318        prune_with_expr(
4319            // try_cast(-i as utf8) >= 0
4320            try_cast(Expr::Negative(Box::new(col("i"))), DataType::Utf8).gt_eq(lit("0")),
4321            &schema,
4322            &statistics,
4323            expected_ret,
4324        );
4325    }
4326
4327    #[test]
4328    fn prune_int32_col_eq_zero() {
4329        let (schema, statistics) = int32_setup();
4330
4331        // Expression "i = 0"
4332        // i [-5, 5] ==> some rows could pass (must keep)
4333        // i [1, 11] ==> no rows can pass (not keep)
4334        // i [-11, -1] ==>  no rows can pass (not keep)
4335        // i [NULL, NULL]  ==> unknown (must keep)
4336        // i [1, NULL]  ==> no rows can pass (not keep)
4337        let expected_ret = &[true, false, false, true, false];
4338
4339        prune_with_expr(
4340            // i = 0
4341            col("i").eq(lit(0)),
4342            &schema,
4343            &statistics,
4344            expected_ret,
4345        );
4346    }
4347
4348    #[test]
4349    fn prune_int32_col_is_not_distinct_from() {
4350        let (schema, statistics) = int32_setup();
4351
4352        // Without null counts, IS NOT DISTINCT FROM a non-null literal can
4353        // still use min/max ranges, but unknown all-null containers must be kept.
4354        let expected_ret = &[true, false, false, true, false];
4355
4356        prune_with_expr(
4357            is_not_distinct_from(col("i"), lit(0)),
4358            &schema,
4359            &statistics,
4360            expected_ret,
4361        );
4362
4363        // The operator is symmetric, so the scalar-left form should prune the
4364        // same row groups.
4365        prune_with_expr(
4366            is_not_distinct_from(lit(0), col("i")),
4367            &schema,
4368            &statistics,
4369            expected_ret,
4370        );
4371
4372        let statistics = statistics
4373            .with_row_counts("i", vec![Some(10), Some(9), None, Some(4), Some(10)])
4374            .with_null_counts("i", vec![Some(0), Some(1), None, Some(4), Some(0)]);
4375
4376        let expected_ret = &[true, false, false, false, false];
4377        prune_with_expr(
4378            is_not_distinct_from(col("i"), lit(0)),
4379            &schema,
4380            &statistics,
4381            expected_ret,
4382        );
4383
4384        let expected_ret = &[false, true, true, true, false];
4385        prune_with_expr(
4386            is_not_distinct_from(col("i"), lit(ScalarValue::Int32(None))),
4387            &schema,
4388            &statistics,
4389            expected_ret,
4390        );
4391    }
4392
4393    #[test]
4394    fn prune_int32_col_is_distinct_from() {
4395        let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
4396        let statistics = TestStatistics::new().with(
4397            "i",
4398            ContainerStats::new_i32(
4399                vec![Some(0), Some(0), Some(5), None],
4400                vec![Some(0), Some(2), Some(5), None],
4401            )
4402            .with_row_counts(vec![Some(2), Some(2), Some(2), Some(2)])
4403            .with_null_counts(vec![Some(0), Some(0), Some(0), Some(2)]),
4404        );
4405
4406        let expected_ret = &[false, true, true, true];
4407        prune_with_expr(
4408            is_distinct_from(col("i"), lit(0)),
4409            &schema,
4410            &statistics,
4411            expected_ret,
4412        );
4413
4414        // The operator is symmetric, so the scalar-left form should prune the
4415        // same row groups.
4416        prune_with_expr(
4417            is_distinct_from(lit(0), col("i")),
4418            &schema,
4419            &statistics,
4420            expected_ret,
4421        );
4422
4423        let expected_ret = &[true, true, true, false];
4424        prune_with_expr(
4425            is_distinct_from(col("i"), lit(ScalarValue::Int32(None))),
4426            &schema,
4427            &statistics,
4428            expected_ret,
4429        );
4430    }
4431
4432    #[test]
4433    fn prune_int32_col_eq_zero_cast() {
4434        let (schema, statistics) = int32_setup();
4435
4436        // Expression "cast(i as int64) = 0"
4437        // i [-5, 5] ==> some rows could pass (must keep)
4438        // i [1, 11] ==> no rows can pass (not keep)
4439        // i [-11, -1] ==>  no rows can pass (not keep)
4440        // i [NULL, NULL]  ==> unknown (must keep)
4441        // i [1, NULL]  ==> no rows can pass (not keep)
4442        let expected_ret = &[true, false, false, true, false];
4443
4444        prune_with_expr(
4445            cast(col("i"), DataType::Int64).eq(lit(0i64)),
4446            &schema,
4447            &statistics,
4448            expected_ret,
4449        );
4450
4451        prune_with_expr(
4452            try_cast(col("i"), DataType::Int64).eq(lit(0i64)),
4453            &schema,
4454            &statistics,
4455            expected_ret,
4456        );
4457    }
4458
4459    #[test]
4460    fn prune_int32_col_eq_zero_cast_as_str() {
4461        let (schema, statistics) = int32_setup();
4462
4463        // Note the cast is to a string where sorting properties are
4464        // not the same as integers
4465        //
4466        // Expression "cast(i as utf8) = '0'"
4467        // i [-5, 5] ==> some rows could pass (keep)
4468        // i [1, 11] ==> no rows can pass  (could keep)
4469        // i [-11, -1] ==>  no rows can pass (could keep)
4470        // i [NULL, NULL]  ==> unknown (keep)
4471        // i [1, NULL]  ==> no rows can pass (could keep)
4472        let expected_ret = &[true, true, true, true, true];
4473
4474        prune_with_expr(
4475            cast(col("i"), DataType::Utf8).eq(lit("0")),
4476            &schema,
4477            &statistics,
4478            expected_ret,
4479        );
4480    }
4481
4482    #[test]
4483    fn prune_int32_col_lt_neg_one() {
4484        let (schema, statistics) = int32_setup();
4485
4486        // Expression "i > -1" and "-i < 1"
4487        // i [-5, 5] ==> some rows could pass (must keep)
4488        // i [1, 11] ==> all rows must pass (must keep)
4489        // i [-11, -1] ==>  no rows can pass (not keep)
4490        // i [NULL, NULL]  ==> unknown (must keep)
4491        // i [1, NULL]  ==> all rows must pass (must keep)
4492        let expected_ret = &[true, true, false, true, true];
4493
4494        prune_with_expr(
4495            // i > -1
4496            col("i").gt(lit(-1)),
4497            &schema,
4498            &statistics,
4499            expected_ret,
4500        );
4501
4502        prune_with_expr(
4503            // -i < 1
4504            Expr::Negative(Box::new(col("i"))).lt(lit(1)),
4505            &schema,
4506            &statistics,
4507            expected_ret,
4508        );
4509    }
4510
4511    #[test]
4512    fn prune_int32_is_null() {
4513        let (schema, statistics) = int32_setup();
4514
4515        // Expression "i IS NULL" when there are no null statistics,
4516        // should all be kept
4517        let expected_ret = &[true, true, true, true, true];
4518
4519        prune_with_expr(
4520            // i IS NULL, no null statistics
4521            col("i").is_null(),
4522            &schema,
4523            &statistics,
4524            expected_ret,
4525        );
4526
4527        // provide null counts for each column
4528        let statistics = statistics.with_null_counts(
4529            "i",
4530            vec![
4531                Some(0), // no nulls (don't keep)
4532                Some(1), // 1 null
4533                None,    // unknown nulls
4534                None, // unknown nulls (min/max are both null too, like no stats at all)
4535                Some(0), // 0 nulls (max=null too which means no known max) (don't keep)
4536            ],
4537        );
4538
4539        let expected_ret = &[false, true, true, true, false];
4540
4541        prune_with_expr(
4542            // i IS NULL, with actual null statistics
4543            col("i").is_null(),
4544            &schema,
4545            &statistics,
4546            expected_ret,
4547        );
4548    }
4549
4550    #[test]
4551    fn prune_int32_column_is_known_all_null() {
4552        let (schema, statistics) = int32_setup();
4553
4554        // Expression "i < 0"
4555        // i [-5, 5] ==> some rows could pass (must keep)
4556        // i [1, 11] ==> no rows can pass (not keep)
4557        // i [-11, -1] ==>  all rows must pass (must keep)
4558        // i [NULL, NULL]  ==> unknown (must keep)
4559        // i [1, NULL]  ==> no rows can pass (not keep)
4560        let expected_ret = &[true, false, true, true, false];
4561
4562        prune_with_expr(
4563            // i < 0
4564            col("i").lt(lit(0)),
4565            &schema,
4566            &statistics,
4567            expected_ret,
4568        );
4569
4570        // provide row counts for each column
4571        let statistics = statistics.with_row_counts(
4572            "i",
4573            vec![
4574                Some(10), // 10 rows of data
4575                Some(9),  // 9 rows of data
4576                None,     // unknown row counts
4577                Some(4),
4578                Some(10),
4579            ],
4580        );
4581
4582        // pruning result is still the same if we only know row counts
4583        prune_with_expr(
4584            // i < 0, with only row counts statistics
4585            col("i").lt(lit(0)),
4586            &schema,
4587            &statistics,
4588            expected_ret,
4589        );
4590
4591        // provide null counts for each column
4592        let statistics = statistics.with_null_counts(
4593            "i",
4594            vec![
4595                Some(0), // no nulls
4596                Some(1), // 1 null
4597                None,    // unknown nulls
4598                Some(4), // 4 nulls, which is the same as the row counts, i.e. this column is all null (don't keep)
4599                Some(0), // 0 nulls (max=null too which means no known max)
4600            ],
4601        );
4602
4603        // Expression "i < 0" with actual null and row counts statistics
4604        // col | min, max     | row counts | null counts |
4605        // ----+--------------+------------+-------------+
4606        //  i  | [-5, 5]      | 10         | 0           | ==> Some rows could pass (must keep)
4607        //  i  | [1, 11]      | 9          | 1           | ==> No rows can pass (not keep)
4608        //  i  | [-11,-1]     | Unknown    | Unknown     | ==> All rows must pass (must keep)
4609        //  i  | [NULL, NULL] | 4          | 4           | ==> The column is all null (not keep)
4610        //  i  | [1, NULL]    | 10         | 0           | ==> No rows can pass (not keep)
4611        let expected_ret = &[true, false, true, false, false];
4612
4613        prune_with_expr(
4614            // i < 0, with actual null and row counts statistics
4615            col("i").lt(lit(0)),
4616            &schema,
4617            &statistics,
4618            expected_ret,
4619        );
4620    }
4621
4622    #[test]
4623    fn prune_cast_scalar() {
4624        // The data type of column i is INT32
4625        let (schema, statistics) = int32_setup();
4626        let expected_ret = &[true, true, false, true, true];
4627
4628        prune_with_expr(
4629            // i > int64(0)
4630            col("i").gt(cast(lit(ScalarValue::Int64(Some(0))), DataType::Int32)),
4631            &schema,
4632            &statistics,
4633            expected_ret,
4634        );
4635
4636        prune_with_expr(
4637            // cast(i as int64) > int64(0)
4638            cast(col("i"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(0)))),
4639            &schema,
4640            &statistics,
4641            expected_ret,
4642        );
4643
4644        prune_with_expr(
4645            // try_cast(i as int64) > int64(0)
4646            try_cast(col("i"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(0)))),
4647            &schema,
4648            &statistics,
4649            expected_ret,
4650        );
4651
4652        prune_with_expr(
4653            // `-cast(i as int64) < 0` convert to `cast(i as int64) > -0`
4654            Expr::Negative(Box::new(cast(col("i"), DataType::Int64)))
4655                .lt(lit(ScalarValue::Int64(Some(0)))),
4656            &schema,
4657            &statistics,
4658            expected_ret,
4659        );
4660    }
4661
4662    #[test]
4663    fn test_increment_utf8() {
4664        // Basic ASCII
4665        assert_eq!(increment_utf8("abc").unwrap(), "abd");
4666        assert_eq!(increment_utf8("abz").unwrap(), "ab{");
4667
4668        // Test around ASCII 127 (DEL)
4669        assert_eq!(increment_utf8("~").unwrap(), "\u{7f}"); // 126 -> 127
4670        assert_eq!(increment_utf8("\u{7f}").unwrap(), "\u{80}"); // 127 -> 128
4671
4672        // Test 2-byte UTF-8 sequences
4673        assert_eq!(increment_utf8("ß").unwrap(), "à"); // U+00DF -> U+00E0
4674
4675        // Test 3-byte UTF-8 sequences
4676        assert_eq!(increment_utf8("℣").unwrap(), "ℤ"); // U+2123 -> U+2124
4677
4678        // Test at UTF-8 boundaries
4679        assert_eq!(increment_utf8("\u{7FF}").unwrap(), "\u{800}"); // 2-byte to 3-byte boundary
4680        assert_eq!(increment_utf8("\u{FFFF}").unwrap(), "\u{10000}"); // 3-byte to 4-byte boundary
4681
4682        // Test that if we can't increment we return None
4683        assert!(increment_utf8("").is_none());
4684        assert!(increment_utf8("\u{10FFFF}").is_none()); // U+10FFFF is the max code point
4685
4686        // Test that if we can't increment the last character we do the previous one and truncate
4687        assert_eq!(increment_utf8("a\u{10FFFF}").unwrap(), "b");
4688
4689        // Test surrogate pair range (0xD800..=0xDFFF)
4690        assert_eq!(increment_utf8("a\u{D7FF}").unwrap(), "b");
4691        assert!(increment_utf8("\u{D7FF}").is_none());
4692
4693        // Test non-characters range (0xFDD0..=0xFDEF)
4694        assert_eq!(increment_utf8("a\u{FDCF}").unwrap(), "b");
4695        assert!(increment_utf8("\u{FDCF}").is_none());
4696
4697        // Test private use area limit (>= 0x110000)
4698        assert_eq!(increment_utf8("a\u{10FFFF}").unwrap(), "b");
4699        assert!(increment_utf8("\u{10FFFF}").is_none()); // Can't increment past max valid codepoint
4700    }
4701
4702    /// Creates a setup for chunk pruning, modeling a utf8 column "s1"
4703    /// with 5 different containers (e.g. RowGroups). They have [min,
4704    /// max]:
4705    /// s1 ["A", "Z"]
4706    /// s1 ["A", "L"]
4707    /// s1 ["N", "Z"]
4708    /// s1 [NULL, NULL]
4709    /// s1 ["A", NULL]
4710    /// s1 ["", "A"]
4711    /// s1 ["", ""]
4712    /// s1 ["AB", "A\u{10ffff}"]
4713    /// s1 ["A\u{10ffff}\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]
4714    fn utf8_setup() -> (SchemaRef, TestStatistics) {
4715        let schema = Arc::new(Schema::new(vec![Field::new("s1", DataType::Utf8, true)]));
4716
4717        let statistics = TestStatistics::new().with(
4718            "s1",
4719            ContainerStats::new_utf8(
4720                vec![
4721                    Some("A"),
4722                    Some("A"),
4723                    Some("N"),
4724                    Some("M"),
4725                    None,
4726                    Some("A"),
4727                    Some(""),
4728                    Some(""),
4729                    Some("AB"),
4730                    Some("A\u{10ffff}\u{10ffff}"),
4731                ], // min
4732                vec![
4733                    Some("Z"),
4734                    Some("L"),
4735                    Some("Z"),
4736                    Some("M"),
4737                    None,
4738                    None,
4739                    Some("A"),
4740                    Some(""),
4741                    Some("A\u{10ffff}\u{10ffff}\u{10ffff}"),
4742                    Some("A\u{10ffff}\u{10ffff}"),
4743                ], // max
4744            ),
4745        );
4746        (schema, statistics)
4747    }
4748
4749    #[test]
4750    fn prune_utf8_eq() {
4751        let (schema, statistics) = utf8_setup();
4752
4753        let expr = col("s1").eq(lit("A"));
4754        #[rustfmt::skip]
4755        let expected_ret = &[
4756            // s1 ["A", "Z"] ==> some rows could pass (must keep)
4757            true,
4758            // s1 ["A", "L"] ==> some rows could pass (must keep)
4759            true,
4760            // s1 ["N", "Z"] ==> no rows can pass (not keep)
4761            false,
4762            // s1 ["M", "M"] ==> no rows can pass (not keep)
4763            false,
4764            // s1 [NULL, NULL]  ==> unknown (must keep)
4765            true,
4766            // s1 ["A", NULL]  ==> unknown (must keep)
4767            true,
4768            // s1 ["", "A"]  ==> some rows could pass (must keep)
4769            true,
4770            // s1 ["", ""]  ==> no rows can pass (not keep)
4771            false,
4772            // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"]  ==> no rows can pass (not keep)
4773            false,
4774            // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]  ==> no rows can pass (not keep)
4775            false,
4776        ];
4777        prune_with_expr(expr, &schema, &statistics, expected_ret);
4778
4779        let expr = col("s1").eq(lit(""));
4780        #[rustfmt::skip]
4781        let expected_ret = &[
4782            // s1 ["A", "Z"] ==> no rows can pass (not keep)
4783            false,
4784            // s1 ["A", "L"] ==> no rows can pass (not keep)
4785            false,
4786            // s1 ["N", "Z"] ==> no rows can pass (not keep)
4787            false,
4788            // s1 ["M", "M"] ==> no rows can pass (not keep)
4789            false,
4790            // s1 [NULL, NULL]  ==> unknown (must keep)
4791            true,
4792            // s1 ["A", NULL]  ==> no rows can pass (not keep)
4793            false,
4794            // s1 ["", "A"]  ==> some rows could pass (must keep)
4795            true,
4796            // s1 ["", ""]  ==> all rows must pass (must keep)
4797            true,
4798            // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"]  ==> no rows can pass (not keep)
4799            false,
4800            // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]  ==> no rows can pass (not keep)
4801            false,
4802        ];
4803        prune_with_expr(expr, &schema, &statistics, expected_ret);
4804    }
4805
4806    #[test]
4807    fn prune_utf8_not_eq() {
4808        let (schema, statistics) = utf8_setup();
4809
4810        let expr = col("s1").not_eq(lit("A"));
4811        #[rustfmt::skip]
4812        let expected_ret = &[
4813            // s1 ["A", "Z"] ==> some rows could pass (must keep)
4814            true,
4815            // s1 ["A", "L"] ==> some rows could pass (must keep)
4816            true,
4817            // s1 ["N", "Z"] ==> all rows must pass (must keep)
4818            true,
4819            // s1 ["M", "M"] ==> all rows must pass (must keep)
4820            true,
4821            // s1 [NULL, NULL]  ==> unknown (must keep)
4822            true,
4823            // s1 ["A", NULL]  ==> unknown (must keep)
4824            true,
4825            // s1 ["", "A"]  ==> some rows could pass (must keep)
4826            true,
4827            // s1 ["", ""]  ==> all rows must pass (must keep)
4828            true,
4829            // s1 ["AB", "A\u{10ffff}\u{10ffff}"]  ==> all rows must pass (must keep)
4830            true,
4831            // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]  ==> all rows must pass (must keep)
4832            true,
4833        ];
4834        prune_with_expr(expr, &schema, &statistics, expected_ret);
4835
4836        let expr = col("s1").not_eq(lit(""));
4837        #[rustfmt::skip]
4838        let expected_ret = &[
4839            // s1 ["A", "Z"] ==> all rows must pass (must keep)
4840            true,
4841            // s1 ["A", "L"] ==> all rows must pass (must keep)
4842            true,
4843            // s1 ["N", "Z"] ==> all rows must pass (must keep)
4844            true,
4845            // s1 ["M", "M"] ==> all rows must pass (must keep)
4846            true,
4847            // s1 [NULL, NULL]  ==> unknown (must keep)
4848            true,
4849            // s1 ["A", NULL]  ==> unknown (must keep)
4850            true,
4851            // s1 ["", "A"]  ==> some rows could pass (must keep)
4852            true,
4853            // s1 ["", ""]  ==> no rows can pass (not keep)
4854            false,
4855            // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"]  ==> all rows must pass (must keep)
4856            true,
4857            // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]  ==> all rows must pass (must keep)
4858            true,
4859        ];
4860        prune_with_expr(expr, &schema, &statistics, expected_ret);
4861    }
4862
4863    #[test]
4864    fn prune_utf8_like_one() {
4865        let (schema, statistics) = utf8_setup();
4866
4867        let expr = col("s1").like(lit("A_"));
4868        #[rustfmt::skip]
4869        let expected_ret = &[
4870            // s1 ["A", "Z"] ==> some rows could pass (must keep)
4871            true,
4872            // s1 ["A", "L"] ==> some rows could pass (must keep)
4873            true,
4874            // s1 ["N", "Z"] ==> no rows can pass (not keep)
4875            false,
4876            // s1 ["M", "M"] ==> no rows can pass (not keep)
4877            false,
4878            // s1 [NULL, NULL]  ==> unknown (must keep)
4879            true,
4880            // s1 ["A", NULL]  ==> unknown (must keep)
4881            true,
4882            // s1 ["", "A"]  ==> some rows could pass (must keep)
4883            true,
4884            // s1 ["", ""]  ==> no rows can pass (not keep)
4885            false,
4886            // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"]  ==> some rows could pass (must keep)
4887            true,
4888            // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]  ==> some rows could pass (must keep)
4889            true,
4890        ];
4891        prune_with_expr(expr, &schema, &statistics, expected_ret);
4892
4893        let expr = col("s1").like(lit("_A_"));
4894        #[rustfmt::skip]
4895        let expected_ret = &[
4896            // s1 ["A", "Z"] ==> some rows could pass (must keep)
4897            true,
4898            // s1 ["A", "L"] ==> some rows could pass (must keep)
4899            true,
4900            // s1 ["N", "Z"] ==> some rows could pass (must keep)
4901            true,
4902            // s1 ["M", "M"] ==> some rows could pass (must keep)
4903            true,
4904            // s1 [NULL, NULL]  ==> unknown (must keep)
4905            true,
4906            // s1 ["A", NULL]  ==> unknown (must keep)
4907            true,
4908            // s1 ["", "A"]  ==> some rows could pass (must keep)
4909            true,
4910            // s1 ["", ""]  ==> some rows could pass (must keep)
4911            true,
4912            // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"]  ==> some rows could pass (must keep)
4913            true,
4914            // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]  ==> some rows could pass (must keep)
4915            true,
4916        ];
4917        prune_with_expr(expr, &schema, &statistics, expected_ret);
4918
4919        let expr = col("s1").like(lit("_"));
4920        #[rustfmt::skip]
4921        let expected_ret = &[
4922            // s1 ["A", "Z"] ==> all rows must pass (must keep)
4923            true,
4924            // s1 ["A", "L"] ==> all rows must pass (must keep)
4925            true,
4926            // s1 ["N", "Z"] ==> all rows must pass (must keep)
4927            true,
4928            // s1 ["M", "M"] ==> all rows must pass (must keep)
4929            true,
4930            // s1 [NULL, NULL]  ==> unknown (must keep)
4931            true,
4932            // s1 ["A", NULL]  ==> unknown (must keep)
4933            true,
4934            // s1 ["", "A"]  ==> all rows must pass (must keep)
4935            true,
4936            // s1 ["", ""]  ==> all rows must pass (must keep)
4937            true,
4938            // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"]  ==> all rows must pass (must keep)
4939            true,
4940            // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]  ==> all rows must pass (must keep)
4941            true,
4942        ];
4943        prune_with_expr(expr, &schema, &statistics, expected_ret);
4944
4945        let expr = col("s1").like(lit(""));
4946        #[rustfmt::skip]
4947        let expected_ret = &[
4948            // s1 ["A", "Z"] ==> no rows can pass (not keep)
4949            false,
4950            // s1 ["A", "L"] ==> no rows can pass (not keep)
4951            false,
4952            // s1 ["N", "Z"] ==> no rows can pass (not keep)
4953            false,
4954            // s1 ["M", "M"] ==> no rows can pass (not keep)
4955            false,
4956            // s1 [NULL, NULL]  ==> unknown (must keep)
4957            true,
4958            // s1 ["A", NULL]  ==> no rows can pass (not keep)
4959            false,
4960            // s1 ["", "A"]  ==> some rows could pass (must keep)
4961            true,
4962            // s1 ["", ""]  ==> all rows must pass (must keep)
4963            true,
4964            // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"]  ==> no rows can pass (not keep)
4965            false,
4966            // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]  ==> no rows can pass (not keep)
4967            false,
4968        ];
4969        prune_with_expr(expr, &schema, &statistics, expected_ret);
4970    }
4971
4972    #[test]
4973    fn prune_utf8_like_many() {
4974        let (schema, statistics) = utf8_setup();
4975
4976        let expr = col("s1").like(lit("A%"));
4977        #[rustfmt::skip]
4978        let expected_ret = &[
4979            // s1 ["A", "Z"] ==> some rows could pass (must keep)
4980            true,
4981            // s1 ["A", "L"] ==> some rows could pass (must keep)
4982            true,
4983            // s1 ["N", "Z"] ==> no rows can pass (not keep)
4984            false,
4985            // s1 ["M", "M"] ==> no rows can pass (not keep)
4986            false,
4987            // s1 [NULL, NULL]  ==> unknown (must keep)
4988            true,
4989            // s1 ["A", NULL]  ==> unknown (must keep)
4990            true,
4991            // s1 ["", "A"]  ==> some rows could pass (must keep)
4992            true,
4993            // s1 ["", ""]  ==> no rows can pass (not keep)
4994            false,
4995            // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"]  ==> some rows could pass (must keep)
4996            true,
4997            // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]  ==> some rows could pass (must keep)
4998            true,
4999        ];
5000        prune_with_expr(expr, &schema, &statistics, expected_ret);
5001
5002        let expr = col("s1").like(lit("%A%"));
5003        #[rustfmt::skip]
5004        let expected_ret = &[
5005            // s1 ["A", "Z"] ==> some rows could pass (must keep)
5006            true,
5007            // s1 ["A", "L"] ==> some rows could pass (must keep)
5008            true,
5009            // s1 ["N", "Z"] ==> some rows could pass (must keep)
5010            true,
5011            // s1 ["M", "M"] ==> some rows could pass (must keep)
5012            true,
5013            // s1 [NULL, NULL]  ==> unknown (must keep)
5014            true,
5015            // s1 ["A", NULL]  ==> unknown (must keep)
5016            true,
5017            // s1 ["", "A"]  ==> some rows could pass (must keep)
5018            true,
5019            // s1 ["", ""]  ==> some rows could pass (must keep)
5020            true,
5021            // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"]  ==> some rows could pass (must keep)
5022            true,
5023            // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]  ==> some rows could pass (must keep)
5024            true,
5025        ];
5026        prune_with_expr(expr, &schema, &statistics, expected_ret);
5027
5028        let expr = col("s1").like(lit("%"));
5029        #[rustfmt::skip]
5030        let expected_ret = &[
5031            // s1 ["A", "Z"] ==> all rows must pass (must keep)
5032            true,
5033            // s1 ["A", "L"] ==> all rows must pass (must keep)
5034            true,
5035            // s1 ["N", "Z"] ==> all rows must pass (must keep)
5036            true,
5037            // s1 ["M", "M"] ==> all rows must pass (must keep)
5038            true,
5039            // s1 [NULL, NULL]  ==> unknown (must keep)
5040            true,
5041            // s1 ["A", NULL]  ==> unknown (must keep)
5042            true,
5043            // s1 ["", "A"]  ==> all rows must pass (must keep)
5044            true,
5045            // s1 ["", ""]  ==> all rows must pass (must keep)
5046            true,
5047            // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"]  ==> all rows must pass (must keep)
5048            true,
5049            // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]  ==> all rows must pass (must keep)
5050            true,
5051        ];
5052        prune_with_expr(expr, &schema, &statistics, expected_ret);
5053
5054        let expr = col("s1").like(lit(""));
5055        #[rustfmt::skip]
5056        let expected_ret = &[
5057            // s1 ["A", "Z"] ==> no rows can pass (not keep)
5058            false,
5059            // s1 ["A", "L"] ==> no rows can pass (not keep)
5060            false,
5061            // s1 ["N", "Z"] ==> no rows can pass (not keep)
5062            false,
5063            // s1 ["M", "M"] ==> no rows can pass (not keep)
5064            false,
5065            // s1 [NULL, NULL]  ==> unknown (must keep)
5066            true,
5067            // s1 ["A", NULL]  ==> no rows can pass (not keep)
5068            false,
5069            // s1 ["", "A"]  ==> some rows could pass (must keep)
5070            true,
5071            // s1 ["", ""]  ==> all rows must pass (must keep)
5072            true,
5073            // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"]  ==> no rows can pass (not keep)
5074            false,
5075            // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]  ==> no rows can pass (not keep)
5076            false,
5077        ];
5078        prune_with_expr(expr, &schema, &statistics, expected_ret);
5079    }
5080
5081    // `build_like_match()` must honor `\` escapes when scanning the pattern for
5082    // wildcards.
5083    #[test]
5084    fn prune_utf8_like_escaped_chars() {
5085        let schema = Arc::new(Schema::new(vec![Field::new("s1", DataType::Utf8, true)]));
5086        let statistics = TestStatistics::new().with(
5087            "s1",
5088            ContainerStats::new_utf8(
5089                vec![
5090                    Some("foo_aaa"),
5091                    Some(r#"foo\aaa"#),
5092                    Some("foo"),
5093                    Some("bar"),
5094                    Some("foo%aaa"),
5095                    Some("%foo_aaa"),
5096                ], // min
5097                vec![
5098                    Some("foo_zzz"),
5099                    Some(r#"foo\zzz"#),
5100                    Some("foozzz"),
5101                    Some("baz"),
5102                    Some("foo%zzz"),
5103                    Some("%foo_zzz"),
5104                ], // max
5105            ),
5106        );
5107
5108        let expr = col("s1").like(lit(r#"foo\_%"#));
5109        #[rustfmt::skip]
5110        let expected_ret = &[
5111            // s1 ["foo_aaa", "foo_zzz"] => every value starts with literal
5112            // "foo_" and matches the pattern; must keep.
5113            true,
5114            // s1 ["foo\aaa", "foo\zzz"] => no rows can pass (not keep)
5115            false,
5116            // s1 ["foo", "foozzz"] => stats don't prove "foo_" is or isn't in
5117            // range; must conservatively keep.
5118            true,
5119            // s1 ["bar", "baz"] => no rows can pass (not keep)
5120            false,
5121            // s1 ["foo%aaa", "foo%zzz"] => no rows can pass (not keep)
5122            false,
5123            // s1 ["%foo_aaa", "%foo_zzz"] => no rows can pass (not keep)
5124            false,
5125        ];
5126        prune_with_expr(expr, &schema, &statistics, expected_ret);
5127
5128        let expr = col("s1").like(lit(r#"foo\\%"#));
5129        #[rustfmt::skip]
5130        let expected_ret = &[
5131            // s1 ["foo_aaa", "foo_zzz"] => no rows can pass (not keep)
5132            false,
5133            // s1 ["foo\aaa", "foo\zzz"] => every value starts with literal
5134            // "foo\" and matches the pattern; must keep.
5135            true,
5136            // s1 ["foo", "foozzz"] => stats don't prove "foo\" is or isn't in
5137            // range; must conservatively keep.
5138            true,
5139            // s1 ["bar", "baz"] => no rows can pass (not keep)
5140            false,
5141            // s1 ["foo%aaa", "foo%zzz"] => no rows can pass (not keep)
5142            false,
5143            // s1 ["%foo_aaa", "%foo_zzz"] => no rows can pass (not keep)
5144            false,
5145        ];
5146        prune_with_expr(expr, &schema, &statistics, expected_ret);
5147
5148        let expr = col("s1").like(lit(r#"foo\%%"#));
5149        #[rustfmt::skip]
5150        let expected_ret = &[
5151            // s1 ["foo_aaa", "foo_zzz"] => no rows can pass (not keep)
5152            false,
5153            // s1 ["foo\aaa", "foo\zzz"] => no rows can pass (not keep)
5154            false,
5155            // s1 ["foo", "foozzz"] => range straddles "foo%"; must keep.
5156            true,
5157            // s1 ["bar", "baz"] => no rows can pass (not keep)
5158            false,
5159            // s1 ["foo%aaa", "foo%zzz"] => every value starts with literal
5160            // "foo%" and matches the pattern; must keep.
5161            true,
5162            // s1 ["%foo_aaa", "%foo_zzz"] => no rows can pass (not keep)
5163            false,
5164        ];
5165        prune_with_expr(expr, &schema, &statistics, expected_ret);
5166
5167        // No wildcard after escapes: pattern reduces to an equality check on
5168        // the literal "foo_".
5169        let expr = col("s1").like(lit(r#"foo\_"#));
5170        #[rustfmt::skip]
5171        let expected_ret = &[
5172            // s1 ["foo_aaa", "foo_zzz"] => no rows can pass (not keep)
5173            false,
5174            // s1 ["foo\aaa", "foo\zzz"] => no rows can pass (not keep)
5175            false,
5176            // s1 ["foo", "foozzz"] => "foo_" is within the range; must keep.
5177            true,
5178            // s1 ["bar", "baz"] => no rows can pass (not keep)
5179            false,
5180            // s1 ["foo%aaa", "foo%zzz"] => no rows can pass (not keep)
5181            false,
5182            // s1 ["%foo_aaa", "%foo_zzz"] => no rows can pass (not keep)
5183            false,
5184        ];
5185        prune_with_expr(expr, &schema, &statistics, expected_ret);
5186
5187        // Leading escaped `%`: prefix is "%foo" (non-empty), so the guard
5188        // for "all wildcards" must NOT bail out here.
5189        let expr = col("s1").like(lit(r#"\%foo%"#));
5190        #[rustfmt::skip]
5191        let expected_ret = &[
5192            // s1 ["foo_aaa", "foo_zzz"] => no rows can pass (not keep)
5193            false,
5194            // s1 ["foo\aaa", "foo\zzz"] => no rows can pass (not keep)
5195            false,
5196            // s1 ["foo", "foozzz"] => no rows can pass (not keep)
5197            false,
5198            // s1 ["bar", "baz"] => no rows can pass (not keep)
5199            false,
5200            // s1 ["foo%aaa", "foo%zzz"] => no rows can pass (not keep)
5201            false,
5202            // s1 ["%foo_aaa", "%foo_zzz"] => every value starts with literal
5203            // "%foo" and matches the pattern; must keep.
5204            true,
5205        ];
5206        prune_with_expr(expr, &schema, &statistics, expected_ret);
5207
5208        // Two escaped wildcards, no real wildcard: equality on "foo%_".
5209        let expr = col("s1").like(lit(r#"foo\%\_"#));
5210        #[rustfmt::skip]
5211        let expected_ret = &[
5212            // s1 ["foo_aaa", "foo_zzz"] => no rows can pass (not keep)
5213            false,
5214            // s1 ["foo\aaa", "foo\zzz"] => no rows can pass (not keep)
5215            false,
5216            // s1 ["foo", "foozzz"] => "foo%_" is within the range; must keep.
5217            true,
5218            // s1 ["bar", "baz"] => no rows can pass (not keep)
5219            false,
5220            // s1 ["foo%aaa", "foo%zzz"] => no rows can pass (not keep)
5221            false,
5222            // s1 ["%foo_aaa", "%foo_zzz"] => no rows can pass (not keep)
5223            false,
5224        ];
5225        prune_with_expr(expr, &schema, &statistics, expected_ret);
5226
5227        // Escaped backslash followed by more literal chars before the
5228        // wildcard: prefix is "foo\bar".
5229        let expr = col("s1").like(lit(r#"foo\\bar%"#));
5230        #[rustfmt::skip]
5231        let expected_ret = &[
5232            // s1 ["foo_aaa", "foo_zzz"] => no rows can pass (not keep)
5233            false,
5234            // s1 ["foo\aaa", "foo\zzz"] => range straddles "foo\bar"; must
5235            // keep.
5236            true,
5237            // s1 ["foo", "foozzz"] => range straddles "foo\bar"; must keep.
5238            true,
5239            // s1 ["bar", "baz"] => no rows can pass (not keep)
5240            false,
5241            // s1 ["foo%aaa", "foo%zzz"] => no rows can pass (not keep)
5242            false,
5243            // s1 ["%foo_aaa", "%foo_zzz"] => no rows can pass (not keep)
5244            false,
5245        ];
5246        prune_with_expr(expr, &schema, &statistics, expected_ret);
5247    }
5248
5249    #[test]
5250    fn prune_utf8_not_like_one() {
5251        let (schema, statistics) = utf8_setup();
5252
5253        let expr = col("s1").not_like(lit("A\u{10ffff}_"));
5254        #[rustfmt::skip]
5255        let expected_ret = &[
5256            // s1 ["A", "Z"] ==> some rows could pass (must keep)
5257            true,
5258            // s1 ["A", "L"] ==> some rows could pass (must keep)
5259            true,
5260            // s1 ["N", "Z"] ==> some rows could pass (must keep)
5261            true,
5262            // s1 ["M", "M"] ==> some rows could pass (must keep)
5263            true,
5264            // s1 [NULL, NULL]  ==> unknown (must keep)
5265            true,
5266            // s1 ["A", NULL]  ==> some rows could pass (must keep)
5267            true,
5268            // s1 ["", "A"]  ==> some rows could pass (must keep)
5269            true,
5270            // s1 ["", ""]  ==> some rows could pass (must keep)
5271            true,
5272            // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"]  ==> some rows could pass (must keep)
5273            true,
5274            // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]  ==> no row match. (min, max) maybe truncate
5275            // original (min, max) maybe ("A\u{10ffff}\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}\u{10ffff}\u{10ffff}")
5276            true,
5277        ];
5278        prune_with_expr(expr, &schema, &statistics, expected_ret);
5279    }
5280
5281    #[test]
5282    fn prune_utf8_not_like_many() {
5283        let (schema, statistics) = utf8_setup();
5284
5285        let expr = col("s1").not_like(lit("A\u{10ffff}%"));
5286        #[rustfmt::skip]
5287        let expected_ret = &[
5288            // s1 ["A", "Z"] ==> some rows could pass (must keep)
5289            true,
5290            // s1 ["A", "L"] ==> some rows could pass (must keep)
5291            true,
5292            // s1 ["N", "Z"] ==> some rows could pass (must keep)
5293            true,
5294            // s1 ["M", "M"] ==> some rows could pass (must keep)
5295            true,
5296            // s1 [NULL, NULL]  ==> unknown (must keep)
5297            true,
5298            // s1 ["A", NULL]  ==> some rows could pass (must keep)
5299            true,
5300            // s1 ["", "A"]  ==> some rows could pass (must keep)
5301            true,
5302            // s1 ["", ""]  ==> some rows could pass (must keep)
5303            true,
5304            // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"]  ==> some rows could pass (must keep)
5305            true,
5306            // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]  ==> no row match
5307            false,
5308        ];
5309        prune_with_expr(expr, &schema, &statistics, expected_ret);
5310
5311        let expr = col("s1").not_like(lit("A\u{10ffff}%\u{10ffff}"));
5312        #[rustfmt::skip]
5313        let expected_ret = &[
5314            // s1 ["A", "Z"] ==> some rows could pass (must keep)
5315            true,
5316            // s1 ["A", "L"] ==> some rows could pass (must keep)
5317            true,
5318            // s1 ["N", "Z"] ==> some rows could pass (must keep)
5319            true,
5320            // s1 ["M", "M"] ==> some rows could pass (must keep)
5321            true,
5322            // s1 [NULL, NULL]  ==> unknown (must keep)
5323            true,
5324            // s1 ["A", NULL]  ==> some rows could pass (must keep)
5325            true,
5326            // s1 ["", "A"]  ==> some rows could pass (must keep)
5327            true,
5328            // s1 ["", ""]  ==> some rows could pass (must keep)
5329            true,
5330            // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"]  ==> some rows could pass (must keep)
5331            true,
5332            // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]  ==> some rows could pass (must keep)
5333            true,
5334        ];
5335        prune_with_expr(expr, &schema, &statistics, expected_ret);
5336
5337        let expr = col("s1").not_like(lit("A\u{10ffff}%\u{10ffff}_"));
5338        #[rustfmt::skip]
5339        let expected_ret = &[
5340            // s1 ["A", "Z"] ==> some rows could pass (must keep)
5341            true,
5342            // s1 ["A", "L"] ==> some rows could pass (must keep)
5343            true,
5344            // s1 ["N", "Z"] ==> some rows could pass (must keep)
5345            true,
5346            // s1 ["M", "M"] ==> some rows could pass (must keep)
5347            true,
5348            // s1 [NULL, NULL]  ==> unknown (must keep)
5349            true,
5350            // s1 ["A", NULL]  ==> some rows could pass (must keep)
5351            true,
5352            // s1 ["", "A"]  ==> some rows could pass (must keep)
5353            true,
5354            // s1 ["", ""]  ==> some rows could pass (must keep)
5355            true,
5356            // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"]  ==> some rows could pass (must keep)
5357            true,
5358            // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]  ==> some rows could pass (must keep)
5359            true,
5360        ];
5361        prune_with_expr(expr, &schema, &statistics, expected_ret);
5362
5363        let expr = col("s1").not_like(lit("A\\%%"));
5364        let statistics = TestStatistics::new().with(
5365            "s1",
5366            ContainerStats::new_utf8(
5367                vec![Some("A%a"), Some("A")],
5368                vec![Some("A%c"), Some("A")],
5369            ),
5370        );
5371        let expected_ret = &[false, true];
5372        prune_with_expr(expr, &schema, &statistics, expected_ret);
5373    }
5374
5375    #[test]
5376    fn test_rewrite_expr_to_prunable() {
5377        let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
5378        let df_schema = DFSchema::try_from(schema.clone()).unwrap();
5379
5380        // column op lit
5381        let left_input = col("a");
5382        let left_input = logical2physical(&left_input, &schema);
5383        let right_input = lit(ScalarValue::Int32(Some(12)));
5384        let right_input = logical2physical(&right_input, &schema);
5385        let (result_left, _, result_right) = rewrite_expr_to_prunable(
5386            &left_input,
5387            Operator::Eq,
5388            &right_input,
5389            df_schema.clone(),
5390        )
5391        .unwrap();
5392        assert_eq!(result_left.to_string(), left_input.to_string());
5393        assert_eq!(result_right.to_string(), right_input.to_string());
5394
5395        // cast op lit
5396        let left_input = cast(col("a"), DataType::Decimal128(20, 3));
5397        let left_input = logical2physical(&left_input, &schema);
5398        let right_input = lit(ScalarValue::Decimal128(Some(12), 20, 3));
5399        let right_input = logical2physical(&right_input, &schema);
5400        let (result_left, _, result_right) = rewrite_expr_to_prunable(
5401            &left_input,
5402            Operator::Gt,
5403            &right_input,
5404            df_schema.clone(),
5405        )
5406        .unwrap();
5407        assert_eq!(result_left.to_string(), left_input.to_string());
5408        assert_eq!(result_right.to_string(), right_input.to_string());
5409
5410        // try_cast op lit
5411        let left_input = try_cast(col("a"), DataType::Int64);
5412        let left_input = logical2physical(&left_input, &schema);
5413        let right_input = lit(ScalarValue::Int64(Some(12)));
5414        let right_input = logical2physical(&right_input, &schema);
5415        let (result_left, _, result_right) =
5416            rewrite_expr_to_prunable(&left_input, Operator::Gt, &right_input, df_schema)
5417                .unwrap();
5418        assert_eq!(result_left.to_string(), left_input.to_string());
5419        assert_eq!(result_right.to_string(), right_input.to_string());
5420
5421        // TODO: add test for other case and op
5422    }
5423
5424    #[test]
5425    fn test_rewrite_expr_to_prunable_custom_unhandled_hook() {
5426        struct CustomUnhandledHook;
5427
5428        impl UnhandledPredicateHook for CustomUnhandledHook {
5429            /// This handles an arbitrary case of a column that doesn't exist in the schema
5430            /// by renaming it to yet another column that doesn't exist in the schema
5431            /// (the transformation is arbitrary, the point is that it can do whatever it wants)
5432            fn handle(&self, _expr: &Arc<dyn PhysicalExpr>) -> Arc<dyn PhysicalExpr> {
5433                Arc::new(phys_expr::Literal::new(ScalarValue::Int32(Some(42))))
5434            }
5435        }
5436
5437        let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
5438        let schema_with_b = Schema::new(vec![
5439            Field::new("a", DataType::Int32, true),
5440            Field::new("b", DataType::Int32, true),
5441        ]);
5442
5443        let rewriter = PredicateRewriter::new()
5444            .with_unhandled_hook(Arc::new(CustomUnhandledHook {}));
5445
5446        let transform_expr = |expr| {
5447            let expr = logical2physical(&expr, &schema_with_b);
5448            rewriter.rewrite_predicate_to_statistics_predicate(&expr, &schema)
5449        };
5450
5451        // transform an arbitrary valid expression that we know is handled
5452        let known_expression = col("a").eq(lit(12));
5453        let known_expression_transformed = PredicateRewriter::new()
5454            .rewrite_predicate_to_statistics_predicate(
5455                &logical2physical(&known_expression, &schema),
5456                &schema,
5457            );
5458
5459        // an expression referencing an unknown column (that is not in the schema) gets passed to the hook
5460        let input = col("b").eq(lit(12));
5461        let expected = logical2physical(&lit(42), &schema);
5462        let transformed = transform_expr(input.clone());
5463        assert_eq!(transformed.to_string(), expected.to_string());
5464
5465        // more complex case with unknown column
5466        let input = known_expression.clone().and(input.clone());
5467        let expected = phys_expr::BinaryExpr::new(
5468            Arc::<dyn PhysicalExpr>::clone(&known_expression_transformed),
5469            Operator::And,
5470            logical2physical(&lit(42), &schema),
5471        );
5472        let transformed = transform_expr(input.clone());
5473        assert_eq!(transformed.to_string(), expected.to_string());
5474
5475        // an unknown expression gets passed to the hook
5476        let input = array_has(make_array(vec![lit(1)]), col("a"));
5477        let expected = logical2physical(&lit(42), &schema);
5478        let transformed = transform_expr(input.clone());
5479        assert_eq!(transformed.to_string(), expected.to_string());
5480
5481        // more complex case with unknown expression
5482        let input = known_expression.and(input);
5483        let expected = phys_expr::BinaryExpr::new(
5484            Arc::<dyn PhysicalExpr>::clone(&known_expression_transformed),
5485            Operator::And,
5486            logical2physical(&lit(42), &schema),
5487        );
5488        let transformed = transform_expr(input.clone());
5489        assert_eq!(transformed.to_string(), expected.to_string());
5490    }
5491
5492    #[test]
5493    fn test_rewrite_expr_to_prunable_error() {
5494        // cast string value to numeric value
5495        // this cast is not supported
5496        let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]);
5497        let df_schema = DFSchema::try_from(schema.clone()).unwrap();
5498        let left_input = cast(col("a"), DataType::Int64);
5499        let left_input = logical2physical(&left_input, &schema);
5500        let right_input = lit(ScalarValue::Int64(Some(12)));
5501        let right_input = logical2physical(&right_input, &schema);
5502        let result = rewrite_expr_to_prunable(
5503            &left_input,
5504            Operator::Gt,
5505            &right_input,
5506            df_schema.clone(),
5507        );
5508        assert!(result.is_err());
5509
5510        // other expr
5511        let left_input = is_null(col("a"));
5512        let left_input = logical2physical(&left_input, &schema);
5513        let right_input = lit(ScalarValue::Int64(Some(12)));
5514        let right_input = logical2physical(&right_input, &schema);
5515        let result =
5516            rewrite_expr_to_prunable(&left_input, Operator::Gt, &right_input, df_schema);
5517        assert!(result.is_err());
5518        // TODO: add other negative test for other case and op
5519    }
5520
5521    #[test]
5522    fn prune_with_contained_one_column() {
5523        let schema = Arc::new(Schema::new(vec![Field::new("s1", DataType::Utf8, true)]));
5524
5525        // Model having information like a bloom filter for s1
5526        let statistics = TestStatistics::new()
5527            .with_contained(
5528                "s1",
5529                [ScalarValue::from("foo")],
5530                [
5531                    // container 0 known to only contain "foo"",
5532                    Some(true),
5533                    // container 1 known to not contain "foo"
5534                    Some(false),
5535                    // container 2 unknown about "foo"
5536                    None,
5537                    // container 3 known to only contain "foo"
5538                    Some(true),
5539                    // container 4 known to not contain "foo"
5540                    Some(false),
5541                    // container 5 unknown about "foo"
5542                    None,
5543                    // container 6 known to only contain "foo"
5544                    Some(true),
5545                    // container 7 known to not contain "foo"
5546                    Some(false),
5547                    // container 8 unknown about "foo"
5548                    None,
5549                ],
5550            )
5551            .with_contained(
5552                "s1",
5553                [ScalarValue::from("bar")],
5554                [
5555                    // containers 0,1,2 known to only contain "bar"
5556                    Some(true),
5557                    Some(true),
5558                    Some(true),
5559                    // container 3,4,5 known to not contain "bar"
5560                    Some(false),
5561                    Some(false),
5562                    Some(false),
5563                    // container 6,7,8 unknown about "bar"
5564                    None,
5565                    None,
5566                    None,
5567                ],
5568            )
5569            .with_contained(
5570                // the way the tests are setup, this data is
5571                // consulted if the "foo" and "bar" are being checked at the same time
5572                "s1",
5573                [ScalarValue::from("foo"), ScalarValue::from("bar")],
5574                [
5575                    // container 0,1,2 unknown about ("foo, "bar")
5576                    None,
5577                    None,
5578                    None,
5579                    // container 3,4,5 known to contain only either "foo" and "bar"
5580                    Some(true),
5581                    Some(true),
5582                    Some(true),
5583                    // container 6,7,8  known to contain  neither "foo" and "bar"
5584                    Some(false),
5585                    Some(false),
5586                    Some(false),
5587                ],
5588            );
5589
5590        // s1 = 'foo'
5591        prune_with_expr(
5592            col("s1").eq(lit("foo")),
5593            &schema,
5594            &statistics,
5595            // rule out containers ('false) where we know foo is not present
5596            &[true, false, true, true, false, true, true, false, true],
5597        );
5598
5599        // s1 = 'bar'
5600        prune_with_expr(
5601            col("s1").eq(lit("bar")),
5602            &schema,
5603            &statistics,
5604            // rule out containers where we know bar is not present
5605            &[true, true, true, false, false, false, true, true, true],
5606        );
5607
5608        // s1 = 'baz' (unknown value)
5609        prune_with_expr(
5610            col("s1").eq(lit("baz")),
5611            &schema,
5612            &statistics,
5613            // can't rule out anything
5614            &[true, true, true, true, true, true, true, true, true],
5615        );
5616
5617        // s1 = 'foo' AND s1 = 'bar'
5618        prune_with_expr(
5619            col("s1").eq(lit("foo")).and(col("s1").eq(lit("bar"))),
5620            &schema,
5621            &statistics,
5622            // logically this predicate can't possibly be true (the column can't
5623            // take on both values) but we could rule it out if the stats tell
5624            // us that both values are not present
5625            &[true, true, true, true, true, true, true, true, true],
5626        );
5627
5628        // s1 = 'foo' OR s1 = 'bar'
5629        prune_with_expr(
5630            col("s1").eq(lit("foo")).or(col("s1").eq(lit("bar"))),
5631            &schema,
5632            &statistics,
5633            // can rule out containers that we know contain neither foo nor bar
5634            &[true, true, true, true, true, true, false, false, false],
5635        );
5636
5637        // s1 = 'foo' OR s1 = 'baz'
5638        prune_with_expr(
5639            col("s1").eq(lit("foo")).or(col("s1").eq(lit("baz"))),
5640            &schema,
5641            &statistics,
5642            // can't rule out anything container
5643            &[true, true, true, true, true, true, true, true, true],
5644        );
5645
5646        // s1 = 'foo' OR s1 = 'bar' OR s1 = 'baz'
5647        prune_with_expr(
5648            col("s1")
5649                .eq(lit("foo"))
5650                .or(col("s1").eq(lit("bar")))
5651                .or(col("s1").eq(lit("baz"))),
5652            &schema,
5653            &statistics,
5654            // can rule out any containers based on knowledge of s1 and `foo`,
5655            // `bar` and (`foo`, `bar`)
5656            &[true, true, true, true, true, true, true, true, true],
5657        );
5658
5659        // s1 != foo
5660        prune_with_expr(
5661            col("s1").not_eq(lit("foo")),
5662            &schema,
5663            &statistics,
5664            // rule out containers we know for sure only contain foo
5665            &[false, true, true, false, true, true, false, true, true],
5666        );
5667
5668        // s1 != bar
5669        prune_with_expr(
5670            col("s1").not_eq(lit("bar")),
5671            &schema,
5672            &statistics,
5673            // rule out when we know for sure s1 has the value bar
5674            &[false, false, false, true, true, true, true, true, true],
5675        );
5676
5677        // s1 != foo AND s1 != bar
5678        prune_with_expr(
5679            col("s1")
5680                .not_eq(lit("foo"))
5681                .and(col("s1").not_eq(lit("bar"))),
5682            &schema,
5683            &statistics,
5684            // can rule out any container where we know s1 does not have either 'foo' or 'bar'
5685            &[true, true, true, false, false, false, true, true, true],
5686        );
5687
5688        // s1 != foo AND s1 != bar AND s1 != baz
5689        prune_with_expr(
5690            col("s1")
5691                .not_eq(lit("foo"))
5692                .and(col("s1").not_eq(lit("bar")))
5693                .and(col("s1").not_eq(lit("baz"))),
5694            &schema,
5695            &statistics,
5696            // can't rule out any container based on  knowledge of s1,s2
5697            &[true, true, true, true, true, true, true, true, true],
5698        );
5699
5700        // s1 != foo OR s1 != bar
5701        prune_with_expr(
5702            col("s1")
5703                .not_eq(lit("foo"))
5704                .or(col("s1").not_eq(lit("bar"))),
5705            &schema,
5706            &statistics,
5707            // cant' rule out anything based on contains information
5708            &[true, true, true, true, true, true, true, true, true],
5709        );
5710
5711        // s1 != foo OR s1 != bar OR s1 != baz
5712        prune_with_expr(
5713            col("s1")
5714                .not_eq(lit("foo"))
5715                .or(col("s1").not_eq(lit("bar")))
5716                .or(col("s1").not_eq(lit("baz"))),
5717            &schema,
5718            &statistics,
5719            // cant' rule out anything based on contains information
5720            &[true, true, true, true, true, true, true, true, true],
5721        );
5722    }
5723
5724    #[test]
5725    fn prune_with_contained_two_columns() {
5726        let schema = Arc::new(Schema::new(vec![
5727            Field::new("s1", DataType::Utf8, true),
5728            Field::new("s2", DataType::Utf8, true),
5729        ]));
5730
5731        // Model having information like bloom filters for s1 and s2
5732        let statistics = TestStatistics::new()
5733            .with_contained(
5734                "s1",
5735                [ScalarValue::from("foo")],
5736                [
5737                    // container 0, s1 known to only contain "foo"",
5738                    Some(true),
5739                    // container 1, s1 known to not contain "foo"
5740                    Some(false),
5741                    // container 2, s1 unknown about "foo"
5742                    None,
5743                    // container 3, s1 known to only contain "foo"
5744                    Some(true),
5745                    // container 4, s1 known to not contain "foo"
5746                    Some(false),
5747                    // container 5, s1 unknown about "foo"
5748                    None,
5749                    // container 6, s1 known to only contain "foo"
5750                    Some(true),
5751                    // container 7, s1 known to not contain "foo"
5752                    Some(false),
5753                    // container 8, s1 unknown about "foo"
5754                    None,
5755                ],
5756            )
5757            .with_contained(
5758                "s2", // for column s2
5759                [ScalarValue::from("bar")],
5760                [
5761                    // containers 0,1,2 s2 known to only contain "bar"
5762                    Some(true),
5763                    Some(true),
5764                    Some(true),
5765                    // container 3,4,5 s2 known to not contain "bar"
5766                    Some(false),
5767                    Some(false),
5768                    Some(false),
5769                    // container 6,7,8 s2 unknown about "bar"
5770                    None,
5771                    None,
5772                    None,
5773                ],
5774            );
5775
5776        // s1 = 'foo'
5777        prune_with_expr(
5778            col("s1").eq(lit("foo")),
5779            &schema,
5780            &statistics,
5781            // rule out containers where we know s1 is not present
5782            &[true, false, true, true, false, true, true, false, true],
5783        );
5784
5785        // s1 = 'foo' OR s2 = 'bar'
5786        let expr = col("s1").eq(lit("foo")).or(col("s2").eq(lit("bar")));
5787        prune_with_expr(
5788            expr,
5789            &schema,
5790            &statistics,
5791            //  can't rule out any container (would need to prove that s1 != foo AND s2 != bar)
5792            &[true, true, true, true, true, true, true, true, true],
5793        );
5794
5795        // s1 = 'foo' AND s2 != 'bar'
5796        prune_with_expr(
5797            col("s1").eq(lit("foo")).and(col("s2").not_eq(lit("bar"))),
5798            &schema,
5799            &statistics,
5800            // can only rule out container where we know either:
5801            // 1. s1 doesn't have the value 'foo` or
5802            // 2. s2 has only the value of 'bar'
5803            &[false, false, false, true, false, true, true, false, true],
5804        );
5805
5806        // s1 != 'foo' AND s2 != 'bar'
5807        prune_with_expr(
5808            col("s1")
5809                .not_eq(lit("foo"))
5810                .and(col("s2").not_eq(lit("bar"))),
5811            &schema,
5812            &statistics,
5813            // Can  rule out any container where we know either
5814            // 1. s1 has only the value 'foo'
5815            // 2. s2 has only the value 'bar'
5816            &[false, false, false, false, true, true, false, true, true],
5817        );
5818
5819        // s1 != 'foo' AND (s2 = 'bar' OR s2 = 'baz')
5820        prune_with_expr(
5821            col("s1")
5822                .not_eq(lit("foo"))
5823                .and(col("s2").eq(lit("bar")).or(col("s2").eq(lit("baz")))),
5824            &schema,
5825            &statistics,
5826            // Can rule out any container where we know s1 has only the value
5827            // 'foo'. Can't use knowledge of s2 and bar to rule out anything
5828            &[false, true, true, false, true, true, false, true, true],
5829        );
5830
5831        // s1 like '%foo%bar%'
5832        prune_with_expr(
5833            col("s1").like(lit("foo%bar%")),
5834            &schema,
5835            &statistics,
5836            // cant rule out anything with information we know
5837            &[true, true, true, true, true, true, true, true, true],
5838        );
5839
5840        // s1 like '%foo%bar%' AND s2 = 'bar'
5841        prune_with_expr(
5842            col("s1")
5843                .like(lit("foo%bar%"))
5844                .and(col("s2").eq(lit("bar"))),
5845            &schema,
5846            &statistics,
5847            // can rule out any container where we know s2 does not have the value 'bar'
5848            &[true, true, true, false, false, false, true, true, true],
5849        );
5850
5851        // s1 like '%foo%bar%' OR s2 = 'bar'
5852        prune_with_expr(
5853            col("s1").like(lit("foo%bar%")).or(col("s2").eq(lit("bar"))),
5854            &schema,
5855            &statistics,
5856            // can't rule out anything (we would have to prove that both the
5857            // like and the equality must be false)
5858            &[true, true, true, true, true, true, true, true, true],
5859        );
5860    }
5861
5862    #[test]
5863    fn prune_with_range_and_contained() {
5864        // Setup mimics range information for i, a bloom filter for s
5865        let schema = Arc::new(Schema::new(vec![
5866            Field::new("i", DataType::Int32, true),
5867            Field::new("s", DataType::Utf8, true),
5868        ]));
5869
5870        let statistics = TestStatistics::new()
5871            .with(
5872                "i",
5873                ContainerStats::new_i32(
5874                    // Container 0, 3, 6: [-5 to 5]
5875                    // Container 1, 4, 7: [10 to 20]
5876                    // Container 2, 5, 9: unknown
5877                    vec![
5878                        Some(-5),
5879                        Some(10),
5880                        None,
5881                        Some(-5),
5882                        Some(10),
5883                        None,
5884                        Some(-5),
5885                        Some(10),
5886                        None,
5887                    ], // min
5888                    vec![
5889                        Some(5),
5890                        Some(20),
5891                        None,
5892                        Some(5),
5893                        Some(20),
5894                        None,
5895                        Some(5),
5896                        Some(20),
5897                        None,
5898                    ], // max
5899                ),
5900            )
5901            // Add contained  information about the s and "foo"
5902            .with_contained(
5903                "s",
5904                [ScalarValue::from("foo")],
5905                [
5906                    // container 0,1,2 known to only contain "foo"
5907                    Some(true),
5908                    Some(true),
5909                    Some(true),
5910                    // container 3,4,5 known to not contain "foo"
5911                    Some(false),
5912                    Some(false),
5913                    Some(false),
5914                    // container 6,7,8 unknown about "foo"
5915                    None,
5916                    None,
5917                    None,
5918                ],
5919            );
5920
5921        // i = 0 and s = 'foo'
5922        prune_with_expr(
5923            col("i").eq(lit(0)).and(col("s").eq(lit("foo"))),
5924            &schema,
5925            &statistics,
5926            // Can rule out container where we know that either:
5927            // 1. 0 is outside the min/max range of i
5928            // 1. s does not contain foo
5929            // (range is false, and contained  is false)
5930            &[true, false, true, false, false, false, true, false, true],
5931        );
5932
5933        // i = 0 and s != 'foo'
5934        prune_with_expr(
5935            col("i").eq(lit(0)).and(col("s").not_eq(lit("foo"))),
5936            &schema,
5937            &statistics,
5938            // Can rule out containers where either:
5939            // 1. 0 is outside the min/max range of i
5940            // 2. s only contains foo
5941            &[false, false, false, true, false, true, true, false, true],
5942        );
5943
5944        // i = 0 OR s = 'foo'
5945        prune_with_expr(
5946            col("i").eq(lit(0)).or(col("s").eq(lit("foo"))),
5947            &schema,
5948            &statistics,
5949            // in theory could rule out containers if we had min/max values for
5950            // s as well. But in this case we don't so we can't rule out anything
5951            &[true, true, true, true, true, true, true, true, true],
5952        );
5953    }
5954
5955    /// prunes the specified expr with the specified schema and statistics, and
5956    /// ensures it returns expected.
5957    ///
5958    /// `expected` is a vector of bools, where true means the row group should
5959    /// be kept, and false means it should be pruned.
5960    // TODO refactor other tests to use this to reduce boiler plate
5961    fn prune_with_expr(
5962        expr: Expr,
5963        schema: &SchemaRef,
5964        statistics: &TestStatistics,
5965        expected: &[bool],
5966    ) {
5967        println!("Pruning with expr: {expr}");
5968        let expr = logical2physical(&expr, schema);
5969        let p = PruningPredicateBuilder::new()
5970            .with_file_schema(Arc::<Schema>::clone(schema))
5971            .try_build(expr)
5972            .unwrap();
5973        let result = p.prune(statistics).unwrap();
5974        assert_eq!(result, expected);
5975    }
5976
5977    fn prune_with_simplified_expr(
5978        expr: Expr,
5979        schema: &SchemaRef,
5980        statistics: &TestStatistics,
5981        expected: &[bool],
5982    ) {
5983        println!("Pruning with expr: {expr}");
5984        let expr = logical2physical(&expr, schema);
5985        let simplifier = PhysicalExprSimplifier::new(schema);
5986        let expr = simplifier.simplify(expr).unwrap();
5987        let p = PruningPredicateBuilder::new()
5988            .with_file_schema(Arc::<Schema>::clone(schema))
5989            .try_build(expr)
5990            .unwrap();
5991        let result = p.prune(statistics).unwrap();
5992        assert_eq!(result, expected);
5993    }
5994
5995    fn is_not_distinct_from(left: Expr, right: Expr) -> Expr {
5996        Expr::BinaryExpr(BinaryExpr::new(
5997            Box::new(left),
5998            Operator::IsNotDistinctFrom,
5999            Box::new(right),
6000        ))
6001    }
6002
6003    fn is_distinct_from(left: Expr, right: Expr) -> Expr {
6004        Expr::BinaryExpr(BinaryExpr::new(
6005            Box::new(left),
6006            Operator::IsDistinctFrom,
6007            Box::new(right),
6008        ))
6009    }
6010
6011    fn test_build_predicate_expression(
6012        expr: &Expr,
6013        schema: &Schema,
6014        required_columns: &mut RequiredColumns,
6015    ) -> Arc<dyn PhysicalExpr> {
6016        let expr = logical2physical(expr, schema);
6017        let unhandled_hook = Arc::new(ConstantUnhandledPredicateHook::default()) as _;
6018        build_predicate_expression(
6019            &expr,
6020            &Arc::new(schema.clone()),
6021            required_columns,
6022            &unhandled_hook,
6023            MAX_IN_LIST_SIZE,
6024        )
6025    }
6026
6027    #[test]
6028    fn test_build_predicate_expression_with_false() {
6029        let expr = lit(ScalarValue::Boolean(Some(false)));
6030        let schema = Schema::empty();
6031        let res =
6032            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
6033        let expected = logical2physical(&expr, &schema);
6034        assert_eq!(&res, &expected);
6035    }
6036
6037    #[test]
6038    fn test_build_predicate_expression_with_and_false() {
6039        let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]);
6040        let expr = and(
6041            col("c1").eq(lit("a")),
6042            lit(ScalarValue::Boolean(Some(false))),
6043        );
6044        let res =
6045            test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
6046        let expected = logical2physical(&lit(ScalarValue::Boolean(Some(false))), &schema);
6047        assert_eq!(&res, &expected);
6048    }
6049
6050    #[test]
6051    fn test_build_predicate_expression_with_or_false() {
6052        let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]);
6053        let left_expr = col("c1").eq(lit("a"));
6054        let right_expr = lit(ScalarValue::Boolean(Some(false)));
6055        let res = test_build_predicate_expression(
6056            &or(left_expr.clone(), right_expr.clone()),
6057            &schema,
6058            &mut RequiredColumns::new(),
6059        );
6060        let expected =
6061            "c1_null_count@2 != row_count@3 AND c1_min@0 <= a AND a <= c1_max@1";
6062        assert_eq!(res.to_string(), expected);
6063    }
6064}