vortex_buffer/
arrow.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use arrow_buffer::{ArrowNativeType, OffsetBuffer};
5use bytes::Bytes;
6use vortex_error::vortex_panic;
7
8use crate::{Alignment, Buffer, ByteBuffer};
9
10impl<T: ArrowNativeType> Buffer<T> {
11    /// Converts the buffer zero-copy into a `arrow_buffer::Buffer`.
12    pub fn into_arrow_scalar_buffer(self) -> arrow_buffer::ScalarBuffer<T> {
13        let buffer = arrow_buffer::Buffer::from(self.into_inner());
14        arrow_buffer::ScalarBuffer::from(buffer)
15    }
16
17    /// Convert an Arrow scalar buffer into a Vortex scalar buffer.
18    ///
19    /// ## Panics
20    ///
21    /// Panics if the Arrow buffer is not aligned to the requested alignment, or if the requested
22    /// alignment is not sufficient for type T.
23    pub fn from_arrow_scalar_buffer(arrow: arrow_buffer::ScalarBuffer<T>) -> Self {
24        let length = arrow.len();
25        let bytes = Bytes::from_owner(ArrowWrapper(arrow.into_inner()));
26
27        let alignment = Alignment::of::<T>();
28        if bytes.as_ptr().align_offset(*alignment) != 0 {
29            vortex_panic!(
30                "Arrow buffer is not aligned to the requested alignment: {}",
31                alignment
32            );
33        }
34
35        Self {
36            bytes,
37            length,
38            alignment,
39            _marker: Default::default(),
40        }
41    }
42
43    /// Converts the buffer zero-copy into a `arrow_buffer::OffsetBuffer`.
44    ///
45    /// SAFETY: The caller should ensure that the buffer contains monotonically increasing values
46    /// greater than or equal to zero.
47    pub fn into_arrow_offset_buffer(self) -> OffsetBuffer<T> {
48        unsafe { OffsetBuffer::new_unchecked(self.into_arrow_scalar_buffer()) }
49    }
50}
51
52impl ByteBuffer {
53    /// Converts the buffer zero-copy into a `arrow_buffer::Buffer`.
54    pub fn into_arrow_buffer(self) -> arrow_buffer::Buffer {
55        arrow_buffer::Buffer::from(self.into_inner())
56    }
57
58    /// Convert an Arrow scalar buffer into a Vortex scalar buffer.
59    ///
60    /// ## Panics
61    ///
62    /// Panics if the Arrow buffer is not sufficiently aligned.
63    pub fn from_arrow_buffer(arrow: arrow_buffer::Buffer, alignment: Alignment) -> Self {
64        let length = arrow.len();
65
66        let bytes = Bytes::from_owner(ArrowWrapper(arrow));
67        if bytes.as_ptr().align_offset(*alignment) != 0 {
68            vortex_panic!(
69                "Arrow buffer is not aligned to the requested alignment: {}",
70                alignment
71            );
72        }
73
74        Self {
75            bytes,
76            length,
77            alignment,
78            _marker: Default::default(),
79        }
80    }
81}
82
83/// A wrapper struct to allow `arrow_buffer::Buffer` to implement `AsRef<[u8]>` for
84/// `Bytes::from_owner`.
85struct ArrowWrapper(arrow_buffer::Buffer);
86
87impl AsRef<[u8]> for ArrowWrapper {
88    fn as_ref(&self) -> &[u8] {
89        self.0.as_slice()
90    }
91}
92
93#[cfg(test)]
94mod test {
95    use arrow_buffer::{Buffer as ArrowBuffer, ScalarBuffer};
96
97    use crate::{Alignment, Buffer, buffer};
98
99    #[test]
100    fn into_arrow_buffer() {
101        let buf = buffer![0u8, 1, 2];
102        let arrow: ArrowBuffer = buf.clone().into_arrow_buffer();
103        assert_eq!(arrow.as_ref(), buf.as_slice(), "Buffer values differ");
104        assert_eq!(arrow.as_ptr(), buf.as_ptr(), "Conversion not zero-copy")
105    }
106
107    #[test]
108    fn into_arrow_scalar_buffer() {
109        let buf = buffer![0i32, 1, 2];
110        let scalar: ScalarBuffer<i32> = buf.clone().into_arrow_scalar_buffer();
111        assert_eq!(scalar.as_ref(), buf.as_slice(), "Buffer values differ");
112        assert_eq!(scalar.as_ptr(), buf.as_ptr(), "Conversion not zero-copy")
113    }
114
115    #[test]
116    fn from_arrow_buffer() {
117        let arrow = ArrowBuffer::from_vec(vec![0i32, 1, 2]);
118        let buf = Buffer::from_arrow_buffer(arrow.clone(), Alignment::of::<i32>());
119        assert_eq!(arrow.as_ref(), buf.as_slice(), "Buffer values differ");
120        assert_eq!(arrow.as_ptr(), buf.as_ptr(), "Conversion not zero-copy");
121    }
122}
vortex_buffer/arrow.rs

vortex_buffer/
arrow.rs
