1use crate::arrow::ProjectionMask;
6use crate::arrow::arrow_reader::{RowSelection, RowSelectionPolicy};
7use crate::schema::types::SchemaDescriptor;
8use arrow_array::{Array, BooleanArray, RecordBatch};
9use arrow_buffer::BooleanBuffer;
10use arrow_schema::ArrowError;
11use arrow_select::filter::{SlicesIterator, filter_record_batch, prep_null_mask_filter};
12use std::fmt::{Debug, Formatter};
13
14pub trait ArrowPredicate: Send + 'static {
21 fn projection(&self) -> &ProjectionMask;
29
30 fn evaluate(&mut self, batch: RecordBatch) -> Result<BooleanArray, ArrowError>;
38}
39
40pub struct ArrowPredicateFn<F> {
98 f: F,
99 projection: ProjectionMask,
100}
101
102impl<F> ArrowPredicateFn<F>
103where
104 F: FnMut(RecordBatch) -> Result<BooleanArray, ArrowError> + Send + 'static,
105{
106 pub fn new(projection: ProjectionMask, f: F) -> Self {
109 Self { f, projection }
110 }
111}
112
113impl<F> ArrowPredicate for ArrowPredicateFn<F>
114where
115 F: FnMut(RecordBatch) -> Result<BooleanArray, ArrowError> + Send + 'static,
116{
117 fn projection(&self) -> &ProjectionMask {
118 &self.projection
119 }
120
121 fn evaluate(&mut self, batch: RecordBatch) -> Result<BooleanArray, ArrowError> {
122 (self.f)(batch)
123 }
124}
125
126pub struct RowFilter {
172 pub(crate) predicates: Vec<Box<dyn ArrowPredicate>>,
174}
175
176impl Debug for RowFilter {
177 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
178 write!(f, "RowFilter {{ {} predicates: }}", self.predicates.len())
179 }
180}
181
182impl RowFilter {
183 pub fn new(predicates: Vec<Box<dyn ArrowPredicate>>) -> Self {
185 Self { predicates }
186 }
187 pub fn predicates(&self) -> &Vec<Box<dyn ArrowPredicate>> {
189 &self.predicates
190 }
191 pub fn into_predicates(self) -> Vec<Box<dyn ArrowPredicate>> {
193 self.predicates
194 }
195
196 pub(crate) fn fuse_same_projection(self, parquet_schema: &SchemaDescriptor) -> Self {
199 let mut predicates: Vec<Box<dyn ArrowPredicate>> =
200 Vec::with_capacity(self.predicates.len());
201 let mut group: Vec<Box<dyn ArrowPredicate>> = Vec::new();
202 let mut flush_group = |group: &mut Vec<Box<dyn ArrowPredicate>>| {
203 if group.len() > 1 && can_fuse_projection(group[0].projection(), parquet_schema) {
204 let group = std::mem::take(group);
205 predicates.push(Box::new(FusedPredicate::new(group)));
206 } else {
207 predicates.append(group);
208 }
209 };
210
211 for predicate in self.predicates {
212 if group
213 .last()
214 .is_some_and(|last| last.projection() != predicate.projection())
215 {
216 flush_group(&mut group);
217 }
218 group.push(predicate);
219 }
220 flush_group(&mut group);
221
222 Self { predicates }
223 }
224}
225
226fn can_fuse_projection(projection: &ProjectionMask, parquet_schema: &SchemaDescriptor) -> bool {
228 let mut leaf_indices =
229 (0..parquet_schema.num_columns()).filter(|idx| projection.leaf_included(*idx));
230 let Some(leaf_idx) = leaf_indices.next() else {
231 return false;
232 };
233 if leaf_indices.next().is_some() {
234 return false;
235 }
236
237 let column = parquet_schema.column(leaf_idx);
238 column.path().parts().len() == 1 && column.max_rep_level() == 0
239}
240
241struct FusedPredicate {
245 predicates: Vec<Box<dyn ArrowPredicate>>,
247 composition_policy: RowSelectionPolicy,
251}
252
253impl FusedPredicate {
254 fn new(predicates: Vec<Box<dyn ArrowPredicate>>) -> Self {
257 debug_assert!(predicates.len() > 1);
258 debug_assert!(
259 predicates
260 .windows(2)
261 .all(|pair| pair[0].projection() == pair[1].projection())
262 );
263 Self {
264 predicates,
265 composition_policy: RowSelectionPolicy::default(),
266 }
267 }
268}
269
270impl Debug for FusedPredicate {
271 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
272 write!(
273 f,
274 "FusedPredicate {{ {} predicates }}",
275 self.predicates.len()
276 )
277 }
278}
279
280impl ArrowPredicate for FusedPredicate {
281 fn projection(&self) -> &ProjectionMask {
282 self.predicates[0].projection()
283 }
284
285 fn evaluate(&mut self, batch: RecordBatch) -> Result<BooleanArray, ArrowError> {
286 let num_rows = batch.num_rows();
287 let mut selection: Option<RowSelection> = None;
289 let mut filtered_batch = batch;
290 let last_predicate_idx = self.predicates.len() - 1;
291
292 for (idx, predicate) in self.predicates.iter_mut().enumerate() {
293 let filter = evaluate_predicate(predicate.as_mut(), filtered_batch.clone())?;
294 if idx == last_predicate_idx && selection.is_none() {
296 return Ok(filter);
297 }
298 let true_count = filter.true_count();
299 if true_count == 0 {
300 return Ok(BooleanArray::new(BooleanBuffer::new_unset(num_rows), None));
301 }
302 if true_count == filter.len() {
303 continue;
304 }
305
306 let predicate_selection = RowSelection::from_boolean_buffer(filter.values().clone());
307 selection = Some(match selection.take() {
308 Some(prev) => self
311 .composition_policy
312 .apply(prev)
313 .and_then(&predicate_selection),
314 None => predicate_selection,
315 });
316 if idx != last_predicate_idx {
317 filtered_batch = narrow_batch(&filtered_batch, &filter, true_count)?;
318 }
319 }
320
321 let mask = match selection {
322 Some(selection) => selection.into_boolean_buffer(),
323 None => BooleanBuffer::new_set(num_rows),
324 };
325 debug_assert_eq!(mask.len(), num_rows);
326 Ok(BooleanArray::new(mask, None))
327 }
328}
329
330fn evaluate_predicate(
332 predicate: &mut dyn ArrowPredicate,
333 batch: RecordBatch,
334) -> Result<BooleanArray, ArrowError> {
335 let input_rows = batch.num_rows();
336 let filter = predicate.evaluate(batch)?;
337 if filter.len() != input_rows {
338 return Err(ArrowError::InvalidArgumentError(format!(
339 "ArrowPredicate predicate returned {} rows, expected {input_rows}",
340 filter.len()
341 )));
342 }
343 Ok(match filter.null_count() {
344 0 => filter,
345 _ => prep_null_mask_filter(&filter),
346 })
347}
348
349fn narrow_batch(
354 batch: &RecordBatch,
355 filter: &BooleanArray,
356 true_count: usize,
357) -> Result<RecordBatch, ArrowError> {
358 let mut slices = SlicesIterator::new(filter);
359 let (start, end) = slices
360 .next()
361 .expect("a partially selected filter has a true slice");
362 if end - start == true_count {
363 Ok(batch.slice(start, end - start))
364 } else {
365 filter_record_batch(batch, filter)
366 }
367}
368
369#[cfg(test)]
370mod tests {
371 use super::*;
372 use crate::arrow::array_reader::ArrayReader;
373 use crate::arrow::array_reader::StructArrayReader;
374 use crate::arrow::array_reader::test_util::make_int32_page_reader;
375 use crate::arrow::arrow_reader::ReadPlanBuilder;
376 use crate::schema::parser::parse_message_type;
377 use arrow_array::Int32Array;
378 use arrow_array::cast::AsArray;
379 use arrow_array::types::Int32Type;
380 use arrow_schema::{DataType, Field, Fields, Schema};
381 use std::sync::Arc;
382
383 fn test_schema() -> SchemaDescriptor {
384 let schema = parse_message_type(
385 "message schema {
386 REQUIRED INT32 a;
387 REQUIRED INT32 b;
388 REQUIRED INT32 c;
389 REQUIRED GROUP nested { REQUIRED INT32 d; }
390 REPEATED INT32 e;
391 }",
392 )
393 .unwrap();
394 SchemaDescriptor::new(Arc::new(schema))
395 }
396
397 #[test]
398 fn can_fuse_projection_requires_one_top_level_leaf() {
399 let schema = test_schema();
400
401 let one_leaf = ProjectionMask::leaves(&schema, [1]);
402 let two_leaves = ProjectionMask::leaves(&schema, [0, 2]);
403 let nested_leaf = ProjectionMask::leaves(&schema, [3]);
404 let repeated_leaf = ProjectionMask::leaves(&schema, [4]);
405 assert!(can_fuse_projection(&one_leaf, &schema));
406 assert!(!can_fuse_projection(&two_leaves, &schema));
407 assert!(!can_fuse_projection(&nested_leaf, &schema));
408 assert!(!can_fuse_projection(&repeated_leaf, &schema));
409 assert!(!can_fuse_projection(
410 &ProjectionMask::none(schema.num_columns()),
411 &schema
412 ));
413 }
414
415 fn accept_all(projection: ProjectionMask) -> Box<dyn ArrowPredicate> {
416 Box::new(ArrowPredicateFn::new(projection, |batch| {
417 Ok(BooleanArray::from(vec![true; batch.num_rows()]))
418 }))
419 }
420
421 #[test]
422 fn fuse_same_projection_groups_consecutive_eligible_predicates() {
423 let schema = test_schema();
424 let a = ProjectionMask::leaves(&schema, [0]);
425 let b = ProjectionMask::leaves(&schema, [1]);
426 let ac = ProjectionMask::leaves(&schema, [0, 2]);
427 let nested = ProjectionMask::leaves(&schema, [3]);
428
429 let filter = RowFilter::new(vec![
430 accept_all(a.clone()),
432 accept_all(a.clone()),
433 accept_all(b.clone()),
435 accept_all(ac.clone()),
437 accept_all(ac.clone()),
438 accept_all(nested.clone()),
440 accept_all(nested.clone()),
441 accept_all(a.clone()),
443 accept_all(a.clone()),
444 accept_all(a.clone()),
445 ])
446 .fuse_same_projection(&schema);
447
448 let projections: Vec<_> = filter
449 .predicates()
450 .iter()
451 .map(|predicate| predicate.projection().clone())
452 .collect();
453 assert_eq!(
454 projections,
455 vec![a.clone(), b, ac.clone(), ac, nested.clone(), nested, a]
456 );
457 assert_eq!(format!("{filter:?}"), "RowFilter { 7 predicates: }");
458 }
459
460 #[test]
461 fn fuse_same_projection_keeps_single_predicates_and_empty_filters() {
462 let schema = test_schema();
463 let a = ProjectionMask::leaves(&schema, [0]);
464 let b = ProjectionMask::leaves(&schema, [1]);
465
466 let filter =
467 RowFilter::new(vec![accept_all(a), accept_all(b)]).fuse_same_projection(&schema);
468 assert_eq!(filter.predicates().len(), 2);
469
470 let filter = RowFilter::new(vec![]).fuse_same_projection(&schema);
471 assert!(filter.predicates().is_empty());
472 }
473
474 fn int_batch(values: impl IntoIterator<Item = i32>) -> RecordBatch {
475 let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)]));
476 let values = Int32Array::from_iter_values(values);
477 RecordBatch::try_new(schema, vec![Arc::new(values)]).unwrap()
478 }
479
480 fn int_predicate(
483 precondition: impl Fn(i32) -> bool + Send + 'static,
484 f: impl Fn(i32) -> Option<bool> + Send + 'static,
485 ) -> Box<dyn ArrowPredicate> {
486 Box::new(ArrowPredicateFn::new(ProjectionMask::all(), move |batch| {
487 let values = batch.column(0).as_primitive::<Int32Type>();
488 assert!(
489 values.values().iter().all(|v| precondition(*v)),
490 "predicate saw rows rejected by a predecessor"
491 );
492 Ok(values.values().iter().map(|v| f(*v)).collect())
493 }))
494 }
495
496 fn even() -> Box<dyn ArrowPredicate> {
498 int_predicate(|_| true, |v| Some(v % 2 == 0))
499 }
500
501 fn in_range() -> Box<dyn ArrowPredicate> {
503 int_predicate(|_| true, |v| Some((10..40).contains(&v)))
504 }
505
506 fn divisible_by_three_after_even() -> Box<dyn ArrowPredicate> {
508 int_predicate(
509 |v| v % 2 == 0,
510 |v| if v % 7 == 0 { None } else { Some(v % 3 == 0) },
511 )
512 }
513
514 fn all_true() -> Box<dyn ArrowPredicate> {
515 int_predicate(|_| true, |_| Some(true))
516 }
517
518 fn all_false() -> Box<dyn ArrowPredicate> {
519 int_predicate(|_| true, |_| Some(false))
520 }
521
522 fn never_called() -> Box<dyn ArrowPredicate> {
523 Box::new(ArrowPredicateFn::new(ProjectionMask::all(), |_| {
524 panic!("predicate evaluated after every row was rejected")
525 }))
526 }
527
528 fn sequential(batch: &RecordBatch, predicates: &mut [Box<dyn ArrowPredicate>]) -> Vec<bool> {
531 let mut kept = vec![true; batch.num_rows()];
532 for predicate in predicates {
533 if !kept.contains(&true) {
535 break;
536 }
537 let survivors = filter_record_batch(batch, &BooleanArray::from(kept.clone())).unwrap();
538 let filter = predicate.evaluate(survivors).unwrap();
539 for (filter_idx, keep) in kept.iter_mut().filter(|keep| **keep).enumerate() {
540 *keep = filter.is_valid(filter_idx) && filter.value(filter_idx);
541 }
542 }
543 kept
544 }
545
546 const POLICIES: [RowSelectionPolicy; 4] = [
547 RowSelectionPolicy::Selectors,
548 RowSelectionPolicy::Mask,
549 RowSelectionPolicy::Auto { threshold: 32 },
550 RowSelectionPolicy::Auto { threshold: 2 },
551 ];
552
553 type PredicateChain = fn() -> Vec<Box<dyn ArrowPredicate>>;
554
555 #[test]
556 fn fused_predicate_matches_sequential_evaluation() {
557 let cases: Vec<PredicateChain> = vec![
558 || vec![even(), divisible_by_three_after_even()],
559 || vec![even(), in_range(), divisible_by_three_after_even()],
560 || vec![in_range(), even()],
561 || vec![even(), all_true(), divisible_by_three_after_even()],
562 || vec![all_true(), even()],
563 || vec![all_true(), all_true()],
564 || vec![even(), all_false(), never_called()],
565 || vec![all_false(), never_called()],
566 ];
567
568 for num_rows in [0, 1, 7, 97, 200] {
569 let batch = int_batch(0..num_rows);
570 for make_predicates in &cases {
571 let expected = sequential(&batch, &mut make_predicates());
572 for policy in POLICIES {
573 let mut fused = FusedPredicate {
574 predicates: make_predicates(),
575 composition_policy: policy,
576 };
577 let actual = fused.evaluate(batch.clone()).unwrap();
578 assert_eq!(actual.null_count(), 0);
579 assert_eq!(
580 actual.values().iter().collect::<Vec<_>>(),
581 expected,
582 "{num_rows} rows, {policy:?}"
583 );
584 }
585 }
586 }
587 }
588
589 #[test]
590 fn fused_predicate_rejects_wrong_result_length() {
591 let short = Box::new(ArrowPredicateFn::new(ProjectionMask::all(), |_| {
592 Ok(BooleanArray::from(vec![true]))
593 }));
594 let mut fused = FusedPredicate::new(vec![all_true(), short]);
595 let err = fused.evaluate(int_batch(0..4)).unwrap_err();
596 assert!(
597 err.to_string()
598 .contains("ArrowPredicate predicate returned 1 rows, expected 4"),
599 "{err}"
600 );
601 }
602
603 #[test]
604 fn narrow_batch_slices_contiguous_survivors() {
605 let batch = int_batch(0..6);
606 let filter = BooleanArray::from(vec![false, false, true, true, true, false]);
607 let narrowed = narrow_batch(&batch, &filter, 3).unwrap();
608 assert_eq!(
609 narrowed.column(0).as_primitive::<Int32Type>().values(),
610 &[2, 3, 4]
611 );
612 assert_eq!(narrowed.column(0).to_data().buffers()[0].as_ptr(), unsafe {
614 batch.column(0).to_data().buffers()[0].as_ptr().add(2 * 4)
615 });
616
617 let filter = BooleanArray::from(vec![false, true, false, true, false, false]);
618 let narrowed = narrow_batch(&batch, &filter, 2).unwrap();
619 assert_eq!(
620 narrowed.column(0).as_primitive::<Int32Type>().values(),
621 &[1, 3]
622 );
623 }
624
625 #[test]
628 fn fused_predicate_matches_predicate_major_read_plan() {
629 let data: Vec<i32> = (0..97).collect();
630 let make_reader = || {
631 let levels = vec![0; data.len()];
632 let leaf = make_int32_page_reader(&data, &levels, &levels, 0, 0, None);
633 let struct_type =
634 DataType::Struct(Fields::from(vec![Field::new("c0", DataType::Int32, false)]));
635 Box::new(StructArrayReader::new(
636 struct_type,
637 vec![leaf],
638 0,
639 0,
640 false,
641 None,
642 )) as Box<dyn ArrayReader>
643 };
644 let make_predicates = || vec![even(), divisible_by_three_after_even()];
645
646 let prior = RowSelection::from_filters(&[BooleanArray::from(
647 (0..data.len()).map(|idx| idx % 5 != 0).collect::<Vec<_>>(),
648 )]);
649 for initial in [None, Some(prior)] {
650 let mut sequential = ReadPlanBuilder::new(7).with_selection(initial.clone());
651 for predicate in &mut make_predicates() {
652 sequential = sequential
653 .with_predicate(make_reader(), predicate.as_mut())
654 .unwrap();
655 }
656
657 for policy in POLICIES {
658 let mut fused = FusedPredicate::new(make_predicates());
659 let plan = ReadPlanBuilder::new(7)
660 .with_selection(initial.clone())
661 .with_row_selection_policy(policy)
662 .with_predicate(make_reader(), &mut fused)
663 .unwrap();
664 assert_eq!(plan.selection(), sequential.selection(), "{policy:?}");
665 }
666 }
667 }
668}