ffi: create data source from memory buffer (#8426)

Match rust API: buffer is borrowed under the condition it won't be
modified
until data source is freed

Signed-off-by: Mikhail Kot <mikhail@spiraldb.com>

Author: myrrc

Cargo.lock

@@ -23,6 +23,7 @@ all-features = true
 arrow-array = { workspace = true, features = ["ffi"] }
 arrow-schema = { workspace = true }
 async-fs = { workspace = true }
+bytes = { workspace = true }
 futures = { workspace = true }
 itertools = { workspace = true }
 mimalloc = { workspace = true, optional = true }

vortex-ffi/Cargo.toml

@@ -884,6 +884,21 @@ void vx_data_source_free(const vx_data_source *ptr);
 const vx_data_source *
 vx_data_source_new(const vx_session *session, const vx_data_source_options *options, vx_error **err);
 
+/**
+ * Create a data source from a single in-memory Vortex file.
+ *
+ * "buffer_len" is the length of "buffer" in bytes.
+ * The bytes are borrowed, not copied: the caller must keep "buffer" alive and
+ * unmodified until the data source is freed.
+ *
+ * The returned pointer is owned by the caller and must be freed with
+ * vx_data_source_free.
+ *
+ * On error, returns NULL and sets "err".
+ */
+const vx_data_source *
+vx_data_source_new_buffer(const vx_session *session, const void *buffer, size_t buffer_len, vx_error **err);
+
 /**
  * Return the schema of the data source as a non-owned dtype.
  * The returned pointer is valid as long as "ds" is alive. Do not free it.

vortex-ffi/cinclude/vortex.h

@@ -5,16 +5,22 @@
 
 use std::ffi::CStr;
 use std::ffi::c_char;
+use std::ffi::c_void;
 use std::ptr;
+use std::slice;
 use std::sync::Arc;
 
+use bytes::Bytes;
+use vortex::buffer::ByteBuffer;
 use vortex::error::VortexResult;
 use vortex::error::vortex_ensure;
 use vortex::expr::stats::Precision::Absent;
 use vortex::expr::stats::Precision::Exact;
 use vortex::expr::stats::Precision::Inexact;
+use vortex::file::OpenOptionsSessionExt;
 use vortex::file::multi::MultiFileDataSource;
 use vortex::io::runtime::BlockingRuntime;
+use vortex::layout::scan::multi::MultiLayoutDataSource;
 use vortex::scan::DataSource;
 use vortex::scan::DataSourceRef;
 
@@ -104,6 +110,38 @@ pub unsafe extern "C-unwind" fn vx_data_source_new(
     })
 }
 
+/// Create a data source from a single in-memory Vortex file.
+///
+/// "buffer_len" is the length of "buffer" in bytes.
+/// The bytes are borrowed, not copied: the caller must keep "buffer" alive and
+/// unmodified until the data source is freed.
+///
+/// The returned pointer is owned by the caller and must be freed with
+/// vx_data_source_free.
+///
+/// On error, returns NULL and sets "err".
+#[unsafe(no_mangle)]
+pub unsafe extern "C-unwind" fn vx_data_source_new_buffer(
+    session: *const vx_session,
+    buffer: *const c_void,
+    buffer_len: usize,
+    err: *mut *mut vx_error,
+) -> *const vx_data_source {
+    try_or(err, ptr::null(), || {
+        vortex_ensure!(!session.is_null());
+        vortex_ensure!(!buffer.is_null());
+
+        let session = vx_session::as_ref(session);
+        let bytes: &'static [u8] =
+            unsafe { slice::from_raw_parts(buffer.cast::<u8>(), buffer_len) };
+        let buffer = ByteBuffer::from(Bytes::from_static(bytes));
+        let file = session.open_options().open_buffer(buffer)?;
+        let ds = MultiLayoutDataSource::new_with_first(file.layout_reader()?, Vec::new(), session);
+
+        Ok(vx_data_source::new(Arc::new(ds) as DataSourceRef))
+    })
+}
+
 /// Return the schema of the data source as a non-owned dtype.
 /// The returned pointer is valid as long as "ds" is alive. Do not free it.
 #[unsafe(no_mangle)]
@@ -140,12 +178,15 @@ pub unsafe extern "C-unwind" fn vx_data_source_get_row_count(
 #[cfg(test)]
 mod tests {
     use std::ffi::CString;
+    use std::ffi::c_void;
+    use std::fs::read;
     use std::ptr;
 
     use crate::data_source::vx_data_source_dtype;
     use crate::data_source::vx_data_source_free;
     use crate::data_source::vx_data_source_get_row_count;
     use crate::data_source::vx_data_source_new;
+    use crate::data_source::vx_data_source_new_buffer;
     use crate::data_source::vx_data_source_options;
     use crate::dtype::vx_dtype;
     use crate::scan::vx_estimate;
@@ -220,4 +261,39 @@ mod tests {
             vx_session_free(session);
         }
     }
+
+    #[test]
+    #[cfg_attr(miri, ignore)]
+    fn test_buffer() {
+        unsafe {
+            let session = vx_session_new();
+            let (sample, struct_array) = write_sample(session);
+
+            let mut error = ptr::null_mut();
+            let ds = vx_data_source_new_buffer(session, ptr::null(), 0, &raw mut error);
+            assert_error(error);
+            assert!(ds.is_null());
+
+            let file = read(sample).unwrap();
+            let ds = vx_data_source_new_buffer(
+                session,
+                file.as_ptr() as *const c_void,
+                file.len(),
+                &raw mut error,
+            );
+            assert_no_error(error);
+            assert!(!ds.is_null());
+
+            let dtype = vx_dtype::as_ref(vx_data_source_dtype(ds));
+            assert_eq!(dtype, struct_array.dtype());
+
+            let mut row_count = vx_estimate::default();
+            vx_data_source_get_row_count(ds, &raw mut row_count);
+            assert_eq!(row_count.r#type, vx_estimate_type::VX_ESTIMATE_EXACT);
+            assert_eq!(row_count.estimate, SAMPLE_ROWS as u64);
+
+            vx_data_source_free(ds);
+            vx_session_free(session);
+        }
+    }
 }

vortex-ffi/src/data_source.rs

@@ -9,6 +9,7 @@
 #include <random>
 #include <thread>
 #include <unistd.h>
+#include <fstream>
 
 using FFI_ArrowArrayStream = ArrowArrayStream;
 using FFI_ArrowArray = ArrowArray;
@@ -218,6 +219,11 @@ TEST_CASE("Creating datasources", "[datasource]") {
     REQUIRE(error != nullptr);
     vx_error_free(error);
 
+    ds = vx_data_source_new_buffer(session, nullptr, 0, &error);
+    REQUIRE(ds == nullptr);
+    REQUIRE(error != nullptr);
+    vx_error_free(error);
+
     TempPath file = write_sample(session);
     opts.paths = file.c_str();
     ds = vx_data_source_new(session, &opts, &error);
@@ -383,24 +389,8 @@ TEST_CASE("Requesting scans", "[datasource]") {
     }
 }
 
-TEST_CASE("Basic scan", "[datasource]") {
-    vx_session *session = vx_session_new();
-    defer {
-        vx_session_free(session);
-    };
-    TempPath path = write_sample(session);
+void basic_scan(const vx_data_source *ds) {
     vx_error *error = nullptr;
-
-    vx_data_source_options ds_options = {};
-    ds_options.paths = path.c_str();
-
-    const vx_data_source *ds = vx_data_source_new(session, &ds_options, &error);
-    require_no_error(error);
-    REQUIRE(ds != nullptr);
-    defer {
-        vx_data_source_free(ds);
-    };
-
     vx_estimate estimate = {};
     vx_scan *scan = vx_data_source_scan(ds, nullptr, &estimate, &error);
     require_no_error(error);
@@ -439,6 +429,50 @@ TEST_CASE("Basic scan", "[datasource]") {
     verify_sample_array(array);
 }
 
+TEST_CASE("Basic scan", "[datasource]") {
+    vx_session *session = vx_session_new();
+    defer {
+        vx_session_free(session);
+    };
+    TempPath path = write_sample(session);
+    vx_error *error = nullptr;
+
+    vx_data_source_options ds_options = {};
+    ds_options.paths = path.c_str();
+    const vx_data_source *ds = vx_data_source_new(session, &ds_options, &error);
+    require_no_error(error);
+    REQUIRE(ds != nullptr);
+    defer {
+        vx_data_source_free(ds);
+    };
+
+    basic_scan(ds);
+}
+
+TEST_CASE("Basic scan from memory", "[datasource]") {
+    vx_session *session = vx_session_new();
+    defer {
+        vx_session_free(session);
+    };
+    TempPath path = write_sample(session);
+
+    std::ifstream file(path, std::ios::binary | std::ios::ate);
+    const std::streamsize size = file.tellg();
+    file.seekg(0, std::ios::beg);
+    std::vector<char> buffer(size);
+    REQUIRE(file.read(buffer.data(), size));
+
+    vx_error *error = nullptr;
+    const vx_data_source *ds = vx_data_source_new_buffer(session, buffer.data(), size, &error);
+    require_no_error(error);
+    REQUIRE(ds != nullptr);
+    defer {
+        vx_data_source_free(ds);
+    };
+
+    basic_scan(ds);
+}
+
 TEST_CASE("Multithreaded scan", "[datasource]") {
     vx_session *session = vx_session_new();
     defer {