Skip to content

Commit

Permalink
[ffi] Fix arrow-array null_count error during conversion from C to Ru…
Browse files Browse the repository at this point in the history
…st (apache#6674)

* Fix arrow-array null_count error during conversion from C to Rust

* add ffi_array.null_count_opt

* add Safety to follow clippy

* add Safety to follow clippy
  • Loading branch information
adbmal authored Nov 5, 2024
1 parent 350ea26 commit 88b296c
Show file tree
Hide file tree
Showing 2 changed files with 50 additions and 3 deletions.
38 changes: 35 additions & 3 deletions arrow-array/src/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -302,8 +302,8 @@ impl ImportedArrowArray<'_> {
let len = self.array.len();
let offset = self.array.offset();
let null_count = match &self.data_type {
DataType::Null => 0,
_ => self.array.null_count(),
DataType::Null => Some(0),
_ => self.array.null_count_opt(),
};

let data_layout = layout(&self.data_type);
Expand All @@ -329,7 +329,7 @@ impl ImportedArrowArray<'_> {
ArrayData::new_unchecked(
self.data_type,
len,
Some(null_count),
null_count,
null_bit_buffer,
offset,
buffers,
Expand Down Expand Up @@ -635,6 +635,38 @@ 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]))
.null_bit_buffer(Some(Buffer::from([0b01011111, 0b00000001])))
.build()
.unwrap();
let mut ffi_array = FFI_ArrowArray::new(&int32_data);
assert_eq!(3, ffi_array.null_count());
assert_eq!(Some(3), ffi_array.null_count_opt());
// Simulating uninitialized state
unsafe {
ffi_array.set_null_count(-1);
}
assert_eq!(None, ffi_array.null_count_opt());
let int32_data = unsafe { from_ffi_and_data_type(ffi_array, DataType::Int32) }.unwrap();
assert_eq!(3, int32_data.null_count());

let null_data = &ArrayData::new_null(&DataType::Null, 10);
let mut ffi_array = FFI_ArrowArray::new(null_data);
assert_eq!(10, ffi_array.null_count());
assert_eq!(Some(10), ffi_array.null_count_opt());
// Simulating uninitialized state
unsafe {
ffi_array.set_null_count(-1);
}
assert_eq!(None, ffi_array.null_count_opt());
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")]);
Expand Down
15 changes: 15 additions & 0 deletions arrow-data/src/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -271,6 +271,21 @@ impl FFI_ArrowArray {
self.null_count as usize
}

/// Returns the null count, checking for validity
#[inline]
pub fn null_count_opt(&self) -> Option<usize> {
usize::try_from(self.null_count).ok()
}

/// Set the null count of the array
///
/// # Safety
/// Null count must match that of null buffer
#[inline]
pub unsafe fn set_null_count(&mut self, null_count: i64) {
self.null_count = null_count;
}

/// Returns the buffer at the provided index
///
/// # Panic
Expand Down

0 comments on commit 88b296c

Please sign in to comment.