Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Avoid memory copies in take_list #3940

Merged
merged 1 commit into from
Mar 25, 2023
Merged

Avoid memory copies in take_list #3940

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions arrow-select/src/take.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -662,15 +662,15 @@ where
IndexType::Native: ToPrimitive,
OffsetType: ArrowPrimitiveType,
OffsetType::Native: ToPrimitive + OffsetSizeTrait,
PrimitiveArray<OffsetType>: From<Vec<Option<OffsetType::Native>>>,
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, None)?;
let value_offsets = Buffer::from_slice_ref(offsets);
let value_offsets = Buffer::from_vec(offsets);
// create a new list with taken data and computed null information
let list_data = ArrayDataBuilder::new(values.data_type().clone())
.len(indices.len())
Expand Down Expand Up @@ -881,7 +881,7 @@ where
IndexType::Native: ToPrimitive,
OffsetType: ArrowPrimitiveType,
OffsetType::Native: OffsetSizeTrait + std::ops::Add + num::Zero + num::One,
PrimitiveArray<OffsetType>: From<Vec<Option<OffsetType::Native>>>,
PrimitiveArray<OffsetType>: From<Vec<OffsetType::Native>>,
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After #3917 this conversion does not copy memory

{
// TODO: benchmark this function, there might be a faster unsafe alternative
let offsets: &[OffsetType::Native] = list.value_offsets();
Expand Down Expand Up @@ -912,7 +912,7 @@ where

// if start == end, this slot is empty
while curr < end {
values.push(Some(curr));
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks redundant previously.

values.push(curr);
curr += num::One::one();
}
if !list.is_valid(ix) {
Expand Down