arrow_buffer/buffer/
offset.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::buffer::ScalarBuffer;
19use crate::{ArrowNativeType, MutableBuffer, NullBuffer, OffsetBufferBuilder, OverflowError};
20use std::ops::Deref;
21
22/// A non-empty buffer of monotonically increasing, positive integers.
23///
24/// [`OffsetBuffer`] are used to represent ranges of offsets. An
25/// `OffsetBuffer` of `N+1` items contains `N` such ranges. The start
26/// offset for element `i` is `offsets[i]` and the end offset is
27/// `offsets[i+1]`. Equal offsets represent an empty range.
28///
29/// # Example
30///
31/// This example shows how 5 distinct ranges, are represented using a
32/// 6 entry `OffsetBuffer`. The first entry `(0, 3)` represents the
33/// three offsets `0, 1, 2`. The entry `(3,3)` represent no offsets
34/// (e.g. an empty list).
35///
36/// ```text
37///   ┌───────┐                ┌───┐
38///   │ (0,3) │                │ 0 │
39///   ├───────┤                ├───┤
40///   │ (3,3) │                │ 3 │
41///   ├───────┤                ├───┤
42///   │ (3,4) │                │ 3 │
43///   ├───────┤                ├───┤
44///   │ (4,5) │                │ 4 │
45///   ├───────┤                ├───┤
46///   │ (5,7) │                │ 5 │
47///   └───────┘                ├───┤
48///                            │ 7 │
49///                            └───┘
50///
51///                        Offsets Buffer
52///    Logical
53///    Offsets
54///
55///  (offsets[i],
56///   offsets[i+1])
57/// ```
58#[derive(Debug, Clone, PartialEq, Eq)]
59pub struct OffsetBuffer<O: ArrowNativeType>(ScalarBuffer<O>);
60
61impl<O: ArrowNativeType> OffsetBuffer<O> {
62    /// Create a new [`OffsetBuffer`] from the provided [`ScalarBuffer`]
63    ///
64    /// # Panics
65    ///
66    /// Panics if `buffer` is not a non-empty buffer containing
67    /// monotonically increasing values greater than or equal to zero
68    pub fn new(buffer: ScalarBuffer<O>) -> Self {
69        assert!(!buffer.is_empty(), "offsets cannot be empty");
70        assert!(
71            buffer[0] >= O::usize_as(0),
72            "offsets must be greater than 0"
73        );
74        assert!(
75            buffer.windows(2).all(|w| w[0] <= w[1]),
76            "offsets must be monotonically increasing"
77        );
78        Self(buffer)
79    }
80
81    /// Create a new [`OffsetBuffer`] from the provided [`ScalarBuffer`]
82    ///
83    /// # Safety
84    ///
85    /// `buffer` must be a non-empty buffer containing monotonically increasing
86    /// values greater than or equal to zero
87    pub unsafe fn new_unchecked(buffer: ScalarBuffer<O>) -> Self {
88        Self(buffer)
89    }
90
91    /// Create a new [`OffsetBuffer`] containing a single 0 value
92    pub fn new_empty() -> Self {
93        let buffer = MutableBuffer::from_len_zeroed(std::mem::size_of::<O>());
94        Self(buffer.into_buffer().into())
95    }
96
97    /// Create a new [`OffsetBuffer`] containing `len + 1` `0` values
98    ///
99    /// # Panics
100    ///
101    /// Panics if `(len + 1) * size_of::<O>()` overflows `usize`.
102    /// Use [`Self::try_new_zeroed`] for a fallible version.
103    pub fn new_zeroed(len: usize) -> Self {
104        Self::try_new_zeroed(len).unwrap_or_else(|err| panic!("{err}"))
105    }
106
107    /// Create a new [`OffsetBuffer`] containing `len + 1` `0` values
108    ///
109    /// # Errors
110    ///
111    /// Errors if `(len + 1) * size_of::<O>()` overflows `usize`
112    pub fn try_new_zeroed(len: usize) -> Result<Self, OverflowError> {
113        let len_bytes = len
114            .checked_add(1)
115            .and_then(|o| o.checked_mul(std::mem::size_of::<O>()))
116            .ok_or_else(|| OverflowError::new::<usize>("buffer length"))?;
117        let buffer = MutableBuffer::from_len_zeroed(len_bytes);
118        Ok(Self(buffer.into_buffer().into()))
119    }
120
121    /// Create a new [`OffsetBuffer`] from the iterator of slice lengths
122    ///
123    /// ```
124    /// # use arrow_buffer::OffsetBuffer;
125    /// let offsets = OffsetBuffer::<i32>::from_lengths([1, 3, 5]);
126    /// assert_eq!(offsets.as_ref(), &[0, 1, 4, 9]);
127    /// ```
128    ///
129    /// If you want to create an [`OffsetBuffer`] where all lengths are the same,
130    /// consider using the faster [`OffsetBuffer::from_repeated_length`] instead.
131    ///
132    /// # Panics
133    ///
134    /// Panics on overflow. Use [`Self::try_from_lengths`] for a fallible version.
135    pub fn from_lengths<I>(lengths: I) -> Self
136    where
137        I: IntoIterator<Item = usize>,
138    {
139        Self::try_from_lengths(lengths).unwrap_or_else(|err| panic!("{err}"))
140    }
141
142    /// Create a new [`OffsetBuffer`] from the iterator of slice lengths
143    ///
144    /// ```
145    /// # use arrow_buffer::OffsetBuffer;
146    /// let offsets = OffsetBuffer::<i32>::try_from_lengths([1, 3, 5]).unwrap();
147    /// assert_eq!(offsets.as_ref(), &[0, 1, 4, 9]);
148    /// ```
149    ///
150    /// # Errors
151    ///
152    /// Errors if the total length overflows `usize` or `O`, e.g. if the lengths
153    /// add up to more than `i32::MAX` for a `OffsetBuffer<i32>`.
154    pub fn try_from_lengths<I>(lengths: I) -> Result<Self, OverflowError>
155    where
156        I: IntoIterator<Item = usize>,
157    {
158        let iter = lengths.into_iter();
159        let mut out = Vec::with_capacity(iter.size_hint().0 + 1);
160        out.push(O::usize_as(0));
161
162        let mut acc = 0_usize;
163        for length in iter {
164            acc = acc
165                .checked_add(length)
166                .ok_or_else(|| OverflowError::new::<usize>("total length"))?;
167            out.push(O::usize_as(acc))
168        }
169        // Check for overflow
170        O::from_usize(acc).ok_or_else(|| OverflowError::new::<O>("offset").with_value(acc))?;
171        Ok(Self(out.into()))
172    }
173
174    /// Create a new [`OffsetBuffer`] where each slice has the same length
175    /// `length`, repeated `n` times.
176    ///
177    ///
178    /// Example
179    /// ```
180    /// # use arrow_buffer::OffsetBuffer;
181    /// let offsets = OffsetBuffer::<i32>::from_repeated_length(4, 3);
182    /// assert_eq!(offsets.as_ref(), &[0, 4, 8, 12]);
183    /// ```
184    ///
185    /// # Panics
186    ///
187    /// Panics on overflow. Use [`Self::try_from_repeated_length`] for a fallible version.
188    pub fn from_repeated_length(length: usize, n: usize) -> Self {
189        Self::try_from_repeated_length(length, n).unwrap_or_else(|err| panic!("{err}"))
190    }
191
192    /// Create a new [`OffsetBuffer`] where each slice has the same length
193    /// `length`, repeated `n` times.
194    ///
195    /// ```
196    /// # use arrow_buffer::OffsetBuffer;
197    /// let offsets = OffsetBuffer::<i32>::try_from_repeated_length(4, 3).unwrap();
198    /// assert_eq!(offsets.as_ref(), &[0, 4, 8, 12]);
199    /// ```
200    ///
201    /// # Errors
202    ///
203    /// Errors if `length * n` overflows `usize` or `O`.
204    pub fn try_from_repeated_length(length: usize, n: usize) -> Result<Self, OverflowError> {
205        if n == 0 {
206            return Ok(Self::new_empty());
207        }
208
209        if length == 0 {
210            return Self::try_new_zeroed(n);
211        }
212
213        // Making sure we don't overflow usize or O when calculating the total length
214        let total_length = length
215            .checked_mul(n)
216            .ok_or_else(|| OverflowError::new::<usize>("total length"))?;
217
218        O::from_usize(total_length)
219            .ok_or_else(|| OverflowError::new::<O>("offset").with_value(total_length))?;
220
221        let offsets = (0..=n)
222            .map(|index| O::usize_as(index * length))
223            .collect::<Vec<O>>();
224
225        Ok(Self(ScalarBuffer::from(offsets)))
226    }
227
228    /// The first offset, i.e. the start of the first range.
229    ///
230    /// An [`OffsetBuffer`] is never empty, so this always returns an offset.
231    ///
232    /// ```
233    /// # use arrow_buffer::OffsetBuffer;
234    /// let offsets = OffsetBuffer::<i32>::from_lengths([1, 3, 5]);
235    /// assert_eq!(offsets.first(), 0);
236    /// assert_eq!(OffsetBuffer::<i32>::new_empty().first(), 0);
237    /// ```
238    #[inline]
239    pub fn first(&self) -> O {
240        self.0
241            .first()
242            .copied()
243            .expect("An `OffsetBuffer` is never empty")
244    }
245
246    /// The last offset, i.e. the end of the last range.
247    ///
248    /// An [`OffsetBuffer`] is never empty, so this always returns an offset.
249    ///
250    /// ```
251    /// # use arrow_buffer::OffsetBuffer;
252    /// let offsets = OffsetBuffer::<i32>::from_lengths([1, 3, 5]);
253    /// assert_eq!(offsets.last(), 9);
254    /// assert_eq!(OffsetBuffer::<i32>::new_empty().last(), 0);
255    /// ```
256    #[inline]
257    pub fn last(&self) -> O {
258        self.0
259            .last()
260            .copied()
261            .expect("An `OffsetBuffer` is never empty")
262    }
263
264    /// Get an Iterator over the lengths of this [`OffsetBuffer`]
265    ///
266    /// ```
267    /// # use arrow_buffer::{OffsetBuffer, ScalarBuffer};
268    /// let offsets = OffsetBuffer::<_>::new(ScalarBuffer::<i32>::from(vec![0, 1, 4, 9]));
269    /// assert_eq!(offsets.lengths().collect::<Vec<usize>>(), vec![1, 3, 5]);
270    /// ```
271    ///
272    /// Empty [`OffsetBuffer`] will return an empty iterator
273    /// ```
274    /// # use arrow_buffer::OffsetBuffer;
275    /// let offsets = OffsetBuffer::<i32>::new_empty();
276    /// assert_eq!(offsets.lengths().count(), 0);
277    /// ```
278    ///
279    /// This can be used to merge multiple [`OffsetBuffer`]s to one
280    /// ```
281    /// # use arrow_buffer::{OffsetBuffer, ScalarBuffer};
282    ///
283    /// let buffer1 = OffsetBuffer::<i32>::from_lengths([2, 6, 3, 7, 2]);
284    /// let buffer2 = OffsetBuffer::<i32>::from_lengths([1, 3, 5, 7, 9]);
285    ///
286    /// let merged = OffsetBuffer::<i32>::from_lengths(
287    ///     vec![buffer1, buffer2].iter().flat_map(|x| x.lengths())
288    /// );
289    ///
290    /// assert_eq!(merged.lengths().collect::<Vec<_>>(), &[2, 6, 3, 7, 2, 1, 3, 5, 7, 9]);
291    /// ```
292    pub fn lengths(&self) -> impl ExactSizeIterator<Item = usize> + '_ {
293        self.0.windows(2).map(|x| x[1].as_usize() - x[0].as_usize())
294    }
295
296    /// Free up unused memory.
297    pub fn shrink_to_fit(&mut self) {
298        self.0.shrink_to_fit();
299    }
300
301    /// Returns the inner [`ScalarBuffer`]
302    pub fn inner(&self) -> &ScalarBuffer<O> {
303        &self.0
304    }
305
306    /// Returns the inner [`ScalarBuffer`], consuming self
307    pub fn into_inner(self) -> ScalarBuffer<O> {
308        self.0
309    }
310
311    /// Claim memory used by this buffer in the provided memory pool.
312    #[cfg(feature = "pool")]
313    pub fn claim(&self, pool: &dyn crate::MemoryPool) {
314        self.0.claim(pool);
315    }
316
317    /// Returns a zero-copy slice of this buffer with length `len` and starting at `offset`
318    ///
319    /// # Panics
320    ///
321    /// Panics if `offset + len > self.len()`
322    pub fn slice(&self, offset: usize, len: usize) -> Self {
323        Self(self.0.slice(offset, len.saturating_add(1)))
324    }
325
326    /// Returns true if this [`OffsetBuffer`] is equal to `other`, using pointer comparisons
327    /// to determine buffer equality. This is cheaper than `PartialEq::eq` but may
328    /// return false when the arrays are logically equal
329    #[inline]
330    pub fn ptr_eq(&self, other: &Self) -> bool {
331        self.0.ptr_eq(&other.0)
332    }
333
334    /// Check if any null positions in the `null_buffer` correspond to
335    /// non-empty ranges in this [`OffsetBuffer`].
336    ///
337    /// In variable-length array types (e.g., `StringArray`, `ListArray`),
338    /// null entries may or may not have empty offset ranges. This method
339    /// detects cases where a null entry has a non-empty range
340    /// (i.e., `offsets[i] != offsets[i+1]`), which means the underlying
341    /// data buffer contains data behind nulls.
342    ///
343    /// This matters because unwrapping (flattening) a list array exposes
344    /// the child values, including those behind null entries. If null
345    /// entries point to non-empty ranges, the unwrapped values will
346    /// contain data that may not be meaningful to operate on and could
347    /// cause errors (e.g., division by zero in the child values).
348    ///
349    /// Returns `false` if `null_buffer` is `None` or contains no nulls.
350    ///
351    /// # Example
352    ///
353    /// ```
354    /// # use arrow_buffer::{OffsetBuffer, ScalarBuffer, NullBuffer};
355    /// // Offsets where null at index 1 has an empty range (3..3)
356    /// let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 3, 6]));
357    /// let nulls = NullBuffer::from(vec![true, false, true]);
358    /// assert!(!offsets.has_non_empty_nulls(Some(&nulls)));
359    ///
360    /// // Offsets where null at index 1 has a non-empty range (3..7)
361    /// let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 7, 10]));
362    /// let nulls = NullBuffer::from(vec![true, false, true]);
363    /// assert!(offsets.has_non_empty_nulls(Some(&nulls)));
364    /// ```
365    ///
366    /// # Panics
367    ///
368    /// Panics if the length of the `null_buffer` does not equal `self.len() - 1`.
369    pub fn has_non_empty_nulls(&self, null_buffer: Option<&NullBuffer>) -> bool {
370        let Some(null_buffer) = null_buffer else {
371            return false;
372        };
373
374        assert_eq!(
375            self.len() - 1,
376            null_buffer.len(),
377            "The length of the offsets should be 1 more than the length of the null buffer"
378        );
379
380        if null_buffer.null_count() == 0 {
381            return false;
382        }
383
384        // Offsets always have at least 1 value
385        let initial_offset = self[0];
386        let last_offset = self[self.len() - 1];
387
388        // If all the values are null (offsets have 1 more value than the length of the array)
389        if null_buffer.null_count() == self.len() - 1 {
390            return last_offset != initial_offset;
391        }
392
393        let mut valid_slices_iter = null_buffer.valid_slices();
394
395        // This is safe as we validated that are at least 1 valid value in the array
396        let (start, end) = valid_slices_iter.next().unwrap();
397
398        // If the nulls before have length greater than 0
399        if self[start] != initial_offset {
400            return true;
401        }
402
403        // End is exclusive, so it already point to the last offset value
404        // This is valid as the length of the array is always 1 less than the length of the offsets
405        let mut end_offset_of_last_valid_value = self[end];
406
407        for (start, end) in valid_slices_iter {
408            // If there is a null value that point to a non-empty value than the start offset of the valid value
409            // will be different that the end offset of the last valid value
410            if self[start] != end_offset_of_last_valid_value {
411                return true;
412            }
413
414            // End is exclusive, so it already point to the last offset value
415            // This is valid as the length of the array is always 1 less than the length of the offsets
416            end_offset_of_last_valid_value = self[end];
417        }
418
419        end_offset_of_last_valid_value != last_offset
420    }
421
422    /// Subtract `rhs` from all offsets
423    /// This will try to reuse the existing allocation as much as possible
424    ///
425    /// # Panics
426    ///
427    /// Panics if `rhs` > the first offset or if `rhs` will lead to overflow (when `rhs` is negative)
428    ///
429    /// # Example
430    ///
431    /// ```
432    /// # use arrow_buffer::OffsetBuffer;
433    /// let offsets = OffsetBuffer::<i32>::from_lengths(vec![4, 1, 5, 6]);
434    /// assert_eq!(offsets.as_ref(), &[0, 4, 5, 10, 16]);
435    ///
436    /// let sliced_offsets = offsets.slice(1, 2);
437    /// assert_eq!(sliced_offsets.as_ref(), &[4, 5, 10]);
438    ///
439    /// let shifted_offsets = sliced_offsets.subtract(4);
440    /// assert_eq!(shifted_offsets.as_ref(), &[0, 1, 6]);
441    /// ```
442    ///
443    pub fn subtract(self, rhs: O) -> Self
444    where
445        O: std::ops::Sub<Output = O> + std::cmp::PartialOrd + num_traits::CheckedSub,
446    {
447        if rhs == O::usize_as(0) {
448            return self;
449        }
450
451        let len = self.len();
452
453        // Offset buffer is guaranteed to be non-empty
454        assert!(
455            self[0] >= rhs,
456            "shifted offsets will become negative which is not allowed"
457        );
458
459        // If negative, make sure that this will not create an overflow
460        if rhs < O::usize_as(0) {
461            self[len - 1].checked_sub(&rhs).expect("must not overflow");
462        }
463
464        // try and reuse buffer
465        let shifted_offsets: Vec<O> = match self.into_inner().into_inner().into_vec() {
466            // If we can reuse the buffer, update in place
467            Ok(mut v) => {
468                for offset in &mut v {
469                    *offset = *offset - rhs;
470                }
471                v
472            }
473            // otherwise, buffer is shared so we need a copy
474            Err(buffer) => {
475                let offsets = ScalarBuffer::<O>::from(buffer);
476                offsets.iter().map(|offset| *offset - rhs).collect()
477            }
478        };
479        let shifted_buffer = ScalarBuffer::from(shifted_offsets);
480        // Safety: offsets are valid as they are coming from a valid
481        // offset buffer and we checked overflow above, and we
482        // subtracted the same value from all offsets, thus keeping the
483        // same properties as the input buffer
484        unsafe { Self::new_unchecked(shifted_buffer) }
485    }
486}
487
488impl<T: ArrowNativeType> Deref for OffsetBuffer<T> {
489    type Target = [T];
490
491    #[inline]
492    fn deref(&self) -> &Self::Target {
493        &self.0
494    }
495}
496
497impl<T: ArrowNativeType> AsRef<[T]> for OffsetBuffer<T> {
498    #[inline]
499    fn as_ref(&self) -> &[T] {
500        self
501    }
502}
503
504impl<O: ArrowNativeType> From<OffsetBufferBuilder<O>> for OffsetBuffer<O> {
505    fn from(value: OffsetBufferBuilder<O>) -> Self {
506        value.finish()
507    }
508}
509
510impl<O: ArrowNativeType> Default for OffsetBuffer<O> {
511    fn default() -> Self {
512        Self::new_empty()
513    }
514}
515
516#[cfg(test)]
517mod tests {
518    use super::*;
519
520    #[test]
521    #[should_panic(expected = "offsets cannot be empty")]
522    fn empty_offsets() {
523        OffsetBuffer::new(Vec::<i32>::new().into());
524    }
525
526    #[test]
527    #[should_panic(expected = "offsets must be greater than 0")]
528    fn negative_offsets() {
529        OffsetBuffer::new(vec![-1, 0, 1].into());
530    }
531
532    #[test]
533    fn try_from_lengths_overflow() {
534        // Fits in an i64, but not in an i32:
535        let lengths = [u32::MAX as usize, 1];
536
537        let err = OffsetBuffer::<i32>::try_from_lengths(lengths).unwrap_err();
538        assert_eq!(
539            err.to_string(),
540            "offset overflow: 4294967296 does not fit in i32"
541        );
542        assert!(OffsetBuffer::<i64>::try_from_lengths(lengths).is_ok());
543
544        let err = OffsetBuffer::<i32>::try_from_lengths([usize::MAX, 1]).unwrap_err();
545        assert_eq!(
546            err.to_string(),
547            "total length overflow: does not fit in usize"
548        );
549
550        // The panicking version agrees:
551        assert!(std::panic::catch_unwind(|| OffsetBuffer::<i32>::from_lengths(lengths)).is_err());
552    }
553
554    #[test]
555    fn try_from_repeated_length_overflow() {
556        assert_eq!(
557            OffsetBuffer::<i32>::try_from_repeated_length(2, usize::MAX)
558                .unwrap_err()
559                .to_string(),
560            "total length overflow: does not fit in usize"
561        );
562        assert_eq!(
563            OffsetBuffer::<i32>::try_from_repeated_length(u32::MAX as usize, 2)
564                .unwrap_err()
565                .to_string(),
566            "offset overflow: 8589934590 does not fit in i32"
567        );
568        assert_eq!(
569            OffsetBuffer::<i32>::try_from_repeated_length(4, 3)
570                .unwrap()
571                .as_ref(),
572            &[0, 4, 8, 12]
573        );
574    }
575
576    #[test]
577    fn try_new_zeroed_overflow() {
578        assert_eq!(
579            OffsetBuffer::<i64>::try_new_zeroed(usize::MAX)
580                .unwrap_err()
581                .to_string(),
582            "buffer length overflow: does not fit in usize"
583        );
584        assert_eq!(
585            OffsetBuffer::<i32>::try_new_zeroed(3).unwrap().as_ref(),
586            &[0; 4]
587        );
588    }
589
590    #[test]
591    fn offsets() {
592        OffsetBuffer::new(vec![0, 1, 2, 3].into());
593
594        let offsets = OffsetBuffer::<i32>::new_zeroed(3);
595        assert_eq!(offsets.as_ref(), &[0; 4]);
596
597        let offsets = OffsetBuffer::<i32>::new_zeroed(0);
598        assert_eq!(offsets.as_ref(), &[0; 1]);
599    }
600
601    #[test]
602    #[should_panic(expected = "overflow")]
603    fn offsets_new_zeroed_overflow() {
604        OffsetBuffer::<i32>::new_zeroed(usize::MAX);
605    }
606
607    #[test]
608    #[should_panic(expected = "offsets must be monotonically increasing")]
609    fn non_monotonic_offsets() {
610        OffsetBuffer::new(vec![1, 2, 0].into());
611    }
612
613    #[test]
614    fn from_lengths() {
615        let buffer = OffsetBuffer::<i32>::from_lengths([2, 6, 3, 7, 2]);
616        assert_eq!(buffer.as_ref(), &[0, 2, 8, 11, 18, 20]);
617
618        let half_max = i32::MAX / 2;
619        let buffer = OffsetBuffer::<i32>::from_lengths([half_max as usize, half_max as usize]);
620        assert_eq!(buffer.as_ref(), &[0, half_max, half_max * 2]);
621    }
622
623    #[test]
624    #[should_panic(expected = "offset overflow")]
625    fn from_lengths_offset_overflow() {
626        OffsetBuffer::<i32>::from_lengths([i32::MAX as usize, 1]);
627    }
628
629    #[test]
630    #[should_panic(expected = "total length overflow: does not fit in usize")]
631    fn from_lengths_usize_overflow() {
632        OffsetBuffer::<i32>::from_lengths([usize::MAX, 1]);
633    }
634
635    #[test]
636    #[should_panic(expected = "offset overflow")]
637    fn from_repeated_lengths_offset_length_overflow() {
638        OffsetBuffer::<i32>::from_repeated_length(i32::MAX as usize / 4, 5);
639    }
640
641    #[test]
642    #[should_panic(expected = "offset overflow")]
643    fn from_repeated_lengths_offset_repeat_overflow() {
644        OffsetBuffer::<i32>::from_repeated_length(1, i32::MAX as usize + 1);
645    }
646
647    #[test]
648    #[should_panic(expected = "offset overflow")]
649    fn from_repeated_lengths_usize_length_overflow() {
650        OffsetBuffer::<i32>::from_repeated_length(usize::MAX, 1);
651    }
652
653    #[test]
654    #[should_panic(expected = "total length overflow: does not fit in usize")]
655    fn from_repeated_lengths_usize_length_usize_overflow() {
656        OffsetBuffer::<i32>::from_repeated_length(usize::MAX, 2);
657    }
658
659    #[test]
660    #[should_panic(expected = "offset overflow")]
661    fn from_repeated_lengths_usize_repeat_overflow() {
662        OffsetBuffer::<i32>::from_repeated_length(1, usize::MAX);
663    }
664
665    #[test]
666    fn get_lengths() {
667        let offsets = OffsetBuffer::<i32>::new(ScalarBuffer::<i32>::from(vec![0, 1, 4, 9]));
668        assert_eq!(offsets.lengths().collect::<Vec<usize>>(), vec![1, 3, 5]);
669    }
670
671    #[test]
672    fn get_lengths_should_be_with_fixed_size() {
673        let offsets = OffsetBuffer::<i32>::new(ScalarBuffer::<i32>::from(vec![0, 1, 4, 9]));
674        let iter = offsets.lengths();
675        assert_eq!(iter.size_hint(), (3, Some(3)));
676        assert_eq!(iter.len(), 3);
677    }
678
679    #[test]
680    fn get_lengths_from_empty_offset_buffer_should_be_empty_iterator() {
681        let offsets = OffsetBuffer::<i32>::new_empty();
682        assert_eq!(offsets.lengths().collect::<Vec<usize>>(), vec![]);
683    }
684
685    #[test]
686    fn impl_eq() {
687        fn are_equal<T: Eq>(a: &T, b: &T) -> bool {
688            a.eq(b)
689        }
690
691        assert!(
692            are_equal(
693                &OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 1, 4, 9])),
694                &OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 1, 4, 9]))
695            ),
696            "OffsetBuffer should implement Eq."
697        );
698    }
699
700    #[test]
701    fn impl_default() {
702        let default = OffsetBuffer::<i32>::default();
703        assert_eq!(default.as_ref(), &[0]);
704    }
705
706    #[test]
707    fn from_repeated_length_basic() {
708        // Basic case with length 4, repeated 3 times
709        let buffer = OffsetBuffer::<i32>::from_repeated_length(4, 3);
710        assert_eq!(buffer.as_ref(), &[0, 4, 8, 12]);
711
712        // Verify the lengths are correct
713        let lengths: Vec<usize> = buffer.lengths().collect();
714        assert_eq!(lengths, vec![4, 4, 4]);
715    }
716
717    #[test]
718    fn from_repeated_length_single_repeat() {
719        // Length 5, repeated once
720        let buffer = OffsetBuffer::<i32>::from_repeated_length(5, 1);
721        assert_eq!(buffer.as_ref(), &[0, 5]);
722
723        let lengths: Vec<usize> = buffer.lengths().collect();
724        assert_eq!(lengths, vec![5]);
725    }
726
727    #[test]
728    fn from_repeated_length_zero_repeats() {
729        let buffer = OffsetBuffer::<i32>::from_repeated_length(10, 0);
730        assert_eq!(buffer, OffsetBuffer::<i32>::new_empty());
731    }
732
733    #[test]
734    fn from_repeated_length_zero_length() {
735        // Zero length, repeated 5 times (all zeros)
736        let buffer = OffsetBuffer::<i32>::from_repeated_length(0, 5);
737        assert_eq!(buffer.as_ref(), &[0, 0, 0, 0, 0, 0]);
738
739        // All lengths should be 0
740        let lengths: Vec<usize> = buffer.lengths().collect();
741        assert_eq!(lengths, vec![0, 0, 0, 0, 0]);
742    }
743
744    #[test]
745    fn from_repeated_length_large_values() {
746        // Test with larger values that don't overflow
747        let buffer = OffsetBuffer::<i32>::from_repeated_length(1000, 100);
748        assert_eq!(buffer[0], 0);
749
750        // Verify all lengths are 1000
751        let lengths: Vec<usize> = buffer.lengths().collect();
752        assert_eq!(lengths.len(), 100);
753        assert!(lengths.iter().all(|&len| len == 1000));
754    }
755
756    #[test]
757    fn from_repeated_length_unit_length() {
758        // Length 1, repeated multiple times
759        let buffer = OffsetBuffer::<i32>::from_repeated_length(1, 10);
760        assert_eq!(buffer.as_ref(), &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
761
762        let lengths: Vec<usize> = buffer.lengths().collect();
763        assert_eq!(lengths, vec![1; 10]);
764    }
765
766    #[test]
767    fn from_repeated_length_max_safe_values() {
768        // Test with maximum safe values for i32
769        // i32::MAX / 3 ensures we don't overflow when repeated twice
770        let third_max = (i32::MAX / 3) as usize;
771        let buffer = OffsetBuffer::<i32>::from_repeated_length(third_max, 2);
772        assert_eq!(
773            buffer.as_ref(),
774            &[0, third_max as i32, (third_max * 2) as i32]
775        );
776    }
777
778    // ---------------------------------------------------------------
779    // Tests for has_non_empty_nulls
780    // ---------------------------------------------------------------
781
782    #[test]
783    fn has_non_empty_nulls_none_null_buffer() {
784        // No null buffer at all -> false
785        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 5, 8]));
786        assert!(!offsets.has_non_empty_nulls(None));
787    }
788
789    #[test]
790    fn has_non_empty_nulls_all_valid() {
791        // Null buffer with zero nulls -> false (early return via filter)
792        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 5, 8]));
793        let nulls = NullBuffer::new_valid(3);
794        assert!(!offsets.has_non_empty_nulls(Some(&nulls)));
795    }
796
797    #[test]
798    fn has_non_empty_nulls_all_null_empty_offsets() {
799        // All values are null and all offsets are equal (no data behind nulls) -> false
800        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 0, 0, 0]));
801        let nulls = NullBuffer::new_null(3);
802        assert!(!offsets.has_non_empty_nulls(Some(&nulls)));
803    }
804
805    #[test]
806    fn has_non_empty_nulls_all_null_non_empty_offsets() {
807        // All values are null but offsets span data -> true
808        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 5, 7]));
809        let nulls = NullBuffer::new_null(3);
810        assert!(offsets.has_non_empty_nulls(Some(&nulls)));
811    }
812
813    #[test]
814    fn has_non_empty_nulls_all_null_nonzero_but_equal_offsets() {
815        // All null, offsets start at non-zero but are all equal -> false
816        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![5, 5, 5]));
817        let nulls = NullBuffer::new_null(2);
818        assert!(!offsets.has_non_empty_nulls(Some(&nulls)));
819    }
820
821    #[test]
822    fn has_non_empty_nulls_leading_nulls_with_data() {
823        // Nulls at the beginning that point to non-empty ranges -> true
824        // offsets: [0, 3, 5, 8]  nulls: [false, true, true]
825        // Index 0 is null with range 0..3 (non-empty)
826        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 5, 8]));
827        let nulls = NullBuffer::from(vec![false, true, true]);
828        assert!(offsets.has_non_empty_nulls(Some(&nulls)));
829    }
830
831    #[test]
832    fn has_non_empty_nulls_leading_nulls_without_data() {
833        // Nulls at the beginning with empty ranges -> continue checking
834        // offsets: [0, 0, 3, 6]  nulls: [false, true, true]
835        // Index 0 is null with range 0..0 (empty)
836        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 0, 3, 6]));
837        let nulls = NullBuffer::from(vec![false, true, true]);
838        assert!(!offsets.has_non_empty_nulls(Some(&nulls)));
839    }
840
841    #[test]
842    fn has_non_empty_nulls_only_trailing_null_has_data() {
843        // Only the trailing null region has data, everything else is clean
844        // offsets: [0, 0, 3, 6, 8]  nulls: [false, true, true, false]
845        // Null at 0 (0..0 empty), valid at 1,2 (0..3, 3..6), null at 3 (6..8 non-empty)
846        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 0, 3, 6, 8]));
847        let nulls = NullBuffer::from(vec![false, true, true, false]);
848        assert!(offsets.has_non_empty_nulls(Some(&nulls)));
849    }
850
851    #[test]
852    fn has_non_empty_nulls_trailing_nulls_without_data() {
853        // Nulls at the end with empty ranges -> false
854        // offsets: [0, 3, 6, 6]  nulls: [true, true, false]
855        // Index 2 is null with range 6..6 (empty)
856        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 6, 6]));
857        let nulls = NullBuffer::from(vec![true, true, false]);
858        assert!(!offsets.has_non_empty_nulls(Some(&nulls)));
859    }
860
861    #[test]
862    fn has_non_empty_nulls_middle_nulls_with_data() {
863        // Null in the middle with non-empty range -> true
864        // offsets: [0, 3, 7, 10]  nulls: [true, false, true]
865        // Index 1 is null with range 3..7 (non-empty)
866        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 7, 10]));
867        let nulls = NullBuffer::from(vec![true, false, true]);
868        assert!(offsets.has_non_empty_nulls(Some(&nulls)));
869    }
870
871    #[test]
872    fn has_non_empty_nulls_middle_nulls_without_data() {
873        // Null in the middle with empty range -> false
874        // offsets: [0, 3, 3, 6]  nulls: [true, false, true]
875        // Index 1 is null with range 3..3 (empty)
876        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 3, 6]));
877        let nulls = NullBuffer::from(vec![true, false, true]);
878        assert!(!offsets.has_non_empty_nulls(Some(&nulls)));
879    }
880
881    #[test]
882    fn has_non_empty_nulls_alternating_null_valid_all_empty() {
883        // Alternating null/valid where every null has an empty range -> false.
884
885        // Ends with null
886        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 0, 3, 3, 6, 6]));
887        let nulls = NullBuffer::from(vec![false, true, false, true, false]);
888        assert!(!offsets.has_non_empty_nulls(Some(&nulls)));
889
890        // Ends with valid
891        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 0, 3, 3, 6, 6, 9]));
892        let nulls = NullBuffer::from(vec![false, true, false, true, false, true]);
893        assert!(!offsets.has_non_empty_nulls(Some(&nulls)));
894    }
895
896    #[test]
897    fn has_non_empty_nulls_multiple_null_regions_second_has_data() {
898        // Two null regions: first empty, second non-empty -> true
899        // offsets: [0, 0, 3, 5, 6]  nulls: [false, true, false, true]
900        // Null at index 0 (0..0 empty), null at index 2 (3..5 non-empty)
901        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 0, 3, 5, 6]));
902        let nulls = NullBuffer::from(vec![false, true, false, true]);
903        assert!(offsets.has_non_empty_nulls(Some(&nulls)));
904    }
905
906    #[test]
907    fn has_non_empty_nulls_multiple_null_regions_later_gap_has_data() {
908        // Three null regions: first two empty, third non-empty -> true
909        // offsets: [0, 0, 3, 3, 6, 8, 10]  nulls: [false, true, false, true, false, true]
910        // valid_slices: (1,2), (3,4), (5,6)
911        // first slice: start=1, self[1]=0 == initial_offset=0 OK, end_offset=self[2]=3
912        // loop iter 1: start=3, self[3]=3 == 3 OK (first gap empty), end_offset=self[4]=6
913        // loop iter 2: start=5, self[5]=8 != 6 -> true (second gap has data)
914        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 0, 3, 3, 6, 8, 10]));
915        let nulls = NullBuffer::from(vec![false, true, false, true, false, true]);
916        assert!(offsets.has_non_empty_nulls(Some(&nulls)));
917    }
918
919    #[test]
920    fn has_non_empty_nulls_single_element_null_empty() {
921        // Single element, null with empty range -> false
922        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 0]));
923        let nulls = NullBuffer::new_null(1);
924        assert!(!offsets.has_non_empty_nulls(Some(&nulls)));
925    }
926
927    #[test]
928    fn has_non_empty_nulls_single_element_null_non_empty() {
929        // Single element, null with non-empty range -> true
930        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 5]));
931        let nulls = NullBuffer::new_null(1);
932        assert!(offsets.has_non_empty_nulls(Some(&nulls)));
933    }
934
935    #[test]
936    fn has_non_empty_nulls_single_element_valid() {
937        // Single element, valid -> false (no nulls at all)
938        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 5]));
939        let nulls = NullBuffer::new_valid(1);
940        assert!(!offsets.has_non_empty_nulls(Some(&nulls)));
941    }
942
943    #[test]
944    fn has_non_empty_nulls_consecutive_nulls_between_valid_slices() {
945        // Multiple consecutive nulls between valid regions
946        // offsets: [0, 2, 2, 2, 5, 8]  nulls: [true, false, false, true, true]
947        // Valid: [0], nulls: [1,2], valid: [3,4]
948        // Null region [1,2] has offsets 2..2..2 (empty) -> false
949        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 2, 2, 5, 8]));
950        let nulls = NullBuffer::from(vec![true, false, false, true, true]);
951        assert!(!offsets.has_non_empty_nulls(Some(&nulls)));
952    }
953
954    #[test]
955    fn has_non_empty_nulls_consecutive_nulls_between_valid_slices_with_data() {
956        // Multiple consecutive nulls between valid regions, nulls have data
957        // offsets: [0, 2, 3, 4, 5, 8]  nulls: [true, false, false, true, true]
958        // valid_slices: (0,1), (3,5)
959        // first slice: start=0, end=1 -> self[0]=0 == initial_offset=0 OK
960        //   end_offset_of_last_valid_value = self[1] = 2
961        // second slice: start=3, end=5 -> self[3]=4 != 2 -> true
962        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3, 4, 5, 8]));
963        let nulls = NullBuffer::from(vec![true, false, false, true, true]);
964        assert!(offsets.has_non_empty_nulls(Some(&nulls)));
965    }
966
967    #[test]
968    fn has_non_empty_nulls_nonzero_initial_offset_all_null_equal() {
969        // Non-zero starting offset, all null, all offsets equal -> false
970        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![10, 10, 10]));
971        let nulls = NullBuffer::new_null(2);
972        assert!(!offsets.has_non_empty_nulls(Some(&nulls)));
973    }
974
975    #[test]
976    fn has_non_empty_nulls_nonzero_initial_offset_with_data() {
977        // Non-zero starting offset, null has data
978        // offsets: [10, 15, 20]  nulls: [false, true]
979        // Null at index 0 with range 10..15 (non-empty) -> true
980        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![10, 15, 20]));
981        let nulls = NullBuffer::from(vec![false, true]);
982        assert!(offsets.has_non_empty_nulls(Some(&nulls)));
983    }
984
985    #[test]
986    fn has_non_empty_nulls_sliced_no_nulls_in_null_region() {
987        // Original: [0, 3, 3, 6, 6, 9]  -> slice(1, 3) -> [3, 3, 6, 6]
988        // initial_offset=3, last_offset=6
989        // nulls: [false, true, false]  (null at index 0 has range 3..3 = empty)
990        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 3, 6, 6, 9]));
991        let sliced = offsets.slice(1, 3);
992        let nulls = NullBuffer::from(vec![false, true, false]);
993        assert!(!sliced.has_non_empty_nulls(Some(&nulls)));
994    }
995
996    #[test]
997    fn has_non_empty_nulls_sliced_null_has_data() {
998        // Original: [0, 3, 7, 10, 15]  -> slice(1, 2) -> [3, 7, 10]
999        // initial_offset=3, last_offset=10
1000        // nulls: [false, true]  (null at index 0 has range 3..7 = non-empty)
1001        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 7, 10, 15]));
1002        let sliced = offsets.slice(1, 2);
1003        let nulls = NullBuffer::from(vec![false, true]);
1004        assert!(sliced.has_non_empty_nulls(Some(&nulls)));
1005    }
1006
1007    #[test]
1008    #[should_panic(
1009        expected = "The length of the offsets should be 1 more than the length of the null buffer"
1010    )]
1011    fn has_non_empty_nulls_all_valid_mismatched_lengths_too_short() {
1012        // All-valid null buffer with wrong length should still panic
1013        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 5, 8]));
1014        let nulls = NullBuffer::new_valid(2); // expects 3
1015        offsets.has_non_empty_nulls(Some(&nulls));
1016    }
1017
1018    #[test]
1019    #[should_panic(
1020        expected = "The length of the offsets should be 1 more than the length of the null buffer"
1021    )]
1022    fn has_non_empty_nulls_all_valid_mismatched_lengths_too_long() {
1023        // All-valid null buffer with wrong length should still panic
1024        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 5, 8]));
1025        let nulls = NullBuffer::new_valid(5); // expects 3
1026        offsets.has_non_empty_nulls(Some(&nulls));
1027    }
1028
1029    #[test]
1030    #[should_panic(expected = "shifted offsets will become negative which is not allowed")]
1031    fn should_panic_for_subtract_by_value_that_will_cause_offsets_to_be_less_than_zero() {
1032        // self[0] = 0, rhs = 1 -> 0 >= 1 is false -> assert fires
1033        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 6]));
1034        offsets.subtract(1);
1035    }
1036
1037    #[test]
1038    fn subtract_by_value_that_will_cause_offsets_to_be_less_than_zero_for_outside_the_slice() {
1039        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 1, 4, 7]));
1040        let sliced = offsets.slice(1, 2); // [1, 4, 7]
1041        drop(offsets);
1042        assert_eq!(sliced.as_ref(), &[1, 4, 7]);
1043
1044        let result = sliced.subtract(1);
1045        assert_eq!(result.as_ref(), &[0, 3, 6]);
1046        assert_eq!(result.len(), 3);
1047    }
1048
1049    #[test]
1050    #[should_panic(expected = "must not overflow")]
1051    fn should_panic_subtract_by_value_that_will_cause_offsets_to_overflow() {
1052        // rhs = -1 (negative). self[0] = 0 >= -1 passes.
1053        // last offset i32::MAX - (-1) = i32::MAX + 1 -> checked_sub returns None -> expect fires
1054        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 5, i32::MAX]));
1055        offsets.subtract(-1);
1056    }
1057
1058    #[test]
1059    fn subtract_by_value_that_will_cause_offsets_to_overflow_outside_the_slice() {
1060        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 6, i32::MAX]));
1061        let sliced = offsets.slice(0, 2); // [0, 3, 6]
1062        assert_eq!(sliced.as_ref(), &[0, 3, 6]);
1063
1064        let result = sliced.subtract(-1);
1065        assert_eq!(result.as_ref(), &[1, 4, 7]);
1066        assert_eq!(result.len(), 3);
1067    }
1068
1069    #[test]
1070    fn when_shift_is_0_subtract_should_reuse_the_buffer_even_when_it_is_shared() {
1071        // subtract(0) hits the early `return self` before any into_mutable,
1072        // so the returned buffer is the exact same allocation even while shared.
1073        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 6]));
1074        let shared = offsets.clone(); // refcount now 2 -> shared
1075        let result = offsets.subtract(0);
1076        assert!(
1077            result.ptr_eq(&shared),
1078            "subtract(0) must return the same underlying buffer, even when shared"
1079        );
1080    }
1081
1082    #[test]
1083    fn should_reuse_the_underline_data_when_the_buffer_is_not_shared() {
1084        // Unique ownership, offset 0 -> into_mutable succeeds -> mutate in place,
1085        // and MutableBuffer -> Buffer keeps the same allocation.
1086        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![2, 5, 8]));
1087        let ptr_before = offsets.as_ptr();
1088        let result = offsets.subtract(2);
1089        assert_eq!(
1090            ptr_before,
1091            result.as_ptr(),
1092            "a non-shared buffer should be mutated in place, reusing the allocation"
1093        );
1094        assert_eq!(result.as_ref(), &[0, 3, 6]);
1095    }
1096
1097    #[test]
1098    fn should_create_a_new_buffer_when_the_buffer_is_shared() {
1099        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![2, 5, 8]));
1100        let shared = offsets.clone();
1101        let ptr_before = offsets.as_ptr();
1102        let result = offsets.subtract(2);
1103        assert_ne!(
1104            ptr_before,
1105            result.as_ptr(),
1106            "a shared buffer must not be mutated in place; a new allocation is created"
1107        );
1108        assert_eq!(result.as_ref(), &[0, 3, 6]);
1109        // The shared view is untouched.
1110        assert_eq!(shared.as_ref(), &[2, 5, 8]);
1111    }
1112
1113    #[test]
1114    fn when_shift_is_negative_it_should_shift_offsets_in_the_right_direction() {
1115        // rhs = -2 -> offset - (-2) = offset + 2, so all offsets move up.
1116        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 6]));
1117        let result = offsets.subtract(-2);
1118        assert_eq!(result.as_ref(), &[2, 5, 8]);
1119    }
1120
1121    // Replace this test with test that assert a reuse after PR #10118 is merged
1122    #[test]
1123    fn for_sliced_unshared_buffer_shift_should_not_reuse_buffer() {
1124        // Underlying [0, 3, 6, 9, 12]; slice -> view [3, 6, 9].
1125        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![1, 3, 6, 9, 12]));
1126        let sliced = offsets.slice(1, 2); // [3, 6, 9]
1127        drop(offsets); // uniquely owned
1128        assert_eq!(sliced.as_ref(), &[3, 6, 9]);
1129
1130        let ptr_before = sliced.as_ptr();
1131        let result = sliced.subtract(1);
1132
1133        assert_ne!(
1134            ptr_before,
1135            result.as_ptr(),
1136            "should not be reused until #10118 is merged"
1137        );
1138
1139        assert_eq!(result.as_ref(), &[2, 5, 8]);
1140    }
1141
1142    #[test]
1143    fn for_sliced_but_start_at_0_unshared_buffer_shift_should_reuse_buffer() {
1144        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![1, 3, 6, 9, 12]));
1145        let sliced = offsets.slice(0, 2);
1146        drop(offsets); // uniquely owned
1147        assert_eq!(sliced.as_ref(), &[1, 3, 6]);
1148
1149        let ptr_before = sliced.as_ptr();
1150        let result = sliced.subtract(1);
1151
1152        assert_eq!(ptr_before, result.as_ptr(), "should be reused");
1153
1154        assert_eq!(result.as_ref(), &[0, 2, 5]);
1155    }
1156
1157    #[test]
1158    fn for_sliced_shared_buffer_shifted_buffer_should_only_include_the_sliced_data() {
1159        // Underlying: [0, 3, 6, 9, 12]; slice(1, 2) -> view [3, 6, 9].
1160        // `offsets` stays alive, so the sliced buffer is shared -> Err branch.
1161        // The Err branch copies `len` (= 3) elements from the *sliced* typed_data,
1162        // so the result contains only the sliced data, shifted.
1163        let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 6, 9, 12]));
1164        let sliced = offsets.slice(1, 2);
1165        assert_eq!(sliced.as_ref(), &[3, 6, 9]);
1166
1167        let result = sliced.subtract(3);
1168
1169        assert_eq!(
1170            result.as_ref(),
1171            &[0, 3, 6],
1172            "shifted result should contain only the sliced data"
1173        );
1174        assert_eq!(result.len(), 3);
1175
1176        // Assert that the underlying buffer of the result is not sliced to make sure it does not include the data outside the slice range from the original buffer
1177        let underlying_buffer = result.inner().inner();
1178        assert_eq!(underlying_buffer.ptr_offset(), 0);
1179        assert_eq!(underlying_buffer.len(), 3 * std::mem::size_of::<i32>());
1180    }
1181}
arrow_buffer/buffer/offset.rs

arrow_buffer/buffer/
offset.rs
