diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index 772f21951ef8..9fd9c97dd0b4 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -97,7 +97,24 @@ 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 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, +) -> 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 +228,7 @@ where } #[inline(never)] -fn take_impl( +fn take_impl( values: &dyn Array, indices: &PrimitiveArray, ) -> Result { @@ -219,31 +236,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 @@ -258,7 +275,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( @@ -274,7 +291,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(); @@ -299,7 +316,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! { @@ -307,20 +324,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. @@ -336,10 +353,10 @@ 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(values, indices)?; + let values = take_impl::<_,VALIDATE>(values, indices)?; children.push(values); } let array = UnionArray::try_new(fields.clone(), type_ids, None, children)?; @@ -348,8 +365,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, _)| { @@ -359,7 +376,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::>()?; @@ -404,7 +421,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> @@ -412,19 +429,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(), @@ -432,7 +449,7 @@ fn take_nulls( } #[inline(never)] -fn take_native( +fn take_native( values: &[T], indices: &PrimitiveArray, ) -> ScalarBuffer { @@ -441,25 +458,44 @@ 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(), } } #[inline(never)] -fn take_bits( +fn take_bits( values: &BooleanBuffer, indices: &PrimitiveArray, ) -> BooleanBuffer { @@ -470,8 +506,17 @@ 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) }; } @@ -481,24 +526,31 @@ fn take_bits( None => { BooleanBuffer::collect_bool(len, |idx: usize| { // SAFETY: idx( +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> { @@ -508,7 +560,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) { @@ -516,8 +568,21 @@ fn take_bytes( 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(); + + let (start, end) = if VALIDATE { + ( + 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( T::Offset::from_usize(capacity) @@ -566,11 +631,23 @@ fn take_bytes( if last_filled < i { offsets[last_filled + 1..=i].fill(current_offset); } - // SAFETY: `i` comes from a validity bitmap over `indices`, so it is in-bounds. 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) = if VALIDATE { + ( + 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[i + 1] = T::Offset::from_usize(capacity) .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?; @@ -629,12 +706,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) @@ -645,7 +722,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> @@ -657,7 +734,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()); @@ -676,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()); } @@ -695,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; @@ -728,7 +823,7 @@ where GenericListArray::::try_new(field, offsets, child, nulls) } -fn take_list_view( +fn take_list_view( values: &GenericListViewArray, indices: &PrimitiveArray, ) -> Result, ArrowError> @@ -737,9 +832,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(), @@ -770,7 +865,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); @@ -801,7 +896,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, @@ -811,15 +906,15 @@ fn take_fixed_size_binary( })?; 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), }; - 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); @@ -875,7 +970,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 { @@ -911,7 +1006,16 @@ 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| { + 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::>(), }; @@ -931,11 +1035,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()) }) } @@ -1125,7 +1229,7 @@ pub fn take_record_batch( let columns = record_batch .columns() .iter() - .map(|c| take(c, indices, None)) + .map(|c| unsafe { take_unchecked(c, indices) }) .collect::, _>>()?; RecordBatch::try_new(record_batch.schema(), columns) } @@ -2252,7 +2356,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!( @@ -2281,7 +2385,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!( @@ -2900,7 +3004,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);