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.

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

Add unsafe/unchecked slice functions #6901

Closed
wants to merge 3 commits into from

Conversation

emilk
Copy link
Contributor

@emilk emilk commented Dec 18, 2024

Which issue does this PR close?

Rationale for this change

Fast slicing when the user has done their own bounds checking before-hand

What changes are included in this PR?

Adds:

  • `Buffer::slice_with_length_unchecked
  • ScalarBuffer::new_unchecked
  • ScalarBuffer::slice_unchecked
  • ScalarBuffer::len_in_elements

Are there any user-facing changes?

Yes, three added methods (see above)

@github-actions github-actions bot added the arrow Changes to the arrow crate label Dec 18, 2024

/// The length of the scalar buffer, in units of `T`.
#[inline]
pub fn len_in_elements(&self) -> usize {
Copy link
Contributor Author

Choose a reason for hiding this comment

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

This is so the user can do their own range checking easily, and also so that the docstring for slice_unchecked has something to refer to.

We should maybe call this function len, but I went with the more explicit, since len usually means "number of bytes" in other contexts.

Copy link
Contributor

Choose a reason for hiding this comment

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

The deref impl already means that x.len() returns this

@tustvold
Copy link
Contributor

tustvold commented Dec 18, 2024

The ticket links to https://docs.rs/arrow2/latest/src/arrow2/array/mod.rs.html#111 and states that unchecked slicing was very important.

However, this PR is changing a method that has a different signature, in particular it is &self and not &mut self. The result is that this method is already performing an atomic increment of the reference count, and so the performance gain of adding an unchecked version is not as obvious to me?

Do you have benchmarks?

Edit:

In general an over-reliance on slicing normally suggests an issue with the way a kernel has been implemented

e.g.

Bad

fn sum_values(a: ListArray<i32>) -> i32 {
    // Iterator creates a new slice for each element
    a.iter().flat_map(|x| x.as_primitive::<Int32Type>().iter().flatten().sum()).sum();
}

Better

fn sum_values(a: ListArray<i32>) -> i32 {
    let offsets = a.offsets();
    a.as_primitive::<Int32Type>().iter()
        .take(offsets.last().unwrap().as_usize())
        .skip(offsets.first().unwrap().as_usize())
        .flatten(|x| x)
        .sum()
}

Better Still

fn sum_values(a: ListArray<i32>) -> i32 {
    let offsets = a.offsets();
    let start= offsets.first().unwrap().as_usize();
    let end = offsets.last().unwrap().as_usize();

    arrow_arith::sum(a.values().as_primitive::<Int32Type>().slice(start, end - start))
}

@emilk
Copy link
Contributor Author

emilk commented Dec 19, 2024

I'll try to get you some benchmarks.

The main thing we want to avoid is doing the same bounds checks twice: first before calling into slice, and then once more inside of slice. One can hope that the optimizer can remove the second bounds check, but that will only happen if the inliner is good enough (and these functions aren't even marked as #[inline]).

After making this PR I started thinking though - and alternative is to add an API like this:

fn try_slice(&self, offset: usize, len: usize) -> Option<Self>

that returns None in case of out-of-bounds.

That way the bounds-check is only done once, though lack of inlining may still hurt perf.

And yes, we have a weird usecase. We are generating code for arrow-deserialization, and the code is not pretty 😓

I'm looking deeper now to see if we can avoid all of this… converting to draft in the meantime.

@emilk emilk marked this pull request as draft December 19, 2024 09:21
@emilk
Copy link
Contributor Author

emilk commented Dec 19, 2024

After some internal review it's become clear that this ugly code-generated arrow-derserialization we do is never on the fast path (anymore), so we don't need this after all. Sorry for the noise!

@emilk emilk closed this Dec 19, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
arrow Changes to the arrow crate
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Add unsafe/unchecked slicing
2 participants