table/
predicate.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use arc_swap::ArcSwap;
18use common_telemetry::{debug, warn};
19use common_time::Timestamp;
20use common_time::range::TimestampRange;
21use common_time::timestamp::TimeUnit;
22use datafusion::common::ScalarValue;
23use datafusion::physical_optimizer::pruning::PruningPredicate;
24use datafusion_common::ToDFSchema;
25use datafusion_common::pruning::PruningStatistics;
26use datafusion_common::tree_node::TreeNode;
27use datafusion_expr::expr::{Expr, InList};
28use datafusion_expr::{Between, BinaryExpr, Operator};
29use datafusion_physical_expr::execution_props::ExecutionProps;
30use datafusion_physical_expr::expressions::{
31    BinaryExpr as PhysicalBinaryExpr, DynamicFilterPhysicalExpr, is_null,
32};
33use datafusion_physical_expr::{PhysicalExpr, create_physical_expr};
34use datatypes::arrow;
35use datatypes::value::scalar_value_to_timestamp;
36use snafu::ResultExt;
37
38use crate::error;
39
40#[cfg(test)]
41mod stats;
42
43/// Assert the scalar value is not utf8. Returns `None` if it's utf8.
44/// In theory, it should be converted to a timestamp scalar value by `TypeConversionRule`.
45macro_rules! return_none_if_utf8 {
46    ($lit: ident) => {
47        if is_string_timestamp_literal($lit) {
48            warn!(
49                "Unexpected ScalarValue::Utf8 in time range predicate: {:?}. Maybe it's an implicit bug, please report it to https://github.com/GreptimeTeam/greptimedb/issues",
50                $lit
51            );
52
53            // Make the predicate ineffective.
54            return None;
55        }
56    };
57}
58
59pub fn is_string_timestamp_literal(scalar: &ScalarValue) -> bool {
60    matches!(
61        scalar,
62        ScalarValue::Utf8(_) | ScalarValue::LargeUtf8(_) | ScalarValue::Utf8View(_)
63    )
64}
65
66/// Reference-counted pointer to a list of logical exprs and a list of dynamic filter physical exprs.
67#[derive(Debug, Clone, Default)]
68pub struct Predicate {
69    /// logical exprs
70    exprs: Arc<Vec<Expr>>,
71    /// dynamic filter physical exprs, only useful if dynamic filtering is enabled
72    ///
73    /// 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
74    dyn_filters: Arc<ArcSwap<Vec<Arc<DynamicFilterPhysicalExpr>>>>,
75}
76
77impl Predicate {
78    /// Creates a new `Predicate` by converting logical exprs to physical exprs that can be
79    /// evaluated against record batches.
80    /// Returns error when failed to convert exprs.
81    pub fn new(exprs: Vec<Expr>) -> Self {
82        Self {
83            exprs: Arc::new(exprs),
84            dyn_filters: Arc::new(ArcSwap::new(Arc::new(vec![]))),
85        }
86    }
87
88    pub fn with_dyn_filters(
89        exprs: Vec<Expr>,
90        dyn_filters: Vec<Arc<DynamicFilterPhysicalExpr>>,
91    ) -> Self {
92        Self {
93            exprs: Arc::new(exprs),
94            dyn_filters: Arc::new(ArcSwap::new(Arc::new(dyn_filters))),
95        }
96    }
97
98    pub fn is_empty(&self) -> bool {
99        self.exprs.is_empty() && self.dyn_filters.load().is_empty()
100    }
101
102    /// Adds dynamic filter physical exprs to the existing list.
103    pub fn add_dyn_filters(&self, dyn_filters: Vec<Arc<DynamicFilterPhysicalExpr>>) {
104        self.dyn_filters.rcu(|existing| {
105            let mut new_filters = existing.as_ref().clone();
106            new_filters.extend(dyn_filters.clone());
107            Arc::new(new_filters)
108        });
109    }
110
111    /// Removes dynamic filters while preserving the static logical expressions.
112    pub fn clear_dyn_filters(&self) {
113        self.dyn_filters.store(Arc::new(vec![]));
114    }
115
116    /// Returns the logical exprs.
117    pub fn exprs(&self) -> &[Expr] {
118        &self.exprs
119    }
120
121    /// Returns the dynamic filter physical exprs. Notice this return a live dynamic filters which
122    /// can change during query execution.
123    pub fn dyn_filters(&self) -> Arc<Vec<Arc<DynamicFilterPhysicalExpr>>> {
124        self.dyn_filters.load_full()
125    }
126
127    /// Returns the dynamic filter as physical exprs. Notice this return a "snapshot" of
128    /// dynamic filters at the time of calling this method.
129    pub fn dyn_filter_phy_exprs(&self) -> error::Result<Vec<Arc<dyn PhysicalExpr>>> {
130        self.dyn_filters
131            .load()
132            .iter()
133            .map(|e| {
134                // Pruning must preserve NULL inputs just like decoded dynamic filtering.
135                e.children()
136                    .into_iter()
137                    .try_fold(e.current()?, |expr, child| {
138                        Ok(Arc::new(PhysicalBinaryExpr::new(
139                            expr,
140                            Operator::Or,
141                            is_null(child.clone())?,
142                        )) as Arc<dyn PhysicalExpr>)
143                    })
144            })
145            .collect::<Result<Vec<_>, _>>()
146            .context(error::DatafusionSnafu)
147    }
148
149    /// Builds a single physical expr according to provided schema.
150    pub fn to_physical_expr(
151        expr: &Expr,
152        schema: &arrow::datatypes::SchemaRef,
153    ) -> error::Result<Arc<dyn PhysicalExpr>> {
154        let df_schema = schema
155            .clone()
156            .to_dfschema_ref()
157            .context(error::DatafusionSnafu)?;
158
159        // TODO(hl): `execution_props` provides variables required by evaluation.
160        // we may reuse the `execution_props` from `SessionState` once we support
161        // registering variables.
162        let execution_props = &ExecutionProps::new();
163
164        create_physical_expr(expr, df_schema.as_ref(), execution_props)
165            .context(error::DatafusionSnafu)
166    }
167
168    /// Builds physical exprs according to provided schema.
169    pub fn to_physical_exprs(
170        &self,
171        schema: &arrow::datatypes::SchemaRef,
172    ) -> error::Result<Vec<Arc<dyn PhysicalExpr>>> {
173        let dyn_filters = self.dyn_filter_phy_exprs()?;
174
175        Ok(self
176            .exprs
177            .iter()
178            .filter_map(|expr| Self::to_physical_expr(expr, schema).ok())
179            .chain(dyn_filters)
180            .collect::<Vec<_>>())
181    }
182
183    /// Evaluates the predicate against the `stats`.
184    /// Returns a vector of boolean values, among which `false` means the row group can be skipped.
185    pub fn prune_with_stats<S: PruningStatistics>(
186        &self,
187        stats: &S,
188        schema: &arrow::datatypes::SchemaRef,
189    ) -> Vec<bool> {
190        let mut res = vec![true; stats.num_containers()];
191        let physical_exprs = match self.to_physical_exprs(schema) {
192            Ok(expr) => expr,
193            Err(e) => {
194                warn!(e; "Failed to build physical expr from predicates: {:?}", &self.exprs);
195                return res;
196            }
197        };
198
199        for expr in &physical_exprs {
200            match PruningPredicate::try_new(expr.clone(), schema.clone()) {
201                Ok(p) => match p.prune(stats) {
202                    Ok(r) => {
203                        for (curr_val, res) in r.into_iter().zip(res.iter_mut()) {
204                            *res &= curr_val
205                        }
206                    }
207                    Err(e) => {
208                        warn!(e; "Failed to prune row groups");
209                    }
210                },
211                Err(e) => {
212                    // since dynamic filter exprs could be complex, it's possible that `PruningPredicate::try_new` fails to prove anything from it. In that case, we just log it and skip pruning with this expr.
213                    debug!("Failed to create pruning predicate for expr: {e:?}");
214                }
215            }
216        }
217        res
218    }
219}
220
221// tests for `build_time_range_predicate` locates in src/query/tests/time_range_filter_test.rs
222// since it requires query engine to convert sql to filters.
223/// `build_time_range_predicate` extracts time range from logical exprs to facilitate fast
224/// time range pruning.
225pub fn build_time_range_predicate(
226    ts_col_name: &str,
227    ts_col_unit: TimeUnit,
228    filters: &[Expr],
229) -> TimestampRange {
230    let mut res = TimestampRange::min_to_max();
231    for expr in filters {
232        if let Some(range) = extract_time_range_from_expr(ts_col_name, ts_col_unit, expr) {
233            res = res.and(&range);
234        }
235    }
236    res
237}
238
239/// The outcome of strictly extracting the time range of `ts_col_name` from a
240/// scan's filters. Unlike [`build_time_range_predicate`], which quietly widens
241/// to `min_to_max` on anything it cannot parse (fine for pruning), this
242/// distinguishes "the column is not filtered at all" from "it is filtered in a
243/// way that cannot be safely turned into a range" — for callers whose contract
244/// forbids silently falling back to a default window.
245#[derive(Debug, Clone, PartialEq, Eq)]
246pub enum TimeRangeExtraction {
247    /// No filter references the column.
248    Absent,
249    /// Every filter referencing the column was folded into this range, which
250    /// over-approximates the filters' satisfying set.
251    Extracted(TimestampRange),
252    /// At least one filter references the column in a shape that cannot be
253    /// safely extracted (e.g. under `OR`/`NOT`, or compared to a non-literal).
254    Unsupported,
255}
256
257/// Strictly extracts the time range of `ts_col_name` from the (implicitly
258/// AND-ed) scan filters. See [`TimeRangeExtraction`].
259pub fn extract_time_range_strict(
260    ts_col_name: &str,
261    ts_col_unit: TimeUnit,
262    filters: &[Expr],
263) -> TimeRangeExtraction {
264    let mut range: Option<TimestampRange> = None;
265    for expr in filters {
266        if !expr
267            .column_refs()
268            .iter()
269            .any(|column| column.name == ts_col_name)
270        {
271            continue;
272        }
273        // Disjunctive shapes (`OR`, `IN`) collapse disjoint ranges into their
274        // convex hull, which is not exactly representable as one contiguous
275        // range; callers deriving synthetic timestamps from the bounds would
276        // emit points inside the gaps.
277        if contains_disjunction_over_column(expr, ts_col_name) {
278            return TimeRangeExtraction::Unsupported;
279        }
280        // `Some` from the lenient extractor over-approximates the expression's
281        // satisfying set (an `AND` side it cannot parse is dropped, which only
282        // widens), so intersecting extracted conjuncts stays an
283        // over-approximation.
284        match extract_time_range_from_expr(ts_col_name, ts_col_unit, expr) {
285            Some(extracted) => {
286                range = Some(match range {
287                    Some(acc) => acc.and(&extracted),
288                    None => extracted,
289                });
290            }
291            None => return TimeRangeExtraction::Unsupported,
292        }
293    }
294    match range {
295        Some(range) => TimeRangeExtraction::Extracted(range),
296        None => TimeRangeExtraction::Absent,
297    }
298}
299
300fn contains_disjunction_over_column(expr: &Expr, ts_col_name: &str) -> bool {
301    let is_disjunction = |expr: &Expr| {
302        matches!(
303            expr,
304            Expr::BinaryExpr(BinaryExpr {
305                op: Operator::Or,
306                ..
307            }) | Expr::InList(_)
308        ) && expr
309            .column_refs()
310            .iter()
311            .any(|column| column.name == ts_col_name)
312    };
313    expr.exists(|expr| Ok(is_disjunction(expr))).unwrap_or(true)
314}
315
316/// Extract time range filter from `WHERE`/`IN (...)`/`BETWEEN` clauses.
317/// Return None if no time range can be found in expr.
318pub fn extract_time_range_from_expr(
319    ts_col_name: &str,
320    ts_col_unit: TimeUnit,
321    expr: &Expr,
322) -> Option<TimestampRange> {
323    match expr {
324        Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
325            extract_from_binary_expr(ts_col_name, ts_col_unit, left, op, right)
326        }
327        Expr::Between(Between {
328            expr,
329            negated,
330            low,
331            high,
332        }) => extract_from_between_expr(ts_col_name, ts_col_unit, expr, negated, low, high),
333        Expr::InList(InList {
334            expr,
335            list,
336            negated,
337        }) => extract_from_in_list_expr(ts_col_name, expr, *negated, list),
338        _ => None,
339    }
340}
341
342fn extract_from_binary_expr(
343    ts_col_name: &str,
344    ts_col_unit: TimeUnit,
345    left: &Expr,
346    op: &Operator,
347    right: &Expr,
348) -> Option<TimestampRange> {
349    match op {
350        Operator::Eq => get_timestamp_filter(ts_col_name, left, right)
351            .and_then(|(ts, _)| ts.convert_to(ts_col_unit))
352            .map(TimestampRange::single),
353        Operator::Lt => {
354            let (ts, reverse) = get_timestamp_filter(ts_col_name, left, right)?;
355            if reverse {
356                // [lit] < ts_col
357                let ts_val = ts.convert_to(ts_col_unit)?.value();
358                Some(TimestampRange::from_start(Timestamp::new(
359                    ts_val + 1,
360                    ts_col_unit,
361                )))
362            } else {
363                // ts_col < [lit]
364                ts.convert_to_ceil(ts_col_unit)
365                    .map(|ts| TimestampRange::until_end(ts, false))
366            }
367        }
368        Operator::LtEq => {
369            let (ts, reverse) = get_timestamp_filter(ts_col_name, left, right)?;
370            if reverse {
371                // [lit] <= ts_col
372                ts.convert_to_ceil(ts_col_unit)
373                    .map(TimestampRange::from_start)
374            } else {
375                // ts_col <= [lit]
376                ts.convert_to(ts_col_unit)
377                    .map(|ts| TimestampRange::until_end(ts, true))
378            }
379        }
380        Operator::Gt => {
381            let (ts, reverse) = get_timestamp_filter(ts_col_name, left, right)?;
382            if reverse {
383                // [lit] > ts_col
384                ts.convert_to_ceil(ts_col_unit)
385                    .map(|t| TimestampRange::until_end(t, false))
386            } else {
387                // ts_col > [lit]
388                let ts_val = ts.convert_to(ts_col_unit)?.value();
389                Some(TimestampRange::from_start(Timestamp::new(
390                    ts_val + 1,
391                    ts_col_unit,
392                )))
393            }
394        }
395        Operator::GtEq => {
396            let (ts, reverse) = get_timestamp_filter(ts_col_name, left, right)?;
397            if reverse {
398                // [lit] >= ts_col
399                ts.convert_to(ts_col_unit)
400                    .map(|t| TimestampRange::until_end(t, true))
401            } else {
402                // ts_col >= [lit]
403                ts.convert_to_ceil(ts_col_unit)
404                    .map(TimestampRange::from_start)
405            }
406        }
407        Operator::And => {
408            // instead of return none when failed to extract time range from left/right, we unwrap the none into
409            // `TimestampRange::min_to_max`.
410            let left = extract_time_range_from_expr(ts_col_name, ts_col_unit, left)
411                .unwrap_or_else(TimestampRange::min_to_max);
412            let right = extract_time_range_from_expr(ts_col_name, ts_col_unit, right)
413                .unwrap_or_else(TimestampRange::min_to_max);
414            Some(left.and(&right))
415        }
416        Operator::Or => {
417            let left = extract_time_range_from_expr(ts_col_name, ts_col_unit, left)?;
418            let right = extract_time_range_from_expr(ts_col_name, ts_col_unit, right)?;
419            Some(left.or(&right))
420        }
421        _ => None,
422    }
423}
424
425fn get_timestamp_filter(ts_col_name: &str, left: &Expr, right: &Expr) -> Option<(Timestamp, bool)> {
426    let (col, lit, reverse) = match (left, right) {
427        (Expr::Column(column), Expr::Literal(scalar, _)) => (column, scalar, false),
428        (Expr::Literal(scalar, _), Expr::Column(column)) => (column, scalar, true),
429        _ => {
430            return None;
431        }
432    };
433    if col.name != ts_col_name {
434        return None;
435    }
436
437    return_none_if_utf8!(lit);
438    scalar_value_to_timestamp(lit, None).map(|t| (t, reverse))
439}
440
441fn extract_from_between_expr(
442    ts_col_name: &str,
443    ts_col_unit: TimeUnit,
444    expr: &Expr,
445    negated: &bool,
446    low: &Expr,
447    high: &Expr,
448) -> Option<TimestampRange> {
449    let Expr::Column(col) = expr else {
450        return None;
451    };
452    if col.name != ts_col_name {
453        return None;
454    }
455
456    if *negated {
457        return None;
458    }
459
460    match (low, high) {
461        (Expr::Literal(low, _), Expr::Literal(high, _)) => {
462            return_none_if_utf8!(low);
463            return_none_if_utf8!(high);
464
465            let low_opt =
466                scalar_value_to_timestamp(low, None).and_then(|ts| ts.convert_to(ts_col_unit));
467            let high_opt = scalar_value_to_timestamp(high, None)
468                .and_then(|ts| ts.convert_to_ceil(ts_col_unit));
469            Some(TimestampRange::new_inclusive(low_opt, high_opt))
470        }
471        _ => None,
472    }
473}
474
475/// Extract time range filter from `IN (...)` expr.
476fn extract_from_in_list_expr(
477    ts_col_name: &str,
478    expr: &Expr,
479    negated: bool,
480    list: &[Expr],
481) -> Option<TimestampRange> {
482    if negated {
483        return None;
484    }
485    let Expr::Column(col) = expr else {
486        return None;
487    };
488    if col.name != ts_col_name {
489        return None;
490    }
491
492    if list.is_empty() {
493        return Some(TimestampRange::empty());
494    }
495    let mut init_range = TimestampRange::empty();
496    for expr in list {
497        if let Expr::Literal(scalar, _) = expr {
498            return_none_if_utf8!(scalar);
499            if let Some(timestamp) = scalar_value_to_timestamp(scalar, None) {
500                init_range = init_range.or(&TimestampRange::single(timestamp))
501            } else {
502                // TODO(hl): maybe we should raise an error here since cannot parse
503                // timestamp value from in list expr
504                return None;
505            }
506        }
507    }
508    Some(init_range)
509}
510
511#[cfg(test)]
512mod tests {
513    use std::sync::Arc;
514
515    use common_test_util::temp_dir::{TempDir, create_temp_dir};
516    use datafusion::parquet::arrow::ArrowWriter;
517    use datafusion_common::{Column, ScalarValue};
518    use datafusion_expr::{BinaryExpr, Literal, Operator, col, lit};
519    use datatypes::arrow::array::Int32Array;
520    use datatypes::arrow::datatypes::{DataType, Field, Schema};
521    use datatypes::arrow::record_batch::RecordBatch;
522    use datatypes::arrow_array::StringArray;
523    use parquet::arrow::ParquetRecordBatchStreamBuilder;
524    use parquet::file::properties::WriterProperties;
525
526    use super::*;
527    use crate::predicate::stats::RowGroupPruningStatistics;
528
529    fn check_build_predicate(expr: Expr, expect: TimestampRange) {
530        assert_eq!(
531            expect,
532            build_time_range_predicate("ts", TimeUnit::Millisecond, &[expr])
533        );
534    }
535
536    #[test]
537    fn test_gt() {
538        // ts > 1ms
539        check_build_predicate(
540            col("ts").gt(lit(ScalarValue::TimestampMillisecond(Some(1), None))),
541            TimestampRange::from_start(Timestamp::new_millisecond(2)),
542        );
543
544        // 1ms > ts
545        check_build_predicate(
546            lit(ScalarValue::TimestampMillisecond(Some(1), None)).gt(col("ts")),
547            TimestampRange::until_end(Timestamp::new_millisecond(1), false),
548        );
549
550        // 1001us > ts
551        check_build_predicate(
552            lit(ScalarValue::TimestampMicrosecond(Some(1001), None)).gt(col("ts")),
553            TimestampRange::until_end(Timestamp::new_millisecond(1), true),
554        );
555
556        // ts > 1001us
557        check_build_predicate(
558            col("ts").gt(lit(ScalarValue::TimestampMicrosecond(Some(1001), None))),
559            TimestampRange::from_start(Timestamp::new_millisecond(2)),
560        );
561
562        // 1s > ts
563        check_build_predicate(
564            lit(ScalarValue::TimestampSecond(Some(1), None)).gt(col("ts")),
565            TimestampRange::until_end(Timestamp::new_millisecond(1000), false),
566        );
567
568        // ts > 1s
569        check_build_predicate(
570            col("ts").gt(lit(ScalarValue::TimestampSecond(Some(1), None))),
571            TimestampRange::from_start(Timestamp::new_millisecond(1001)),
572        );
573    }
574
575    #[test]
576    fn test_gt_eq() {
577        // ts >= 1ms
578        check_build_predicate(
579            col("ts").gt_eq(lit(ScalarValue::TimestampMillisecond(Some(1), None))),
580            TimestampRange::from_start(Timestamp::new_millisecond(1)),
581        );
582
583        // 1ms >= ts
584        check_build_predicate(
585            lit(ScalarValue::TimestampMillisecond(Some(1), None)).gt_eq(col("ts")),
586            TimestampRange::until_end(Timestamp::new_millisecond(1), true),
587        );
588
589        // 1001us >= ts
590        check_build_predicate(
591            lit(ScalarValue::TimestampMicrosecond(Some(1001), None)).gt_eq(col("ts")),
592            TimestampRange::until_end(Timestamp::new_millisecond(1), true),
593        );
594
595        // ts >= 1001us
596        check_build_predicate(
597            col("ts").gt_eq(lit(ScalarValue::TimestampMicrosecond(Some(1001), None))),
598            TimestampRange::from_start(Timestamp::new_millisecond(2)),
599        );
600
601        // 1s >= ts
602        check_build_predicate(
603            lit(ScalarValue::TimestampSecond(Some(1), None)).gt_eq(col("ts")),
604            TimestampRange::until_end(Timestamp::new_millisecond(1000), true),
605        );
606
607        // ts >= 1s
608        check_build_predicate(
609            col("ts").gt_eq(lit(ScalarValue::TimestampSecond(Some(1), None))),
610            TimestampRange::from_start(Timestamp::new_millisecond(1000)),
611        );
612    }
613
614    #[test]
615    fn test_lt() {
616        // ts < 1ms
617        check_build_predicate(
618            col("ts").lt(lit(ScalarValue::TimestampMillisecond(Some(1), None))),
619            TimestampRange::until_end(Timestamp::new_millisecond(1), false),
620        );
621
622        // 1ms < ts
623        check_build_predicate(
624            lit(ScalarValue::TimestampMillisecond(Some(1), None)).lt(col("ts")),
625            TimestampRange::from_start(Timestamp::new_millisecond(2)),
626        );
627
628        // 1001us < ts
629        check_build_predicate(
630            lit(ScalarValue::TimestampMicrosecond(Some(1001), None)).lt(col("ts")),
631            TimestampRange::from_start(Timestamp::new_millisecond(2)),
632        );
633
634        // ts < 1001us
635        check_build_predicate(
636            col("ts").lt(lit(ScalarValue::TimestampMicrosecond(Some(1001), None))),
637            TimestampRange::until_end(Timestamp::new_millisecond(1), true),
638        );
639
640        // 1s < ts
641        check_build_predicate(
642            lit(ScalarValue::TimestampSecond(Some(1), None)).lt(col("ts")),
643            TimestampRange::from_start(Timestamp::new_millisecond(1001)),
644        );
645
646        // ts < 1s
647        check_build_predicate(
648            col("ts").lt(lit(ScalarValue::TimestampSecond(Some(1), None))),
649            TimestampRange::until_end(Timestamp::new_millisecond(1000), false),
650        );
651    }
652
653    #[test]
654    fn test_lt_eq() {
655        // ts <= 1ms
656        check_build_predicate(
657            col("ts").lt_eq(lit(ScalarValue::TimestampMillisecond(Some(1), None))),
658            TimestampRange::until_end(Timestamp::new_millisecond(1), true),
659        );
660
661        // 1ms <= ts
662        check_build_predicate(
663            lit(ScalarValue::TimestampMillisecond(Some(1), None)).lt_eq(col("ts")),
664            TimestampRange::from_start(Timestamp::new_millisecond(1)),
665        );
666
667        // 1001us <= ts
668        check_build_predicate(
669            lit(ScalarValue::TimestampMicrosecond(Some(1001), None)).lt_eq(col("ts")),
670            TimestampRange::from_start(Timestamp::new_millisecond(2)),
671        );
672
673        // ts <= 1001us
674        check_build_predicate(
675            col("ts").lt_eq(lit(ScalarValue::TimestampMicrosecond(Some(1001), None))),
676            TimestampRange::until_end(Timestamp::new_millisecond(1), true),
677        );
678
679        // 1s <= ts
680        check_build_predicate(
681            lit(ScalarValue::TimestampSecond(Some(1), None)).lt_eq(col("ts")),
682            TimestampRange::from_start(Timestamp::new_millisecond(1000)),
683        );
684
685        // ts <= 1s
686        check_build_predicate(
687            col("ts").lt_eq(lit(ScalarValue::TimestampSecond(Some(1), None))),
688            TimestampRange::until_end(Timestamp::new_millisecond(1000), true),
689        );
690    }
691
692    #[test]
693    fn test_extract_time_range_strict() {
694        fn ts_lit(ms: i64) -> Expr {
695            lit(ScalarValue::TimestampMillisecond(Some(ms), None))
696        }
697        let extract =
698            |filters: &[Expr]| extract_time_range_strict("ts", TimeUnit::Millisecond, filters);
699        let range = |start: i64, end: i64| {
700            TimestampRange::new(
701                Timestamp::new_millisecond(start),
702                Timestamp::new_millisecond(end),
703            )
704            .unwrap()
705        };
706
707        // No filter references the column.
708        assert_eq!(extract(&[]), TimeRangeExtraction::Absent);
709        assert_eq!(
710            extract(&[col("host").eq(lit("a"))]),
711            TimeRangeExtraction::Absent
712        );
713
714        // Both bounds across conjuncts; unrelated filters are ignored.
715        assert_eq!(
716            extract(&[
717                col("ts").gt_eq(ts_lit(1000)),
718                col("ts").lt(ts_lit(2000)),
719                col("host").eq(lit("a")),
720            ]),
721            TimeRangeExtraction::Extracted(range(1000, 2000))
722        );
723
724        // Lower bound only / upper bound only.
725        assert_eq!(
726            extract(&[col("ts").gt_eq(ts_lit(1000))]),
727            TimeRangeExtraction::Extracted(TimestampRange::from_start(Timestamp::new_millisecond(
728                1000
729            )))
730        );
731        assert_eq!(
732            extract(&[col("ts").lt(ts_lit(2000))]),
733            TimeRangeExtraction::Extracted(TimestampRange::until_end(
734                Timestamp::new_millisecond(2000),
735                false
736            ))
737        );
738
739        // BETWEEN is inclusive on both ends.
740        assert_eq!(
741            extract(&[col("ts").between(ts_lit(1000), ts_lit(2000))]),
742            TimeRangeExtraction::Extracted(range(1000, 2001))
743        );
744
745        // Equality pins a single point.
746        assert_eq!(
747            extract(&[col("ts").eq(ts_lit(1500))]),
748            TimeRangeExtraction::Extracted(TimestampRange::single(Timestamp::new_millisecond(
749                1500
750            )))
751        );
752
753        // An unparsable side under AND only widens; the extraction stays safe.
754        assert_eq!(
755            extract(&[col("ts").gt_eq(ts_lit(1000)).and(col("ts").lt(col("t2")))]),
756            TimeRangeExtraction::Extracted(TimestampRange::from_start(Timestamp::new_millisecond(
757                1000
758            )))
759        );
760
761        // Contradictory bounds collapse to the empty range, not an error.
762        let TimeRangeExtraction::Extracted(empty) =
763            extract(&[col("ts").gt_eq(ts_lit(2000)), col("ts").lt(ts_lit(1000))])
764        else {
765            panic!("expected an extraction");
766        };
767        assert!(empty.is_empty());
768
769        // Shapes that could widen the satisfying set beyond what is extractable
770        // must be refused: OR with an unparsable side, NOT, non-literal bounds.
771        assert_eq!(
772            extract(&[col("ts").gt(ts_lit(1000)).or(col("host").eq(lit("a")))]),
773            TimeRangeExtraction::Unsupported
774        );
775        assert_eq!(
776            extract(&[!col("ts").gt(ts_lit(1000))]),
777            TimeRangeExtraction::Unsupported
778        );
779        assert_eq!(
780            extract(&[col("ts").gt_eq(col("t2"))]),
781            TimeRangeExtraction::Unsupported
782        );
783    }
784
785    async fn gen_test_parquet_file(dir: &TempDir, cnt: usize) -> (String, Arc<Schema>) {
786        let path = dir
787            .path()
788            .join("test-prune.parquet")
789            .to_string_lossy()
790            .to_string();
791
792        let name_field = Field::new("name", DataType::Utf8, true);
793        let count_field = Field::new("cnt", DataType::Int32, true);
794        let schema = Arc::new(Schema::new(vec![name_field, count_field]));
795
796        let file = std::fs::OpenOptions::new()
797            .write(true)
798            .create(true)
799            .truncate(true)
800            .open(path.clone())
801            .unwrap();
802
803        let write_props = WriterProperties::builder()
804            .set_max_row_group_row_count(Some(10))
805            .build();
806        let mut writer = ArrowWriter::try_new(file, schema.clone(), Some(write_props)).unwrap();
807
808        for i in (0..cnt).step_by(10) {
809            let name_array = Arc::new(StringArray::from(
810                (i..(i + 10).min(cnt))
811                    .map(|i| i.to_string())
812                    .collect::<Vec<_>>(),
813            )) as Arc<_>;
814            let count_array = Arc::new(Int32Array::from(
815                (i..(i + 10).min(cnt)).map(|i| i as i32).collect::<Vec<_>>(),
816            )) as Arc<_>;
817            let rb = RecordBatch::try_new(schema.clone(), vec![name_array, count_array]).unwrap();
818            writer.write(&rb).unwrap();
819        }
820        let _ = writer.close().unwrap();
821        (path, schema)
822    }
823
824    async fn assert_prune(array_cnt: usize, filters: Vec<Expr>, expect: Vec<bool>) {
825        let dir = create_temp_dir("prune_parquet");
826        let (path, arrow_schema) = gen_test_parquet_file(&dir, array_cnt).await;
827        let schema = Arc::new(datatypes::schema::Schema::try_from(arrow_schema.clone()).unwrap());
828        let arrow_predicate = Predicate::new(filters);
829        let builder = ParquetRecordBatchStreamBuilder::new(
830            tokio::fs::OpenOptions::new()
831                .read(true)
832                .open(path)
833                .await
834                .unwrap(),
835        )
836        .await
837        .unwrap();
838        let metadata = builder.metadata().clone();
839        let row_groups = metadata.row_groups();
840
841        let stats = RowGroupPruningStatistics::new(row_groups, &schema);
842        let res = arrow_predicate.prune_with_stats(&stats, &arrow_schema);
843        assert_eq!(expect, res);
844    }
845
846    #[test]
847    fn test_clear_dyn_filters_preserves_static_predicates() {
848        use datafusion_physical_expr::expressions::lit as physical_lit;
849
850        let static_exprs = vec![col("a").eq(lit(1_i32))];
851        let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(vec![], physical_lit(true)));
852        let predicate =
853            Predicate::with_dyn_filters(static_exprs.clone(), vec![dynamic_filter.clone()]);
854
855        predicate.clear_dyn_filters();
856        // An update from the old producer must not put its wrapper back into this execution.
857        dynamic_filter.update(physical_lit(false)).unwrap();
858
859        assert_eq!(predicate.exprs(), static_exprs);
860        assert!(predicate.dyn_filters().is_empty());
861        assert!(predicate.dyn_filter_phy_exprs().unwrap().is_empty());
862    }
863
864    #[tokio::test]
865    async fn test_dynamic_pruning_keeps_null_row_group() {
866        use datafusion_physical_expr::expressions::{
867            Column as PhysicalColumn, lit as physical_lit,
868        };
869
870        let dir = create_temp_dir("dynamic_pruning_nulls");
871        let path = dir.path().join("nullable.parquet");
872        let arrow_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
873        let file = std::fs::File::create(&path).unwrap();
874        let mut writer = ArrowWriter::try_new(file, arrow_schema.clone(), None).unwrap();
875        for values in [[None, Some(1)], [Some(1), Some(1)], [None, None]] {
876            let batch = RecordBatch::try_new(
877                arrow_schema.clone(),
878                vec![Arc::new(Int32Array::from(values.to_vec()))],
879            )
880            .unwrap();
881            writer.write(&batch).unwrap();
882            writer.flush().unwrap();
883        }
884        writer.close().unwrap();
885        let builder =
886            ParquetRecordBatchStreamBuilder::new(tokio::fs::File::open(path).await.unwrap())
887                .await
888                .unwrap();
889        let schema = Arc::new(datatypes::schema::Schema::try_from(arrow_schema.clone()).unwrap());
890        let stats = RowGroupPruningStatistics::new(builder.metadata().row_groups(), &schema);
891        let filter = Arc::new(DynamicFilterPhysicalExpr::new(
892            vec![Arc::new(PhysicalColumn::new("a", 0))],
893            physical_lit(true),
894        ));
895        let predicate = Predicate::with_dyn_filters(vec![], vec![filter.clone()]);
896        assert_eq!(
897            predicate.prune_with_stats(&stats, &arrow_schema),
898            vec![true; 3]
899        );
900        filter
901            .update(
902                Predicate::to_physical_expr(
903                    &col("a").gt_eq(lit(10_i32)).and(col("a").lt_eq(lit(10_i32))),
904                    &arrow_schema,
905                )
906                .unwrap(),
907            )
908            .unwrap();
909        // NULL inputs must reach decoded filtering; non-NULL misses can still be pruned.
910        assert_eq!(
911            predicate.prune_with_stats(&stats, &arrow_schema),
912            vec![true, false, true],
913        );
914    }
915
916    #[tokio::test]
917    async fn test_clear_dyn_filters_restores_static_row_group_pruning() {
918        use datafusion_physical_expr::expressions::{
919            Column as PhysicalColumn, lit as physical_lit,
920        };
921
922        let dir = create_temp_dir("dynamic_pruning_reset");
923        let (path, arrow_schema) = gen_test_parquet_file(&dir, 30).await;
924        let schema = Arc::new(datatypes::schema::Schema::try_from(arrow_schema.clone()).unwrap());
925        let builder =
926            ParquetRecordBatchStreamBuilder::new(tokio::fs::File::open(path).await.unwrap())
927                .await
928                .unwrap();
929        let stats = RowGroupPruningStatistics::new(builder.metadata().row_groups(), &schema);
930        let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(
931            vec![Arc::new(PhysicalColumn::new("cnt", 1))],
932            physical_lit(true),
933        ));
934        let predicate = Predicate::with_dyn_filters(
935            vec![col("cnt").gt_eq(lit(10_i32))],
936            vec![dynamic_filter.clone()],
937        );
938
939        dynamic_filter
940            .update(
941                Predicate::to_physical_expr(&col("cnt").gt(lit(100_i32)), &arrow_schema).unwrap(),
942            )
943            .unwrap();
944        assert_eq!(
945            predicate.prune_with_stats(&stats, &arrow_schema),
946            vec![false; 3]
947        );
948
949        // Reset after the prior stream is dropped: no dynamic filter remains, while the
950        // static predicate still excludes only the first row group.
951        predicate.clear_dyn_filters();
952        assert_eq!(
953            predicate.prune_with_stats(&stats, &arrow_schema),
954            vec![false, true, true]
955        );
956    }
957
958    fn gen_predicate(max_val: i32, op: Operator) -> Vec<Expr> {
959        vec![datafusion_expr::Expr::BinaryExpr(BinaryExpr {
960            left: Box::new(datafusion_expr::Expr::Column(Column::from_name("cnt"))),
961            op,
962            right: Box::new(max_val.lit()),
963        })]
964    }
965
966    #[tokio::test]
967    async fn test_prune_empty() {
968        assert_prune(3, vec![], vec![true]).await;
969    }
970
971    #[tokio::test]
972    async fn test_prune_all_match() {
973        let p = gen_predicate(3, Operator::Gt);
974        assert_prune(2, p, vec![false]).await;
975    }
976
977    #[tokio::test]
978    async fn test_prune_gt() {
979        let p = gen_predicate(29, Operator::Gt);
980        assert_prune(
981            100,
982            p,
983            vec![
984                false, false, false, true, true, true, true, true, true, true,
985            ],
986        )
987        .await;
988    }
989
990    #[tokio::test]
991    async fn test_prune_eq_expr() {
992        let p = gen_predicate(30, Operator::Eq);
993        assert_prune(40, p, vec![false, false, false, true]).await;
994    }
995
996    #[tokio::test]
997    async fn test_prune_neq_expr() {
998        let p = gen_predicate(30, Operator::NotEq);
999        assert_prune(40, p, vec![true, true, true, true]).await;
1000    }
1001
1002    #[tokio::test]
1003    async fn test_prune_gteq_expr() {
1004        let p = gen_predicate(29, Operator::GtEq);
1005        assert_prune(40, p, vec![false, false, true, true]).await;
1006    }
1007
1008    #[tokio::test]
1009    async fn test_prune_lt_expr() {
1010        let p = gen_predicate(30, Operator::Lt);
1011        assert_prune(40, p, vec![true, true, true, false]).await;
1012    }
1013
1014    #[tokio::test]
1015    async fn test_prune_lteq_expr() {
1016        let p = gen_predicate(30, Operator::LtEq);
1017        assert_prune(40, p, vec![true, true, true, true]).await;
1018    }
1019
1020    #[tokio::test]
1021    async fn test_prune_between_expr() {
1022        let p = gen_predicate(30, Operator::LtEq);
1023        assert_prune(40, p, vec![true, true, true, true]).await;
1024    }
1025
1026    #[tokio::test]
1027    async fn test_or() {
1028        // cnt > 30 or cnt < 20
1029        let e = datafusion_expr::Expr::Column(Column::from_name("cnt"))
1030            .gt(30.lit())
1031            .or(datafusion_expr::Expr::Column(Column::from_name("cnt")).lt(20.lit()));
1032        assert_prune(40, vec![e], vec![true, true, false, true]).await;
1033    }
1034
1035    #[tokio::test]
1036    async fn test_to_physical_expr() {
1037        let predicate = Predicate::new(vec![
1038            col("host").eq(lit("host_a")),
1039            col("ts").gt(lit(ScalarValue::TimestampMicrosecond(Some(123), None))),
1040        ]);
1041
1042        let schema = Arc::new(arrow::datatypes::Schema::new(vec![Field::new(
1043            "host",
1044            arrow::datatypes::DataType::Utf8,
1045            false,
1046        )]));
1047
1048        let predicates = predicate.to_physical_exprs(&schema).unwrap();
1049        assert!(!predicates.is_empty());
1050
1051        let physical_expr = Predicate::to_physical_expr(&col("host").eq(lit("host_a")), &schema);
1052        assert!(physical_expr.is_ok());
1053    }
1054}
table/predicate.rs

table/
predicate.rs
