This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 274f433115 refactor(coalesce): mark unchecked append helpers unsafe
(#11993)
274f433115 is described below
commit 274f433115806af8a5ba557131f1c3fb9fc91630
Author: RIchard Baah <[email protected]>
AuthorDate: Sat Oct 31 11:31:46 2026 +0000
refactor(coalesce): mark unchecked append helpers unsafe (#11993)
# Which issue does this PR close?
- Relates to #11174
- Supersedes #11782, originally started by @alamb
# Rationale for this change
Some coalesce append helpers use unchecked indexing and raw pointer
writes, but were exposed as safe functions even though their correctness
depends on caller-provided bounds and byte-view invariants.
# What changes are included in this PR?
- Mark the primitive and byte-view append helpers unsafe
- Document their bounds and `B::validate` requirements
- Add safety comments at the existing call sites
- Use checked slicing in the safe `copy_rows` implementations
- Have the primitive append helpers reserve their own capacity
- Validate primitive filter lengths before selection dispatch
- Test an all-true filter longer than its source
# Are there any significant user-facing changes?
No.
---------
Co-authored-by: Andrew Lamb <[email protected]>
---
arrow-select/src/coalesce.rs | 15 ++++++
arrow-select/src/coalesce/byte_view.rs | 85 ++++++++++++++++++++++++++--------
arrow-select/src/coalesce/primitive.rs | 83 ++++++++++++++++++++++++---------
3 files changed, 143 insertions(+), 40 deletions(-)
diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs
index 1f6da98708..7c583c1319 100644
--- a/arrow-select/src/coalesce.rs
+++ b/arrow-select/src/coalesce.rs
@@ -848,6 +848,21 @@ fn create_in_progress_array(data_type: &DataType,
batch_size: usize) -> Box<dyn
}
}
+/// Return an error if `filter` is longer than the source array it is applied
to.
+///
+/// A [`FilterPredicate`] yields positions less than its length, so once this
+/// check passes, the filter is guaranteed to produce valid indices less than
+/// `source_len`.
+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 468aab0b52..e5d796eba1 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;
@@ -41,13 +41,13 @@ pub(crate) struct InProgressByteViewArray<B: ByteViewType> {
source: Option<Source>,
/// the target batch size (and thus size for views allocation)
batch_size: usize,
- /// The in progress views
+ /// The in progress views (guaranteed to be valid)
views: Vec<u128>,
/// In progress nulls
nulls: NullBufferBuilder,
/// current buffer
current: Option<Vec<u8>>,
- /// completed buffers
+ /// completed buffers (corresponding to content of views)
completed: Vec<Buffer>,
/// Allocates new buffers of increasing size as needed
buffer_source: BufferSource,
@@ -121,13 +121,19 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
self.completed.push(buffer);
}
- fn append_views_by_filter(&mut self, views: &[u128], filter:
&FilterPredicate) {
+ /// # Safety
+ ///
+ /// `filter.filter_len()` must not exceed `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 that selected rows are in bounds,
+ // `FilterPredicate` yields `selected_count` rows, and capacity was
+ // reserved above.
unsafe {
let mut out =
self.views.spare_capacity_mut().as_mut_ptr().cast::<u128>();
@@ -173,8 +179,16 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
}
/// Append views to self.views, updating the buffer index if necessary
+ ///
+ /// # Safety
+ ///
+ /// [`B::validate`](ByteViewType::validate) must succeed for `views` and
+ /// `buffers`.
+ ///
+ /// Note: this function performs no unsafe operations itself, but it
assumes
+ /// `views` are valid views for self.views and `buffers` for self.buffers
#[inline(never)]
- fn append_views_and_update_buffer_index(
+ unsafe fn append_views_and_update_buffer_index(
&mut self,
views: &[u128],
buffers: &[Buffer],
@@ -233,9 +247,14 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
/// - `views` - the views to append
/// - `view_buffer_size` - the total number of bytes pointed to by all
/// views (used to allocate new buffers if needed)
- /// - `buffers` - the buffers the reviews point to
+ /// - `buffers` - the buffers the views point to
+ ///
+ /// # Safety
+ ///
+ /// [`B::validate`](ByteViewType::validate) must succeed for `views` and
+ /// `buffers`.
#[inline(never)]
- fn append_views_and_copy_strings(
+ unsafe fn append_views_and_copy_strings(
&mut self,
views: &[u128],
view_buffer_size: usize,
@@ -248,7 +267,8 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
// If there is no current buffer, allocate a new one
let Some(current) = self.current.take() else {
let new_buffer = self.buffer_source.next_buffer(view_buffer_size);
- self.append_views_and_copy_strings_inner(views, new_buffer,
buffers);
+ // SAFETY: the caller guarantees the views and buffers are valid.
+ unsafe { self.append_views_and_copy_strings_inner(views,
new_buffer, buffers) };
return;
};
@@ -256,7 +276,8 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
// copy the strings into the existing buffer.
let mut remaining_capacity = current.capacity() - current.len();
if view_buffer_size <= remaining_capacity {
- self.append_views_and_copy_strings_inner(views, current, buffers);
+ // SAFETY: the caller guarantees the views and buffers are valid.
+ unsafe { self.append_views_and_copy_strings_inner(views, current,
buffers) };
return;
}
@@ -282,7 +303,8 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
let string_bytes_to_copy = current.capacity() - current.len() -
remaining_capacity;
let remaining_view_buffer_size = view_buffer_size -
string_bytes_to_copy;
- self.append_views_and_copy_strings_inner(first_views, current,
buffers);
+ // SAFETY: `first_views` is a subset of the caller-validated views.
+ unsafe { self.append_views_and_copy_strings_inner(first_views,
current, buffers) };
let completed: Buffer = self.current.take().expect("completed").into();
self.completed_buffers_size += completed.capacity();
self.completed.push(completed);
@@ -290,7 +312,8 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
// Copy any remaining views into a new buffer
let remaining_views = &views[num_view_to_current..];
let new_buffer =
self.buffer_source.next_buffer(remaining_view_buffer_size);
- self.append_views_and_copy_strings_inner(remaining_views, new_buffer,
buffers);
+ // SAFETY: `remaining_views` is a subset of the caller-validated views.
+ unsafe { self.append_views_and_copy_strings_inner(remaining_views,
new_buffer, buffers) };
}
/// Append views to self.views, copying data from the buffers into
@@ -299,9 +322,14 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
/// # Panics:
/// If `self.current` is `Some`
///
+ /// # Safety
+ ///
+ /// [`B::validate`](ByteViewType::validate) must succeed for `views` and
+ /// `buffers`.
+ ///
/// See `append_views_and_copy_strings` for more details
#[inline(never)]
- fn append_views_and_copy_strings_inner(
+ unsafe fn append_views_and_copy_strings_inner(
&mut self,
views: &[u128],
mut dst_buffer: Vec<u8>,
@@ -427,8 +455,7 @@ impl<B: ByteViewType> InProgressArray for
InProgressByteViewArray<B> {
}
let buffers = s.data_buffers();
- // SAFETY: copy_rows is called with ranges derived from the source
array.
- let views = unsafe { s.views().as_ref().get_unchecked(offset..offset +
len) };
+ let views = &s.views().as_ref()[offset..offset + len];
// If there are no data buffers in s (all inlined views), can append
the
// views/nulls and done
@@ -441,9 +468,11 @@ impl<B: ByteViewType> InProgressArray for
InProgressByteViewArray<B> {
// Copying the strings into a buffer can be time-consuming so
// only do it if the array is sparse
if source.need_gc {
- self.append_views_and_copy_strings(views,
source.ideal_buffer_size, buffers);
+ // SAFETY: `views` and `buffers` come from the same valid
byte-view array.
+ unsafe { self.append_views_and_copy_strings(views,
source.ideal_buffer_size, buffers) };
} else {
- self.append_views_and_update_buffer_index(views, buffers, true);
+ // SAFETY: `views` and `buffers` come from the same valid
byte-view array.
+ unsafe { self.append_views_and_update_buffer_index(views, buffers,
true) };
}
self.source = Some(source);
Ok(())
@@ -468,8 +497,16 @@ impl<B: ByteViewType> InProgressArray for
InProgressByteViewArray<B> {
));
}
+ if let Err(e) = check_filter_len(filter, s.len()) {
+ // Restore the source taken above before returning the error.
+ 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 the filter is no longer than
+ // `s.views()`.
+ unsafe { self.append_views_by_filter(s.views(), filter) };
self.source = Some(source);
Ok(())
@@ -484,9 +521,12 @@ 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 the filter is no longer than
+ // `s.views()`.
+ unsafe { self.append_views_by_filter(s.views(), filter) };
return Ok(());
}
@@ -502,7 +542,14 @@ impl<B: ByteViewType> InProgressArray for
InProgressByteViewArray<B> {
} else {
self.nulls.append_n_non_nulls(filter.count());
}
- self.append_views_and_update_buffer_index(filtered.views(),
filtered.data_buffers(), false);
+ // SAFETY: the views and buffers come from the same valid filtered
array.
+ unsafe {
+ self.append_views_and_update_buffer_index(
+ filtered.views(),
+ filtered.data_buffers(),
+ false,
+ )
+ };
Ok(())
}
diff --git a/arrow-select/src/coalesce/primitive.rs
b/arrow-select/src/coalesce/primitive.rs
index 62460ab677..e2818eff1a 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,15 +61,23 @@ impl<T: ArrowPrimitiveType> InProgressPrimitiveArray<T> {
}
}
- fn append_values_by_indices(
+ /// # Safety
+ ///
+ /// Every index yielded by `indices` must be less than `values.len()`, and
+ /// `indices` must yield exactly `selected_count` indices.
+ unsafe fn append_values_by_indices(
current: &mut Vec<T::Native>,
values: &[T::Native],
indices: FilterIndices<'_>,
selected_count: usize,
) {
+ current.reserve(selected_count);
let current_len = current.len();
let mut written = 0;
+ // SAFETY: the caller guarantees that indices are in bounds and that
+ // exactly `selected_count` values are written. Capacity was reserved
+ // above.
unsafe {
let mut out = current
.spare_capacity_mut()
@@ -77,7 +85,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,15 +96,24 @@ impl<T: ArrowPrimitiveType> InProgressPrimitiveArray<T> {
debug_assert_eq!(written, selected_count);
}
- fn append_values_by_slices(
+ /// # Safety
+ ///
+ /// Every range yielded by `slices` must satisfy
+ /// `start <= end <= values.len()`, and the ranges must contain exactly
+ /// `selected_count` values in total.
+ unsafe fn append_values_by_slices(
current: &mut Vec<T::Native>,
values: &[T::Native],
slices: FilterSlices<'_>,
selected_count: usize,
) {
+ current.reserve(selected_count);
let current_len = current.len();
let mut written = 0;
+ // SAFETY: the caller guarantees that slices are in bounds and that
+ // exactly `selected_count` values are written. Capacity was reserved
+ // above.
unsafe {
let mut out = current
.spare_capacity_mut()
@@ -106,7 +122,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;
@@ -164,27 +179,32 @@ impl<T: ArrowPrimitiveType + Debug> InProgressArray for
InProgressPrimitiveArray
// Copy the values
let values = s.values();
- // SAFETY: copy_rows is called with ranges derived from the source
array.
self.current
- .extend_from_slice(unsafe { values.get_unchecked(offset..offset +
len) });
+ .extend_from_slice(&values[offset..offset + len]);
Ok(())
}
fn copy_rows_by_filter(&mut self, filter: &FilterPredicate) -> Result<(),
ArrowError> {
+ let source_len = primitive_source::<T>(self.source.as_ref())?.len();
+ check_filter_len(filter, source_len)?;
+
match filter.selection() {
FilterSelection::Indices(indices) => {
self.ensure_capacity();
let s = primitive_source::<T>(self.source.as_ref())?;
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 the selected indices are
+ // in bounds, and the filter yields `filter.count()` indices.
+ unsafe {
+ Self::append_values_by_indices(
+ &mut self.current,
+ s.values(),
+ indices,
+ filter.count(),
+ )
+ };
Ok(())
}
FilterSelection::Slices(slices) => {
@@ -192,13 +212,16 @@ impl<T: ArrowPrimitiveType + Debug> InProgressArray for
InProgressPrimitiveArray
let s = primitive_source::<T>(self.source.as_ref())?;
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 the selected slices are
+ // in bounds, and the filter yields `filter.count()` values.
+ unsafe {
+ Self::append_values_by_slices(
+ &mut self.current,
+ s.values(),
+ slices,
+ filter.count(),
+ )
+ };
Ok(())
}
// Other selection shapes reuse the generic copy_rows path.
@@ -305,6 +328,24 @@ mod tests {
assert_eq!(result, &expected);
}
+ #[test]
+ fn test_copy_rows_by_filter_rejects_all_filter_longer_than_source() {
+ let source: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3]));
+ let mut in_progress = InProgressPrimitiveArray::<Int32Type>::new(4,
DataType::Int32);
+ in_progress.set_source(Some(source));
+
+ let filter = BooleanArray::from(vec![true; 4]);
+ let predicate = FilterBuilder::new(&filter).build();
+ assert!(matches!(predicate.selection(), FilterSelection::All { .. }));
+
+ let err = in_progress.copy_rows_by_filter(&predicate).unwrap_err();
+ assert!(matches!(
+ err,
+ ArrowError::InvalidArgumentError(message)
+ if message == "Filter predicate of length 4 is larger than
source array of length 3"
+ ));
+ }
+
#[test]
fn test_size_empty() {
// A fresh in-progress array has allocated nothing yet
