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

fix: #5373 

I think we could have these do RAII but we need something better than
just scalar_at

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

---------

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

Author: robert3005

vortex-ffi/cinclude/vortex.h

@@ -273,6 +273,11 @@ typedef struct vx_array_iterator vx_array_iterator;
  */
 typedef struct vx_array_sink vx_array_sink;
 
+/**
+ * Strings for use within Vortex.
+ */
+typedef struct vx_binary vx_binary;
+
 /**
  * A Vortex data type.
  *
@@ -461,16 +466,18 @@ double vx_array_get_f64(const vx_array *array, uint32_t index);
 double vx_array_get_storage_f64(const vx_array *array, uint32_t index);
 
 /**
- * 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.
  */
-void vx_array_get_utf8(const vx_array *array, uint32_t index, void *dst, int *len);
+const vx_string *vx_array_get_utf8(const vx_array *array,
+                                   uint32_t index);
 
 /**
- * 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.
  */
-void vx_array_get_binary(const vx_array *array, uint32_t index, void *dst, int *len);
+const vx_binary *vx_array_get_binary(const vx_array *array,
+                                     uint32_t index);
 
 /**
  * Free an owned [`vx_array_iterator`] object.
@@ -488,6 +495,34 @@ void vx_array_iterator_free(vx_array_iterator *ptr);
 const vx_array *vx_array_iterator_next(vx_array_iterator *iter,
                                        vx_error **error_out);
 
+/**
+ * Clone a borrowed [`vx_binary`], returning an owned [`vx_binary`].
+ *
+ *
+ * Must be released with [`vx_binary_free`].
+ */
+const vx_binary *vx_binary_clone(const vx_binary *ptr);
+
+/**
+ * Free an owned [`vx_binary`] object.
+ */
+void vx_binary_free(const vx_binary *ptr);
+
+/**
+ * Create a new Vortex UTF-8 string by copying from a pointer and length.
+ */
+const vx_binary *vx_binary_new(const char *ptr, size_t len);
+
+/**
+ * Return the length of the string in bytes.
+ */
+size_t vx_binary_len(const vx_binary *ptr);
+
+/**
+ * Return the pointer to the string data.
+ */
+const char *vx_binary_ptr(const vx_binary *ptr);
+
 /**
  * Clone a borrowed [`vx_dtype`], returning an owned [`vx_dtype`].
  *
@@ -629,9 +664,9 @@ bool vx_dtype_is_timestamp(const DType *dtype);
 uint8_t vx_dtype_time_unit(const DType *dtype);
 
 /**
- * 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.
  */
-void vx_dtype_time_zone(const DType *dtype, void *dst, int *len);
+const vx_string *vx_dtype_time_zone(const DType *dtype);
 
 /**
  * Free an owned [`vx_error`] object.

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);
+        }
+    }
+}