## Coverage Report

#### Created: 2025-11-17 14:14

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /Users/andrewlamb/Software/arrow-rs/arrow-buffer/src/buffer/immutable.rs

    Line
    Count
    Source
    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
    18
    use std::alloc::Layout;
    19
    use std::fmt::Debug;
    20
    use std::ptr::NonNull;
    21
    use std::sync::Arc;
    22
    23
    use crate::BufferBuilder;
    24
    use crate::alloc::{Allocation, Deallocation};
    25
    use crate::util::bit_chunk_iterator::{BitChunks, UnalignedBitChunk};
    26
    use crate::{bit_util, bytes::Bytes, native::ArrowNativeType};
    27
    28
    #[cfg(feature = "pool")]
    29
    use crate::pool::MemoryPool;
    30
    31
    use super::ops::bitwise_unary_op_helper;
    32
    use super::{MutableBuffer, ScalarBuffer};
    33
    34
    /// A contiguous memory region that can be shared with other buffers and across
    35
    /// thread boundaries that stores Arrow data.
    36
    ///
    37
    /// `Buffer`s can be sliced and cloned without copying the underlying data and can
    38
    /// be created from memory allocated by non-Rust sources such as C/C++.
    39
    ///
    40
    /// # Example: Create a `Buffer` from a `Vec` (without copying)
    41
    /// ```
    42
    /// # use arrow_buffer::Buffer;
    43
    /// let vec: Vec<u32> = vec![1, 2, 3];
    44
    /// let buffer = Buffer::from(vec);
    45
    /// ```
    46
    ///
    47
    /// # Example: Convert a `Buffer` to a `Vec` (without copying)
    48
    ///
    49
    /// Use [`Self::into_vec`] to convert a `Buffer` back into a `Vec` if there are
    50
    /// no other references and the types are aligned correctly.
    51
    /// ```
    52
    /// # use arrow_buffer::Buffer;
    53
    /// # let vec: Vec<u32> = vec![1, 2, 3];
    54
    /// # let buffer = Buffer::from(vec);
    55
    /// // convert the buffer back into a Vec of u32
    56
    /// // note this will fail if the buffer is shared or not aligned correctly
    57
    /// let vec: Vec<u32> = buffer.into_vec().unwrap();
    58
    /// ```
    59
    ///
    60
    /// # Example: Create a `Buffer` from a [`bytes::Bytes`] (without copying)
    61
    ///
    62
    /// [`bytes::Bytes`] is a common type in the Rust ecosystem for shared memory
    63
    /// regions. You can create a buffer from a `Bytes` instance using the `From`
    64
    /// implementation, also without copying.
    65
    ///
    66
    /// ```
    67
    /// # use arrow_buffer::Buffer;
    68
    /// let bytes = bytes::Bytes::from("hello");
    69
    /// let buffer = Buffer::from(bytes);
    70
    ///```
    71
    #[derive(Clone, Debug)]
    72
    pub struct Buffer {
    73
        /// the internal byte buffer.
    74
        data: Arc<Bytes>,
    75
    76
        /// Pointer into `data` valid
    77
        ///
    78
        /// We store a pointer instead of an offset to avoid pointer arithmetic
    79
        /// which causes LLVM to fail to vectorise code correctly
    80
        ptr: *const u8,
    81
    82
        /// Byte length of the buffer.
    83
        ///
    84
        /// Must be less than or equal to `data.len()`
    85
        length: usize,
    86
    }
    87
    88
    impl Default for Buffer {
    89
        #[inline]
    90
        fn default() -> Self {
    91
            MutableBuffer::default().into()
    92
        }
    93
    }
    94
    95
    impl PartialEq for Buffer {
    96
    0
        fn eq(&self, other: &Self) -> bool {
    97
    0
            self.as_slice().eq(other.as_slice())
    98
    0
        }
    99
    }
    100
    101
    impl Eq for Buffer {}
    102
    103
    unsafe impl Send for Buffer where Bytes: Send {}
    104
    unsafe impl Sync for Buffer where Bytes: Sync {}
    105
    106
    impl Buffer {
    107
        /// Create a new Buffer from a (internal) `Bytes`
    108
        ///
    109
        /// NOTE despite the same name, `Bytes` is an internal struct in arrow-rs
    110
        /// and is different than [`bytes::Bytes`].
    111
        ///
    112
        /// See examples on [`Buffer`] for ways to create a buffer from a [`bytes::Bytes`].
    113
        #[deprecated(since = "54.1.0", note = "Use Buffer::from instead")]
    114
    0
        pub fn from_bytes(bytes: Bytes) -> Self {
    115
    0
            Self::from(bytes)
    116
    0
        }
    117
    118
        /// Returns the offset, in bytes, of `Self::ptr` to `Self::data`
    119
        ///
    120
        /// self.ptr and self.data can be different after slicing or advancing the buffer.
    121
    0
        pub fn ptr_offset(&self) -> usize {
    122
            // Safety: `ptr` is always in bounds of `data`.
    123
    0
            unsafe { self.ptr.offset_from(self.data.ptr().as_ptr()) as usize }
    124
    0
        }
    125
    126
        /// Returns the pointer to the start of the buffer without the offset.
    127
    0
        pub fn data_ptr(&self) -> NonNull<u8> {
    128
    0
            self.data.ptr()
    129
    0
        }
    130
    131
        /// Returns the number of strong references to the buffer.
    132
        ///
    133
        /// This method is safe but if the buffer is shared across multiple threads
    134
        /// the underlying value could change between calling this method and using
    135
        /// the result.
    136
    0
        pub fn strong_count(&self) -> usize {
    137
    0
            Arc::strong_count(&self.data)
    138
    0
        }
    139
    140
        /// Create a [`Buffer`] from the provided [`Vec`] without copying
    141
        #[inline]
    142
    324k
        pub fn from_vec<T: ArrowNativeType>(vec: Vec<T>) -> Self {
    143
    324k
            MutableBuffer::from(vec).into()
    144
    324k
        }
    145
    146
        /// Initializes a [Buffer] from a slice of items.
    147
    58
        pub fn from_slice_ref<U: ArrowNativeType, T: AsRef<[U]>>(items: T) -> Self {
    148
    58
            let slice = items.as_ref();
    149
    58
            let capacity = std::mem::size_of_val(slice);
    150
    58
            let mut buffer = MutableBuffer::with_capacity(capacity);
    151
    58
            buffer.extend_from_slice(slice);
    152
    58
            buffer.into()
    153
    58
        }
    154
    155
        /// Creates a buffer from an existing memory region.
    156
        ///
    157
        /// Ownership of the memory is tracked via reference counting
    158
        /// and the memory will be freed using the `drop` method of
    159
        /// [crate::alloc::Allocation] when the reference count reaches zero.
    160
        ///
    161
        /// # Arguments
    162
        ///
    163
        /// * `ptr` - Pointer to raw parts
    164
        /// * `len` - Length of raw parts in **bytes**
    165
        /// * `owner` - A [crate::alloc::Allocation] which is responsible for freeing that data
    166
        ///
    167
        /// # Safety
    168
        ///
    169
        /// This function is unsafe as there is no guarantee that the given pointer is valid for `len` bytes
    170
    0
        pub unsafe fn from_custom_allocation(
    171
    0
            ptr: NonNull<u8>,
    172
    0
            len: usize,
    173
    0
            owner: Arc<dyn Allocation>,
    174
    0
        ) -> Self {
    175
    0
            unsafe { Buffer::build_with_arguments(ptr, len, Deallocation::Custom(owner, len)) }
    176
    0
        }
    177
    178
        /// Auxiliary method to create a new Buffer
    179
    0
        unsafe fn build_with_arguments(
    180
    0
            ptr: NonNull<u8>,
    181
    0
            len: usize,
    182
    0
            deallocation: Deallocation,
    183
    0
        ) -> Self {
    184
    0
            let bytes = unsafe { Bytes::new(ptr, len, deallocation) };
    185
    0
            let ptr = bytes.as_ptr();
    186
    0
            Buffer {
    187
    0
                ptr,
    188
    0
                data: Arc::new(bytes),
    189
    0
                length: len,
    190
    0
            }
    191
    0
        }
    192
    193
        /// Returns the number of bytes in the buffer
    194
        #[inline]
    195
    2.49M
        pub fn len(&self) -> usize {
    196
    2.49M
            self.length
    197
    2.49M
        }
    198
    199
        /// Returns the capacity of this buffer.
    200
        /// For externally owned buffers, this returns zero
    201
        #[inline]
    202
    174
        pub fn capacity(&self) -> usize {
    203
    174
            self.data.capacity()
    204
    174
        }
    205
    206
        /// Tries to shrink the capacity of the buffer as much as possible, freeing unused memory.
    207
        ///
    208
        /// If the buffer is shared, this is a no-op.
    209
        ///
    210
        /// If the memory was allocated with a custom allocator, this is a no-op.
    211
        ///
    212
        /// If the capacity is already less than or equal to the desired capacity, this is a no-op.
    213
        ///
    214
        /// The memory region will be reallocated using `std::alloc::realloc`.
    215
    0
        pub fn shrink_to_fit(&mut self) {
    216
    0
            let offset = self.ptr_offset();
    217
    0
            let is_empty = self.is_empty();
    218
    0
            let desired_capacity = if is_empty {
    219
    0
                0
    220
            } else {
    221
                // For realloc to work, we cannot free the elements before the offset
    222
    0
                offset + self.len()
    223
            };
    224
    0
            if desired_capacity < self.capacity() {
    225
    0
                if let Some(bytes) = Arc::get_mut(&mut self.data) {
    226
    0
                    if bytes.try_realloc(desired_capacity).is_ok() {
    227
                        // Realloc complete - update our pointer into `bytes`:
    228
    0
                        self.ptr = if is_empty {
    229
    0
                            bytes.as_ptr()
    230
                        } else {
    231
                            // SAFETY: we kept all elements leading up to the offset
    232
    0
                            unsafe { bytes.as_ptr().add(offset) }
    233
                        }
    234
    0
                    } else {
    235
    0
                        // Failure to reallocate is fine; we just failed to free up memory.
    236
    0
                    }
    237
    0
                }
    238
    0
            }
    239
    0
        }
    240
    241
        /// Returns true if the buffer is empty.
    242
        #[inline]
    243
    39
        pub fn is_empty(&self) -> bool {
    244
    39
            self.length == 0
    245
    39
        }
    246
    247
        /// Returns the byte slice stored in this buffer
    248
    571k
        pub fn as_slice(&self) -> &[u8] {
    249
    571k
            unsafe { std::slice::from_raw_parts(self.ptr, self.length) }
    250
    571k
        }
    251
    252
    565k
        pub(crate) fn deallocation(&self) -> &Deallocation {
    253
    565k
            self.data.deallocation()
    254
    565k
        }
    255
    256
        /// Returns a new [Buffer] that is a slice of this buffer starting at `offset`.
    257
        ///
    258
        /// This function is `O(1)` and does not copy any data, allowing the
    259
        /// same memory region to be shared between buffers.
    260
        ///
    261
        /// # Panics
    262
        ///
    263
        /// Panics iff `offset` is larger than `len`.
    264
    0
        pub fn slice(&self, offset: usize) -> Self {
    265
    0
            let mut s = self.clone();
    266
    0
            s.advance(offset);
    267
    0
            s
    268
    0
        }
    269
    270
        /// Increases the offset of this buffer by `offset`
    271
        ///
    272
        /// # Panics
    273
        ///
    274
        /// Panics iff `offset` is larger than `len`.
    275
        #[inline]
    276
    0
        pub fn advance(&mut self, offset: usize) {
    277
    0
            assert!(
    278
    0
                offset <= self.length,
    279
    0
                "the offset of the new Buffer cannot exceed the existing length: offset={} length={}",
    280
                offset,
    281
                self.length
    282
            );
    283
    0
            self.length -= offset;
    284
            // Safety:
    285
            // This cannot overflow as
    286
            // `self.offset + self.length < self.data.len()`
    287
            // `offset < self.length`
    288
    0
            self.ptr = unsafe { self.ptr.add(offset) };
    289
    0
        }
    290
    291
        /// Returns a new [Buffer] that is a slice of this buffer starting at `offset`,
    292
        /// with `length` bytes.
    293
        ///
    294
        /// This function is `O(1)` and does not copy any data, allowing the same
    295
        /// memory region to be shared between buffers.
    296
        ///
    297
        /// # Panics
    298
        /// Panics iff `(offset + length)` is larger than the existing length.
    299
    485k
        pub fn slice_with_length(&self, offset: usize, length: usize) -> Self {
    300
    485k
            assert!(
    301
    485k
                offset.saturating_add(length) <= self.length,
    302
    0
                "the offset of the new Buffer cannot exceed the existing length: slice offset={offset} length={length} selflen={}",
    303
                self.length
    304
            );
    305
            // Safety:
    306
            // offset + length <= self.length
    307
    485k
            let ptr = unsafe { self.ptr.add(offset) };
    308
    485k
            Self {
    309
    485k
                data: self.data.clone(),
    310
    485k
                ptr,
    311
    485k
                length,
    312
    485k
            }
    313
    485k
        }
    314
    315
        /// Returns a pointer to the start of this buffer.
    316
        ///
    317
        /// Note that this should be used cautiously, and the returned pointer should not be
    318
        /// stored anywhere, to avoid dangling pointers.
    319
        #[inline]
    320
    2.79M
        pub fn as_ptr(&self) -> *const u8 {
    321
    2.79M
            self.ptr
    322
    2.79M
        }
    323
    324
        /// View buffer as a slice of a specific type.
    325
        ///
    326
        /// # Panics
    327
        ///
    328
        /// This function panics if the underlying buffer is not aligned
    329
        /// correctly for type `T`.
    330
    320k
        pub fn typed_data<T: ArrowNativeType>(&self) -> &[T] {
    331
            // SAFETY
    332
            // ArrowNativeType is trivially transmutable, is sealed to prevent potentially incorrect
    333
            // implementation outside this crate, and this method checks alignment
    334
    320k
            let (prefix, offsets, suffix) = unsafe { self.as_slice().align_to::<T>() };
    335
    320k
            assert!(prefix.is_empty() && suffix.is_empty());
    336
    320k
            offsets
    337
    320k
        }
    338
    339
        /// Returns a slice of this buffer starting at a certain bit offset.
    340
        /// If the offset is byte-aligned the returned buffer is a shallow clone,
    341
        /// otherwise a new buffer is allocated and filled with a copy of the bits in the range.
    342
    0
        pub fn bit_slice(&self, offset: usize, len: usize) -> Self {
    343
    0
            if offset % 8 == 0 {
    344
    0
                return self.slice_with_length(offset / 8, bit_util::ceil(len, 8));
    345
    0
            }
    346
    347
    0
            bitwise_unary_op_helper(self, offset, len, |a| a)
    348
    0
        }
    349
    350
        /// Returns a `BitChunks` instance which can be used to iterate over this buffers bits
    351
        /// in larger chunks and starting at arbitrary bit offsets.
    352
        /// Note that both `offset` and `length` are measured in bits.
    353
    1.84k
        pub fn bit_chunks(&self, offset: usize, len: usize) -> BitChunks<'_> {
    354
    1.84k
            BitChunks::new(self.as_slice(), offset, len)
    355
    1.84k
        }
    356
    357
        /// Returns the number of 1-bits in this buffer, starting from `offset` with `length` bits
    358
        /// inspected. Note that both `offset` and `length` are measured in bits.
    359
    86.6k
        pub fn count_set_bits_offset(&self, offset: usize, len: usize) -> usize {
    360
    86.6k
            UnalignedBitChunk::new(self.as_slice(), offset, len).count_ones()
    361
    86.6k
        }
    362
    363
        /// Returns `MutableBuffer` for mutating the buffer if this buffer is not shared.
    364
        /// Returns `Err` if this is shared or its allocation is from an external source or
    365
        /// it is not allocated with alignment [`ALIGNMENT`]
    366
        ///
    367
        /// [`ALIGNMENT`]: crate::alloc::ALIGNMENT
    368
    0
        pub fn into_mutable(self) -> Result<MutableBuffer, Self> {
    369
    0
            let ptr = self.ptr;
    370
    0
            let length = self.length;
    371
    0
            Arc::try_unwrap(self.data)
    372
    0
                .and_then(|bytes| {
    373
                    // The pointer of underlying buffer should not be offset.
    374
    0
                    assert_eq!(ptr, bytes.ptr().as_ptr());
    375
    0
                    MutableBuffer::from_bytes(bytes).map_err(Arc::new)
    376
    0
                })
    377
    0
                .map_err(|bytes| Buffer {
    378
    0
                    data: bytes,
    379
    0
                    ptr,
    380
    0
                    length,
    381
    0
                })
    382
    0
        }
    383
    384
        /// Converts self into a `Vec`, if possible.
    385
        ///
    386
        /// This can be used to reuse / mutate the underlying data.
    387
        ///
    388
        /// # Errors
    389
        ///
    390
        /// Returns `Err(self)` if
    391
        /// 1. The buffer does not have the same [`Layout`] as the destination Vec
    392
        /// 2. The buffer contains a non-zero offset
    393
        /// 3. The buffer is shared
    394
        pub fn into_vec<T: ArrowNativeType>(self) -> Result<Vec<T>, Self> {
    395
            let layout = match self.data.deallocation() {
    396
                Deallocation::Standard(l) => l,
    397
                _ => return Err(self), // Custom allocation
    398
            };
    399
    400
            if self.ptr != self.data.as_ptr() {
    401
                return Err(self); // Data is offset
    402
            }
    403
    404
            let v_capacity = layout.size() / std::mem::size_of::<T>();
    405
            match Layout::array::<T>(v_capacity) {
    406
                Ok(expected) if layout == &expected => {}
    407
                _ => return Err(self), // Incorrect layout
    408
            }
    409
    410
            let length = self.length;
    411
            let ptr = self.ptr;
    412
            let v_len = self.length / std::mem::size_of::<T>();
    413
    414
            Arc::try_unwrap(self.data)
    415
                .map(|bytes| unsafe {
    416
                    let ptr = bytes.ptr().as_ptr() as _;
    417
                    std::mem::forget(bytes);
    418
                    // Safety
    419
                    // Verified that bytes layout matches that of Vec
    420
                    Vec::from_raw_parts(ptr, v_len, v_capacity)
    421
                })
    422
                .map_err(|bytes| Buffer {
    423
                    data: bytes,
    424
                    ptr,
    425
                    length,
    426
                })
    427
        }
    428
    429
        /// Returns true if this [`Buffer`] is equal to `other`, using pointer comparisons
    430
        /// to determine buffer equality. This is cheaper than `PartialEq::eq` but may
    431
        /// return false when the arrays are logically equal
    432
        #[inline]
    433
    18
        pub fn ptr_eq(&self, other: &Self) -> bool {
    434
    18
            self.ptr == other.ptr && self.length == other.length4
    435
    18
        }
    436
    437
        /// Register this [`Buffer`] with the provided [`MemoryPool`]
    438
        ///
    439
        /// This claims the memory used by this buffer in the pool, allowing for
    440
        /// accurate accounting of memory usage. Any prior reservation will be
    441
        /// released so this works well when the buffer is being shared among
    442
        /// multiple arrays.
    443
        #[cfg(feature = "pool")]
    444
        pub fn claim(&self, pool: &dyn MemoryPool) {
    445
            self.data.claim(pool)
    446
        }
    447
    }
    448
    449
    /// Note that here we deliberately do not implement
    450
    /// `impl<T: AsRef<[u8]>> From<T> for Buffer`
    451
    /// As it would accept `Buffer::from(vec![...])` that would cause an unexpected copy.
    452
    /// Instead, we ask user to be explicit when copying is occurring, e.g., `Buffer::from(vec![...].to_byte_slice())`.
    453
    /// For zero-copy conversion, user should use `Buffer::from_vec(vec![...])`.
    454
    ///
    455
    /// Since we removed impl for `AsRef<u8>`, we added the following three specific implementations to reduce API breakage.
    456
    /// See <https://github.com/apache/arrow-rs/issues/6033> for more discussion on this.
    457
    impl From<&[u8]> for Buffer {
    458
    25
        fn from(p: &[u8]) -> Self {
    459
    25
            Self::from_slice_ref(p)
    460
    25
        }
    461
    }
    462
    463
    impl<const N: usize> From<[u8; N]> for Buffer {
    464
    9
        fn from(p: [u8; N]) -> Self {
    465
    9
            Self::from_slice_ref(p)
    466
    9
        }
    467
    }
    468
    469
    impl<const N: usize> From<&[u8; N]> for Buffer {
    470
    4
        fn from(p: &[u8; N]) -> Self {
    471
    4
            Self::from_slice_ref(p)
    472
    4
        }
    473
    }
    474
    475
    impl<T: ArrowNativeType> From<Vec<T>> for Buffer {
    476
    241k
        fn from(value: Vec<T>) -> Self {
    477
    241k
            Self::from_vec(value)
    478
    241k
        }
    479
    }
    480
    481
    impl<T: ArrowNativeType> From<ScalarBuffer<T>> for Buffer {
    482
    16
        fn from(value: ScalarBuffer<T>) -> Self {
    483
    16
            value.into_inner()
    484
    16
        }
    485
    }
    486
    487
    /// Convert from internal `Bytes` (not [`bytes::Bytes`]) to `Buffer`
    488
    impl From<Bytes> for Buffer {
    489
        #[inline]
    490
    413k
        fn from(bytes: Bytes) -> Self {
    491
    413k
            let length = bytes.len();
    492
    413k
            let ptr = bytes.as_ptr();
    493
    413k
            Self {
    494
    413k
                data: Arc::new(bytes),
    495
    413k
                ptr,
    496
    413k
                length,
    497
    413k
            }
    498
    413k
        }
    499
    }
    500
    501
    /// Convert from [`bytes::Bytes`], not internal `Bytes` to `Buffer`
    502
    impl From<bytes::Bytes> for Buffer {
    503
    0
        fn from(bytes: bytes::Bytes) -> Self {
    504
    0
            let bytes: Bytes = bytes.into();
    505
    0
            Self::from(bytes)
    506
    0
        }
    507
    }
    508
    509
    /// Create a `Buffer` instance by storing the boolean values into the buffer
    510
    impl FromIterator<bool> for Buffer {
    511
    112
        fn from_iter<I>(iter: I) -> Self
    512
    112
        where
    513
    112
            I: IntoIterator<Item = bool>,
    514
        {
    515
    112
            MutableBuffer::from_iter(iter).into()
    516
    112
        }
    517
    }
    518
    519
    impl std::ops::Deref for Buffer {
    520
        type Target = [u8];
    521
    522
    130k
        fn deref(&self) -> &[u8] {
    523
    130k
            unsafe { std::slice::from_raw_parts(self.as_ptr(), self.len()) }
    524
    130k
        }
    525
    }
    526
    527
    impl AsRef<[u8]> for &Buffer {
    528
    0
        fn as_ref(&self) -> &[u8] {
    529
    0
            self.as_slice()
    530
    0
        }
    531
    }
    532
    533
    impl From<MutableBuffer> for Buffer {
    534
        #[inline]
    535
    413k
        fn from(buffer: MutableBuffer) -> Self {
    536
    413k
            buffer.into_buffer()
    537
    413k
        }
    538
    }
    539
    540
    impl<T: ArrowNativeType> From<BufferBuilder<T>> for Buffer {
    541
        fn from(mut value: BufferBuilder<T>) -> Self {
    542
            value.finish()
    543
        }
    544
    }
    545
    546
    impl Buffer {
    547
        /// Creates a [`Buffer`] from an [`Iterator`] with a trusted (upper) length.
    548
        ///
    549
        /// Prefer this to `collect` whenever possible, as it is ~60% faster.
    550
        ///
    551
        /// # Example
    552
        /// ```
    553
        /// # use arrow_buffer::buffer::Buffer;
    554
        /// let v = vec![1u32];
    555
        /// let iter = v.iter().map(|x| x * 2);
    556
        /// let buffer = unsafe { Buffer::from_trusted_len_iter(iter) };
    557
        /// assert_eq!(buffer.len(), 4) // u32 has 4 bytes
    558
        /// ```
    559
        /// # Safety
    560
        /// This method assumes that the iterator's size is correct and is undefined behavior
    561
        /// to use it on an iterator that reports an incorrect length.
    562
        // This implementation is required for two reasons:
    563
        // 1. there is no trait `TrustedLen` in stable rust and therefore
    564
        //    we can't specialize `extend` for `TrustedLen` like `Vec` does.
    565
        // 2. `from_trusted_len_iter` is faster.
    566
        #[inline]
    567
        pub unsafe fn from_trusted_len_iter<T: ArrowNativeType, I: Iterator<Item = T>>(
    568
            iterator: I,
    569
        ) -> Self {
    570
            unsafe { MutableBuffer::from_trusted_len_iter(iterator).into() }
    571
        }
    572
    573
        /// Creates a [`Buffer`] from an [`Iterator`] with a trusted (upper) length or errors
    574
        /// if any of the items of the iterator is an error.
    575
        /// Prefer this to `collect` whenever possible, as it is ~60% faster.
    576
        /// # Safety
    577
        /// This method assumes that the iterator's size is correct and is undefined behavior
    578
        /// to use it on an iterator that reports an incorrect length.
    579
        #[inline]
    580
        pub unsafe fn try_from_trusted_len_iter<
    581
            E,
    582
            T: ArrowNativeType,
    583
            I: Iterator<Item = Result<T, E>>,
    584
        >(
    585
            iterator: I,
    586
        ) -> Result<Self, E> {
    587
            unsafe { Ok(MutableBuffer::try_from_trusted_len_iter(iterator)?.into()) }
    588
        }
    589
    }
    590
    591
    impl<T: ArrowNativeType> FromIterator<T> for Buffer {
    592
    1.25k
        fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
    593
    1.25k
            let vec = Vec::from_iter(iter);
    594
    1.25k
            Buffer::from_vec(vec)
    595
    1.25k
        }
    596
    }
    597
    598
    #[cfg(test)]
    599
    mod tests {
    600
        use crate::i256;
    601
        use std::panic::{RefUnwindSafe, UnwindSafe};
    602
        use std::thread;
    603
    604
        use super::*;
    605
    606
        #[test]
    607
        fn test_buffer_data_equality() {
    608
            let buf1 = Buffer::from(&[0, 1, 2, 3, 4]);
    609
            let buf2 = Buffer::from(&[0, 1, 2, 3, 4]);
    610
            assert_eq!(buf1, buf2);
    611
    612
            // slice with same offset and same length should still preserve equality
    613
            let buf3 = buf1.slice(2);
    614
            assert_ne!(buf1, buf3);
    615
            let buf4 = buf2.slice_with_length(2, 3);
    616
            assert_eq!(buf3, buf4);
    617
    618
            // Different capacities should still preserve equality
    619
            let mut buf2 = MutableBuffer::new(65);
    620
            buf2.extend_from_slice(&[0u8, 1, 2, 3, 4]);
    621
    622
            let buf2 = buf2.into();
    623
            assert_eq!(buf1, buf2);
    624
    625
            // unequal because of different elements
    626
            let buf2 = Buffer::from(&[0, 0, 2, 3, 4]);
    627
            assert_ne!(buf1, buf2);
    628
    629
            // unequal because of different length
    630
            let buf2 = Buffer::from(&[0, 1, 2, 3]);
    631
            assert_ne!(buf1, buf2);
    632
        }
    633
    634
        #[test]
    635
        fn test_from_raw_parts() {
    636
            let buf = Buffer::from(&[0, 1, 2, 3, 4]);
    637
            assert_eq!(5, buf.len());
    638
            assert!(!buf.as_ptr().is_null());
    639
            assert_eq!([0, 1, 2, 3, 4], buf.as_slice());
    640
        }
    641
    642
        #[test]
    643
        fn test_from_vec() {
    644
            let buf = Buffer::from(&[0, 1, 2, 3, 4]);
    645
            assert_eq!(5, buf.len());
    646
            assert!(!buf.as_ptr().is_null());
    647
            assert_eq!([0, 1, 2, 3, 4], buf.as_slice());
    648
        }
    649
    650
        #[test]
    651
        fn test_copy() {
    652
            let buf = Buffer::from(&[0, 1, 2, 3, 4]);
    653
            let buf2 = buf;
    654
            assert_eq!(5, buf2.len());
    655
            assert_eq!(64, buf2.capacity());
    656
            assert!(!buf2.as_ptr().is_null());
    657
            assert_eq!([0, 1, 2, 3, 4], buf2.as_slice());
    658
        }
    659
    660
        #[test]
    661
        fn test_slice() {
    662
            let buf = Buffer::from(&[2, 4, 6, 8, 10]);
    663
            let buf2 = buf.slice(2);
    664
    665
            assert_eq!([6, 8, 10], buf2.as_slice());
    666
            assert_eq!(3, buf2.len());
    667
            assert_eq!(unsafe { buf.as_ptr().offset(2) }, buf2.as_ptr());
    668
    669
            let buf3 = buf2.slice_with_length(1, 2);
    670
            assert_eq!([8, 10], buf3.as_slice());
    671
            assert_eq!(2, buf3.len());
    672
            assert_eq!(unsafe { buf.as_ptr().offset(3) }, buf3.as_ptr());
    673
    674
            let buf4 = buf.slice(5);
    675
            let empty_slice: [u8; 0] = [];
    676
            assert_eq!(empty_slice, buf4.as_slice());
    677
            assert_eq!(0, buf4.len());
    678
            assert!(buf4.is_empty());
    679
            assert_eq!(buf2.slice_with_length(2, 1).as_slice(), &[10]);
    680
        }
    681
    682
        #[test]
    683
        fn test_shrink_to_fit() {
    684
            let original = Buffer::from(&[0, 1, 2, 3, 4, 5, 6, 7]);
    685
            assert_eq!(original.as_slice(), &[0, 1, 2, 3, 4, 5, 6, 7]);
    686
            assert_eq!(original.capacity(), 64);
    687
    688
            let slice = original.slice_with_length(2, 3);
    689
            drop(original); // Make sure the buffer isn't shared (or shrink_to_fit won't work)
    690
            assert_eq!(slice.as_slice(), &[2, 3, 4]);
    691
            assert_eq!(slice.capacity(), 64);
    692
    693
            let mut shrunk = slice;
    694
            shrunk.shrink_to_fit();
    695
            assert_eq!(shrunk.as_slice(), &[2, 3, 4]);
    696
            assert_eq!(shrunk.capacity(), 5); // shrink_to_fit is allowed to keep the elements before the offset
    697
    698
            // Test that we can handle empty slices:
    699
            let empty_slice = shrunk.slice_with_length(1, 0);
    700
            drop(shrunk); // Make sure the buffer isn't shared (or shrink_to_fit won't work)
    701
            assert_eq!(empty_slice.as_slice(), &[]);
    702
            assert_eq!(empty_slice.capacity(), 5);
    703
    704
            let mut shrunk_empty = empty_slice;
    705
            shrunk_empty.shrink_to_fit();
    706
            assert_eq!(shrunk_empty.as_slice(), &[]);
    707
            assert_eq!(shrunk_empty.capacity(), 0);
    708
        }
    709
    710
        #[test]
    711
        #[should_panic(expected = "the offset of the new Buffer cannot exceed the existing length")]
    712
        fn test_slice_offset_out_of_bound() {
    713
            let buf = Buffer::from(&[2, 4, 6, 8, 10]);
    714
            buf.slice(6);
    715
        }
    716
    717
        #[test]
    718
        fn test_access_concurrently() {
    719
            let buffer = Buffer::from([1, 2, 3, 4, 5]);
    720
            let buffer2 = buffer.clone();
    721
            assert_eq!([1, 2, 3, 4, 5], buffer.as_slice());
    722
    723
            let buffer_copy = thread::spawn(move || {
    724
                // access buffer in another thread.
    725
                buffer
    726
            })
    727
            .join();
    728
    729
            assert!(buffer_copy.is_ok());
    730
            assert_eq!(buffer2, buffer_copy.ok().unwrap());
    731
        }
    732
    733
        macro_rules! check_as_typed_data {
    734
            ($input: expr, $native_t: ty) => {{
    735
                let buffer = Buffer::from_slice_ref($input);
    736
                let slice: &[$native_t] = buffer.typed_data::<$native_t>();
    737
                assert_eq!($input, slice);
    738
            }};
    739
        }
    740
    741
        #[test]
    742
        #[allow(clippy::float_cmp)]
    743
        fn test_as_typed_data() {
    744
            check_as_typed_data!(&[1i8, 3i8, 6i8], i8);
    745
            check_as_typed_data!(&[1u8, 3u8, 6u8], u8);
    746
            check_as_typed_data!(&[1i16, 3i16, 6i16], i16);
    747
            check_as_typed_data!(&[1i32, 3i32, 6i32], i32);
    748
            check_as_typed_data!(&[1i64, 3i64, 6i64], i64);
    749
            check_as_typed_data!(&[1u16, 3u16, 6u16], u16);
    750
            check_as_typed_data!(&[1u32, 3u32, 6u32], u32);
    751
            check_as_typed_data!(&[1u64, 3u64, 6u64], u64);
    752
            check_as_typed_data!(&[1f32, 3f32, 6f32], f32);
    753
            check_as_typed_data!(&[1f64, 3f64, 6f64], f64);
    754
        }
    755
    756
        #[test]
    757
        fn test_count_bits() {
    758
            assert_eq!(0, Buffer::from(&[0b00000000]).count_set_bits_offset(0, 8));
    759
            assert_eq!(8, Buffer::from(&[0b11111111]).count_set_bits_offset(0, 8));
    760
            assert_eq!(3, Buffer::from(&[0b00001101]).count_set_bits_offset(0, 8));
    761
            assert_eq!(
    762
                6,
    763
                Buffer::from(&[0b01001001, 0b01010010]).count_set_bits_offset(0, 16)
    764
            );
    765
            assert_eq!(
    766
                16,
    767
                Buffer::from(&[0b11111111, 0b11111111]).count_set_bits_offset(0, 16)
    768
            );
    769
        }
    770
    771
        #[test]
    772
        fn test_count_bits_slice() {
    773
            assert_eq!(
    774
                0,
    775
                Buffer::from(&[0b11111111, 0b00000000])
    776
                    .slice(1)
    777
                    .count_set_bits_offset(0, 8)
    778
            );
    779
            assert_eq!(
    780
                8,
    781
                Buffer::from(&[0b11111111, 0b11111111])
    782
                    .slice_with_length(1, 1)
    783
                    .count_set_bits_offset(0, 8)
    784
            );
    785
            assert_eq!(
    786
                3,
    787
                Buffer::from(&[0b11111111, 0b11111111, 0b00001101])
    788
                    .slice(2)
    789
                    .count_set_bits_offset(0, 8)
    790
            );
    791
            assert_eq!(
    792
                6,
    793
                Buffer::from(&[0b11111111, 0b01001001, 0b01010010])
    794
                    .slice_with_length(1, 2)
    795
                    .count_set_bits_offset(0, 16)
    796
            );
    797
            assert_eq!(
    798
                16,
    799
                Buffer::from(&[0b11111111, 0b11111111, 0b11111111, 0b11111111])
    800
                    .slice(2)
    801
                    .count_set_bits_offset(0, 16)
    802
            );
    803
        }
    804
    805
        #[test]
    806
        fn test_count_bits_offset_slice() {
    807
            assert_eq!(8, Buffer::from(&[0b11111111]).count_set_bits_offset(0, 8));
    808
            assert_eq!(3, Buffer::from(&[0b11111111]).count_set_bits_offset(0, 3));
    809
            assert_eq!(5, Buffer::from(&[0b11111111]).count_set_bits_offset(3, 5));
    810
            assert_eq!(1, Buffer::from(&[0b11111111]).count_set_bits_offset(3, 1));
    811
            assert_eq!(0, Buffer::from(&[0b11111111]).count_set_bits_offset(8, 0));
    812
            assert_eq!(2, Buffer::from(&[0b01010101]).count_set_bits_offset(0, 3));
    813
            assert_eq!(
    814
                16,
    815
                Buffer::from(&[0b11111111, 0b11111111]).count_set_bits_offset(0, 16)
    816
            );
    817
            assert_eq!(
    818
                10,
    819
                Buffer::from(&[0b11111111, 0b11111111]).count_set_bits_offset(0, 10)
    820
            );
    821
            assert_eq!(
    822
                10,
    823
                Buffer::from(&[0b11111111, 0b11111111]).count_set_bits_offset(3, 10)
    824
            );
    825
            assert_eq!(
    826
                8,
    827
                Buffer::from(&[0b11111111, 0b11111111]).count_set_bits_offset(8, 8)
    828
            );
    829
            assert_eq!(
    830
                5,
    831
                Buffer::from(&[0b11111111, 0b11111111]).count_set_bits_offset(11, 5)
    832
            );
    833
            assert_eq!(
    834
                0,
    835
                Buffer::from(&[0b11111111, 0b11111111]).count_set_bits_offset(16, 0)
    836
            );
    837
            assert_eq!(
    838
                2,
    839
                Buffer::from(&[0b01101101, 0b10101010]).count_set_bits_offset(7, 5)
    840
            );
    841
            assert_eq!(
    842
                4,
    843
                Buffer::from(&[0b01101101, 0b10101010]).count_set_bits_offset(7, 9)
    844
            );
    845
        }
    846
    847
        #[test]
    848
        fn test_unwind_safe() {
    849
            fn assert_unwind_safe<T: RefUnwindSafe + UnwindSafe>() {}
    850
            assert_unwind_safe::<Buffer>()
    851
        }
    852
    853
        #[test]
    854
        fn test_from_foreign_vec() {
    855
            let mut vector = vec![1_i32, 2, 3, 4, 5];
    856
            let buffer = unsafe {
    857
                Buffer::from_custom_allocation(
    858
                    NonNull::new_unchecked(vector.as_mut_ptr() as *mut u8),
    859
                    vector.len() * std::mem::size_of::<i32>(),
    860
                    Arc::new(vector),
    861
                )
    862
            };
    863
    864
            let slice = buffer.typed_data::<i32>();
    865
            assert_eq!(slice, &[1, 2, 3, 4, 5]);
    866
    867
            let buffer = buffer.slice(std::mem::size_of::<i32>());
    868
    869
            let slice = buffer.typed_data::<i32>();
    870
            assert_eq!(slice, &[2, 3, 4, 5]);
    871
        }
    872
    873
        #[test]
    874
        #[should_panic(expected = "the offset of the new Buffer cannot exceed the existing length")]
    875
        fn slice_overflow() {
    876
            let buffer = Buffer::from(MutableBuffer::from_len_zeroed(12));
    877
            buffer.slice_with_length(2, usize::MAX);
    878
        }
    879
    880
        #[test]
    881
        fn test_vec_interop() {
    882
            // Test empty vec
    883
            let a: Vec<i128> = Vec::new();
    884
            let b = Buffer::from_vec(a);
    885
            b.into_vec::<i128>().unwrap();
    886
    887
            // Test vec with capacity
    888
            let a: Vec<i128> = Vec::with_capacity(20);
    889
            let b = Buffer::from_vec(a);
    890
            let back = b.into_vec::<i128>().unwrap();
    891
            assert_eq!(back.len(), 0);
    892
            assert_eq!(back.capacity(), 20);
    893
    894
            // Test vec with values
    895
            let mut a: Vec<i128> = Vec::with_capacity(3);
    896
            a.extend_from_slice(&[1, 2, 3]);
    897
            let b = Buffer::from_vec(a);
    898
            let back = b.into_vec::<i128>().unwrap();
    899
            assert_eq!(back.len(), 3);
    900
            assert_eq!(back.capacity(), 3);
    901
    902
            // Test vec with values and spare capacity
    903
            let mut a: Vec<i128> = Vec::with_capacity(20);
    904
            a.extend_from_slice(&[1, 4, 7, 8, 9, 3, 6]);
    905
            let b = Buffer::from_vec(a);
    906
            let back = b.into_vec::<i128>().unwrap();
    907
            assert_eq!(back.len(), 7);
    908
            assert_eq!(back.capacity(), 20);
    909
    910
            // Test incorrect alignment
    911
            let a: Vec<i128> = Vec::new();
    912
            let b = Buffer::from_vec(a);
    913
            let b = b.into_vec::<i32>().unwrap_err();
    914
            b.into_vec::<i8>().unwrap_err();
    915
    916
            // Test convert between types with same alignment
    917
            // This is an implementation quirk, but isn't harmful
    918
            // as ArrowNativeType are trivially transmutable
    919
            let a: Vec<i64> = vec![1, 2, 3, 4];
    920
            let b = Buffer::from_vec(a);
    921
            let back = b.into_vec::<u64>().unwrap();
    922
            assert_eq!(back.len(), 4);
    923
            assert_eq!(back.capacity(), 4);
    924
    925
            // i256 has the same layout as i128 so this is valid
    926
            let mut b: Vec<i128> = Vec::with_capacity(4);
    927
            b.extend_from_slice(&[1, 2, 3, 4]);
    928
            let b = Buffer::from_vec(b);
    929
            let back = b.into_vec::<i256>().unwrap();
    930
            assert_eq!(back.len(), 2);
    931
            assert_eq!(back.capacity(), 2);
    932
    933
            // Invalid layout
    934
            let b: Vec<i128> = vec![1, 2, 3];
    935
            let b = Buffer::from_vec(b);
    936
            b.into_vec::<i256>().unwrap_err();
    937
    938
            // Invalid layout
    939
            let mut b: Vec<i128> = Vec::with_capacity(5);
    940
            b.extend_from_slice(&[1, 2, 3, 4]);
    941
            let b = Buffer::from_vec(b);
    942
            b.into_vec::<i256>().unwrap_err();
    943
    944
            // Truncates length
    945
            // This is an implementation quirk, but isn't harmful
    946
            let mut b: Vec<i128> = Vec::with_capacity(4);
    947
            b.extend_from_slice(&[1, 2, 3]);
    948
            let b = Buffer::from_vec(b);
    949
            let back = b.into_vec::<i256>().unwrap();
    950
            assert_eq!(back.len(), 1);
    951
            assert_eq!(back.capacity(), 2);
    952
    953
            // Cannot use aligned allocation
    954
            let b = Buffer::from(MutableBuffer::new(10));
    955
            let b = b.into_vec::<u8>().unwrap_err();
    956
            b.into_vec::<u64>().unwrap_err();
    957
    958
            // Test slicing
    959
            let mut a: Vec<i128> = Vec::with_capacity(20);
    960
            a.extend_from_slice(&[1, 4, 7, 8, 9, 3, 6]);
    961
            let b = Buffer::from_vec(a);
    962
            let slice = b.slice_with_length(0, 64);
    963
    964
            // Shared reference fails
    965
            let slice = slice.into_vec::<i128>().unwrap_err();
    966
            drop(b);
    967
    968
            // Succeeds as no outstanding shared reference
    969
            let back = slice.into_vec::<i128>().unwrap();
    970
            assert_eq!(&back, &[1, 4, 7, 8]);
    971
            assert_eq!(back.capacity(), 20);
    972
    973
            // Slicing by non-multiple length truncates
    974
            let mut a: Vec<i128> = Vec::with_capacity(8);
    975
            a.extend_from_slice(&[1, 4, 7, 3]);
    976
    977
            let b = Buffer::from_vec(a);
    978
            let slice = b.slice_with_length(0, 34);
    979
            drop(b);
    980
    981
            let back = slice.into_vec::<i128>().unwrap();
    982
            assert_eq!(&back, &[1, 4]);
    983
            assert_eq!(back.capacity(), 8);
    984
    985
            // Offset prevents conversion
    986
            let a: Vec<u32> = vec![1, 3, 4, 6];
    987
            let b = Buffer::from_vec(a).slice(2);
    988
            b.into_vec::<u32>().unwrap_err();
    989
    990
            let b = MutableBuffer::new(16).into_buffer();
    991
            let b = b.into_vec::<u8>().unwrap_err(); // Invalid layout
    992
            let b = b.into_vec::<u32>().unwrap_err(); // Invalid layout
    993
            b.into_mutable().unwrap();
    994
    995
            let b = Buffer::from_vec(vec![1_u32, 3, 5]);
    996
            let b = b.into_mutable().unwrap();
    997
            let b = Buffer::from(b);
    998
            let b = b.into_vec::<u32>().unwrap();
    999
            assert_eq!(b, &[1, 3, 5]);
    1000
        }
    1001
    1002
        #[test]
    1003
        #[should_panic(expected = "capacity overflow")]
    1004
        fn test_from_iter_overflow() {
    1005
            let iter_len = usize::MAX / std::mem::size_of::<u64>() + 1;
    1006
            let _ = Buffer::from_iter(std::iter::repeat_n(0_u64, iter_len));
    1007
        }
    1008
    1009
        #[test]
    1010
        fn bit_slice_length_preserved() {
    1011
            // Create a boring buffer
    1012
            let buf = Buffer::from_iter(std::iter::repeat_n(true, 64));
    1013
    1014
            let assert_preserved = |offset: usize, len: usize| {
    1015
                let new_buf = buf.bit_slice(offset, len);
    1016
                assert_eq!(new_buf.len(), bit_util::ceil(len, 8));
    1017
    1018
                // if the offset is not byte-aligned, we have to create a deep copy to a new buffer
    1019
                // (since the `offset` value inside a Buffer is byte-granular, not bit-granular), so
    1020
                // checking the offset should always return 0 if so. If the offset IS byte-aligned, we
    1021
                // want to make sure it doesn't unnecessarily create a deep copy.
    1022
                if offset % 8 == 0 {
    1023
                    assert_eq!(new_buf.ptr_offset(), offset / 8);
    1024
                } else {
    1025
                    assert_eq!(new_buf.ptr_offset(), 0);
    1026
                }
    1027
            };
    1028
    1029
            // go through every available value for offset
    1030
            for o in 0..=64 {
    1031
                // and go through every length that could accompany that offset - we can't have a
    1032
                // situation where offset + len > 64, because that would go past the end of the buffer,
    1033
                // so we use the map to ensure it's in range.
    1034
                for l in (o..=64).map(|l| l - o) {
    1035
                    // and we just want to make sure every one of these keeps its offset and length
    1036
                    // when neeeded
    1037
                    assert_preserved(o, l);
    1038
                }
    1039
            }
    1040
        }
    1041
    1042
        #[test]
    1043
        fn test_strong_count() {
    1044
            let buffer = Buffer::from_iter(std::iter::repeat_n(0_u8, 100));
    1045
            assert_eq!(buffer.strong_count(), 1);
    1046
    1047
            let buffer2 = buffer.clone();
    1048
            assert_eq!(buffer.strong_count(), 2);
    1049
    1050
            let buffer3 = buffer2.clone();
    1051
            assert_eq!(buffer.strong_count(), 3);
    1052
    1053
            drop(buffer);
    1054
            assert_eq!(buffer2.strong_count(), 2);
    1055
            assert_eq!(buffer3.strong_count(), 2);
    1056
    1057
            // Strong count does not increase on move
    1058
            let capture = move || {
    1059
                assert_eq!(buffer3.strong_count(), 2);
    1060
            };
    1061
    1062
            capture();
    1063
            assert_eq!(buffer2.strong_count(), 2);
    1064
    1065
            drop(capture);
    1066
            assert_eq!(buffer2.strong_count(), 1);
    1067
        }
    1068
    }