chore[duckdb]: set max threads (#5062)

This allows setting the max thread count used by the vortex-duckdb
extension, also query_bench `-t X` will set the thread count. This is
useful for a single threaded run of a query.

---------

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

bench-vortex/src/benchmark_driver.rs

@@ -106,6 +106,7 @@ pub fn run_benchmark<B: Benchmark>(benchmark: B, config: DriverConfig) -> Result
             config.disable_datafusion_cache,
             config.emit_plan,
             config.delete_duckdb_database,
+            config.threads,
         )?;
 
         tokio_runtime.block_on(benchmark.register_tables(&engine_ctx, target.format()))?;
@@ -346,6 +347,7 @@ fn run_explain_query<B: Benchmark>(
             config.disable_datafusion_cache,
             config.emit_plan,
             config.delete_duckdb_database,
+            config.threads,
         )?;
 
         tokio_runtime.block_on(benchmark.register_tables(&engine_ctx, target.format()))?;

bench-vortex/src/benchmark_trait.rs

@@ -20,6 +20,7 @@ pub trait Benchmark {
         disable_datafusion_cache: bool,
         emit_plan: bool,
         delete_duckdb_database: bool,
+        threads: Option<usize>,
     ) -> Result<EngineCtx> {
         let engine = target.engine();
         let format = target.format();
@@ -37,6 +38,7 @@ pub trait Benchmark {
                     self.dataset(),
                     format,
                     delete_duckdb_database,
+                    threads,
                 )?)
             }
             _ => unreachable!("engine not supported"),

bench-vortex/src/engines/ddb/mod.rs

@@ -9,6 +9,7 @@ use log::trace;
 use url::Url;
 use vortex::error::VortexExpect;
 use vortex_duckdb::duckdb::{Config, Connection, Database};
+use vortex_duckdb::register_extension_options;
 
 use crate::statpopgen::StatPopGenBenchmark;
 use crate::{BenchmarkDataset, Format, IdempotentPath};
@@ -33,10 +34,16 @@ pub struct DuckDBCtx {
     pub db: Database,
     pub connection: Connection,
     pub db_path: Option<PathBuf>,
+    pub threads: Option<usize>,
 }
 
 impl DuckDBCtx {
-    pub fn new(dataset: BenchmarkDataset, format: Format, delete_database: bool) -> Result<Self> {
+    pub fn new(
+        dataset: BenchmarkDataset,
+        format: Format,
+        delete_database: bool,
+        threads: Option<usize>,
+    ) -> Result<Self> {
         let dir = match dataset {
             BenchmarkDataset::ClickBench { flavor, .. } => {
                 format!("clickbench_{}/{}", flavor, format.name()).to_data_path()
@@ -60,22 +67,30 @@ impl DuckDBCtx {
             std::fs::remove_file(&db_path)?;
         }
 
-        let (db, connection) = Self::open_and_setup_database(Some(db_path.clone()))?;
+        let (db, connection) = Self::open_and_setup_database(Some(db_path.clone()), threads)?;
 
         Ok(Self {
             db,
             connection,
             db_path: Some(db_path),
+            threads,
         })
     }
 
-    pub fn open_and_setup_database(path: Option<PathBuf>) -> Result<(Database, Connection)> {
+    pub fn open_and_setup_database(
+        path: Option<PathBuf>,
+        threads: Option<usize>,
+    ) -> Result<(Database, Connection)> {
         let config = Config::new().vortex_expect("failed to create duckdb config");
 
+        // Register Vortex extension options before creating connection
+        register_extension_options(&config);
+
         let db = match path {
             Some(path) => Database::open_with_config(path, config),
             None => Database::open_in_memory_with_config(config),
         }?;
+
         let connection = db.connect()?;
         vortex_duckdb::register_table_functions(&connection)?;
 
@@ -90,6 +105,11 @@ impl DuckDBCtx {
         // parquet_metadata_cache" when running DuckDB in debug mode.
         connection.query("SET parquet_metadata_cache = true")?;
 
+        // Set vortex_max_threads if specified
+        if let Some(thread_count) = threads {
+            connection.query(&format!("SET vortex_max_threads = {}", thread_count))?;
+        }
+
         Ok((db, connection))
     }
 
@@ -105,7 +125,8 @@ impl DuckDBCtx {
         drop(connection);
         drop(db);
 
-        let (mut db, mut connection) = Self::open_and_setup_database(self.db_path.clone())?;
+        let (mut db, mut connection) =
+            Self::open_and_setup_database(self.db_path.clone(), self.threads)?;
 
         std::mem::swap(&mut self.connection, &mut connection);
         std::mem::swap(&mut self.db, &mut db);
@@ -114,13 +135,12 @@ impl DuckDBCtx {
     }
 
     pub fn new_in_memory() -> Result<Self> {
-        let db = Database::open_in_memory()?;
-        let connection = db.connect()?;
-        vortex_duckdb::register_table_functions(&connection)?;
+        let (db, connection) = Self::open_and_setup_database(None, None)?;
         Ok(Self {
             db,
             connection,
             db_path: None,
+            threads: None,
         })
     }
 

bench-vortex/src/engines/mod.rs

@@ -32,11 +32,13 @@ impl EngineCtx {
         dataset: BenchmarkDataset,
         format: Format,
         delete_duckdb_database: bool,
+        threads: Option<usize>,
     ) -> anyhow::Result<Self> {
         Ok(EngineCtx::DuckDB(ddb::DuckDBCtx::new(
             dataset,
             format,
             delete_duckdb_database,
+            threads,
         )?))
     }
 

vortex-duckdb/cpp/client_context.cpp

@@ -26,4 +26,25 @@ duckdb_vx_client_context_get_object_cache(duckdb_vx_client_context context) {
     } catch (...) {
         return nullptr;
     }
+}
+
+extern "C" duckdb_value duckdb_vx_client_context_try_get_current_setting(duckdb_vx_client_context context,
+                                                                           const char *key) {
+    if (!context || !key) {
+        return nullptr;
+    }
+
+    try {
+        auto client_context = reinterpret_cast<duckdb::ClientContext *>(context);
+        duckdb::Value result;
+        auto lookup_result = client_context->TryGetCurrentSetting(key, result);
+
+        if (lookup_result) {
+            return reinterpret_cast<duckdb_value>(new duckdb::Value(result));
+        }
+
+        return nullptr;
+    } catch (...) {
+        return nullptr;
+    }
 }

vortex-duckdb/cpp/config.cpp

@@ -3,10 +3,12 @@
 
 #include "include/duckdb_vx/config.h"
 #include "duckdb.hpp"
+#include "duckdb/main/config.hpp"
 #include <string>
 #include <memory>
 #include <cstdlib>
 #include <cstring>
+#include <thread>
 
 using namespace duckdb;
 
@@ -111,4 +113,35 @@ char *duckdb_vx_value_to_string(duckdb_value value) {
     }
 }
 
+duckdb_state duckdb_vx_add_extension_option(duckdb_config config,
+                                             const char *name,
+                                             const char *description,
+                                             duckdb_logical_type logical_type,
+                                             duckdb_value default_value) {
+    if (!name || !description || !logical_type || !default_value) {
+        return DuckDBError;
+    }
+
+    try {
+        auto *db_config = reinterpret_cast<DBConfig *>(config);
+        if (!db_config) {
+            return DuckDBError;
+        }
+
+        auto *type = reinterpret_cast<LogicalType *>(logical_type);
+        auto *value = reinterpret_cast<Value *>(default_value);
+
+        db_config->AddExtensionOption(
+            string(name),
+            string(description),
+            *type,
+            *value
+        );
+
+        return DuckDBSuccess;
+    } catch (...) {
+        return DuckDBError;
+    }
+}
+
 } // extern "C"

vortex-duckdb/cpp/include/duckdb_vx/client_context.h

@@ -15,6 +15,11 @@ typedef struct duckdb_vx_client_context_ *duckdb_vx_client_context;
 // This reference is valid as long as the connection is valid.
 duckdb_vx_client_context duckdb_vx_connection_get_client_context(duckdb_connection conn);
 
+// Try to get the current value of a configuration setting.
+// Returns a duckdb_value if the setting exists, or NULL if it doesn't.
+// The caller is responsible for freeing the returned value using duckdb_destroy_value.
+duckdb_value duckdb_vx_client_context_try_get_current_setting(duckdb_vx_client_context context, const char *key);
+
 #ifdef __cplusplus /* End C ABI */
 }
 #endif

vortex-duckdb/cpp/include/duckdb_vx/config.h

@@ -33,6 +33,21 @@ int duckdb_vx_config_has_key(duckdb_config config, const char* key);
 /// @return A newly allocated string containing the value's string representation, or NULL on error
 char* duckdb_vx_value_to_string(duckdb_value value);
 
+/// Add an extension-specific configuration option to DuckDB.
+/// This allows extensions to register custom SET variables.
+///
+/// @param config The DuckDB config instance
+/// @param name The name of the configuration option
+/// @param description A description of what the option does
+/// @param logical_type The DuckDB logical type for the option
+/// @param default_value The default value for the option
+/// @return DuckDBSuccess on success, DuckDBError on error
+duckdb_state duckdb_vx_add_extension_option(duckdb_config config,
+                                             const char *name,
+                                             const char *description,
+                                             duckdb_logical_type logical_type,
+                                             duckdb_value default_value);
+
 #ifdef __cplusplus
 }
 #endif

vortex-duckdb/src/duckdb/client_context.rs

@@ -1,9 +1,11 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::ffi::CStr;
+
 use vortex::error::vortex_panic;
 
-use crate::duckdb::ObjectCacheRef;
+use crate::duckdb::{ObjectCacheRef, Value};
 use crate::{cpp, wrapper};
 
 wrapper!(
@@ -24,4 +26,18 @@ impl ClientContext {
             ObjectCacheRef::borrow(cache)
         }
     }
+
+    /// Try to get the current value of a configuration setting.
+    /// Returns None if the setting doesn't exist.
+    pub fn try_get_current_setting(&self, key: &CStr) -> Option<Value> {
+        unsafe {
+            let value_ptr =
+                cpp::duckdb_vx_client_context_try_get_current_setting(self.as_ptr(), key.as_ptr());
+            if value_ptr.is_null() {
+                None
+            } else {
+                Some(Value::own(value_ptr))
+            }
+        }
+    }
 }

vortex-duckdb/src/duckdb/logical_type.rs

@@ -175,6 +175,10 @@ impl LogicalType {
         Self::new(DUCKDB_TYPE::DUCKDB_TYPE_BIGINT)
     }
 
+    pub fn uint64() -> Self {
+        Self::new(DUCKDB_TYPE::DUCKDB_TYPE_UBIGINT)
+    }
+
     pub fn int32() -> Self {
         Self::new(DUCKDB_TYPE::DUCKDB_TYPE_INTEGER)
     }