Added support for ffi for FixedSizeList and FixedSizeBinary

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

@@ -85,6 +85,14 @@ def test_string_sliced(self):
         assert a.to_pylist() == b.to_pylist()
         assert a.type == b.type
 
+    def test_fixed_binary(self):
+        a = pyarrow.array([b"aa", None, b"cc"], pyarrow.binary(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
@@ -179,6 +187,35 @@ def test_list_list_array(self):
             assert a.to_pylist() == b.to_pylist()
             assert a.type == b.type
 
+    def test_fixed_list(self):
+        """
+        Python -> Rust -> Python
+        """
+        a = pyarrow.array(
+            [None, [1, 2], [4, 5]],
+            pyarrow.list_(pyarrow.int64(), 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
+
+    # same as https://issues.apache.org/jira/browse/ARROW-14383
+    def _test_fixed_list_sliced(self):
+        """
+        Python -> Rust -> Python
+        """
+        a = pyarrow.array(
+            [None, [1, 2], [4, 5]],
+            pyarrow.list_(pyarrow.int64(), 2),
+        ).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_dict(self):
         """
         Python -> Rust -> Python

src/array/binary/ffi.rs

@@ -55,8 +55,8 @@ impl<O: Offset, A: ffi::ArrowArrayRef> FromFfi<A> for BinaryArray<O> {
         let data_type = array.field().data_type().clone();
 
         let validity = unsafe { array.validity() }?;
-        let offsets = unsafe { array.buffer::<O>(0) }?;
-        let values = unsafe { array.buffer::<u8>(1) }?;
+        let offsets = unsafe { array.buffer::<O>(1) }?;
+        let values = unsafe { array.buffer::<u8>(2) }?;
 
         Ok(Self::from_data_unchecked(
             data_type, offsets, values, validity,

src/array/boolean/ffi.rs

@@ -52,7 +52,7 @@ 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();
         let validity = unsafe { array.validity() }?;
-        let values = unsafe { array.bitmap(0) }?;
+        let values = unsafe { array.bitmap(1) }?;
         Ok(Self::from_data(data_type, values, validity))
     }
 }

src/array/dictionary/ffi.rs

@@ -28,7 +28,7 @@ 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 validity = unsafe { array.validity() }?;
-        let values = unsafe { array.buffer::<K>(0) }?;
+        let values = unsafe { array.buffer::<K>(1) }?;
 
         let keys = PrimitiveArray::<K>::from_data(K::DATA_TYPE, values, validity);
         let values = array.dictionary()?.unwrap();

src/array/ffi.rs

@@ -27,7 +27,7 @@ pub(crate) unsafe trait ToFfi {
 
 /// Trait describing how a struct imports into itself from the
 /// [C data interface](https://arrow.apache.org/docs/format/CDataInterface.html) (FFI).
-pub trait FromFfi<T: ffi::ArrowArrayRef>: Sized {
+pub(crate) trait FromFfi<T: ffi::ArrowArrayRef>: Sized {
     /// Convert itself from FFI.
     /// # Safety
     /// This function is intrinsically `unsafe` as it requires the FFI to be made according

src/array/fixed_size_binary/ffi.rs

@@ -1,4 +1,9 @@
-use crate::{array::ToFfi, bitmap::align};
+use crate::{
+    array::{FromFfi, ToFfi},
+    bitmap::align,
+    error::Result,
+    ffi,
+};
 
 use super::FixedSizeBinaryArray;
 
@@ -42,3 +47,13 @@ unsafe impl ToFfi for FixedSizeBinaryArray {
         }
     }
 }
+
+impl<A: ffi::ArrowArrayRef> FromFfi<A> for FixedSizeBinaryArray {
+    unsafe fn try_from_ffi(array: A) -> Result<Self> {
+        let data_type = array.field().data_type().clone();
+        let validity = unsafe { array.validity() }?;
+        let values = unsafe { array.buffer::<u8>(1) }?;
+
+        Ok(Self::from_data(data_type, values, validity))
+    }
+}

src/array/fixed_size_list/ffi.rs

@@ -1,7 +1,14 @@
 use std::sync::Arc;
 
-use super::super::{ffi::ToFfi, Array};
 use super::FixedSizeListArray;
+use crate::{
+    array::{
+        ffi::{FromFfi, ToFfi},
+        Array,
+    },
+    error::Result,
+    ffi,
+};
 
 unsafe impl ToFfi for FixedSizeListArray {
     fn buffers(&self) -> Vec<Option<std::ptr::NonNull<u8>>> {
@@ -25,3 +32,14 @@ unsafe impl ToFfi for FixedSizeListArray {
         self.clone()
     }
 }
+
+impl<A: ffi::ArrowArrayRef> FromFfi<A> for FixedSizeListArray {
+    unsafe fn try_from_ffi(array: A) -> Result<Self> {
+        let data_type = array.field().data_type().clone();
+        let validity = unsafe { array.validity() }?;
+        let child = unsafe { array.child(0)? };
+        let values = ffi::try_from(child)?.into();
+
+        Ok(Self::from_data(data_type, values, validity))
+    }
+}

src/array/list/ffi.rs

@@ -54,7 +54,7 @@ impl<O: Offset, A: ffi::ArrowArrayRef> FromFfi<A> for ListArray<O> {
     unsafe fn try_from_ffi(array: A) -> Result<Self> {
         let data_type = array.field().data_type().clone();
         let validity = unsafe { array.validity() }?;
-        let offsets = unsafe { array.buffer::<O>(0) }?;
+        let offsets = unsafe { array.buffer::<O>(1) }?;
         let child = unsafe { array.child(0)? };
         let values = ffi::try_from(child)?.into();
 

src/array/list/mutable.rs

@@ -42,15 +42,6 @@ impl<O: Offset, M: MutableArray + Default> MutableListArray<O, M> {
             validity: None,
         }
     }
-
-    /// Shrinks the capacity of the [`MutableListArray`] to fit its current length.
-    pub fn shrink_to_fit(&mut self) {
-        self.values.shrink_to_fit();
-        self.offsets.shrink_to_fit();
-        if let Some(validity) = &mut self.validity {
-            validity.shrink_to_fit()
-        }
-    }
 }
 
 impl<O: Offset, M: MutableArray + Default> Default for MutableListArray<O, M> {
@@ -188,6 +179,15 @@ impl<O: Offset, M: MutableArray> MutableListArray<O, M> {
         let a: ListArray<O> = self.into();
         Arc::new(a)
     }
+
+    /// Shrinks the capacity of the [`MutableListArray`] to fit its current length.
+    pub fn shrink_to_fit(&mut self) {
+        self.values.shrink_to_fit();
+        self.offsets.shrink_to_fit();
+        if let Some(validity) = &mut self.validity {
+            validity.shrink_to_fit()
+        }
+    }
 }
 
 impl<O: Offset, M: MutableArray + Default + 'static> MutableArray for MutableListArray<O, M> {

src/array/map/ffi.rs

@@ -54,7 +54,7 @@ impl<A: ffi::ArrowArrayRef> FromFfi<A> for MapArray {
     unsafe fn try_from_ffi(array: A) -> Result<Self> {
         let data_type = array.field().data_type().clone();
         let validity = unsafe { array.validity() }?;
-        let offsets = unsafe { array.buffer::<i32>(0) }?;
+        let offsets = unsafe { array.buffer::<i32>(1) }?;
         let child = array.child(0)?;
         let values = ffi::try_from(child)?.into();
 

src/array/primitive/ffi.rs

@@ -53,7 +53,7 @@ impl<T: NativeType, A: ffi::ArrowArrayRef> FromFfi<A> for PrimitiveArray<T> {
     unsafe fn try_from_ffi(array: A) -> Result<Self> {
         let data_type = array.field().data_type().clone();
         let validity = unsafe { array.validity() }?;
-        let values = unsafe { array.buffer::<T>(0) }?;
+        let values = unsafe { array.buffer::<T>(1) }?;
 
         Ok(Self::from_data(data_type, values, validity))
     }

src/array/union/ffi.rs

@@ -37,11 +37,11 @@ impl<A: ffi::ArrowArrayRef> FromFfi<A> for UnionArray {
         let data_type = field.data_type().clone();
         let fields = Self::get_fields(field.data_type());
 
-        let mut types = unsafe { array.buffer::<i8>(0) }?;
+        let mut types = unsafe { array.buffer::<i8>(1) }?;
         let offsets = if Self::is_sparse(&data_type) {
             None
         } else {
-            Some(unsafe { array.buffer::<i32>(1) }?)
+            Some(unsafe { array.buffer::<i32>(2) }?)
         };
 
         let length = array.array().len();

src/array/utf8/ffi.rs

@@ -53,8 +53,8 @@ impl<O: Offset, A: ffi::ArrowArrayRef> FromFfi<A> for Utf8Array<O> {
     unsafe fn try_from_ffi(array: A) -> Result<Self> {
         let data_type = array.field().data_type().clone();
         let validity = unsafe { array.validity() }?;
-        let offsets = unsafe { array.buffer::<O>(0) }?;
-        let values = unsafe { array.buffer::<u8>(1)? };
+        let offsets = unsafe { array.buffer::<O>(1) }?;
+        let values = unsafe { array.buffer::<u8>(2)? };
 
         Ok(Self::from_data_unchecked(
             data_type, offsets, values, validity,

src/ffi/array.rs

@@ -21,8 +21,10 @@ pub unsafe fn try_from<A: ArrowArrayRef>(array: A) -> Result<Box<dyn Array>> {
         LargeUtf8 => Box::new(Utf8Array::<i64>::try_from_ffi(array)?),
         Binary => Box::new(BinaryArray::<i32>::try_from_ffi(array)?),
         LargeBinary => Box::new(BinaryArray::<i64>::try_from_ffi(array)?),
+        FixedSizeBinary => Box::new(FixedSizeBinaryArray::try_from_ffi(array)?),
         List => Box::new(ListArray::<i32>::try_from_ffi(array)?),
         LargeList => Box::new(ListArray::<i64>::try_from_ffi(array)?),
+        FixedSizeList => Box::new(FixedSizeListArray::try_from_ffi(array)?),
         Struct => Box::new(StructArray::try_from_ffi(array)?),
         Dictionary(key_type) => {
             with_match_physical_dictionary_key_type!(key_type, |$T| {

src/ffi/ffi.rs

@@ -252,6 +252,20 @@ fn buffer_offset(array: &Ffi_ArrowArray, data_type: &DataType, i: usize) -> usiz
 // to fetch offset buffer's len to build the second buffer.
 fn buffer_len(array: &Ffi_ArrowArray, data_type: &DataType, i: usize) -> Result<usize> {
     Ok(match (data_type.to_physical_type(), i) {
+        (PhysicalType::FixedSizeBinary, 1) => {
+            if let DataType::FixedSizeBinary(size) = data_type.to_logical_type() {
+                *size * (array.offset as usize + array.length as usize)
+            } else {
+                unreachable!()
+            }
+        }
+        (PhysicalType::FixedSizeList, 1) => {
+            if let DataType::FixedSizeList(_, size) = data_type.to_logical_type() {
+                *size * (array.offset as usize + array.length as usize)
+            } else {
+                unreachable!()
+            }
+        }
         (PhysicalType::Utf8, 1)
         | (PhysicalType::LargeUtf8, 1)
         | (PhysicalType::Binary, 1)
@@ -321,7 +335,7 @@ fn create_dictionary(
     }
 }
 
-pub trait ArrowArrayRef {
+pub trait ArrowArrayRef: std::fmt::Debug {
     fn deallocation(&self) -> Deallocation {
         Deallocation::Foreign(self.parent().clone())
     }
@@ -344,12 +358,11 @@ pub trait ArrowArrayRef {
     /// The caller must guarantee that the buffer `index` corresponds to a bitmap.
     /// This function assumes that the bitmap created from FFI is valid; this is impossible to prove.
     unsafe fn buffer<T: NativeType>(&self, index: usize) -> Result<Buffer<T>> {
-        // +1 to ignore null bitmap
         create_buffer::<T>(
             self.array(),
             self.field().data_type(),
             self.deallocation(),
-            index + 1,
+            index,
         )
     }
 
@@ -358,7 +371,7 @@ pub trait ArrowArrayRef {
     /// This function assumes that the bitmap created from FFI is valid; this is impossible to prove.
     unsafe fn bitmap(&self, index: usize) -> Result<Bitmap> {
         // +1 to ignore null bitmap
-        create_bitmap(self.array(), self.deallocation(), index + 1)
+        create_bitmap(self.array(), self.deallocation(), index)
     }
 
     /// # Safety
@@ -371,6 +384,8 @@ pub trait ArrowArrayRef {
         create_dictionary(self.array(), self.field(), self.parent().clone())
     }
 
+    fn n_buffers(&self) -> usize;
+
     fn parent(&self) -> &Arc<ArrowArray>;
     fn array(&self) -> &Ffi_ArrowArray;
     fn field(&self) -> &Field;
@@ -420,6 +435,10 @@ impl ArrowArrayRef for Arc<ArrowArray> {
     fn array(&self) -> &Ffi_ArrowArray {
         self.array.as_ref()
     }
+
+    fn n_buffers(&self) -> usize {
+        self.array.n_buffers as usize
+    }
 }
 
 #[derive(Debug)]
@@ -442,6 +461,10 @@ impl<'a> ArrowArrayRef for ArrowArrayChild<'a> {
     fn array(&self) -> &Ffi_ArrowArray {
         self.array
     }
+
+    fn n_buffers(&self) -> usize {
+        self.array.n_buffers as usize
+    }
 }
 
 impl<'a> ArrowArrayChild<'a> {

src/ffi/schema.rs

@@ -63,6 +63,9 @@ impl Ffi_ArrowSchema {
             DataType::List(field) => {
                 vec![Box::new(Ffi_ArrowSchema::new(field.as_ref()))]
             }
+            DataType::FixedSizeList(field, _) => {
+                vec![Box::new(Ffi_ArrowSchema::new(field.as_ref()))]
+            }
             DataType::LargeList(field) => {
                 vec![Box::new(Ffi_ArrowSchema::new(field.as_ref()))]
             }
@@ -290,6 +293,17 @@ unsafe fn to_data_type(schema: &Ffi_ArrowSchema) -> Result<DataType> {
                 DataType::Timestamp(TimeUnit::Microsecond, Some(parts[1].to_string()))
             } else if parts.len() == 2 && parts[0] == "tsn" {
                 DataType::Timestamp(TimeUnit::Nanosecond, Some(parts[1].to_string()))
+            } else if parts.len() == 2 && parts[0] == "w" {
+                let size = parts[1]
+                    .parse::<usize>()
+                    .map_err(|_| ArrowError::Ffi("size is not a valid integer".to_string()))?;
+                DataType::FixedSizeBinary(size)
+            } else if parts.len() == 2 && parts[0] == "+w" {
+                let size = parts[1]
+                    .parse::<usize>()
+                    .map_err(|_| ArrowError::Ffi("size is not a valid integer".to_string()))?;
+                let child = to_field(schema.child(0))?;
+                DataType::FixedSizeList(Box::new(child), size)
             } else if parts.len() == 2 && parts[0] == "d" {
                 let parts = parts[1].split(',').collect::<Vec<_>>();
                 if parts.len() < 2 || parts.len() > 3 {
@@ -395,7 +409,7 @@ fn to_format(data_type: &DataType) -> String {
         DataType::List(_) => "+l".to_string(),
         DataType::LargeList(_) => "+L".to_string(),
         DataType::Struct(_) => "+s".to_string(),
-        DataType::FixedSizeBinary(size) => format!("w{}", size),
+        DataType::FixedSizeBinary(size) => format!("w:{}", size),
         DataType::FixedSizeList(_, size) => format!("+w:{}", size),
         DataType::Union(f, ids, mode) => {
             let sparsness = if mode.is_sparse() { 's' } else { 'd' };
@@ -419,6 +433,7 @@ fn to_format(data_type: &DataType) -> String {
 pub(super) fn get_field_child(field: &Field, index: usize) -> Result<Field> {
     match (index, field.data_type()) {
         (0, DataType::List(field)) => Ok(field.as_ref().clone()),
+        (0, DataType::FixedSizeList(field, _)) => Ok(field.as_ref().clone()),
         (0, DataType::LargeList(field)) => Ok(field.as_ref().clone()),
         (0, DataType::Map(field, _)) => Ok(field.as_ref().clone()),
         (index, DataType::Struct(fields)) => Ok(fields[index].clone()),

tests/it/ffi.rs

@@ -125,6 +125,22 @@ fn list() -> Result<()> {
     test_round_trip(array)
 }
 
+#[test]
+fn fixed_list() -> Result<()> {
+    let data = vec![
+        Some(vec![Some(1i32), Some(2), Some(3)]),
+        None,
+        Some(vec![Some(4), None, Some(6)]),
+    ];
+
+    let mut array = MutableFixedSizeListArray::new(MutablePrimitiveArray::<i32>::new(), 3);
+    array.try_extend(data)?;
+
+    let array: FixedSizeListArray = array.into();
+
+    test_round_trip(array)
+}
+
 #[test]
 fn list_list() -> Result<()> {
     let data = vec![