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

arrow_buffer/buffer/
offset.rs
