Sketch out cached filter result API

Author: alamb

parquet/src/arrow/arrow_reader/mod.rs

@@ -694,7 +694,11 @@ impl<T: ChunkReader + 'static> ParquetRecordBatchReaderBuilder<T> {
                 let array_reader =
                     build_array_reader(self.fields.as_deref(), predicate.projection(), &reader)?;
 
-                plan_builder = plan_builder.with_predicate(array_reader, predicate.as_mut())?;
+                plan_builder = plan_builder.with_predicate(
+                    array_reader,
+                    predicate.as_mut(),
+                    &self.projection,
+                )?;
             }
         }
 

parquet/src/arrow/arrow_reader/read_plan.rs

@@ -22,9 +22,12 @@ use crate::arrow::array_reader::ArrayReader;
 use crate::arrow::arrow_reader::{
     ArrowPredicate, ParquetRecordBatchReader, RowSelection, RowSelector,
 };
+use crate::arrow::ProjectionMask;
 use crate::errors::{ParquetError, Result};
-use arrow_array::Array;
-use arrow_select::filter::prep_null_mask_filter;
+use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, RecordBatchReader};
+use arrow_schema::{DataType, Schema};
+use arrow_select::concat::concat;
+use arrow_select::filter::{filter, prep_null_mask_filter};
 use std::collections::VecDeque;
 
 /// A builder for [`ReadPlan`]
@@ -35,7 +38,8 @@ pub(crate) struct ReadPlanBuilder {
     batch_size: usize,
     /// Current to apply, includes all filters
     selection: Option<RowSelection>,
-    // TODO: Cached result of evaluating some columns with the RowSelection
+    /// Cached result of evaluating some columns with the RowSelection
+    cached_predicate_result: Option<CachedPredicateResult>,
 }
 
 impl ReadPlanBuilder {
@@ -46,6 +50,7 @@ impl ReadPlanBuilder {
         Self {
             batch_size,
             selection: None,
+            cached_predicate_result: None,
         }
     }
 
@@ -88,41 +93,60 @@ impl ReadPlanBuilder {
 
     /// Evaluates an [`ArrowPredicate`], updating the read plan's selection
     ///
+    /// # Arguments
+    ///
+    /// * `array_reader`: The array reader to use for evaluating the predicate.
+    ///   must be configured with the projection mask specified by
+    ///   [`ArrowPredicate::projection`] for the `predicate`.
+    ///
+    /// * `predicate`: The predicate to evaluate
+    ///
+    /// * `projection`: The projection mask that will be selected. This code will
+    ///   potentially cache the results of filtering columns that also appear in the
+    ///   projection mask.
+    ///
     /// If `this.selection` is `Some`, the resulting  [`RowSelection`] will be
-    /// the conjunction of it and the rows selected by `predicate`.
+    /// the conjunction of it and the rows selected by `predicate` (they will be
+    /// `AND`ed).
     ///
-    /// Note: A pre-existing selection may come from evaluating a previous predicate
-    /// or if the [`ParquetRecordBatchReader`] specified an explicit
+    /// Note: A pre-existing selection may come from evaluating a previous
+    /// predicate or if the [`ParquetRecordBatchReader`] specifies an explicit
     /// [`RowSelection`] in addition to one or more predicates.
     pub(crate) fn with_predicate(
         mut self,
         array_reader: Box<dyn ArrayReader>,
         predicate: &mut dyn ArrowPredicate,
+        projection_mask: &ProjectionMask,
     ) -> Result<Self> {
+        // Prepare to decode all rows in the selection to evaluate the predicate
         let reader = ParquetRecordBatchReader::new(array_reader, self.clone().build());
-        let mut filters = vec![];
+        let mut cached_results_builder = CachedPredicateResultBuilder::new(
+            &reader.schema(),
+            predicate.projection(),
+            projection_mask,
+        );
         for maybe_batch in reader {
-            let maybe_batch = maybe_batch?;
-            let input_rows = maybe_batch.num_rows();
-            let filter = predicate.evaluate(maybe_batch)?;
+            let batch = maybe_batch?;
+            let input_rows = batch.num_rows();
+            let filter = predicate.evaluate(batch.clone())?;
             // Since user supplied predicate, check error here to catch bugs quickly
             if filter.len() != input_rows {
                 return Err(arrow_err!(
                     "ArrowPredicate predicate returned {} rows, expected {input_rows}",
                     filter.len()
                 ));
             }
-            match filter.null_count() {
-                0 => filters.push(filter),
-                _ => filters.push(prep_null_mask_filter(&filter)),
-            };
+            cached_results_builder.add(batch, filter)?;
         }
 
-        let raw = RowSelection::from_filters(&filters);
+        let (raw, cached_predicate_result) =
+            cached_results_builder.build(self.batch_size, predicate.projection())?;
         self.selection = match self.selection.take() {
             Some(selection) => Some(selection.and_then(&raw)),
             None => Some(raw),
         };
+
+        self.cached_predicate_result = cached_predicate_result;
         Ok(self)
     }
 
@@ -135,13 +159,15 @@ impl ReadPlanBuilder {
         let Self {
             batch_size,
             selection,
+            cached_predicate_result,
         } = self;
 
         let selection = selection.map(|s| s.trim().into());
 
         ReadPlan {
             batch_size,
             selection,
+            cached_predicate_result,
         }
     }
 }
@@ -238,7 +264,8 @@ pub(crate) struct ReadPlan {
     batch_size: usize,
     /// Row ranges to be selected from the data source
     selection: Option<VecDeque<RowSelector>>,
-    // TODO: Cached result of evaluating some columns with the RowSelection
+    /// Cached result of evaluating some columns with the RowSelection
+    cached_predicate_result: Option<CachedPredicateResult>,
 }
 
 impl ReadPlan {
@@ -252,3 +279,237 @@ impl ReadPlan {
         self.batch_size
     }
 }
+
+/// Incrementally builds the result of evaluating a ArrowPredicate on
+/// a RowGroup.
+struct CachedPredicateResultBuilder {
+    /// The entire result of the predicate evaluation in memory
+    ///
+    /// TODO: potentially incrementally build the result of the predicate
+    /// evaluation without holding all the batches in memory. See
+    /// <https://github.com/apache/arrow-rs/issues/6692>
+    in_progress_arrays: Vec<Box<dyn InProgressArray>>,
+    filters: Vec<BooleanArray>,
+}
+
+impl CachedPredicateResultBuilder {
+    /// Create a new CachedPredicateResultBuilder
+    ///
+    /// # Arguments:
+    /// * `schema`: The schema of the filter record batch
+    /// * `filter_mask`: which columns of the original parquet schema did the filter columns come from?
+    /// * `projection_mask`: which columns of the original parquet schema are in the final projection?
+    ///
+    /// This structure does not cache filter results for the columns that are not
+    /// in the projection mask. This is because the filter results are not needed
+    fn new(
+        schema: &Schema,
+        filter_mask: &ProjectionMask,
+        projection_mask: &ProjectionMask,
+    ) -> Self {
+        let mut field_iter = schema.fields.iter();
+
+        let (filter_mask_inner, projection_mask_inner) =
+            match (filter_mask.mask(), projection_mask.mask()) {
+                (Some(filter_mask), Some(projection_mask)) => (filter_mask, projection_mask),
+                // NB, None means all columns and we just want the intersection of the two
+                (Some(filter_mask), None) => (filter_mask, filter_mask),
+                (None, Some(projection_mask)) => (projection_mask, projection_mask),
+                (None, None) => {
+                    // this means all columns are in the projection and filter so cache them all when possible
+                    let in_progress_arrays = field_iter
+                        .into_iter()
+                        .map(|field| create_in_progress_array(true, field.data_type()))
+                        .collect();
+                    return {
+                        Self {
+                            in_progress_arrays,
+                            filters: vec![],
+                        }
+                    };
+                }
+            };
+
+        let mut in_progress_arrays = Vec::with_capacity(filter_mask_inner.len());
+
+        for (&in_filter, &in_projection) in
+            filter_mask_inner.iter().zip(projection_mask_inner.iter())
+        {
+            if !in_filter {
+                continue;
+            }
+            // field is in the filter
+            let field = field_iter.next().expect("mismatch in field lengths");
+            in_progress_arrays.push(create_in_progress_array(in_projection, field.data_type()));
+        }
+        assert_eq!(in_progress_arrays.len(), schema.fields().len());
+
+        Self {
+            in_progress_arrays,
+            filters: vec![],
+        }
+    }
+
+    /// Add a new batch and filter to the builder
+    fn add(&mut self, batch: RecordBatch, mut filter: BooleanArray) -> Result<()> {
+        if filter.null_count() > 0 {
+            filter = prep_null_mask_filter(&filter);
+        }
+
+        let (_schema, columns, _row_count) = batch.into_parts();
+
+        for (in_progress, array) in self.in_progress_arrays.iter_mut().zip(columns.into_iter()) {
+            in_progress.append(array, &filter)?;
+        }
+
+        self.filters.push(filter);
+        Ok(())
+    }
+
+    /// Return (selection, maybe_cached_predicate_result) that represents the rows
+    /// that were selected and batches that were evaluated.
+    fn build(
+        self,
+        _batch_size: usize,
+        filter_mask: &ProjectionMask,
+    ) -> Result<(RowSelection, Option<CachedPredicateResult>)> {
+        let Self {
+            in_progress_arrays,
+            filters,
+        } = self;
+
+        let new_selection = RowSelection::from_filters(&filters);
+
+        let Some(mask) = filter_mask.mask() else {
+            return Ok((new_selection, None));
+        };
+
+        let mut arrays: Vec<Option<ArrayRef>> = vec![None; mask.len()];
+        let mut in_progress_arrays = VecDeque::from(in_progress_arrays);
+
+        // Now find the location in the original parquet schema of the filter columns
+        for i in 0..mask.len() {
+            if mask[i] {
+                let mut in_progress = in_progress_arrays
+                    .pop_front()
+                    .expect("insufficient inprogress arrays");
+                arrays[i] = in_progress.try_build()?;
+            }
+        }
+
+        let cached_result = CachedPredicateResult { arrays, filters };
+
+        Ok((new_selection, Some(cached_result)))
+    }
+}
+
+/// The result of evaluating a predicate on a RowGroup with a specific
+/// RowSelection
+#[derive(Clone)]
+struct CachedPredicateResult {
+    /// Map of parquet schema column index to the result of evaluating the predicate
+    /// on that column.
+    ///
+    /// NOTE each array already has the corresponding filters applied
+    ///
+    /// TODO: store as Vec<Vec<ArrayRef>> to avoid having to have one large
+    /// array for each column
+    arrays: Vec<Option<ArrayRef>>,
+    /// The results of evaluating the predicate (this has already been applied to the
+    /// cached results)
+    filters: Vec<BooleanArray>,
+}
+
+impl CachedPredicateResult {
+    fn empty() -> Self {
+        Self {
+            arrays: vec![],
+            filters: vec![],
+        }
+    }
+
+    /// Apply the results of other to self
+    /// Updates cached filter result and filters potentially
+    pub fn merge(self, other: Self) -> Self {
+        // TODO do something with self
+        other
+    }
+}
+
+/// Progressively creates array from filtered values
+///
+/// TODO avoid buffering the input memory
+trait InProgressArray {
+    /// Appends all values of the array to the in progress array at locations where filter[i] is true
+    /// to the in progress array
+    fn append(&mut self, _array: ArrayRef, filter: &BooleanArray) -> Result<()>;
+
+    /// Builds the final array, consuming all state from self. Returns None if the array
+    /// cannot be created (e.g. data type not supported or out of buffer space)
+    fn try_build(&mut self) -> Result<Option<ArrayRef>>;
+}
+
+/// Return a new InProgressArray for the given data type
+///
+/// if `in_projection` is false then a NoOpInProgressArray is returned (will not
+/// actually cache arrays results)
+fn create_in_progress_array(
+    in_projection: bool,
+    _data_type: &DataType,
+) -> Box<dyn InProgressArray> {
+    if in_projection {
+        Box::new(GenericInProgressArray::new())
+    } else {
+        // column is not in the projection, so no need to cache
+        Box::new(NoOpInProgressArray {})
+    }
+}
+
+/// Placeholder that does nothing until we support the entire set of datatypes
+struct NoOpInProgressArray {}
+
+impl InProgressArray for NoOpInProgressArray {
+    fn append(&mut self, _array: ArrayRef, _filter: &BooleanArray) -> Result<()> {
+        // do nothing
+        Ok(())
+    }
+    fn try_build(&mut self) -> Result<Option<ArrayRef>> {
+        // do nothing
+        Ok(None)
+    }
+}
+
+/// a generic implementation of InProgressArray that uses filter and concat kernels
+/// to create the final array
+///
+/// TODO: make this better with per type implementations
+/// <https://github.com/apache/arrow-rs/issues/6692>
+struct GenericInProgressArray {
+    /// previously filtered arrays
+    arrays: Vec<ArrayRef>,
+}
+
+impl GenericInProgressArray {
+    fn new() -> Self {
+        Self { arrays: vec![] }
+    }
+}
+impl InProgressArray for GenericInProgressArray {
+    fn append(&mut self, array: ArrayRef, filter_array: &BooleanArray) -> Result<()> {
+        self.arrays.push(filter(&array, filter_array)?);
+        Ok(())
+    }
+
+    fn try_build(&mut self) -> Result<Option<ArrayRef>> {
+        if self.arrays.is_empty() {
+            return Ok(None);
+        }
+        if self.arrays.len() == 1 {
+            return Ok(Some(self.arrays.pop().unwrap()));
+        }
+        // Vomit: need to copy to a new Vec to get dyn array
+        let arrays: Vec<&dyn Array> = self.arrays.iter().map(|a| a.as_ref()).collect();
+        let array = concat(&arrays)?;
+        Ok(Some(array))
+    }
+}

parquet/src/arrow/async_reader/mod.rs

@@ -618,7 +618,8 @@ where
                 let array_reader =
                     build_array_reader(self.fields.as_deref(), predicate.projection(), &row_group)?;
 
-                plan_builder = plan_builder.with_predicate(array_reader, predicate.as_mut())?;
+                plan_builder =
+                    plan_builder.with_predicate(array_reader, predicate.as_mut(), &projection)?;
             }
         }