1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

//! Contains [`Buffer`], an immutable container for all Arrow physical types (e.g. i32, f64).

mod immutable;
mod iterator;

use crate::ffi::InternalArrowArray;
use std::ops::Deref;

pub(crate) enum BytesAllocator {
    #[allow(dead_code)]
    InternalArrowArray(InternalArrowArray),

    #[cfg(feature = "arrow")]
    Arrow(arrow_buffer::Buffer),
}
pub(crate) type BytesInner<T> = foreign_vec::ForeignVec<BytesAllocator, T>;

/// Bytes representation.
#[repr(transparent)]
pub struct Bytes<T>(BytesInner<T>);

impl<T> Bytes<T> {
    /// Takes ownership of an allocated memory region.
    /// # Panics
    /// This function panics if and only if pointer is not null
    /// # Safety
    /// This function is safe if and only if `ptr` is valid for `length`
    /// # Implementation
    /// This function leaks if and only if `owner` does not deallocate
    /// the region `[ptr, ptr+length[` when dropped.
    #[inline]
    pub(crate) unsafe fn from_foreign(ptr: *const T, length: usize, owner: BytesAllocator) -> Self {
        Self(BytesInner::from_foreign(ptr, length, owner))
    }

    /// Returns a `Some` mutable reference of [`Vec<T>`] iff this was initialized
    /// from a [`Vec<T>`] and `None` otherwise.
    #[inline]
    pub(crate) fn get_vec(&mut self) -> Option<&mut Vec<T>> {
        self.0.get_vec()
    }
}

impl<T> Deref for Bytes<T> {
    type Target = [T];

    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T> From<Vec<T>> for Bytes<T> {
    #[inline]
    fn from(data: Vec<T>) -> Self {
        let inner: BytesInner<T> = data.into();
        Bytes(inner)
    }
}

impl<T> From<BytesInner<T>> for Bytes<T> {
    #[inline]
    fn from(value: BytesInner<T>) -> Self {
        Self(value)
    }
}

#[cfg(feature = "arrow")]
pub(crate) fn to_buffer<T: crate::types::NativeType>(
    value: std::sync::Arc<Bytes<T>>,
) -> arrow_buffer::Buffer {
    // This should never panic as ForeignVec pointer must be non-null
    let ptr = std::ptr::NonNull::new(value.as_ptr() as _).unwrap();
    let len = value.len() * std::mem::size_of::<T>();
    // Safety: allocation is guaranteed to be valid for `len` bytes
    unsafe { arrow_buffer::Buffer::from_custom_allocation(ptr, len, value) }
}

#[cfg(feature = "arrow")]
pub(crate) fn to_bytes<T: crate::types::NativeType>(value: arrow_buffer::Buffer) -> Bytes<T> {
    let ptr = value.as_ptr();
    let align = ptr.align_offset(std::mem::align_of::<T>());
    assert_eq!(align, 0, "not aligned");
    let len = value.len() / std::mem::size_of::<T>();

    // Valid as `NativeType: Pod` and checked alignment above
    let ptr = value.as_ptr() as *const T;

    let owner = crate::buffer::BytesAllocator::Arrow(value);

    // Safety: slice is valid for len elements of T
    unsafe { Bytes::from_foreign(ptr, len, owner) }
}

pub(super) use iterator::IntoIter;

pub use immutable::Buffer;