This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch alamb/coalesce-copy-rows-safe
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/alamb/coalesce-copy-rows-safe
by this push:
new cd3e3cf602 Mark filter-based append helpers unsafe and check filter
length in callers
cd3e3cf602 is described below
commit cd3e3cf60229cb593a154ddf838f83b2c44437bc
Author: Andrew Lamb <[email protected]>
AuthorDate: Wed Oct 21 15:46:01 2026 -0400
Mark filter-based append helpers unsafe and check filter length in callers
append_values_by_indices, append_values_by_slices and
append_views_by_filter write through raw pointers using indices taken from a
FilterPredicate, relying on the predicate matching the source array. Mark them
unsafe fn with documented preconditions, and have the safe copy_rows_by_filter
callers verify filter_len <= source len before calling them, returning an error
otherwise.
Co-Authored-By: Claude Fable 5.1 <[email protected]>
---
arrow-select/src/coalesce.rs | 16 ++++++
arrow-select/src/coalesce/byte_view.rs | 53 ++++++++++++++++++--
arrow-select/src/coalesce/primitive.rs | 92 +++++++++++++++++++++++++++-------
3 files changed, 140 insertions(+), 21 deletions(-)
diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs
index 6f5378a8f5..3687243287 100644
--- a/arrow-select/src/coalesce.rs
+++ b/arrow-select/src/coalesce.rs
@@ -790,6 +790,22 @@ fn create_in_progress_array(data_type: &DataType,
batch_size: usize) -> Box<dyn
}
}
+/// Returns an error if `filter` selects rows beyond `source_len`.
+///
+/// Specialized [`InProgressArray`] implementations rely on this check before
+/// copying rows selected by a [`FilterPredicate`] with unchecked indexing:
+/// every index the predicate yields is less than its filter length, so once
+/// `filter.filter_len() <= source_len` all selected rows are in bounds.
+fn check_filter_len(filter: &FilterPredicate, source_len: usize) -> Result<(),
ArrowError> {
+ let filter_len = filter.filter_len();
+ if filter_len > source_len {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "Filter predicate of length {filter_len} is larger than source
array of length {source_len}"
+ )));
+ }
+ Ok(())
+}
+
/// Incrementally builds up arrays
///
/// [`GenericInProgressArray`] is the default implementation that buffers
diff --git a/arrow-select/src/coalesce/byte_view.rs
b/arrow-select/src/coalesce/byte_view.rs
index 74c10b3fd7..d86d6810e0 100644
--- a/arrow-select/src/coalesce/byte_view.rs
+++ b/arrow-select/src/coalesce/byte_view.rs
@@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-use crate::coalesce::InProgressArray;
+use crate::coalesce::{InProgressArray, check_filter_len};
use crate::filter::{FilterPredicate, FilterSelection};
use arrow_array::cast::AsArray;
use arrow_array::types::ByteViewType;
@@ -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();
let current_len = self.views.len();
self.views.reserve(selected_count);
let mut written = 0;
+ // SAFETY: the caller guarantees all selected rows are in bounds,
+ // `FilterPredicate` guarantees exactly `selected_count` rows are
+ // selected, and capacity was reserved above.
unsafe {
let mut out =
self.views.spare_capacity_mut().as_mut_ptr().cast::<u128>();
@@ -448,8 +457,14 @@ impl<B: ByteViewType> InProgressArray for
InProgressByteViewArray<B> {
));
}
+ if let Err(e) = check_filter_len(filter, s.len()) {
+ self.source = Some(source);
+ return Err(e);
+ }
+
self.append_nulls_by_filter(filter, s.nulls());
- self.append_views_by_filter(s.views(), filter);
+ // SAFETY: `check_filter_len` verified all selected rows are in bounds
+ unsafe { self.append_views_by_filter(s.views(), filter) };
self.source = Some(source);
Ok(())
@@ -464,9 +479,11 @@ impl<B: ByteViewType> InProgressArray for
InProgressByteViewArray<B> {
// The source views reference no external buffers, so they must all be
// inline and we can copy just the nulls and views.
if s.data_buffers().is_empty() {
+ check_filter_len(filter, s.len())?;
self.ensure_capacity();
self.append_nulls_by_filter(filter, s.nulls());
- self.append_views_by_filter(s.views(), filter);
+ // SAFETY: `check_filter_len` verified all selected rows are in
bounds
+ unsafe { self.append_views_by_filter(s.views(), filter) };
return Ok(());
}
@@ -736,4 +753,32 @@ mod tests {
assert_eq!(in_progress_size_with_split,
in_progress_size_without_split);
}
+
+ #[test]
+ fn test_copy_rows_by_filter_rejects_filter_longer_than_source() {
+ let values: Vec<Option<&[u8]>> = vec![Some(b"a"), Some(b"b"),
Some(b"c")];
+ let array = BinaryViewArray::from_iter(values);
+ assert!(array.data_buffers().is_empty());
+ let array: ArrayRef = Arc::new(array);
+
+ let filter = BooleanArray::from_iter((0..8).map(|idx| Some(idx != 3)));
+ let predicate = FilterBuilder::new(&filter).build();
+
+ let mut in_progress =
InProgressByteViewArray::<BinaryViewType>::new(8);
+ in_progress.set_source(Some(Arc::clone(&array)));
+ let err = in_progress.copy_rows_by_filter(&predicate).unwrap_err();
+ assert!(
+ err.to_string().contains("larger than source array"),
+ "unexpected error: {err}"
+ );
+
+ let mut in_progress =
InProgressByteViewArray::<BinaryViewType>::new(8);
+ let err = in_progress
+ .copy_rows_by_filter_from(array, &predicate)
+ .unwrap_err();
+ assert!(
+ err.to_string().contains("larger than source array"),
+ "unexpected error: {err}"
+ );
+ }
}
diff --git a/arrow-select/src/coalesce/primitive.rs
b/arrow-select/src/coalesce/primitive.rs
index b547409c75..159bc084a9 100644
--- a/arrow-select/src/coalesce/primitive.rs
+++ b/arrow-select/src/coalesce/primitive.rs
@@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-use crate::coalesce::InProgressArray;
+use crate::coalesce::{InProgressArray, check_filter_len};
use crate::filter::{FilterIndices, FilterPredicate, FilterSelection,
FilterSlices};
use arrow_array::cast::AsArray;
use arrow_array::{Array, ArrayRef, ArrowPrimitiveType, PrimitiveArray};
@@ -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(
current: &mut Vec<T::Native>,
values: &[T::Native],
indices: FilterIndices<'_>,
@@ -70,6 +78,8 @@ impl<T: ArrowPrimitiveType> InProgressPrimitiveArray<T> {
let current_len = current.len();
let mut written = 0;
+ // SAFETY: the caller guarantees all indices are in bounds, and that
+ // `current` has capacity for `selected_count` more elements.
unsafe {
let mut out = current
.spare_capacity_mut()
@@ -77,7 +87,6 @@ impl<T: ArrowPrimitiveType> InProgressPrimitiveArray<T> {
.cast::<T::Native>();
indices.for_each(|idx| {
- // SAFETY: indices are derived from the filter predicate for
this source.
out.write(*values.get_unchecked(idx));
out = out.add(1);
written += 1;
@@ -89,7 +98,16 @@ impl<T: ArrowPrimitiveType> InProgressPrimitiveArray<T> {
debug_assert_eq!(written, selected_count);
}
- fn append_values_by_slices(
+ /// Append the `values` in each `start..end` range of `slices` to
`current`.
+ ///
+ /// # Safety
+ ///
+ /// * Every `(start, end)` yielded by `slices` must satisfy
+ /// `start <= end <= values.len()`
+ /// * `slices` must cover exactly `selected_count` elements in total
+ /// * `current` must have spare capacity for at least `selected_count`
+ /// additional elements
+ unsafe fn append_values_by_slices(
current: &mut Vec<T::Native>,
values: &[T::Native],
slices: FilterSlices<'_>,
@@ -98,6 +116,8 @@ impl<T: ArrowPrimitiveType> InProgressPrimitiveArray<T> {
let current_len = current.len();
let mut written = 0;
+ // SAFETY: the caller guarantees all slices are in bounds, and that
+ // `current` has capacity for `selected_count` more elements.
unsafe {
let mut out = current
.spare_capacity_mut()
@@ -106,7 +126,6 @@ impl<T: ArrowPrimitiveType> InProgressPrimitiveArray<T> {
slices.for_each(|(start, end)| {
let len = end - start;
- // SAFETY: slices are derived from the filter predicate for
this source.
std::ptr::copy_nonoverlapping(values.as_ptr().add(start), out,
len);
out = out.add(len);
written += len;
@@ -175,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.
@@ -372,4 +403,31 @@ mod tests {
"null buffer must be included in size(): with_nulls={with_nulls}
all_valid={all_valid}"
);
}
+
+ #[test]
+ fn test_copy_rows_by_filter_rejects_filter_longer_than_source() {
+ let source = Int32Array::from(vec![1, 2, 3]);
+ let mut in_progress = InProgressPrimitiveArray::<Int32Type>::new(8,
DataType::Int32);
+ in_progress.set_source(Some(Arc::new(source)));
+
+ // Index iterator: a sparse filter longer than the source
+ let filter = BooleanArray::from_iter((0..21).map(|idx| Some(idx ==
20)));
+ let predicate = FilterBuilder::new(&filter).build();
+ assert!(matches!(predicate.selection(), FilterSelection::Indices(_)));
+ let err = in_progress.copy_rows_by_filter(&predicate).unwrap_err();
+ assert!(
+ err.to_string().contains("larger than source array"),
+ "unexpected error: {err}"
+ );
+
+ // Slice iterator: a dense filter longer than the source
+ let filter = BooleanArray::from_iter((0..8).map(|idx| Some(idx != 3)));
+ let predicate = FilterBuilder::new(&filter).build();
+ assert!(matches!(predicate.selection(), FilterSelection::Slices(_)));
+ let err = in_progress.copy_rows_by_filter(&predicate).unwrap_err();
+ assert!(
+ err.to_string().contains("larger than source array"),
+ "unexpected error: {err}"
+ );
+ }
}
