Jefffrey commented on code in PR #10441:
URL: https://github.com/apache/arrow-rs/pull/10441#discussion_r3867996416
##########
arrow-select/src/take.rs:
##########
@@ -766,31 +766,87 @@ fn take_fixed_size_list<IndexType: ArrowPrimitiveType>(
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 field = match values.data_type() {
+ DataType::FixedSizeList(field, _) => field.clone(),
+ d => unreachable!("take_fixed_size_list called with
non-fixed-size-list data type {d}"),
+ };
- // determine null count and null buffer, which are a function of `values`
and `indices`
- let num_bytes = bit_util::ceil(indices.len(), 8);
- let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes,
true);
- let null_slice = null_buf.as_slice_mut();
+ let child = values.values();
+ let nulls = take_nulls(values.nulls(), indices);
- for i in 0..indices.len() {
- let index = indices
- .value(i)
- .to_usize()
- .ok_or_else(|| ArrowError::ComputeError("Cast to usize
failed".to_string()))?;
- if !indices.is_valid(i) || values.is_null(index) {
- bit_util::unset_bit(null_slice, i);
+ // Fast path: primitive child with no nulls copy row-sized byte blocks
directly,
+ let taken_child = if child.null_count() == 0 {
Review Comment:
```rust
// Fast path: primitive child with no nulls copy row-sized byte blocks
directly,
let taken_child = if child.null_count() == 0
&& let Some(element_size) = child.data_type().primitive_width()
{
take_fixed_size_list_primitive(child, indices, length as usize,
element_size)
} else {
let list_indices = take_value_indices_from_fixed_size_list(values,
indices, length)?;
take_impl::<UInt32Type>(child.as_ref(), &list_indices)?
};
```
##########
arrow-select/src/take.rs:
##########
@@ -766,31 +766,87 @@ fn take_fixed_size_list<IndexType: ArrowPrimitiveType>(
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 field = match values.data_type() {
+ DataType::FixedSizeList(field, _) => field.clone(),
+ d => unreachable!("take_fixed_size_list called with
non-fixed-size-list data type {d}"),
+ };
- // determine null count and null buffer, which are a function of `values`
and `indices`
- let num_bytes = bit_util::ceil(indices.len(), 8);
- let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes,
true);
- let null_slice = null_buf.as_slice_mut();
+ let child = values.values();
+ let nulls = take_nulls(values.nulls(), indices);
- for i in 0..indices.len() {
- let index = indices
- .value(i)
- .to_usize()
- .ok_or_else(|| ArrowError::ComputeError("Cast to usize
failed".to_string()))?;
- if !indices.is_valid(i) || values.is_null(index) {
- bit_util::unset_bit(null_slice, i);
+ // Fast path: primitive child with no nulls copy row-sized byte blocks
directly,
+ let taken_child = if child.null_count() == 0 {
+ if let Some(element_size) = child.data_type().primitive_width() {
+ take_fixed_size_list_primitive(child, indices, length as usize,
element_size)
+ } else {
+ let list_indices = take_value_indices_from_fixed_size_list(values,
indices, length)?;
+ take_impl::<UInt32Type>(child.as_ref(), &list_indices)?
}
- }
+ } else {
+ let list_indices = take_value_indices_from_fixed_size_list(values,
indices, length)?;
+ take_impl::<UInt32Type>(child.as_ref(), &list_indices)?
+ };
- let field = match values.data_type() {
- DataType::FixedSizeList(field, _) => field.clone(),
- d => unreachable!("take_fixed_size_list called with
non-fixed-size-list data type {d}"),
+ FixedSizeListArray::try_new(field, length as i32, taken_child, nulls)
+}
+
+#[inline(never)]
+fn take_fixed_size_list_primitive<IndexType: ArrowPrimitiveType>(
+ child: &ArrayRef,
+ indices: &PrimitiveArray<IndexType>,
+ list_size: usize,
+ element_size: usize,
+) -> ArrayRef {
+ let row_bytes = list_size * element_size;
+ let child_data = child.to_data();
+ let src = child_data.buffers()[0].as_slice();
+ let child_byte_offset = child_data.offset() * element_size;
+
+ debug_assert!(
+ indices.len().checked_mul(row_bytes).is_some(),
+ "take_fixed_size_list_primitive: output buffer size overflows usize"
+ );
+ let out_len = indices.len() * list_size;
+
+ let mut out = MutableBuffer::from_len_zeroed(indices.len() * row_bytes);
+ let out_slice = out.as_slice_mut();
+
+ let child_null_buf = if indices.null_count() == 0 {
+ for (out_row, index) in indices.values().iter().enumerate() {
+ let src_start = child_byte_offset + index.as_usize() * row_bytes;
+ out_slice[out_row * row_bytes..(out_row + 1) * row_bytes]
+ .copy_from_slice(&src[src_start..src_start + row_bytes]);
+ }
+ None
+ } else {
+ let mut null_buf =
MutableBuffer::from_len_zeroed(bit_util::ceil(out_len, 8));
+ let null_slice = null_buf.as_slice_mut();
+ for (out_row, index) in indices.values().iter().enumerate() {
+ if indices.is_valid(out_row) {
+ let src_start = child_byte_offset + index.as_usize() *
row_bytes;
+ out_slice[out_row * row_bytes..(out_row + 1) * row_bytes]
+ .copy_from_slice(&src[src_start..src_start + row_bytes]);
+ for j in 0..list_size {
+ bit_util::set_bit(null_slice, out_row * list_size + j);
Review Comment:
for null buffer, is it valid to grab the nulls buffer of the fixedsizelist
array after the take (`nulls` in `take_fixed_size_list()`) then use
[`expand`](https://docs.rs/arrow/latest/arrow/buffer/struct.NullBuffer.html#method.expand)
to calculate it, instead of doing it in this loop?
e.g.
```rust
#[inline(never)]
fn take_fixed_size_list_primitive<IndexType: ArrowPrimitiveType>(
child: &ArrayRef,
indices: &PrimitiveArray<IndexType>,
list_size: usize,
element_size: usize,
parent_nulls: Option<&NullBuffer>,
) -> ArrayRef {
let row_bytes = list_size * element_size;
...
} else {
for (out_row, index) in indices.values().iter().enumerate() {
if indices.is_valid(out_row) {
let src_start = child_byte_offset + index.as_usize() *
row_bytes;
out_slice[out_row * row_bytes..(out_row + 1) * row_bytes]
.copy_from_slice(&src[src_start..src_start + row_bytes]);
}
}
if let Some(a) = parent_nulls {
Some(a.expand(list_size).buffer().clone())
} else {
None
}
};
...
```
##########
arrow-select/src/take.rs:
##########
@@ -766,31 +766,87 @@ fn take_fixed_size_list<IndexType: ArrowPrimitiveType>(
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 field = match values.data_type() {
+ DataType::FixedSizeList(field, _) => field.clone(),
+ d => unreachable!("take_fixed_size_list called with
non-fixed-size-list data type {d}"),
+ };
- // determine null count and null buffer, which are a function of `values`
and `indices`
- let num_bytes = bit_util::ceil(indices.len(), 8);
- let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes,
true);
- let null_slice = null_buf.as_slice_mut();
+ let child = values.values();
+ let nulls = take_nulls(values.nulls(), indices);
- for i in 0..indices.len() {
- let index = indices
- .value(i)
- .to_usize()
- .ok_or_else(|| ArrowError::ComputeError("Cast to usize
failed".to_string()))?;
- if !indices.is_valid(i) || values.is_null(index) {
- bit_util::unset_bit(null_slice, i);
+ // Fast path: primitive child with no nulls copy row-sized byte blocks
directly,
+ let taken_child = if child.null_count() == 0 {
+ if let Some(element_size) = child.data_type().primitive_width() {
+ take_fixed_size_list_primitive(child, indices, length as usize,
element_size)
+ } else {
+ let list_indices = take_value_indices_from_fixed_size_list(values,
indices, length)?;
+ take_impl::<UInt32Type>(child.as_ref(), &list_indices)?
}
- }
+ } else {
+ let list_indices = take_value_indices_from_fixed_size_list(values,
indices, length)?;
+ take_impl::<UInt32Type>(child.as_ref(), &list_indices)?
+ };
- let field = match values.data_type() {
- DataType::FixedSizeList(field, _) => field.clone(),
- d => unreachable!("take_fixed_size_list called with
non-fixed-size-list data type {d}"),
+ FixedSizeListArray::try_new(field, length as i32, taken_child, nulls)
+}
+
+#[inline(never)]
+fn take_fixed_size_list_primitive<IndexType: ArrowPrimitiveType>(
+ child: &ArrayRef,
+ indices: &PrimitiveArray<IndexType>,
+ list_size: usize,
+ element_size: usize,
+) -> ArrayRef {
+ let row_bytes = list_size * element_size;
+ let child_data = child.to_data();
+ let src = child_data.buffers()[0].as_slice();
+ let child_byte_offset = child_data.offset() * element_size;
+
+ debug_assert!(
+ indices.len().checked_mul(row_bytes).is_some(),
+ "take_fixed_size_list_primitive: output buffer size overflows usize"
+ );
+ let out_len = indices.len() * list_size;
+
+ let mut out = MutableBuffer::from_len_zeroed(indices.len() * row_bytes);
+ let out_slice = out.as_slice_mut();
+
+ let child_null_buf = if indices.null_count() == 0 {
+ for (out_row, index) in indices.values().iter().enumerate() {
+ let src_start = child_byte_offset + index.as_usize() * row_bytes;
+ out_slice[out_row * row_bytes..(out_row + 1) * row_bytes]
+ .copy_from_slice(&src[src_start..src_start + row_bytes]);
+ }
+ None
+ } else {
+ let mut null_buf =
MutableBuffer::from_len_zeroed(bit_util::ceil(out_len, 8));
+ let null_slice = null_buf.as_slice_mut();
+ for (out_row, index) in indices.values().iter().enumerate() {
+ if indices.is_valid(out_row) {
+ let src_start = child_byte_offset + index.as_usize() *
row_bytes;
+ out_slice[out_row * row_bytes..(out_row + 1) * row_bytes]
+ .copy_from_slice(&src[src_start..src_start + row_bytes]);
+ for j in 0..list_size {
+ bit_util::set_bit(null_slice, out_row * list_size + j);
+ }
+ }
+ }
+ Some(null_buf.into())
};
- let nulls = NullBuffer::from_unsliced_buffer(null_buf, indices.len());
- FixedSizeListArray::try_new(field, length as i32, taken, nulls)
+ // SAFETY:
+ // - The data buffer has `indices.len() * row_bytes` bytes = `out_len`
elements of a
+ // primitive type, matching `.len(out_len)`.
+ // - The null buffer (when present) is `ceil(out_len, 8)` bytes via
`from_len_zeroed`.
+ // - The data type is primitive: exactly one data buffer, no children
required.
+ let array_data = unsafe {
+ ArrayData::builder(child.data_type().clone())
Review Comment:
i suppose this is better than needing to monomorphize for all primitive
types 👍
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
