From 1e0e5a4b3e6190897b166576f3c5b3cfcde745e0 Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Sun, 26 Jul 2026 21:53:00 -0400 Subject: [PATCH 1/8] avoided repeated bounds checks --- arrow-select/src/take.rs | 25 +++++++++++++++++++------ 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index 772f21951ef8..53fc9f154290 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -1122,12 +1122,25 @@ pub fn take_record_batch( record_batch: &RecordBatch, indices: &dyn Array, ) -> Result { - let columns = record_batch - .columns() - .iter() - .map(|c| take(c, indices, None)) - .collect::, _>>()?; - RecordBatch::try_new(record_batch.schema(), columns) + let mut columns = record_batch.columns().iter(); + + let mut taken = Vec::with_capacity(record_batch.num_columns()); + if let Some(first) = columns.next() { + // the only call that actually validates. a bad index surfaces + // as an ArrowError here instead of panicking later + taken.push(take(first, indices, None)?); + } + + // every column in a RecordBatch has the same length, so if the first + // column's indices were in bounds, so are everyone else's + let unchecked = Some(TakeOptions { + check_bounds: false, + }); + for column in columns { + taken.push(take(column, indices, unchecked.clone())?); + } + + RecordBatch::try_new(record_batch.schema(), taken) } #[cfg(test)] From 1a626c58c524e30c945937be9baeac73932abad5 Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Mon, 27 Jul 2026 15:05:34 -0400 Subject: [PATCH 2/8] replace get bounds checks with unsafe blocks --- arrow-select/src/take.rs | 53 ++++++++++++++++++++-------------------- 1 file changed, 26 insertions(+), 27 deletions(-) diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index 53fc9f154290..862db926d157 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -471,7 +471,8 @@ fn take_bits( let output_slice = output_buffer.as_slice_mut(); nulls.valid_indices().for_each(|idx| { // SAFETY: idx is a valid index in indices.nulls() --> idx( } None => { BooleanBuffer::collect_bool(len, |idx: usize| { - // SAFETY: idx( None => { for index in indices.values() { let index = index.as_usize(); - let start = input_offsets[index].as_usize(); - let end = input_offsets[index + 1].as_usize(); + // SAFETY: caller guarantees index < values.len(), and input_offsets has len values.len()+1 + let (start, end) = unsafe { + ( + input_offsets.get_unchecked(index).as_usize(), + input_offsets.get_unchecked(index + 1).as_usize(), + ) + }; capacity += end - start; offsets.push( T::Offset::from_usize(capacity) @@ -568,9 +574,14 @@ fn take_bytes( } // SAFETY: `i` comes from a validity bitmap over `indices`, so it is in-bounds. + // SAFETY: caller guarantees index < values.len(), and input_offsets has len values.len()+1 let index = unsafe { indices.value_unchecked(i) }.as_usize(); - let start = input_offsets[index].as_usize(); - let end = input_offsets[index + 1].as_usize(); + let (start, end) = unsafe { + ( + input_offsets.get_unchecked(index).as_usize(), + input_offsets.get_unchecked(index + 1).as_usize(), + ) + }; capacity += end - start; offsets[i + 1] = T::Offset::from_usize(capacity) .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?; @@ -911,7 +922,8 @@ fn take_fixed_size( None => indices .values() .iter() - .map(|index| buffer[index.as_usize()]) + // SAFETY: caller guarantees all indices are in-bounds (check_bounds or trusted source) + .map(|index| unsafe { *buffer.get_unchecked(index.as_usize()) }) .collect::>(), }; @@ -1122,25 +1134,12 @@ pub fn take_record_batch( record_batch: &RecordBatch, indices: &dyn Array, ) -> Result { - let mut columns = record_batch.columns().iter(); - - let mut taken = Vec::with_capacity(record_batch.num_columns()); - if let Some(first) = columns.next() { - // the only call that actually validates. a bad index surfaces - // as an ArrowError here instead of panicking later - taken.push(take(first, indices, None)?); - } - - // every column in a RecordBatch has the same length, so if the first - // column's indices were in bounds, so are everyone else's - let unchecked = Some(TakeOptions { - check_bounds: false, - }); - for column in columns { - taken.push(take(column, indices, unchecked.clone())?); - } - - RecordBatch::try_new(record_batch.schema(), taken) + let columns = record_batch + .columns() + .iter() + .map(|c| take(c, indices, None)) + .collect::, _>>()?; + RecordBatch::try_new(record_batch.schema(), columns) } #[cfg(test)] From 5887b5b19547458a03384ef7ee2ced395b6cde00 Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Wed, 29 Jul 2026 19:22:49 -0400 Subject: [PATCH 3/8] introduce unsafe alternative for --- arrow-select/src/take.rs | 29 ++++++++++++++++++++++------- 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index 862db926d157..13324dfd59cb 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -97,7 +97,21 @@ pub fn take( check_bounds(values.len(), indices)?; } let indices = indices.to_indices(); - take_impl(values, &indices) + take_impl::<_,true>(values, &indices) + }, + d => Err(ArrowError::InvalidArgumentError(format!("Take only supported for integers, got {d:?}"))) + ) +} +/// unsafe version of take that using unsafe accessor methods. +/// Saftey : caller must gaurentee the indices are valid +pub unsafe fn take_unchecked( + values: &dyn Array, + indices: &dyn Array, +) -> Result { + downcast_integer_array!( + indices => { + let indices = indices.to_indices(); + take_impl::<_,false>(values, &indices) }, d => Err(ArrowError::InvalidArgumentError(format!("Take only supported for integers, got {d:?}"))) ) @@ -211,7 +225,7 @@ where } #[inline(never)] -fn take_impl( +fn take_impl( values: &dyn Array, indices: &PrimitiveArray, ) -> Result { @@ -274,7 +288,7 @@ fn take_impl( let arrays = array .columns() .iter() - .map(|a| take_impl(a.as_ref(), indices)) + .map(|a| take_impl::<_,VALIDATE>(a.as_ref(), indices)) .collect::, _>>()?; let fields: Vec<(FieldRef, ArrayRef)> = fields.iter().cloned().zip(arrays).collect(); @@ -339,7 +353,7 @@ fn take_impl( let type_ids = take_native(values.type_ids(), indices); for (type_id, _field) in fields.iter() { let values = values.child(type_id); - let values = take_impl(values, indices)?; + let values = take_impl::<_,VALIDATE>(values, indices)?; children.push(values); } let array = UnionArray::try_new(fields.clone(), type_ids, None, children)?; @@ -359,7 +373,7 @@ fn take_impl( let values = values.child(field_type_id); - take_impl(values, indices.as_primitive::()) + take_impl::<_,VALIDATE>(values, indices.as_primitive::()) }) .collect::>()?; @@ -781,7 +795,7 @@ fn take_fixed_size_list( length: ::Native, ) -> Result { let list_indices = take_value_indices_from_fixed_size_list(values, indices, length)?; - let taken = take_impl::(values.values().as_ref(), &list_indices)?; + let taken = take_impl::(values.values().as_ref(), &list_indices)?; // determine null count and null buffer, which are a function of `values` and `indices` let num_bytes = bit_util::ceil(indices.len(), 8); @@ -2912,7 +2926,8 @@ mod tests { let logical_indices: PrimitiveArray = PrimitiveArray::from(Vec::::new()); - let result = take_impl(&run_array, &logical_indices).expect("take_run with empty indices"); + let result = take_impl::<_, false>(&run_array, &logical_indices) + .expect("take_run with empty indices"); // Verify the result is a valid empty RunArray assert_eq!(result.len(), 0); From d9fe5cfaabc4c343d68a0283be4dd2d43e34c8b4 Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Wed, 29 Jul 2026 23:57:54 -0400 Subject: [PATCH 4/8] checkmark --- arrow-select/src/take.rs | 93 ++++++++++++++++++++-------------------- 1 file changed, 46 insertions(+), 47 deletions(-) diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index 13324dfd59cb..89b50bfa585b 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -102,7 +102,7 @@ pub fn take( d => Err(ArrowError::InvalidArgumentError(format!("Take only supported for integers, got {d:?}"))) ) } -/// unsafe version of take that using unsafe accessor methods. +/// unsafe version of `take` that using unsafe accessor methods. /// Saftey : caller must gaurentee the indices are valid pub unsafe fn take_unchecked( values: &dyn Array, @@ -233,31 +233,31 @@ fn take_impl( return Ok(new_empty_array(values.data_type())); } downcast_primitive_array! { - values => Ok(Arc::new(take_primitive(values, indices)?)), + values => Ok(Arc::new(take_primitive::<_,_,VALIDATE>(values, indices)?)), DataType::Boolean => { let values = values.as_any().downcast_ref::().unwrap(); - Ok(Arc::new(take_boolean(values, indices))) + Ok(Arc::new(take_boolean::<_,VALIDATE>(values, indices))) } DataType::Utf8 => { - Ok(Arc::new(take_bytes(values.as_string::(), indices)?)) + Ok(Arc::new(take_bytes::<_,_,VALIDATE>(values.as_string::(), indices)?)) } DataType::LargeUtf8 => { - Ok(Arc::new(take_bytes(values.as_string::(), indices)?)) + Ok(Arc::new(take_bytes::<_,_,VALIDATE>(values.as_string::(), indices)?)) } DataType::Utf8View => { - Ok(Arc::new(take_byte_view(values.as_string_view(), indices)?)) + Ok(Arc::new(take_byte_view::<_,_,VALIDATE>(values.as_string_view(), indices)?)) } DataType::List(_) => { - Ok(Arc::new(take_list::<_, Int32Type>(values.as_list(), indices)?)) + Ok(Arc::new(take_list::<_, Int32Type,VALIDATE>(values.as_list(), indices)?)) } DataType::LargeList(_) => { - Ok(Arc::new(take_list::<_, Int64Type>(values.as_list(), indices)?)) + Ok(Arc::new(take_list::<_, Int64Type,VALIDATE>(values.as_list(), indices)?)) } DataType::ListView(_) => { - Ok(Arc::new(take_list_view::<_, Int32Type>(values.as_list_view(), indices)?)) + Ok(Arc::new(take_list_view::<_, Int32Type,VALIDATE>(values.as_list_view(), indices)?)) } DataType::LargeListView(_) => { - Ok(Arc::new(take_list_view::<_, Int64Type>(values.as_list_view(), indices)?)) + Ok(Arc::new(take_list_view::<_, Int64Type,VALIDATE>(values.as_list_view(), indices)?)) } DataType::FixedSizeList(_, length) => { let values = values @@ -272,7 +272,7 @@ fn take_impl( } DataType::Map(field, ordered) => { let list_arr = ListArray::from(values.as_map().clone()); - let list_data = take_list::<_, Int32Type>(&list_arr, indices)?; + let list_data = take_list::<_, Int32Type,VALIDATE>(&list_arr, indices)?; let (_, offsets, entries, nulls) = list_data.into_parts(); let entries = entries.as_struct().clone(); Ok(Arc::new(MapArray::try_new( @@ -313,7 +313,7 @@ fn take_impl( } } DataType::Dictionary(_, _) => downcast_dictionary_array! { - values => Ok(Arc::new(take_dict(values, indices)?)), + values => Ok(Arc::new(take_dict::<_,_,VALIDATE>(values, indices)?)), t => unimplemented!("Take not supported for dictionary type {:?}", t) } DataType::RunEndEncoded(_, _) => downcast_run_array! { @@ -321,20 +321,20 @@ fn take_impl( t => unimplemented!("Take not supported for run type {:?}", t) } DataType::Binary => { - Ok(Arc::new(take_bytes(values.as_binary::(), indices)?)) + Ok(Arc::new(take_bytes::<_,_,VALIDATE>(values.as_binary::(), indices)?)) } DataType::LargeBinary => { - Ok(Arc::new(take_bytes(values.as_binary::(), indices)?)) + Ok(Arc::new(take_bytes::<_,_,VALIDATE>(values.as_binary::(), indices)?)) } DataType::BinaryView => { - Ok(Arc::new(take_byte_view(values.as_binary_view(), indices)?)) + Ok(Arc::new(take_byte_view::<_,_,VALIDATE>(values.as_binary_view(), indices)?)) } DataType::FixedSizeBinary(size) => { let values = values .as_any() .downcast_ref::() .unwrap(); - Ok(Arc::new(take_fixed_size_binary(values, indices, *size)?)) + Ok(Arc::new(take_fixed_size_binary::<_,VALIDATE>(values, indices, *size)?)) } DataType::Null => { // Take applied to a null array produces a null array. @@ -350,7 +350,7 @@ fn take_impl( DataType::Union(fields, UnionMode::Sparse) => { let mut children = Vec::with_capacity(fields.len()); let values = values.as_any().downcast_ref::().unwrap(); - let type_ids = take_native(values.type_ids(), indices); + let type_ids = take_native::<_,_,VALIDATE>(values.type_ids(), indices); for (type_id, _field) in fields.iter() { let values = values.child(type_id); let values = take_impl::<_,VALIDATE>(values, indices)?; @@ -362,8 +362,8 @@ fn take_impl( DataType::Union(fields, UnionMode::Dense) => { let values = values.as_any().downcast_ref::().unwrap(); - let type_ids = >::try_new(take_native(values.type_ids(), indices), None)?; - let offsets = >::try_new(take_native(values.offsets().unwrap(), indices), None)?; + let type_ids = >::try_new(take_native::<_,_,VALIDATE>(values.type_ids(), indices), None)?; + let offsets = >::try_new(take_native::<_,_,VALIDATE>(values.offsets().unwrap(), indices), None)?; let children = fields.iter() .map(|(field_type_id, _)| { @@ -418,7 +418,7 @@ pub struct TakeOptions { /// values: [1, 2, 3, null, 5] /// indices: [0, null, 4, 3] /// The result is: [1 (slot 0), null (null slot), 5 (slot 4), null (slot 3)] -fn take_primitive( +fn take_primitive( values: &PrimitiveArray, indices: &PrimitiveArray, ) -> Result, ArrowError> @@ -426,19 +426,19 @@ where T: ArrowPrimitiveType, I: ArrowPrimitiveType, { - let values_buf = take_native(values.values(), indices); - let nulls = take_nulls(values.nulls(), indices); + let values_buf = take_native::<_, _, VALIDATE>(values.values(), indices); + let nulls = take_nulls::<_, VALIDATE>(values.nulls(), indices); Ok(PrimitiveArray::try_new(values_buf, nulls)?.with_data_type(values.data_type().clone())) } #[inline(never)] -fn take_nulls( +fn take_nulls( values: Option<&NullBuffer>, indices: &PrimitiveArray, ) -> Option { match values.filter(|n| n.null_count() > 0) { Some(n) => NullBuffer::from_unsliced_buffer( - take_bits(n.inner(), indices).into_inner(), + take_bits::<_, VALIDATE>(n.inner(), indices).into_inner(), indices.len(), ), None => indices.nulls().cloned(), @@ -446,7 +446,7 @@ fn take_nulls( } #[inline(never)] -fn take_native( +fn take_native( values: &[T], indices: &PrimitiveArray, ) -> ScalarBuffer { @@ -473,7 +473,7 @@ fn take_native( } #[inline(never)] -fn take_bits( +fn take_bits( values: &BooleanBuffer, indices: &PrimitiveArray, ) -> BooleanBuffer { @@ -485,7 +485,6 @@ fn take_bits( let output_slice = output_buffer.as_slice_mut(); nulls.valid_indices().for_each(|idx| { // SAFETY: idx is a valid index in indices.nulls() --> idx( } /// `take` implementation for boolean arrays -fn take_boolean( +fn take_boolean( values: &BooleanArray, indices: &PrimitiveArray, ) -> BooleanArray { - let val_buf = take_bits(values.values(), indices); - let null_buf = take_nulls(values.nulls(), indices); + let val_buf = take_bits::<_, VALIDATE>(values.values(), indices); + let null_buf = take_nulls::<_, VALIDATE>(values.nulls(), indices); BooleanArray::new(val_buf, null_buf) } /// `take` implementation for string arrays -fn take_bytes( +fn take_bytes( array: &GenericByteArray, indices: &PrimitiveArray, ) -> Result, ArrowError> { @@ -523,7 +522,7 @@ fn take_bytes( let input_offsets = array.value_offsets(); let mut capacity = 0; - let nulls = take_nulls(array.nulls(), indices); + let nulls = take_nulls::<_, VALIDATE>(array.nulls(), indices); // Branch on output nulls — `None` means every output slot is valid. match nulls.as_ref().filter(|n| n.null_count() > 0) { @@ -654,12 +653,12 @@ fn take_bytes( } /// `take` implementation for byte view arrays -fn take_byte_view( +fn take_byte_view( array: &GenericByteViewArray, indices: &PrimitiveArray, ) -> Result, ArrowError> { - let new_views = take_native(array.views(), indices); - let new_nulls = take_nulls(array.nulls(), indices); + let new_views = take_native::<_, _, VALIDATE>(array.views(), indices); + let new_nulls = take_nulls::<_, VALIDATE>(array.nulls(), indices); // Safety: array.views was valid, and take_native copies only valid values, and verifies bounds Ok(unsafe { GenericByteViewArray::new_unchecked(new_views, array.data_buffers().to_vec(), new_nulls) @@ -670,7 +669,7 @@ fn take_byte_view( /// /// Copies the selected list entries' child slices into a new child array /// via `MutableArrayData`, then reconstructs a list array with new offsets -fn take_list( +fn take_list( values: &GenericListArray, indices: &PrimitiveArray, ) -> Result, ArrowError> @@ -682,7 +681,7 @@ where { let list_offsets = values.value_offsets(); let child_data = values.values().to_data(); - let nulls = take_nulls(values.nulls(), indices); + let nulls = take_nulls::<_, VALIDATE>(values.nulls(), indices); let mut new_offsets = Vec::with_capacity(indices.len() + 1); new_offsets.push(OffsetType::Native::zero()); @@ -753,7 +752,7 @@ where GenericListArray::::try_new(field, offsets, child, nulls) } -fn take_list_view( +fn take_list_view( values: &GenericListViewArray, indices: &PrimitiveArray, ) -> Result, ArrowError> @@ -762,9 +761,9 @@ where OffsetType: ArrowPrimitiveType, OffsetType::Native: OffsetSizeTrait, { - let taken_offsets = take_native(values.offsets(), indices); - let taken_sizes = take_native(values.sizes(), indices); - let nulls = take_nulls(values.nulls(), indices); + let taken_offsets = take_native::<_, _, VALIDATE>(values.offsets(), indices); + let taken_sizes = take_native::<_, _, VALIDATE>(values.sizes(), indices); + let nulls = take_nulls::<_, VALIDATE>(values.nulls(), indices); let field = match values.data_type() { DataType::ListView(field) | DataType::LargeListView(field) => field.clone(), @@ -826,7 +825,7 @@ fn take_fixed_size_list( /// The computation is done in two steps: /// - Compute the values buffer /// - Compute the null buffer -fn take_fixed_size_binary( +fn take_fixed_size_binary( values: &FixedSizeBinaryArray, indices: &PrimitiveArray, size: i32, @@ -844,7 +843,7 @@ fn take_fixed_size_binary( _ => take_fixed_size_binary_buffer_dynamic_length(values, indices, size_usize), }; - let value_nulls = take_nulls(values.nulls(), indices); + let value_nulls = take_nulls::<_, VALIDATE>(values.nulls(), indices); let final_nulls = NullBuffer::union(value_nulls.as_ref(), indices.nulls()); return FixedSizeBinaryArray::try_new(size, result_buffer, final_nulls); @@ -957,11 +956,11 @@ fn take_fixed_size( /// /// applies `take` to the keys of the dictionary array and returns a new dictionary array /// with the same dictionary values and reordered keys -fn take_dict( +fn take_dict( values: &DictionaryArray, indices: &PrimitiveArray, ) -> Result, ArrowError> { - let new_keys = take_primitive(values.keys(), indices)?; + let new_keys = take_primitive::<_, _, VALIDATE>(values.keys(), indices)?; Ok(unsafe { DictionaryArray::new_unchecked(new_keys, values.values().clone()) }) } @@ -2278,7 +2277,7 @@ mod tests { // The two middle indices are null -> Should be null in the output. let indices = UInt32Array::from(vec![Some(0), None, None, Some(3)]); - let result = take_fixed_size_binary(&fsb, &indices, 4).unwrap(); + let result = take_fixed_size_binary::<_, true>(&fsb, &indices, 4).unwrap(); assert_eq!(result.len(), 4); assert_eq!(result.null_count(), 2); assert_eq!( @@ -2307,7 +2306,7 @@ mod tests { // The two middle indices are null -> Should be null in the output. let indices = UInt32Array::from(vec![Some(0), None, None, Some(3)]); - let result = take_fixed_size_binary(&fsb, &indices, 5).unwrap(); + let result = take_fixed_size_binary::<_, true>(&fsb, &indices, 5).unwrap(); assert_eq!(result.len(), 4); assert_eq!(result.null_count(), 2); assert_eq!( From 099b4c052cb9478e878c7ecc31fd3f6898c92e1b Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Wed, 29 Jul 2026 23:59:28 -0400 Subject: [PATCH 5/8] fix git --- arrow-select/src/take.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index 89b50bfa585b..1b449fded8e7 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -490,6 +490,7 @@ fn take_bits( unsafe { bit_util::set_bit_raw(output_slice.as_mut_ptr(), idx) }; } }); + BooleanBuffer::new(output_buffer.into(), 0, len) } None => { From e550818ba3e42b9cd7ab38ae03d605d7eed1c1f8 Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Thu, 30 Jul 2026 00:12:30 -0400 Subject: [PATCH 6/8] branch on const bool --- arrow-select/src/take.rs | 80 +++++++++++++++++++++++++++++----------- 1 file changed, 58 insertions(+), 22 deletions(-) diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index 1b449fded8e7..4e6fec4856b7 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -484,19 +484,33 @@ fn take_bits( let mut output_buffer = MutableBuffer::new_null(len); let output_slice = output_buffer.as_slice_mut(); nulls.valid_indices().for_each(|idx| { - // SAFETY: idx is a valid index in indices.nulls() --> idx idx < indices.len() + let index = unsafe { indices.value_unchecked(idx).as_usize() }; + + let value = if VALIDATE { + values.value(index) + } else { + // SAFETY: caller guarantees all valid indices are in-bounds + unsafe { values.value_unchecked(index) } + }; + + if value { // SAFETY: MutableBuffer was created with space for indices.len() bit, and idx < indices.len() unsafe { bit_util::set_bit_raw(output_slice.as_mut_ptr(), idx) }; } }); - BooleanBuffer::new(output_buffer.into(), 0, len) } None => { BooleanBuffer::collect_bool(len, |idx: usize| { - // SAFETY: idx { for index in indices.values() { let index = index.as_usize(); - // SAFETY: caller guarantees index < values.len(), and input_offsets has len values.len()+1 - let (start, end) = unsafe { + + let (start, end) = if VALIDATE { ( - input_offsets.get_unchecked(index).as_usize(), - input_offsets.get_unchecked(index + 1).as_usize(), + input_offsets[index].as_usize(), + input_offsets[index + 1].as_usize(), ) + } else { + // SAFETY: caller guarantees index < values.len(), and input_offsets has len values.len()+1 + unsafe { + ( + input_offsets.get_unchecked(index).as_usize(), + input_offsets.get_unchecked(index + 1).as_usize(), + ) + } }; capacity += end - start; offsets.push( @@ -586,15 +608,21 @@ fn take_bytes( })?; let result_buffer = match size_usize { - 1 => take_fixed_size::(values.values(), indices), - 2 => take_fixed_size::(values.values(), indices), - 4 => take_fixed_size::(values.values(), indices), - 8 => take_fixed_size::(values.values(), indices), - 16 => take_fixed_size::(values.values(), indices), + 1 => take_fixed_size::(values.values(), indices), + 2 => take_fixed_size::(values.values(), indices), + 4 => take_fixed_size::(values.values(), indices), + 8 => take_fixed_size::(values.values(), indices), + 16 => take_fixed_size::(values.values(), indices), _ => take_fixed_size_binary_buffer_dynamic_length(values, indices, size_usize), }; @@ -900,7 +928,7 @@ fn take_fixed_size_binary( /// [feature(generic_const_exprs)](https://github.com/rust-lang/rust/issues/76560) for calling /// `take_fixed_size () } >(...)`. Once this feature has been stabilized, /// we can use this function also in the primitive kernels. -fn take_fixed_size( +fn take_fixed_size( buffer: &Buffer, indices: &PrimitiveArray, ) -> Buffer { @@ -937,7 +965,15 @@ fn take_fixed_size( .values() .iter() // SAFETY: caller guarantees all indices are in-bounds (check_bounds or trusted source) - .map(|index| unsafe { *buffer.get_unchecked(index.as_usize()) }) + .map(|index| { + let index = index.as_usize(); + if VALIDATE { + buffer[index] + } else { + // SAFETY: caller guarantees all indices are in-bounds (check_bounds or trusted source) + unsafe { *buffer.get_unchecked(index) } + } + }) .collect::>(), }; From 6181df3d10a17d642904d511c7c6007cd05967ca Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Thu, 30 Jul 2026 11:31:30 -0400 Subject: [PATCH 7/8] wire up take_record_batch to use take_unchecked --- arrow-select/src/take.rs | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index 4e6fec4856b7..d6996fb56e49 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -102,8 +102,11 @@ pub fn take( d => Err(ArrowError::InvalidArgumentError(format!("Take only supported for integers, got {d:?}"))) ) } -/// unsafe version of `take` that using unsafe accessor methods. -/// Saftey : caller must gaurentee the indices are valid +/// Unsafe version of `take` that uses unsafe accessor methods. +/// +/// # Safety +/// +/// Caller must guarantee that all `indices` are valid (in-bounds) for `values`. pub unsafe fn take_unchecked( values: &dyn Array, indices: &dyn Array, @@ -503,6 +506,7 @@ fn take_bits( } None => { BooleanBuffer::collect_bool(len, |idx: usize| { + // SAFETY: idx, _>>()?; RecordBatch::try_new(record_batch.schema(), columns) } From 2947869ac8b330768c504862f22d18f019076adc Mon Sep 17 00:00:00 2001 From: rich-T-kid Date: Thu, 30 Jul 2026 13:36:39 -0400 Subject: [PATCH 8/8] update all callsites --- arrow-select/src/take.rs | 61 ++++++++++++++++++++++++++++++++-------- 1 file changed, 49 insertions(+), 12 deletions(-) diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index d6996fb56e49..9fd9c97dd0b4 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -458,19 +458,38 @@ fn take_native( .values() .iter() .enumerate() - .map(|(idx, index)| match values.get(index.as_usize()) { - Some(v) => *v, - // SAFETY: idx match unsafe { n.inner().value_unchecked(idx) } { - false => T::default(), - true => panic!("Out-of-bounds index {index:?}"), - }, + .map(|(idx, index)| { + if VALIDATE { + match values.get(index.as_usize()) { + Some(v) => *v, + // SAFETY: idx match unsafe { n.inner().value_unchecked(idx) } { + false => T::default(), + true => panic!("Out-of-bounds index {index:?}"), + }, + } + } else { + // SAFETY: idx < indices.len() + if unsafe { n.inner().value_unchecked(idx) } { + // SAFETY: caller guarantees all valid (non-null) indices are in-bounds + unsafe { *values.get_unchecked(index.as_usize()) } + } else { + T::default() + } + } }) .collect(), None => indices .values() .iter() - .map(|index| values[index.as_usize()]) + .map(|index| { + if VALIDATE { + values[index.as_usize()] + } else { + // SAFETY: caller guarantees all indices are in-bounds + unsafe { *values.get_unchecked(index.as_usize()) } + } + }) .collect(), } } @@ -734,8 +753,17 @@ where None => { for index in indices.values() { let ix = index.as_usize(); - let start = list_offsets[ix].as_usize(); - let end = list_offsets[ix + 1].as_usize(); + let (start, end) = if VALIDATE { + (list_offsets[ix].as_usize(), list_offsets[ix + 1].as_usize()) + } else { + // SAFETY: caller guarantees all indices are in-bounds; list_offsets has len values.len()+1 + unsafe { + ( + list_offsets.get_unchecked(ix).as_usize(), + list_offsets.get_unchecked(ix + 1).as_usize(), + ) + } + }; array_data.try_extend(0, start, end)?; new_offsets.push(OffsetType::Native::from_usize(array_data.len()).unwrap()); } @@ -753,8 +781,17 @@ where // SAFETY: `i` comes from validity bitmap over `indices`, so in-bounds. let ix = unsafe { indices.value_unchecked(i) }.as_usize(); - let start = list_offsets[ix].as_usize(); - let end = list_offsets[ix + 1].as_usize(); + let (start, end) = if VALIDATE { + (list_offsets[ix].as_usize(), list_offsets[ix + 1].as_usize()) + } else { + // SAFETY: caller guarantees all valid indices are in-bounds; list_offsets has len values.len()+1 + unsafe { + ( + list_offsets.get_unchecked(ix).as_usize(), + list_offsets.get_unchecked(ix + 1).as_usize(), + ) + } + }; array_data.try_extend(0, start, end)?; new_offsets.push(OffsetType::Native::from_usize(array_data.len()).unwrap()); last_filled = i + 1;