For review

apache · Jun 17, 2024 · 017f6cb · 017f6cb
1 parent eb51190
commit 017f6cb
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 11 deletions.
diff --git 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
@@ -71,9 +71,16 @@ impl Buffer {
         }
     }
 
-    /// Returns the internal byte buffer.
-    pub fn get_bytes(&self) -> Arc<Bytes> {
-        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<u8> {
+        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
@@ -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::<i32>()
+                    b.ptr_offset() / std::mem::size_of::<i32>()
                 })
             }
             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::<i64>()
+                    b.ptr_offset() / std::mem::size_of::<i64>()
                 })
             }
             _ => 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),