slice pushdown optimization

polars/polars-lazy/src/frame/csv.rs

@@ -203,20 +203,29 @@ impl<'a> LazyCsvReader<'a> {
         if self.path.contains('*') {
             let paths = glob::glob(&self.path)
                 .map_err(|_| PolarsError::ValueError("invalid glob pattern given".into()))?;
+
             let lfs = paths
                 .map(|r| {
                     let path = r.map_err(|e| PolarsError::ComputeError(format!("{e}").into()))?;
                     let path_string = path.to_string_lossy().into_owned();
                     let mut builder = self.clone();
                     builder.path = path_string;
+                    builder.skip_rows = 0;
+                    builder.n_rows = None;
                     // do no rechunk yet.
                     builder.rechunk = false;
                     builder.finish_impl()
                 })
                 .collect::<Result<Vec<_>>>()?;
-
             concat(&lfs, self.rechunk)
                 .map_err(|_| PolarsError::ComputeError("no matching files found".into()))
+                .map(|lf| {
+                    if self.skip_rows != 0 || self.n_rows.is_some() {
+                        lf.slice(self.skip_rows as i64, self.n_rows.unwrap() as u32)
+                    } else {
+                        lf
+                    }
+                })
         } else {
             self.finish_impl()
         }

polars/polars-lazy/src/frame/ipc.rs

@@ -41,12 +41,21 @@ impl LazyFrame {
                 .map(|r| {
                     let path = r.map_err(|e| PolarsError::ComputeError(format!("{e}").into()))?;
                     let path_string = path.to_string_lossy().into_owned();
+                    let mut args = args;
+                    args.n_rows = None;
                     Self::scan_ipc_impl(path_string, args)
                 })
                 .collect::<Result<Vec<_>>>()?;
 
             concat(&lfs, args.rechunk)
                 .map_err(|_| PolarsError::ComputeError("no matching files found".into()))
+                .map(|lf| {
+                    if let Some(n_rows) = args.n_rows {
+                        lf.slice(0, n_rows as u32)
+                    } else {
+                        lf
+                    }
+                })
         } else {
             Self::scan_ipc_impl(path, args)
         }

polars/polars-lazy/src/frame/mod.rs

@@ -36,6 +36,7 @@ use crate::prelude::aggregate_scan_projections::agg_projection;
 use crate::prelude::drop_nulls::ReplaceDropNulls;
 use crate::prelude::fast_projection::FastProjection;
 use crate::prelude::simplify_expr::SimplifyBooleanRule;
+use crate::prelude::slice_pushdown::SlicePushDown;
 use crate::utils::{combine_predicates_expr, expr_to_root_column_names};
 use crate::{logical_plan::FETCH_ROWS, prelude::*};
 use polars_arrow::prelude::QuantileInterpolOptions;
@@ -104,7 +105,7 @@ pub struct OptState {
     pub agg_scan_projection: bool,
     pub aggregate_pushdown: bool,
     pub global_string_cache: bool,
-    pub join_pruning: bool,
+    pub slice_pushdown: bool,
 }
 
 impl Default for OptState {
@@ -115,7 +116,7 @@ impl Default for OptState {
             type_coercion: true,
             simplify_expr: true,
             global_string_cache: true,
-            join_pruning: true,
+            slice_pushdown: true,
             // will be toggled by a scan operation such as csv scan or parquet scan
             agg_scan_projection: false,
             aggregate_pushdown: false,
@@ -192,9 +193,9 @@ impl LazyFrame {
         self
     }
 
-    /// Toggle join pruning optimization
-    pub fn with_join_pruning(mut self, toggle: bool) -> Self {
-        self.opt_state.join_pruning = toggle;
+    /// Toggle slice pushdown optimization
+    pub fn with_slice_pushdown(mut self, toggle: bool) -> Self {
+        self.opt_state.slice_pushdown = toggle;
         self
     }
 
@@ -399,6 +400,7 @@ impl LazyFrame {
         let projection_pushdown = self.opt_state.projection_pushdown;
         let type_coercion = self.opt_state.type_coercion;
         let simplify_expr = self.opt_state.simplify_expr;
+        let slice_pushdown = self.opt_state.slice_pushdown;
 
         #[cfg(any(feature = "parquet", feature = "csv-file"))]
         let agg_scan_projection = self.opt_state.agg_scan_projection;
@@ -409,16 +411,14 @@ impl LazyFrame {
         // gradually fill the rules passed to the optimizer
         let mut rules: Vec<Box<dyn OptimizationRule>> = Vec::with_capacity(8);
 
-        let predicate_pushdown_opt = PredicatePushDown::default();
-        let projection_pushdown_opt = ProjectionPushDown {};
-
         // during debug we check if the optimizations have not modified the final schema
         #[cfg(debug_assertions)]
         let prev_schema = logical_plan.schema().clone();
 
         let mut lp_top = to_alp(logical_plan, expr_arena, lp_arena);
 
         if projection_pushdown {
+            let projection_pushdown_opt = ProjectionPushDown {};
             let alp = lp_arena.take(lp_top);
             let alp = projection_pushdown_opt
                 .optimize(alp, lp_arena, expr_arena)
@@ -427,13 +427,24 @@ impl LazyFrame {
         }
 
         if predicate_pushdown {
+            let predicate_pushdown_opt = PredicatePushDown::default();
             let alp = lp_arena.take(lp_top);
             let alp = predicate_pushdown_opt
                 .optimize(alp, lp_arena, expr_arena)
                 .expect("predicate pushdown failed");
             lp_arena.replace(lp_top, alp);
         }
 
+        if slice_pushdown {
+            let slice_pushdown_opt = SlicePushDown {};
+            let alp = lp_arena.take(lp_top);
+            let alp = slice_pushdown_opt
+                .optimize(alp, lp_arena, expr_arena)
+                .expect("slice pushdown failed");
+
+            lp_arena.replace(lp_top, alp);
+        }
+
         if type_coercion {
             rules.push(Box::new(TypeCoercionRule {}))
         }
@@ -852,7 +863,7 @@ impl LazyFrame {
     }
 
     /// Slice the DataFrame.
-    pub fn slice(self, offset: i64, len: usize) -> LazyFrame {
+    pub fn slice(self, offset: i64, len: u32) -> LazyFrame {
         let opt_state = self.get_opt_state();
         let lp = self.get_plan_builder().slice(offset, len).build();
         Self::from_logical_plan(lp, opt_state)
@@ -869,7 +880,7 @@ impl LazyFrame {
     }
 
     /// Get the n last rows
-    pub fn tail(self, n: usize) -> LazyFrame {
+    pub fn tail(self, n: u32) -> LazyFrame {
         let neg_tail = -(n as i64);
         self.slice(neg_tail, n)
     }
@@ -886,7 +897,7 @@ impl LazyFrame {
 
     /// Limit the DataFrame to the first `n` rows. Note if you don't want the rows to be scanned,
     /// use [fetch](LazyFrame::fetch).
-    pub fn limit(self, n: usize) -> LazyFrame {
+    pub fn limit(self, n: u32) -> LazyFrame {
         self.slice(0, n)
     }
 

polars/polars-lazy/src/frame/parquet.rs

@@ -45,12 +45,19 @@ impl LazyFrame {
                 .map(|r| {
                     let path = r.map_err(|e| PolarsError::ComputeError(format!("{e}").into()))?;
                     let path_string = path.to_string_lossy().into_owned();
-                    Self::scan_parquet_impl(path_string, args.n_rows, args.cache, false)
+                    Self::scan_parquet_impl(path_string, None, args.cache, false)
                 })
                 .collect::<Result<Vec<_>>>()?;
 
             concat(&lfs, args.rechunk)
                 .map_err(|_| PolarsError::ComputeError("no matching files found".into()))
+                .map(|lf| {
+                    if let Some(n_rows) = args.n_rows {
+                        lf.slice(0, n_rows as u32)
+                    } else {
+                        lf
+                    }
+                })
         } else {
             Self::scan_parquet_impl(path, args.n_rows, args.cache, args.parallel)
         }

polars/polars-lazy/src/functions.rs

@@ -326,7 +326,10 @@ pub fn concat<L: AsRef<[LazyFrame]>>(inputs: L, rechunk: bool) -> Result<LazyFra
         lps.push(lp)
     }
 
-    let lp = LogicalPlan::Union { inputs: lps };
+    let lp = LogicalPlan::Union {
+        inputs: lps,
+        options: Default::default(),
+    };
     let mut lf = LazyFrame::from(lp);
     lf.opt_state = opt_state;
 

polars/polars-lazy/src/logical_plan/alp.rs

@@ -26,7 +26,7 @@ pub enum ALogicalPlan {
     Slice {
         input: Node,
         offset: i64,
-        len: usize,
+        len: u32,
     },
     Selection {
         input: Node,
@@ -130,6 +130,7 @@ pub enum ALogicalPlan {
     },
     Union {
         inputs: Vec<Node>,
+        options: UnionOptions,
     },
 }
 
@@ -194,7 +195,10 @@ impl ALogicalPlan {
         use ALogicalPlan::*;
 
         match self {
-            Union { .. } => Union { inputs },
+            Union { options, .. } => Union {
+                inputs,
+                options: *options,
+            },
             Melt {
                 id_vars,
                 value_vars,

polars/polars-lazy/src/logical_plan/builder.rs

@@ -357,7 +357,7 @@ impl LogicalPlanBuilder {
         .into()
     }
 
-    pub fn slice(self, offset: i64, len: usize) -> Self {
+    pub fn slice(self, offset: i64, len: u32) -> Self {
         LogicalPlan::Slice {
             input: Box::new(self.0),
             offset,

polars/polars-lazy/src/logical_plan/conversion.rs

@@ -157,12 +157,12 @@ pub(crate) fn to_alp(
     lp_arena: &mut Arena<ALogicalPlan>,
 ) -> Node {
     let v = match lp {
-        LogicalPlan::Union { inputs } => {
+        LogicalPlan::Union { inputs, options } => {
             let inputs = inputs
                 .into_iter()
                 .map(|lp| to_alp(lp, expr_arena, lp_arena))
                 .collect();
-            ALogicalPlan::Union { inputs }
+            ALogicalPlan::Union { inputs, options }
         }
         LogicalPlan::Selection { input, predicate } => {
             let i = to_alp(*input, expr_arena, lp_arena);
@@ -642,12 +642,12 @@ pub(crate) fn node_to_lp(
     let lp = std::mem::take(lp);
 
     match lp {
-        ALogicalPlan::Union { inputs } => {
+        ALogicalPlan::Union { inputs, options } => {
             let inputs = inputs
                 .into_iter()
                 .map(|node| node_to_lp(node, expr_arena, lp_arena))
                 .collect();
-            LogicalPlan::Union { inputs }
+            LogicalPlan::Union { inputs, options }
         }
         ALogicalPlan::Slice { input, offset, len } => {
             let lp = node_to_lp(input, expr_arena, lp_arena);

polars/polars-lazy/src/logical_plan/mod.rs

@@ -17,15 +17,15 @@ mod format;
 pub(crate) mod iterator;
 mod lit;
 pub(crate) mod optimizer;
+mod options;
 mod projection;
-mod scans;
 
 use polars_core::frame::groupby::DynamicGroupOptions;
 
 pub(crate) use apply::*;
 pub(crate) use builder::*;
 pub use lit::*;
-pub(crate) use scans::*;
+pub(crate) use options::*;
 
 // Will be set/ unset in the fetch operation to communicate overwriting the number of rows to scan.
 thread_local! {pub(crate) static FETCH_ROWS: Cell<Option<usize>> = Cell::new(None)}
@@ -47,9 +47,7 @@ pub enum LogicalPlan {
         predicate: Expr,
     },
     /// Cache the input at this point in the LP
-    Cache {
-        input: Box<LogicalPlan>,
-    },
+    Cache { input: Box<LogicalPlan> },
     /// Scan a CSV file
     #[cfg(feature = "csv-file")]
     CsvScan {
@@ -147,7 +145,7 @@ pub enum LogicalPlan {
     Slice {
         input: Box<LogicalPlan>,
         offset: i64,
-        len: usize,
+        len: u32,
     },
     /// A Melt operation
     Melt {
@@ -168,6 +166,7 @@ pub enum LogicalPlan {
     },
     Union {
         inputs: Vec<LogicalPlan>,
+        options: UnionOptions,
     },
 }
 

polars/polars-lazy/src/logical_plan/optimizer/mod.rs

@@ -9,6 +9,7 @@ pub(crate) mod fast_projection;
 pub(crate) mod predicate_pushdown;
 pub(crate) mod projection_pushdown;
 pub(crate) mod simplify_expr;
+pub(crate) mod slice_pushdown;
 pub(crate) mod stack_opt;
 pub(crate) mod type_coercion;
 

polars/polars-lazy/src/logical_plan/optimizer/predicate_pushdown/mod.rs

@@ -291,6 +291,7 @@ impl PredicatePushDown {
                 maintain_order,
                 dynamic_options,
             } => {
+                // start with a new empty predicate aggregator
                 self.pushdown_and_assign(input, optimizer::init_hashmap(), lp_arena, expr_arena)?;
 
                 // dont push down predicates. An aggregation needs all rows