polars_io/
predicates.rs

1use std::fmt;
2
3use arrow::array::Array;
4use arrow::bitmap::{Bitmap, BitmapBuilder};
5use arrow::datatypes::ArrowDataType;
6use polars_core::prelude::*;
7#[cfg(feature = "parquet")]
8use polars_parquet::read::expr::{ParquetColumnExpr, ParquetScalar, SpecializedParquetColumnExpr};
9use polars_utils::format_pl_smallstr;
10#[cfg(feature = "serde")]
11use serde::{Deserialize, Serialize};
12
13pub trait PhysicalIoExpr: Send + Sync {
14    /// Take a [`DataFrame`] and produces a boolean [`Series`] that serves
15    /// as a predicate mask
16    fn evaluate_io(&self, df: &DataFrame) -> PolarsResult<Series>;
17}
18
19#[derive(Debug, Clone)]
20pub enum SpecializedColumnPredicate {
21    Equal(Scalar),
22    /// A closed (inclusive) range.
23    Between(Scalar, Scalar),
24    EqualOneOf(Box<[Scalar]>),
25    StartsWith(Box<[u8]>),
26    EndsWith(Box<[u8]>),
27    RegexMatch(regex::bytes::Regex),
28}
29
30#[derive(Clone)]
31pub struct ColumnPredicateExpr {
32    column_name: PlSmallStr,
33    dtype: DataType,
34    source_arrow_dtype: ArrowDataType,
35    #[cfg(feature = "parquet")]
36    specialized: Option<SpecializedParquetColumnExpr>,
37    expr: Arc<dyn PhysicalIoExpr>,
38}
39
40impl ColumnPredicateExpr {
41    pub fn new(
42        column_name: PlSmallStr,
43        dtype: DataType,
44        source_arrow_dtype: ArrowDataType,
45        expr: Arc<dyn PhysicalIoExpr>,
46        specialized: Option<SpecializedColumnPredicate>,
47    ) -> Self {
48        use SpecializedColumnPredicate as S;
49        #[cfg(feature = "parquet")]
50        use SpecializedParquetColumnExpr as P;
51        #[cfg(feature = "parquet")]
52        let specialized = specialized.and_then(|s| {
53            Some(match s {
54                S::Equal(s) => P::Equal(cast_to_parquet_scalar(s)?),
55                S::Between(low, high) => {
56                    P::Between(cast_to_parquet_scalar(low)?, cast_to_parquet_scalar(high)?)
57                },
58                S::EqualOneOf(scalars) => P::EqualOneOf(
59                    scalars
60                        .into_iter()
61                        .map(|s| cast_to_parquet_scalar(s).ok_or(()))
62                        .collect::<Result<Box<_>, ()>>()
63                        .ok()?,
64                ),
65                S::StartsWith(s) => P::StartsWith(s),
66                S::EndsWith(s) => P::EndsWith(s),
67                S::RegexMatch(s) => P::RegexMatch(s),
68            })
69        });
70
71        Self {
72            column_name,
73            dtype,
74            source_arrow_dtype,
75            #[cfg(feature = "parquet")]
76            specialized,
77            expr,
78        }
79    }
80}
81
82#[cfg(feature = "parquet")]
83impl ParquetColumnExpr for ColumnPredicateExpr {
84    fn evaluate_mut(&self, values: &dyn Array, bm: &mut BitmapBuilder) {
85        // We should never evaluate nulls with this.
86        assert!(values.validity().is_none_or(|v| v.set_bits() == 0));
87
88        // @TODO: Probably these unwraps should be removed.
89        let series = predicate_values_to_series(
90            self.column_name.clone(),
91            values,
92            &self.dtype,
93            &self.source_arrow_dtype,
94        )
95        .unwrap();
96        let column = series.into_column();
97        let df = unsafe { DataFrame::new_unchecked(values.len(), vec![column]) };
98
99        // @TODO: Probably these unwraps should be removed.
100        let true_mask = self.expr.evaluate_io(&df).unwrap();
101        let true_mask = true_mask.bool().unwrap();
102
103        bm.reserve(true_mask.len());
104        for chunk in true_mask.downcast_iter() {
105            match chunk.validity() {
106                None => bm.extend_from_bitmap(chunk.values()),
107                Some(v) => bm.extend_from_bitmap(&(chunk.values() & v)),
108            }
109        }
110    }
111    fn evaluate_null(&self) -> bool {
112        let column = Column::full_null(self.column_name.clone(), 1, &self.dtype);
113        let df = unsafe { DataFrame::new_unchecked(1, vec![column]) };
114
115        // @TODO: Probably these unwraps should be removed.
116        let true_mask = self.expr.evaluate_io(&df).unwrap();
117        let true_mask = true_mask.bool().unwrap();
118
119        true_mask.get(0).unwrap_or(false)
120    }
121
122    fn as_specialized(&self) -> Option<&SpecializedParquetColumnExpr> {
123        self.specialized.as_ref()
124    }
125}
126
127#[cfg(feature = "parquet")]
128fn predicate_values_to_series(
129    name: PlSmallStr,
130    values: &dyn Array,
131    dtype: &DataType,
132    source_arrow_dtype: &ArrowDataType,
133) -> PolarsResult<Series> {
134    // For example, Arrow seconds will be stored as Polars milliseconds, so we can not just
135    // zero-copy construct the predicate series
136    let timestamp_units_differ = matches!(
137        (source_arrow_dtype, dtype),
138        (
139            ArrowDataType::Timestamp(source_unit, _),
140            DataType::Datetime(target_unit, _),
141        ) if source_unit != &target_unit.to_arrow()
142    );
143
144    if timestamp_units_differ {
145        let values = polars_compute::cast::cast(
146            values,
147            source_arrow_dtype,
148            polars_compute::cast::CastOptionsImpl::default(),
149        )?;
150        Series::try_from((name, values))
151    } else {
152        Series::from_chunk_and_dtype(name, values.to_boxed(), dtype)
153    }
154}
155
156#[cfg(feature = "parquet")]
157fn cast_to_parquet_scalar(scalar: Scalar) -> Option<ParquetScalar> {
158    use AnyValue as A;
159    use ParquetScalar as P;
160
161    Some(match scalar.into_value() {
162        A::Null => P::Null,
163        A::Boolean(v) => P::Boolean(v),
164
165        A::UInt8(v) => P::UInt8(v),
166        A::UInt16(v) => P::UInt16(v),
167        A::UInt32(v) => P::UInt32(v),
168        A::UInt64(v) => P::UInt64(v),
169
170        A::Int8(v) => P::Int8(v),
171        A::Int16(v) => P::Int16(v),
172        A::Int32(v) => P::Int32(v),
173        A::Int64(v) => P::Int64(v),
174
175        #[cfg(feature = "dtype-time")]
176        A::Date(v) => P::Int32(v),
177        #[cfg(feature = "dtype-datetime")]
178        A::Datetime(v, _, _) | A::DatetimeOwned(v, _, _) => P::Int64(v),
179        #[cfg(feature = "dtype-duration")]
180        A::Duration(v, _) => P::Int64(v),
181        #[cfg(feature = "dtype-time")]
182        A::Time(v) => P::Int64(v),
183
184        A::Float32(v) => P::Float32(v),
185        A::Float64(v) => P::Float64(v),
186
187        // @TODO: Cast to string
188        #[cfg(feature = "dtype-categorical")]
189        A::Categorical(_, _) | A::CategoricalOwned(_, _) | A::Enum(_, _) | A::EnumOwned(_, _) => {
190            return None;
191        },
192
193        A::String(v) => P::String(v.into()),
194        A::StringOwned(v) => P::String(v.as_str().into()),
195        A::Binary(v) => P::Binary(v.into()),
196        A::BinaryOwned(v) => P::Binary(v.into()),
197        _ => return None,
198    })
199}
200
201#[cfg(any(feature = "parquet", feature = "ipc"))]
202pub fn apply_predicate(
203    df: &mut DataFrame,
204    predicate: Option<&dyn PhysicalIoExpr>,
205    parallel: bool,
206) -> PolarsResult<()> {
207    if let (Some(predicate), false) = (&predicate, df.columns().is_empty()) {
208        let s = predicate.evaluate_io(df)?;
209        let mask = s.bool().expect("filter predicates was not of type boolean");
210
211        if parallel {
212            *df = df.filter(mask)?;
213        } else {
214            *df = df.filter_seq(mask)?;
215        }
216    }
217    Ok(())
218}
219
220/// Statistics of the values in a column.
221///
222/// The following statistics are tracked for each row group:
223/// - Null count
224/// - Minimum value
225/// - Maximum value
226#[derive(Debug, Clone)]
227#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
228pub struct ColumnStats {
229    field: Field,
230    // Each Series contains the stats for each row group.
231    null_count: Option<Series>,
232    min_value: Option<Series>,
233    max_value: Option<Series>,
234}
235
236impl ColumnStats {
237    /// Constructs a new [`ColumnStats`].
238    pub fn new(
239        field: Field,
240        null_count: Option<Series>,
241        min_value: Option<Series>,
242        max_value: Option<Series>,
243    ) -> Self {
244        Self {
245            field,
246            null_count,
247            min_value,
248            max_value,
249        }
250    }
251
252    /// Constructs a new [`ColumnStats`] with only the [`Field`] information and no statistics.
253    pub fn from_field(field: Field) -> Self {
254        Self {
255            field,
256            null_count: None,
257            min_value: None,
258            max_value: None,
259        }
260    }
261
262    /// Constructs a new [`ColumnStats`] from a single-value Series.
263    pub fn from_column_literal(s: Series) -> Self {
264        debug_assert_eq!(s.len(), 1);
265        Self {
266            field: s.field().into_owned(),
267            null_count: None,
268            min_value: Some(s.clone()),
269            max_value: Some(s),
270        }
271    }
272
273    pub fn field_name(&self) -> &PlSmallStr {
274        self.field.name()
275    }
276
277    /// Returns the [`DataType`] of the column.
278    pub fn dtype(&self) -> &DataType {
279        self.field.dtype()
280    }
281
282    /// Returns the null count of each row group of the column.
283    pub fn get_null_count_state(&self) -> Option<&Series> {
284        self.null_count.as_ref()
285    }
286
287    /// Returns the minimum value of each row group of the column.
288    pub fn get_min_state(&self) -> Option<&Series> {
289        self.min_value.as_ref()
290    }
291
292    /// Returns the maximum value of each row group of the column.
293    pub fn get_max_state(&self) -> Option<&Series> {
294        self.max_value.as_ref()
295    }
296
297    /// Returns the null count of the column.
298    pub fn null_count(&self) -> Option<usize> {
299        match self.dtype() {
300            #[cfg(feature = "dtype-struct")]
301            DataType::Struct(_) => None,
302            _ => {
303                let s = self.get_null_count_state()?;
304                // if all null, there are no statistics.
305                if s.null_count() != s.len() {
306                    s.sum().ok()
307                } else {
308                    None
309                }
310            },
311        }
312    }
313
314    /// Returns the minimum and maximum values of the column as a single [`Series`].
315    pub fn to_min_max(&self) -> Option<Series> {
316        let min_val = self.get_min_state()?;
317        let max_val = self.get_max_state()?;
318        let dtype = self.dtype();
319
320        if !use_min_max(dtype) {
321            return None;
322        }
323
324        let mut min_max_values = min_val.clone();
325        min_max_values.append(max_val).unwrap();
326        if min_max_values.null_count() > 0 {
327            None
328        } else {
329            Some(min_max_values)
330        }
331    }
332
333    /// Returns the minimum value of the column as a single-value [`Series`].
334    ///
335    /// Returns `None` if no maximum value is available.
336    pub fn to_min(&self) -> Option<&Series> {
337        // @scalar-opt
338        let min_val = self.min_value.as_ref()?;
339        let dtype = min_val.dtype();
340
341        if !use_min_max(dtype) || min_val.len() != 1 {
342            return None;
343        }
344
345        if min_val.null_count() > 0 {
346            None
347        } else {
348            Some(min_val)
349        }
350    }
351
352    /// Returns the maximum value of the column as a single-value [`Series`].
353    ///
354    /// Returns `None` if no maximum value is available.
355    pub fn to_max(&self) -> Option<&Series> {
356        // @scalar-opt
357        let max_val = self.max_value.as_ref()?;
358        let dtype = max_val.dtype();
359
360        if !use_min_max(dtype) || max_val.len() != 1 {
361            return None;
362        }
363
364        if max_val.null_count() > 0 {
365            None
366        } else {
367            Some(max_val)
368        }
369    }
370}
371
372/// Returns whether the [`DataType`] supports minimum/maximum operations.
373fn use_min_max(dtype: &DataType) -> bool {
374    dtype.is_primitive_numeric()
375        || dtype.is_temporal()
376        || matches!(
377            dtype,
378            DataType::String | DataType::Binary | DataType::Boolean
379        )
380}
381
382pub struct ColumnStatistics {
383    pub dtype: DataType,
384    pub min: AnyValue<'static>,
385    pub max: AnyValue<'static>,
386    pub null_count: Option<IdxSize>,
387}
388
389pub trait SkipBatchPredicate: Send + Sync {
390    fn schema(&self) -> &SchemaRef;
391
392    fn can_skip_batch(
393        &self,
394        batch_size: IdxSize,
395        live_columns: &PlIndexSet<PlSmallStr>,
396        mut statistics: PlIndexMap<PlSmallStr, ColumnStatistics>,
397    ) -> PolarsResult<bool> {
398        let mut columns = Vec::with_capacity(1 + live_columns.len() * 3);
399
400        columns.push(Column::new_scalar(
401            PlSmallStr::from_static("len"),
402            Scalar::new(IDX_DTYPE, batch_size.into()),
403            1,
404        ));
405
406        for col in live_columns.iter() {
407            let dtype = self.schema().get(col).unwrap();
408            let (min, max, nc) = match statistics.swap_remove(col) {
409                None => (
410                    Scalar::null(dtype.clone()),
411                    Scalar::null(dtype.clone()),
412                    Scalar::null(IDX_DTYPE),
413                ),
414                Some(stat) => (
415                    Scalar::new(dtype.clone(), stat.min),
416                    Scalar::new(dtype.clone(), stat.max),
417                    Scalar::new(
418                        IDX_DTYPE,
419                        stat.null_count.map_or(AnyValue::Null, |nc| nc.into()),
420                    ),
421                ),
422            };
423            columns.extend([
424                Column::new_scalar(format_pl_smallstr!("{col}_min"), min, 1),
425                Column::new_scalar(format_pl_smallstr!("{col}_max"), max, 1),
426                Column::new_scalar(format_pl_smallstr!("{col}_nc"), nc, 1),
427            ]);
428        }
429
430        // SAFETY:
431        // * Each column is length = 1
432        // * We have an IndexSet, so each column name is unique
433        let df = unsafe { DataFrame::new_unchecked(1, columns) };
434        Ok(self.evaluate_with_stat_df(&df)?.get_bit(0))
435    }
436    fn evaluate_with_stat_df(&self, df: &DataFrame) -> PolarsResult<Bitmap>;
437}
438
439#[derive(Clone)]
440pub struct ColumnPredicates {
441    pub predicates:
442        PlHashMap<PlSmallStr, (Arc<dyn PhysicalIoExpr>, Option<SpecializedColumnPredicate>)>,
443    pub is_sumwise_complete: bool,
444}
445
446// I want to be explicit here.
447#[allow(clippy::derivable_impls)]
448impl Default for ColumnPredicates {
449    fn default() -> Self {
450        Self {
451            predicates: PlHashMap::default(),
452            is_sumwise_complete: false,
453        }
454    }
455}
456
457pub struct PhysicalExprWithConstCols<T> {
458    constants: Vec<(PlSmallStr, Scalar)>,
459    child: T,
460}
461
462impl SkipBatchPredicate for PhysicalExprWithConstCols<Arc<dyn SkipBatchPredicate>> {
463    fn schema(&self) -> &SchemaRef {
464        self.child.schema()
465    }
466
467    fn evaluate_with_stat_df(&self, df: &DataFrame) -> PolarsResult<Bitmap> {
468        let mut df = df.clone();
469        for (name, scalar) in self.constants.iter() {
470            df.with_column(Column::new_scalar(
471                name.clone(),
472                scalar.clone(),
473                df.height(),
474            ))?;
475        }
476        self.child.evaluate_with_stat_df(&df)
477    }
478}
479
480impl PhysicalIoExpr for PhysicalExprWithConstCols<Arc<dyn PhysicalIoExpr>> {
481    fn evaluate_io(&self, df: &DataFrame) -> PolarsResult<Series> {
482        let mut df = df.clone();
483        for (name, scalar) in self.constants.iter() {
484            df.with_column(Column::new_scalar(
485                name.clone(),
486                scalar.clone(),
487                df.height(),
488            ))?;
489        }
490
491        self.child.evaluate_io(&df)
492    }
493}
494
495#[derive(Clone)]
496pub struct ScanIOPredicate {
497    pub predicate: Arc<dyn PhysicalIoExpr>,
498
499    /// Column names that are used in the predicate.
500    pub live_columns: Arc<PlIndexSet<PlSmallStr>>,
501
502    /// A predicate that gets given statistics and evaluates whether a batch can be skipped.
503    pub skip_batch_predicate: Option<Arc<dyn SkipBatchPredicate>>,
504
505    /// A predicate that gets given statistics and evaluates whether a batch can be skipped.
506    pub column_predicates: Arc<ColumnPredicates>,
507
508    /// Predicate parts only referring to hive columns.
509    pub hive_predicate: Option<Arc<dyn PhysicalIoExpr>>,
510
511    pub hive_predicate_is_full_predicate: bool,
512}
513
514impl ScanIOPredicate {
515    pub fn set_external_constant_columns(&mut self, constant_columns: Vec<(PlSmallStr, Scalar)>) {
516        if constant_columns.is_empty() {
517            return;
518        }
519
520        let mut live_columns = self.live_columns.as_ref().clone();
521        for (c, _) in constant_columns.iter() {
522            live_columns.swap_remove(c);
523        }
524        self.live_columns = Arc::new(live_columns);
525
526        if let Some(skip_batch_predicate) = self.skip_batch_predicate.take() {
527            let mut sbp_constant_columns = Vec::with_capacity(constant_columns.len() * 3);
528            for (c, v) in constant_columns.iter() {
529                sbp_constant_columns.push((format_pl_smallstr!("{c}_min"), v.clone()));
530                sbp_constant_columns.push((format_pl_smallstr!("{c}_max"), v.clone()));
531                let nc = if v.is_null() {
532                    AnyValue::Null
533                } else {
534                    (0 as IdxSize).into()
535                };
536                sbp_constant_columns
537                    .push((format_pl_smallstr!("{c}_nc"), Scalar::new(IDX_DTYPE, nc)));
538            }
539            self.skip_batch_predicate = Some(Arc::new(PhysicalExprWithConstCols {
540                constants: sbp_constant_columns,
541                child: skip_batch_predicate,
542            }));
543        }
544
545        let mut column_predicates = self.column_predicates.as_ref().clone();
546        for (c, _) in constant_columns.iter() {
547            column_predicates.predicates.remove(c);
548        }
549        self.column_predicates = Arc::new(column_predicates);
550
551        self.predicate = Arc::new(PhysicalExprWithConstCols {
552            constants: constant_columns,
553            child: self.predicate.clone(),
554        });
555    }
556}
557
558impl fmt::Debug for ScanIOPredicate {
559    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
560        f.write_str("scan_io_predicate")
561    }
562}
polars_io/predicates.rs

polars_io/
predicates.rs
