vortex_buffer/
arrow.rs

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

vortex_buffer/
arrow.rs
