feat: Enable global runtime for duckdb (#16)

Author: Xuanwo

rust/lib.rs

@@ -8,8 +8,8 @@ use arrow::array::{Array, RecordBatch, StructArray};
 use arrow::datatypes::{DataType, Schema};
 use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema};
 use lance::Dataset;
-use tokio::runtime::Runtime;
 
+mod runtime;
 mod scanner;
 
 use scanner::LanceStream;
@@ -32,14 +32,9 @@ pub unsafe extern "C" fn lance_open_dataset(path: *const c_char) -> *mut c_void
         }
     };
 
-    let runtime = match Runtime::new() {
-        Ok(rt) => rt,
-        Err(_) => return ptr::null_mut(),
-    };
-
-    let dataset = match runtime.block_on(Dataset::open(path_str)) {
-        Ok(ds) => Arc::new(ds),
-        Err(_) => return ptr::null_mut(),
+    let dataset = match runtime::block_on(Dataset::open(path_str)) {
+        Ok(Ok(ds)) => Arc::new(ds),
+        _ => return ptr::null_mut(),
     };
 
     let handle = Box::new(DatasetHandle { dataset });
@@ -110,13 +105,7 @@ pub unsafe extern "C" fn lance_create_stream(dataset: *mut c_void) -> *mut c_voi
 
     let handle = unsafe { &*(dataset as *const DatasetHandle) };
 
-    // Create a new runtime for this stream to avoid conflicts
-    let runtime = match Runtime::new() {
-        Ok(rt) => rt,
-        Err(_) => return ptr::null_mut(),
-    };
-
-    match LanceStream::new(&handle.dataset, runtime) {
+    match LanceStream::new(&handle.dataset) {
         Ok(stream) => Box::into_raw(Box::new(stream)) as *mut c_void,
         Err(_) => ptr::null_mut(),
     }

rust/runtime.rs

@@ -0,0 +1,22 @@
+use std::future::Future;
+use std::io;
+use std::sync::OnceLock;
+
+use tokio::runtime::{Handle, Runtime};
+
+static RUNTIME: OnceLock<Result<Runtime, io::Error>> = OnceLock::new();
+
+pub fn runtime() -> Result<&'static Runtime, io::Error> {
+    match RUNTIME.get_or_init(Runtime::new) {
+        Ok(rt) => Ok(rt),
+        Err(err) => Err(io::Error::new(err.kind(), err.to_string())),
+    }
+}
+
+pub fn handle() -> Result<Handle, io::Error> {
+    Ok(runtime()?.handle().clone())
+}
+
+pub fn block_on<F: Future>(future: F) -> Result<F::Output, io::Error> {
+    Ok(runtime()?.block_on(future))
+}

rust/scanner.rs

@@ -3,23 +3,24 @@ use std::pin::Pin;
 use arrow::array::RecordBatch;
 use futures::stream::Stream;
 use lance::Dataset;
-use tokio::runtime::Runtime;
+use tokio::runtime::Handle;
 
 /// A stream wrapper that holds the Lance RecordBatchStream
 pub struct LanceStream {
-    runtime: Runtime,
+    handle: Handle,
     stream: Pin<Box<dyn Stream<Item = Result<RecordBatch, lance::Error>> + Send>>,
 }
 
 impl LanceStream {
     /// Create a new stream from a dataset path
-    pub fn new(dataset: &Dataset, runtime: Runtime) -> Result<Self, Box<dyn std::error::Error>> {
+    pub fn new(dataset: &Dataset) -> Result<Self, Box<dyn std::error::Error>> {
+        let handle = crate::runtime::handle()?;
         let scanner = dataset.scan();
 
-        let stream = runtime.block_on(async { scanner.try_into_stream().await })?;
+        let stream = handle.block_on(async { scanner.try_into_stream().await })?;
 
         Ok(Self {
-            runtime,
+            handle,
             stream: Box::pin(stream),
         })
     }
@@ -28,7 +29,7 @@ impl LanceStream {
     pub fn next(&mut self) -> Option<RecordBatch> {
         use futures::StreamExt;
 
-        self.runtime.block_on(async {
+        self.handle.block_on(async {
             match self.stream.next().await {
                 Some(Ok(batch)) => Some(batch),
                 _ => None,