arrow_buffer/buffer/
scalar.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::alloc::Deallocation;
19use crate::buffer::Buffer;
20use crate::native::ArrowNativeType;
21use crate::{BufferBuilder, MutableBuffer, OffsetBuffer};
22use std::fmt::Formatter;
23use std::marker::PhantomData;
24use std::ops::Deref;
25
26/// A strongly-typed [`Buffer`] supporting zero-copy cloning and slicing
27///
28/// The easiest way to think about `ScalarBuffer<T>` is being equivalent to a `Arc<Vec<T>>`,
29/// with the following differences:
30///
31/// - slicing and cloning is O(1).
32/// - support for external allocated memory (e.g. via FFI).
33///
34/// See [`Buffer`] for more low-level memory management details.
35///
36/// # Example: Convert to/from Vec (without copies)
37///
38/// (See [`Buffer::from_vec`] and [`Buffer::into_vec`] for a lower level API)
39/// ```
40/// # use arrow_buffer::ScalarBuffer;
41/// // Zero-copy conversion from Vec
42/// let buffer = ScalarBuffer::from(vec![1, 2, 3]);
43/// assert_eq!(&buffer, &[1, 2, 3]);
44/// // convert the buffer back to Vec without copy assuming:
45/// // 1. the inner buffer is not sliced
46/// // 2. the inner buffer uses standard allocation
47/// // 3. there are no other references to the inner buffer
48/// let vec: Vec<i32> = buffer.into();
49/// assert_eq!(&vec, &[1, 2, 3]);
50/// ```
51///
52/// # Example: Zero copy slicing
53/// ```
54/// # use arrow_buffer::ScalarBuffer;
55/// let buffer = ScalarBuffer::from(vec![1, 2, 3]);
56/// assert_eq!(&buffer, &[1, 2, 3]);
57/// // Zero-copy slicing
58/// let sliced = buffer.slice(1, 2);
59/// assert_eq!(&sliced, &[2, 3]);
60/// // Original buffer is unchanged
61/// assert_eq!(&buffer, &[1, 2, 3]);
62/// // converting the sliced buffer back to Vec incurs a copy
63/// let vec: Vec<i32> = sliced.into();
64/// ```
65#[derive(Clone, Default)]
66pub struct ScalarBuffer<T: ArrowNativeType> {
67    /// Underlying data buffer
68    buffer: Buffer,
69    phantom: PhantomData<T>,
70}
71
72impl<T: ArrowNativeType> std::fmt::Debug for ScalarBuffer<T> {
73    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
74        f.debug_tuple("ScalarBuffer").field(&self.as_ref()).finish()
75    }
76}
77
78impl<T: ArrowNativeType> ScalarBuffer<T> {
79    /// Create a new [`ScalarBuffer`] from a [`Buffer`], and an `offset`
80    /// and `length` in units of `T`
81    ///
82    /// # Panics
83    ///
84    /// This method will panic if
85    ///
86    /// * `offset` or `len` would result in overflow
87    /// * `buffer` is not aligned to a multiple of `std::mem::align_of::<T>`
88    /// * `bytes` is not large enough for the requested slice
89    pub fn new(buffer: Buffer, offset: usize, len: usize) -> Self {
90        let size = std::mem::size_of::<T>();
91        let byte_offset = offset.checked_mul(size).expect("offset overflow");
92        let byte_len = len.checked_mul(size).expect("length overflow");
93        buffer.slice_with_length(byte_offset, byte_len).into()
94    }
95
96    /// Unsafe function to create a new [`ScalarBuffer`] from a [`Buffer`].
97    /// Only use for testing purpose.
98    ///
99    /// # Safety
100    ///
101    /// This function is unsafe because it does not check if the `buffer` is aligned
102    pub unsafe fn new_unchecked(buffer: Buffer) -> Self {
103        Self {
104            buffer,
105            phantom: Default::default(),
106        }
107    }
108
109    /// Free up unused memory.
110    pub fn shrink_to_fit(&mut self) {
111        self.buffer.shrink_to_fit();
112    }
113
114    /// Returns a zero-copy slice of this buffer with length `len` and starting at `offset`
115    ///
116    /// # Panics
117    ///
118    /// Panics for the same reasons as [`Self::new`]
119    pub fn slice(&self, offset: usize, len: usize) -> Self {
120        Self::new(self.buffer.clone(), offset, len)
121    }
122
123    /// Returns the inner [`Buffer`]
124    pub fn inner(&self) -> &Buffer {
125        &self.buffer
126    }
127
128    /// Returns the inner [`Buffer`], consuming self
129    pub fn into_inner(self) -> Buffer {
130        self.buffer
131    }
132
133    /// Claim memory used by this buffer in the provided memory pool.
134    ///
135    /// See [`Buffer::claim`] for details.
136    #[cfg(feature = "pool")]
137    pub fn claim(&self, pool: &dyn crate::MemoryPool) {
138        self.buffer.claim(pool);
139    }
140
141    /// Returns true if this [`ScalarBuffer`] is equal to `other`, using pointer comparisons
142    /// to determine buffer equality. This is cheaper than `PartialEq::eq` but may
143    /// return false when the arrays are logically equal
144    #[inline]
145    pub fn ptr_eq(&self, other: &Self) -> bool {
146        self.buffer.ptr_eq(&other.buffer)
147    }
148
149    /// Returns the number of elements in the buffer
150    pub fn len(&self) -> usize {
151        self.buffer.len() / std::mem::size_of::<T>()
152    }
153
154    /// Returns if the buffer is empty
155    pub fn is_empty(&self) -> bool {
156        self.len() == 0
157    }
158}
159
160impl<T: ArrowNativeType> Deref for ScalarBuffer<T> {
161    type Target = [T];
162
163    #[inline]
164    fn deref(&self) -> &Self::Target {
165        // SAFETY: Verified alignment in From<Buffer>
166        unsafe {
167            std::slice::from_raw_parts(
168                self.buffer.as_ptr().cast::<T>(),
169                self.buffer.len() / std::mem::size_of::<T>(),
170            )
171        }
172    }
173}
174
175impl<T: ArrowNativeType> AsRef<[T]> for ScalarBuffer<T> {
176    #[inline]
177    fn as_ref(&self) -> &[T] {
178        self
179    }
180}
181
182impl<T: ArrowNativeType> From<MutableBuffer> for ScalarBuffer<T> {
183    fn from(value: MutableBuffer) -> Self {
184        Buffer::from(value).into()
185    }
186}
187
188impl<T: ArrowNativeType> From<Buffer> for ScalarBuffer<T> {
189    fn from(buffer: Buffer) -> Self {
190        let align = std::mem::align_of::<T>();
191        let is_aligned = buffer.as_ptr().align_offset(align) == 0;
192
193        match buffer.deallocation() {
194            Deallocation::Standard(_) => assert!(
195                is_aligned,
196                "Memory pointer is not aligned with the specified scalar type"
197            ),
198            Deallocation::Custom(_, _) => assert!(
199                is_aligned,
200                "Memory pointer from external source (e.g, FFI) is not aligned with the specified scalar type. Before importing buffer through FFI, please make sure the allocation is aligned."
201            ),
202        }
203
204        Self {
205            buffer,
206            phantom: Default::default(),
207        }
208    }
209}
210
211impl<T: ArrowNativeType> From<OffsetBuffer<T>> for ScalarBuffer<T> {
212    fn from(value: OffsetBuffer<T>) -> Self {
213        value.into_inner()
214    }
215}
216
217impl<T: ArrowNativeType> From<Vec<T>> for ScalarBuffer<T> {
218    fn from(value: Vec<T>) -> Self {
219        Self {
220            buffer: Buffer::from_vec(value),
221            phantom: Default::default(),
222        }
223    }
224}
225
226impl<T: ArrowNativeType> From<ScalarBuffer<T>> for Vec<T> {
227    fn from(value: ScalarBuffer<T>) -> Self {
228        value
229            .buffer
230            .into_vec()
231            .unwrap_or_else(|buffer| buffer.typed_data::<T>().into())
232    }
233}
234
235impl<T: ArrowNativeType> From<BufferBuilder<T>> for ScalarBuffer<T> {
236    fn from(mut value: BufferBuilder<T>) -> Self {
237        let len = value.len();
238        Self::new(value.finish(), 0, len)
239    }
240}
241
242impl<T: ArrowNativeType> FromIterator<T> for ScalarBuffer<T> {
243    #[inline]
244    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
245        iter.into_iter().collect::<Vec<_>>().into()
246    }
247}
248
249impl<'a, T: ArrowNativeType> IntoIterator for &'a ScalarBuffer<T> {
250    type Item = &'a T;
251    type IntoIter = std::slice::Iter<'a, T>;
252
253    fn into_iter(self) -> Self::IntoIter {
254        self.as_ref().iter()
255    }
256}
257
258impl<T: ArrowNativeType, S: AsRef<[T]> + ?Sized> PartialEq<S> for ScalarBuffer<T> {
259    fn eq(&self, other: &S) -> bool {
260        self.as_ref().eq(other.as_ref())
261    }
262}
263
264impl<T: ArrowNativeType, const N: usize> PartialEq<ScalarBuffer<T>> for [T; N] {
265    fn eq(&self, other: &ScalarBuffer<T>) -> bool {
266        self.as_ref().eq(other.as_ref())
267    }
268}
269
270impl<T: ArrowNativeType> PartialEq<ScalarBuffer<T>> for [T] {
271    fn eq(&self, other: &ScalarBuffer<T>) -> bool {
272        self.as_ref().eq(other.as_ref())
273    }
274}
275
276impl<T: ArrowNativeType> PartialEq<ScalarBuffer<T>> for Vec<T> {
277    fn eq(&self, other: &ScalarBuffer<T>) -> bool {
278        self.as_slice().eq(other.as_ref())
279    }
280}
281
282/// If T implements Eq, then so does ScalarBuffer.
283impl<T: ArrowNativeType + Eq> Eq for ScalarBuffer<T> {}
284
285#[cfg(test)]
286mod tests {
287    use std::{ptr::NonNull, sync::Arc};
288
289    use super::*;
290
291    #[test]
292    fn test_basic() {
293        let expected = [0_i32, 1, 2];
294        let buffer = Buffer::from_iter(expected.iter().copied());
295        let typed = ScalarBuffer::<i32>::new(buffer.clone(), 0, 3);
296        assert_eq!(*typed, expected);
297
298        let typed = ScalarBuffer::<i32>::new(buffer.clone(), 1, 2);
299        assert_eq!(*typed, expected[1..]);
300
301        let typed = ScalarBuffer::<i32>::new(buffer.clone(), 1, 0);
302        assert!(typed.is_empty());
303
304        let typed = ScalarBuffer::<i32>::new(buffer, 3, 0);
305        assert!(typed.is_empty());
306    }
307
308    #[test]
309    fn test_debug() {
310        let buffer = ScalarBuffer::from(vec![1, 2, 3]);
311        assert_eq!(format!("{buffer:?}"), "ScalarBuffer([1, 2, 3])");
312    }
313
314    #[test]
315    #[should_panic(expected = "Memory pointer is not aligned with the specified scalar type")]
316    fn test_unaligned() {
317        let expected = [0_i32, 1, 2];
318        let buffer = Buffer::from_iter(expected.iter().copied());
319        let buffer = buffer.slice(1);
320        ScalarBuffer::<i32>::new(buffer, 0, 2);
321    }
322
323    #[test]
324    #[should_panic(expected = "the offset of the new Buffer cannot exceed the existing length")]
325    fn test_length_out_of_bounds() {
326        let buffer = Buffer::from_iter([0_i32, 1, 2]);
327        ScalarBuffer::<i32>::new(buffer, 1, 3);
328    }
329
330    #[test]
331    #[should_panic(expected = "the offset of the new Buffer cannot exceed the existing length")]
332    fn test_offset_out_of_bounds() {
333        let buffer = Buffer::from_iter([0_i32, 1, 2]);
334        ScalarBuffer::<i32>::new(buffer, 4, 0);
335    }
336
337    #[test]
338    #[should_panic(expected = "offset overflow")]
339    fn test_length_overflow() {
340        let buffer = Buffer::from_iter([0_i32, 1, 2]);
341        ScalarBuffer::<i32>::new(buffer, usize::MAX, 1);
342    }
343
344    #[test]
345    #[should_panic(expected = "offset overflow")]
346    fn test_start_overflow() {
347        let buffer = Buffer::from_iter([0_i32, 1, 2]);
348        ScalarBuffer::<i32>::new(buffer, usize::MAX / 4 + 1, 0);
349    }
350
351    #[test]
352    #[should_panic(expected = "length overflow")]
353    fn test_end_overflow() {
354        let buffer = Buffer::from_iter([0_i32, 1, 2]);
355        ScalarBuffer::<i32>::new(buffer, 0, usize::MAX / 4 + 1);
356    }
357
358    #[test]
359    fn convert_from_buffer_builder() {
360        let input = vec![1, 2, 3, 4];
361        let buffer_builder = BufferBuilder::from(input.clone());
362        let scalar_buffer = ScalarBuffer::from(buffer_builder);
363        assert_eq!(scalar_buffer.as_ref(), input);
364    }
365
366    #[test]
367    fn into_vec() {
368        let input = vec![1u8, 2, 3, 4];
369
370        // No copy
371        let input_buffer = Buffer::from_vec(input.clone());
372        let input_ptr = input_buffer.as_ptr();
373        let input_len = input_buffer.len();
374        let scalar_buffer = ScalarBuffer::<u8>::new(input_buffer, 0, input_len);
375        let vec = Vec::from(scalar_buffer);
376        assert_eq!(vec.as_slice(), input.as_slice());
377        assert_eq!(vec.as_ptr(), input_ptr);
378
379        // Custom allocation - makes a copy
380        let mut input_clone = input.clone();
381        let input_ptr = NonNull::new(input_clone.as_mut_ptr()).unwrap();
382        let dealloc = Arc::new(());
383        let buffer =
384            unsafe { Buffer::from_custom_allocation(input_ptr, input_clone.len(), dealloc as _) };
385        let scalar_buffer = ScalarBuffer::<u8>::new(buffer, 0, input.len());
386        let vec = Vec::from(scalar_buffer);
387        assert_eq!(vec, input.as_slice());
388        assert_ne!(vec.as_ptr(), input_ptr.as_ptr());
389
390        // Offset - makes a copy
391        let input_buffer = Buffer::from_vec(input.clone());
392        let input_ptr = input_buffer.as_ptr();
393        let input_len = input_buffer.len();
394        let scalar_buffer = ScalarBuffer::<u8>::new(input_buffer, 1, input_len - 1);
395        let vec = Vec::from(scalar_buffer);
396        assert_eq!(vec.as_slice(), &input[1..]);
397        assert_ne!(vec.as_ptr(), input_ptr);
398
399        // Inner buffer Arc ref count != 0 - makes a copy
400        let buffer = Buffer::from_slice_ref(input.as_slice());
401        let scalar_buffer = ScalarBuffer::<u8>::new(buffer, 0, input.len());
402        let vec = Vec::from(scalar_buffer);
403        assert_eq!(vec, input.as_slice());
404        assert_ne!(vec.as_ptr(), input.as_ptr());
405    }
406
407    #[test]
408    fn scalar_buffer_impl_eq() {
409        fn are_equal<T: Eq>(a: &T, b: &T) -> bool {
410            a.eq(b)
411        }
412
413        assert!(
414            are_equal(
415                &ScalarBuffer::<i16>::from(vec![23]),
416                &ScalarBuffer::<i16>::from(vec![23])
417            ),
418            "ScalarBuffer should implement Eq if the inner type does"
419        );
420    }
421}
arrow_buffer/buffer/scalar.rs

arrow_buffer/buffer/
scalar.rs
