feat: add OffsetBuffer::subtract to allow to shift offsets by value

[rluvaton]

Which issue does this PR close?

N/A

Rationale for this change

When dealing with sliced variable size arrays it is common to subtract the offset to shift them to 0, so this is a helper for that

What changes are included in this PR?

added OffsetBuffer::subtract method

Are these changes tested?

yes

Are there any user-facing changes?

new function

[rluvaton]

this is sealed since ArrowNativeType is sealed, so it is safe to add this here

[rluvaton]

After pr #10118 is merged we should replace this test with the following test that verify reuse

(this test need some cleanup but the idea is the same)

#[test]
fn for_sliced_unshared_buffer_shift_should_reuse_buffer_but_only_modify_the_data_in_slice() {
    // Underlying [0, 3, 6, 9, 12]; slice -> view [3, 6, 9].
    let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![1, 3, 6, 9, 12]));
    let sliced = offsets.slice(1, 2); // [3, 6, 9]
    drop(offsets); // uniquely owned -> eligible for in-place reuse
    assert_eq!(sliced.as_ref(), &[3, 6, 9]);

    let ptr_before = sliced.as_ptr();
    let result = sliced.subtract(1);

    // Reuses the allocation...
    assert_eq!(
        ptr_before,
        result.as_ptr(),
        "uniquely-owned sliced buffer should be mutated in place"
    );

    // ...but only the slice is shifted, and the length stays the slice length.
    assert_eq!(result.as_ref(), &[2, 5, 8]);
    assert_eq!(result.len(), 3);
    let underlying_buffer = result.inner().inner();
    // data should be shifted by 1
    assert_eq!(underlying_buffer.ptr_offset(), std::mem::size_of::<i32>());
    let original_value_ptr = unsafe {underlying_buffer.data_ptr().as_ptr() as *mut i32};
    let value_in_ptr = unsafe {*original_value_ptr};
    assert_eq!(value_in_ptr, 1);
    let last_value_ptr = unsafe {(underlying_buffer.data_ptr().as_ptr() as *mut i32).add(5)};
    assert_eq!(unsafe {*last_value_ptr}, 12);
}

[alamb]

Thanks @rluvaton

The PR says

When dealing with sliced variable size arrays it is common to subtract the offset to shift them to 0, so this is a helper for that

But this PR doesn't have any uses of this API - are there places in the arrow crate that subtract the offsets that we could update to use this API (and thus get additional test coverage)?

[rluvaton]

@alamb

In datafusion code:

• https://github.com/apache/datafusion/blob/b8998c762bb864a9f3607a518384b03dcf40eb61/datafusion/common/src/utils/mod.rs#L1231-L1248
• https://github.com/apache/datafusion/blob/b8998c762bb864a9f3607a518384b03dcf40eb61/datafusion/functions-nested/src/sort.rs#L468-L480
• https://github.com/apache/datafusion/blob/b8998c762bb864a9f3607a518384b03dcf40eb61/datafusion/functions/src/string/common.rs#L601-L609
• https://github.com/apache/datafusion/blob/b8998c762bb864a9f3607a518384b03dcf40eb61/datafusion/functions/src/unicode/initcap.rs#L220-L230

and many more in my project as well

[rluvaton]

And I want to optimize cast to only cast visible list values which would require me to shift the offsets

there are many more cases, but did not find in arrow uses that still keep the offset buffer

[alamb]

Thank you for this PR @rluvaton

I think the API is good, but I have some suggestions to simplify the implementation

[alamb]

I don't follow the issue described in #10117

I also think we can accomplish the same thing using into_vec which is less code and possibly faster (as Vec is crazy fast):

    /// Subtract `rhs` from all offsets
    /// This will try to reuse the existing allocation as much as possible
    ///
    /// Panics: this will panic if `rhs` > the first offset or if `rhs` will lead to overflow (when `rhs` is negative)
    pub fn subtract(self, rhs: O) -> Self
    where
        O: std::ops::Sub<Output = O> + std::cmp::PartialOrd + num_traits::CheckedSub,
    {
        if rhs == O::usize_as(0) {
            return self;
        }

        let len = self.len();

        // Offset buffer is guaranteed to be non-empty
        assert!(
            self[0] >= rhs,
            "shifted offsets will become negative which is not allowed"
        );

        // If negative, make sure that this will not create an overflow
        if rhs < O::usize_as(0) {
            self[len - 1].checked_sub(&rhs).expect("must not overflow");
        }

        // try and reuse buffer:
        let shifted_offsets: Vec<O> = match self.into_inner().into_inner().into_vec() {
            // If we can reuse the buffer, update in place
            Ok(mut v) => {
                for offset in v.iter_mut() {
                    *offset = *offset - rhs;
                }
                v
            }
            // otherwise, buffer is shared so we need a copy
            Err(buffer) => {
                let offsets = ScalarBuffer::<O>::from(buffer);
                offsets.iter().map(|offset| *offset - rhs).collect()
            }
        };
        let shifted_buffer = ScalarBuffer::from(shifted_offsets);
        // Safety: offsets are valid as they are coming from a valid
        // offset buffer and we checked overflow above, and we
        // subtracted the same value from all offsets, thus keeping the
        // same properties as the input buffer
        unsafe { Self::new_unchecked(shifted_buffer) }
    }

[rluvaton]

I was not aware that into vec reused the buffer, thanks!

[rluvaton]

done, thanks

[alamb]

Thanks @rluvaton -- this looks good to me

[rluvaton]

Thanks, can you please merge