alamb commented on code in PR #11206:
URL: https://github.com/apache/arrow-rs/pull/11206#discussion_r4104930313
##########
arrow-select/src/filter.rs:
##########
@@ -199,7 +199,38 @@ pub fn prep_null_mask_filter(filter: &BooleanArray) ->
BooleanArray {
/// assert_eq!(c, &Int32Array::from(vec![5, 8]));
/// ```
pub fn filter(values: &dyn Array, predicate: &BooleanArray) ->
Result<ArrayRef, ArrowError> {
- let mut filter_builder = FilterBuilder::new(predicate);
+ // SAFETY: the count is computed from `predicate` itself.
+ unsafe { filter_with_count(values, predicate, predicate.true_count()) }
+}
+
+/// [`filter`] for a `predicate` whose number of selected rows is already
known.
+///
+/// For callers that already know the total number of selected rows, this is
+/// more efficient than calling `filter` directly.
+///
+/// # Safety
+///
+/// `count` must equal [`BooleanArray::true_count`] of `predicate`; see
+/// [`FilterBuilder::new_with_count`].
+///
+/// # Example
+/// ```rust
+/// # use arrow_array::{Int32Array, BooleanArray};
+/// # use arrow_select::filter::filter_with_count;
+/// let array = Int32Array::from(vec![5, 6, 7, 8, 9]);
+/// let filter_array = BooleanArray::from(vec![true, false, false, true,
false]);
+/// // SAFETY: the count matches the mask.
+/// let c = unsafe { filter_with_count(&array, &filter_array, 2) }.unwrap();
+/// let c = c.as_any().downcast_ref::<Int32Array>().unwrap();
+/// assert_eq!(c, &Int32Array::from(vec![5, 8]));
+/// ```
+pub unsafe fn filter_with_count(
Review Comment:
rather than a new API here, I would prefer to direct people to use
FilterBuilder / FilterPredicate directly if they want to use known counts,
rather than a free function. So maybe we can add a doc comment or something
##########
arrow-select/src/filter.rs:
##########
@@ -254,10 +305,40 @@ pub struct FilterBuilder {
impl FilterBuilder {
/// Create a new [`FilterBuilder`] that can be used to construct a
[`FilterPredicate`]
pub fn new(filter: &BooleanArray) -> Self {
- Self::new_with_count(filter, filter.true_count())
+ // SAFETY: the count is computed from `filter` itself.
+ unsafe { Self::new_with_count(filter, filter.true_count()) }
}
- pub(crate) fn new_with_count(filter: &BooleanArray, count: usize) -> Self {
+ /// Create a new [`FilterBuilder`] from a mask whose number of selected
rows
Review Comment:
what do you think about instead adding `unsafe fn
FilterBuilder::with_count(mut self, count)` ?
So then calls would look like
```rust
let filter = unsafe {
FilterBuilder::new(filter)
.with_count(count) // Safety: count is known
.build()
}
```
That would require deferring the calculation of the count and strategy until
build() but I think that would be ok
--
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]
