feat: Add projection pushdown to new streaming multiscan

pola-rs · Feb 7, 2025 · 2647bf6 · 2647bf6
1 parent 1145ec0
commit 2647bf6
Show file tree

Hide file tree

Showing 7 changed files with 413 additions and 204 deletions.
diff --git a/crates/polars-core/src/schema.rs b/crates/polars-core/src/schema.rs
@@ -1,5 +1,6 @@
 use std::fmt::Debug;
 
+use arrow::bitmap::Bitmap;
 use polars_utils::pl_str::PlSmallStr;
 
 use crate::prelude::*;
@@ -20,6 +21,9 @@ pub trait SchemaExt {
     fn iter_fields(&self) -> impl ExactSizeIterator<Item = Field> + '_;
 
     fn to_supertype(&mut self, other: &Schema) -> PolarsResult<bool>;
+
+    /// Select fields using a bitmap.
+    fn project_select(&self, select: &Bitmap) -> Self;
 }
 
 impl SchemaExt for Schema {
@@ -88,6 +92,15 @@ impl SchemaExt for Schema {
         }
         Ok(changed)
     }
+
+    fn project_select(&self, select: &Bitmap) -> Self {
+        assert_eq!(self.len(), select.len());
+        self.iter()
+            .zip(select.iter())
+            .filter(|(_, select)| *select)
+            .map(|((n, dt), _)| (n.clone(), dt.clone()))
+            .collect()
+    }
 }
 
 pub trait SchemaNamesAndDtypes {