alamb commented on code in PR #10859:
URL: https://github.com/apache/arrow-rs/pull/10859#discussion_r4123679478
##########
parquet/src/arrow/arrow_reader/filter.rs:
##########
@@ -198,4 +205,477 @@ impl RowFilter {
pub fn into_predicates(self) -> Vec<Box<dyn ArrowPredicate>> {
self.predicates
}
+
+ /// Fuse consecutive predicates on the same single top-level, non-repeated
leaf.
+ /// This avoids repeated decoding or predicate-cache replay of that column.
+ pub(crate) fn fuse_same_projection(self, parquet_schema:
&SchemaDescriptor) -> Self {
+ let mut predicates: Vec<Box<dyn ArrowPredicate>> =
+ Vec::with_capacity(self.predicates.len());
+ let mut group: Vec<Box<dyn ArrowPredicate>> = Vec::new();
+ let mut flush_group = |group: &mut Vec<Box<dyn ArrowPredicate>>| {
+ if group.len() > 1 && can_fuse_projection(group[0].projection(),
parquet_schema) {
+ let group = std::mem::take(group);
+ predicates.push(Box::new(FusedPredicate::new(group)));
+ } else {
+ predicates.append(group);
+ }
+ };
+
+ for predicate in self.predicates {
+ if group
+ .last()
+ .is_some_and(|last| last.projection() !=
predicate.projection())
+ {
+ flush_group(&mut group);
+ }
+ group.push(predicate);
+ }
+ flush_group(&mut group);
+
+ Self { predicates }
+ }
+}
+
+/// Restrict fusion to one top-level, non-repeated leaf to limit compaction
costs.
+fn can_fuse_projection(projection: &ProjectionMask, parquet_schema:
&SchemaDescriptor) -> bool {
+ let mut leaf_indices =
+ (0..parquet_schema.num_columns()).filter(|idx|
projection.leaf_included(*idx));
+ let Some(leaf_idx) = leaf_indices.next() else {
+ return false;
+ };
+ if leaf_indices.next().is_some() {
+ return false;
+ }
+
+ let column = parquet_schema.column(leaf_idx);
+ column.path().parts().len() == 1 && column.max_rep_level() == 0
+}
+
+/// Evaluate same-projection predicates in order on one decoded batch.
+/// Later predicates see only surviving rows, sliced when contiguous or
compacted
+/// otherwise.
+struct FusedPredicate {
+ /// At least two predicates, all with the same projection.
+ predicates: Vec<Box<dyn ArrowPredicate>>,
+ /// Chooses how intermediate selections are composed. The output is always
+ /// a [`BooleanArray`], so this is independent of the reader's
+ /// [`RowSelectionPolicy`], which governs decoding.
+ composition_policy: RowSelectionPolicy,
+}
+
+impl FusedPredicate {
+ /// Create a fused predicate from at least two predicates that share one
+ /// projection.
+ fn new(predicates: Vec<Box<dyn ArrowPredicate>>) -> Self {
+ debug_assert!(predicates.len() > 1);
+ debug_assert!(
+ predicates
+ .windows(2)
+ .all(|pair| pair[0].projection() == pair[1].projection())
+ );
+ Self {
+ predicates,
+ composition_policy: RowSelectionPolicy::default(),
+ }
+ }
+}
+
+impl Debug for FusedPredicate {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(
+ f,
+ "FusedPredicate {{ {} predicates }}",
+ self.predicates.len()
+ )
+ }
+}
+
+impl ArrowPredicate for FusedPredicate {
+ fn projection(&self) -> &ProjectionMask {
+ self.predicates[0].projection()
+ }
+
+ fn evaluate(&mut self, batch: RecordBatch) -> Result<BooleanArray,
ArrowError> {
+ let num_rows = batch.num_rows();
+ // Positions in the original batch; None until a predicate rejects
rows.
+ let mut selection: Option<RowSelection> = None;
+ let mut filtered_batch = batch;
+ let last_predicate_idx = self.predicates.len() - 1;
+
+ for (idx, predicate) in self.predicates.iter_mut().enumerate() {
+ let filter = evaluate_predicate(predicate.as_mut(),
filtered_batch.clone())?;
+ // No mapping is needed if all preceding predicates accepted every
row.
+ if idx == last_predicate_idx && selection.is_none() {
+ return Ok(filter);
+ }
+ let true_count = filter.true_count();
+ if true_count == 0 {
+ return
Ok(BooleanArray::new(BooleanBuffer::new_unset(num_rows), None));
+ }
+ if true_count == filter.len() {
+ continue;
+ }
+
+ let predicate_selection =
RowSelection::from_boolean_buffer(filter.values().clone());
+ selection = Some(match selection.take() {
+ // Only the accumulated selection drives the composition
+ // algorithm, so adapt it once, right before it is used.
+ Some(prev) => self
+ .composition_policy
+ .apply(prev)
+ .and_then(&predicate_selection),
+ None => predicate_selection,
+ });
+ if idx != last_predicate_idx {
+ filtered_batch = narrow_batch(&filtered_batch, &filter,
true_count)?;
+ }
+ }
+
+ let mask = match selection {
+ Some(selection) => selection.into_boolean_buffer(),
+ None => BooleanBuffer::new_set(num_rows),
+ };
+ debug_assert_eq!(mask.len(), num_rows);
+ Ok(BooleanArray::new(mask, None))
+ }
+}
+
+/// Validate the predicate result length and treat nulls as false.
+fn evaluate_predicate(
+ predicate: &mut dyn ArrowPredicate,
+ batch: RecordBatch,
+) -> Result<BooleanArray, ArrowError> {
+ let input_rows = batch.num_rows();
+ let filter = predicate.evaluate(batch)?;
+ if filter.len() != input_rows {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "ArrowPredicate predicate returned {} rows, expected {input_rows}",
+ filter.len()
+ )));
+ }
+ Ok(match filter.null_count() {
+ 0 => filter,
+ _ => prep_null_mask_filter(&filter),
+ })
+}
+
+/// Restrict `batch` to the rows `filter` accepts, slicing zero-copy when they
+/// form one contiguous range and compacting otherwise.
+///
+/// `filter` must be null-free and accept some but not all rows.
+fn narrow_batch(
+ batch: &RecordBatch,
+ filter: &BooleanArray,
+ true_count: usize,
+) -> Result<RecordBatch, ArrowError> {
+ let mut slices = SlicesIterator::new(filter);
Review Comment:
I think we could potentially add this optimization to `filter` and then just
use `filter_record_batch`:
- Filed https://github.com/apache/arrow-rs/issues/11255
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
