reduce peak memory of reading parquet by row groups (#4006)

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

@@ -164,13 +164,22 @@ impl DataFrame {
         POOL.install(|| self.columns.par_iter().map(|s| func(s)).collect())
     }
 
+    // reduce monomorphization
     fn try_apply_columns_par(
         &self,
         func: &(dyn Fn(&Series) -> Result<Series> + Send + Sync),
     ) -> Result<Vec<Series>> {
         POOL.install(|| self.columns.par_iter().map(|s| func(s)).collect())
     }
 
+    // reduce monomorphization
+    fn try_apply_columns(
+        &self,
+        func: &(dyn Fn(&Series) -> Result<Series> + Send + Sync),
+    ) -> Result<Vec<Series>> {
+        self.columns.iter().map(|s| func(s)).collect()
+    }
+
     /// Get the index of the column.
     fn check_name_to_idx(&self, name: &str) -> Result<usize> {
         self.find_idx_by_name(name)
@@ -1473,6 +1482,12 @@ impl DataFrame {
         Ok(DataFrame::new_no_checks(new_col))
     }
 
+    /// Same as `filter` but does not parallelize.
+    pub fn _filter_seq(&self, mask: &BooleanChunked) -> Result<Self> {
+        let new_col = self.try_apply_columns(&|s| s.filter(mask))?;
+        Ok(DataFrame::new_no_checks(new_col))
+    }
+
     /// Take `DataFrame` value by indexes from an iterator.
     ///
     /// # Example

polars/polars-io/src/parquet/read_impl.rs

@@ -136,7 +136,7 @@ fn rg_to_dfs(
             df.with_row_count_mut(&rc.name, Some(previous_row_count + rc.offset));
         }
 
-        apply_predicate(&mut df, predicate.as_deref())?;
+        apply_predicate(&mut df, predicate.as_deref(), true)?;
         apply_aggregations(&mut df, aggregate)?;
 
         previous_row_count += current_row_count;
@@ -201,7 +201,7 @@ fn rg_to_dfs_par(
                 df.with_row_count_mut(&rc.name, Some(row_count_start as IdxSize + rc.offset));
             }
 
-            apply_predicate(&mut df, predicate.as_deref())?;
+            apply_predicate(&mut df, predicate.as_deref(), false)?;
             apply_aggregations(&mut df, aggregate)?;
 
             Ok(Some(df))

polars/polars-io/src/predicates.rs

@@ -38,11 +38,17 @@ pub(crate) fn arrow_schema_to_empty_df(schema: &ArrowSchema) -> DataFrame {
 pub(crate) fn apply_predicate(
     df: &mut DataFrame,
     predicate: Option<&dyn PhysicalIoExpr>,
+    parallel: bool,
 ) -> Result<()> {
     if let (Some(predicate), false) = (&predicate, df.is_empty()) {
         let s = predicate.evaluate(df)?;
         let mask = s.bool().expect("filter predicates was not of type boolean");
-        *df = df.filter(mask)?;
+
+        if parallel {
+            *df = df.filter(mask)?;
+        } else {
+            *df = df._filter_seq(mask)?;
+        }
     }
     Ok(())
 }