diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index 4efa8ca8518b..f38387c0c391 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -841,6 +841,17 @@ fn create_in_progress_array(data_type: &DataType, batch_size: usize) -> Box 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 468aab0b52b9..34bc11c40a06 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 InProgressByteViewArray { 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::(); @@ -427,8 +436,7 @@ impl InProgressArray for InProgressByteViewArray { } 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 @@ -468,8 +476,14 @@ impl InProgressArray for InProgressByteViewArray { )); } + 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(()) @@ -484,9 +498,11 @@ impl InProgressArray for InProgressByteViewArray { // 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(()); } @@ -756,4 +772,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> = 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::::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::::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 62460ab677bd..159bc084a940 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 InProgressPrimitiveArray { } } - 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, values: &[T::Native], indices: FilterIndices<'_>, @@ -70,6 +78,8 @@ impl InProgressPrimitiveArray { 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 InProgressPrimitiveArray { .cast::(); 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 InProgressPrimitiveArray { 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, values: &[T::Native], slices: FilterSlices<'_>, @@ -98,6 +116,8 @@ impl InProgressPrimitiveArray { 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 InProgressPrimitiveArray { 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,9 +183,8 @@ impl 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(()) } @@ -176,29 +194,41 @@ impl InProgressArray for InProgressPrimitiveArray FilterSelection::Indices(indices) => { self.ensure_capacity(); let s = primitive_source::(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::(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. @@ -373,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::::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}" + ); + } }