AdamGS commented on code in PR #9643:
URL: https://github.com/apache/arrow-rs/pull/9643#discussion_r3093312262
##########
arrow-select/src/take.rs:
##########
@@ -624,23 +623,65 @@ where
OffsetType::Native: OffsetSizeTrait,
PrimitiveArray<OffsetType>: From<Vec<OffsetType::Native>>,
{
- // TODO: Some optimizations can be done here such as if it is
- // taking the whole list or a contiguous sublist
- let (list_indices, offsets, null_buf) =
- take_value_indices_from_list::<IndexType, OffsetType>(values,
indices)?;
-
- let taken = take_impl::<OffsetType>(values.values().as_ref(),
&list_indices)?;
- let value_offsets = Buffer::from_vec(offsets);
- // create a new list with taken data and computed null information
+ let list_offsets = values.value_offsets();
+ let child_data = values.values().to_data();
+ let nulls = take_nulls(values.nulls(), indices);
+
+ let mut new_offsets = Vec::with_capacity(indices.len() + 1);
+ new_offsets.push(OffsetType::Native::zero());
+
+ let use_nulls = child_data.null_count() > 0;
+
+ let capacity = child_data
+ .len()
+ .checked_div(values.len())
+ .map(|v| v * indices.len())
+ .unwrap_or_default();
+
+ let mut array_data = MutableArrayData::new(vec![&child_data], use_nulls,
capacity);
+
+ match nulls.as_ref().filter(|n| n.null_count() > 0) {
+ 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();
+ array_data.extend(0, start, end);
+
new_offsets.push(OffsetType::Native::from_usize(array_data.len()).unwrap());
+ }
+ }
+ Some(output_nulls) => {
+ new_offsets.resize(indices.len() + 1, OffsetType::Native::zero());
Review Comment:
I assumed that this would be a nice way to avoid `unsafe` while keeping
performance, but locally perf seems to be as good
--
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]
