mbutrovich commented on code in PR #11782:
URL: https://github.com/apache/arrow-rs/pull/11782#discussion_r4342214714
##########
arrow-select/src/coalesce/primitive.rs:
##########
@@ -176,29 +194,41 @@ impl<T: ArrowPrimitiveType + Debug> InProgressArray for
InProgressPrimitiveArray
FilterSelection::Indices(indices) => {
self.ensure_capacity();
let s = primitive_source::<T>(self.source.as_ref())?;
+ check_filter_len(filter, s.len())?;
append_filtered_nulls(&mut self.nulls, s.nulls(), filter);
self.current.reserve(filter.count());
- Self::append_values_by_indices(
- &mut self.current,
- s.values(),
- indices,
- filter.count(),
- );
+ // SAFETY: `check_filter_len` verified all selected indices are
+ // in bounds, `FilterPredicate` guarantees `filter.count()`
+ // rows are selected, and capacity was reserved above.
+ unsafe {
+ Self::append_values_by_indices(
+ &mut self.current,
+ s.values(),
+ indices,
+ filter.count(),
+ )
+ };
Ok(())
}
FilterSelection::Slices(slices) => {
self.ensure_capacity();
let s = primitive_source::<T>(self.source.as_ref())?;
+ check_filter_len(filter, s.len())?;
append_filtered_nulls(&mut self.nulls, s.nulls(), filter);
self.current.reserve(filter.count());
- Self::append_values_by_slices(
- &mut self.current,
- s.values(),
- slices,
- filter.count(),
- );
+ // SAFETY: `check_filter_len` verified all selected slices are
+ // in bounds, `FilterPredicate` guarantees `filter.count()`
+ // rows are selected, and capacity was reserved above.
+ unsafe {
+ Self::append_values_by_slices(
+ &mut self.current,
+ s.values(),
+ slices,
+ filter.count(),
+ )
+ };
Ok(())
}
// Other selection shapes reuse the generic copy_rows path.
Review Comment:
Should `check_filter_len` run once before the `match`? An all-true filter
longer than the source takes this fallback to `copy_rows`. That was UB on
`main`
([L169](https://github.com/apache/arrow-rs/blob/01a9917658fecd639991e92a5a5810be710b539c/arrow-select/src/coalesce/primitive.rs#L169)),
and this PR makes it a panic
([L187](https://github.com/apache/arrow-rs/blob/209188fe7dc568259b5352cc86430aab65334b08/arrow-select/src/coalesce/primitive.rs#L187)).
The other two arms and the byte view impl return the new error. A case for
this shape in `test_copy_rows_by_filter_rejects_filter_longer_than_source`
would cover it.
##########
arrow-select/src/coalesce/primitive.rs:
##########
@@ -61,7 +61,15 @@ impl<T: ArrowPrimitiveType> InProgressPrimitiveArray<T> {
}
}
- fn append_values_by_indices(
+ /// Append the `values` at `indices` to `current`.
+ ///
+ /// # Safety
+ ///
+ /// * Every index yielded by `indices` must be less than `values.len()`
+ /// * `indices` must yield exactly `selected_count` indices
+ /// * `current` must have spare capacity for at least `selected_count`
+ /// additional elements
+ unsafe fn append_values_by_indices(
Review Comment:
What do you think about having `append_values_by_indices` and
`append_values_by_slices` call `current.reserve(selected_count)` themselves,
the way `append_views_by_filter` does ([byte_view.rs
L133](https://github.com/apache/arrow-rs/blob/209188fe7dc568259b5352cc86430aab65334b08/arrow-select/src/coalesce/byte_view.rs#L133))?
That drops the capacity precondition from both `# Safety` sections, and the
callers at
[L200](https://github.com/apache/arrow-rs/blob/209188fe7dc568259b5352cc86430aab65334b08/arrow-select/src/coalesce/primitive.rs#L200)
and
[L220](https://github.com/apache/arrow-rs/blob/209188fe7dc568259b5352cc86430aab65334b08/arrow-select/src/coalesce/primitive.rs#L220)
no longer need their own `reserve`.
##########
arrow-select/src/coalesce/byte_view.rs:
##########
@@ -121,13 +121,22 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
self.completed.push(buffer);
}
- fn append_views_by_filter(&mut self, views: &[u128], filter:
&FilterPredicate) {
+ /// Append the `views` selected by `filter` to `self.views`.
+ ///
+ /// # Safety
+ ///
+ /// Every row selected by `filter` must be less than `views.len()`, i.e.
+ /// `filter.filter_len() <= views.len()`.
+ unsafe fn append_views_by_filter(&mut self, views: &[u128], filter:
&FilterPredicate) {
let selected_count = filter.count();
Review Comment:
Should this check its precondition with a `debug_assert!`?
```suggestion
unsafe fn append_views_by_filter(&mut self, views: &[u128], filter:
&FilterPredicate) {
debug_assert!(filter.filter_len() <= views.len());
let selected_count = filter.count();
```
I don't see this as a pattern in the repo currently (enforcing unsafe
invariants with `debug_assert`) so more of a question than a suggestion.
--
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]
