wire duckdb

Signed-off-by: Onur Satici <[email protected]>

Author: onursatici

vortex-duckdb/src/copy.rs

@@ -25,7 +25,6 @@ use vortex::io::runtime::Task;
 use vortex::io::runtime::current::CurrentThreadWorkerPool;
 use vortex::io::session::RuntimeSessionExt;
 
-use crate::RUNTIME;
 use crate::SESSION;
 use crate::convert::data_chunk_to_vortex;
 use crate::convert::from_duckdb_table;
@@ -53,7 +52,7 @@ pub struct GlobalState {
     // Note that this is optional and without it, we would only drive the task when DuckDB calls
     // into us, and we call `RUNTIME.block_on`.
     #[allow(dead_code)]
-    worker_pool: CurrentThreadWorkerPool,
+    worker_pool: Option<CurrentThreadWorkerPool>,
 }
 
 impl CopyFunction for VortexCopyFunction {
@@ -86,7 +85,7 @@ impl CopyFunction for VortexCopyFunction {
         chunk: &mut DataChunk,
     ) -> VortexResult<()> {
         let chunk = data_chunk_to_vortex(bind_data.fields.names(), chunk);
-        RUNTIME.block_on(async {
+        crate::blocking_runtime().block_on(async {
             init_global
                 .sink
                 .as_mut()
@@ -103,7 +102,7 @@ impl CopyFunction for VortexCopyFunction {
         _bind_data: &Self::BindData,
         init_global: &mut Self::GlobalState,
     ) -> VortexResult<()> {
-        RUNTIME.block_on(async {
+        crate::blocking_runtime().block_on(async {
             if let Some(sink) = init_global.sink.take() {
                 drop(sink)
             }
@@ -131,11 +130,8 @@ impl CopyFunction for VortexCopyFunction {
             SESSION.write_options().write(&mut file, array_stream).await
         });
 
-        let worker_pool = RUNTIME.new_pool();
-        worker_pool.set_workers_to_available_parallelism();
-
         Ok(GlobalState {
-            worker_pool,
+            worker_pool: None,
             write_task: Mutex::new(Some(writer)),
             sink: Some(sink),
         })

vortex-duckdb/src/lib.rs

@@ -6,12 +6,15 @@
 use std::ffi::CStr;
 use std::ffi::c_char;
 use std::sync::LazyLock;
+use std::sync::OnceLock;
 
 use vortex::VortexSessionDefault;
 use vortex::error::VortexExpect;
 use vortex::error::VortexResult;
 use vortex::io::runtime::BlockingRuntime;
+use vortex::io::runtime::Handle;
 use vortex::io::runtime::current::CurrentThreadRuntime;
+pub use vortex::io::runtime::current::CurrentThreadRuntime as DuckdbDefaultRuntime;
 use vortex::io::session::RuntimeSessionExt;
 use vortex::session::VortexSession;
 
@@ -40,9 +43,35 @@ mod copy;
 mod e2e_test;
 
 // A global runtime for Vortex operations within DuckDB.
-static RUNTIME: LazyLock<CurrentThreadRuntime> = LazyLock::new(CurrentThreadRuntime::new);
-static SESSION: LazyLock<VortexSession> =
-    LazyLock::new(|| VortexSession::default().with_handle(RUNTIME.handle()));
+static DEFAULT_RUNTIME: LazyLock<CurrentThreadRuntime> = LazyLock::new(CurrentThreadRuntime::new);
+static RUNTIME_OVERRIDE: OnceLock<Handle> = OnceLock::new();
+static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+    let handle = RUNTIME_OVERRIDE
+        .get()
+        .cloned()
+        .unwrap_or_else(|| DEFAULT_RUNTIME.handle());
+    VortexSession::default().with_handle(handle)
+});
+
+/// Configure the runtime handle used by the DuckDB extension.
+///
+/// Must be called before the session is first accessed (e.g., before registering table functions).
+pub fn configure_runtime(handle: Handle) {
+    drop(RUNTIME_OVERRIDE.set(handle));
+}
+
+/// Returns the effective runtime handle used by the DuckDB extension.
+pub fn runtime_handle() -> Handle {
+    RUNTIME_OVERRIDE
+        .get()
+        .cloned()
+        .unwrap_or_else(|| DEFAULT_RUNTIME.handle())
+}
+
+/// Returns the blocking runtime used for synchronous work.
+pub fn blocking_runtime() -> &'static CurrentThreadRuntime {
+    &DEFAULT_RUNTIME
+}
 
 /// Register Vortex extension configuration options with DuckDB.
 /// This must be called before `register_table_functions` to take effect.

vortex-duckdb/src/scan.rs

@@ -43,8 +43,9 @@ use vortex::file::VortexOpenOptions;
 use vortex::io::runtime::BlockingRuntime;
 use vortex::io::runtime::current::ThreadSafeIterator;
 
-use crate::RUNTIME;
 use crate::SESSION;
+use crate::blocking_runtime;
+use crate::runtime_handle;
 use crate::convert::try_from_bound_expression;
 use crate::convert::try_from_table_filter;
 use crate::duckdb;
@@ -274,9 +275,8 @@ impl TableFunction for VortexTableFunction {
 
         log::trace!("running scan with max_threads {max_threads}");
 
-        let (file_urls, _metadata) = RUNTIME.block_on(Compat::new(expand_glob(
-            file_glob_string.as_ref().as_string(),
-        )))?;
+        let (file_urls, _metadata) = blocking_runtime()
+            .block_on(Compat::new(expand_glob(file_glob_string.as_ref().as_string())))?;
 
         // The first file is skipped in `create_file_paths_queue`.
         let Some(first_file_url) = file_urls.first() else {
@@ -285,7 +285,7 @@ impl TableFunction for VortexTableFunction {
 
         let footer_cache = FooterCache::new(ctx.object_cache());
         let entry = footer_cache.entry(first_file_url.as_ref());
-        let first_file = RUNTIME.block_on(async move {
+        let first_file = blocking_runtime().block_on(async move {
             let options = entry.apply_to_file(SESSION.open_options());
             let file = open_file(first_file_url.clone(), options).await?;
             entry.put_if_absent(|| file.footer().clone());
@@ -373,7 +373,8 @@ impl TableFunction for VortexTableFunction {
         let client_context = init_input.client_context()?;
         let object_cache = client_context.object_cache();
 
-        let handle = RUNTIME.handle();
+        let scan_handle = runtime_handle();
+        let handle = scan_handle.clone();
         let first_file = bind_data.first_file.clone();
         let scan_streams = stream::iter(bind_data.file_urls.clone())
             .enumerate()
@@ -423,7 +424,7 @@ impl TableFunction for VortexTableFunction {
             .filter_map(|result| async move { result.transpose() });
 
         Ok(VortexGlobalData {
-            iterator: RUNTIME.block_on_stream_thread_safe(move |_| MultiScan {
+            iterator: blocking_runtime().block_on_stream_thread_safe(move |_| MultiScan {
                 streams: scan_streams.boxed(),
                 streams_finished: false,
                 select_all: Default::default(),