From 017f6cb383d7c29d483b32a4770ac036d65cb4bb Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Mon, 17 Jun 2024 12:43:38 -0700 Subject: [PATCH] For review --- arrow-array/src/ffi.rs | 6 +++--- arrow-buffer/src/buffer/immutable.rs | 13 ++++++++++--- arrow-data/src/ffi.rs | 8 +++----- 3 files changed, 16 insertions(+), 11 deletions(-) diff --git a/arrow-array/src/ffi.rs b/arrow-array/src/ffi.rs index ebc5427b8951..fe755b91f765 100644 --- a/arrow-array/src/ffi.rs +++ b/arrow-array/src/ffi.rs @@ -1461,7 +1461,7 @@ mod tests_from_ffi { test_round_trip(&imported_array.consume()?) } - fn export_string(array: StringArray) -> StringArray { + fn roundtrip_string_array(array: StringArray) -> StringArray { let data = array.into_data(); let array = FFI_ArrowArray::new(&data); @@ -1490,7 +1490,7 @@ mod tests_from_ffi { let string_array = StringArray::from(strings); - let imported = export_string(string_array.clone()); + let imported = roundtrip_string_array(string_array.clone()); assert_eq!(imported.len(), 1000); assert_eq!(imported.value(0), "string: 0"); assert_eq!(imported.value(499), "string: 499"); @@ -1503,7 +1503,7 @@ mod tests_from_ffi { let slice = string_array.slice(500, 500); - let imported = export_string(slice); + let imported = roundtrip_string_array(slice); assert_eq!(imported.len(), 500); assert_eq!(imported.value(0), "string: 500"); assert_eq!(imported.value(499), "string: 999"); diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index 8f2fe226e107..5eeea76db980 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -71,9 +71,16 @@ impl Buffer { } } - /// Returns the internal byte buffer. - pub fn get_bytes(&self) -> Arc { - self.data.clone() + /// Returns the offset, in bytes, of `Self::ptr` to `Self::data` + /// + /// self.ptr and self.data can be different after slicing or advancing the buffer. + pub unsafe fn ptr_offset(&self) -> usize { + self.ptr.offset_from(self.data.ptr().as_ptr()) as usize + } + + /// Returns the pointer to the start of the buffer without the offset. + pub fn data_ptr(&self) -> NonNull { + self.data.ptr() } /// Create a [`Buffer`] from the provided [`Vec`] without copying diff --git a/arrow-data/src/ffi.rs b/arrow-data/src/ffi.rs index 729a0c59c83f..a3e22e10ed3d 100644 --- a/arrow-data/src/ffi.rs +++ b/arrow-data/src/ffi.rs @@ -141,8 +141,7 @@ impl FFI_ArrowArray { // the `FFI_ArrowArray` offset field. Some(unsafe { let b = &data.buffers()[0]; - b.as_ptr().offset_from(b.get_bytes().ptr().as_ptr()) as usize - / std::mem::size_of::() + b.ptr_offset() / std::mem::size_of::() }) } DataType::LargeUtf8 | DataType::LargeBinary => { @@ -153,8 +152,7 @@ impl FFI_ArrowArray { // the `FFI_ArrowArray` offset field. Some(unsafe { let b = &data.buffers()[0]; - b.as_ptr().offset_from(b.get_bytes().ptr().as_ptr()) as usize - / std::mem::size_of::() + b.ptr_offset() / std::mem::size_of::() }) } _ => None, @@ -195,7 +193,7 @@ impl FFI_ArrowArray { ) => { // For offset buffer, take original pointer without offset. // Buffer offset should be handled by `FFI_ArrowArray` offset field. - Some(b.get_bytes().ptr().as_ptr() as *const c_void) + Some(b.data_ptr().as_ptr() as *const c_void) } // For other buffers, note that `raw_data` takes into account the buffer's offset _ => Some(b.as_ptr() as *const c_void),