feat(rust): add predicate pushdown to anonymous_scan (#5467)

pola-rs · Nov 11, 2022 · 0c94ae5 · 0c94ae5
1 parent e7403bc
commit 0c94ae5
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 4 deletions.
diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/builder.rs b/polars/polars-lazy/polars-plan/src/logical_plan/builder.rs
@@ -88,6 +88,7 @@ impl LogicalPlanBuilder {
                 n_rows,
                 output_schema: None,
                 with_columns: None,
+                predicate: None,
             },
         }
         .into())

diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/options.rs b/polars/polars-lazy/polars-plan/src/logical_plan/options.rs
@@ -6,6 +6,8 @@ use polars_time::{DynamicGroupOptions, RollingGroupOptions};
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
 
+use crate::prelude::Expr;
+
 pub type FileCount = u32;
 
 #[derive(Clone, Debug, PartialEq, Eq)]
@@ -226,5 +228,6 @@ pub struct AnonymousScanOptions {
     pub skip_rows: Option<usize>,
     pub n_rows: Option<usize>,
     pub with_columns: Option<Arc<Vec<String>>>,
+    pub predicate: Option<Expr>,
     pub fmt_str: &'static str,
 }
diff --git a/polars/polars-lazy/src/physical_plan/executors/scan/mod.rs b/polars/polars-lazy/src/physical_plan/executors/scan/mod.rs
@@ -109,16 +109,22 @@ pub(crate) struct AnonymousScanExec {
 impl Executor for AnonymousScanExec {
     fn execute(&mut self, state: &mut ExecutionState) -> PolarsResult<DataFrame> {
         state.record(
-            || {
-                let mut df = self.function.scan(self.options.clone())?;
-                if let Some(predicate) = &self.predicate {
+            || match (self.function.allows_predicate_pushdown(), &self.predicate) {
+                (true, Some(predicate)) => {
+                    self.options.predicate = predicate.as_expression().cloned();
+                    self.function.scan(self.options.clone())
+                }
+                (false, Some(predicate)) => {
+                    let mut df = self.function.scan(self.options.clone())?;
                     let s = predicate.evaluate(&df, state)?;
                     let mask = s.bool().map_err(|_| {
                         PolarsError::ComputeError("filter predicate was not of type boolean".into())
                     })?;
                     df = df.filter(mask)?;
+
+                    Ok(df)
                 }
-                Ok(df)
+                _ => self.function.scan(self.options.clone()),
             },
             "anonymous_scan".into(),
         )