arrow_data/transform/
utils.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use arrow_buffer::{ArrowNativeType, MutableBuffer, bit_util};
6use arrow_schema::ArrowError;
7use num_integer::Integer;
8use num_traits::CheckedAdd;
9
10/// extends the `buffer` to be able to hold `len` bits, setting all bits of the new size to zero.
11#[inline]
12pub(super) fn resize_for_bits(buffer: &mut MutableBuffer, len: usize) {
13    let needed_bytes = bit_util::ceil(len, 8);
14    if buffer.len() < needed_bytes {
15        buffer.resize(needed_bytes, 0);
16    }
17}
18
19/// Extends `buffer` with the re-based offsets from `offsets`, returning an error on overflow.
20pub(super) fn try_extend_offsets<T: ArrowNativeType + Integer + CheckedAdd>(
21    buffer: &mut MutableBuffer,
22    mut last_offset: T,
23    offsets: &[T],
24) -> Result<(), ArrowError> {
25    buffer
26        .try_reserve(std::mem::size_of_val(offsets))
27        .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
28    // Snapshot the length so we can roll back partial writes on overflow.
29    let original_len = buffer.len();
30    for window in offsets.windows(2) {
31        let length = window[1] - window[0];
32        match last_offset.checked_add(&length) {
33            Some(new_offset) => {
34                last_offset = new_offset;
35                buffer.push(last_offset);
36            }
37            None => {
38                // Restore the buffer to its state before this call so the
39                // caller is not left with a partially-written offset sequence.
40                buffer.resize(original_len, 0);
41                return Err(ArrowError::InvalidArgumentError(
42                    "offset overflow: data exceeds the capacity of the offset type. \
43                     Try splitting into smaller batches or using a larger type \
44                     (e.g. LargeStringArray / LargeBinaryArray instead of StringArray / BinaryArray)"
45                        .to_string(),
46                ));
47            }
48        }
49    }
50    Ok(())
51}
52
53#[inline]
54pub(super) unsafe fn get_last_offset<T: ArrowNativeType>(offset_buffer: &MutableBuffer) -> T {
55    // JUSTIFICATION
56    //  Benefit
57    //      20% performance improvement extend of variable sized arrays (see bench `mutable_array`)
58    //  Soundness
59    //      * offset buffer is always extended in slices of T and aligned accordingly.
60    //      * Buffer[0] is initialized with one element, 0, and thus `mutable_offsets.len() - 1` is always valid.
61    let (prefix, offsets, suffix) = unsafe { offset_buffer.as_slice().align_to::<T>() };
62    debug_assert!(prefix.is_empty() && suffix.is_empty());
63    *unsafe { offsets.get_unchecked(offsets.len() - 1) }
64}
65
66#[cfg(test)]
67mod tests {
68    use crate::transform::utils::try_extend_offsets;
69    use arrow_buffer::MutableBuffer;
70
71    #[test]
72    fn test_overflow_returns_error() {
73        let mut buffer = MutableBuffer::new(10);
74        let err = try_extend_offsets(&mut buffer, i32::MAX - 4, &[0i32, 5]).unwrap_err();
75        assert!(
76            err.to_string().contains("offset overflow"),
77            "unexpected error: {err}"
78        );
79    }
80
81    #[test]
82    fn test_overflow_restores_buffer() {
83        // Pre-populate the buffer with a known-good offset so we can verify
84        // it is unchanged after a failed extend.
85        let mut buffer = MutableBuffer::new(16);
86        buffer.push(0i32);
87        buffer.push(10i32);
88        let len_before = buffer.len();
89
90        // Offsets [0, 3, i32::MAX]: the second window (3 → i32::MAX) will overflow
91        // because last_offset (i32::MAX - 4 + 3 = i32::MAX - 1) + (i32::MAX - 3) overflows.
92        // Use a simpler case: start near MAX so the very first window overflows.
93        let err = try_extend_offsets(&mut buffer, i32::MAX - 2, &[0i32, 5]).unwrap_err();
94        assert!(
95            err.to_string().contains("offset overflow"),
96            "unexpected error: {err}"
97        );
98        // Buffer must be exactly as it was before the failed call.
99        assert_eq!(
100            buffer.len(),
101            len_before,
102            "buffer length changed after overflow rollback"
103        );
104    }
105}
arrow_data/transform/utils.rs

arrow_data/transform/
utils.rs
