sunchao commented on a change in pull request #1248:
URL: https://github.com/apache/arrow-rs/pull/1248#discussion_r800221300
##########
File path: arrow/src/compute/kernels/filter.rs
##########
@@ -185,79 +287,585 @@ pub fn prep_null_mask_filter(filter: &BooleanArray) ->
BooleanArray {
/// # Ok(())
/// # }
/// ```
-pub fn filter(array: &dyn Array, predicate: &BooleanArray) -> Result<ArrayRef>
{
- if predicate.null_count() > 0 {
- // this greatly simplifies subsequent filtering code
- // now we only have a boolean mask to deal with
- let predicate = prep_null_mask_filter(predicate);
- return filter(array, &predicate);
+pub fn filter(values: &dyn Array, predicate: &BooleanArray) ->
Result<ArrayRef> {
+ let predicate = FilterBuilder::new(predicate).build();
+ filter_array(values, &predicate)
+}
+
+/// Returns a new [RecordBatch] with arrays containing only values matching
the filter.
+pub fn filter_record_batch(
+ record_batch: &RecordBatch,
+ predicate: &BooleanArray,
+) -> Result<RecordBatch> {
+ let mut filter_builder = FilterBuilder::new(predicate);
+ if record_batch.num_columns() > 1 {
+ // Only optimize if filtering more than one column
+ filter_builder = filter_builder.optimize();
+ }
+ let filter = filter_builder.build();
+
+ let filtered_arrays = record_batch
+ .columns()
+ .iter()
+ .map(|a| filter_array(a, &filter))
+ .collect::<Result<Vec<_>>>()?;
+
+ RecordBatch::try_new(record_batch.schema(), filtered_arrays)
+}
+
+/// A builder to construct [`FilterPredicate`]
+#[derive(Debug)]
+pub struct FilterBuilder {
+ filter: BooleanArray,
+ count: usize,
+ strategy: IterationStrategy,
+}
+
+impl FilterBuilder {
+ /// Create a new [`FilterBuilder`] that can be used to construct a
[`FilterPredicate`]
+ pub fn new(filter: &BooleanArray) -> Self {
+ let filter = match filter.null_count() {
+ 0 => BooleanArray::from(filter.data().clone()),
+ _ => prep_null_mask_filter(filter),
+ };
+
+ let count = filter_count(&filter);
Review comment:
For the special case where `count` is equal to the length of original
array (i.e., all values are selected), is there a faster path?
##########
File path: arrow/src/compute/kernels/filter.rs
##########
@@ -119,17 +155,83 @@ impl<'a> Iterator for SlicesIterator<'a> {
}
}
+/// An iterator of `usize` whose index in [`BooleanArray`] is true
+///
+/// This provides the best performance on all but the least selective
predicates (which keep most
+/// / all rows), where the benefits of copying large runs instead favours
[`SlicesIterator`]
Review comment:
`keep most / all rows`?
##########
File path: arrow/src/compute/kernels/filter.rs
##########
@@ -185,79 +287,585 @@ pub fn prep_null_mask_filter(filter: &BooleanArray) ->
BooleanArray {
/// # Ok(())
/// # }
/// ```
-pub fn filter(array: &dyn Array, predicate: &BooleanArray) -> Result<ArrayRef>
{
- if predicate.null_count() > 0 {
- // this greatly simplifies subsequent filtering code
- // now we only have a boolean mask to deal with
- let predicate = prep_null_mask_filter(predicate);
- return filter(array, &predicate);
+pub fn filter(values: &dyn Array, predicate: &BooleanArray) ->
Result<ArrayRef> {
Review comment:
Not quite related to this PR, but I'm wondering that in the case there
are multiple predicates on the same input array, maybe it's beneficial to only
materialize the final output array after all the predicates have been
evaluated. It seems we don't have such API yet?
--
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]
