feat(python): push down predicates to pyarrow datasets (#5780)

ritchie46 · web-flow · commit 0f36679b48b8 · 2022-12-12T11:11:28.000+01:00
diff --git a/polars/polars-lazy/polars-plan/src/dsl/expr.rs b/polars/polars-lazy/polars-plan/src/dsl/expr.rs
@@ -1,4 +1,4 @@
-use std::fmt::{Debug, Formatter};
+use std::fmt::{Debug, Display, Formatter};
 use std::hash::{Hash, Hasher};
 use std::ops::Deref;
 
@@ -432,6 +432,31 @@ pub enum Operator {
     Xor,
 }
 
+impl Display for Operator {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        use Operator::*;
+        let tkn = match self {
+            Eq => "==",
+            NotEq => "!=",
+            Lt => "<",
+            LtEq => "<=",
+            Gt => ">",
+            GtEq => ">=",
+            Plus => "+",
+            Minus => "-",
+            Multiply => "*",
+            Divide => "//",
+            TrueDivide => "/",
+            FloorDivide => "floor_div",
+            Modulus => "%",
+            And => "&",
+            Or => "|",
+            Xor => "^",
+        };
+        write!(f, "{}", tkn)
+    }
+}
+
 impl Operator {
     pub(crate) fn is_comparison(&self) -> bool {
         matches!(
diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/alp.rs b/polars/polars-lazy/polars-plan/src/logical_plan/alp.rs
@@ -30,6 +30,7 @@ pub enum ALogicalPlan {
     #[cfg(feature = "python")]
     PythonScan {
         options: PythonOptions,
+        predicate: Option<Node>,
     },
     Melt {
         input: Node,
@@ -164,7 +165,7 @@ impl ALogicalPlan {
         use ALogicalPlan::*;
         match self {
             #[cfg(feature = "python")]
-            PythonScan { options } => &options.schema,
+            PythonScan { options, .. } => &options.schema,
             #[cfg(feature = "csv-file")]
             CsvScan { file_info, .. } => &file_info.schema,
             #[cfg(feature = "parquet")]
@@ -181,7 +182,7 @@ impl ALogicalPlan {
         use ALogicalPlan::*;
         let schema = match self {
             #[cfg(feature = "python")]
-            PythonScan { options } => &options.schema,
+            PythonScan { options, .. } => &options.schema,
             Union { inputs, .. } => return arena.get(inputs[0]).schema(arena),
             Cache { input, .. } => return arena.get(*input).schema(arena),
             Sort { input, .. } => return arena.get(*input).schema(arena),
@@ -249,8 +250,9 @@ impl ALogicalPlan {
 
         match self {
             #[cfg(feature = "python")]
-            PythonScan { options } => PythonScan {
+            PythonScan { options, predicate } => PythonScan {
                 options: options.clone(),
+                predicate: *predicate,
             },
             Union { options, .. } => Union {
                 inputs,
diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/conversion.rs b/polars/polars-lazy/polars-plan/src/logical_plan/conversion.rs
@@ -175,7 +175,10 @@ pub fn to_alp(
             options,
         },
         #[cfg(feature = "python")]
-        LogicalPlan::PythonScan { options } => ALogicalPlan::PythonScan { options },
+        LogicalPlan::PythonScan { options } => ALogicalPlan::PythonScan {
+            options,
+            predicate: None,
+        },
         LogicalPlan::Union { inputs, options } => {
             let inputs = inputs
                 .into_iter()
@@ -668,7 +671,7 @@ impl ALogicalPlan {
                 options,
             },
             #[cfg(feature = "python")]
-            ALogicalPlan::PythonScan { options } => LogicalPlan::PythonScan { options },
+            ALogicalPlan::PythonScan { options, .. } => LogicalPlan::PythonScan { options },
             ALogicalPlan::Union { inputs, options } => {
                 let inputs = inputs
                     .into_iter()
diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/mod.rs b/polars/polars-lazy/polars-plan/src/logical_plan/mod.rs
@@ -24,6 +24,8 @@ mod lit;
 pub(crate) mod optimizer;
 pub(crate) mod options;
 mod projection;
+#[cfg(feature = "python")]
+mod pyarrow;
 mod schema;
 
 pub use aexpr::*;
diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/optimizer/cse.rs b/polars/polars-lazy/polars-plan/src/logical_plan/optimizer/cse.rs
@@ -230,7 +230,7 @@ fn lp_node_equal(a: &ALogicalPlan, b: &ALogicalPlan, expr_arena: &Arena<AExpr>)
                 && expr_nodes_equal(agg_l, agg_r, expr_arena)
         }
         #[cfg(feature = "python")]
-        (PythonScan { options: l }, PythonScan { options: r, .. }) => l == r,
+        (PythonScan { options: l, .. }, PythonScan { options: r, .. }) => l == r,
         _ => {
             // joins and unions are also false
             // they do not originate from a single trail
diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/optimizer/predicate_pushdown/mod.rs b/polars/polars-lazy/polars-plan/src/logical_plan/optimizer/predicate_pushdown/mod.rs
@@ -544,9 +544,28 @@ impl PredicatePushDown {
                 self.no_pushdown_restart_opt(lp, acc_predicates, lp_arena, expr_arena)
             }
             #[cfg(feature = "python")]
-            // python node does not yet support predicates
-             lp @ PythonScan {..} => {
-                self.no_pushdown_restart_opt(lp, acc_predicates, lp_arena, expr_arena)
+             PythonScan {mut options, predicate} => {
+                if options.pyarrow {
+                    let predicate = predicate_at_scan(acc_predicates, predicate, expr_arena);
+
+                    if let Some(predicate) = predicate {
+                        match super::super::pyarrow::predicate_to_pa(predicate, expr_arena) {
+                            // we we able to create a pyarrow string, mutate the options
+                            Some(eval_str) => {
+                                options.predicate = Some(eval_str)
+                            },
+                            // we were not able to translate the predicate
+                            // apply here
+                            None => {
+                                let lp = PythonScan { options, predicate: None };
+                                return Ok(self.optional_apply_predicate(lp, vec![predicate], lp_arena, expr_arena))
+                            }
+                        }
+                    }
+                    Ok(PythonScan {options, predicate})
+                } else {
+                    self.no_pushdown_restart_opt(PythonScan {options, predicate}, acc_predicates, lp_arena, expr_arena)
+                }
             }
 
         }
diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/optimizer/projection_pushdown/mod.rs b/polars/polars-lazy/polars-plan/src/logical_plan/optimizer/projection_pushdown/mod.rs
@@ -455,7 +455,10 @@ impl ProjectionPushDown {
                 Ok(lp)
             }
             #[cfg(feature = "python")]
-            PythonScan { mut options } => {
+            PythonScan {
+                mut options,
+                predicate,
+            } => {
                 options.with_columns = get_scan_columns(&mut acc_projections, expr_arena);
 
                 options.output_schema = if options.with_columns.is_none() {
@@ -468,7 +471,7 @@ impl ProjectionPushDown {
                         true,
                     )?))
                 };
-                Ok(PythonScan { options })
+                Ok(PythonScan { options, predicate })
             }
             #[cfg(feature = "csv-file")]
             CsvScan {
diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/options.rs b/polars/polars-lazy/polars-plan/src/logical_plan/options.rs
@@ -219,6 +219,10 @@ pub struct PythonOptions {
     pub schema: SchemaRef,
     pub output_schema: Option<SchemaRef>,
     pub with_columns: Option<Arc<Vec<String>>>,
+    pub pyarrow: bool,
+    // a pyarrow predicate python expression
+    // can be evaluated with python.eval
+    pub predicate: Option<String>,
 }
 
 #[derive(Clone, PartialEq, Eq, Debug, Default)]
diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/pyarrow.rs b/polars/polars-lazy/polars-plan/src/logical_plan/pyarrow.rs
@@ -0,0 +1,59 @@
+use crate::prelude::*;
+
+// convert to a pyarrow expression that can be evaluated with pythons eval
+pub(super) fn predicate_to_pa(predicate: Node, expr_arena: &Arena<AExpr>) -> Option<String> {
+    match expr_arena.get(predicate) {
+        AExpr::BinaryExpr { left, right, op } => {
+            if op.is_comparison() {
+                let left = predicate_to_pa(*left, expr_arena)?;
+                let right = predicate_to_pa(*right, expr_arena)?;
+                Some(format!("({} {} {})", left, op, right))
+            } else {
+                None
+            }
+        }
+        AExpr::Column(name) => Some(format!("pa.dataset.field('{}')", name.as_ref())),
+        AExpr::Alias(input, _) => predicate_to_pa(*input, expr_arena),
+        AExpr::Literal(lv) => {
+            let av = lv.to_anyvalue()?;
+            let dtype = av.dtype();
+            if dtype.is_float() {
+                let val = av.extract::<f64>()?;
+                Some(format!("{}", val))
+            } else if dtype.is_integer() {
+                let val = av.extract::<i64>()?;
+                Some(format!("{}", val))
+            } else {
+                None
+            }
+        }
+        AExpr::Function {
+            function: FunctionExpr::Not,
+            input,
+            ..
+        } => {
+            let input = input.first().unwrap();
+            let input = predicate_to_pa(*input, expr_arena)?;
+            Some(format!("~({})", input))
+        }
+        AExpr::Function {
+            function: FunctionExpr::IsNull,
+            input,
+            ..
+        } => {
+            let input = input.first().unwrap();
+            let input = predicate_to_pa(*input, expr_arena)?;
+            Some(format!("({}).is_null()", input))
+        }
+        AExpr::Function {
+            function: FunctionExpr::IsNotNull,
+            input,
+            ..
+        } => {
+            let input = input.first().unwrap();
+            let input = predicate_to_pa(*input, expr_arena)?;
+            Some(format!("~({}).is_null()", input))
+        }
+        _ => None,
+    }
+}
diff --git a/polars/polars-lazy/src/frame/python.rs b/polars/polars-lazy/src/frame/python.rs
@@ -3,11 +3,12 @@ use polars_core::prelude::*;
 use crate::prelude::*;
 
 impl LazyFrame {
-    pub fn scan_from_python_function(schema: Schema, scan_fn: Vec<u8>) -> Self {
+    pub fn scan_from_python_function(schema: Schema, scan_fn: Vec<u8>, pyarrow: bool) -> Self {
         LogicalPlan::PythonScan {
             options: PythonOptions {
                 scan_fn,
                 schema: Arc::new(schema),
+                pyarrow,
                 ..Default::default()
             },
         }
diff --git a/polars/polars-lazy/src/physical_plan/executors/python_scan.rs b/polars/polars-lazy/src/physical_plan/executors/python_scan.rs
@@ -16,6 +16,7 @@ impl Executor for PythonScanExec {
             }
         }
         let with_columns = self.options.with_columns.take();
+        let pyarrow_predicate = self.options.predicate.take();
         Python::with_gil(|py| {
             let pl = PyModule::import(py, "polars").unwrap();
             let pli = pl.getattr("internals").unwrap();
@@ -27,7 +28,7 @@ impl Executor for PythonScanExec {
                 with_columns.map(|mut cols| std::mem::take(Arc::make_mut(&mut cols)));
 
             let out = deser_and_exec
-                .call1((bytes, with_columns))
+                .call1((bytes, with_columns, pyarrow_predicate))
                 .map_err(|err| PolarsError::ComputeError(format!("{:?}", err).into()))?;
             let pydf = out.getattr("_df").unwrap();
             let raw_parts = pydf.call_method0("into_raw_parts").unwrap();
diff --git a/polars/polars-lazy/src/physical_plan/planner/lp.rs b/polars/polars-lazy/src/physical_plan/planner/lp.rs
@@ -149,7 +149,7 @@ pub fn create_physical_plan(
     let logical_plan = lp_arena.take(root);
     match logical_plan {
         #[cfg(feature = "python")]
-        PythonScan { options } => Ok(Box::new(executors::PythonScanExec { options })),
+        PythonScan { options, .. } => Ok(Box::new(executors::PythonScanExec { options })),
         Union { inputs, options } => {
             let inputs = inputs
                 .into_iter()
diff --git a/py-polars/Cargo.lock b/py-polars/Cargo.lock
diff --git a/py-polars/polars/internals/anonymous_scan.py b/py-polars/polars/internals/anonymous_scan.py
diff --git a/py-polars/polars/internals/lazyframe/frame.py b/py-polars/polars/internals/lazyframe/frame.py
diff --git a/py-polars/polars/io.py b/py-polars/polars/io.py
diff --git a/py-polars/src/lazy/dataframe.rs b/py-polars/src/lazy/dataframe.rs
diff --git a/py-polars/tests/unit/io/test_pyarrow_dataset.py b/py-polars/tests/unit/io/test_pyarrow_dataset.py

Original file line number	Diff line number	Diff line change
`@@ -230,7 +230,7 @@ fn lp_node_equal(a: &ALogicalPlan, b: &ALogicalPlan, expr_arena: &Arena<AExpr>)`
`230`	`230`	`&& expr_nodes_equal(agg_l, agg_r, expr_arena)`
`231`	`231`	`}`
`232`	`232`	`#[cfg(feature = "python")]`
`233`		`- (PythonScan { options: l }, PythonScan { options: r, .. }) => l == r,`
	`233`	`+ (PythonScan { options: l, .. }, PythonScan { options: r, .. }) => l == r,`
`234`	`234`	`_ => {`
`235`	`235`	`// joins and unions are also false`
`236`	`236`	`// they do not originate from a single trail`