[ffi] Fix arrow-array null_count error during conversion from C to Rust

apache · Nov 2, 2024 · 72b7f0b · 72b7f0b
1 parent 37cd34d
commit 72b7f0b
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 7 deletions.
diff --git a/arrow-array/src/ffi.rs b/arrow-array/src/ffi.rs
@@ -301,10 +301,7 @@ impl ImportedArrowArray<'_> {
     fn consume(self) -> Result<ArrayData> {
         let len = self.array.len();
         let offset = self.array.offset();
-        let null_count = match &self.data_type {
-            DataType::Null => 0,
-            _ => self.array.null_count(),
-        };
+        let null_count = self.array.null_count();
 
         let data_layout = layout(&self.data_type);
         let buffers = self.buffers(data_layout.can_contain_null_mask, data_layout.variadic)?;
@@ -329,7 +326,7 @@ impl ImportedArrowArray<'_> {
             ArrayData::new_unchecked(
                 self.data_type,
                 len,
-                Some(null_count),
+                null_count,
                 null_bit_buffer,
                 offset,
                 buffers,
@@ -635,6 +632,31 @@ mod tests_to_then_from_ffi {
     }
     // case with nulls is tested in the docs, through the example on this module.
 
+    #[test]
+    fn test_null_count_handling() {
+        let int32_data = ArrayData::builder(DataType::Int32)
+            .len(10)
+            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))
+            .build()
+            .unwrap();
+        let mut ffi_array = FFI_ArrowArray::new(&int32_data);
+        assert_eq!(Some(0), ffi_array.null_count());
+        // Simulating uninitialized state
+        ffi_array.set_null_count(-1);
+        assert_eq!(None, ffi_array.null_count());
+        let int32_data = unsafe { from_ffi_and_data_type(ffi_array, DataType::Int32) }.unwrap();
+        assert_eq!(0, int32_data.null_count());
+
+        let null_data = &ArrayData::new_null(&DataType::Null, 10);
+        let mut ffi_array = FFI_ArrowArray::new(null_data);
+        assert_eq!(Some(10), ffi_array.null_count());
+        // Simulating uninitialized state
+        ffi_array.set_null_count(-1);
+        assert_eq!(None, ffi_array.null_count());
+        let null_data = unsafe { from_ffi_and_data_type(ffi_array, DataType::Null) }.unwrap();
+        assert_eq!(0, null_data.null_count());
+    }
+
     fn test_generic_string<Offset: OffsetSizeTrait>() -> Result<()> {
         // create an array natively
         let array = GenericStringArray::<Offset>::from(vec![Some("a"), None, Some("aaa")]);

diff --git a/arrow-data/src/ffi.rs b/arrow-data/src/ffi.rs
@@ -267,8 +267,19 @@ impl FFI_ArrowArray {
 
     /// the null count of the array
     #[inline]
-    pub fn null_count(&self) -> usize {
-        self.null_count as usize
+    pub fn null_count(&self) -> Option<usize> {
+        // MAY be -1 if not yet computed
+        if self.null_count >= 0 {
+            Some(self.null_count as usize)
+        } else {
+            None
+        }
+    }
+
+    /// set the null count of the array
+    #[inline]
+    pub fn set_null_count(&mut self, null_count: i64) {
+        self.null_count = null_count;
     }
 
     /// Returns the buffer at the provided index