fix: FFI returns a string/binary instead of requiring the caller to guess the necessary memory

Signed-off-by: Robert Kruszewski <[email protected]>

Author: robert3005

vortex-ffi/src/array.rs

@@ -2,16 +2,18 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 //! FFI interface for working with Vortex Arrays.
-use std::ffi::{c_int, c_void};
-use std::slice;
+use std::ptr;
+use std::sync::Arc;
 
 use vortex::dtype::half::f16;
-use vortex::error::{VortexExpect, VortexUnwrap, vortex_err};
+use vortex::error::{VortexExpect, vortex_err};
 use vortex::{Array, ToCanonical};
 
 use crate::arc_dyn_wrapper;
+use crate::binary::vx_binary;
 use crate::dtype::vx_dtype;
 use crate::error::{try_or_default, vx_error};
+use crate::string::vx_string;
 
 arc_dyn_wrapper!(
     /// Base type for all Vortex arrays.
@@ -133,48 +135,42 @@ ffiarray_get_ptype!(f16);
 ffiarray_get_ptype!(f32);
 ffiarray_get_ptype!(f64);
 
-/// Write the UTF-8 string at `index` in the array into the provided destination buffer, recording
-/// the length in `len`.
+/// Return the utf-8 string at `index` in the array. The pointer will be null if the value at `index` is null.
+/// The caller must free the returned pointer.
 #[unsafe(no_mangle)]
 pub unsafe extern "C-unwind" fn vx_array_get_utf8(
     array: *const vx_array,
     index: u32,
-    dst: *mut c_void,
-    len: *mut c_int,
-) {
+) -> *const vx_string {
     let array = vx_array::as_ref(array);
     let value = array.scalar_at(index as usize);
     let utf8_scalar = value.as_utf8();
     if let Some(buffer) = utf8_scalar.value() {
-        let bytes = buffer.as_bytes();
-        let dst = unsafe { slice::from_raw_parts_mut(dst as *mut u8, bytes.len()) };
-        dst.copy_from_slice(bytes);
-        unsafe { *len = bytes.len().try_into().vortex_unwrap() };
+        vx_string::new(Arc::from(buffer.as_str()))
+    } else {
+        ptr::null()
     }
 }
 
-/// Write the UTF-8 string at `index` in the array into the provided destination buffer, recording
-/// the length in `len`.
+/// Return the binary at `index` in the array. The pointer will be null if the value at `index` is null.
+/// The caller must free the returned pointer.
 #[unsafe(no_mangle)]
 pub unsafe extern "C-unwind" fn vx_array_get_binary(
     array: *const vx_array,
     index: u32,
-    dst: *mut c_void,
-    len: *mut c_int,
-) {
+) -> *const vx_binary {
     let array = vx_array::as_ref(array);
     let value = array.scalar_at(index as usize);
-    let utf8_scalar = value.as_binary();
-    if let Some(bytes) = utf8_scalar.value() {
-        let dst = unsafe { slice::from_raw_parts_mut(dst as *mut u8, bytes.len()) };
-        dst.copy_from_slice(&bytes);
-        unsafe { *len = bytes.len().try_into().vortex_unwrap() };
+    let binary_scalar = value.as_binary();
+    if let Some(bytes) = binary_scalar.value() {
+        vx_binary::new(Arc::from(bytes.as_bytes()))
+    } else {
+        ptr::null()
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use std::ffi::{c_int, c_void};
     use std::ptr;
 
     use vortex::IntoArray;
@@ -185,8 +181,10 @@ mod tests {
     use vortex::validity::Validity;
 
     use crate::array::*;
+    use crate::binary::vx_binary_free;
     use crate::dtype::{vx_dtype_get_variant, vx_dtype_variant};
     use crate::error::vx_error_free;
+    use crate::string::vx_string_free;
 
     #[test]
     fn test_simple() {
@@ -349,35 +347,17 @@ mod tests {
             let utf8_array = VarBinViewArray::from_iter_str(["hello", "world", "test"]);
             let ffi_array = vx_array::new(utf8_array.into_array());
 
-            let mut buffer = vec![0u8; 10];
-            let mut len: c_int = 0;
+            let vx_str1 = vx_array_get_utf8(ffi_array, 0);
+            assert_eq!(vx_string::as_str(vx_str1), "hello");
+            vx_string_free(vx_str1);
 
-            vx_array_get_utf8(
-                ffi_array,
-                0,
-                buffer.as_mut_ptr() as *mut c_void,
-                &raw mut len,
-            );
-            assert_eq!(len, 5);
-            assert_eq!(&buffer[..5], b"hello");
-
-            vx_array_get_utf8(
-                ffi_array,
-                1,
-                buffer.as_mut_ptr() as *mut c_void,
-                &raw mut len,
-            );
-            assert_eq!(len, 5);
-            assert_eq!(&buffer[..5], b"world");
-
-            vx_array_get_utf8(
-                ffi_array,
-                2,
-                buffer.as_mut_ptr() as *mut c_void,
-                &raw mut len,
-            );
-            assert_eq!(len, 4);
-            assert_eq!(&buffer[..4], b"test");
+            let vx_str2 = vx_array_get_utf8(ffi_array, 1);
+            assert_eq!(vx_string::as_str(vx_str2), "world");
+            vx_string_free(vx_str2);
+
+            let vx_str3 = vx_array_get_utf8(ffi_array, 2);
+            assert_eq!(vx_string::as_str(vx_str3), "test");
+            vx_string_free(vx_str3);
 
             vx_array_free(ffi_array);
         }
@@ -393,35 +373,17 @@ mod tests {
             ]);
             let ffi_array = vx_array::new(binary_array.into_array());
 
-            let mut buffer = vec![0u8; 10];
-            let mut len: c_int = 0;
+            let vx_bin1 = vx_array_get_binary(ffi_array, 0);
+            assert_eq!(vx_binary::as_slice(vx_bin1), &[0x01, 0x02, 0x03]);
+            vx_binary_free(vx_bin1);
 
-            vx_array_get_binary(
-                ffi_array,
-                0,
-                buffer.as_mut_ptr() as *mut c_void,
-                &raw mut len,
-            );
-            assert_eq!(len, 3);
-            assert_eq!(&buffer[..3], &[0x01, 0x02, 0x03]);
-
-            vx_array_get_binary(
-                ffi_array,
-                1,
-                buffer.as_mut_ptr() as *mut c_void,
-                &raw mut len,
-            );
-            assert_eq!(len, 2);
-            assert_eq!(&buffer[..2], &[0xFF, 0xEE]);
-
-            vx_array_get_binary(
-                ffi_array,
-                2,
-                buffer.as_mut_ptr() as *mut c_void,
-                &raw mut len,
-            );
-            assert_eq!(len, 4);
-            assert_eq!(&buffer[..4], &[0xAA, 0xBB, 0xCC, 0xDD]);
+            let vx_bin2 = vx_array_get_binary(ffi_array, 1);
+            assert_eq!(vx_binary::as_slice(vx_bin2), &[0xFF, 0xEE]);
+            vx_binary_free(vx_bin2);
+
+            let vx_bin3 = vx_array_get_binary(ffi_array, 2);
+            assert_eq!(vx_binary::as_slice(vx_bin3), &[0xAA, 0xBB, 0xCC, 0xDD]);
+            vx_binary_free(vx_bin3);
 
             vx_array_free(ffi_array);
         }

vortex-ffi/src/binary.rs

@@ -0,0 +1,104 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::ffi::c_char;
+use std::slice;
+
+use crate::arc_dyn_wrapper;
+
+arc_dyn_wrapper!(
+    /// Strings for use within Vortex.
+    [u8],
+    vx_binary
+);
+
+impl vx_binary {
+    #[allow(dead_code)]
+    pub(crate) fn as_slice(ptr: *const vx_binary) -> &'static [u8] {
+        unsafe { slice::from_raw_parts(vx_binary_ptr(ptr).cast(), vx_binary_len(ptr)) }
+    }
+}
+
+/// Create a new Vortex UTF-8 string by copying from a pointer and length.
+#[unsafe(no_mangle)]
+pub unsafe extern "C-unwind" fn vx_binary_new(ptr: *const c_char, len: usize) -> *const vx_binary {
+    let slice = unsafe { slice::from_raw_parts(ptr.cast(), len) };
+    vx_binary::new(slice.into())
+}
+
+/// Return the length of the string in bytes.
+#[unsafe(no_mangle)]
+pub unsafe extern "C-unwind" fn vx_binary_len(ptr: *const vx_binary) -> usize {
+    vx_binary::as_ref(ptr).len()
+}
+
+/// Return the pointer to the string data.
+#[unsafe(no_mangle)]
+pub unsafe extern "C-unwind" fn vx_binary_ptr(ptr: *const vx_binary) -> *const c_char {
+    vx_binary::as_ref(ptr).as_ptr().cast()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_string_new() {
+        unsafe {
+            let test_str = "hello world";
+            let ptr = test_str.as_ptr() as *const c_char;
+            let len = test_str.len();
+
+            let vx_str = vx_binary_new(ptr, len);
+            assert_eq!(vx_binary_len(vx_str), 11);
+            assert_eq!(vx_binary::as_slice(vx_str), "hello world".as_bytes());
+
+            vx_binary_free(vx_str);
+        }
+    }
+
+    #[test]
+    fn test_string_ptr() {
+        unsafe {
+            let test_str = "testing".as_bytes();
+            let vx_str = vx_binary::new(test_str.into());
+
+            let ptr = vx_binary_ptr(vx_str);
+            let len = vx_binary_len(vx_str);
+
+            let slice = slice::from_raw_parts(ptr as *const u8, len);
+            assert_eq!(slice, "testing".as_bytes());
+
+            vx_binary_free(vx_str);
+        }
+    }
+
+    #[test]
+    fn test_empty_string() {
+        unsafe {
+            let empty = "";
+            let ptr = empty.as_ptr() as *const c_char;
+            let vx_str = vx_binary_new(ptr, 0);
+
+            assert_eq!(vx_binary_len(vx_str), 0);
+            assert_eq!(vx_binary::as_slice(vx_str), "".as_bytes());
+
+            vx_binary_free(vx_str);
+        }
+    }
+
+    #[test]
+    fn test_unicode_string() {
+        unsafe {
+            let unicode_str = "Hello 世界 🌍";
+            let ptr = unicode_str.as_ptr() as *const c_char;
+            let len = unicode_str.len();
+
+            let vx_str = vx_binary_new(ptr, len);
+            assert_eq!(vx_binary_len(vx_str), unicode_str.len());
+            assert_eq!(vx_binary::as_slice(vx_str), unicode_str.as_bytes());
+
+            vx_binary_free(vx_str);
+        }
+    }
+}

vortex-ffi/src/dtype.rs

@@ -1,15 +1,16 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use std::ffi::{c_int, c_void};
+use std::ptr;
 use std::sync::Arc;
 
 use vortex::dtype::datetime::{DATE_ID, TIME_ID, TIMESTAMP_ID, TemporalMetadata};
 use vortex::dtype::{DType, DecimalDType};
-use vortex::error::{VortexExpect, VortexUnwrap, vortex_panic};
+use vortex::error::{VortexExpect, vortex_panic};
 
 use crate::arc_wrapper;
 use crate::ptype::vx_ptype;
+use crate::string::vx_string;
 use crate::struct_fields::vx_struct_fields;
 
 arc_wrapper!(
@@ -286,13 +287,9 @@ pub unsafe extern "C-unwind" fn vx_dtype_time_unit(dtype: *const DType) -> u8 {
     metadata.as_ref()[0]
 }
 
-/// Returns the time zone, assuming the type is time.
+/// Returns the time zone, assuming the type is time. Caller is responsible for freeing the returned pointer.
 #[unsafe(no_mangle)]
-pub unsafe extern "C-unwind" fn vx_dtype_time_zone(
-    dtype: *const DType,
-    dst: *mut c_void,
-    len: *mut c_int,
-) {
+pub unsafe extern "C-unwind" fn vx_dtype_time_zone(dtype: *const DType) -> *const vx_string {
     let dtype = unsafe { dtype.as_ref() }.vortex_expect("dtype null");
 
     let DType::Extension(ext_dtype) = dtype else {
@@ -302,13 +299,9 @@ pub unsafe extern "C-unwind" fn vx_dtype_time_zone(
     match TemporalMetadata::try_from(ext_dtype).vortex_expect("timestamp") {
         TemporalMetadata::Timestamp(_, zone) => {
             if let Some(zone) = zone {
-                let bytes = zone.as_bytes();
-                let dst = unsafe { std::slice::from_raw_parts_mut(dst as *mut u8, bytes.len()) };
-                dst.copy_from_slice(bytes);
-                unsafe { *len = bytes.len().try_into().vortex_unwrap() };
+                vx_string::new(zone.into())
             } else {
-                // No time zone, using local timestamps.
-                unsafe { *len = 0 };
+                ptr::null()
             }
         }
         _ => vortex_panic!("DType_time_zone: not a timestamp metadata: {ext_dtype:?}"),

vortex-ffi/src/lib.rs

@@ -8,6 +8,7 @@
 
 mod array;
 mod array_iterator;
+mod binary;
 mod dtype;
 mod error;
 mod file;

vortex-ffi/src/string.rs

@@ -16,7 +16,7 @@ arc_dyn_wrapper!(
 
 impl vx_string {
     #[allow(dead_code)]
-    pub(crate) fn as_str<'a>(ptr: *const vx_string) -> &'a str {
+    pub(crate) fn as_str(ptr: *const vx_string) -> &'static str {
         unsafe {
             str::from_utf8_unchecked(slice::from_raw_parts(
                 vx_string_ptr(ptr).cast(),