Added Offsets and OffsetsBuffer

jorgecarleitao · Dec 6, 2022 · 832a02f · 832a02f
1 parent 9ea25f0
commit 832a02f
Show file tree

Hide file tree

Showing 72 changed files with 1,160 additions and 1,301 deletions.
diff --git a/benches/iter_list.rs b/benches/iter_list.rs
@@ -16,8 +16,7 @@ fn add_benchmark(c: &mut Criterion) {
         let values = Buffer::from_iter(0..size as i32);
         let values = PrimitiveArray::<i32>::from_data(DataType::Int32, values, None);
 
-        let mut offsets = (0..size as i32).step_by(2).collect::<Vec<_>>();
-        offsets.push(size as i32);
+        let offsets = (0..=size as i32).step_by(2).collect::<Vec<_>>();
 
         let validity = (0..(offsets.len() - 1))
             .map(|i| i % 4 == 0)
@@ -26,7 +25,7 @@ fn add_benchmark(c: &mut Criterion) {
         let data_type = ListArray::<i32>::default_datatype(DataType::Int32);
         let array = ListArray::<i32>::from_data(
             data_type,
-            offsets.into(),
+            offsets.try_into().unwrap(),
             Box::new(values),
             Some(validity),
         );

diff --git a/src/array/binary/ffi.rs b/src/array/binary/ffi.rs
@@ -2,7 +2,7 @@ use crate::{
     array::{FromFfi, ToFfi},
     bitmap::align,
     ffi,
-    offset::Offset,
+    offset::{Offset, OffsetsBuffer},
 };
 
 use crate::error::Result;
@@ -13,13 +13,13 @@ unsafe impl<O: Offset> ToFfi for BinaryArray<O> {
     fn buffers(&self) -> Vec<Option<*const u8>> {
         vec![
             self.validity.as_ref().map(|x| x.as_ptr()),
-            Some(self.offsets.as_ptr().cast::<u8>()),
+            Some(self.offsets.buffer().as_ptr().cast::<u8>()),
             Some(self.values.as_ptr().cast::<u8>()),
         ]
     }
 
     fn offset(&self) -> Option<usize> {
-        let offset = self.offsets.offset();
+        let offset = self.offsets.buffer().offset();
         if let Some(bitmap) = self.validity.as_ref() {
             if bitmap.offset() == offset {
                 Some(offset)
@@ -32,7 +32,7 @@ unsafe impl<O: Offset> ToFfi for BinaryArray<O> {
     }
 
     fn to_ffi_aligned(&self) -> Self {
-        let offset = self.offsets.offset();
+        let offset = self.offsets.buffer().offset();
 
         let validity = self.validity.as_ref().map(|bitmap| {
             if bitmap.offset() == offset {
@@ -59,8 +59,9 @@ impl<O: Offset, A: ffi::ArrowArrayRef> FromFfi<A> for BinaryArray<O> {
         let offsets = unsafe { array.buffer::<O>(1) }?;
         let values = unsafe { array.buffer::<u8>(2) }?;
 
-        Ok(Self::from_data_unchecked(
-            data_type, offsets, values, validity,
-        ))
+        // assumption that data from FFI is well constructed
+        let offsets = unsafe { OffsetsBuffer::new_unchecked(offsets) };
+
+        Ok(Self::new(data_type, offsets, values, validity))
     }
 }