Fixed error in passing sliced arrays via FFI

arrow-pyarrow-integration-testing/tests/test_sql.py

@@ -39,15 +39,52 @@ def tearDown(self):
         # No leak of C++ memory
         self.assertEqual(self.old_allocated_cpp, pyarrow.total_allocated_bytes())
 
-    def test_string_roundtrip(self):
-        """
-        Python -> Rust -> Python
-        """
+    def test_primitive(self):
+        a = pyarrow.array([0, None, 2, 3, 4])
+        b = arrow_pyarrow_integration_testing.round_trip_array(a)
+
+        b.validate(full=True)
+        assert a.to_pylist() == b.to_pylist()
+        assert a.type == b.type
+
+    def test_primitive_sliced(self):
+        a = pyarrow.array([0, None, 2, 3, 4]).slice(1, 2)
+        b = arrow_pyarrow_integration_testing.round_trip_array(a)
+
+        b.validate(full=True)
+        assert a.to_pylist() == b.to_pylist()
+        assert a.type == b.type
+
+    def test_boolean(self):
+        a = pyarrow.array([True, None, False, True, False])
+        b = arrow_pyarrow_integration_testing.round_trip_array(a)
+
+        b.validate(full=True)
+        assert a.to_pylist() == b.to_pylist()
+        assert a.type == b.type
+
+    def test_boolean_sliced(self):
+        a = pyarrow.array([True, None, False, True, False]).slice(1, 2)
+        b = arrow_pyarrow_integration_testing.round_trip_array(a)
+
+        b.validate(full=True)
+        assert a.to_pylist() == b.to_pylist()
+        assert a.type == b.type
+
+    def test_string(self):
         a = pyarrow.array(["a", None, "ccc"])
         b = arrow_pyarrow_integration_testing.round_trip_array(a)
         c = pyarrow.array(["a", None, "ccc"])
         self.assertEqual(b, c)
 
+    def test_string_sliced(self):
+        a = pyarrow.array(["a", None, "ccc"]).slice(1, 2)
+        b = arrow_pyarrow_integration_testing.round_trip_array(a)
+
+        b.validate(full=True)
+        assert a.to_pylist() == b.to_pylist()
+        assert a.type == b.type
+
     def test_decimal_roundtrip(self):
         """
         Python -> Rust -> Python
@@ -75,10 +112,17 @@ def test_list_array(self):
             assert a.to_pylist() == b.to_pylist()
             assert a.type == b.type
 
-    def test_struct_array(self):
-        """
-        Python -> Rust -> Python
-        """
+    def test_list_sliced(self):
+        a = pyarrow.array(
+            [[], None, [1, 2], [4, 5, 6]], pyarrow.list_(pyarrow.int64())
+        ).slice(1, 2)
+        b = arrow_pyarrow_integration_testing.round_trip_array(a)
+
+        b.validate(full=True)
+        assert a.to_pylist() == b.to_pylist()
+        assert a.type == b.type
+
+    def test_struct(self):
         fields = [
             ("f1", pyarrow.int32()),
             ("f2", pyarrow.string()),
@@ -99,6 +143,27 @@ def test_struct_array(self):
         assert a.to_pylist() == b.to_pylist()
         assert a.type == b.type
 
+    # see https://issues.apache.org/jira/browse/ARROW-14383
+    def _test_struct_sliced(self):
+        fields = [
+            ("f1", pyarrow.int32()),
+            ("f2", pyarrow.string()),
+        ]
+        a = pyarrow.array(
+            [
+                {"f1": 1, "f2": "a"},
+                None,
+                {"f1": 3, "f2": None},
+                {"f1": None, "f2": "d"},
+                {"f1": None, "f2": None},
+            ],
+            pyarrow.struct(fields),
+        ).slice(1, 2)
+        b = arrow_pyarrow_integration_testing.round_trip_array(a)
+        b.validate(full=True)
+        assert a.to_pylist() == b.to_pylist()
+        assert a.type == b.type
+
     def test_list_list_array(self):
         """
         Python -> Rust -> Python

src/array/binary/ffi.rs

@@ -1,6 +1,6 @@
 use crate::{
     array::{FromFfi, Offset, ToFfi},
-    datatypes::DataType,
+    bitmap::align,
     ffi,
 };
 
@@ -12,43 +12,54 @@ unsafe impl<O: Offset> ToFfi for BinaryArray<O> {
     fn buffers(&self) -> Vec<Option<std::ptr::NonNull<u8>>> {
         vec![
             self.validity.as_ref().map(|x| x.as_ptr()),
-            std::ptr::NonNull::new(self.offsets.as_ptr() as *mut u8),
-            std::ptr::NonNull::new(self.values.as_ptr() as *mut u8),
+            Some(self.offsets.as_ptr().cast::<u8>()),
+            Some(self.values.as_ptr().cast::<u8>()),
         ]
     }
 
-    #[inline]
-    fn offset(&self) -> usize {
-        self.offset
+    fn offset(&self) -> Option<usize> {
+        let offset = self.offsets.offset();
+        if let Some(bitmap) = self.validity.as_ref() {
+            if bitmap.offset() == offset {
+                Some(offset)
+            } else {
+                None
+            }
+        } else {
+            Some(offset)
+        }
+    }
+
+    fn to_ffi_aligned(&self) -> Self {
+        let offset = self.offsets.offset();
+
+        let validity = self.validity.as_ref().map(|bitmap| {
+            if bitmap.offset() == offset {
+                bitmap.clone()
+            } else {
+                align(bitmap, offset)
+            }
+        });
+
+        Self {
+            data_type: self.data_type.clone(),
+            validity,
+            offsets: self.offsets.clone(),
+            values: self.values.clone(),
+        }
     }
 }
 
 impl<O: Offset, A: ffi::ArrowArrayRef> FromFfi<A> for BinaryArray<O> {
     unsafe fn try_from_ffi(array: A) -> Result<Self> {
         let data_type = array.field().data_type().clone();
-        let expected = if O::is_large() {
-            DataType::LargeBinary
-        } else {
-            DataType::Binary
-        };
-        assert_eq!(data_type, expected);
-
-        let length = array.array().len();
-        let offset = array.array().offset();
-        let mut validity = unsafe { array.validity() }?;
-        let mut offsets = unsafe { array.buffer::<O>(0) }?;
-        let values = unsafe { array.buffer::<u8>(1) }?;
 
-        if offset > 0 {
-            offsets = offsets.slice(offset, length);
-            validity = validity.map(|x| x.slice(offset, length))
-        }
+        let validity = unsafe { array.validity() }?;
+        let offsets = unsafe { array.buffer::<O>(0) }?;
+        let values = unsafe { array.buffer::<u8>(1) }?;
 
         Ok(Self::from_data_unchecked(
-            Self::default_data_type(),
-            offsets,
-            values,
-            validity,
+            data_type, offsets, values, validity,
         ))
     }
 }

src/array/binary/mod.rs

@@ -23,7 +23,6 @@ pub struct BinaryArray<O: Offset> {
     offsets: Buffer<O>,
     values: Buffer<u8>,
     validity: Option<Bitmap>,
-    offset: usize,
 }
 
 // constructors
@@ -71,7 +70,6 @@ impl<O: Offset> BinaryArray<O> {
             offsets,
             values,
             validity,
-            offset: 0,
         }
     }
 
@@ -113,7 +111,6 @@ impl<O: Offset> BinaryArray<O> {
             offsets,
             values,
             validity,
-            offset: 0,
         }
     }
 
@@ -146,7 +143,6 @@ impl<O: Offset> BinaryArray<O> {
             offsets,
             values: self.values.clone(),
             validity,
-            offset: self.offset + offset,
         }
     }
 

src/array/boolean/ffi.rs

@@ -1,6 +1,6 @@
 use crate::{
     array::{FromFfi, ToFfi},
-    datatypes::DataType,
+    bitmap::align,
     ffi,
 };
 
@@ -16,24 +16,43 @@ unsafe impl ToFfi for BooleanArray {
         ]
     }
 
-    fn offset(&self) -> usize {
-        self.offset
+    fn offset(&self) -> Option<usize> {
+        let offset = self.values.offset();
+        if let Some(bitmap) = self.validity.as_ref() {
+            if bitmap.offset() == offset {
+                Some(offset)
+            } else {
+                None
+            }
+        } else {
+            Some(offset)
+        }
+    }
+
+    fn to_ffi_aligned(&self) -> Self {
+        let offset = self.values.offset();
+
+        let validity = self.validity.as_ref().map(|bitmap| {
+            if bitmap.offset() == offset {
+                bitmap.clone()
+            } else {
+                align(bitmap, offset)
+            }
+        });
+
+        Self {
+            data_type: self.data_type.clone(),
+            validity,
+            values: self.values.clone(),
+        }
     }
 }
 
 impl<A: ffi::ArrowArrayRef> FromFfi<A> for BooleanArray {
     unsafe fn try_from_ffi(array: A) -> Result<Self> {
         let data_type = array.field().data_type().clone();
-        assert_eq!(data_type, DataType::Boolean);
-        let length = array.array().len();
-        let offset = array.array().offset();
-        let mut validity = unsafe { array.validity() }?;
-        let mut values = unsafe { array.bitmap(0) }?;
-
-        if offset > 0 {
-            values = values.slice(offset, length);
-            validity = validity.map(|x| x.slice(offset, length))
-        }
+        let validity = unsafe { array.validity() }?;
+        let values = unsafe { array.bitmap(0) }?;
         Ok(Self::from_data(data_type, values, validity))
     }
 }

src/array/boolean/mod.rs

@@ -20,7 +20,6 @@ pub struct BooleanArray {
     data_type: DataType,
     values: Bitmap,
     validity: Option<Bitmap>,
-    offset: usize,
 }
 
 impl BooleanArray {
@@ -50,7 +49,6 @@ impl BooleanArray {
             data_type,
             values,
             validity,
-            offset: 0,
         }
     }
 
@@ -83,7 +81,6 @@ impl BooleanArray {
             data_type: self.data_type.clone(),
             values: self.values.clone().slice_unchecked(offset, length),
             validity,
-            offset: self.offset + offset,
         }
     }
 

src/array/dictionary/ffi.rs

@@ -11,26 +11,26 @@ unsafe impl<K: DictionaryKey> ToFfi for DictionaryArray<K> {
         self.keys.buffers()
     }
 
-    #[inline]
-    fn offset(&self) -> usize {
-        self.offset
+    fn offset(&self) -> Option<usize> {
+        self.keys.offset()
+    }
+
+    fn to_ffi_aligned(&self) -> Self {
+        Self {
+            data_type: self.data_type.clone(),
+            keys: self.keys.to_ffi_aligned(),
+            values: self.values.clone(),
+        }
     }
 }
 
 impl<K: DictionaryKey, A: ffi::ArrowArrayRef> FromFfi<A> for DictionaryArray<K> {
     unsafe fn try_from_ffi(array: A) -> Result<Self> {
         // keys: similar to PrimitiveArray, but the datatype is the inner one
-        let length = array.array().len();
-        let offset = array.array().offset();
-        let mut validity = unsafe { array.validity() }?;
-        let mut values = unsafe { array.buffer::<K>(0) }?;
+        let validity = unsafe { array.validity() }?;
+        let values = unsafe { array.buffer::<K>(0) }?;
 
-        if offset > 0 {
-            values = values.slice(offset, length);
-            validity = validity.map(|x| x.slice(offset, length))
-        }
         let keys = PrimitiveArray::<K>::from_data(K::DATA_TYPE, values, validity);
-        // values
         let values = array.dictionary()?.unwrap();
         let values = ffi::try_from(values)?.into();
 

src/array/dictionary/mod.rs

@@ -37,7 +37,6 @@ pub struct DictionaryArray<K: DictionaryKey> {
     data_type: DataType,
     keys: PrimitiveArray<K>,
     values: Arc<dyn Array>,
-    offset: usize,
 }
 
 impl<K: DictionaryKey> DictionaryArray<K> {
@@ -69,7 +68,6 @@ impl<K: DictionaryKey> DictionaryArray<K> {
             data_type,
             keys,
             values,
-            offset: 0,
         }
     }
 
@@ -81,7 +79,6 @@ impl<K: DictionaryKey> DictionaryArray<K> {
             data_type: self.data_type.clone(),
             keys: self.keys.clone().slice(offset, length),
             values: self.values.clone(),
-            offset: self.offset + offset,
         }
     }
 
@@ -93,7 +90,6 @@ impl<K: DictionaryKey> DictionaryArray<K> {
             data_type: self.data_type.clone(),
             keys: self.keys.clone().slice_unchecked(offset, length),
             values: self.values.clone(),
-            offset: self.offset + offset,
         }
     }