blob: 380ada10df6e104d30b49b5340066759c1186286 [file]
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! [`PruningPredicate`] to apply filter [`Expr`] to prune "containers"
//! based on statistics (e.g. Parquet Row Groups)
//!
//! [`Expr`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/enum.Expr.html
use std::collections::HashSet;
use std::sync::Arc;
use arrow::array::AsArray;
use arrow::{
array::{new_null_array, ArrayRef, BooleanArray},
datatypes::{DataType, Field, Schema, SchemaRef},
record_batch::{RecordBatch, RecordBatchOptions},
};
// pub use for backwards compatibility
pub use datafusion_common::pruning::PruningStatistics;
use datafusion_physical_expr::simplifier::PhysicalExprSimplifier;
use datafusion_physical_plan::metrics::Count;
use log::{debug, trace};
use datafusion_common::error::Result;
use datafusion_common::tree_node::TransformedResult;
use datafusion_common::{
internal_datafusion_err, internal_err, plan_datafusion_err, plan_err,
tree_node::{Transformed, TreeNode},
ScalarValue,
};
use datafusion_common::{Column, DFSchema};
use datafusion_expr_common::operator::Operator;
use datafusion_physical_expr::utils::{collect_columns, Guarantee, LiteralGuarantee};
use datafusion_physical_expr::{expressions as phys_expr, PhysicalExprRef};
use datafusion_physical_expr_common::physical_expr::snapshot_physical_expr;
use datafusion_physical_plan::{ColumnarValue, PhysicalExpr};
/// Used to prove that arbitrary predicates (boolean expression) can not
/// possibly evaluate to `true` given information about a column provided by
/// [`PruningStatistics`].
///
/// # Introduction
///
/// `PruningPredicate` analyzes filter expressions using statistics such as
/// min/max values and null counts, attempting to prove a "container" (e.g.
/// Parquet Row Group) can be skipped without reading the actual data,
/// potentially leading to significant performance improvements.
///
/// For example, `PruningPredicate`s are used to prune Parquet Row Groups based
/// on the min/max values found in the Parquet metadata. If the
/// `PruningPredicate` can prove that the filter can never evaluate to `true`
/// for any row in the Row Group, the entire Row Group is skipped during query
/// execution.
///
/// The `PruningPredicate` API is general, and can be used for pruning other
/// types of containers (e.g. files) based on statistics that may be known from
/// external catalogs (e.g. Delta Lake) or other sources. How this works is a
/// subtle topic. See the Background and Implementation section for details.
///
/// `PruningPredicate` supports:
///
/// 1. Arbitrary expressions (including user defined functions)
///
/// 2. Vectorized evaluation (provide more than one set of statistics at a time)
/// so it is suitable for pruning 1000s of containers.
///
/// 3. Any source of information that implements the [`PruningStatistics`] trait
/// (not just Parquet metadata).
///
/// # Example
///
/// See the [`pruning.rs` example in the `datafusion-examples`] for a complete
/// example of how to use `PruningPredicate` to prune files based on min/max
/// values.
///
/// [`pruning.rs` example in the `datafusion-examples`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/pruning.rs
///
/// Given an expression like `x = 5` and statistics for 3 containers (Row
/// Groups, files, etc) `A`, `B`, and `C`:
///
/// ```text
/// A: {x_min = 0, x_max = 4}
/// B: {x_min = 2, x_max = 10}
/// C: {x_min = 5, x_max = 8}
/// ```
///
/// `PruningPredicate` will conclude that the rows in container `A` can never
/// be true (as the maximum value is only `4`), so it can be pruned:
///
/// ```text
/// A: false (no rows could possibly match x = 5)
/// B: true (rows might match x = 5)
/// C: true (rows might match x = 5)
/// ```
///
/// See [`PruningPredicate::try_new`] and [`PruningPredicate::prune`] for more information.
///
/// # Background
///
/// ## Boolean Tri-state logic
///
/// To understand the details of the rest of this documentation, it is important
/// to understand how the tri-state boolean logic in SQL works. As this is
/// somewhat esoteric, we review it here.
///
/// SQL has a notion of `NULL` that represents the value is `“unknown”` and this
/// uncertainty propagates through expressions. SQL `NULL` behaves very
/// differently than the `NULL` in most other languages where it is a special,
/// sentinel value (e.g. `0` in `C/C++`). While representing uncertainty with
/// `NULL` is powerful and elegant, SQL `NULL`s are often deeply confusing when
/// first encountered as they behave differently than most programmers may
/// expect.
///
/// In most other programming languages,
/// * `a == NULL` evaluates to `true` if `a` also had the value `NULL`
/// * `a == NULL` evaluates to `false` if `a` has any other value
///
/// However, in SQL `a = NULL` **always** evaluates to `NULL` (never `true` or
/// `false`):
///
/// Expression | Result
/// ------------- | ---------
/// `1 = NULL` | `NULL`
/// `NULL = NULL` | `NULL`
///
/// Also important is how `AND` and `OR` works with tri-state boolean logic as
/// (perhaps counterintuitively) the result is **not** always NULL. While
/// consistent with the notion of `NULL` representing “unknown”, this is again,
/// often deeply confusing 🤯 when first encountered.
///
/// Expression | Result | Intuition
/// --------------- | --------- | -----------
/// `NULL AND true` | `NULL` | The `NULL` stands for “unknown” and if it were `true` or `false` the overall expression value could change
/// `NULL AND false` | `false` | If the `NULL` was either `true` or `false` the overall expression is still `false`
/// `NULL AND NULL` | `NULL` |
///
/// Expression | Result | Intuition
/// --------------- | --------- | ----------
/// `NULL OR true` | `true` | If the `NULL` was either `true` or `false` the overall expression is still `true`
/// `NULL OR false` | `NULL` | The `NULL` stands for “unknown” and if it were `true` or `false` the overall expression value could change
/// `NULL OR NULL` | `NULL` |
///
/// ## SQL Filter Semantics
///
/// The SQL `WHERE` clause has a boolean expression, often called a filter or
/// predicate. The semantics of this predicate are that the query evaluates the
/// predicate for each row in the input tables and:
///
/// * Rows that evaluate to `true` are returned in the query results
///
/// * Rows that evaluate to `false` are not returned (“filtered out” or “pruned” or “skipped”).
///
/// * Rows that evaluate to `NULL` are **NOT** returned (also “filtered out”).
/// Note: *this treatment of `NULL` is **DIFFERENT** than how `NULL` is treated
/// in the rewritten predicate described below.*
///
/// # `PruningPredicate` Implementation
///
/// Armed with the information in the Background section, we can now understand
/// how the `PruningPredicate` logic works.
///
/// ## Interface
///
/// **Inputs**
/// 1. An input schema describing what columns exist
///
/// 2. A predicate (expression that evaluates to a boolean)
///
/// 3. [`PruningStatistics`] that provides information about columns in that
/// schema, for multiple “containers”. For each column in each container, it
/// provides optional information on contained values, min_values, max_values,
/// null_counts counts, and row_counts counts.
///
/// **Outputs**:
/// A (non null) boolean value for each container:
/// * `true`: There MAY be rows that match the predicate
///
/// * `false`: There are no rows that could possibly match the predicate (the
/// predicate can never possibly be true). The container can be pruned (skipped)
/// entirely.
///
/// While `PruningPredicate` will never return a `NULL` value, the
/// rewritten predicate (as returned by `build_predicate_expression` and used internally
/// by `PruningPredicate`) may evaluate to `NULL` when some of the min/max values
/// or null / row counts are not known.
///
/// In order to be correct, `PruningPredicate` must return false
/// **only** if it can determine that for all rows in the container, the
/// predicate could never evaluate to `true` (always evaluates to either `NULL`
/// or `false`).
///
/// ## Contains Analysis and Min/Max Rewrite
///
/// `PruningPredicate` works by first analyzing the predicate to see what
/// [`LiteralGuarantee`] must hold for the predicate to be true.
///
/// Then, the `PruningPredicate` rewrites the original predicate into an
/// expression that references the min/max values of each column in the original
/// predicate.
///
/// When the min/max values are actually substituted in to this expression and
/// evaluated, the result means
///
/// * `true`: there MAY be rows that pass the predicate, **KEEPS** the container
///
/// * `NULL`: there MAY be rows that pass the predicate, **KEEPS** the container
/// Note that rewritten predicate can evaluate to NULL when some of
/// the min/max values are not known. *Note that this is different than
/// the SQL filter semantics where `NULL` means the row is filtered
/// out.*
///
/// * `false`: there are no rows that could possibly match the predicate,
/// **PRUNES** the container
///
/// For example, given a column `x`, the `x_min`, `x_max`, `x_null_count`, and
/// `x_row_count` represent the minimum and maximum values, the null count of
/// column `x`, and the row count of column `x`, provided by the `PruningStatistics`.
/// `x_null_count` and `x_row_count` are used to handle the case where the column `x`
/// is known to be all `NULL`s. Note this is different from knowing nothing about
/// the column `x`, which confusingly is encoded by returning `NULL` for the min/max
/// values from [`PruningStatistics::max_values`] and [`PruningStatistics::min_values`].
///
/// Here are some examples of the rewritten predicates:
///
/// Original Predicate | Rewritten Predicate
/// ------------------ | --------------------
/// `x = 5` | `x_null_count != x_row_count AND (x_min <= 5 AND 5 <= x_max)`
/// `x < 5` | `x_null_count != x_row_count THEN false (x_max < 5)`
/// `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)`
/// `x IS NULL` | `x_null_count > 0`
/// `x IS NOT NULL` | `x_null_count != row_count`
/// `CAST(x as int) = 5` | `x_null_count != x_row_count (CAST(x_min as int) <= 5 AND 5 <= CAST(x_max as int))`
///
/// ## Predicate Evaluation
/// The PruningPredicate works in two passes
///
/// **First pass**: For each `LiteralGuarantee` calls
/// [`PruningStatistics::contained`] and rules out containers where the
/// LiteralGuarantees are not satisfied
///
/// **Second Pass**: Evaluates the rewritten expression using the
/// min/max/null_counts/row_counts values for each column for each container. For any
/// container that this expression evaluates to `false`, it rules out those
/// containers.
///
///
/// ### Example 1
///
/// Given the predicate, `x = 5 AND y = 10`, the rewritten predicate would look like:
///
/// ```sql
/// x_null_count != x_row_count AND (x_min <= 5 AND 5 <= x_max)
/// AND
/// y_null_count != y_row_count AND (y_min <= 10 AND 10 <= y_max)
/// ```
///
/// If we know that for a given container, `x` is between `1 and 100` and we know that
/// `y` is between `4` and `7`, we know nothing about the null count and row count of
/// `x` and `y`, the input statistics might look like:
///
/// Column | Value
/// -------- | -----
/// `x_min` | `1`
/// `x_max` | `100`
/// `x_null_count` | `null`
/// `x_row_count` | `null`
/// `y_min` | `4`
/// `y_max` | `7`
/// `y_null_count` | `null`
/// `y_row_count` | `null`
///
/// When these statistics values are substituted in to the rewritten predicate and
/// simplified, the result is `false`:
///
/// * `null != null AND (1 <= 5 AND 5 <= 100) AND null != null AND (4 <= 10 AND 10 <= 7)`
/// * `null = null` is `null` which is not true, so the AND moves on to the next clause
/// * `null and (1 <= 5 AND 5 <= 100) AND null AND (4 <= 10 AND 10 <= 7)`
/// * evaluating the clauses further we get:
/// * `null and true and null and false`
/// * `null and false`
/// * `false`
///
/// Returning `false` means the container can be pruned, which matches the
/// intuition that `x = 5 AND y = 10` can’t be true for any row if all values of `y`
/// are `7` or less.
///
/// Note that if we had ended up with `null AND true AND null AND true` the result
/// would have been `null`.
/// `null` is treated the same as`true`, because we can't prove that the predicate is `false.`
///
/// If, for some other container, we knew `y` was between the values `4` and
/// `15`, then the rewritten predicate evaluates to `true` (verifying this is
/// left as an exercise to the reader -- are you still here?), and the container
/// **could not** be pruned. The intuition is that there may be rows where the
/// predicate *might* evaluate to `true`, and the only way to find out is to do
/// more analysis, for example by actually reading the data and evaluating the
/// predicate row by row.
///
/// ### Example 2
///
/// Given the same predicate, `x = 5 AND y = 10`, the rewritten predicate would
/// look like the same as example 1:
///
/// ```sql
/// x_null_count != x_row_count AND (x_min <= 5 AND 5 <= x_max)
/// AND
/// y_null_count != y_row_count AND (y_min <= 10 AND 10 <= y_max)
/// ```
///
/// If we know that for another given container, `x_min` is NULL and `x_max` is
/// NULL (the min/max values are unknown), `x_null_count` is `100` and `x_row_count`
/// is `100`; we know that `y` is between `4` and `7`, but we know nothing about
/// the null count and row count of `y`. The input statistics might look like:
///
/// Column | Value
/// -------- | -----
/// `x_min` | `null`
/// `x_max` | `null`
/// `x_null_count` | `100`
/// `x_row_count` | `100`
/// `y_min` | `4`
/// `y_max` | `7`
/// `y_null_count` | `null`
/// `y_row_count` | `null`
///
/// When these statistics values are substituted in to the rewritten predicate and
/// simplified, the result is `false`:
///
/// * `100 != 100 AND (null <= 5 AND 5 <= null) AND null = null AND (4 <= 10 AND 10 <= 7)`
/// * `false AND null AND null AND false`
/// * `false AND false`
/// * `false`
///
/// Returning `false` means the container can be pruned, which matches the
/// intuition that `x = 5 AND y = 10` can’t be true because all values in `x`
/// are known to be NULL.
///
/// # Related Work
///
/// [`PruningPredicate`] implements the type of min/max pruning described in
/// Section `3.3.3` of the [`Snowflake SIGMOD Paper`]. The technique is
/// described by various research such as [small materialized aggregates], [zone
/// maps], and [data skipping].
///
/// [`Snowflake SIGMOD Paper`]: https://dl.acm.org/doi/10.1145/2882903.2903741
/// [small materialized aggregates]: https://www.vldb.org/conf/1998/p476.pdf
/// [zone maps]: https://dl.acm.org/doi/10.1007/978-3-642-03730-6_10
/// [data skipping]: https://dl.acm.org/doi/10.1145/2588555.2610515
#[derive(Debug, Clone)]
pub struct PruningPredicate {
/// The input schema against which the predicate will be evaluated
schema: SchemaRef,
/// A min/max pruning predicate (rewritten in terms of column min/max
/// values, which are supplied by statistics)
predicate_expr: Arc<dyn PhysicalExpr>,
/// Description of which statistics are required to evaluate `predicate_expr`
required_columns: RequiredColumns,
/// Original physical predicate from which this predicate expr is derived
/// (required for serialization)
orig_expr: Arc<dyn PhysicalExpr>,
/// [`LiteralGuarantee`]s used to try and prove a predicate can not possibly
/// evaluate to `true`.
///
/// See [`PruningPredicate::literal_guarantees`] for more details.
literal_guarantees: Vec<LiteralGuarantee>,
}
/// Build a pruning predicate from an optional predicate expression.
/// If the predicate is None or the predicate cannot be converted to a pruning
/// predicate, return None.
/// If there is an error creating the pruning predicate it is recorded by incrementing
/// the `predicate_creation_errors` counter.
pub fn build_pruning_predicate(
predicate: Arc<dyn PhysicalExpr>,
file_schema: &SchemaRef,
predicate_creation_errors: &Count,
) -> Option<Arc<PruningPredicate>> {
match PruningPredicate::try_new(predicate, Arc::clone(file_schema)) {
Ok(pruning_predicate) => {
if !pruning_predicate.always_true() {
return Some(Arc::new(pruning_predicate));
}
}
Err(e) => {
debug!("Could not create pruning predicate for: {e}");
predicate_creation_errors.add(1);
}
}
None
}
/// Rewrites predicates that [`PredicateRewriter`] can not handle, e.g. certain
/// complex expressions or predicates that reference columns that are not in the
/// schema.
pub trait UnhandledPredicateHook {
/// Called when a predicate can not be rewritten in terms of statistics or
/// references a column that is not in the schema.
fn handle(&self, expr: &Arc<dyn PhysicalExpr>) -> Arc<dyn PhysicalExpr>;
}
/// The default handling for unhandled predicates is to return a constant `true`
/// (meaning don't prune the container)
#[derive(Debug, Clone)]
struct ConstantUnhandledPredicateHook {
default: Arc<dyn PhysicalExpr>,
}
impl Default for ConstantUnhandledPredicateHook {
fn default() -> Self {
Self {
default: Arc::new(phys_expr::Literal::new(ScalarValue::from(true))),
}
}
}
impl UnhandledPredicateHook for ConstantUnhandledPredicateHook {
fn handle(&self, _expr: &Arc<dyn PhysicalExpr>) -> Arc<dyn PhysicalExpr> {
Arc::clone(&self.default)
}
}
impl PruningPredicate {
/// Try to create a new instance of [`PruningPredicate`]
///
/// This will translate the provided `expr` filter expression into
/// a *pruning predicate*.
///
/// A pruning predicate is one that has been rewritten in terms of
/// the min and max values of column references and that evaluates
/// to FALSE if the filter predicate would evaluate FALSE *for
/// every row* whose values fell within the min / max ranges (aka
/// could be pruned).
///
/// The pruning predicate evaluates to TRUE or NULL
/// if the filter predicate *might* evaluate to TRUE for at least
/// one row whose values fell within the min/max ranges (in other
/// words they might pass the predicate)
///
/// For example, the filter expression `(column / 2) = 4` becomes
/// the pruning predicate
/// `(column_min / 2) <= 4 && 4 <= (column_max / 2))`
///
/// See the struct level documentation on [`PruningPredicate`] for more
/// details.
pub fn try_new(expr: Arc<dyn PhysicalExpr>, schema: SchemaRef) -> Result<Self> {
// Get a (simpler) snapshot of the physical expr here to use with `PruningPredicate`
// which does not handle dynamic exprs in general
let expr = snapshot_physical_expr(expr)?;
let unhandled_hook = Arc::new(ConstantUnhandledPredicateHook::default()) as _;
// build predicate expression once
let mut required_columns = RequiredColumns::new();
let predicate_expr = build_predicate_expression(
&expr,
&schema,
&mut required_columns,
&unhandled_hook,
);
let predicate_schema = required_columns.schema();
// Simplify the newly created predicate to get rid of redundant casts, comparisons, etc.
let predicate_expr =
PhysicalExprSimplifier::new(&predicate_schema).simplify(predicate_expr)?;
let literal_guarantees = LiteralGuarantee::analyze(&expr);
Ok(Self {
schema,
predicate_expr,
required_columns,
orig_expr: expr,
literal_guarantees,
})
}
/// For each set of statistics, evaluates the pruning predicate
/// and returns a `bool` with the following meaning for a
/// all rows whose values match the statistics:
///
/// `true`: There MAY be rows that match the predicate
///
/// `false`: There are no rows that could possibly match the predicate
///
/// Note: the predicate passed to `prune` should already be simplified as
/// much as possible (e.g. this pass doesn't handle some
/// expressions like `b = false`, but it does handle the
/// simplified version `b`. See [`ExprSimplifier`] to simplify expressions.
///
/// [`ExprSimplifier`]: https://docs.rs/datafusion/latest/datafusion/optimizer/simplify_expressions/struct.ExprSimplifier.html
pub fn prune<S: PruningStatistics + ?Sized>(
&self,
statistics: &S,
) -> Result<Vec<bool>> {
let mut builder = BoolVecBuilder::new(statistics.num_containers());
// Try to prove the predicate can't be true for the containers based on
// literal guarantees
for literal_guarantee in &self.literal_guarantees {
let LiteralGuarantee {
column,
guarantee,
literals,
} = literal_guarantee;
if let Some(results) = statistics.contained(column, literals) {
match guarantee {
// `In` means the values in the column must be one of the
// values in the set for the predicate to evaluate to true.
// If `contained` returns false, that means the column is
// not any of the values so we can prune the container
Guarantee::In => builder.combine_array(&results),
// `NotIn` means the values in the column must not be
// any of the values in the set for the predicate to
// evaluate to true. If `contained` returns true, it means the
// column is only in the set of values so we can prune the
// container
Guarantee::NotIn => {
builder.combine_array(&arrow::compute::not(&results)?)
}
}
// if all containers are pruned (has rows that DEFINITELY DO NOT pass the predicate)
// can return early without evaluating the rest of predicates.
if builder.check_all_pruned() {
return Ok(builder.build());
}
}
}
// Next, try to prove the predicate can't be true for the containers based
// on min/max values
// build a RecordBatch that contains the min/max values in the
// appropriate statistics columns for the min/max predicate
let statistics_batch =
build_statistics_record_batch(statistics, &self.required_columns)?;
// Evaluate the pruning predicate on that record batch and append any results to the builder
builder.combine_value(self.predicate_expr.evaluate(&statistics_batch)?);
Ok(builder.build())
}
/// Return a reference to the input schema
pub fn schema(&self) -> &SchemaRef {
&self.schema
}
/// Returns a reference to the physical expr used to construct this pruning predicate
pub fn orig_expr(&self) -> &Arc<dyn PhysicalExpr> {
&self.orig_expr
}
/// Returns a reference to the predicate expr
pub fn predicate_expr(&self) -> &Arc<dyn PhysicalExpr> {
&self.predicate_expr
}
/// Returns a reference to the literal guarantees
///
/// Note that **All** `LiteralGuarantee`s must be satisfied for the
/// expression to possibly be `true`. If any is not satisfied, the
/// expression is guaranteed to be `null` or `false`.
pub fn literal_guarantees(&self) -> &[LiteralGuarantee] {
&self.literal_guarantees
}
/// Returns true if this pruning predicate can not prune anything.
///
/// This happens if the predicate is a literal `true` and
/// literal_guarantees is empty.
///
/// This can happen when a predicate is simplified to a constant `true`
pub fn always_true(&self) -> bool {
is_always_true(&self.predicate_expr) && self.literal_guarantees.is_empty()
}
// this is only used by `parquet` feature right now
#[allow(dead_code)]
pub fn required_columns(&self) -> &RequiredColumns {
&self.required_columns
}
/// Names of the columns that are known to be / not be in a set
/// of literals (constants). These are the columns the that may be passed to
/// [`PruningStatistics::contained`] during pruning.
///
/// This is useful to avoid fetching statistics for columns that will not be
/// used in the predicate. For example, it can be used to avoid reading
/// unneeded bloom filters (a non trivial operation).
pub fn literal_columns(&self) -> Vec<String> {
let mut seen = HashSet::new();
self.literal_guarantees
.iter()
.map(|e| &e.column.name)
// avoid duplicates
.filter(|name| seen.insert(*name))
.map(|s| s.to_string())
.collect()
}
}
/// Builds the return `Vec` for [`PruningPredicate::prune`].
#[derive(Debug)]
struct BoolVecBuilder {
/// One element per container. Each element is
/// * `true`: if the container has row that may pass the predicate
/// * `false`: if the container has rows that DEFINITELY DO NOT pass the predicate
inner: Vec<bool>,
}
impl BoolVecBuilder {
/// Create a new `BoolVecBuilder` with `num_containers` elements
fn new(num_containers: usize) -> Self {
Self {
// assume by default all containers may pass the predicate
inner: vec![true; num_containers],
}
}
/// Combines result `array` for a conjunct (e.g. `AND` clause) of a
/// predicate into the currently in progress array.
///
/// Each `array` element is:
/// * `true`: container has row that may pass the predicate
/// * `false`: all container rows DEFINITELY DO NOT pass the predicate
/// * `null`: container may or may not have rows that pass the predicate
fn combine_array(&mut self, array: &BooleanArray) {
assert_eq!(array.len(), self.inner.len());
for (cur, new) in self.inner.iter_mut().zip(array.iter()) {
// `false` for this conjunct means we know for sure no rows could
// pass the predicate and thus we set the corresponding container
// location to false.
if let Some(false) = new {
*cur = false;
}
}
}
/// Combines the results in the [`ColumnarValue`] to the currently in
/// progress array, following the same rules as [`Self::combine_array`].
///
/// # Panics
/// If `value` is not boolean
fn combine_value(&mut self, value: ColumnarValue) {
match value {
ColumnarValue::Array(array) => {
self.combine_array(array.as_boolean());
}
ColumnarValue::Scalar(ScalarValue::Boolean(Some(false))) => {
// False means all containers can not pass the predicate
self.inner = vec![false; self.inner.len()];
}
_ => {
// Null or true means the rows in container may pass this
// conjunct so we can't prune any containers based on that
}
}
}
/// Convert this builder into a Vec of bools
fn build(self) -> Vec<bool> {
self.inner
}
/// Check all containers has rows that DEFINITELY DO NOT pass the predicate
fn check_all_pruned(&self) -> bool {
self.inner.iter().all(|&x| !x)
}
}
fn is_always_true(expr: &Arc<dyn PhysicalExpr>) -> bool {
expr.as_any()
.downcast_ref::<phys_expr::Literal>()
.map(|l| matches!(l.value(), ScalarValue::Boolean(Some(true))))
.unwrap_or_default()
}
fn is_always_false(expr: &Arc<dyn PhysicalExpr>) -> bool {
expr.as_any()
.downcast_ref::<phys_expr::Literal>()
.map(|l| matches!(l.value(), ScalarValue::Boolean(Some(false))))
.unwrap_or_default()
}
/// Describes which columns statistics are necessary to evaluate a
/// [`PruningPredicate`].
///
/// This structure permits reading and creating the minimum number statistics,
/// which is important since statistics may be non trivial to read (e.g. large
/// strings or when there are 1000s of columns).
///
/// Handles creating references to the min/max statistics
/// for columns as well as recording which statistics are needed
#[derive(Debug, Default, Clone)]
pub struct RequiredColumns {
/// The statistics required to evaluate this predicate:
/// * The unqualified column in the input schema
/// * Statistics type (e.g. Min or Max or Null_Count)
/// * The field the statistics value should be placed in for
/// pruning predicate evaluation (e.g. `min_value` or `max_value`)
columns: Vec<(phys_expr::Column, StatisticsType, Field)>,
}
impl RequiredColumns {
fn new() -> Self {
Self::default()
}
/// Returns Some(column) if this is a single column predicate.
///
/// Returns None if this is a multi-column predicate.
///
/// Examples:
/// * `a > 5 OR a < 10` returns `Some(a)`
/// * `a > 5 OR b < 10` returns `None`
/// * `true` returns None
#[allow(dead_code)]
// this fn is only used by `parquet` feature right now, thus the `allow(dead_code)`
pub fn single_column(&self) -> Option<&phys_expr::Column> {
if self.columns.windows(2).all(|w| {
// check if all columns are the same (ignoring statistics and field)
let c1 = &w[0].0;
let c2 = &w[1].0;
c1 == c2
}) {
self.columns.first().map(|r| &r.0)
} else {
None
}
}
/// Returns a schema that describes the columns required to evaluate this
/// pruning predicate.
/// The schema contains the fields for each column in `self.columns` with
/// the appropriate data type for the statistics.
/// Order matters, this same order is used to evaluate the
/// pruning predicate.
fn schema(&self) -> Schema {
let fields = self
.columns
.iter()
.map(|(_c, _t, f)| f.clone())
.collect::<Vec<_>>();
Schema::new(fields)
}
/// Returns an iterator over items in columns (see doc on
/// `self.columns` for details)
pub(crate) fn iter(
&self,
) -> impl Iterator<Item = &(phys_expr::Column, StatisticsType, Field)> {
self.columns.iter()
}
fn find_stat_column(
&self,
column: &phys_expr::Column,
statistics_type: StatisticsType,
) -> Option<usize> {
match statistics_type {
StatisticsType::RowCount => {
// Use the first row count we find, if any
self.columns
.iter()
.enumerate()
.find(|(_i, (_c, t, _f))| t == &statistics_type)
.map(|(i, (_c, _t, _f))| i)
}
_ => self
.columns
.iter()
.enumerate()
.find(|(_i, (c, t, _f))| c == column && t == &statistics_type)
.map(|(i, (_c, _t, _f))| i),
}
}
/// Rewrites column_expr so that all appearances of column
/// are replaced with a reference to either the min or max
/// statistics column, while keeping track that a reference to the statistics
/// column is required
///
/// for example, an expression like `col("foo") > 5`, when called
/// with Max would result in an expression like `col("foo_max") >
/// 5` with the appropriate entry noted in self.columns
fn stat_column_expr(
&mut self,
column: &phys_expr::Column,
column_expr: &Arc<dyn PhysicalExpr>,
field: &Field,
stat_type: StatisticsType,
) -> Result<Arc<dyn PhysicalExpr>> {
let (idx, need_to_insert) = match self.find_stat_column(column, stat_type) {
Some(idx) => (idx, false),
None => (self.columns.len(), true),
};
let column_name = column.name();
let stat_column_name = match stat_type {
StatisticsType::Min => format!("{column_name}_min"),
StatisticsType::Max => format!("{column_name}_max"),
StatisticsType::NullCount => format!("{column_name}_null_count"),
StatisticsType::RowCount => "row_count".to_string(),
};
let stat_column = phys_expr::Column::new(&stat_column_name, idx);
// only add statistics column if not previously added
if need_to_insert {
// may be null if statistics are not present
let nullable = true;
let stat_field =
Field::new(stat_column.name(), field.data_type().clone(), nullable);
self.columns.push((column.clone(), stat_type, stat_field));
}
rewrite_column_expr(Arc::clone(column_expr), column, &stat_column)
}
/// rewrite col --> col_min
fn min_column_expr(
&mut self,
column: &phys_expr::Column,
column_expr: &Arc<dyn PhysicalExpr>,
field: &Field,
) -> Result<Arc<dyn PhysicalExpr>> {
self.stat_column_expr(column, column_expr, field, StatisticsType::Min)
}
/// rewrite col --> col_max
fn max_column_expr(
&mut self,
column: &phys_expr::Column,
column_expr: &Arc<dyn PhysicalExpr>,
field: &Field,
) -> Result<Arc<dyn PhysicalExpr>> {
self.stat_column_expr(column, column_expr, field, StatisticsType::Max)
}
/// rewrite col --> col_null_count
fn null_count_column_expr(
&mut self,
column: &phys_expr::Column,
column_expr: &Arc<dyn PhysicalExpr>,
field: &Field,
) -> Result<Arc<dyn PhysicalExpr>> {
self.stat_column_expr(column, column_expr, field, StatisticsType::NullCount)
}
/// rewrite col --> col_row_count
fn row_count_column_expr(
&mut self,
column: &phys_expr::Column,
column_expr: &Arc<dyn PhysicalExpr>,
field: &Field,
) -> Result<Arc<dyn PhysicalExpr>> {
self.stat_column_expr(column, column_expr, field, StatisticsType::RowCount)
}
}
impl From<Vec<(phys_expr::Column, StatisticsType, Field)>> for RequiredColumns {
fn from(columns: Vec<(phys_expr::Column, StatisticsType, Field)>) -> Self {
Self { columns }
}
}
/// Build a RecordBatch from a list of statistics, creating arrays,
/// with one row for each PruningStatistics and columns specified in
/// the required_columns parameter.
///
/// For example, if the requested columns are
/// ```text
/// ("s1", Min, Field:s1_min)
/// ("s2", Max, field:s2_max)
/// ```
///
/// And the input statistics had
/// ```text
/// S1(Min: 5, Max: 10)
/// S2(Min: 99, Max: 1000)
/// S3(Min: 1, Max: 2)
/// ```
///
/// Then this function would build a record batch with 2 columns and
/// one row s1_min and s2_max as follows (s3 is not requested):
///
/// ```text
/// s1_min | s2_max
/// -------+--------
/// 5 | 1000
/// ```
fn build_statistics_record_batch<S: PruningStatistics + ?Sized>(
statistics: &S,
required_columns: &RequiredColumns,
) -> Result<RecordBatch> {
let mut arrays = Vec::<ArrayRef>::new();
// For each needed statistics column:
for (column, statistics_type, stat_field) in required_columns.iter() {
let column = Column::from_name(column.name());
let data_type = stat_field.data_type();
let num_containers = statistics.num_containers();
let array = match statistics_type {
StatisticsType::Min => statistics.min_values(&column),
StatisticsType::Max => statistics.max_values(&column),
StatisticsType::NullCount => statistics.null_counts(&column),
StatisticsType::RowCount => statistics.row_counts(&column),
};
let array = array.unwrap_or_else(|| new_null_array(data_type, num_containers));
if num_containers != array.len() {
return internal_err!(
"mismatched statistics length. Expected {}, got {}",
num_containers,
array.len()
);
}
// cast statistics array to required data type (e.g. parquet
// provides timestamp statistics as "Int64")
let array = arrow::compute::cast(&array, data_type)?;
arrays.push(array);
}
let schema = Arc::new(required_columns.schema());
// provide the count in case there were no needed statistics
let mut options = RecordBatchOptions::default();
options.row_count = Some(statistics.num_containers());
trace!("Creating statistics batch for {required_columns:#?} with {arrays:#?}");
RecordBatch::try_new_with_options(schema, arrays, &options).map_err(|err| {
plan_datafusion_err!("Can not create statistics record batch: {err}")
})
}
struct PruningExpressionBuilder<'a> {
column: phys_expr::Column,
column_expr: Arc<dyn PhysicalExpr>,
op: Operator,
scalar_expr: Arc<dyn PhysicalExpr>,
field: &'a Field,
required_columns: &'a mut RequiredColumns,
}
impl<'a> PruningExpressionBuilder<'a> {
fn try_new(
left: &'a Arc<dyn PhysicalExpr>,
right: &'a Arc<dyn PhysicalExpr>,
op: Operator,
schema: &'a SchemaRef,
required_columns: &'a mut RequiredColumns,
) -> Result<Self> {
// find column name; input could be a more complicated expression
let left_columns = collect_columns(left);
let right_columns = collect_columns(right);
let (column_expr, scalar_expr, columns, correct_operator) =
match (left_columns.len(), right_columns.len()) {
(1, 0) => (left, right, left_columns, op),
(0, 1) => (right, left, right_columns, reverse_operator(op)?),
_ => {
// if more than one column used in expression - not supported
return plan_err!(
"Multi-column expressions are not currently supported"
);
}
};
let df_schema = DFSchema::try_from(Arc::clone(schema))?;
let (column_expr, correct_operator, scalar_expr) = rewrite_expr_to_prunable(
column_expr,
correct_operator,
scalar_expr,
df_schema,
)?;
let column = columns.iter().next().unwrap().clone();
let field = match schema.column_with_name(column.name()) {
Some((_, f)) => f,
_ => {
return plan_err!("Field not found in schema");
}
};
Ok(Self {
column,
column_expr,
op: correct_operator,
scalar_expr,
field,
required_columns,
})
}
fn op(&self) -> Operator {
self.op
}
fn scalar_expr(&self) -> &Arc<dyn PhysicalExpr> {
&self.scalar_expr
}
fn min_column_expr(&mut self) -> Result<Arc<dyn PhysicalExpr>> {
self.required_columns
.min_column_expr(&self.column, &self.column_expr, self.field)
}
fn max_column_expr(&mut self) -> Result<Arc<dyn PhysicalExpr>> {
self.required_columns
.max_column_expr(&self.column, &self.column_expr, self.field)
}
/// This function is to simply retune the `null_count` physical expression no matter what the
/// predicate expression is
///
/// i.e., x > 5 => x_null_count,
/// cast(x as int) < 10 => x_null_count,
/// try_cast(x as float) < 10.0 => x_null_count
fn null_count_column_expr(&mut self) -> Result<Arc<dyn PhysicalExpr>> {
// Retune to [`phys_expr::Column`]
let column_expr = Arc::new(self.column.clone()) as _;
// null_count is DataType::UInt64, which is different from the column's data type (i.e. self.field)
let null_count_field = &Field::new(self.field.name(), DataType::UInt64, true);
self.required_columns.null_count_column_expr(
&self.column,
&column_expr,
null_count_field,
)
}
/// This function is to simply retune the `row_count` physical expression no matter what the
/// predicate expression is
///
/// i.e., x > 5 => x_row_count,
/// cast(x as int) < 10 => x_row_count,
/// try_cast(x as float) < 10.0 => x_row_count
fn row_count_column_expr(&mut self) -> Result<Arc<dyn PhysicalExpr>> {
// Retune to [`phys_expr::Column`]
let column_expr = Arc::new(self.column.clone()) as _;
// row_count is DataType::UInt64, which is different from the column's data type (i.e. self.field)
let row_count_field = &Field::new(self.field.name(), DataType::UInt64, true);
self.required_columns.row_count_column_expr(
&self.column,
&column_expr,
row_count_field,
)
}
}
/// This function is designed to rewrite the column_expr to
/// ensure the column_expr is monotonically increasing.
///
/// For example,
/// 1. `col > 10`
/// 2. `-col > 10` should be rewritten to `col < -10`
/// 3. `!col = true` would be rewritten to `col = !true`
/// 4. `abs(a - 10) > 0` not supported
/// 5. `cast(can_prunable_expr) > 10`
/// 6. `try_cast(can_prunable_expr) > 10`
///
/// More rewrite rules are still in progress.
fn rewrite_expr_to_prunable(
column_expr: &PhysicalExprRef,
op: Operator,
scalar_expr: &PhysicalExprRef,
schema: DFSchema,
) -> Result<(PhysicalExprRef, Operator, PhysicalExprRef)> {
if !is_compare_op(op) {
return plan_err!("rewrite_expr_to_prunable only support compare expression");
}
let column_expr_any = column_expr.as_any();
if column_expr_any
.downcast_ref::<phys_expr::Column>()
.is_some()
{
// `col op lit()`
Ok((Arc::clone(column_expr), op, Arc::clone(scalar_expr)))
} else if let Some(cast) = column_expr_any.downcast_ref::<phys_expr::CastExpr>() {
// `cast(col) op lit()`
let arrow_schema = schema.as_arrow();
let from_type = cast.expr().data_type(arrow_schema)?;
verify_support_type_for_prune(&from_type, cast.cast_type())?;
let (left, op, right) =
rewrite_expr_to_prunable(cast.expr(), op, scalar_expr, schema)?;
let left = Arc::new(phys_expr::CastExpr::new(
left,
cast.cast_type().clone(),
None,
));
Ok((left, op, right))
} else if let Some(try_cast) =
column_expr_any.downcast_ref::<phys_expr::TryCastExpr>()
{
// `try_cast(col) op lit()`
let arrow_schema = schema.as_arrow();
let from_type = try_cast.expr().data_type(arrow_schema)?;
verify_support_type_for_prune(&from_type, try_cast.cast_type())?;
let (left, op, right) =
rewrite_expr_to_prunable(try_cast.expr(), op, scalar_expr, schema)?;
let left = Arc::new(phys_expr::TryCastExpr::new(
left,
try_cast.cast_type().clone(),
));
Ok((left, op, right))
} else if let Some(neg) = column_expr_any.downcast_ref::<phys_expr::NegativeExpr>() {
// `-col > lit()` --> `col < -lit()`
let (left, op, right) =
rewrite_expr_to_prunable(neg.arg(), op, scalar_expr, schema)?;
let right = Arc::new(phys_expr::NegativeExpr::new(right));
Ok((left, reverse_operator(op)?, right))
} else if let Some(not) = column_expr_any.downcast_ref::<phys_expr::NotExpr>() {
// `!col = true` --> `col = !true`
if op != Operator::Eq && op != Operator::NotEq {
return plan_err!("Not with operator other than Eq / NotEq is not supported");
}
if not
.arg()
.as_any()
.downcast_ref::<phys_expr::Column>()
.is_some()
{
let left = Arc::clone(not.arg());
let right = Arc::new(phys_expr::NotExpr::new(Arc::clone(scalar_expr)));
Ok((left, reverse_operator(op)?, right))
} else {
plan_err!("Not with complex expression {column_expr:?} is not supported")
}
} else {
plan_err!("column expression {column_expr:?} is not supported")
}
}
fn is_compare_op(op: Operator) -> bool {
matches!(
op,
Operator::Eq
| Operator::NotEq
| Operator::Lt
| Operator::LtEq
| Operator::Gt
| Operator::GtEq
| Operator::LikeMatch
| Operator::NotLikeMatch
)
}
fn is_string_type(data_type: &DataType) -> bool {
matches!(
data_type,
DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View
)
}
// The pruning logic is based on the comparing the min/max bounds.
// Must make sure the two type has order.
// For example, casts from string to numbers is not correct.
// Because the "13" is less than "3" with UTF8 comparison order.
fn verify_support_type_for_prune(from_type: &DataType, to_type: &DataType) -> Result<()> {
// Dictionary casts are always supported as long as the value types are supported
let from_type = match from_type {
DataType::Dictionary(_, t) => {
return verify_support_type_for_prune(t.as_ref(), to_type)
}
_ => from_type,
};
let to_type = match to_type {
DataType::Dictionary(_, t) => {
return verify_support_type_for_prune(from_type, t.as_ref())
}
_ => to_type,
};
// If both types are strings or both are not strings (number, timestamp, etc)
// then we can compare them.
// PruningPredicate does not support casting of strings to numbers and such.
if is_string_type(from_type) == is_string_type(to_type) {
Ok(())
} else {
plan_err!(
"Try Cast/Cast with from type {from_type} to type {to_type} is not supported"
)
}
}
/// replaces a column with an old name with a new name in an expression
fn rewrite_column_expr(
e: Arc<dyn PhysicalExpr>,
column_old: &phys_expr::Column,
column_new: &phys_expr::Column,
) -> Result<Arc<dyn PhysicalExpr>> {
e.transform(|expr| {
if let Some(column) = expr.as_any().downcast_ref::<phys_expr::Column>() {
if column == column_old {
return Ok(Transformed::yes(Arc::new(column_new.clone())));
}
}
Ok(Transformed::no(expr))
})
.data()
}
fn reverse_operator(op: Operator) -> Result<Operator> {
op.swap().ok_or_else(|| {
internal_datafusion_err!(
"Could not reverse operator {op} while building pruning predicate"
)
})
}
/// Given a column reference to `column`, returns a pruning
/// expression in terms of the min and max that will evaluate to true
/// if the column may contain values, and false if definitely does not
/// contain values
fn build_single_column_expr(
column: &phys_expr::Column,
schema: &Schema,
required_columns: &mut RequiredColumns,
is_not: bool, // if true, treat as !col
) -> Option<Arc<dyn PhysicalExpr>> {
let field = schema.field_with_name(column.name()).ok()?;
if matches!(field.data_type(), &DataType::Boolean) {
let col_ref = Arc::new(column.clone()) as _;
let min = required_columns
.min_column_expr(column, &col_ref, field)
.ok()?;
let max = required_columns
.max_column_expr(column, &col_ref, field)
.ok()?;
// remember -- we want an expression that is:
// TRUE: if there may be rows that match
// FALSE: if there are no rows that match
if is_not {
// The only way we know a column couldn't match is if both the min and max are true
// !(min && max)
Some(Arc::new(phys_expr::NotExpr::new(Arc::new(
phys_expr::BinaryExpr::new(min, Operator::And, max),
))))
} else {
// the only way we know a column couldn't match is if both the min and max are false
// !(!min && !max) --> min || max
Some(Arc::new(phys_expr::BinaryExpr::new(min, Operator::Or, max)))
}
} else {
None
}
}
/// Given an expression reference to `expr`, if `expr` is a column expression,
/// returns a pruning expression in terms of IsNull that will evaluate to true
/// if the column may contain null, and false if definitely does not
/// contain null.
/// If `with_not` is true, build a pruning expression for `col IS NOT NULL`: `col_count != col_null_count`
/// The pruning expression evaluates to true ONLY if the column definitely CONTAINS
/// at least one NULL value. In this case we can know that `IS NOT NULL` can not be true and
/// thus can prune the row group / value
fn build_is_null_column_expr(
expr: &Arc<dyn PhysicalExpr>,
schema: &Schema,
required_columns: &mut RequiredColumns,
with_not: bool,
) -> Option<Arc<dyn PhysicalExpr>> {
if let Some(col) = expr.as_any().downcast_ref::<phys_expr::Column>() {
let field = schema.field_with_name(col.name()).ok()?;
let null_count_field = &Field::new(field.name(), DataType::UInt64, true);
if with_not {
if let Ok(row_count_expr) =
required_columns.row_count_column_expr(col, expr, null_count_field)
{
required_columns
.null_count_column_expr(col, expr, null_count_field)
.map(|null_count_column_expr| {
// IsNotNull(column) => null_count != row_count
Arc::new(phys_expr::BinaryExpr::new(
null_count_column_expr,
Operator::NotEq,
row_count_expr,
)) as _
})
.ok()
} else {
None
}
} else {
required_columns
.null_count_column_expr(col, expr, null_count_field)
.map(|null_count_column_expr| {
// IsNull(column) => null_count > 0
Arc::new(phys_expr::BinaryExpr::new(
null_count_column_expr,
Operator::Gt,
Arc::new(phys_expr::Literal::new(ScalarValue::UInt64(Some(0)))),
)) as _
})
.ok()
}
} else {
None
}
}
/// The maximum number of entries in an `InList` that might be rewritten into
/// an OR chain
const MAX_LIST_VALUE_SIZE_REWRITE: usize = 20;
/// Rewrite a predicate expression in terms of statistics (min/max/null_counts)
/// for use as a [`PruningPredicate`].
pub struct PredicateRewriter {
unhandled_hook: Arc<dyn UnhandledPredicateHook>,
}
impl Default for PredicateRewriter {
fn default() -> Self {
Self {
unhandled_hook: Arc::new(ConstantUnhandledPredicateHook::default()),
}
}
}
impl PredicateRewriter {
/// Create a new `PredicateRewriter`
pub fn new() -> Self {
Self::default()
}
/// Set the unhandled hook to be used when a predicate can not be rewritten
pub fn with_unhandled_hook(
self,
unhandled_hook: Arc<dyn UnhandledPredicateHook>,
) -> Self {
Self { unhandled_hook }
}
/// Translate logical filter expression into pruning predicate
/// expression that will evaluate to FALSE if it can be determined no
/// rows between the min/max values could pass the predicates.
///
/// Any predicates that can not be translated will be passed to `unhandled_hook`.
///
/// Returns the pruning predicate as an [`PhysicalExpr`]
///
/// Notice: Does not handle [`phys_expr::InListExpr`] greater than 20, which will fall back to calling `unhandled_hook`
pub fn rewrite_predicate_to_statistics_predicate(
&self,
expr: &Arc<dyn PhysicalExpr>,
schema: &Schema,
) -> Arc<dyn PhysicalExpr> {
let mut required_columns = RequiredColumns::new();
build_predicate_expression(
expr,
&Arc::new(schema.clone()),
&mut required_columns,
&self.unhandled_hook,
)
}
}
/// Translate logical filter expression into pruning predicate
/// expression that will evaluate to FALSE if it can be determined no
/// rows between the min/max values could pass the predicates.
///
/// Any predicates that can not be translated will be passed to `unhandled_hook`.
///
/// Returns the pruning predicate as an [`PhysicalExpr`]
///
/// Notice: Does not handle [`phys_expr::InListExpr`] greater than 20, which will fall back to calling `unhandled_hook`
fn build_predicate_expression(
expr: &Arc<dyn PhysicalExpr>,
schema: &SchemaRef,
required_columns: &mut RequiredColumns,
unhandled_hook: &Arc<dyn UnhandledPredicateHook>,
) -> Arc<dyn PhysicalExpr> {
if is_always_false(expr) {
// Shouldn't return `unhandled_hook.handle(expr)`
// Because it will transfer false to true.
return Arc::clone(expr);
}
// predicate expression can only be a binary expression
let expr_any = expr.as_any();
if let Some(is_null) = expr_any.downcast_ref::<phys_expr::IsNullExpr>() {
return build_is_null_column_expr(is_null.arg(), schema, required_columns, false)
.unwrap_or_else(|| unhandled_hook.handle(expr));
}
if let Some(is_not_null) = expr_any.downcast_ref::<phys_expr::IsNotNullExpr>() {
return build_is_null_column_expr(
is_not_null.arg(),
schema,
required_columns,
true,
)
.unwrap_or_else(|| unhandled_hook.handle(expr));
}
if let Some(col) = expr_any.downcast_ref::<phys_expr::Column>() {
return build_single_column_expr(col, schema, required_columns, false)
.unwrap_or_else(|| unhandled_hook.handle(expr));
}
if let Some(not) = expr_any.downcast_ref::<phys_expr::NotExpr>() {
// match !col (don't do so recursively)
if let Some(col) = not.arg().as_any().downcast_ref::<phys_expr::Column>() {
return build_single_column_expr(col, schema, required_columns, true)
.unwrap_or_else(|| unhandled_hook.handle(expr));
} else {
return unhandled_hook.handle(expr);
}
}
if let Some(in_list) = expr_any.downcast_ref::<phys_expr::InListExpr>() {
if !in_list.list().is_empty()
&& in_list.list().len() <= MAX_LIST_VALUE_SIZE_REWRITE
{
let eq_op = if in_list.negated() {
Operator::NotEq
} else {
Operator::Eq
};
let re_op = if in_list.negated() {
Operator::And
} else {
Operator::Or
};
let change_expr = in_list
.list()
.iter()
.map(|e| {
Arc::new(phys_expr::BinaryExpr::new(
Arc::clone(in_list.expr()),
eq_op,
Arc::clone(e),
)) as _
})
.reduce(|a, b| Arc::new(phys_expr::BinaryExpr::new(a, re_op, b)) as _)
.unwrap();
return build_predicate_expression(
&change_expr,
schema,
required_columns,
unhandled_hook,
);
} else {
return unhandled_hook.handle(expr);
}
}
let (left, op, right) = {
if let Some(bin_expr) = expr_any.downcast_ref::<phys_expr::BinaryExpr>() {
(
Arc::clone(bin_expr.left()),
*bin_expr.op(),
Arc::clone(bin_expr.right()),
)
} else if let Some(like_expr) = expr_any.downcast_ref::<phys_expr::LikeExpr>() {
if like_expr.case_insensitive() {
return unhandled_hook.handle(expr);
}
let op = match (like_expr.negated(), like_expr.case_insensitive()) {
(false, false) => Operator::LikeMatch,
(true, false) => Operator::NotLikeMatch,
(false, true) => Operator::ILikeMatch,
(true, true) => Operator::NotILikeMatch,
};
(
Arc::clone(like_expr.expr()),
op,
Arc::clone(like_expr.pattern()),
)
} else {
return unhandled_hook.handle(expr);
}
};
if op == Operator::And || op == Operator::Or {
let left_expr =
build_predicate_expression(&left, schema, required_columns, unhandled_hook);
let right_expr =
build_predicate_expression(&right, schema, required_columns, unhandled_hook);
// simplify boolean expression if applicable
let expr = match (&left_expr, op, &right_expr) {
(left, Operator::And, right)
if is_always_false(left) || is_always_false(right) =>
{
Arc::new(phys_expr::Literal::new(ScalarValue::Boolean(Some(false))))
}
(left, Operator::And, _) if is_always_true(left) => right_expr,
(_, Operator::And, right) if is_always_true(right) => left_expr,
(left, Operator::Or, right)
if is_always_true(left) || is_always_true(right) =>
{
Arc::new(phys_expr::Literal::new(ScalarValue::Boolean(Some(true))))
}
(left, Operator::Or, _) if is_always_false(left) => right_expr,
(_, Operator::Or, right) if is_always_false(right) => left_expr,
_ => Arc::new(phys_expr::BinaryExpr::new(left_expr, op, right_expr)),
};
return expr;
}
let expr_builder =
PruningExpressionBuilder::try_new(&left, &right, op, schema, required_columns);
let mut expr_builder = match expr_builder {
Ok(builder) => builder,
// allow partial failure in predicate expression generation
// this can still produce a useful predicate when multiple conditions are joined using AND
Err(e) => {
debug!("Error building pruning expression: {e}");
return unhandled_hook.handle(expr);
}
};
build_statistics_expr(&mut expr_builder)
.unwrap_or_else(|_| unhandled_hook.handle(expr))
}
fn build_statistics_expr(
expr_builder: &mut PruningExpressionBuilder,
) -> Result<Arc<dyn PhysicalExpr>> {
let statistics_expr: Arc<dyn PhysicalExpr> = match expr_builder.op() {
Operator::NotEq => {
// column != literal => (min, max) = literal =>
// !(min != literal && max != literal) ==>
// min != literal || literal != max
let min_column_expr = expr_builder.min_column_expr()?;
let max_column_expr = expr_builder.max_column_expr()?;
Arc::new(phys_expr::BinaryExpr::new(
Arc::new(phys_expr::BinaryExpr::new(
min_column_expr,
Operator::NotEq,
Arc::clone(expr_builder.scalar_expr()),
)),
Operator::Or,
Arc::new(phys_expr::BinaryExpr::new(
Arc::clone(expr_builder.scalar_expr()),
Operator::NotEq,
max_column_expr,
)),
))
}
Operator::Eq => {
// column = literal => (min, max) = literal => min <= literal && literal <= max
// (column / 2) = 4 => (column_min / 2) <= 4 && 4 <= (column_max / 2)
let min_column_expr = expr_builder.min_column_expr()?;
let max_column_expr = expr_builder.max_column_expr()?;
Arc::new(phys_expr::BinaryExpr::new(
Arc::new(phys_expr::BinaryExpr::new(
min_column_expr,
Operator::LtEq,
Arc::clone(expr_builder.scalar_expr()),
)),
Operator::And,
Arc::new(phys_expr::BinaryExpr::new(
Arc::clone(expr_builder.scalar_expr()),
Operator::LtEq,
max_column_expr,
)),
))
}
Operator::NotLikeMatch => build_not_like_match(expr_builder)?,
Operator::LikeMatch => build_like_match(expr_builder).ok_or_else(|| {
plan_datafusion_err!(
"LIKE expression with wildcard at the beginning is not supported"
)
})?,
Operator::Gt => {
// column > literal => (min, max) > literal => max > literal
Arc::new(phys_expr::BinaryExpr::new(
expr_builder.max_column_expr()?,
Operator::Gt,
Arc::clone(expr_builder.scalar_expr()),
))
}
Operator::GtEq => {
// column >= literal => (min, max) >= literal => max >= literal
Arc::new(phys_expr::BinaryExpr::new(
expr_builder.max_column_expr()?,
Operator::GtEq,
Arc::clone(expr_builder.scalar_expr()),
))
}
Operator::Lt => {
// column < literal => (min, max) < literal => min < literal
Arc::new(phys_expr::BinaryExpr::new(
expr_builder.min_column_expr()?,
Operator::Lt,
Arc::clone(expr_builder.scalar_expr()),
))
}
Operator::LtEq => {
// column <= literal => (min, max) <= literal => min <= literal
Arc::new(phys_expr::BinaryExpr::new(
expr_builder.min_column_expr()?,
Operator::LtEq,
Arc::clone(expr_builder.scalar_expr()),
))
}
// other expressions are not supported
_ => {
return plan_err!(
"expressions other than (neq, eq, gt, gteq, lt, lteq) are not supported"
);
}
};
let statistics_expr = wrap_null_count_check_expr(statistics_expr, expr_builder)?;
Ok(statistics_expr)
}
/// returns the string literal of the scalar value if it is a string
fn unpack_string(s: &ScalarValue) -> Option<&str> {
s.try_as_str().flatten()
}
fn extract_string_literal(expr: &Arc<dyn PhysicalExpr>) -> Option<&str> {
if let Some(lit) = expr.as_any().downcast_ref::<phys_expr::Literal>() {
let s = unpack_string(lit.value())?;
return Some(s);
}
None
}
/// Convert `column LIKE literal` where P is a constant prefix of the literal
/// to a range check on the column: `P <= column && column < P'`, where P' is the
/// lowest string after all P* strings.
fn build_like_match(
expr_builder: &mut PruningExpressionBuilder,
) -> Option<Arc<dyn PhysicalExpr>> {
// column LIKE literal => (min, max) LIKE literal split at % => min <= split literal && split literal <= max
// column LIKE 'foo%' => min <= 'foo' && 'foo' <= max
// column LIKE '%foo' => min <= '' && '' <= max => true
// column LIKE '%foo%' => min <= '' && '' <= max => true
// column LIKE 'foo' => min <= 'foo' && 'foo' <= max
// TODO Handle ILIKE perhaps by making the min lowercase and max uppercase
// this may involve building the physical expressions that call lower() and upper()
let min_column_expr = expr_builder.min_column_expr().ok()?;
let max_column_expr = expr_builder.max_column_expr().ok()?;
let scalar_expr = expr_builder.scalar_expr();
// check that the scalar is a string literal
let s = extract_string_literal(scalar_expr)?;
// ANSI SQL specifies two wildcards: % and _. % matches zero or more characters, _ matches exactly one character.
let first_wildcard_index = s.find(['%', '_']);
if first_wildcard_index == Some(0) {
// there's no filtering we could possibly do, return an error and have this be handled by the unhandled hook
return None;
}
let (lower_bound, upper_bound) = if let Some(wildcard_index) = first_wildcard_index {
let prefix = &s[..wildcard_index];
let lower_bound_lit = Arc::new(phys_expr::Literal::new(ScalarValue::Utf8(Some(
prefix.to_string(),
))));
let upper_bound_lit = Arc::new(phys_expr::Literal::new(ScalarValue::Utf8(Some(
increment_utf8(prefix)?,
))));
(lower_bound_lit, upper_bound_lit)
} else {
// the like expression is a literal and can be converted into a comparison
let bound = Arc::new(phys_expr::Literal::new(ScalarValue::Utf8(Some(
s.to_string(),
))));
(Arc::clone(&bound), bound)
};
let lower_bound_expr = Arc::new(phys_expr::BinaryExpr::new(
lower_bound,
Operator::LtEq,
Arc::clone(&max_column_expr),
));
let upper_bound_expr = Arc::new(phys_expr::BinaryExpr::new(
Arc::clone(&min_column_expr),
Operator::LtEq,
upper_bound,
));
let combined = Arc::new(phys_expr::BinaryExpr::new(
upper_bound_expr,
Operator::And,
lower_bound_expr,
));
Some(combined)
}
// For predicate `col NOT LIKE 'const_prefix%'`, we rewrite it as `(col_min NOT LIKE 'const_prefix%' OR col_max NOT LIKE 'const_prefix%')`.
//
// The intuition is that if both `col_min` and `col_max` begin with `const_prefix` that means
// **all** data in this row group begins with `const_prefix` as well (and therefore the predicate
// looking for rows that don't begin with `const_prefix` can never be true)
fn build_not_like_match(
expr_builder: &mut PruningExpressionBuilder<'_>,
) -> Result<Arc<dyn PhysicalExpr>> {
// 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%')
let min_column_expr = expr_builder.min_column_expr()?;
let max_column_expr = expr_builder.max_column_expr()?;
let scalar_expr = expr_builder.scalar_expr();
let pattern = extract_string_literal(scalar_expr).ok_or_else(|| {
plan_datafusion_err!("cannot extract literal from NOT LIKE expression")
})?;
let (const_prefix, remaining) = split_constant_prefix(pattern);
if const_prefix.is_empty() || remaining != "%" {
// we can not handle `%` at the beginning or in the middle of the pattern
// Example: For pattern "foo%bar", the row group might include values like
// ["foobar", "food", "foodbar"], making it unsafe to prune.
// Even if the min/max values in the group (e.g., "foobar" and "foodbar")
// match the pattern, intermediate values like "food" may not
// match the full pattern "foo%bar", making pruning unsafe.
// (truncate foo%bar to foo% have same problem)
// we can not handle pattern containing `_`
// Example: For pattern "foo_", row groups might contain ["fooa", "fooaa", "foob"],
// which means not every row is guaranteed to match the pattern.
return Err(plan_datafusion_err!(
"NOT LIKE expressions only support constant_prefix+wildcard`%`"
));
}
let min_col_not_like_epxr = Arc::new(phys_expr::LikeExpr::new(
true,
false,
Arc::clone(&min_column_expr),
Arc::clone(scalar_expr),
));
let max_col_not_like_expr = Arc::new(phys_expr::LikeExpr::new(
true,
false,
Arc::clone(&max_column_expr),
Arc::clone(scalar_expr),
));
Ok(Arc::new(phys_expr::BinaryExpr::new(
min_col_not_like_epxr,
Operator::Or,
max_col_not_like_expr,
)))
}
/// Returns unescaped constant prefix of a LIKE pattern (possibly empty) and the remaining pattern (possibly empty)
fn split_constant_prefix(pattern: &str) -> (&str, &str) {
let char_indices = pattern.char_indices().collect::<Vec<_>>();
for i in 0..char_indices.len() {
let (idx, char) = char_indices[i];
if char == '%' || char == '_' {
if i != 0 && char_indices[i - 1].1 == '\\' {
// ecsaped by `\`
continue;
}
return (&pattern[..idx], &pattern[idx..]);
}
}
(pattern, "")
}
/// Increment a UTF8 string by one, returning `None` if it can't be incremented.
/// This makes it so that the returned string will always compare greater than the input string
/// or any other string with the same prefix.
/// This is necessary since the statistics may have been truncated: if we have a min statistic
/// of "fo" that may have originally been "foz" or anything else with the prefix "fo".
/// E.g. `increment_utf8("foo") >= "foo"` and `increment_utf8("foo") >= "fooz"`
/// In this example `increment_utf8("foo") == "fop"
fn increment_utf8(data: &str) -> Option<String> {
// Helper function to check if a character is valid to use
fn is_valid_unicode(c: char) -> bool {
let cp = c as u32;
// Filter out non-characters (https://www.unicode.org/versions/corrigendum9.html)
if [0xFFFE, 0xFFFF].contains(&cp) || (0xFDD0..=0xFDEF).contains(&cp) {
return false;
}
// Filter out private use area
if cp >= 0x110000 {
return false;
}
true
}
// Convert string to vector of code points
let mut code_points: Vec<char> = data.chars().collect();
// Work backwards through code points
for idx in (0..code_points.len()).rev() {
let original = code_points[idx] as u32;
// Try incrementing the code point
if let Some(next_char) = char::from_u32(original + 1) {
if is_valid_unicode(next_char) {
code_points[idx] = next_char;
// truncate the string to the current index
code_points.truncate(idx + 1);
return Some(code_points.into_iter().collect());
}
}
}
None
}
/// Wrap the statistics expression in a check that skips the expression if the column is all nulls.
///
/// This is important not only as an optimization but also because statistics may not be
/// accurate for columns that are all nulls.
/// For example, for an `int` column `x` with all nulls, the min/max/null_count statistics
/// might be set to 0 and evaluating `x = 0` would incorrectly include the column.
///
/// For example:
///
/// `x_min <= 10 AND 10 <= x_max`
///
/// will become
///
/// ```sql
/// x_null_count != x_row_count AND (x_min <= 10 AND 10 <= x_max)
/// ````
///
/// If the column is known to be all nulls, then the expression
/// `x_null_count = x_row_count` will be true, which will cause the
/// boolean expression to return false. Therefore, prune out the container.
fn wrap_null_count_check_expr(
statistics_expr: Arc<dyn PhysicalExpr>,
expr_builder: &mut PruningExpressionBuilder,
) -> Result<Arc<dyn PhysicalExpr>> {
// x_null_count != x_row_count
let not_when_null_count_eq_row_count = Arc::new(phys_expr::BinaryExpr::new(
expr_builder.null_count_column_expr()?,
Operator::NotEq,
expr_builder.row_count_column_expr()?,
));
// (x_null_count != x_row_count) AND (<statistics_expr>)
Ok(Arc::new(phys_expr::BinaryExpr::new(
not_when_null_count_eq_row_count,
Operator::And,
statistics_expr,
)))
}
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub(crate) enum StatisticsType {
Min,
Max,
NullCount,
RowCount,
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use std::ops::{Not, Rem};
use super::*;
use datafusion_common::test_util::batches_to_string;
use datafusion_expr::{and, col, lit, or};
use insta::assert_snapshot;
use arrow::array::Decimal128Array;
use arrow::{
array::{BinaryArray, Int32Array, Int64Array, StringArray, UInt64Array},
datatypes::TimeUnit,
};
use datafusion_expr::expr::InList;
use datafusion_expr::{cast, is_null, try_cast, Expr};
use datafusion_functions_nested::expr_fn::{array_has, make_array};
use datafusion_physical_expr::expressions as phys_expr;
use datafusion_physical_expr::planner::logical2physical;
#[derive(Debug, Default)]
/// Mock statistic provider for tests
///
/// Each row represents the statistics for a "container" (which
/// might represent an entire parquet file, or directory of files,
/// or some other collection of data for which we had statistics)
///
/// Note All `ArrayRefs` must be the same size.
struct ContainerStats {
min: Option<ArrayRef>,
max: Option<ArrayRef>,
/// Optional values
null_counts: Option<ArrayRef>,
row_counts: Option<ArrayRef>,
/// Optional known values (e.g. mimic a bloom filter)
/// (value, contained)
/// If present, all BooleanArrays must be the same size as min/max
contained: Vec<(HashSet<ScalarValue>, BooleanArray)>,
}
impl ContainerStats {
fn new() -> Self {
Default::default()
}
fn new_decimal128(
min: impl IntoIterator<Item = Option<i128>>,
max: impl IntoIterator<Item = Option<i128>>,
precision: u8,
scale: i8,
) -> Self {
Self::new()
.with_min(Arc::new(
min.into_iter()
.collect::<Decimal128Array>()
.with_precision_and_scale(precision, scale)
.unwrap(),
))
.with_max(Arc::new(
max.into_iter()
.collect::<Decimal128Array>()
.with_precision_and_scale(precision, scale)
.unwrap(),
))
}
fn new_i64(
min: impl IntoIterator<Item = Option<i64>>,
max: impl IntoIterator<Item = Option<i64>>,
) -> Self {
Self::new()
.with_min(Arc::new(min.into_iter().collect::<Int64Array>()))
.with_max(Arc::new(max.into_iter().collect::<Int64Array>()))
}
fn new_i32(
min: impl IntoIterator<Item = Option<i32>>,
max: impl IntoIterator<Item = Option<i32>>,
) -> Self {
Self::new()
.with_min(Arc::new(min.into_iter().collect::<Int32Array>()))
.with_max(Arc::new(max.into_iter().collect::<Int32Array>()))
}
fn new_utf8<'a>(
min: impl IntoIterator<Item = Option<&'a str>>,
max: impl IntoIterator<Item = Option<&'a str>>,
) -> Self {
Self::new()
.with_min(Arc::new(min.into_iter().collect::<StringArray>()))
.with_max(Arc::new(max.into_iter().collect::<StringArray>()))
}
fn new_bool(
min: impl IntoIterator<Item = Option<bool>>,
max: impl IntoIterator<Item = Option<bool>>,
) -> Self {
Self::new()
.with_min(Arc::new(min.into_iter().collect::<BooleanArray>()))
.with_max(Arc::new(max.into_iter().collect::<BooleanArray>()))
}
fn min(&self) -> Option<ArrayRef> {
self.min.clone()
}
fn max(&self) -> Option<ArrayRef> {
self.max.clone()
}
fn null_counts(&self) -> Option<ArrayRef> {
self.null_counts.clone()
}
fn row_counts(&self) -> Option<ArrayRef> {
self.row_counts.clone()
}
/// return an iterator over all arrays in this statistics
fn arrays(&self) -> Vec<ArrayRef> {
let contained_arrays = self
.contained
.iter()
.map(|(_values, contained)| Arc::new(contained.clone()) as ArrayRef);
[
self.min.as_ref().cloned(),
self.max.as_ref().cloned(),
self.null_counts.as_ref().cloned(),
self.row_counts.as_ref().cloned(),
]
.into_iter()
.flatten()
.chain(contained_arrays)
.collect()
}
/// Returns the number of containers represented by this statistics This
/// picks the length of the first array as all arrays must have the same
/// length (which is verified by `assert_invariants`).
fn len(&self) -> usize {
// pick the first non zero length
self.arrays().iter().map(|a| a.len()).next().unwrap_or(0)
}
/// Ensure that the lengths of all arrays are consistent
fn assert_invariants(&self) {
let mut prev_len = None;
for len in self.arrays().iter().map(|a| a.len()) {
// Get a length, if we don't already have one
match prev_len {
None => {
prev_len = Some(len);
}
Some(prev_len) => {
assert_eq!(prev_len, len);
}
}
}
}
/// Add min values
fn with_min(mut self, min: ArrayRef) -> Self {
self.min = Some(min);
self
}
/// Add max values
fn with_max(mut self, max: ArrayRef) -> Self {
self.max = Some(max);
self
}
/// Add null counts. There must be the same number of null counts as
/// there are containers
fn with_null_counts(
mut self,
counts: impl IntoIterator<Item = Option<u64>>,
) -> Self {
let null_counts: ArrayRef =
Arc::new(counts.into_iter().collect::<UInt64Array>());
self.assert_invariants();
self.null_counts = Some(null_counts);
self
}
/// Add row counts. There must be the same number of row counts as
/// there are containers
fn with_row_counts(
mut self,
counts: impl IntoIterator<Item = Option<u64>>,
) -> Self {
let row_counts: ArrayRef =
Arc::new(counts.into_iter().collect::<UInt64Array>());
self.assert_invariants();
self.row_counts = Some(row_counts);
self
}
/// Add contained information.
pub fn with_contained(
mut self,
values: impl IntoIterator<Item = ScalarValue>,
contained: impl IntoIterator<Item = Option<bool>>,
) -> Self {
let contained: BooleanArray = contained.into_iter().collect();
let values: HashSet<_> = values.into_iter().collect();
self.contained.push((values, contained));
self.assert_invariants();
self
}
/// get any contained information for the specified values
fn contained(&self, find_values: &HashSet<ScalarValue>) -> Option<BooleanArray> {
// find the one with the matching values
self.contained
.iter()
.find(|(values, _contained)| values == find_values)
.map(|(_values, contained)| contained.clone())
}
}
#[derive(Debug, Default)]
struct TestStatistics {
// key: column name
stats: HashMap<Column, ContainerStats>,
}
impl TestStatistics {
fn new() -> Self {
Self::default()
}
fn with(
mut self,
name: impl Into<String>,
container_stats: ContainerStats,
) -> Self {
let col = Column::from_name(name.into());
self.stats.insert(col, container_stats);
self
}
/// Add null counts for the specified column.
/// There must be the same number of null counts as
/// there are containers
fn with_null_counts(
mut self,
name: impl Into<String>,
counts: impl IntoIterator<Item = Option<u64>>,
) -> Self {
let col = Column::from_name(name.into());
// take stats out and update them
let container_stats = self
.stats
.remove(&col)
.unwrap_or_default()
.with_null_counts(counts);
// put stats back in
self.stats.insert(col, container_stats);
self
}
/// Add row counts for the specified column.
/// There must be the same number of row counts as
/// there are containers
fn with_row_counts(
mut self,
name: impl Into<String>,
counts: impl IntoIterator<Item = Option<u64>>,
) -> Self {
let col = Column::from_name(name.into());
// take stats out and update them
let container_stats = self
.stats
.remove(&col)
.unwrap_or_default()
.with_row_counts(counts);
// put stats back in
self.stats.insert(col, container_stats);
self
}
/// Add contained information for the specified column.
fn with_contained(
mut self,
name: impl Into<String>,
values: impl IntoIterator<Item = ScalarValue>,
contained: impl IntoIterator<Item = Option<bool>>,
) -> Self {
let col = Column::from_name(name.into());
// take stats out and update them
let container_stats = self
.stats
.remove(&col)
.unwrap_or_default()
.with_contained(values, contained);
// put stats back in
self.stats.insert(col, container_stats);
self
}
}
impl PruningStatistics for TestStatistics {
fn min_values(&self, column: &Column) -> Option<ArrayRef> {
self.stats
.get(column)
.map(|container_stats| container_stats.min())
.unwrap_or(None)
}
fn max_values(&self, column: &Column) -> Option<ArrayRef> {
self.stats
.get(column)
.map(|container_stats| container_stats.max())
.unwrap_or(None)
}
fn num_containers(&self) -> usize {
self.stats
.values()
.next()
.map(|container_stats| container_stats.len())
.unwrap_or(0)
}
fn null_counts(&self, column: &Column) -> Option<ArrayRef> {
self.stats
.get(column)
.map(|container_stats| container_stats.null_counts())
.unwrap_or(None)
}
fn row_counts(&self, column: &Column) -> Option<ArrayRef> {
self.stats
.get(column)
.map(|container_stats| container_stats.row_counts())
.unwrap_or(None)
}
fn contained(
&self,
column: &Column,
values: &HashSet<ScalarValue>,
) -> Option<BooleanArray> {
self.stats
.get(column)
.and_then(|container_stats| container_stats.contained(values))
}
}
/// Returns the specified min/max container values
struct OneContainerStats {
min_values: Option<ArrayRef>,
max_values: Option<ArrayRef>,
num_containers: usize,
}
impl PruningStatistics for OneContainerStats {
fn min_values(&self, _column: &Column) -> Option<ArrayRef> {
self.min_values.clone()
}
fn max_values(&self, _column: &Column) -> Option<ArrayRef> {
self.max_values.clone()
}
fn num_containers(&self) -> usize {
self.num_containers
}
fn null_counts(&self, _column: &Column) -> Option<ArrayRef> {
None
}
fn row_counts(&self, _column: &Column) -> Option<ArrayRef> {
None
}
fn contained(
&self,
_column: &Column,
_values: &HashSet<ScalarValue>,
) -> Option<BooleanArray> {
None
}
}
/// Row count should only be referenced once in the pruning expression, even if we need the row count
/// for multiple columns.
#[test]
fn test_unique_row_count_field_and_column() {
// c1 = 100 AND c2 = 200
let schema: SchemaRef = Arc::new(Schema::new(vec![
Field::new("c1", DataType::Int32, true),
Field::new("c2", DataType::Int32, true),
]));
let expr = col("c1").eq(lit(100)).and(col("c2").eq(lit(200)));
let expr = logical2physical(&expr, &schema);
let p = PruningPredicate::try_new(expr, Arc::clone(&schema)).unwrap();
// note pruning expression refers to row_count twice
assert_eq!(
"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",
p.predicate_expr.to_string()
);
// Fields in required schema should be unique, otherwise when creating batches
// it will fail because of duplicate field names
let mut fields = HashSet::new();
for (_col, _ty, field) in p.required_columns().iter() {
let was_new = fields.insert(field);
if !was_new {
panic!(
"Duplicate field in required schema: {field:?}. Previous fields:\n{fields:#?}"
);
}
}
}
#[test]
fn prune_all_rows_null_counts() {
// if null_count = row_count then we should prune the container for i = 0
// regardless of the statistics
let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
let statistics = TestStatistics::new().with(
"i",
ContainerStats::new_i32(
vec![Some(0)], // min
vec![Some(0)], // max
)
.with_null_counts(vec![Some(1)])
.with_row_counts(vec![Some(1)]),
);
let expected_ret = &[false];
prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
// this should be true even if the container stats are missing
let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
let container_stats = ContainerStats {
min: Some(Arc::new(Int32Array::from(vec![None]))),
max: Some(Arc::new(Int32Array::from(vec![None]))),
null_counts: Some(Arc::new(UInt64Array::from(vec![Some(1)]))),
row_counts: Some(Arc::new(UInt64Array::from(vec![Some(1)]))),
..ContainerStats::default()
};
let statistics = TestStatistics::new().with("i", container_stats);
let expected_ret = &[false];
prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
// If the null counts themselves are missing we should be able to fall back to the stats
let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
let container_stats = ContainerStats {
min: Some(Arc::new(Int32Array::from(vec![Some(0)]))),
max: Some(Arc::new(Int32Array::from(vec![Some(0)]))),
null_counts: Some(Arc::new(UInt64Array::from(vec![None]))),
row_counts: Some(Arc::new(UInt64Array::from(vec![Some(1)]))),
..ContainerStats::default()
};
let statistics = TestStatistics::new().with("i", container_stats);
let expected_ret = &[true];
prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
let expected_ret = &[false];
prune_with_expr(col("i").gt(lit(0)), &schema, &statistics, expected_ret);
// Same for the row counts
let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
let container_stats = ContainerStats {
min: Some(Arc::new(Int32Array::from(vec![Some(0)]))),
max: Some(Arc::new(Int32Array::from(vec![Some(0)]))),
null_counts: Some(Arc::new(UInt64Array::from(vec![Some(1)]))),
row_counts: Some(Arc::new(UInt64Array::from(vec![None]))),
..ContainerStats::default()
};
let statistics = TestStatistics::new().with("i", container_stats);
let expected_ret = &[true];
prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
let expected_ret = &[false];
prune_with_expr(col("i").gt(lit(0)), &schema, &statistics, expected_ret);
}
#[test]
fn prune_missing_statistics() {
// If the min or max stats are missing we should not prune
// (unless we know all rows are null, see `prune_all_rows_null_counts`)
let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
let container_stats = ContainerStats {
min: Some(Arc::new(Int32Array::from(vec![None, Some(0)]))),
max: Some(Arc::new(Int32Array::from(vec![Some(0), None]))),
null_counts: Some(Arc::new(UInt64Array::from(vec![Some(0), Some(0)]))),
row_counts: Some(Arc::new(UInt64Array::from(vec![Some(1), Some(1)]))),
..ContainerStats::default()
};
let statistics = TestStatistics::new().with("i", container_stats);
let expected_ret = &[true, true];
prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
let expected_ret = &[false, true];
prune_with_expr(col("i").gt(lit(0)), &schema, &statistics, expected_ret);
let expected_ret = &[true, false];
prune_with_expr(col("i").lt(lit(0)), &schema, &statistics, expected_ret);
}
#[test]
fn prune_null_stats() {
// if null_count = row_count then we should prune the container for i = 0
// regardless of the statistics
let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
let statistics = TestStatistics::new().with(
"i",
ContainerStats::new_i32(
vec![Some(0)], // min
vec![Some(0)], // max
)
.with_null_counts(vec![Some(1)])
.with_row_counts(vec![Some(1)]),
);
let expected_ret = &[false];
// i = 0
prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
}
#[test]
fn test_build_statistics_record_batch() {
// Request a record batch with of s1_min, s2_max, s3_max, s3_min
let required_columns = RequiredColumns::from(vec![
// min of original column s1, named s1_min
(
phys_expr::Column::new("s1", 1),
StatisticsType::Min,
Field::new("s1_min", DataType::Int32, true),
),
// max of original column s2, named s2_max
(
phys_expr::Column::new("s2", 2),
StatisticsType::Max,
Field::new("s2_max", DataType::Int32, true),
),
// max of original column s3, named s3_max
(
phys_expr::Column::new("s3", 3),
StatisticsType::Max,
Field::new("s3_max", DataType::Utf8, true),
),
// min of original column s3, named s3_min
(
phys_expr::Column::new("s3", 3),
StatisticsType::Min,
Field::new("s3_min", DataType::Utf8, true),
),
]);
let statistics = TestStatistics::new()
.with(
"s1",
ContainerStats::new_i32(
vec![None, None, Some(9), None], // min
vec![Some(10), None, None, None], // max
),
)
.with(
"s2",
ContainerStats::new_i32(
vec![Some(2), None, None, None], // min
vec![Some(20), None, None, None], // max
),
)
.with(
"s3",
ContainerStats::new_utf8(
vec![Some("a"), None, None, None], // min
vec![Some("q"), None, Some("r"), None], // max
),
);
let batch =
build_statistics_record_batch(&statistics, &required_columns).unwrap();
assert_snapshot!(batches_to_string(&[batch]), @r"
+--------+--------+--------+--------+
| s1_min | s2_max | s3_max | s3_min |
+--------+--------+--------+--------+
| | 20 | q | a |
| | | | |
| 9 | | r | |
| | | | |
+--------+--------+--------+--------+
");
}
#[test]
fn test_build_statistics_casting() {
// Test requesting a Timestamp column, but getting statistics as Int64
// which is what Parquet does
// Request a record batch with of s1_min as a timestamp
let required_columns = RequiredColumns::from(vec![(
phys_expr::Column::new("s3", 3),
StatisticsType::Min,
Field::new(
"s1_min",
DataType::Timestamp(TimeUnit::Nanosecond, None),
true,
),
)]);
// Note the statistics pass back i64 (not timestamp)
let statistics = OneContainerStats {
min_values: Some(Arc::new(Int64Array::from(vec![Some(10)]))),
max_values: Some(Arc::new(Int64Array::from(vec![Some(20)]))),
num_containers: 1,
};
let batch =
build_statistics_record_batch(&statistics, &required_columns).unwrap();
assert_snapshot!(batches_to_string(&[batch]), @r"
+-------------------------------+
| s1_min |
+-------------------------------+
| 1970-01-01T00:00:00.000000010 |
+-------------------------------+
");
}
#[test]
fn test_build_statistics_no_required_stats() {
let required_columns = RequiredColumns::new();
let statistics = OneContainerStats {
min_values: Some(Arc::new(Int64Array::from(vec![Some(10)]))),
max_values: Some(Arc::new(Int64Array::from(vec![Some(20)]))),
num_containers: 1,
};
let batch =
build_statistics_record_batch(&statistics, &required_columns).unwrap();
assert_eq!(batch.num_rows(), 1); // had 1 container
}
#[test]
fn test_build_statistics_inconsistent_types() {
// Test requesting a Utf8 column when the stats return some other type
// Request a record batch with of s1_min as a timestamp
let required_columns = RequiredColumns::from(vec![(
phys_expr::Column::new("s3", 3),
StatisticsType::Min,
Field::new("s1_min", DataType::Utf8, true),
)]);
// Note the statistics return an invalid UTF-8 sequence which will be converted to null
let statistics = OneContainerStats {
min_values: Some(Arc::new(BinaryArray::from(vec![&[255u8] as &[u8]]))),
max_values: None,
num_containers: 1,
};
let batch =
build_statistics_record_batch(&statistics, &required_columns).unwrap();
assert_snapshot!(batches_to_string(&[batch]), @r"
+--------+
| s1_min |
+--------+
| |
+--------+
");
}
#[test]
fn test_build_statistics_inconsistent_length() {
// return an inconsistent length to the actual statistics arrays
let required_columns = RequiredColumns::from(vec![(
phys_expr::Column::new("s1", 3),
StatisticsType::Min,
Field::new("s1_min", DataType::Int64, true),
)]);
// Note the statistics pass back i64 (not timestamp)
let statistics = OneContainerStats {
min_values: Some(Arc::new(Int64Array::from(vec![Some(10)]))),
max_values: Some(Arc::new(Int64Array::from(vec![Some(20)]))),
num_containers: 3,
};
let result =
build_statistics_record_batch(&statistics, &required_columns).unwrap_err();
assert!(
result
.to_string()
.contains("mismatched statistics length. Expected 3, got 1"),
"{}",
result
);
}
#[test]
fn row_group_predicate_eq() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
let expected_expr =
"c1_null_count@2 != row_count@3 AND c1_min@0 <= 1 AND 1 <= c1_max@1";
// test column on the left
let expr = col("c1").eq(lit(1));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
// test column on the right
let expr = lit(1).eq(col("c1"));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_not_eq() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
let expected_expr =
"c1_null_count@2 != row_count@3 AND (c1_min@0 != 1 OR 1 != c1_max@1)";
// test column on the left
let expr = col("c1").not_eq(lit(1));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
// test column on the right
let expr = lit(1).not_eq(col("c1"));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_gt() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
let expected_expr = "c1_null_count@1 != row_count@2 AND c1_max@0 > 1";
// test column on the left
let expr = col("c1").gt(lit(1));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
// test column on the right
let expr = lit(1).lt(col("c1"));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_gt_eq() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
let expected_expr = "c1_null_count@1 != row_count@2 AND c1_max@0 >= 1";
// test column on the left
let expr = col("c1").gt_eq(lit(1));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
// test column on the right
let expr = lit(1).lt_eq(col("c1"));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_lt() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 < 1";
// test column on the left
let expr = col("c1").lt(lit(1));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
// test column on the right
let expr = lit(1).gt(col("c1"));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_lt_eq() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 <= 1";
// test column on the left
let expr = col("c1").lt_eq(lit(1));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
// test column on the right
let expr = lit(1).gt_eq(col("c1"));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_and() -> Result<()> {
let schema = Schema::new(vec![
Field::new("c1", DataType::Int32, false),
Field::new("c2", DataType::Int32, false),
Field::new("c3", DataType::Int32, false),
]);
// test AND operator joining supported c1 < 1 expression and unsupported c2 > c3 expression
let expr = col("c1").lt(lit(1)).and(col("c2").lt(col("c3")));
let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 < 1";
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_or() -> Result<()> {
let schema = Schema::new(vec![
Field::new("c1", DataType::Int32, false),
Field::new("c2", DataType::Int32, false),
]);
// test OR operator joining supported c1 < 1 expression and unsupported c2 % 2 = 0 expression
let expr = col("c1").lt(lit(1)).or(col("c2").rem(lit(2)).eq(lit(0)));
let expected_expr = "true";
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_not() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
let expected_expr = "true";
let expr = col("c1").not();
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_not_bool() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Boolean, false)]);
let expected_expr = "NOT c1_min@0 AND c1_max@1";
let expr = col("c1").not();
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_bool() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Boolean, false)]);
let expected_expr = "c1_min@0 OR c1_max@1";
let expr = col("c1");
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_lt_bool() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Boolean, false)]);
let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 < true";
// DF doesn't support arithmetic on boolean columns so
// this predicate will error when evaluated
let expr = col("c1").lt(lit(true));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_required_columns() -> Result<()> {
let schema = Schema::new(vec![
Field::new("c1", DataType::Int32, false),
Field::new("c2", DataType::Int32, false),
]);
let mut required_columns = RequiredColumns::new();
// c1 < 1 and (c2 = 2 or c2 = 3)
let expr = col("c1")
.lt(lit(1))
.and(col("c2").eq(lit(2)).or(col("c2").eq(lit(3))));
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)";
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut required_columns);
assert_eq!(predicate_expr.to_string(), expected_expr);
println!("required_columns: {required_columns:#?}"); // for debugging assertions below
// c1 < 1 should add c1_min
let c1_min_field = Field::new("c1_min", DataType::Int32, false);
assert_eq!(
required_columns.columns[0],
(
phys_expr::Column::new("c1", 0),
StatisticsType::Min,
c1_min_field.with_nullable(true) // could be nullable if stats are not present
)
);
// c1 < 1 should add c1_null_count
let c1_null_count_field = Field::new("c1_null_count", DataType::UInt64, false);
assert_eq!(
required_columns.columns[1],
(
phys_expr::Column::new("c1", 0),
StatisticsType::NullCount,
c1_null_count_field.with_nullable(true) // could be nullable if stats are not present
)
);
// c1 < 1 should add row_count
let row_count_field = Field::new("row_count", DataType::UInt64, false);
assert_eq!(
required_columns.columns[2],
(
phys_expr::Column::new("c1", 0),
StatisticsType::RowCount,
row_count_field.with_nullable(true) // could be nullable if stats are not present
)
);
// c2 = 2 should add c2_min and c2_max
let c2_min_field = Field::new("c2_min", DataType::Int32, false);
assert_eq!(
required_columns.columns[3],
(
phys_expr::Column::new("c2", 1),
StatisticsType::Min,
c2_min_field.with_nullable(true) // could be nullable if stats are not present
)
);
let c2_max_field = Field::new("c2_max", DataType::Int32, false);
assert_eq!(
required_columns.columns[4],
(
phys_expr::Column::new("c2", 1),
StatisticsType::Max,
c2_max_field.with_nullable(true) // could be nullable if stats are not present
)
);
// c2 = 2 should add c2_null_count
let c2_null_count_field = Field::new("c2_null_count", DataType::UInt64, false);
assert_eq!(
required_columns.columns[5],
(
phys_expr::Column::new("c2", 1),
StatisticsType::NullCount,
c2_null_count_field.with_nullable(true) // could be nullable if stats are not present
)
);
// c2 = 1 should add row_count
let row_count_field = Field::new("row_count", DataType::UInt64, false);
assert_eq!(
required_columns.columns[2],
(
phys_expr::Column::new("c1", 0),
StatisticsType::RowCount,
row_count_field.with_nullable(true) // could be nullable if stats are not present
)
);
// c2 = 3 shouldn't add any new statistics fields
assert_eq!(required_columns.columns.len(), 6);
Ok(())
}
#[test]
fn row_group_predicate_in_list() -> Result<()> {
let schema = Schema::new(vec![
Field::new("c1", DataType::Int32, false),
Field::new("c2", DataType::Int32, false),
]);
// test c1 in(1, 2, 3)
let expr = Expr::InList(InList::new(
Box::new(col("c1")),
vec![lit(1), lit(2), lit(3)],
false,
));
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";
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_in_list_empty() -> Result<()> {
let schema = Schema::new(vec![
Field::new("c1", DataType::Int32, false),
Field::new("c2", DataType::Int32, false),
]);
// test c1 in()
let expr = Expr::InList(InList::new(Box::new(col("c1")), vec![], false));
let expected_expr = "true";
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_in_list_negated() -> Result<()> {
let schema = Schema::new(vec![
Field::new("c1", DataType::Int32, false),
Field::new("c2", DataType::Int32, false),
]);
// test c1 not in(1, 2, 3)
let expr = Expr::InList(InList::new(
Box::new(col("c1")),
vec![lit(1), lit(2), lit(3)],
true,
));
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)";
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_between() -> Result<()> {
let schema = Schema::new(vec![
Field::new("c1", DataType::Int32, false),
Field::new("c2", DataType::Int32, false),
]);
// test c1 BETWEEN 1 AND 5
let expr1 = col("c1").between(lit(1), lit(5));
// test 1 <= c1 <= 5
let expr2 = col("c1").gt_eq(lit(1)).and(col("c1").lt_eq(lit(5)));
let predicate_expr1 =
test_build_predicate_expression(&expr1, &schema, &mut RequiredColumns::new());
let predicate_expr2 =
test_build_predicate_expression(&expr2, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr1.to_string(), predicate_expr2.to_string());
Ok(())
}
#[test]
fn row_group_predicate_between_with_in_list() -> Result<()> {
let schema = Schema::new(vec![
Field::new("c1", DataType::Int32, false),
Field::new("c2", DataType::Int32, false),
]);
// test c1 in(1, 2)
let expr1 = col("c1").in_list(vec![lit(1), lit(2)], false);
// test c2 BETWEEN 4 AND 5
let expr2 = col("c2").between(lit(4), lit(5));
// test c1 in(1, 2) and c2 BETWEEN 4 AND 5
let expr3 = expr1.and(expr2);
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";
let predicate_expr =
test_build_predicate_expression(&expr3, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_in_list_to_many_values() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
// test c1 in(1..21)
// in pruning.rs has MAX_LIST_VALUE_SIZE_REWRITE = 20, more than this value will be rewrite
// always true
let expr = col("c1").in_list((1..=21).map(lit).collect(), false);
let expected_expr = "true";
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_cast_int_int() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
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)";
// test cast(c1 as int64) = 1
// test column on the left
let expr = cast(col("c1"), DataType::Int64).eq(lit(ScalarValue::Int64(Some(1))));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
// test column on the right
let expr = lit(ScalarValue::Int64(Some(1))).eq(cast(col("c1"), DataType::Int64));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
let expected_expr =
"c1_null_count@1 != row_count@2 AND TRY_CAST(c1_max@0 AS Int64) > 1";
// test column on the left
let expr =
try_cast(col("c1"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(1))));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
// test column on the right
let expr =
lit(ScalarValue::Int64(Some(1))).lt(try_cast(col("c1"), DataType::Int64));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_cast_string_string() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]);
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)";
// test column on the left
let expr = cast(col("c1"), DataType::Utf8)
.eq(lit(ScalarValue::Utf8(Some("1".to_string()))));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
// test column on the right
let expr = lit(ScalarValue::Utf8(Some("1".to_string())))
.eq(cast(col("c1"), DataType::Utf8));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_cast_string_int() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]);
let expected_expr = "true";
// test column on the left
let expr = cast(col("c1"), DataType::Int32).eq(lit(ScalarValue::Int32(Some(1))));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
// test column on the right
let expr = lit(ScalarValue::Int32(Some(1))).eq(cast(col("c1"), DataType::Int32));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_cast_int_string() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
let expected_expr = "true";
// test column on the left
let expr = cast(col("c1"), DataType::Utf8)
.eq(lit(ScalarValue::Utf8(Some("1".to_string()))));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
// test column on the right
let expr = lit(ScalarValue::Utf8(Some("1".to_string())))
.eq(cast(col("c1"), DataType::Utf8));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_date_date() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Date32, false)]);
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)";
// test column on the left
let expr =
cast(col("c1"), DataType::Date64).eq(lit(ScalarValue::Date64(Some(123))));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
// test column on the right
let expr =
lit(ScalarValue::Date64(Some(123))).eq(cast(col("c1"), DataType::Date64));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_dict_string_date() -> Result<()> {
// Test with Dictionary<UInt8, Utf8> for the literal
let schema = Schema::new(vec![Field::new("c1", DataType::Date32, false)]);
let expected_expr = "true";
// test column on the left
let expr = cast(
col("c1"),
DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
)
.eq(lit(ScalarValue::Utf8(Some("2024-01-01".to_string()))));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
// test column on the right
let expr = lit(ScalarValue::Utf8(Some("2024-01-01".to_string()))).eq(cast(
col("c1"),
DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_date_dict_string() -> Result<()> {
// Test with Dictionary<UInt8, Utf8> for the column
let schema = Schema::new(vec![Field::new(
"c1",
DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
false,
)]);
let expected_expr = "true";
// test column on the left
let expr =
cast(col("c1"), DataType::Date32).eq(lit(ScalarValue::Date32(Some(123))));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
// test column on the right
let expr =
lit(ScalarValue::Date32(Some(123))).eq(cast(col("c1"), DataType::Date32));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_dict_dict_same_value_type() -> Result<()> {
// Test with Dictionary types that have the same value type but different key types
let schema = Schema::new(vec![Field::new(
"c1",
DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
false,
)]);
// Direct comparison with no cast
let expr = col("c1").eq(lit(ScalarValue::Utf8(Some("test".to_string()))));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
let expected_expr =
"c1_null_count@2 != row_count@3 AND c1_min@0 <= test AND test <= c1_max@1";
assert_eq!(predicate_expr.to_string(), expected_expr);
// Test with column cast to a dictionary with different key type
let expr = cast(
col("c1"),
DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
)
.eq(lit(ScalarValue::Utf8(Some("test".to_string()))));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
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))";
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_dict_dict_different_value_type() -> Result<()> {
// Test with Dictionary types that have different value types
let schema = Schema::new(vec![Field::new(
"c1",
DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Int32)),
false,
)]);
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)";
// Test with literal of a different type
let expr =
cast(col("c1"), DataType::Int64).eq(lit(ScalarValue::Int64(Some(123))));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_nested_dict() -> Result<()> {
// Test with nested Dictionary types
let schema = Schema::new(vec![Field::new(
"c1",
DataType::Dictionary(
Box::new(DataType::UInt8),
Box::new(DataType::Dictionary(
Box::new(DataType::UInt16),
Box::new(DataType::Utf8),
)),
),
false,
)]);
let expected_expr =
"c1_null_count@2 != row_count@3 AND c1_min@0 <= test AND test <= c1_max@1";
// Test with a simple literal
let expr = col("c1").eq(lit(ScalarValue::Utf8(Some("test".to_string()))));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_dict_date_dict_date() -> Result<()> {
// Test with dictionary-wrapped date types for both sides
let schema = Schema::new(vec![Field::new(
"c1",
DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Date32)),
false,
)]);
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))";
// Test with a cast to a different date type
let expr = cast(
col("c1"),
DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Date64)),
)
.eq(lit(ScalarValue::Date64(Some(123))));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_date_string() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Utf8, false)]);
let expected_expr = "true";
// test column on the left
let expr =
cast(col("c1"), DataType::Date32).eq(lit(ScalarValue::Date32(Some(123))));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
// test column on the right
let expr =
lit(ScalarValue::Date32(Some(123))).eq(cast(col("c1"), DataType::Date32));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_string_date() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Date32, false)]);
let expected_expr = "true";
// test column on the left
let expr = cast(col("c1"), DataType::Utf8)
.eq(lit(ScalarValue::Utf8(Some("2024-01-01".to_string()))));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
// test column on the right
let expr = lit(ScalarValue::Utf8(Some("2024-01-01".to_string())))
.eq(cast(col("c1"), DataType::Utf8));
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn row_group_predicate_cast_list() -> Result<()> {
let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
// test cast(c1 as int64) in int64(1, 2, 3)
let expr = Expr::InList(InList::new(
Box::new(cast(col("c1"), DataType::Int64)),
vec![
lit(ScalarValue::Int64(Some(1))),
lit(ScalarValue::Int64(Some(2))),
lit(ScalarValue::Int64(Some(3))),
],
false,
));
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)";
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
let expr = Expr::InList(InList::new(
Box::new(cast(col("c1"), DataType::Int64)),
vec![
lit(ScalarValue::Int64(Some(1))),
lit(ScalarValue::Int64(Some(2))),
lit(ScalarValue::Int64(Some(3))),
],
true,
));
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))";
let predicate_expr =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
assert_eq!(predicate_expr.to_string(), expected_expr);
Ok(())
}
#[test]
fn prune_decimal_data() {
// decimal(9,2)
let schema = Arc::new(Schema::new(vec![Field::new(
"s1",
DataType::Decimal128(9, 2),
true,
)]));
prune_with_expr(
// s1 > 5
col("s1").gt(lit(ScalarValue::Decimal128(Some(500), 9, 2))),
&schema,
// If the data is written by spark, the physical data type is INT32 in the parquet
// So we use the INT32 type of statistic.
&TestStatistics::new().with(
"s1",
ContainerStats::new_i32(
vec![Some(0), Some(4), None, Some(3)], // min
vec![Some(5), Some(6), Some(4), None], // max
),
),
&[false, true, false, true],
);
prune_with_expr(
// with cast column to other type
cast(col("s1"), DataType::Decimal128(14, 3))
.gt(lit(ScalarValue::Decimal128(Some(5000), 14, 3))),
&schema,
&TestStatistics::new().with(
"s1",
ContainerStats::new_i32(
vec![Some(0), Some(4), None, Some(3)], // min
vec![Some(5), Some(6), Some(4), None], // max
),
),
&[false, true, false, true],
);
prune_with_expr(
// with try cast column to other type
try_cast(col("s1"), DataType::Decimal128(14, 3))
.gt(lit(ScalarValue::Decimal128(Some(5000), 14, 3))),
&schema,
&TestStatistics::new().with(
"s1",
ContainerStats::new_i32(
vec![Some(0), Some(4), None, Some(3)], // min
vec![Some(5), Some(6), Some(4), None], // max
),
),
&[false, true, false, true],
);
// decimal(18,2)
let schema = Arc::new(Schema::new(vec![Field::new(
"s1",
DataType::Decimal128(18, 2),
true,
)]));
prune_with_expr(
// s1 > 5
col("s1").gt(lit(ScalarValue::Decimal128(Some(500), 18, 2))),
&schema,
// If the data is written by spark, the physical data type is INT64 in the parquet
// So we use the INT32 type of statistic.
&TestStatistics::new().with(
"s1",
ContainerStats::new_i64(
vec![Some(0), Some(4), None, Some(3)], // min
vec![Some(5), Some(6), Some(4), None], // max
),
),
&[false, true, false, true],
);
// decimal(23,2)
let schema = Arc::new(Schema::new(vec![Field::new(
"s1",
DataType::Decimal128(23, 2),
true,
)]));
prune_with_expr(
// s1 > 5
col("s1").gt(lit(ScalarValue::Decimal128(Some(500), 23, 2))),
&schema,
&TestStatistics::new().with(
"s1",
ContainerStats::new_decimal128(
vec![Some(0), Some(400), None, Some(300)], // min
vec![Some(500), Some(600), Some(400), None], // max
23,
2,
),
),
&[false, true, false, true],
);
}
#[test]
fn prune_api() {
let schema = Arc::new(Schema::new(vec![
Field::new("s1", DataType::Utf8, true),
Field::new("s2", DataType::Int32, true),
]));
let statistics = TestStatistics::new().with(
"s2",
ContainerStats::new_i32(
vec![Some(0), Some(4), None, Some(3)], // min
vec![Some(5), Some(6), None, None], // max
),
);
prune_with_expr(
// Prune using s2 > 5
col("s2").gt(lit(5)),
&schema,
&statistics,
// s2 [0, 5] ==> no rows should pass
// s2 [4, 6] ==> some rows could pass
// No stats for s2 ==> some rows could pass
// s2 [3, None] (null max) ==> some rows could pass
&[false, true, true, true],
);
prune_with_expr(
// filter with cast
cast(col("s2"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(5)))),
&schema,
&statistics,
&[false, true, true, true],
);
}
#[test]
fn prune_not_eq_data() {
let schema = Arc::new(Schema::new(vec![Field::new("s1", DataType::Utf8, true)]));
prune_with_expr(
// Prune using s2 != 'M'
col("s1").not_eq(lit("M")),
&schema,
&TestStatistics::new().with(
"s1",
ContainerStats::new_utf8(
vec![Some("A"), Some("A"), Some("N"), Some("M"), None, Some("A")], // min
vec![Some("Z"), Some("L"), Some("Z"), Some("M"), None, None], // max
),
),
// s1 [A, Z] ==> might have values that pass predicate
// s1 [A, L] ==> all rows pass the predicate
// s1 [N, Z] ==> all rows pass the predicate
// s1 [M, M] ==> all rows do not pass the predicate
// No stats for s2 ==> some rows could pass
// s2 [3, None] (null max) ==> some rows could pass
&[true, true, true, false, true, true],
);
}
/// Creates setup for boolean chunk pruning
///
/// For predicate "b1" (boolean expr)
/// b1 [false, false] ==> no rows can pass (not keep)
/// b1 [false, true] ==> some rows could pass (must keep)
/// b1 [true, true] ==> all rows must pass (must keep)
/// b1 [NULL, NULL] ==> unknown (must keep)
/// b1 [false, NULL] ==> unknown (must keep)
///
/// For predicate "!b1" (boolean expr)
/// b1 [false, false] ==> all rows pass (must keep)
/// b1 [false, true] ==> some rows could pass (must keep)
/// b1 [true, true] ==> no rows can pass (not keep)
/// b1 [NULL, NULL] ==> unknown (must keep)
/// b1 [false, NULL] ==> unknown (must keep)
fn bool_setup() -> (SchemaRef, TestStatistics, Vec<bool>, Vec<bool>) {
let schema =
Arc::new(Schema::new(vec![Field::new("b1", DataType::Boolean, true)]));
let statistics = TestStatistics::new().with(
"b1",
ContainerStats::new_bool(
vec![Some(false), Some(false), Some(true), None, Some(false)], // min
vec![Some(false), Some(true), Some(true), None, None], // max
),
);
let expected_true = vec![false, true, true, true, true];
let expected_false = vec![true, true, false, true, true];
(schema, statistics, expected_true, expected_false)
}
#[test]
fn prune_bool_const_expr() {
let (schema, statistics, _, _) = bool_setup();
prune_with_expr(
// true
lit(true),
&schema,
&statistics,
&[true, true, true, true, true],
);
prune_with_expr(
// false
lit(false),
&schema,
&statistics,
&[false, false, false, false, false],
);
}
#[test]
fn prune_bool_column() {
let (schema, statistics, expected_true, _) = bool_setup();
prune_with_expr(
// b1
col("b1"),
&schema,
&statistics,
&expected_true,
);
}
#[test]
fn prune_bool_not_column() {
let (schema, statistics, _, expected_false) = bool_setup();
prune_with_expr(
// !b1
col("b1").not(),
&schema,
&statistics,
&expected_false,
);
}
#[test]
fn prune_bool_column_eq_true() {
let (schema, statistics, expected_true, _) = bool_setup();
prune_with_expr(
// b1 = true
col("b1").eq(lit(true)),
&schema,
&statistics,
&expected_true,
);
}
#[test]
fn prune_bool_not_column_eq_true() {
let (schema, statistics, _, expected_false) = bool_setup();
prune_with_expr(
// !b1 = true
col("b1").not().eq(lit(true)),
&schema,
&statistics,
&expected_false,
);
}
/// Creates a setup for chunk pruning, modeling a int32 column "i"
/// with 5 different containers (e.g. RowGroups). They have [min,
/// max]:
///
/// i [-5, 5]
/// i [1, 11]
/// i [-11, -1]
/// i [NULL, NULL]
/// i [1, NULL]
fn int32_setup() -> (SchemaRef, TestStatistics) {
let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
let statistics = TestStatistics::new().with(
"i",
ContainerStats::new_i32(
vec![Some(-5), Some(1), Some(-11), None, Some(1)], // min
vec![Some(5), Some(11), Some(-1), None, None], // max
),
);
(schema, statistics)
}
#[test]
fn prune_int32_col_gt_zero() {
let (schema, statistics) = int32_setup();
// Expression "i > 0" and "-i < 0"
// i [-5, 5] ==> some rows could pass (must keep)
// i [1, 11] ==> all rows must pass (must keep)
// i [-11, -1] ==> no rows can pass (not keep)
// i [NULL, NULL] ==> unknown (must keep)
// i [1, NULL] ==> unknown (must keep)
let expected_ret = &[true, true, false, true, true];
// i > 0
prune_with_expr(col("i").gt(lit(0)), &schema, &statistics, expected_ret);
// -i < 0
prune_with_expr(
Expr::Negative(Box::new(col("i"))).lt(lit(0)),
&schema,
&statistics,
expected_ret,
);
}
#[test]
fn prune_int32_col_lte_zero() {
let (schema, statistics) = int32_setup();
// Expression "i <= 0" and "-i >= 0"
// i [-5, 5] ==> some rows could pass (must keep)
// i [1, 11] ==> no rows can pass (not keep)
// i [-11, -1] ==> all rows must pass (must keep)
// i [NULL, NULL] ==> unknown (must keep)
// i [1, NULL] ==> no rows can pass (not keep)
let expected_ret = &[true, false, true, true, false];
prune_with_expr(
// i <= 0
col("i").lt_eq(lit(0)),
&schema,
&statistics,
expected_ret,
);
prune_with_expr(
// -i >= 0
Expr::Negative(Box::new(col("i"))).gt_eq(lit(0)),
&schema,
&statistics,
expected_ret,
);
}
#[test]
fn prune_int32_col_lte_zero_cast() {
let (schema, statistics) = int32_setup();
// Expression "cast(i as utf8) <= '0'"
// i [-5, 5] ==> some rows could pass (must keep)
// i [1, 11] ==> no rows can pass in theory, -0.22 (conservatively keep)
// i [-11, -1] ==> no rows could pass in theory (conservatively keep)
// i [NULL, NULL] ==> unknown (must keep)
// i [1, NULL] ==> no rows can pass (conservatively keep)
let expected_ret = &[true, true, true, true, true];
prune_with_expr(
// cast(i as utf8) <= 0
cast(col("i"), DataType::Utf8).lt_eq(lit("0")),
&schema,
&statistics,
expected_ret,
);
prune_with_expr(
// try_cast(i as utf8) <= 0
try_cast(col("i"), DataType::Utf8).lt_eq(lit("0")),
&schema,
&statistics,
expected_ret,
);
prune_with_expr(
// cast(-i as utf8) >= 0
cast(Expr::Negative(Box::new(col("i"))), DataType::Utf8).gt_eq(lit("0")),
&schema,
&statistics,
expected_ret,
);
prune_with_expr(
// try_cast(-i as utf8) >= 0
try_cast(Expr::Negative(Box::new(col("i"))), DataType::Utf8).gt_eq(lit("0")),
&schema,
&statistics,
expected_ret,
);
}
#[test]
fn prune_int32_col_eq_zero() {
let (schema, statistics) = int32_setup();
// Expression "i = 0"
// i [-5, 5] ==> some rows could pass (must keep)
// i [1, 11] ==> no rows can pass (not keep)
// i [-11, -1] ==> no rows can pass (not keep)
// i [NULL, NULL] ==> unknown (must keep)
// i [1, NULL] ==> no rows can pass (not keep)
let expected_ret = &[true, false, false, true, false];
prune_with_expr(
// i = 0
col("i").eq(lit(0)),
&schema,
&statistics,
expected_ret,
);
}
#[test]
fn prune_int32_col_eq_zero_cast() {
let (schema, statistics) = int32_setup();
// Expression "cast(i as int64) = 0"
// i [-5, 5] ==> some rows could pass (must keep)
// i [1, 11] ==> no rows can pass (not keep)
// i [-11, -1] ==> no rows can pass (not keep)
// i [NULL, NULL] ==> unknown (must keep)
// i [1, NULL] ==> no rows can pass (not keep)
let expected_ret = &[true, false, false, true, false];
prune_with_expr(
cast(col("i"), DataType::Int64).eq(lit(0i64)),
&schema,
&statistics,
expected_ret,
);
prune_with_expr(
try_cast(col("i"), DataType::Int64).eq(lit(0i64)),
&schema,
&statistics,
expected_ret,
);
}
#[test]
fn prune_int32_col_eq_zero_cast_as_str() {
let (schema, statistics) = int32_setup();
// Note the cast is to a string where sorting properties are
// not the same as integers
//
// Expression "cast(i as utf8) = '0'"
// i [-5, 5] ==> some rows could pass (keep)
// i [1, 11] ==> no rows can pass (could keep)
// i [-11, -1] ==> no rows can pass (could keep)
// i [NULL, NULL] ==> unknown (keep)
// i [1, NULL] ==> no rows can pass (could keep)
let expected_ret = &[true, true, true, true, true];
prune_with_expr(
cast(col("i"), DataType::Utf8).eq(lit("0")),
&schema,
&statistics,
expected_ret,
);
}
#[test]
fn prune_int32_col_lt_neg_one() {
let (schema, statistics) = int32_setup();
// Expression "i > -1" and "-i < 1"
// i [-5, 5] ==> some rows could pass (must keep)
// i [1, 11] ==> all rows must pass (must keep)
// i [-11, -1] ==> no rows can pass (not keep)
// i [NULL, NULL] ==> unknown (must keep)
// i [1, NULL] ==> all rows must pass (must keep)
let expected_ret = &[true, true, false, true, true];
prune_with_expr(
// i > -1
col("i").gt(lit(-1)),
&schema,
&statistics,
expected_ret,
);
prune_with_expr(
// -i < 1
Expr::Negative(Box::new(col("i"))).lt(lit(1)),
&schema,
&statistics,
expected_ret,
);
}
#[test]
fn prune_int32_is_null() {
let (schema, statistics) = int32_setup();
// Expression "i IS NULL" when there are no null statistics,
// should all be kept
let expected_ret = &[true, true, true, true, true];
prune_with_expr(
// i IS NULL, no null statistics
col("i").is_null(),
&schema,
&statistics,
expected_ret,
);
// provide null counts for each column
let statistics = statistics.with_null_counts(
"i",
vec![
Some(0), // no nulls (don't keep)
Some(1), // 1 null
None, // unknown nulls
None, // unknown nulls (min/max are both null too, like no stats at all)
Some(0), // 0 nulls (max=null too which means no known max) (don't keep)
],
);
let expected_ret = &[false, true, true, true, false];
prune_with_expr(
// i IS NULL, with actual null statistics
col("i").is_null(),
&schema,
&statistics,
expected_ret,
);
}
#[test]
fn prune_int32_column_is_known_all_null() {
let (schema, statistics) = int32_setup();
// Expression "i < 0"
// i [-5, 5] ==> some rows could pass (must keep)
// i [1, 11] ==> no rows can pass (not keep)
// i [-11, -1] ==> all rows must pass (must keep)
// i [NULL, NULL] ==> unknown (must keep)
// i [1, NULL] ==> no rows can pass (not keep)
let expected_ret = &[true, false, true, true, false];
prune_with_expr(
// i < 0
col("i").lt(lit(0)),
&schema,
&statistics,
expected_ret,
);
// provide row counts for each column
let statistics = statistics.with_row_counts(
"i",
vec![
Some(10), // 10 rows of data
Some(9), // 9 rows of data
None, // unknown row counts
Some(4),
Some(10),
],
);
// pruning result is still the same if we only know row counts
prune_with_expr(
// i < 0, with only row counts statistics
col("i").lt(lit(0)),
&schema,
&statistics,
expected_ret,
);
// provide null counts for each column
let statistics = statistics.with_null_counts(
"i",
vec![
Some(0), // no nulls
Some(1), // 1 null
None, // unknown nulls
Some(4), // 4 nulls, which is the same as the row counts, i.e. this column is all null (don't keep)
Some(0), // 0 nulls (max=null too which means no known max)
],
);
// Expression "i < 0" with actual null and row counts statistics
// col | min, max | row counts | null counts |
// ----+--------------+------------+-------------+
// i | [-5, 5] | 10 | 0 | ==> Some rows could pass (must keep)
// i | [1, 11] | 9 | 1 | ==> No rows can pass (not keep)
// i | [-11,-1] | Unknown | Unknown | ==> All rows must pass (must keep)
// i | [NULL, NULL] | 4 | 4 | ==> The column is all null (not keep)
// i | [1, NULL] | 10 | 0 | ==> No rows can pass (not keep)
let expected_ret = &[true, false, true, false, false];
prune_with_expr(
// i < 0, with actual null and row counts statistics
col("i").lt(lit(0)),
&schema,
&statistics,
expected_ret,
);
}
#[test]
fn prune_cast_column_scalar() {
// The data type of column i is INT32
let (schema, statistics) = int32_setup();
let expected_ret = &[true, true, false, true, true];
prune_with_expr(
// i > int64(0)
col("i").gt(cast(lit(ScalarValue::Int64(Some(0))), DataType::Int32)),
&schema,
&statistics,
expected_ret,
);
prune_with_expr(
// cast(i as int64) > int64(0)
cast(col("i"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(0)))),
&schema,
&statistics,
expected_ret,
);
prune_with_expr(
// try_cast(i as int64) > int64(0)
try_cast(col("i"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(0)))),
&schema,
&statistics,
expected_ret,
);
prune_with_expr(
// `-cast(i as int64) < 0` convert to `cast(i as int64) > -0`
Expr::Negative(Box::new(cast(col("i"), DataType::Int64)))
.lt(lit(ScalarValue::Int64(Some(0)))),
&schema,
&statistics,
expected_ret,
);
}
#[test]
fn test_increment_utf8() {
// Basic ASCII
assert_eq!(increment_utf8("abc").unwrap(), "abd");
assert_eq!(increment_utf8("abz").unwrap(), "ab{");
// Test around ASCII 127 (DEL)
assert_eq!(increment_utf8("~").unwrap(), "\u{7f}"); // 126 -> 127
assert_eq!(increment_utf8("\u{7f}").unwrap(), "\u{80}"); // 127 -> 128
// Test 2-byte UTF-8 sequences
assert_eq!(increment_utf8("ß").unwrap(), "à"); // U+00DF -> U+00E0
// Test 3-byte UTF-8 sequences
assert_eq!(increment_utf8("℣").unwrap(), "ℤ"); // U+2123 -> U+2124
// Test at UTF-8 boundaries
assert_eq!(increment_utf8("\u{7FF}").unwrap(), "\u{800}"); // 2-byte to 3-byte boundary
assert_eq!(increment_utf8("\u{FFFF}").unwrap(), "\u{10000}"); // 3-byte to 4-byte boundary
// Test that if we can't increment we return None
assert!(increment_utf8("").is_none());
assert!(increment_utf8("\u{10FFFF}").is_none()); // U+10FFFF is the max code point
// Test that if we can't increment the last character we do the previous one and truncate
assert_eq!(increment_utf8("a\u{10FFFF}").unwrap(), "b");
// Test surrogate pair range (0xD800..=0xDFFF)
assert_eq!(increment_utf8("a\u{D7FF}").unwrap(), "b");
assert!(increment_utf8("\u{D7FF}").is_none());
// Test non-characters range (0xFDD0..=0xFDEF)
assert_eq!(increment_utf8("a\u{FDCF}").unwrap(), "b");
assert!(increment_utf8("\u{FDCF}").is_none());
// Test private use area limit (>= 0x110000)
assert_eq!(increment_utf8("a\u{10FFFF}").unwrap(), "b");
assert!(increment_utf8("\u{10FFFF}").is_none()); // Can't increment past max valid codepoint
}
/// Creates a setup for chunk pruning, modeling a utf8 column "s1"
/// with 5 different containers (e.g. RowGroups). They have [min,
/// max]:
/// s1 ["A", "Z"]
/// s1 ["A", "L"]
/// s1 ["N", "Z"]
/// s1 [NULL, NULL]
/// s1 ["A", NULL]
/// s1 ["", "A"]
/// s1 ["", ""]
/// s1 ["AB", "A\u{10ffff}"]
/// s1 ["A\u{10ffff}\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"]
fn utf8_setup() -> (SchemaRef, TestStatistics) {
let schema = Arc::new(Schema::new(vec![Field::new("s1", DataType::Utf8, true)]));
let statistics = TestStatistics::new().with(
"s1",
ContainerStats::new_utf8(
vec![
Some("A"),
Some("A"),
Some("N"),
Some("M"),
None,
Some("A"),
Some(""),
Some(""),
Some("AB"),
Some("A\u{10ffff}\u{10ffff}"),
], // min
vec![
Some("Z"),
Some("L"),
Some("Z"),
Some("M"),
None,
None,
Some("A"),
Some(""),
Some("A\u{10ffff}\u{10ffff}\u{10ffff}"),
Some("A\u{10ffff}\u{10ffff}"),
], // max
),
);
(schema, statistics)
}
#[test]
fn prune_utf8_eq() {
let (schema, statistics) = utf8_setup();
let expr = col("s1").eq(lit("A"));
#[rustfmt::skip]
let expected_ret = &[
// s1 ["A", "Z"] ==> some rows could pass (must keep)
true,
// s1 ["A", "L"] ==> some rows could pass (must keep)
true,
// s1 ["N", "Z"] ==> no rows can pass (not keep)
false,
// s1 ["M", "M"] ==> no rows can pass (not keep)
false,
// s1 [NULL, NULL] ==> unknown (must keep)
true,
// s1 ["A", NULL] ==> unknown (must keep)
true,
// s1 ["", "A"] ==> some rows could pass (must keep)
true,
// s1 ["", ""] ==> no rows can pass (not keep)
false,
// s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> no rows can pass (not keep)
false,
// s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> no rows can pass (not keep)
false,
];
prune_with_expr(expr, &schema, &statistics, expected_ret);
let expr = col("s1").eq(lit(""));
#[rustfmt::skip]
let expected_ret = &[
// s1 ["A", "Z"] ==> no rows can pass (not keep)
false,
// s1 ["A", "L"] ==> no rows can pass (not keep)
false,
// s1 ["N", "Z"] ==> no rows can pass (not keep)
false,
// s1 ["M", "M"] ==> no rows can pass (not keep)
false,
// s1 [NULL, NULL] ==> unknown (must keep)
true,
// s1 ["A", NULL] ==> no rows can pass (not keep)
false,
// s1 ["", "A"] ==> some rows could pass (must keep)
true,
// s1 ["", ""] ==> all rows must pass (must keep)
true,
// s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> no rows can pass (not keep)
false,
// s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> no rows can pass (not keep)
false,
];
prune_with_expr(expr, &schema, &statistics, expected_ret);
}
#[test]
fn prune_utf8_not_eq() {
let (schema, statistics) = utf8_setup();
let expr = col("s1").not_eq(lit("A"));
#[rustfmt::skip]
let expected_ret = &[
// s1 ["A", "Z"] ==> some rows could pass (must keep)
true,
// s1 ["A", "L"] ==> some rows could pass (must keep)
true,
// s1 ["N", "Z"] ==> all rows must pass (must keep)
true,
// s1 ["M", "M"] ==> all rows must pass (must keep)
true,
// s1 [NULL, NULL] ==> unknown (must keep)
true,
// s1 ["A", NULL] ==> unknown (must keep)
true,
// s1 ["", "A"] ==> some rows could pass (must keep)
true,
// s1 ["", ""] ==> all rows must pass (must keep)
true,
// s1 ["AB", "A\u{10ffff}\u{10ffff}"] ==> all rows must pass (must keep)
true,
// s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> all rows must pass (must keep)
true,
];
prune_with_expr(expr, &schema, &statistics, expected_ret);
let expr = col("s1").not_eq(lit(""));
#[rustfmt::skip]
let expected_ret = &[
// s1 ["A", "Z"] ==> all rows must pass (must keep)
true,
// s1 ["A", "L"] ==> all rows must pass (must keep)
true,
// s1 ["N", "Z"] ==> all rows must pass (must keep)
true,
// s1 ["M", "M"] ==> all rows must pass (must keep)
true,
// s1 [NULL, NULL] ==> unknown (must keep)
true,
// s1 ["A", NULL] ==> unknown (must keep)
true,
// s1 ["", "A"] ==> some rows could pass (must keep)
true,
// s1 ["", ""] ==> no rows can pass (not keep)
false,
// s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> all rows must pass (must keep)
true,
// s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> all rows must pass (must keep)
true,
];
prune_with_expr(expr, &schema, &statistics, expected_ret);
}
#[test]
fn prune_utf8_like_one() {
let (schema, statistics) = utf8_setup();
let expr = col("s1").like(lit("A_"));
#[rustfmt::skip]
let expected_ret = &[
// s1 ["A", "Z"] ==> some rows could pass (must keep)
true,
// s1 ["A", "L"] ==> some rows could pass (must keep)
true,
// s1 ["N", "Z"] ==> no rows can pass (not keep)
false,
// s1 ["M", "M"] ==> no rows can pass (not keep)
false,
// s1 [NULL, NULL] ==> unknown (must keep)
true,
// s1 ["A", NULL] ==> unknown (must keep)
true,
// s1 ["", "A"] ==> some rows could pass (must keep)
true,
// s1 ["", ""] ==> no rows can pass (not keep)
false,
// s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep)
true,
// s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep)
true,
];
prune_with_expr(expr, &schema, &statistics, expected_ret);
let expr = col("s1").like(lit("_A_"));
#[rustfmt::skip]
let expected_ret = &[
// s1 ["A", "Z"] ==> some rows could pass (must keep)
true,
// s1 ["A", "L"] ==> some rows could pass (must keep)
true,
// s1 ["N", "Z"] ==> some rows could pass (must keep)
true,
// s1 ["M", "M"] ==> some rows could pass (must keep)
true,
// s1 [NULL, NULL] ==> unknown (must keep)
true,
// s1 ["A", NULL] ==> unknown (must keep)
true,
// s1 ["", "A"] ==> some rows could pass (must keep)
true,
// s1 ["", ""] ==> some rows could pass (must keep)
true,
// s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep)
true,
// s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep)
true,
];
prune_with_expr(expr, &schema, &statistics, expected_ret);
let expr = col("s1").like(lit("_"));
#[rustfmt::skip]
let expected_ret = &[
// s1 ["A", "Z"] ==> all rows must pass (must keep)
true,
// s1 ["A", "L"] ==> all rows must pass (must keep)
true,
// s1 ["N", "Z"] ==> all rows must pass (must keep)
true,
// s1 ["M", "M"] ==> all rows must pass (must keep)
true,
// s1 [NULL, NULL] ==> unknown (must keep)
true,
// s1 ["A", NULL] ==> unknown (must keep)
true,
// s1 ["", "A"] ==> all rows must pass (must keep)
true,
// s1 ["", ""] ==> all rows must pass (must keep)
true,
// s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> all rows must pass (must keep)
true,
// s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> all rows must pass (must keep)
true,
];
prune_with_expr(expr, &schema, &statistics, expected_ret);
let expr = col("s1").like(lit(""));
#[rustfmt::skip]
let expected_ret = &[
// s1 ["A", "Z"] ==> no rows can pass (not keep)
false,
// s1 ["A", "L"] ==> no rows can pass (not keep)
false,
// s1 ["N", "Z"] ==> no rows can pass (not keep)
false,
// s1 ["M", "M"] ==> no rows can pass (not keep)
false,
// s1 [NULL, NULL] ==> unknown (must keep)
true,
// s1 ["A", NULL] ==> no rows can pass (not keep)
false,
// s1 ["", "A"] ==> some rows could pass (must keep)
true,
// s1 ["", ""] ==> all rows must pass (must keep)
true,
// s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> no rows can pass (not keep)
false,
// s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> no rows can pass (not keep)
false,
];
prune_with_expr(expr, &schema, &statistics, expected_ret);
}
#[test]
fn prune_utf8_like_many() {
let (schema, statistics) = utf8_setup();
let expr = col("s1").like(lit("A%"));
#[rustfmt::skip]
let expected_ret = &[
// s1 ["A", "Z"] ==> some rows could pass (must keep)
true,
// s1 ["A", "L"] ==> some rows could pass (must keep)
true,
// s1 ["N", "Z"] ==> no rows can pass (not keep)
false,
// s1 ["M", "M"] ==> no rows can pass (not keep)
false,
// s1 [NULL, NULL] ==> unknown (must keep)
true,
// s1 ["A", NULL] ==> unknown (must keep)
true,
// s1 ["", "A"] ==> some rows could pass (must keep)
true,
// s1 ["", ""] ==> no rows can pass (not keep)
false,
// s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep)
true,
// s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep)
true,
];
prune_with_expr(expr, &schema, &statistics, expected_ret);
let expr = col("s1").like(lit("%A%"));
#[rustfmt::skip]
let expected_ret = &[
// s1 ["A", "Z"] ==> some rows could pass (must keep)
true,
// s1 ["A", "L"] ==> some rows could pass (must keep)
true,
// s1 ["N", "Z"] ==> some rows could pass (must keep)
true,
// s1 ["M", "M"] ==> some rows could pass (must keep)
true,
// s1 [NULL, NULL] ==> unknown (must keep)
true,
// s1 ["A", NULL] ==> unknown (must keep)
true,
// s1 ["", "A"] ==> some rows could pass (must keep)
true,
// s1 ["", ""] ==> some rows could pass (must keep)
true,
// s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep)
true,
// s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep)
true,
];
prune_with_expr(expr, &schema, &statistics, expected_ret);
let expr = col("s1").like(lit("%"));
#[rustfmt::skip]
let expected_ret = &[
// s1 ["A", "Z"] ==> all rows must pass (must keep)
true,
// s1 ["A", "L"] ==> all rows must pass (must keep)
true,
// s1 ["N", "Z"] ==> all rows must pass (must keep)
true,
// s1 ["M", "M"] ==> all rows must pass (must keep)
true,
// s1 [NULL, NULL] ==> unknown (must keep)
true,
// s1 ["A", NULL] ==> unknown (must keep)
true,
// s1 ["", "A"] ==> all rows must pass (must keep)
true,
// s1 ["", ""] ==> all rows must pass (must keep)
true,
// s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> all rows must pass (must keep)
true,
// s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> all rows must pass (must keep)
true,
];
prune_with_expr(expr, &schema, &statistics, expected_ret);
let expr = col("s1").like(lit(""));
#[rustfmt::skip]
let expected_ret = &[
// s1 ["A", "Z"] ==> no rows can pass (not keep)
false,
// s1 ["A", "L"] ==> no rows can pass (not keep)
false,
// s1 ["N", "Z"] ==> no rows can pass (not keep)
false,
// s1 ["M", "M"] ==> no rows can pass (not keep)
false,
// s1 [NULL, NULL] ==> unknown (must keep)
true,
// s1 ["A", NULL] ==> no rows can pass (not keep)
false,
// s1 ["", "A"] ==> some rows could pass (must keep)
true,
// s1 ["", ""] ==> all rows must pass (must keep)
true,
// s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> no rows can pass (not keep)
false,
// s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> no rows can pass (not keep)
false,
];
prune_with_expr(expr, &schema, &statistics, expected_ret);
}
#[test]
fn prune_utf8_not_like_one() {
let (schema, statistics) = utf8_setup();
let expr = col("s1").not_like(lit("A\u{10ffff}_"));
#[rustfmt::skip]
let expected_ret = &[
// s1 ["A", "Z"] ==> some rows could pass (must keep)
true,
// s1 ["A", "L"] ==> some rows could pass (must keep)
true,
// s1 ["N", "Z"] ==> some rows could pass (must keep)
true,
// s1 ["M", "M"] ==> some rows could pass (must keep)
true,
// s1 [NULL, NULL] ==> unknown (must keep)
true,
// s1 ["A", NULL] ==> some rows could pass (must keep)
true,
// s1 ["", "A"] ==> some rows could pass (must keep)
true,
// s1 ["", ""] ==> some rows could pass (must keep)
true,
// s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep)
true,
// s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> no row match. (min, max) maybe truncate
// original (min, max) maybe ("A\u{10ffff}\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}\u{10ffff}\u{10ffff}")
true,
];
prune_with_expr(expr, &schema, &statistics, expected_ret);
}
#[test]
fn prune_utf8_not_like_many() {
let (schema, statistics) = utf8_setup();
let expr = col("s1").not_like(lit("A\u{10ffff}%"));
#[rustfmt::skip]
let expected_ret = &[
// s1 ["A", "Z"] ==> some rows could pass (must keep)
true,
// s1 ["A", "L"] ==> some rows could pass (must keep)
true,
// s1 ["N", "Z"] ==> some rows could pass (must keep)
true,
// s1 ["M", "M"] ==> some rows could pass (must keep)
true,
// s1 [NULL, NULL] ==> unknown (must keep)
true,
// s1 ["A", NULL] ==> some rows could pass (must keep)
true,
// s1 ["", "A"] ==> some rows could pass (must keep)
true,
// s1 ["", ""] ==> some rows could pass (must keep)
true,
// s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep)
true,
// s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> no row match
false,
];
prune_with_expr(expr, &schema, &statistics, expected_ret);
let expr = col("s1").not_like(lit("A\u{10ffff}%\u{10ffff}"));
#[rustfmt::skip]
let expected_ret = &[
// s1 ["A", "Z"] ==> some rows could pass (must keep)
true,
// s1 ["A", "L"] ==> some rows could pass (must keep)
true,
// s1 ["N", "Z"] ==> some rows could pass (must keep)
true,
// s1 ["M", "M"] ==> some rows could pass (must keep)
true,
// s1 [NULL, NULL] ==> unknown (must keep)
true,
// s1 ["A", NULL] ==> some rows could pass (must keep)
true,
// s1 ["", "A"] ==> some rows could pass (must keep)
true,
// s1 ["", ""] ==> some rows could pass (must keep)
true,
// s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep)
true,
// s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep)
true,
];
prune_with_expr(expr, &schema, &statistics, expected_ret);
let expr = col("s1").not_like(lit("A\u{10ffff}%\u{10ffff}_"));
#[rustfmt::skip]
let expected_ret = &[
// s1 ["A", "Z"] ==> some rows could pass (must keep)
true,
// s1 ["A", "L"] ==> some rows could pass (must keep)
true,
// s1 ["N", "Z"] ==> some rows could pass (must keep)
true,
// s1 ["M", "M"] ==> some rows could pass (must keep)
true,
// s1 [NULL, NULL] ==> unknown (must keep)
true,
// s1 ["A", NULL] ==> some rows could pass (must keep)
true,
// s1 ["", "A"] ==> some rows could pass (must keep)
true,
// s1 ["", ""] ==> some rows could pass (must keep)
true,
// s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep)
true,
// s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep)
true,
];
prune_with_expr(expr, &schema, &statistics, expected_ret);
let expr = col("s1").not_like(lit("A\\%%"));
let statistics = TestStatistics::new().with(
"s1",
ContainerStats::new_utf8(
vec![Some("A%a"), Some("A")],
vec![Some("A%c"), Some("A")],
),
);
let expected_ret = &[false, true];
prune_with_expr(expr, &schema, &statistics, expected_ret);
}
#[test]
fn test_rewrite_expr_to_prunable() {
let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
let df_schema = DFSchema::try_from(schema.clone()).unwrap();
// column op lit
let left_input = col("a");
let left_input = logical2physical(&left_input, &schema);
let right_input = lit(ScalarValue::Int32(Some(12)));
let right_input = logical2physical(&right_input, &schema);
let (result_left, _, result_right) = rewrite_expr_to_prunable(
&left_input,
Operator::Eq,
&right_input,
df_schema.clone(),
)
.unwrap();
assert_eq!(result_left.to_string(), left_input.to_string());
assert_eq!(result_right.to_string(), right_input.to_string());
// cast op lit
let left_input = cast(col("a"), DataType::Decimal128(20, 3));
let left_input = logical2physical(&left_input, &schema);
let right_input = lit(ScalarValue::Decimal128(Some(12), 20, 3));
let right_input = logical2physical(&right_input, &schema);
let (result_left, _, result_right) = rewrite_expr_to_prunable(
&left_input,
Operator::Gt,
&right_input,
df_schema.clone(),
)
.unwrap();
assert_eq!(result_left.to_string(), left_input.to_string());
assert_eq!(result_right.to_string(), right_input.to_string());
// try_cast op lit
let left_input = try_cast(col("a"), DataType::Int64);
let left_input = logical2physical(&left_input, &schema);
let right_input = lit(ScalarValue::Int64(Some(12)));
let right_input = logical2physical(&right_input, &schema);
let (result_left, _, result_right) =
rewrite_expr_to_prunable(&left_input, Operator::Gt, &right_input, df_schema)
.unwrap();
assert_eq!(result_left.to_string(), left_input.to_string());
assert_eq!(result_right.to_string(), right_input.to_string());
// TODO: add test for other case and op
}
#[test]
fn test_rewrite_expr_to_prunable_custom_unhandled_hook() {
struct CustomUnhandledHook;
impl UnhandledPredicateHook for CustomUnhandledHook {
/// This handles an arbitrary case of a column that doesn't exist in the schema
/// by renaming it to yet another column that doesn't exist in the schema
/// (the transformation is arbitrary, the point is that it can do whatever it wants)
fn handle(&self, _expr: &Arc<dyn PhysicalExpr>) -> Arc<dyn PhysicalExpr> {
Arc::new(phys_expr::Literal::new(ScalarValue::Int32(Some(42))))
}
}
let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
let schema_with_b = Schema::new(vec![
Field::new("a", DataType::Int32, true),
Field::new("b", DataType::Int32, true),
]);
let rewriter = PredicateRewriter::new()
.with_unhandled_hook(Arc::new(CustomUnhandledHook {}));
let transform_expr = |expr| {
let expr = logical2physical(&expr, &schema_with_b);
rewriter.rewrite_predicate_to_statistics_predicate(&expr, &schema)
};
// transform an arbitrary valid expression that we know is handled
let known_expression = col("a").eq(lit(12));
let known_expression_transformed = PredicateRewriter::new()
.rewrite_predicate_to_statistics_predicate(
&logical2physical(&known_expression, &schema),
&schema,
);
// an expression referencing an unknown column (that is not in the schema) gets passed to the hook
let input = col("b").eq(lit(12));
let expected = logical2physical(&lit(42), &schema);
let transformed = transform_expr(input.clone());
assert_eq!(transformed.to_string(), expected.to_string());
// more complex case with unknown column
let input = known_expression.clone().and(input.clone());
let expected = phys_expr::BinaryExpr::new(
Arc::<dyn PhysicalExpr>::clone(&known_expression_transformed),
Operator::And,
logical2physical(&lit(42), &schema),
);
let transformed = transform_expr(input.clone());
assert_eq!(transformed.to_string(), expected.to_string());
// an unknown expression gets passed to the hook
let input = array_has(make_array(vec![lit(1)]), col("a"));
let expected = logical2physical(&lit(42), &schema);
let transformed = transform_expr(input.clone());
assert_eq!(transformed.to_string(), expected.to_string());
// more complex case with unknown expression
let input = known_expression.and(input);
let expected = phys_expr::BinaryExpr::new(
Arc::<dyn PhysicalExpr>::clone(&known_expression_transformed),
Operator::And,
logical2physical(&lit(42), &schema),
);
let transformed = transform_expr(input.clone());
assert_eq!(transformed.to_string(), expected.to_string());
}
#[test]
fn test_rewrite_expr_to_prunable_error() {
// cast string value to numeric value
// this cast is not supported
let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]);
let df_schema = DFSchema::try_from(schema.clone()).unwrap();
let left_input = cast(col("a"), DataType::Int64);
let left_input = logical2physical(&left_input, &schema);
let right_input = lit(ScalarValue::Int64(Some(12)));
let right_input = logical2physical(&right_input, &schema);
let result = rewrite_expr_to_prunable(
&left_input,
Operator::Gt,
&right_input,
df_schema.clone(),
);
assert!(result.is_err());
// other expr
let left_input = is_null(col("a"));
let left_input = logical2physical(&left_input, &schema);
let right_input = lit(ScalarValue::Int64(Some(12)));
let right_input = logical2physical(&right_input, &schema);
let result =
rewrite_expr_to_prunable(&left_input, Operator::Gt, &right_input, df_schema);
assert!(result.is_err());
// TODO: add other negative test for other case and op
}
#[test]
fn prune_with_contained_one_column() {
let schema = Arc::new(Schema::new(vec![Field::new("s1", DataType::Utf8, true)]));
// Model having information like a bloom filter for s1
let statistics = TestStatistics::new()
.with_contained(
"s1",
[ScalarValue::from("foo")],
[
// container 0 known to only contain "foo"",
Some(true),
// container 1 known to not contain "foo"
Some(false),
// container 2 unknown about "foo"
None,
// container 3 known to only contain "foo"
Some(true),
// container 4 known to not contain "foo"
Some(false),
// container 5 unknown about "foo"
None,
// container 6 known to only contain "foo"
Some(true),
// container 7 known to not contain "foo"
Some(false),
// container 8 unknown about "foo"
None,
],
)
.with_contained(
"s1",
[ScalarValue::from("bar")],
[
// containers 0,1,2 known to only contain "bar"
Some(true),
Some(true),
Some(true),
// container 3,4,5 known to not contain "bar"
Some(false),
Some(false),
Some(false),
// container 6,7,8 unknown about "bar"
None,
None,
None,
],
)
.with_contained(
// the way the tests are setup, this data is
// consulted if the "foo" and "bar" are being checked at the same time
"s1",
[ScalarValue::from("foo"), ScalarValue::from("bar")],
[
// container 0,1,2 unknown about ("foo, "bar")
None,
None,
None,
// container 3,4,5 known to contain only either "foo" and "bar"
Some(true),
Some(true),
Some(true),
// container 6,7,8 known to contain neither "foo" and "bar"
Some(false),
Some(false),
Some(false),
],
);
// s1 = 'foo'
prune_with_expr(
col("s1").eq(lit("foo")),
&schema,
&statistics,
// rule out containers ('false) where we know foo is not present
&[true, false, true, true, false, true, true, false, true],
);
// s1 = 'bar'
prune_with_expr(
col("s1").eq(lit("bar")),
&schema,
&statistics,
// rule out containers where we know bar is not present
&[true, true, true, false, false, false, true, true, true],
);
// s1 = 'baz' (unknown value)
prune_with_expr(
col("s1").eq(lit("baz")),
&schema,
&statistics,
// can't rule out anything
&[true, true, true, true, true, true, true, true, true],
);
// s1 = 'foo' AND s1 = 'bar'
prune_with_expr(
col("s1").eq(lit("foo")).and(col("s1").eq(lit("bar"))),
&schema,
&statistics,
// logically this predicate can't possibly be true (the column can't
// take on both values) but we could rule it out if the stats tell
// us that both values are not present
&[true, true, true, true, true, true, true, true, true],
);
// s1 = 'foo' OR s1 = 'bar'
prune_with_expr(
col("s1").eq(lit("foo")).or(col("s1").eq(lit("bar"))),
&schema,
&statistics,
// can rule out containers that we know contain neither foo nor bar
&[true, true, true, true, true, true, false, false, false],
);
// s1 = 'foo' OR s1 = 'baz'
prune_with_expr(
col("s1").eq(lit("foo")).or(col("s1").eq(lit("baz"))),
&schema,
&statistics,
// can't rule out anything container
&[true, true, true, true, true, true, true, true, true],
);
// s1 = 'foo' OR s1 = 'bar' OR s1 = 'baz'
prune_with_expr(
col("s1")
.eq(lit("foo"))
.or(col("s1").eq(lit("bar")))
.or(col("s1").eq(lit("baz"))),
&schema,
&statistics,
// can rule out any containers based on knowledge of s1 and `foo`,
// `bar` and (`foo`, `bar`)
&[true, true, true, true, true, true, true, true, true],
);
// s1 != foo
prune_with_expr(
col("s1").not_eq(lit("foo")),
&schema,
&statistics,
// rule out containers we know for sure only contain foo
&[false, true, true, false, true, true, false, true, true],
);
// s1 != bar
prune_with_expr(
col("s1").not_eq(lit("bar")),
&schema,
&statistics,
// rule out when we know for sure s1 has the value bar
&[false, false, false, true, true, true, true, true, true],
);
// s1 != foo AND s1 != bar
prune_with_expr(
col("s1")
.not_eq(lit("foo"))
.and(col("s1").not_eq(lit("bar"))),
&schema,
&statistics,
// can rule out any container where we know s1 does not have either 'foo' or 'bar'
&[true, true, true, false, false, false, true, true, true],
);
// s1 != foo AND s1 != bar AND s1 != baz
prune_with_expr(
col("s1")
.not_eq(lit("foo"))
.and(col("s1").not_eq(lit("bar")))
.and(col("s1").not_eq(lit("baz"))),
&schema,
&statistics,
// can't rule out any container based on knowledge of s1,s2
&[true, true, true, true, true, true, true, true, true],
);
// s1 != foo OR s1 != bar
prune_with_expr(
col("s1")
.not_eq(lit("foo"))
.or(col("s1").not_eq(lit("bar"))),
&schema,
&statistics,
// cant' rule out anything based on contains information
&[true, true, true, true, true, true, true, true, true],
);
// s1 != foo OR s1 != bar OR s1 != baz
prune_with_expr(
col("s1")
.not_eq(lit("foo"))
.or(col("s1").not_eq(lit("bar")))
.or(col("s1").not_eq(lit("baz"))),
&schema,
&statistics,
// cant' rule out anything based on contains information
&[true, true, true, true, true, true, true, true, true],
);
}
#[test]
fn prune_with_contained_two_columns() {
let schema = Arc::new(Schema::new(vec![
Field::new("s1", DataType::Utf8, true),
Field::new("s2", DataType::Utf8, true),
]));
// Model having information like bloom filters for s1 and s2
let statistics = TestStatistics::new()
.with_contained(
"s1",
[ScalarValue::from("foo")],
[
// container 0, s1 known to only contain "foo"",
Some(true),
// container 1, s1 known to not contain "foo"
Some(false),
// container 2, s1 unknown about "foo"
None,
// container 3, s1 known to only contain "foo"
Some(true),
// container 4, s1 known to not contain "foo"
Some(false),
// container 5, s1 unknown about "foo"
None,
// container 6, s1 known to only contain "foo"
Some(true),
// container 7, s1 known to not contain "foo"
Some(false),
// container 8, s1 unknown about "foo"
None,
],
)
.with_contained(
"s2", // for column s2
[ScalarValue::from("bar")],
[
// containers 0,1,2 s2 known to only contain "bar"
Some(true),
Some(true),
Some(true),
// container 3,4,5 s2 known to not contain "bar"
Some(false),
Some(false),
Some(false),
// container 6,7,8 s2 unknown about "bar"
None,
None,
None,
],
);
// s1 = 'foo'
prune_with_expr(
col("s1").eq(lit("foo")),
&schema,
&statistics,
// rule out containers where we know s1 is not present
&[true, false, true, true, false, true, true, false, true],
);
// s1 = 'foo' OR s2 = 'bar'
let expr = col("s1").eq(lit("foo")).or(col("s2").eq(lit("bar")));
prune_with_expr(
expr,
&schema,
&statistics,
// can't rule out any container (would need to prove that s1 != foo AND s2 != bar)
&[true, true, true, true, true, true, true, true, true],
);
// s1 = 'foo' AND s2 != 'bar'
prune_with_expr(
col("s1").eq(lit("foo")).and(col("s2").not_eq(lit("bar"))),
&schema,
&statistics,
// can only rule out container where we know either:
// 1. s1 doesn't have the value 'foo` or
// 2. s2 has only the value of 'bar'
&[false, false, false, true, false, true, true, false, true],
);
// s1 != 'foo' AND s2 != 'bar'
prune_with_expr(
col("s1")
.not_eq(lit("foo"))
.and(col("s2").not_eq(lit("bar"))),
&schema,
&statistics,
// Can rule out any container where we know either
// 1. s1 has only the value 'foo'
// 2. s2 has only the value 'bar'
&[false, false, false, false, true, true, false, true, true],
);
// s1 != 'foo' AND (s2 = 'bar' OR s2 = 'baz')
prune_with_expr(
col("s1")
.not_eq(lit("foo"))
.and(col("s2").eq(lit("bar")).or(col("s2").eq(lit("baz")))),
&schema,
&statistics,
// Can rule out any container where we know s1 has only the value
// 'foo'. Can't use knowledge of s2 and bar to rule out anything
&[false, true, true, false, true, true, false, true, true],
);
// s1 like '%foo%bar%'
prune_with_expr(
col("s1").like(lit("foo%bar%")),
&schema,
&statistics,
// cant rule out anything with information we know
&[true, true, true, true, true, true, true, true, true],
);
// s1 like '%foo%bar%' AND s2 = 'bar'
prune_with_expr(
col("s1")
.like(lit("foo%bar%"))
.and(col("s2").eq(lit("bar"))),
&schema,
&statistics,
// can rule out any container where we know s2 does not have the value 'bar'
&[true, true, true, false, false, false, true, true, true],
);
// s1 like '%foo%bar%' OR s2 = 'bar'
prune_with_expr(
col("s1").like(lit("foo%bar%")).or(col("s2").eq(lit("bar"))),
&schema,
&statistics,
// can't rule out anything (we would have to prove that both the
// like and the equality must be false)
&[true, true, true, true, true, true, true, true, true],
);
}
#[test]
fn prune_with_range_and_contained() {
// Setup mimics range information for i, a bloom filter for s
let schema = Arc::new(Schema::new(vec![
Field::new("i", DataType::Int32, true),
Field::new("s", DataType::Utf8, true),
]));
let statistics = TestStatistics::new()
.with(
"i",
ContainerStats::new_i32(
// Container 0, 3, 6: [-5 to 5]
// Container 1, 4, 7: [10 to 20]
// Container 2, 5, 9: unknown
vec![
Some(-5),
Some(10),
None,
Some(-5),
Some(10),
None,
Some(-5),
Some(10),
None,
], // min
vec![
Some(5),
Some(20),
None,
Some(5),
Some(20),
None,
Some(5),
Some(20),
None,
], // max
),
)
// Add contained information about the s and "foo"
.with_contained(
"s",
[ScalarValue::from("foo")],
[
// container 0,1,2 known to only contain "foo"
Some(true),
Some(true),
Some(true),
// container 3,4,5 known to not contain "foo"
Some(false),
Some(false),
Some(false),
// container 6,7,8 unknown about "foo"
None,
None,
None,
],
);
// i = 0 and s = 'foo'
prune_with_expr(
col("i").eq(lit(0)).and(col("s").eq(lit("foo"))),
&schema,
&statistics,
// Can rule out container where we know that either:
// 1. 0 is outside the min/max range of i
// 1. s does not contain foo
// (range is false, and contained is false)
&[true, false, true, false, false, false, true, false, true],
);
// i = 0 and s != 'foo'
prune_with_expr(
col("i").eq(lit(0)).and(col("s").not_eq(lit("foo"))),
&schema,
&statistics,
// Can rule out containers where either:
// 1. 0 is outside the min/max range of i
// 2. s only contains foo
&[false, false, false, true, false, true, true, false, true],
);
// i = 0 OR s = 'foo'
prune_with_expr(
col("i").eq(lit(0)).or(col("s").eq(lit("foo"))),
&schema,
&statistics,
// in theory could rule out containers if we had min/max values for
// s as well. But in this case we don't so we can't rule out anything
&[true, true, true, true, true, true, true, true, true],
);
}
/// prunes the specified expr with the specified schema and statistics, and
/// ensures it returns expected.
///
/// `expected` is a vector of bools, where true means the row group should
/// be kept, and false means it should be pruned.
// TODO refactor other tests to use this to reduce boiler plate
fn prune_with_expr(
expr: Expr,
schema: &SchemaRef,
statistics: &TestStatistics,
expected: &[bool],
) {
println!("Pruning with expr: {expr}");
let expr = logical2physical(&expr, schema);
let p = PruningPredicate::try_new(expr, Arc::<Schema>::clone(schema)).unwrap();
let result = p.prune(statistics).unwrap();
assert_eq!(result, expected);
}
fn test_build_predicate_expression(
expr: &Expr,
schema: &Schema,
required_columns: &mut RequiredColumns,
) -> Arc<dyn PhysicalExpr> {
let expr = logical2physical(expr, schema);
let unhandled_hook = Arc::new(ConstantUnhandledPredicateHook::default()) as _;
build_predicate_expression(
&expr,
&Arc::new(schema.clone()),
required_columns,
&unhandled_hook,
)
}
#[test]
fn test_build_predicate_expression_with_false() {
let expr = lit(ScalarValue::Boolean(Some(false)));
let schema = Schema::empty();
let res =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
let expected = logical2physical(&expr, &schema);
assert_eq!(&res, &expected);
}
#[test]
fn test_build_predicate_expression_with_and_false() {
let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]);
let expr = and(
col("c1").eq(lit("a")),
lit(ScalarValue::Boolean(Some(false))),
);
let res =
test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
let expected = logical2physical(&lit(ScalarValue::Boolean(Some(false))), &schema);
assert_eq!(&res, &expected);
}
#[test]
fn test_build_predicate_expression_with_or_false() {
let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]);
let left_expr = col("c1").eq(lit("a"));
let right_expr = lit(ScalarValue::Boolean(Some(false)));
let res = test_build_predicate_expression(
&or(left_expr.clone(), right_expr.clone()),
&schema,
&mut RequiredColumns::new(),
);
let expected =
"c1_null_count@2 != row_count@3 AND c1_min@0 <= a AND a <= c1_max@1";
assert_eq!(res.to_string(), expected);
}
}