table/

predicate.rs

// Copyright 2023 Greptime Team
//
// Licensed 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.

use std::sync::Arc;

use arc_swap::ArcSwap;
use common_telemetry::{debug, warn};
use common_time::Timestamp;
use common_time::range::TimestampRange;
use common_time::timestamp::TimeUnit;
use datafusion::common::ScalarValue;
use datafusion::physical_optimizer::pruning::PruningPredicateBuilder;
use datafusion_common::ToDFSchema;
use datafusion_common::pruning::PruningStatistics;
use datafusion_common::tree_node::TreeNode;
use datafusion_expr::expr::{Expr, InList};
use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
use datafusion_expr::{Between, BinaryExpr, Operator};
use datafusion_physical_expr::execution_props::ExecutionProps;
use datafusion_physical_expr::expressions::{
    BinaryExpr as PhysicalBinaryExpr, DynamicFilterPhysicalExpr, is_null,
};
use datafusion_physical_expr::{PhysicalExpr, create_physical_expr};
use datatypes::arrow;
use datatypes::value::scalar_value_to_timestamp;
use snafu::ResultExt;

use crate::error;

#[cfg(test)]
mod stats;

/// Assert the scalar value is not utf8. Returns `None` if it's utf8.
/// In theory, it should be converted to a timestamp scalar value by `TypeConversionRule`.
macro_rules! return_none_if_utf8 {
    ($lit: ident) => {
        if is_string_timestamp_literal($lit) {
            warn!(
                "Unexpected ScalarValue::Utf8 in time range predicate: {:?}. Maybe it's an implicit bug, please report it to https://github.com/GreptimeTeam/greptimedb/issues",
                $lit
            );

            // Make the predicate ineffective.
            return None;
        }
    };
}

pub fn is_string_timestamp_literal(scalar: &ScalarValue) -> bool {
    matches!(
        scalar,
        ScalarValue::Utf8(_) | ScalarValue::LargeUtf8(_) | ScalarValue::Utf8View(_)
    )
}

/// Reference-counted pointer to a list of logical exprs and a list of dynamic filter physical exprs.
#[derive(Debug, Clone, Default)]
pub struct Predicate {
    /// logical exprs
    exprs: Arc<Vec<Expr>>,
    /// dynamic filter physical exprs, only useful if dynamic filtering is enabled
    ///
    /// They are usually from `TopK` or `Join` operators, and can dynamically filter data during query execution by using current runtime information to further reduce data scanning
    dyn_filters: Arc<ArcSwap<Vec<Arc<DynamicFilterPhysicalExpr>>>>,
}

impl Predicate {
    /// Creates a new `Predicate` by converting logical exprs to physical exprs that can be
    /// evaluated against record batches.
    /// Returns error when failed to convert exprs.
    pub fn new(exprs: Vec<Expr>) -> Self {
        Self {
            exprs: Arc::new(exprs),
            dyn_filters: Arc::new(ArcSwap::new(Arc::new(vec![]))),
        }
    }

    pub fn with_dyn_filters(
        exprs: Vec<Expr>,
        dyn_filters: Vec<Arc<DynamicFilterPhysicalExpr>>,
    ) -> Self {
        Self {
            exprs: Arc::new(exprs),
            dyn_filters: Arc::new(ArcSwap::new(Arc::new(dyn_filters))),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.exprs.is_empty() && self.dyn_filters.load().is_empty()
    }

    /// Adds dynamic filter physical exprs to the existing list.
    pub fn add_dyn_filters(&self, dyn_filters: Vec<Arc<DynamicFilterPhysicalExpr>>) {
        self.dyn_filters.rcu(|existing| {
            let mut new_filters = existing.as_ref().clone();
            new_filters.extend(dyn_filters.clone());
            Arc::new(new_filters)
        });
    }

    /// Removes dynamic filters while preserving the static logical expressions.
    pub fn clear_dyn_filters(&self) {
        self.dyn_filters.store(Arc::new(vec![]));
    }

    /// Returns the logical exprs.
    pub fn exprs(&self) -> &[Expr] {
        &self.exprs
    }

    /// Returns the dynamic filter physical exprs. Notice this return a live dynamic filters which
    /// can change during query execution.
    pub fn dyn_filters(&self) -> Arc<Vec<Arc<DynamicFilterPhysicalExpr>>> {
        self.dyn_filters.load_full()
    }

    /// Returns the dynamic filter as physical exprs. Notice this return a "snapshot" of
    /// dynamic filters at the time of calling this method.
    pub fn dyn_filter_phy_exprs(&self) -> error::Result<Vec<Arc<dyn PhysicalExpr>>> {
        self.dyn_filters
            .load()
            .iter()
            .map(|e| {
                // Pruning must preserve NULL inputs just like decoded dynamic filtering.
                e.children()
                    .into_iter()
                    .try_fold(e.current()?, |expr, child| {
                        Ok(Arc::new(PhysicalBinaryExpr::new(
                            expr,
                            Operator::Or,
                            is_null(child.clone())?,
                        )) as Arc<dyn PhysicalExpr>)
                    })
            })
            .collect::<Result<Vec<_>, _>>()
            .context(error::DatafusionSnafu)
    }

    /// Builds a single physical expr according to provided schema.
    pub fn to_physical_expr(
        expr: &Expr,
        schema: &arrow::datatypes::SchemaRef,
    ) -> error::Result<Arc<dyn PhysicalExpr>> {
        let df_schema = schema
            .clone()
            .to_dfschema_ref()
            .context(error::DatafusionSnafu)?;

        // TODO(hl): `execution_props` provides variables required by evaluation.
        // we may reuse the `execution_props` from `SessionState` once we support
        // registering variables.
        let execution_props = &ExecutionProps::new();

        create_physical_expr(
            expr,
            df_schema.as_ref(),
            execution_props,
            &PhysicalPlanningContext::default(),
        )
        .context(error::DatafusionSnafu)
    }

    /// Builds physical exprs according to provided schema.
    pub fn to_physical_exprs(
        &self,
        schema: &arrow::datatypes::SchemaRef,
    ) -> error::Result<Vec<Arc<dyn PhysicalExpr>>> {
        let dyn_filters = self.dyn_filter_phy_exprs()?;

        Ok(self
            .exprs
            .iter()
            .filter_map(|expr| Self::to_physical_expr(expr, schema).ok())
            .chain(dyn_filters)
            .collect::<Vec<_>>())
    }

    /// Evaluates the predicate against the `stats`.
    /// Returns a vector of boolean values, among which `false` means the row group can be skipped.
    pub fn prune_with_stats<S: PruningStatistics>(
        &self,
        stats: &S,
        schema: &arrow::datatypes::SchemaRef,
    ) -> Vec<bool> {
        let mut res = vec![true; stats.num_containers()];
        let physical_exprs = match self.to_physical_exprs(schema) {
            Ok(expr) => expr,
            Err(e) => {
                warn!(e; "Failed to build physical expr from predicates: {:?}", &self.exprs);
                return res;
            }
        };

        for expr in &physical_exprs {
            match PruningPredicateBuilder::new()
                .with_file_schema(schema.clone())
                .try_build(expr.clone())
            {
                Ok(p) => match p.prune(stats) {
                    Ok(r) => {
                        for (curr_val, res) in r.into_iter().zip(res.iter_mut()) {
                            *res &= curr_val
                        }
                    }
                    Err(e) => {
                        warn!(e; "Failed to prune row groups");
                    }
                },
                Err(e) => {
                    // since dynamic filter exprs could be complex, it's possible that the pruning predicate builder fails to prove anything from it. In that case, we just log it and skip pruning with this expr.
                    debug!("Failed to create pruning predicate for expr: {e:?}");
                }
            }
        }
        res
    }
}

// tests for `build_time_range_predicate` locates in src/query/tests/time_range_filter_test.rs
// since it requires query engine to convert sql to filters.
/// `build_time_range_predicate` extracts time range from logical exprs to facilitate fast
/// time range pruning.
pub fn build_time_range_predicate(
    ts_col_name: &str,
    ts_col_unit: TimeUnit,
    filters: &[Expr],
) -> TimestampRange {
    let mut res = TimestampRange::min_to_max();
    for expr in filters {
        if let Some(range) = extract_time_range_from_expr(ts_col_name, ts_col_unit, expr) {
            res = res.and(&range);
        }
    }
    res
}

/// The outcome of strictly extracting the time range of `ts_col_name` from a
/// scan's filters. Unlike [`build_time_range_predicate`], which quietly widens
/// to `min_to_max` on anything it cannot parse (fine for pruning), this
/// distinguishes "the column is not filtered at all" from "it is filtered in a
/// way that cannot be safely turned into a range" — for callers whose contract
/// forbids silently falling back to a default window.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TimeRangeExtraction {
    /// No filter references the column.
    Absent,
    /// Every filter referencing the column was folded into this range, which
    /// over-approximates the filters' satisfying set.
    Extracted(TimestampRange),
    /// At least one filter references the column in a shape that cannot be
    /// safely extracted (e.g. under `OR`/`NOT`, or compared to a non-literal).
    Unsupported,
}

/// Strictly extracts the time range of `ts_col_name` from the (implicitly
/// AND-ed) scan filters. See [`TimeRangeExtraction`].
pub fn extract_time_range_strict(
    ts_col_name: &str,
    ts_col_unit: TimeUnit,
    filters: &[Expr],
) -> TimeRangeExtraction {
    let mut range: Option<TimestampRange> = None;
    for expr in filters {
        if !expr
            .column_refs()
            .iter()
            .any(|column| column.name == ts_col_name)
        {
            continue;
        }
        // Disjunctive shapes (`OR`, `IN`) collapse disjoint ranges into their
        // convex hull, which is not exactly representable as one contiguous
        // range; callers deriving synthetic timestamps from the bounds would
        // emit points inside the gaps.
        if contains_disjunction_over_column(expr, ts_col_name) {
            return TimeRangeExtraction::Unsupported;
        }
        // `Some` from the lenient extractor over-approximates the expression's
        // satisfying set (an `AND` side it cannot parse is dropped, which only
        // widens), so intersecting extracted conjuncts stays an
        // over-approximation.
        match extract_time_range_from_expr(ts_col_name, ts_col_unit, expr) {
            Some(extracted) => {
                range = Some(match range {
                    Some(acc) => acc.and(&extracted),
                    None => extracted,
                });
            }
            None => return TimeRangeExtraction::Unsupported,
        }
    }
    match range {
        Some(range) => TimeRangeExtraction::Extracted(range),
        None => TimeRangeExtraction::Absent,
    }
}

fn contains_disjunction_over_column(expr: &Expr, ts_col_name: &str) -> bool {
    let is_disjunction = |expr: &Expr| {
        matches!(
            expr,
            Expr::BinaryExpr(BinaryExpr {
                op: Operator::Or,
                ..
            }) | Expr::InList(_)
        ) && expr
            .column_refs()
            .iter()
            .any(|column| column.name == ts_col_name)
    };
    expr.exists(|expr| Ok(is_disjunction(expr))).unwrap_or(true)
}

/// Extract time range filter from `WHERE`/`IN (...)`/`BETWEEN` clauses.
/// Return None if no time range can be found in expr.
pub fn extract_time_range_from_expr(
    ts_col_name: &str,
    ts_col_unit: TimeUnit,
    expr: &Expr,
) -> Option<TimestampRange> {
    match expr {
        Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
            extract_from_binary_expr(ts_col_name, ts_col_unit, left, op, right)
        }
        Expr::Between(Between {
            expr,
            negated,
            low,
            high,
        }) => extract_from_between_expr(ts_col_name, ts_col_unit, expr, negated, low, high),
        Expr::InList(InList {
            expr,
            list,
            negated,
        }) => extract_from_in_list_expr(ts_col_name, expr, *negated, list),
        _ => None,
    }
}

fn extract_from_binary_expr(
    ts_col_name: &str,
    ts_col_unit: TimeUnit,
    left: &Expr,
    op: &Operator,
    right: &Expr,
) -> Option<TimestampRange> {
    match op {
        Operator::Eq => get_timestamp_filter(ts_col_name, left, right)
            .and_then(|(ts, _)| ts.convert_to(ts_col_unit))
            .map(TimestampRange::single),
        Operator::Lt => {
            let (ts, reverse) = get_timestamp_filter(ts_col_name, left, right)?;
            if reverse {
                // [lit] < ts_col
                let ts_val = ts.convert_to(ts_col_unit)?.value();
                Some(TimestampRange::from_start(Timestamp::new(
                    ts_val + 1,
                    ts_col_unit,
                )))
            } else {
                // ts_col < [lit]
                ts.convert_to_ceil(ts_col_unit)
                    .map(|ts| TimestampRange::until_end(ts, false))
            }
        }
        Operator::LtEq => {
            let (ts, reverse) = get_timestamp_filter(ts_col_name, left, right)?;
            if reverse {
                // [lit] <= ts_col
                ts.convert_to_ceil(ts_col_unit)
                    .map(TimestampRange::from_start)
            } else {
                // ts_col <= [lit]
                ts.convert_to(ts_col_unit)
                    .map(|ts| TimestampRange::until_end(ts, true))
            }
        }
        Operator::Gt => {
            let (ts, reverse) = get_timestamp_filter(ts_col_name, left, right)?;
            if reverse {
                // [lit] > ts_col
                ts.convert_to_ceil(ts_col_unit)
                    .map(|t| TimestampRange::until_end(t, false))
            } else {
                // ts_col > [lit]
                let ts_val = ts.convert_to(ts_col_unit)?.value();
                Some(TimestampRange::from_start(Timestamp::new(
                    ts_val + 1,
                    ts_col_unit,
                )))
            }
        }
        Operator::GtEq => {
            let (ts, reverse) = get_timestamp_filter(ts_col_name, left, right)?;
            if reverse {
                // [lit] >= ts_col
                ts.convert_to(ts_col_unit)
                    .map(|t| TimestampRange::until_end(t, true))
            } else {
                // ts_col >= [lit]
                ts.convert_to_ceil(ts_col_unit)
                    .map(TimestampRange::from_start)
            }
        }
        Operator::And => {
            // instead of return none when failed to extract time range from left/right, we unwrap the none into
            // `TimestampRange::min_to_max`.
            let left = extract_time_range_from_expr(ts_col_name, ts_col_unit, left)
                .unwrap_or_else(TimestampRange::min_to_max);
            let right = extract_time_range_from_expr(ts_col_name, ts_col_unit, right)
                .unwrap_or_else(TimestampRange::min_to_max);
            Some(left.and(&right))
        }
        Operator::Or => {
            let left = extract_time_range_from_expr(ts_col_name, ts_col_unit, left)?;
            let right = extract_time_range_from_expr(ts_col_name, ts_col_unit, right)?;
            Some(left.or(&right))
        }
        _ => None,
    }
}

fn get_timestamp_filter(ts_col_name: &str, left: &Expr, right: &Expr) -> Option<(Timestamp, bool)> {
    let (col, lit, reverse) = match (left, right) {
        (Expr::Column(column), Expr::Literal(scalar, _)) => (column, scalar, false),
        (Expr::Literal(scalar, _), Expr::Column(column)) => (column, scalar, true),
        _ => {
            return None;
        }
    };
    if col.name != ts_col_name {
        return None;
    }

    return_none_if_utf8!(lit);
    scalar_value_to_timestamp(lit, None).map(|t| (t, reverse))
}

fn extract_from_between_expr(
    ts_col_name: &str,
    ts_col_unit: TimeUnit,
    expr: &Expr,
    negated: &bool,
    low: &Expr,
    high: &Expr,
) -> Option<TimestampRange> {
    let Expr::Column(col) = expr else {
        return None;
    };
    if col.name != ts_col_name {
        return None;
    }

    if *negated {
        return None;
    }

    match (low, high) {
        (Expr::Literal(low, _), Expr::Literal(high, _)) => {
            return_none_if_utf8!(low);
            return_none_if_utf8!(high);

            let low_opt =
                scalar_value_to_timestamp(low, None).and_then(|ts| ts.convert_to(ts_col_unit));
            let high_opt = scalar_value_to_timestamp(high, None)
                .and_then(|ts| ts.convert_to_ceil(ts_col_unit));
            Some(TimestampRange::new_inclusive(low_opt, high_opt))
        }
        _ => None,
    }
}

/// Extract time range filter from `IN (...)` expr.
fn extract_from_in_list_expr(
    ts_col_name: &str,
    expr: &Expr,
    negated: bool,
    list: &[Expr],
) -> Option<TimestampRange> {
    if negated {
        return None;
    }
    let Expr::Column(col) = expr else {
        return None;
    };
    if col.name != ts_col_name {
        return None;
    }

    if list.is_empty() {
        return Some(TimestampRange::empty());
    }
    let mut init_range = TimestampRange::empty();
    for expr in list {
        if let Expr::Literal(scalar, _) = expr {
            return_none_if_utf8!(scalar);
            // TODO(hl): maybe we should raise an error here since cannot parse
            // timestamp value from in list expr
            let timestamp = scalar_value_to_timestamp(scalar, None)?;
            init_range = init_range.or(&TimestampRange::single(timestamp));
        }
    }
    Some(init_range)
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use common_test_util::temp_dir::{TempDir, create_temp_dir};
    use datafusion::parquet::arrow::ArrowWriter;
    use datafusion_common::{Column, ScalarValue};
    use datafusion_expr::{BinaryExpr, Literal, Operator, col, lit};
    use datatypes::arrow::array::Int32Array;
    use datatypes::arrow::datatypes::{DataType, Field, Schema};
    use datatypes::arrow::record_batch::RecordBatch;
    use datatypes::arrow_array::StringArray;
    use parquet::arrow::ParquetRecordBatchStreamBuilder;
    use parquet::file::properties::WriterProperties;

    use super::*;
    use crate::predicate::stats::RowGroupPruningStatistics;

    fn check_build_predicate(expr: Expr, expect: TimestampRange) {
        assert_eq!(
            expect,
            build_time_range_predicate("ts", TimeUnit::Millisecond, &[expr])
        );
    }

    #[test]
    fn test_gt() {
        // ts > 1ms
        check_build_predicate(
            col("ts").gt(lit(ScalarValue::TimestampMillisecond(Some(1), None))),
            TimestampRange::from_start(Timestamp::new_millisecond(2)),
        );

        // 1ms > ts
        check_build_predicate(
            lit(ScalarValue::TimestampMillisecond(Some(1), None)).gt(col("ts")),
            TimestampRange::until_end(Timestamp::new_millisecond(1), false),
        );

        // 1001us > ts
        check_build_predicate(
            lit(ScalarValue::TimestampMicrosecond(Some(1001), None)).gt(col("ts")),
            TimestampRange::until_end(Timestamp::new_millisecond(1), true),
        );

        // ts > 1001us
        check_build_predicate(
            col("ts").gt(lit(ScalarValue::TimestampMicrosecond(Some(1001), None))),
            TimestampRange::from_start(Timestamp::new_millisecond(2)),
        );

        // 1s > ts
        check_build_predicate(
            lit(ScalarValue::TimestampSecond(Some(1), None)).gt(col("ts")),
            TimestampRange::until_end(Timestamp::new_millisecond(1000), false),
        );

        // ts > 1s
        check_build_predicate(
            col("ts").gt(lit(ScalarValue::TimestampSecond(Some(1), None))),
            TimestampRange::from_start(Timestamp::new_millisecond(1001)),
        );
    }

    #[test]
    fn test_gt_eq() {
        // ts >= 1ms
        check_build_predicate(
            col("ts").gt_eq(lit(ScalarValue::TimestampMillisecond(Some(1), None))),
            TimestampRange::from_start(Timestamp::new_millisecond(1)),
        );

        // 1ms >= ts
        check_build_predicate(
            lit(ScalarValue::TimestampMillisecond(Some(1), None)).gt_eq(col("ts")),
            TimestampRange::until_end(Timestamp::new_millisecond(1), true),
        );

        // 1001us >= ts
        check_build_predicate(
            lit(ScalarValue::TimestampMicrosecond(Some(1001), None)).gt_eq(col("ts")),
            TimestampRange::until_end(Timestamp::new_millisecond(1), true),
        );

        // ts >= 1001us
        check_build_predicate(
            col("ts").gt_eq(lit(ScalarValue::TimestampMicrosecond(Some(1001), None))),
            TimestampRange::from_start(Timestamp::new_millisecond(2)),
        );

        // 1s >= ts
        check_build_predicate(
            lit(ScalarValue::TimestampSecond(Some(1), None)).gt_eq(col("ts")),
            TimestampRange::until_end(Timestamp::new_millisecond(1000), true),
        );

        // ts >= 1s
        check_build_predicate(
            col("ts").gt_eq(lit(ScalarValue::TimestampSecond(Some(1), None))),
            TimestampRange::from_start(Timestamp::new_millisecond(1000)),
        );
    }

    #[test]
    fn test_lt() {
        // ts < 1ms
        check_build_predicate(
            col("ts").lt(lit(ScalarValue::TimestampMillisecond(Some(1), None))),
            TimestampRange::until_end(Timestamp::new_millisecond(1), false),
        );

        // 1ms < ts
        check_build_predicate(
            lit(ScalarValue::TimestampMillisecond(Some(1), None)).lt(col("ts")),
            TimestampRange::from_start(Timestamp::new_millisecond(2)),
        );

        // 1001us < ts
        check_build_predicate(
            lit(ScalarValue::TimestampMicrosecond(Some(1001), None)).lt(col("ts")),
            TimestampRange::from_start(Timestamp::new_millisecond(2)),
        );

        // ts < 1001us
        check_build_predicate(
            col("ts").lt(lit(ScalarValue::TimestampMicrosecond(Some(1001), None))),
            TimestampRange::until_end(Timestamp::new_millisecond(1), true),
        );

        // 1s < ts
        check_build_predicate(
            lit(ScalarValue::TimestampSecond(Some(1), None)).lt(col("ts")),
            TimestampRange::from_start(Timestamp::new_millisecond(1001)),
        );

        // ts < 1s
        check_build_predicate(
            col("ts").lt(lit(ScalarValue::TimestampSecond(Some(1), None))),
            TimestampRange::until_end(Timestamp::new_millisecond(1000), false),
        );
    }

    #[test]
    fn test_lt_eq() {
        // ts <= 1ms
        check_build_predicate(
            col("ts").lt_eq(lit(ScalarValue::TimestampMillisecond(Some(1), None))),
            TimestampRange::until_end(Timestamp::new_millisecond(1), true),
        );

        // 1ms <= ts
        check_build_predicate(
            lit(ScalarValue::TimestampMillisecond(Some(1), None)).lt_eq(col("ts")),
            TimestampRange::from_start(Timestamp::new_millisecond(1)),
        );

        // 1001us <= ts
        check_build_predicate(
            lit(ScalarValue::TimestampMicrosecond(Some(1001), None)).lt_eq(col("ts")),
            TimestampRange::from_start(Timestamp::new_millisecond(2)),
        );

        // ts <= 1001us
        check_build_predicate(
            col("ts").lt_eq(lit(ScalarValue::TimestampMicrosecond(Some(1001), None))),
            TimestampRange::until_end(Timestamp::new_millisecond(1), true),
        );

        // 1s <= ts
        check_build_predicate(
            lit(ScalarValue::TimestampSecond(Some(1), None)).lt_eq(col("ts")),
            TimestampRange::from_start(Timestamp::new_millisecond(1000)),
        );

        // ts <= 1s
        check_build_predicate(
            col("ts").lt_eq(lit(ScalarValue::TimestampSecond(Some(1), None))),
            TimestampRange::until_end(Timestamp::new_millisecond(1000), true),
        );
    }

    #[test]
    fn test_extract_time_range_strict() {
        fn ts_lit(ms: i64) -> Expr {
            lit(ScalarValue::TimestampMillisecond(Some(ms), None))
        }
        let extract =
            |filters: &[Expr]| extract_time_range_strict("ts", TimeUnit::Millisecond, filters);
        let range = |start: i64, end: i64| {
            TimestampRange::new(
                Timestamp::new_millisecond(start),
                Timestamp::new_millisecond(end),
            )
            .unwrap()
        };

        // No filter references the column.
        assert_eq!(extract(&[]), TimeRangeExtraction::Absent);
        assert_eq!(
            extract(&[col("host").eq(lit("a"))]),
            TimeRangeExtraction::Absent
        );

        // Both bounds across conjuncts; unrelated filters are ignored.
        assert_eq!(
            extract(&[
                col("ts").gt_eq(ts_lit(1000)),
                col("ts").lt(ts_lit(2000)),
                col("host").eq(lit("a")),
            ]),
            TimeRangeExtraction::Extracted(range(1000, 2000))
        );

        // Lower bound only / upper bound only.
        assert_eq!(
            extract(&[col("ts").gt_eq(ts_lit(1000))]),
            TimeRangeExtraction::Extracted(TimestampRange::from_start(Timestamp::new_millisecond(
                1000
            )))
        );
        assert_eq!(
            extract(&[col("ts").lt(ts_lit(2000))]),
            TimeRangeExtraction::Extracted(TimestampRange::until_end(
                Timestamp::new_millisecond(2000),
                false
            ))
        );

        // BETWEEN is inclusive on both ends.
        assert_eq!(
            extract(&[col("ts").between(ts_lit(1000), ts_lit(2000))]),
            TimeRangeExtraction::Extracted(range(1000, 2001))
        );

        // Equality pins a single point.
        assert_eq!(
            extract(&[col("ts").eq(ts_lit(1500))]),
            TimeRangeExtraction::Extracted(TimestampRange::single(Timestamp::new_millisecond(
                1500
            )))
        );

        // An unparsable side under AND only widens; the extraction stays safe.
        assert_eq!(
            extract(&[col("ts").gt_eq(ts_lit(1000)).and(col("ts").lt(col("t2")))]),
            TimeRangeExtraction::Extracted(TimestampRange::from_start(Timestamp::new_millisecond(
                1000
            )))
        );

        // Contradictory bounds collapse to the empty range, not an error.
        let TimeRangeExtraction::Extracted(empty) =
            extract(&[col("ts").gt_eq(ts_lit(2000)), col("ts").lt(ts_lit(1000))])
        else {
            panic!("expected an extraction");
        };
        assert!(empty.is_empty());

        // Shapes that could widen the satisfying set beyond what is extractable
        // must be refused: OR with an unparsable side, NOT, non-literal bounds.
        assert_eq!(
            extract(&[col("ts").gt(ts_lit(1000)).or(col("host").eq(lit("a")))]),
            TimeRangeExtraction::Unsupported
        );
        assert_eq!(
            extract(&[!col("ts").gt(ts_lit(1000))]),
            TimeRangeExtraction::Unsupported
        );
        assert_eq!(
            extract(&[col("ts").gt_eq(col("t2"))]),
            TimeRangeExtraction::Unsupported
        );
    }

    async fn gen_test_parquet_file(dir: &TempDir, cnt: usize) -> (String, Arc<Schema>) {
        let path = dir
            .path()
            .join("test-prune.parquet")
            .to_string_lossy()
            .to_string();

        let name_field = Field::new("name", DataType::Utf8, true);
        let count_field = Field::new("cnt", DataType::Int32, true);
        let schema = Arc::new(Schema::new(vec![name_field, count_field]));

        let file = std::fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(path.clone())
            .unwrap();

        let write_props = WriterProperties::builder()
            .set_max_row_group_row_count(Some(10))
            .build();
        let mut writer = ArrowWriter::try_new(file, schema.clone(), Some(write_props)).unwrap();

        for i in (0..cnt).step_by(10) {
            let name_array = Arc::new(StringArray::from(
                (i..(i + 10).min(cnt))
                    .map(|i| i.to_string())
                    .collect::<Vec<_>>(),
            )) as Arc<_>;
            let count_array = Arc::new(Int32Array::from(
                (i..(i + 10).min(cnt)).map(|i| i as i32).collect::<Vec<_>>(),
            )) as Arc<_>;
            let rb = RecordBatch::try_new(schema.clone(), vec![name_array, count_array]).unwrap();
            writer.write(&rb).unwrap();
        }
        let _ = writer.close().unwrap();
        (path, schema)
    }

    async fn assert_prune(array_cnt: usize, filters: Vec<Expr>, expect: Vec<bool>) {
        let dir = create_temp_dir("prune_parquet");
        let (path, arrow_schema) = gen_test_parquet_file(&dir, array_cnt).await;
        let schema = Arc::new(datatypes::schema::Schema::try_from(arrow_schema.clone()).unwrap());
        let arrow_predicate = Predicate::new(filters);
        let builder = ParquetRecordBatchStreamBuilder::new(
            tokio::fs::OpenOptions::new()
                .read(true)
                .open(path)
                .await
                .unwrap(),
        )
        .await
        .unwrap();
        let metadata = builder.metadata().clone();
        let row_groups = metadata.row_groups();

        let stats = RowGroupPruningStatistics::new(row_groups, &schema);
        let res = arrow_predicate.prune_with_stats(&stats, &arrow_schema);
        assert_eq!(expect, res);
    }

    #[test]
    fn test_clear_dyn_filters_preserves_static_predicates() {
        use datafusion_physical_expr::expressions::lit as physical_lit;

        let static_exprs = vec![col("a").eq(lit(1_i32))];
        let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(vec![], physical_lit(true)));
        let predicate =
            Predicate::with_dyn_filters(static_exprs.clone(), vec![dynamic_filter.clone()]);

        predicate.clear_dyn_filters();
        // An update from the old producer must not put its wrapper back into this execution.
        dynamic_filter.update(physical_lit(false)).unwrap();

        assert_eq!(predicate.exprs(), static_exprs);
        assert!(predicate.dyn_filters().is_empty());
        assert!(predicate.dyn_filter_phy_exprs().unwrap().is_empty());
    }

    #[tokio::test]
    async fn test_dynamic_pruning_keeps_null_row_group() {
        use datafusion_physical_expr::expressions::{
            Column as PhysicalColumn, lit as physical_lit,
        };

        let dir = create_temp_dir("dynamic_pruning_nulls");
        let path = dir.path().join("nullable.parquet");
        let arrow_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
        let file = std::fs::File::create(&path).unwrap();
        let mut writer = ArrowWriter::try_new(file, arrow_schema.clone(), None).unwrap();
        for values in [[None, Some(1)], [Some(1), Some(1)], [None, None]] {
            let batch = RecordBatch::try_new(
                arrow_schema.clone(),
                vec![Arc::new(Int32Array::from(values.to_vec()))],
            )
            .unwrap();
            writer.write(&batch).unwrap();
            writer.flush().unwrap();
        }
        writer.close().unwrap();
        let builder =
            ParquetRecordBatchStreamBuilder::new(tokio::fs::File::open(path).await.unwrap())
                .await
                .unwrap();
        let schema = Arc::new(datatypes::schema::Schema::try_from(arrow_schema.clone()).unwrap());
        let stats = RowGroupPruningStatistics::new(builder.metadata().row_groups(), &schema);
        let filter = Arc::new(DynamicFilterPhysicalExpr::new(
            vec![Arc::new(PhysicalColumn::new("a", 0))],
            physical_lit(true),
        ));
        let predicate = Predicate::with_dyn_filters(vec![], vec![filter.clone()]);
        assert_eq!(
            predicate.prune_with_stats(&stats, &arrow_schema),
            vec![true; 3]
        );
        filter
            .update(
                Predicate::to_physical_expr(
                    &col("a").gt_eq(lit(10_i32)).and(col("a").lt_eq(lit(10_i32))),
                    &arrow_schema,
                )
                .unwrap(),
            )
            .unwrap();
        // NULL inputs must reach decoded filtering; non-NULL misses can still be pruned.
        assert_eq!(
            predicate.prune_with_stats(&stats, &arrow_schema),
            vec![true, false, true],
        );
    }

    #[tokio::test]
    async fn test_clear_dyn_filters_restores_static_row_group_pruning() {
        use datafusion_physical_expr::expressions::{
            Column as PhysicalColumn, lit as physical_lit,
        };

        let dir = create_temp_dir("dynamic_pruning_reset");
        let (path, arrow_schema) = gen_test_parquet_file(&dir, 30).await;
        let schema = Arc::new(datatypes::schema::Schema::try_from(arrow_schema.clone()).unwrap());
        let builder =
            ParquetRecordBatchStreamBuilder::new(tokio::fs::File::open(path).await.unwrap())
                .await
                .unwrap();
        let stats = RowGroupPruningStatistics::new(builder.metadata().row_groups(), &schema);
        let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(
            vec![Arc::new(PhysicalColumn::new("cnt", 1))],
            physical_lit(true),
        ));
        let predicate = Predicate::with_dyn_filters(
            vec![col("cnt").gt_eq(lit(10_i32))],
            vec![dynamic_filter.clone()],
        );

        dynamic_filter
            .update(
                Predicate::to_physical_expr(&col("cnt").gt(lit(100_i32)), &arrow_schema).unwrap(),
            )
            .unwrap();
        assert_eq!(
            predicate.prune_with_stats(&stats, &arrow_schema),
            vec![false; 3]
        );

        // Reset after the prior stream is dropped: no dynamic filter remains, while the
        // static predicate still excludes only the first row group.
        predicate.clear_dyn_filters();
        assert_eq!(
            predicate.prune_with_stats(&stats, &arrow_schema),
            vec![false, true, true]
        );
    }

    fn gen_predicate(max_val: i32, op: Operator) -> Vec<Expr> {
        vec![datafusion_expr::Expr::BinaryExpr(BinaryExpr {
            left: Box::new(datafusion_expr::Expr::Column(Column::from_name("cnt"))),
            op,
            right: Box::new(max_val.lit()),
        })]
    }

    #[tokio::test]
    async fn test_prune_empty() {
        assert_prune(3, vec![], vec![true]).await;
    }

    #[tokio::test]
    async fn test_prune_all_match() {
        let p = gen_predicate(3, Operator::Gt);
        assert_prune(2, p, vec![false]).await;
    }

    #[tokio::test]
    async fn test_prune_gt() {
        let p = gen_predicate(29, Operator::Gt);
        assert_prune(
            100,
            p,
            vec![
                false, false, false, true, true, true, true, true, true, true,
            ],
        )
        .await;
    }

    #[tokio::test]
    async fn test_prune_eq_expr() {
        let p = gen_predicate(30, Operator::Eq);
        assert_prune(40, p, vec![false, false, false, true]).await;
    }

    #[tokio::test]
    async fn test_prune_neq_expr() {
        let p = gen_predicate(30, Operator::NotEq);
        assert_prune(40, p, vec![true, true, true, true]).await;
    }

    #[tokio::test]
    async fn test_prune_gteq_expr() {
        let p = gen_predicate(29, Operator::GtEq);
        assert_prune(40, p, vec![false, false, true, true]).await;
    }

    #[tokio::test]
    async fn test_prune_lt_expr() {
        let p = gen_predicate(30, Operator::Lt);
        assert_prune(40, p, vec![true, true, true, false]).await;
    }

    #[tokio::test]
    async fn test_prune_lteq_expr() {
        let p = gen_predicate(30, Operator::LtEq);
        assert_prune(40, p, vec![true, true, true, true]).await;
    }

    #[tokio::test]
    async fn test_prune_between_expr() {
        let p = gen_predicate(30, Operator::LtEq);
        assert_prune(40, p, vec![true, true, true, true]).await;
    }

    #[tokio::test]
    async fn test_or() {
        // cnt > 30 or cnt < 20
        let e = datafusion_expr::Expr::Column(Column::from_name("cnt"))
            .gt(30.lit())
            .or(datafusion_expr::Expr::Column(Column::from_name("cnt")).lt(20.lit()));
        assert_prune(40, vec![e], vec![true, true, false, true]).await;
    }

    #[tokio::test]
    async fn test_to_physical_expr() {
        let predicate = Predicate::new(vec![
            col("host").eq(lit("host_a")),
            col("ts").gt(lit(ScalarValue::TimestampMicrosecond(Some(123), None))),
        ]);

        let schema = Arc::new(arrow::datatypes::Schema::new(vec![Field::new(
            "host",
            arrow::datatypes::DataType::Utf8,
            false,
        )]));

        let predicates = predicate.to_physical_exprs(&schema).unwrap();
        assert!(!predicates.is_empty());

        let physical_expr = Predicate::to_physical_expr(&col("host").eq(lit("host_a")), &schema);
        assert!(physical_expr.is_ok());
    }
}