polars_io/
predicates.rs

1use std::fmt;
2
3use polars_arrow::array::Array;
4use polars_arrow::bitmap::{Bitmap, BitmapBuilder};
5use polars_arrow::datatypes::ArrowDataType;
6use polars_core::chunked_array::cast::CastOptions;
7use polars_core::prelude::*;
8#[cfg(feature = "parquet")]
9use polars_parquet::read::expr::{ParquetColumnExpr, ParquetScalar, SpecializedParquetColumnExpr};
10use polars_utils::format_pl_smallstr;
11
12pub trait PhysicalIoExpr: Send + Sync {
13    /// Take a [`DataFrame`] and produces a boolean [`Series`] that serves
14    /// as a predicate mask
15    fn evaluate_io(&self, df: &DataFrame) -> PolarsResult<Series>;
16}
17
18#[derive(Debug, Clone)]
19pub enum SpecializedColumnPredicate {
20    Equal(Scalar),
21    /// A closed (inclusive) range.
22    Between(Scalar, Scalar),
23    EqualOneOf(Box<[Scalar]>),
24    StartsWith(Box<[u8]>),
25    EndsWith(Box<[u8]>),
26    RegexMatch(regex::bytes::Regex),
27}
28
29#[derive(Clone)]
30pub struct ColumnPredicateExpr {
31    column_name: PlSmallStr,
32    dtype: DataType,
33    source_arrow_dtype: ArrowDataType,
34    #[cfg(feature = "parquet")]
35    specialized: Option<SpecializedParquetColumnExpr>,
36    expr: Arc<dyn PhysicalIoExpr>,
37}
38
39impl ColumnPredicateExpr {
40    pub fn new(
41        column_name: PlSmallStr,
42        dtype: DataType,
43        source_arrow_dtype: ArrowDataType,
44        expr: Arc<dyn PhysicalIoExpr>,
45        specialized: Option<SpecializedColumnPredicate>,
46    ) -> Self {
47        use SpecializedColumnPredicate as S;
48        #[cfg(feature = "parquet")]
49        use SpecializedParquetColumnExpr as P;
50        // A specialized predicate compares its scalars with the values as the file
51        // stores them, which polars scales for some Arrow types.
52        #[cfg(feature = "parquet")]
53        let specialized = specialized.and_then(|s| {
54            if DataType::arrow_value_scale(&source_arrow_dtype) != 1 {
55                return None;
56            }
57            Some(match s {
58                S::Equal(s) => P::Equal(cast_to_parquet_scalar(s)?),
59                S::Between(low, high) => {
60                    P::Between(cast_to_parquet_scalar(low)?, cast_to_parquet_scalar(high)?)
61                },
62                S::EqualOneOf(scalars) => P::EqualOneOf(
63                    scalars
64                        .into_iter()
65                        .map(|s| cast_to_parquet_scalar(s).ok_or(()))
66                        .collect::<Result<Box<_>, ()>>()
67                        .ok()?,
68                ),
69                S::StartsWith(s) => P::StartsWith(s),
70                S::EndsWith(s) => P::EndsWith(s),
71                S::RegexMatch(s) => P::RegexMatch(s),
72            })
73        });
74
75        Self {
76            column_name,
77            dtype,
78            source_arrow_dtype,
79            #[cfg(feature = "parquet")]
80            specialized,
81            expr,
82        }
83    }
84}
85
86#[cfg(feature = "parquet")]
87impl ParquetColumnExpr for ColumnPredicateExpr {
88    fn evaluate_mut(&self, values: &dyn Array, bm: &mut BitmapBuilder) {
89        // We should never evaluate nulls with this.
90        assert!(values.validity().is_none_or(|v| v.set_bits() == 0));
91
92        // @TODO: Probably these unwraps should be removed.
93        let series = predicate_values_to_series(
94            self.column_name.clone(),
95            values,
96            &self.dtype,
97            &self.source_arrow_dtype,
98        )
99        .unwrap();
100        let column = series.into_column();
101        let df = unsafe { DataFrame::new_unchecked(values.len(), vec![column]) };
102
103        // @TODO: Probably these unwraps should be removed.
104        let true_mask = self.expr.evaluate_io(&df).unwrap();
105        let true_mask = true_mask.bool().unwrap();
106
107        bm.reserve(true_mask.len());
108        for chunk in true_mask.downcast_iter() {
109            match chunk.validity() {
110                None => bm.extend_from_bitmap(chunk.values()),
111                Some(v) => bm.extend_from_bitmap(&(chunk.values() & v)),
112            }
113        }
114    }
115    fn evaluate_null(&self) -> bool {
116        let column = Column::full_null(self.column_name.clone(), 1, &self.dtype);
117        let df = unsafe { DataFrame::new_unchecked(1, vec![column]) };
118
119        // @TODO: Probably these unwraps should be removed.
120        let true_mask = self.expr.evaluate_io(&df).unwrap();
121        let true_mask = true_mask.bool().unwrap();
122
123        true_mask.get(0).unwrap_or(false)
124    }
125
126    fn as_specialized(&self) -> Option<&SpecializedParquetColumnExpr> {
127        self.specialized.as_ref()
128    }
129}
130
131#[cfg(feature = "parquet")]
132fn predicate_values_to_series(
133    name: PlSmallStr,
134    values: &dyn Array,
135    dtype: &DataType,
136    source_arrow_dtype: &ArrowDataType,
137) -> PolarsResult<Series> {
138    // Polars stores the values of some Arrow types scaled, e.g. Arrow seconds as
139    // milliseconds, so the predicate series cannot be constructed zero-copy.
140    if DataType::arrow_value_scale(source_arrow_dtype) != 1 {
141        let values = polars_compute::cast::cast(
142            values,
143            source_arrow_dtype,
144            polars_compute::cast::CastOptionsImpl::default(),
145        )?;
146        Series::try_from((name, values))
147    } else {
148        Series::from_chunk_and_dtype(name, values.to_boxed(), dtype)
149    }
150}
151
152#[cfg(feature = "parquet")]
153fn cast_to_parquet_scalar(scalar: Scalar) -> Option<ParquetScalar> {
154    use AnyValue as A;
155    use ParquetScalar as P;
156
157    Some(match scalar.into_value() {
158        A::Null => P::Null,
159        A::Boolean(v) => P::Boolean(v),
160
161        A::UInt8(v) => P::UInt8(v),
162        A::UInt16(v) => P::UInt16(v),
163        A::UInt32(v) => P::UInt32(v),
164        A::UInt64(v) => P::UInt64(v),
165
166        A::Int8(v) => P::Int8(v),
167        A::Int16(v) => P::Int16(v),
168        A::Int32(v) => P::Int32(v),
169        A::Int64(v) => P::Int64(v),
170
171        #[cfg(feature = "dtype-date")]
172        A::Date(v) => P::Int32(v),
173        #[cfg(feature = "dtype-datetime")]
174        A::Datetime(v, _, _) | A::DatetimeOwned(v, _, _) => P::Int64(v),
175        #[cfg(feature = "dtype-duration")]
176        A::Duration(v, _) => P::Int64(v),
177        #[cfg(feature = "dtype-time")]
178        A::Time(v) => P::Int64(v),
179
180        A::Float32(v) => P::Float32(v),
181        A::Float64(v) => P::Float64(v),
182
183        // @TODO: Cast to string
184        #[cfg(feature = "dtype-categorical")]
185        A::Categorical(_, _) | A::CategoricalOwned(_, _) | A::Enum(_, _) | A::EnumOwned(_, _) => {
186            return None;
187        },
188
189        A::String(v) => P::String(v.into()),
190        A::StringOwned(v) => P::String(v.as_str().into()),
191        A::Binary(v) => P::Binary(v.into()),
192        A::BinaryOwned(v) => P::Binary(v.into()),
193        _ => return None,
194    })
195}
196
197#[cfg(any(feature = "parquet", feature = "ipc"))]
198pub fn apply_predicate(
199    df: &mut DataFrame,
200    predicate: Option<&dyn PhysicalIoExpr>,
201    parallel: bool,
202) -> PolarsResult<()> {
203    if let (Some(predicate), false) = (&predicate, df.columns().is_empty()) {
204        let s = predicate.evaluate_io(df)?;
205        let mask = s.bool().expect("filter predicates was not of type boolean");
206
207        if parallel {
208            *df = df.filter(mask)?;
209        } else {
210            *df = df.filter_seq(mask)?;
211        }
212    }
213    Ok(())
214}
215
216pub struct ColumnStatistics {
217    pub dtype: DataType,
218    pub min: AnyValue<'static>,
219    pub max: AnyValue<'static>,
220    pub null_count: Option<IdxSize>,
221}
222
223pub trait SkipBatchPredicate: Send + Sync {
224    fn schema(&self) -> &SchemaRef;
225
226    fn can_skip_batch(
227        &self,
228        batch_size: IdxSize,
229        live_columns: &PlIndexSet<PlSmallStr>,
230        mut statistics: PlIndexMap<PlSmallStr, ColumnStatistics>,
231    ) -> PolarsResult<bool> {
232        let mut columns = Vec::with_capacity(1 + live_columns.len() * 3);
233
234        columns.push(Column::new_scalar(
235            PlSmallStr::from_static("len"),
236            Scalar::new(IDX_DTYPE, batch_size.into()),
237            1,
238        ));
239
240        for col in live_columns.iter() {
241            let dtype = self.schema().get(col).unwrap();
242            let (min, max, nc) = match statistics.swap_remove(col) {
243                None => (
244                    Scalar::null(dtype.clone()),
245                    Scalar::null(dtype.clone()),
246                    Scalar::null(IDX_DTYPE),
247                ),
248                Some(stat) => (
249                    Scalar::new(dtype.clone(), stat.min),
250                    Scalar::new(dtype.clone(), stat.max),
251                    Scalar::new(
252                        IDX_DTYPE,
253                        stat.null_count.map_or(AnyValue::Null, |nc| nc.into()),
254                    ),
255                ),
256            };
257            columns.extend([
258                Column::new_scalar(format_pl_smallstr!("{col}_min"), min, 1),
259                Column::new_scalar(format_pl_smallstr!("{col}_max"), max, 1),
260                Column::new_scalar(format_pl_smallstr!("{col}_nc"), nc, 1),
261            ]);
262        }
263
264        // SAFETY:
265        // * Each column is length = 1
266        // * We have an IndexSet, so each column name is unique
267        let df = unsafe { DataFrame::new_unchecked(1, columns) };
268        Ok(self.evaluate_with_stat_df(&df)?.get_bit(0))
269    }
270    fn evaluate_with_stat_df(&self, df: &DataFrame) -> PolarsResult<Bitmap>;
271}
272
273/// The conjuncts of a row predicate that read one column, conjoined.
274#[derive(Clone)]
275pub struct ColumnPredicate {
276    /// The static conjuncts, conjoined. `None` when the column only has dynamic ones.
277    pub predicate: Option<Arc<dyn PhysicalIoExpr>>,
278    pub specialized: Option<SpecializedColumnPredicate>,
279    /// The conjuncts a producer sets at run time, each on its own.
280    pub dynamic: Vec<DynamicColumnPredicate>,
281}
282
283impl ColumnPredicate {
284    /// Every conjunct, static and dynamic, conjoined.
285    pub fn conjoined(&self) -> Arc<dyn PhysicalIoExpr> {
286        self.predicate
287            .iter()
288            .chain(self.dynamic.iter().map(|d| &d.predicate))
289            .cloned()
290            .reduce(|a, b| Arc::new(AndIoExpr(a, b)))
291            .unwrap()
292    }
293}
294
295/// A conjunct on one column that a producer sets at run time. It keeps every
296/// row until `source` says it filters rows.
297#[derive(Clone)]
298pub struct DynamicColumnPredicate {
299    pub predicate: Arc<dyn PhysicalIoExpr>,
300    pub source: Arc<dyn DynamicPredicateSource>,
301}
302
303/// `a AND b`.
304struct AndIoExpr(Arc<dyn PhysicalIoExpr>, Arc<dyn PhysicalIoExpr>);
305
306impl PhysicalIoExpr for AndIoExpr {
307    fn evaluate_io(&self, df: &DataFrame) -> PolarsResult<Series> {
308        let a = self.0.evaluate_io(df)?;
309        let b = self.1.evaluate_io(df)?;
310        Ok((a.bool()? & b.bool()?).into_series())
311    }
312}
313
314pub struct PhysicalExprWithConstCols<T> {
315    constants: Vec<(PlSmallStr, Scalar)>,
316    child: T,
317}
318
319impl SkipBatchPredicate for PhysicalExprWithConstCols<Arc<dyn SkipBatchPredicate>> {
320    fn schema(&self) -> &SchemaRef {
321        self.child.schema()
322    }
323
324    fn evaluate_with_stat_df(&self, df: &DataFrame) -> PolarsResult<Bitmap> {
325        let mut df = df.clone();
326        for (name, scalar) in self.constants.iter() {
327            df.with_column(Column::new_scalar(
328                name.clone(),
329                scalar.clone(),
330                df.height(),
331            ))?;
332        }
333        self.child.evaluate_with_stat_df(&df)
334    }
335}
336
337impl PhysicalIoExpr for PhysicalExprWithConstCols<Arc<dyn PhysicalIoExpr>> {
338    fn evaluate_io(&self, df: &DataFrame) -> PolarsResult<Series> {
339        let mut df = df.clone();
340        for (name, scalar) in self.constants.iter() {
341            df.with_column(Column::new_scalar(
342                name.clone(),
343                scalar.clone(),
344                df.height(),
345            ))?;
346        }
347
348        self.child.evaluate_io(&df)
349    }
350}
351
352/// The row predicate split into the conjuncts that read a single column and the rest.
353#[derive(Clone)]
354pub struct StagedScanIOPredicate {
355    pub column_predicates: Arc<PlIndexMap<PlSmallStr, ColumnPredicate>>,
356    /// The conjuncts that read no or several columns, conjoined.
357    pub rest: Option<Arc<dyn PhysicalIoExpr>>,
358}
359
360impl StagedScanIOPredicate {
361    /// Every conjunct on a constant column becomes part of `rest`.
362    fn with_constant_columns(&self, constants: &[(PlSmallStr, Scalar)]) -> Self {
363        let mut column_predicates = self.column_predicates.as_ref().clone();
364        let mut rest = self.rest.clone();
365        for (c, _) in constants {
366            if let Some(p) = column_predicates.shift_remove(c) {
367                let p = p.conjoined();
368                rest = Some(match rest {
369                    None => p,
370                    Some(rest) => Arc::new(AndIoExpr(rest, p)),
371                });
372            }
373        }
374        Self {
375            column_predicates: Arc::new(column_predicates),
376            rest: rest.map(|rest| {
377                Arc::new(PhysicalExprWithConstCols {
378                    constants: constants.to_vec(),
379                    child: rest,
380                }) as _
381            }),
382        }
383    }
384}
385
386/// What a producer has published for a column, read once per file: a reader
387/// skips the batches whose statistics fall outside the range.
388#[derive(Clone, Debug)]
389pub enum RuntimeRange {
390    /// Not published yet. Every batch is kept; a later file may see a range.
391    Pending,
392    /// Never published. Every batch is kept.
393    Disabled,
394    /// No value can match. Every batch is skipped.
395    Empty,
396    /// Only values in `lo..=hi` can match.
397    Range { lo: Scalar, hi: Scalar },
398}
399
400/// A reader's view of a predicate that a producer sets at run time.
401pub trait DynamicPredicateSource: Send + Sync {
402    fn runtime_range(&self) -> RuntimeRange;
403
404    /// Whether the producer has published a predicate that rejects rows.
405    fn filters_rows(&self) -> bool;
406
407    /// Whether a reader may stop evaluating the predicate when it rejects too
408    /// little: it stays as it is once set, and the producer checks every row
409    /// again.
410    fn can_bypass(&self) -> bool;
411}
412
413/// A column whose batches a reader may skip by a [`RuntimeRange`]. It is never
414/// evaluated per row.
415#[derive(Clone)]
416pub struct RuntimeRangeHint {
417    pub column: PlSmallStr,
418    pub source: Arc<dyn DynamicPredicateSource>,
419    /// The column's value in this file when it is not stored in the file, such as
420    /// a hive column or a missing column with a default.
421    pub constant: Option<Scalar>,
422}
423
424/// A range bound as `dtype`, or `None` when it does not survive the cast, in
425/// which case it bounds nothing.
426pub fn cast_bound(bound: &Scalar, dtype: &DataType) -> Option<Scalar> {
427    bound
428        .clone()
429        .cast_with_options(dtype, CastOptions::NonStrict)
430        .ok()
431        .filter(|b| !b.is_null())
432}
433
434impl RuntimeRangeHint {
435    /// Whether a file whose column is the constant `value` can hold a match.
436    /// `None` when the range does not settle it.
437    pub fn constant_matches(range: &RuntimeRange, value: &Scalar) -> Option<bool> {
438        match range {
439            RuntimeRange::Pending | RuntimeRange::Disabled => None,
440            RuntimeRange::Empty => Some(false),
441            RuntimeRange::Range { lo, hi } => {
442                if value.is_null() {
443                    return Some(false);
444                }
445                let lo = cast_bound(lo, value.dtype())?;
446                let hi = cast_bound(hi, value.dtype())?;
447                let value = value.value();
448                Some(value >= lo.value() && value <= hi.value())
449            },
450        }
451    }
452
453    /// Bind the hints of column `name` to its constant value in a file.
454    pub fn set_constant(hints: &mut [Self], name: &str, value: &Scalar) {
455        for hint in hints.iter_mut().filter(|h| h.column == name) {
456            hint.constant = Some(value.clone());
457        }
458    }
459}
460
461#[derive(Clone)]
462pub struct ScanIOPredicate {
463    pub predicate: Arc<dyn PhysicalIoExpr>,
464
465    /// `predicate` split for readers that filter while decoding.
466    pub staged: Option<StagedScanIOPredicate>,
467
468    /// Whether `predicate` filters rows at all. False when the scan only has
469    /// runtime ranges to skip batches by.
470    pub filters_rows: bool,
471
472    /// Column names that are used in the predicate.
473    pub live_columns: Arc<PlIndexSet<PlSmallStr>>,
474
475    /// A predicate that gets given statistics and evaluates whether a batch can be skipped.
476    pub skip_batch_predicate: Option<Arc<dyn SkipBatchPredicate>>,
477
478    /// Columns whose batches are skipped by a range published at run time.
479    pub runtime_ranges: Vec<RuntimeRangeHint>,
480
481    /// Predicate parts only referring to hive columns.
482    pub hive_predicate: Option<Arc<dyn PhysicalIoExpr>>,
483
484    pub hive_predicate_is_full_predicate: bool,
485}
486
487impl ScanIOPredicate {
488    /// Whether the predicate or a range hint reads the column.
489    pub fn reads_column(&self, name: &str) -> bool {
490        self.live_columns.contains(name) || self.runtime_ranges.iter().any(|h| h.column == name)
491    }
492
493    pub fn set_external_constant_columns(&mut self, constant_columns: Vec<(PlSmallStr, Scalar)>) {
494        if constant_columns.is_empty() {
495            return;
496        }
497
498        let mut live_columns = self.live_columns.as_ref().clone();
499        for (c, _) in constant_columns.iter() {
500            live_columns.swap_remove(c);
501        }
502        self.live_columns = Arc::new(live_columns);
503
504        for (name, value) in constant_columns.iter() {
505            RuntimeRangeHint::set_constant(&mut self.runtime_ranges, name, value);
506        }
507
508        if let Some(skip_batch_predicate) = self.skip_batch_predicate.take() {
509            let mut sbp_constant_columns = Vec::with_capacity(constant_columns.len() * 3);
510            for (c, v) in constant_columns.iter() {
511                sbp_constant_columns.push((format_pl_smallstr!("{c}_min"), v.clone()));
512                sbp_constant_columns.push((format_pl_smallstr!("{c}_max"), v.clone()));
513                let nc = if v.is_null() {
514                    AnyValue::Null
515                } else {
516                    (0 as IdxSize).into()
517                };
518                sbp_constant_columns
519                    .push((format_pl_smallstr!("{c}_nc"), Scalar::new(IDX_DTYPE, nc)));
520            }
521            self.skip_batch_predicate = Some(Arc::new(PhysicalExprWithConstCols {
522                constants: sbp_constant_columns,
523                child: skip_batch_predicate,
524            }));
525        }
526
527        if let Some(staged) = self.staged.as_mut() {
528            *staged = staged.with_constant_columns(&constant_columns);
529        }
530
531        self.predicate = Arc::new(PhysicalExprWithConstCols {
532            constants: constant_columns,
533            child: self.predicate.clone(),
534        });
535    }
536}
537
538impl fmt::Debug for ScanIOPredicate {
539    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
540        f.write_str("scan_io_predicate")
541    }
542}
polars_io/predicates.rs

polars_io/
predicates.rs
