1use crate::alloc::Deallocation;
19use crate::buffer::Buffer;
20use crate::native::ArrowNativeType;
21use crate::{BufferBuilder, MutableBuffer, OffsetBuffer};
22use std::fmt::Formatter;
23use std::marker::PhantomData;
24use std::ops::Deref;
25
26#[derive(Clone, Default)]
45pub struct ScalarBuffer<T: ArrowNativeType> {
46 buffer: Buffer,
48 phantom: PhantomData<T>,
49}
50
51impl<T: ArrowNativeType> std::fmt::Debug for ScalarBuffer<T> {
52 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
53 f.debug_tuple("ScalarBuffer").field(&self.as_ref()).finish()
54 }
55}
56
57impl<T: ArrowNativeType> ScalarBuffer<T> {
58 pub fn new(buffer: Buffer, offset: usize, len: usize) -> Self {
69 let size = std::mem::size_of::<T>();
70 let byte_offset = offset.checked_mul(size).expect("offset overflow");
71 let byte_len = len.checked_mul(size).expect("length overflow");
72 buffer.slice_with_length(byte_offset, byte_len).into()
73 }
74
75 pub fn shrink_to_fit(&mut self) {
77 self.buffer.shrink_to_fit();
78 }
79
80 pub fn slice(&self, offset: usize, len: usize) -> Self {
82 Self::new(self.buffer.clone(), offset, len)
83 }
84
85 pub fn inner(&self) -> &Buffer {
87 &self.buffer
88 }
89
90 pub fn into_inner(self) -> Buffer {
92 self.buffer
93 }
94
95 #[inline]
99 pub fn ptr_eq(&self, other: &Self) -> bool {
100 self.buffer.ptr_eq(&other.buffer)
101 }
102}
103
104impl<T: ArrowNativeType> Deref for ScalarBuffer<T> {
105 type Target = [T];
106
107 #[inline]
108 fn deref(&self) -> &Self::Target {
109 unsafe {
111 std::slice::from_raw_parts(
112 self.buffer.as_ptr() as *const T,
113 self.buffer.len() / std::mem::size_of::<T>(),
114 )
115 }
116 }
117}
118
119impl<T: ArrowNativeType> AsRef<[T]> for ScalarBuffer<T> {
120 #[inline]
121 fn as_ref(&self) -> &[T] {
122 self
123 }
124}
125
126impl<T: ArrowNativeType> From<MutableBuffer> for ScalarBuffer<T> {
127 fn from(value: MutableBuffer) -> Self {
128 Buffer::from(value).into()
129 }
130}
131
132impl<T: ArrowNativeType> From<Buffer> for ScalarBuffer<T> {
133 fn from(buffer: Buffer) -> Self {
134 let align = std::mem::align_of::<T>();
135 let is_aligned = buffer.as_ptr().align_offset(align) == 0;
136
137 match buffer.deallocation() {
138 Deallocation::Standard(_) => assert!(
139 is_aligned,
140 "Memory pointer is not aligned with the specified scalar type"
141 ),
142 Deallocation::Custom(_, _) =>
143 assert!(is_aligned, "Memory pointer from external source (e.g, FFI) is not aligned with the specified scalar type. Before importing buffer through FFI, please make sure the allocation is aligned."),
144 }
145
146 Self {
147 buffer,
148 phantom: Default::default(),
149 }
150 }
151}
152
153impl<T: ArrowNativeType> From<OffsetBuffer<T>> for ScalarBuffer<T> {
154 fn from(value: OffsetBuffer<T>) -> Self {
155 value.into_inner()
156 }
157}
158
159impl<T: ArrowNativeType> From<Vec<T>> for ScalarBuffer<T> {
160 fn from(value: Vec<T>) -> Self {
161 Self {
162 buffer: Buffer::from_vec(value),
163 phantom: Default::default(),
164 }
165 }
166}
167
168impl<T: ArrowNativeType> From<ScalarBuffer<T>> for Vec<T> {
169 fn from(value: ScalarBuffer<T>) -> Self {
170 value
171 .buffer
172 .into_vec()
173 .unwrap_or_else(|buffer| buffer.typed_data::<T>().into())
174 }
175}
176
177impl<T: ArrowNativeType> From<BufferBuilder<T>> for ScalarBuffer<T> {
178 fn from(mut value: BufferBuilder<T>) -> Self {
179 let len = value.len();
180 Self::new(value.finish(), 0, len)
181 }
182}
183
184impl<T: ArrowNativeType> FromIterator<T> for ScalarBuffer<T> {
185 #[inline]
186 fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
187 iter.into_iter().collect::<Vec<_>>().into()
188 }
189}
190
191impl<'a, T: ArrowNativeType> IntoIterator for &'a ScalarBuffer<T> {
192 type Item = &'a T;
193 type IntoIter = std::slice::Iter<'a, T>;
194
195 fn into_iter(self) -> Self::IntoIter {
196 self.as_ref().iter()
197 }
198}
199
200impl<T: ArrowNativeType, S: AsRef<[T]> + ?Sized> PartialEq<S> for ScalarBuffer<T> {
201 fn eq(&self, other: &S) -> bool {
202 self.as_ref().eq(other.as_ref())
203 }
204}
205
206impl<T: ArrowNativeType, const N: usize> PartialEq<ScalarBuffer<T>> for [T; N] {
207 fn eq(&self, other: &ScalarBuffer<T>) -> bool {
208 self.as_ref().eq(other.as_ref())
209 }
210}
211
212impl<T: ArrowNativeType> PartialEq<ScalarBuffer<T>> for [T] {
213 fn eq(&self, other: &ScalarBuffer<T>) -> bool {
214 self.as_ref().eq(other.as_ref())
215 }
216}
217
218impl<T: ArrowNativeType> PartialEq<ScalarBuffer<T>> for Vec<T> {
219 fn eq(&self, other: &ScalarBuffer<T>) -> bool {
220 self.as_slice().eq(other.as_ref())
221 }
222}
223
224impl<T: ArrowNativeType + Eq> Eq for ScalarBuffer<T> {}
226
227#[cfg(test)]
228mod tests {
229 use std::{ptr::NonNull, sync::Arc};
230
231 use super::*;
232
233 #[test]
234 fn test_basic() {
235 let expected = [0_i32, 1, 2];
236 let buffer = Buffer::from_iter(expected.iter().cloned());
237 let typed = ScalarBuffer::<i32>::new(buffer.clone(), 0, 3);
238 assert_eq!(*typed, expected);
239
240 let typed = ScalarBuffer::<i32>::new(buffer.clone(), 1, 2);
241 assert_eq!(*typed, expected[1..]);
242
243 let typed = ScalarBuffer::<i32>::new(buffer.clone(), 1, 0);
244 assert!(typed.is_empty());
245
246 let typed = ScalarBuffer::<i32>::new(buffer, 3, 0);
247 assert!(typed.is_empty());
248 }
249
250 #[test]
251 fn test_debug() {
252 let buffer = ScalarBuffer::from(vec![1, 2, 3]);
253 assert_eq!(format!("{buffer:?}"), "ScalarBuffer([1, 2, 3])");
254 }
255
256 #[test]
257 #[should_panic(expected = "Memory pointer is not aligned with the specified scalar type")]
258 fn test_unaligned() {
259 let expected = [0_i32, 1, 2];
260 let buffer = Buffer::from_iter(expected.iter().cloned());
261 let buffer = buffer.slice(1);
262 ScalarBuffer::<i32>::new(buffer, 0, 2);
263 }
264
265 #[test]
266 #[should_panic(expected = "the offset of the new Buffer cannot exceed the existing length")]
267 fn test_length_out_of_bounds() {
268 let buffer = Buffer::from_iter([0_i32, 1, 2]);
269 ScalarBuffer::<i32>::new(buffer, 1, 3);
270 }
271
272 #[test]
273 #[should_panic(expected = "the offset of the new Buffer cannot exceed the existing length")]
274 fn test_offset_out_of_bounds() {
275 let buffer = Buffer::from_iter([0_i32, 1, 2]);
276 ScalarBuffer::<i32>::new(buffer, 4, 0);
277 }
278
279 #[test]
280 #[should_panic(expected = "offset overflow")]
281 fn test_length_overflow() {
282 let buffer = Buffer::from_iter([0_i32, 1, 2]);
283 ScalarBuffer::<i32>::new(buffer, usize::MAX, 1);
284 }
285
286 #[test]
287 #[should_panic(expected = "offset overflow")]
288 fn test_start_overflow() {
289 let buffer = Buffer::from_iter([0_i32, 1, 2]);
290 ScalarBuffer::<i32>::new(buffer, usize::MAX / 4 + 1, 0);
291 }
292
293 #[test]
294 #[should_panic(expected = "length overflow")]
295 fn test_end_overflow() {
296 let buffer = Buffer::from_iter([0_i32, 1, 2]);
297 ScalarBuffer::<i32>::new(buffer, 0, usize::MAX / 4 + 1);
298 }
299
300 #[test]
301 fn convert_from_buffer_builder() {
302 let input = vec![1, 2, 3, 4];
303 let buffer_builder = BufferBuilder::from(input.clone());
304 let scalar_buffer = ScalarBuffer::from(buffer_builder);
305 assert_eq!(scalar_buffer.as_ref(), input);
306 }
307
308 #[test]
309 fn into_vec() {
310 let input = vec![1u8, 2, 3, 4];
311
312 let input_buffer = Buffer::from_vec(input.clone());
314 let input_ptr = input_buffer.as_ptr();
315 let input_len = input_buffer.len();
316 let scalar_buffer = ScalarBuffer::<u8>::new(input_buffer, 0, input_len);
317 let vec = Vec::from(scalar_buffer);
318 assert_eq!(vec.as_slice(), input.as_slice());
319 assert_eq!(vec.as_ptr(), input_ptr);
320
321 let mut input_clone = input.clone();
323 let input_ptr = NonNull::new(input_clone.as_mut_ptr()).unwrap();
324 let dealloc = Arc::new(());
325 let buffer =
326 unsafe { Buffer::from_custom_allocation(input_ptr, input_clone.len(), dealloc as _) };
327 let scalar_buffer = ScalarBuffer::<u8>::new(buffer, 0, input.len());
328 let vec = Vec::from(scalar_buffer);
329 assert_eq!(vec, input.as_slice());
330 assert_ne!(vec.as_ptr(), input_ptr.as_ptr());
331
332 let input_buffer = Buffer::from_vec(input.clone());
334 let input_ptr = input_buffer.as_ptr();
335 let input_len = input_buffer.len();
336 let scalar_buffer = ScalarBuffer::<u8>::new(input_buffer, 1, input_len - 1);
337 let vec = Vec::from(scalar_buffer);
338 assert_eq!(vec.as_slice(), &input[1..]);
339 assert_ne!(vec.as_ptr(), input_ptr);
340
341 let buffer = Buffer::from_slice_ref(input.as_slice());
343 let scalar_buffer = ScalarBuffer::<u8>::new(buffer, 0, input.len());
344 let vec = Vec::from(scalar_buffer);
345 assert_eq!(vec, input.as_slice());
346 assert_ne!(vec.as_ptr(), input.as_ptr());
347 }
348
349 #[test]
350 fn scalar_buffer_impl_eq() {
351 fn are_equal<T: Eq>(a: &T, b: &T) -> bool {
352 a.eq(b)
353 }
354
355 assert!(
356 are_equal(
357 &ScalarBuffer::<i16>::from(vec![23]),
358 &ScalarBuffer::<i16>::from(vec![23])
359 ),
360 "ScalarBuffer should implement Eq if the inner type does"
361 );
362 }
363}