haohuaijin commented on code in PR #10859:
URL: https://github.com/apache/arrow-rs/pull/10859#discussion_r4110204003
##########
parquet/src/arrow/arrow_reader/filter.rs:
##########
@@ -198,4 +206,540 @@ 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,
+ row_selection_policy: RowSelectionPolicy,
+ ) -> 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,
row_selection_policy)));
+ } 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. Selections follow the reader's row selection policy.
+struct FusedPredicate {
+ /// At least two predicates, all with the same projection.
+ predicates: Vec<Box<dyn ArrowPredicate>>,
+ row_selection_policy: RowSelectionPolicy,
+}
+
+impl FusedPredicate {
+ /// Create a fused predicate from at least two predicates that share one
+ /// projection.
+ fn new(
+ predicates: Vec<Box<dyn ArrowPredicate>>,
+ row_selection_policy: RowSelectionPolicy,
+ ) -> Self {
+ debug_assert!(predicates.len() > 1);
+ debug_assert!(
+ predicates
+ .windows(2)
+ .all(|pair| pair[0].projection() == pair[1].projection())
+ );
+ Self {
+ predicates,
+ row_selection_policy,
+ }
+ }
+}
+
+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) => adapt_fusion_selection(prev,
self.row_selection_policy)
+ .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))
+ }
+}
+
+/// Apply the reader's row selection policy to a fusion selection.
+fn adapt_fusion_selection(selection: RowSelection, policy: RowSelectionPolicy)
-> RowSelection {
Review Comment:
update in
[6984860](https://github.com/apache/arrow-rs/pull/10859/commits/6984860b5dc44453043189d00d49d4fe239ff82c)
--
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]
