Added Offsets and OffsetsBuffer (#1316)

jorgecarleitao · Dec 10, 2022 · 4ed9b91 · 4ed9b91
1 parent 1da33ac
commit 4ed9b91
Show file tree

Hide file tree

Showing 124 changed files with 1,449 additions and 1,470 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
@@ -1,7 +1,8 @@
 use crate::{
-    array::{FromFfi, Offset, ToFfi},
+    array::{FromFfi, ToFfi},
     bitmap::align,
     ffi,
+    offset::{Offset, OffsetsBuffer},
 };
 
 use crate::error::Result;
@@ -12,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)
@@ -31,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 {
@@ -58,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))
     }
 }
diff --git a/src/array/binary/fmt.rs b/src/array/binary/fmt.rs
@@ -1,7 +1,8 @@
 use std::fmt::{Debug, Formatter, Result, Write};
 
+use crate::offset::Offset;
+
 use super::super::fmt::write_vec;
-use super::super::Offset;
 use super::BinaryArray;
 
 pub fn write_value<O: Offset, W: Write>(array: &BinaryArray<O>, index: usize, f: &mut W) -> Result {

diff --git a/src/array/binary/from.rs b/src/array/binary/from.rs
@@ -1,6 +1,6 @@
 use std::iter::FromIterator;
 
-use crate::array::Offset;
+use crate::offset::Offset;
 
 use super::{BinaryArray, MutableBinaryArray};
 

diff --git a/src/array/binary/iterator.rs b/src/array/binary/iterator.rs
@@ -1,6 +1,7 @@
 use crate::{
-    array::{ArrayAccessor, ArrayValuesIter, Offset},
+    array::{ArrayAccessor, ArrayValuesIter},
     bitmap::utils::{BitmapIter, ZipValidity},
+    offset::Offset,
 };
 
 use super::{BinaryArray, MutableBinaryValuesArray};