This is an automated email from the ASF dual-hosted git repository.
etseidl 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 ef94f4fa20 add const generic to take kernel (#10820)
ef94f4fa20 is described below
commit ef94f4fa200dc7f821743240dbd912892929ecb9
Author: RIchard Baah <[email protected]>
AuthorDate: Thu Aug 27 16:13:30 2026 -0400
add const generic to take kernel (#10820)
# Which issue does this PR close?
- first step in closing #8879.
- this PR is blocking #10813
# Rationale for this change
Several hot paths in take perform redundant bounds checks on indices
that have already been validated upfront via `check_bounds`. Threading a
compile-time const CHECKED: bool through the call chain gives later PRs
a zero-cost hook to switch to unchecked accessors (e.g.
`value_unchecked`) when bounds are guaranteed, without changing the
public API or default behavior.
Why is this zero-cost (compute wise, not in terms of binary size)? :
https://rustc-dev-guide.rust-lang.org/backend/monomorph.html
# What changes are included in this PR?
- Adds const CHECKED: bool to `take_impl`, `take_primitive`,
`take_nulls`, `take_bits`, `take_boolean`, `take_bytes`,
`take_byte_view`, `take_list`, `take_list_view`, `take_fixed_size_list`,
`take_fixed_size_binary`, and `take_dict`
- All callers thread the parameter through; the public take function
always passes true for now
- benchmarks show no regressions across all take kernel variants
# Are these changes tested?
nothing new to test.
n/a
# Are there any user-facing changes?
no
---
arrow-select/src/take.rs | 98 ++++++++++++++++++++++++------------------------
1 file changed, 50 insertions(+), 48 deletions(-)
diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs
index b66dbf5463..9ddbf38083 100644
--- a/arrow-select/src/take.rs
+++ b/arrow-select/src/take.rs
@@ -98,7 +98,7 @@ 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:?}")))
)
@@ -209,7 +209,7 @@ where
}
#[inline(never)]
-fn take_impl<IndexType: ArrowPrimitiveType>(
+fn take_impl<IndexType: ArrowPrimitiveType, const CHECKED: bool>(
values: &dyn Array,
indices: &PrimitiveArray<IndexType>,
) -> Result<ArrayRef, ArrowError> {
@@ -217,38 +217,38 @@ fn take_impl<IndexType: ArrowPrimitiveType>(
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::<_, _, CHECKED>(values,
indices)?)),
DataType::Boolean => {
let values =
values.as_any().downcast_ref::<BooleanArray>().unwrap();
- Ok(Arc::new(take_boolean(values, indices)))
+ Ok(Arc::new(take_boolean::<_, CHECKED>(values, indices)))
}
DataType::Utf8 => {
- Ok(Arc::new(take_bytes(values.as_string::<i32>(), indices)?))
+ Ok(Arc::new(take_bytes::<_, _, CHECKED>(values.as_string::<i32>(),
indices)?))
}
DataType::LargeUtf8 => {
- Ok(Arc::new(take_bytes(values.as_string::<i64>(), indices)?))
+ Ok(Arc::new(take_bytes::<_, _, CHECKED>(values.as_string::<i64>(),
indices)?))
}
DataType::Utf8View => {
- Ok(Arc::new(take_byte_view(values.as_string_view(), indices)?))
+ Ok(Arc::new(take_byte_view::<_, _,
CHECKED>(values.as_string_view(), indices)?))
}
DataType::List(_) => {
- Ok(Arc::new(take_list::<_, Int32Type>(values.as_list(), indices)?))
+ Ok(Arc::new(take_list::<_, Int32Type, CHECKED>(values.as_list(),
indices)?))
}
DataType::LargeList(_) => {
- Ok(Arc::new(take_list::<_, Int64Type>(values.as_list(), indices)?))
+ Ok(Arc::new(take_list::<_, Int64Type, CHECKED>(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,
CHECKED>(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,
CHECKED>(values.as_list_view(), indices)?))
}
DataType::FixedSizeList(_, length) => {
let values = values
.as_any()
.downcast_ref::<FixedSizeListArray>()
.unwrap();
- Ok(Arc::new(take_fixed_size_list(
+ Ok(Arc::new(take_fixed_size_list::<_, CHECKED>(
values,
indices,
*length as u32,
@@ -256,7 +256,7 @@ fn take_impl<IndexType: ArrowPrimitiveType>(
}
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, CHECKED>(&list_arr,
indices)?;
let (_, offsets, entries, nulls) = list_data.into_parts();
let entries = entries.as_struct().clone();
Ok(Arc::new(MapArray::try_new(
@@ -272,7 +272,7 @@ fn take_impl<IndexType: ArrowPrimitiveType>(
let arrays = array
.columns()
.iter()
- .map(|a| take_impl(a.as_ref(), indices))
+ .map(|a| take_impl::<_, CHECKED>(a.as_ref(), indices))
.collect::<Result<Vec<ArrayRef>, _>>()?;
let fields: Vec<(FieldRef, ArrayRef)> =
fields.iter().cloned().zip(arrays).collect();
@@ -297,7 +297,7 @@ fn take_impl<IndexType: ArrowPrimitiveType>(
}
}
DataType::Dictionary(_, _) => downcast_dictionary_array! {
- values => Ok(Arc::new(take_dict(values, indices)?)),
+ values => Ok(Arc::new(take_dict::<_, _, CHECKED>(values,
indices)?)),
t => unimplemented!("Take not supported for dictionary type {:?}",
t)
}
DataType::RunEndEncoded(_, _) => downcast_run_array! {
@@ -305,20 +305,20 @@ fn take_impl<IndexType: ArrowPrimitiveType>(
t => unimplemented!("Take not supported for run type {:?}", t)
}
DataType::Binary => {
- Ok(Arc::new(take_bytes(values.as_binary::<i32>(), indices)?))
+ Ok(Arc::new(take_bytes::<_, _, CHECKED>(values.as_binary::<i32>(),
indices)?))
}
DataType::LargeBinary => {
- Ok(Arc::new(take_bytes(values.as_binary::<i64>(), indices)?))
+ Ok(Arc::new(take_bytes::<_, _, CHECKED>(values.as_binary::<i64>(),
indices)?))
}
DataType::BinaryView => {
- Ok(Arc::new(take_byte_view(values.as_binary_view(), indices)?))
+ Ok(Arc::new(take_byte_view::<_, _,
CHECKED>(values.as_binary_view(), indices)?))
}
DataType::FixedSizeBinary(size) => {
let values = values
.as_any()
.downcast_ref::<FixedSizeBinaryArray>()
.unwrap();
- Ok(Arc::new(take_fixed_size_binary(values, indices, *size)?))
+ Ok(Arc::new(take_fixed_size_binary::<_, CHECKED>(values, indices,
*size)?))
}
DataType::Null => {
// Take applied to a null array produces a null array.
@@ -337,7 +337,7 @@ fn take_impl<IndexType: ArrowPrimitiveType>(
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::<_, CHECKED>(values, indices)?;
children.push(values);
}
let array = UnionArray::try_new(fields.clone(), type_ids, None,
children)?;
@@ -357,7 +357,7 @@ fn take_impl<IndexType: ArrowPrimitiveType>(
let values = values.child(field_type_id);
- take_impl(values, indices.as_primitive::<Int32Type>())
+ take_impl::<_, CHECKED>(values,
indices.as_primitive::<Int32Type>())
})
.collect::<Result<_, _>>()?;
@@ -402,7 +402,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<T, I>(
+fn take_primitive<T, I, const CHECKED: bool>(
values: &PrimitiveArray<T>,
indices: &PrimitiveArray<I>,
) -> Result<PrimitiveArray<T>, ArrowError>
@@ -411,18 +411,18 @@ where
I: ArrowPrimitiveType,
{
let values_buf = take_native(values.values(), indices);
- let nulls = take_nulls(values.nulls(), indices);
+ let nulls = take_nulls::<_, CHECKED>(values.nulls(), indices);
Ok(PrimitiveArray::try_new(values_buf,
nulls)?.with_data_type(values.data_type().clone()))
}
#[inline(never)]
-fn take_nulls<I: ArrowPrimitiveType>(
+fn take_nulls<I: ArrowPrimitiveType, const CHECKED: bool>(
values: Option<&NullBuffer>,
indices: &PrimitiveArray<I>,
) -> Option<NullBuffer> {
match values.filter(|n| n.null_count() > 0) {
Some(n) => NullBuffer::from_unsliced_buffer(
- take_bits(n.inner(), indices).into_inner(),
+ take_bits::<_, CHECKED>(n.inner(), indices).into_inner(),
indices.len(),
),
None => indices.nulls().cloned(),
@@ -457,7 +457,7 @@ fn take_native<T: ArrowNativeType, I: ArrowPrimitiveType>(
}
#[inline(never)]
-fn take_bits<I: ArrowPrimitiveType>(
+fn take_bits<I: ArrowPrimitiveType, const CHECKED: bool>(
values: &BooleanBuffer,
indices: &PrimitiveArray<I>,
) -> BooleanBuffer {
@@ -469,7 +469,7 @@ fn take_bits<I: ArrowPrimitiveType>(
let output_slice = output_buffer.as_slice_mut();
nulls.valid_indices().for_each(|idx| {
// SAFETY: idx is a valid index in indices.nulls() -->
idx<indices.len()
- if values.value(unsafe {
indices.value_unchecked(idx).as_usize() }) {
+ if unsafe {
values.value(indices.value_unchecked(idx).as_usize()) } {
// 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) };
}
@@ -486,17 +486,17 @@ fn take_bits<I: ArrowPrimitiveType>(
}
/// `take` implementation for boolean arrays
-fn take_boolean<IndexType: ArrowPrimitiveType>(
+fn take_boolean<IndexType: ArrowPrimitiveType, const CHECKED: bool>(
values: &BooleanArray,
indices: &PrimitiveArray<IndexType>,
) -> BooleanArray {
- let val_buf = take_bits(values.values(), indices);
- let null_buf = take_nulls(values.nulls(), indices);
+ let val_buf = take_bits::<_, CHECKED>(values.values(), indices);
+ let null_buf = take_nulls::<_, CHECKED>(values.nulls(), indices);
BooleanArray::new(val_buf, null_buf)
}
/// `take` implementation for string arrays
-fn take_bytes<T: ByteArrayType, IndexType: ArrowPrimitiveType>(
+fn take_bytes<T: ByteArrayType, IndexType: ArrowPrimitiveType, const CHECKED:
bool>(
array: &GenericByteArray<T>,
indices: &PrimitiveArray<IndexType>,
) -> Result<GenericByteArray<T>, ArrowError> {
@@ -506,7 +506,7 @@ fn take_bytes<T: ByteArrayType, IndexType:
ArrowPrimitiveType>(
let input_offsets = array.value_offsets();
let mut capacity = 0;
- let nulls = take_nulls(array.nulls(), indices);
+ let nulls = take_nulls::<_, CHECKED>(array.nulls(), indices);
// Branch on output nulls — `None` means every output slot is valid.
match nulls.as_ref().filter(|n| n.null_count() > 0) {
@@ -627,12 +627,12 @@ fn take_bytes<T: ByteArrayType, IndexType:
ArrowPrimitiveType>(
}
/// `take` implementation for byte view arrays
-fn take_byte_view<T: ByteViewType, IndexType: ArrowPrimitiveType>(
+fn take_byte_view<T: ByteViewType, IndexType: ArrowPrimitiveType, const
CHECKED: bool>(
array: &GenericByteViewArray<T>,
indices: &PrimitiveArray<IndexType>,
) -> Result<GenericByteViewArray<T>, ArrowError> {
let new_views = take_native(array.views(), indices);
- let new_nulls = take_nulls(array.nulls(), indices);
+ let new_nulls = take_nulls::<_, CHECKED>(array.nulls(), indices);
let buffers = Arc::clone(array.data_buffers());
// Safety: array.views was valid, and take_native copies only valid
values, and verifies bounds
Ok(unsafe { GenericByteViewArray::new_unchecked(new_views, buffers,
new_nulls) })
@@ -642,7 +642,7 @@ fn take_byte_view<T: ByteViewType, IndexType:
ArrowPrimitiveType>(
///
/// 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<IndexType, OffsetType>(
+fn take_list<IndexType, OffsetType, const CHECKED: bool>(
values: &GenericListArray<OffsetType::Native>,
indices: &PrimitiveArray<IndexType>,
) -> Result<GenericListArray<OffsetType::Native>, ArrowError>
@@ -654,7 +654,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::<_, CHECKED>(values.nulls(), indices);
let mut new_offsets = Vec::with_capacity(indices.len() + 1);
new_offsets.push(OffsetType::Native::zero());
@@ -725,7 +725,7 @@ where
GenericListArray::<OffsetType::Native>::try_new(field, offsets, child,
nulls)
}
-fn take_list_view<IndexType, OffsetType>(
+fn take_list_view<IndexType, OffsetType, const CHECKED: bool>(
values: &GenericListViewArray<OffsetType::Native>,
indices: &PrimitiveArray<IndexType>,
) -> Result<GenericListViewArray<OffsetType::Native>, ArrowError>
@@ -736,7 +736,7 @@ where
{
let taken_offsets = take_native(values.offsets(), indices);
let taken_sizes = take_native(values.sizes(), indices);
- let nulls = take_nulls(values.nulls(), indices);
+ let nulls = take_nulls::<_, CHECKED>(values.nulls(), indices);
let field = match values.data_type() {
DataType::ListView(field) | DataType::LargeListView(field) =>
field.clone(),
@@ -761,13 +761,13 @@ where
/// Calculates the index and indexed offset for the inner array,
/// applying `take` on the inner array, then reconstructing a list array
/// with the indexed offsets
-fn take_fixed_size_list<IndexType: ArrowPrimitiveType>(
+fn take_fixed_size_list<IndexType: ArrowPrimitiveType, const CHECKED: bool>(
values: &FixedSizeListArray,
indices: &PrimitiveArray<IndexType>,
length: <UInt32Type as ArrowPrimitiveType>::Native,
) -> Result<FixedSizeListArray, ArrowError> {
let list_indices = take_value_indices_from_fixed_size_list(values,
indices, length)?;
- let taken = take_impl::<UInt32Type>(values.values().as_ref(),
&list_indices)?;
+ let taken = take_impl::<UInt32Type, CHECKED>(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);
@@ -798,7 +798,7 @@ fn take_fixed_size_list<IndexType: ArrowPrimitiveType>(
/// The computation is done in two steps:
/// - Compute the values buffer
/// - Compute the null buffer
-fn take_fixed_size_binary<IndexType: ArrowPrimitiveType>(
+fn take_fixed_size_binary<IndexType: ArrowPrimitiveType, const CHECKED: bool>(
values: &FixedSizeBinaryArray,
indices: &PrimitiveArray<IndexType>,
size: i32,
@@ -816,7 +816,7 @@ fn take_fixed_size_binary<IndexType: ArrowPrimitiveType>(
_ => take_fixed_size_binary_buffer_dynamic_length(values, indices,
size_usize),
};
- let value_nulls = take_nulls(values.nulls(), indices);
+ let value_nulls = take_nulls::<_, CHECKED>(values.nulls(), indices);
let final_nulls = NullBuffer::union(value_nulls.as_ref(), indices.nulls());
return FixedSizeBinaryArray::try_new(size, result_buffer, final_nulls);
@@ -928,11 +928,11 @@ fn take_fixed_size<IndexType: ArrowPrimitiveType, const
N: usize>(
///
/// 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<T: ArrowDictionaryKeyType, I: ArrowPrimitiveType>(
+fn take_dict<T: ArrowDictionaryKeyType, I: ArrowPrimitiveType, const CHECKED:
bool>(
values: &DictionaryArray<T>,
indices: &PrimitiveArray<I>,
) -> Result<DictionaryArray<T>, ArrowError> {
- let new_keys = take_primitive(values.keys(), indices)?;
+ let new_keys = take_primitive::<_, _, CHECKED>(values.keys(), indices)?;
Ok(unsafe { DictionaryArray::new_unchecked(new_keys,
values.values().clone()) })
}
@@ -2081,7 +2081,8 @@ mod tests {
let input_array = FixedSizeListArray::from_iter_primitive::<T, _,
_>(input_data, length);
- let output = take_fixed_size_list(&input_array, &indices, length as
u32).unwrap();
+ let output =
+ take_fixed_size_list::<_, true>(&input_array, &indices, length as
u32).unwrap();
let expected = FixedSizeListArray::from_iter_primitive::<T, _,
_>(expected_data, length);
@@ -2257,7 +2258,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!(
@@ -2286,7 +2287,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!(
@@ -2918,7 +2919,8 @@ mod tests {
let logical_indices: PrimitiveArray<Int32Type> =
PrimitiveArray::from(Vec::<i32>::new());
- let result = take_impl(&run_array, &logical_indices).expect("take_run
with empty indices");
+ let result = take_impl::<_, true>(&run_array, &logical_indices)
+ .expect("take_run with empty indices");
// Verify the result is a valid empty RunArray
assert_eq!(result.len(), 0);
