This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11993-ed478d7b795d5e03b916213e40b7a03e860e5366 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
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
