diff --git a/vortex-file/src/pruning.rs b/vortex-file/src/pruning.rs index e19d64d3e8..b022149331 100644 --- a/vortex-file/src/pruning.rs +++ b/vortex-file/src/pruning.rs @@ -15,6 +15,8 @@ use vortex_error::{VortexExpect as _, VortexResult}; use vortex_expr::{BinaryExpr, Column, ExprRef, Literal, Not, Operator}; use vortex_scalar::Scalar; +use crate::RowFilter; + #[derive(Debug, Clone)] pub struct Relation { map: HashMap>, @@ -39,6 +41,17 @@ impl Relation { } } + pub fn union(mut iter: impl Iterator>) -> Relation { + if let Some(mut x) = iter.next() { + for y in iter { + x.extend(y) + } + x + } else { + Relation::new() + } + } + pub fn extend(&mut self, other: Relation) { for (l, rs) in other.map.into_iter() { self.map.entry(l).or_default().extend(rs.into_iter()) @@ -183,6 +196,21 @@ fn convert_to_pruning_expression(expr: &ExprRef) -> PruningPredicateStats { } } + if let Some(RowFilter { conjunction }) = expr.as_any().downcast_ref::() { + let (rewritten_conjunction, refses): (Vec, Vec>) = + conjunction + .iter() + .map(convert_to_pruning_expression) + .unzip(); + + let refs = Relation::union(refses.into_iter()); + + return ( + RowFilter::from_conjunction_expr(rewritten_conjunction), + refs, + ); + } + not_prunable() } diff --git a/vortex-file/src/read/filtering.rs b/vortex-file/src/read/filtering.rs index ec5e0c94ab..87fe00c945 100644 --- a/vortex-file/src/read/filtering.rs +++ b/vortex-file/src/read/filtering.rs @@ -18,7 +18,7 @@ use crate::read::expr_project::expr_project; #[derive(Debug, Clone)] pub struct RowFilter { - conjunction: Vec, + pub(crate) conjunction: Vec, } impl RowFilter { @@ -37,6 +37,11 @@ impl RowFilter { Self { conjunction } } + /// Create a new row filter from a conjunction. The conjunction **must** have length > 0. + pub fn from_conjunction_expr(conjunction: Vec) -> Arc { + Arc::new(Self::from_conjunction(conjunction)) + } + pub fn only_fields(&self, fields: &[Field]) -> Option { let conj = self .conjunction