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 vortex_error::vortex_panic;
7
8use crate::Alignment;
9use crate::Buffer;
10use crate::ByteBuffer;
11
12impl<T: ArrowNativeType> Buffer<T> {
13    /// Converts the buffer zero-copy into a `arrow_buffer::Buffer`.
14    pub fn into_arrow_scalar_buffer(self) -> arrow_buffer::ScalarBuffer<T> {
15        if self.is_empty() {
16            return Vec::new().into();
17        }
18        let buffer = self.into_byte_buffer().into_arrow_buffer();
19        arrow_buffer::ScalarBuffer::from(buffer)
20    }
21
22    /// Convert an Arrow scalar buffer into a Vortex scalar buffer.
23    ///
24    /// ## Panics
25    ///
26    /// Panics if the Arrow buffer is not aligned to the requested alignment, or if the requested
27    /// alignment is not sufficient for type T.
28    pub fn from_arrow_scalar_buffer(arrow: arrow_buffer::ScalarBuffer<T>) -> Self {
29        let length = arrow.len();
30        let arrow = arrow.into_inner();
31
32        let alignment = Alignment::of::<T>();
33        if arrow.as_ptr().align_offset(alignment.as_usize()) != 0 {
34            vortex_panic!(
35                "Arrow buffer is not aligned to the requested alignment: {}",
36                alignment
37            );
38        }
39
40        debug_assert_eq!(length, arrow.len() / size_of::<T>());
41        Self::from_arrow_owner(arrow, length, alignment)
42    }
43
44    /// Converts the buffer zero-copy into a `arrow_buffer::OffsetBuffer`.
45    ///
46    /// SAFETY: The caller should ensure that the buffer contains monotonically increasing values
47    /// greater than or equal to zero.
48    pub fn into_arrow_offset_buffer(self) -> OffsetBuffer<T> {
49        unsafe { OffsetBuffer::new_unchecked(self.into_arrow_scalar_buffer()) }
50    }
51}
52
53impl ByteBuffer {
54    /// Converts the buffer zero-copy into a `arrow_buffer::Buffer`.
55    pub fn into_arrow_buffer(self) -> arrow_buffer::Buffer {
56        if let Some(crate::BufferBacking::Arrow(arrow)) = self.backing.as_deref() {
57            let offset = self.ptr.addr().get() - arrow.as_ptr().addr();
58            return arrow.slice_with_length(offset, self.length);
59        }
60        arrow_buffer::Buffer::from(self.into_bytes())
61    }
62
63    /// Convert an Arrow scalar buffer into a Vortex scalar buffer.
64    ///
65    /// ## Panics
66    ///
67    /// Panics if the Arrow buffer is not sufficiently aligned.
68    pub fn from_arrow_buffer(arrow: arrow_buffer::Buffer, alignment: Alignment) -> Self {
69        let length = arrow.len();
70
71        if arrow.as_ptr().align_offset(alignment.as_usize()) != 0 {
72            vortex_panic!(
73                "Arrow buffer is not aligned to the requested alignment: {}",
74                alignment
75            );
76        }
77
78        Self::from_arrow_owner(arrow, length, alignment)
79    }
80}
81
82#[cfg(test)]
83mod test {
84    use arrow_buffer::Buffer as ArrowBuffer;
85    use arrow_buffer::ScalarBuffer;
86
87    use crate::Alignment;
88    use crate::Buffer;
89    use crate::buffer;
90
91    #[test]
92    fn into_arrow_buffer() {
93        let buf = buffer![0u8, 1, 2];
94        let arrow: ArrowBuffer = buf.clone().into_arrow_buffer();
95        assert_eq!(arrow.as_ref(), buf.as_slice(), "Buffer values differ");
96        assert_eq!(arrow.as_ptr(), buf.as_ptr(), "Conversion not zero-copy")
97    }
98
99    #[test]
100    fn into_arrow_scalar_buffer() {
101        let buf = buffer![0i32, 1, 2];
102        let scalar: ScalarBuffer<i32> = buf.clone().into_arrow_scalar_buffer();
103        assert_eq!(scalar.as_ref(), buf.as_slice(), "Buffer values differ");
104        assert_eq!(scalar.as_ptr(), buf.as_ptr(), "Conversion not zero-copy")
105    }
106
107    #[test]
108    fn empty_into_arrow_scalar_buffer() {
109        let scalar = Buffer::<i64>::empty().into_arrow_scalar_buffer();
110
111        assert!(scalar.is_empty());
112        assert_eq!(scalar.as_ptr().align_offset(align_of::<i64>()), 0);
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        let round_trip = buf.into_arrow_buffer();
123        assert_eq!(round_trip.as_ptr(), arrow.as_ptr());
124    }
125}
vortex_buffer/arrow.rs

vortex_buffer/
arrow.rs
