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

joseph-isaacs · web-flow · commit ccc56ebae68c · 2025-10-21T14:30:44.000+01:00
Set max threads on the rust side

---------

Signed-off-by: Joe Isaacs &lt;joe.isaacs@live.co.uk&gt;
diff --git a/vortex-duckdb/cpp/config.cpp b/vortex-duckdb/cpp/config.cpp
@@ -12,28 +12,28 @@ using namespace duckdb;
 
 extern "C" {
 
-duckdb_state duckdb_vx_get_config_value(duckdb_config config, const char* key, duckdb_value* out_value) {
+duckdb_state duckdb_vx_get_config_value(duckdb_config config, const char *key, duckdb_value *out_value) {
     if (!config || !key || !out_value) {
         return DuckDBError;
     }
 
     try {
         // Cast the config to DuckDB's internal config type
-        auto* db_config = reinterpret_cast<DBConfig*>(config);
-        
+        auto *db_config = reinterpret_cast<DBConfig *>(config);
+
         if (!db_config) {
             return DuckDBError;
         }
 
         std::string key_str(key);
-        
+
         // First check set_variables (the primary location for config values)
         auto set_it = db_config->options.set_variables.find(key_str);
         if (set_it != db_config->options.set_variables.end()) {
             *out_value = reinterpret_cast<duckdb_value>(new Value(set_it->second));
             return DuckDBSuccess;
         }
-        
+
         // Then check user_options
         auto user_it = db_config->options.user_options.find(key_str);
         if (user_it != db_config->options.user_options.end()) {
@@ -43,69 +43,69 @@ duckdb_state duckdb_vx_get_config_value(duckdb_config config, const char* key, d
 
         // Key not found in any map
         return DuckDBError;
-        
+
     } catch (...) {
         return DuckDBError;
     }
 }
 
-int duckdb_vx_config_has_key(duckdb_config config, const char* key) {
+int duckdb_vx_config_has_key(duckdb_config config, const char *key) {
     if (!config || !key) {
         return 0;
     }
 
     try {
-        auto* db_config = reinterpret_cast<DBConfig*>(config);
+        auto *db_config = reinterpret_cast<DBConfig *>(config);
         if (!db_config) {
             return 0;
         }
-        
+
         std::string key_str(key);
-        
+
         // Check if the key exists in set_variables (primary location)
         if (db_config->options.set_variables.find(key_str) != db_config->options.set_variables.end()) {
             return 1;
         }
-        
+
         // Check if the key exists in user_options
         if (db_config->options.user_options.find(key_str) != db_config->options.user_options.end()) {
             return 1;
         }
 
         return 0;
-        
+
     } catch (...) {
         return 0;
     }
 }
 
-char* duckdb_vx_value_to_string(duckdb_value value) {
+char *duckdb_vx_value_to_string(duckdb_value value) {
     if (!value) {
         return nullptr;
     }
 
     try {
         // Cast the value to DuckDB's internal Value type
-        auto* ddb_value = reinterpret_cast<Value*>(value);
-        
+        auto *ddb_value = reinterpret_cast<Value *>(value);
+
         if (!ddb_value) {
             return nullptr;
         }
 
         // Use the ToString method to get the string representation
         std::string str_value = ddb_value->ToString();
-        
+
         // Allocate memory for the C string using malloc (compatible with duckdb_free)
         size_t str_len = str_value.length() + 1;
-        char* result = static_cast<char*>(malloc(str_len));
+        char *result = static_cast<char *>(malloc(str_len));
         if (!result) {
             return nullptr;
         }
-        
+
         // Copy the string and null terminate
         std::memcpy(result, str_value.c_str(), str_len);
         return result;
-        
+
     } catch (...) {
         return nullptr;
     }
diff --git a/vortex-duckdb/cpp/include/duckdb_vx/table_function.h b/vortex-duckdb/cpp/include/duckdb_vx/table_function.h
@@ -140,6 +140,7 @@ typedef struct {
     bool filter_prune;
     bool sampling_pushdown;
     bool late_materialization;
+    idx_t max_threads;
 } duckdb_vx_tfunc_vtab_t;
 
 // A single function for configuring the DuckDB table function vtable.
diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp
@@ -16,10 +16,12 @@ using namespace duckdb;
 
 namespace vortex {
 struct CTableFunctionInfo final : TableFunctionInfo {
-    explicit CTableFunctionInfo(const duckdb_vx_tfunc_vtab_t &vtab) : vtab(vtab) {
+    explicit CTableFunctionInfo(const duckdb_vx_tfunc_vtab_t &vtab)
+        : vtab(vtab), max_threads(vtab.max_threads) {
     }
 
     duckdb_vx_tfunc_vtab_t vtab;
+    idx_t max_threads;
 };
 
 struct CTableBindData final : TableFunctionData {
@@ -44,13 +46,15 @@ struct CTableBindData final : TableFunctionData {
 };
 
 struct CTableGlobalData final : GlobalTableFunctionState {
-    explicit CTableGlobalData(unique_ptr<vortex::CData> ffi_data_p) : ffi_data(std::move(ffi_data_p)) {
+    explicit CTableGlobalData(unique_ptr<vortex::CData> ffi_data_p, idx_t max_threads_p)
+        : ffi_data(std::move(ffi_data_p)), max_threads(max_threads_p) {
     }
 
     unique_ptr<vortex::CData> ffi_data;
+    idx_t max_threads;
 
     idx_t MaxThreads() const override {
-        return GlobalTableFunctionState::MAX_THREADS;
+        return max_threads;
     }
 };
 
@@ -111,7 +115,8 @@ unique_ptr<GlobalTableFunctionState> c_init_global(ClientContext &context, Table
     }
 
     return make_uniq<CTableGlobalData>(
-        unique_ptr<vortex::CData>(reinterpret_cast<vortex::CData *>(ffi_global_data)));
+        unique_ptr<vortex::CData>(reinterpret_cast<vortex::CData *>(ffi_global_data)),
+        bind.info->max_threads);
 }
 
 unique_ptr<LocalTableFunctionState> c_init_local(ExecutionContext &context, TableFunctionInitInput &input,
@@ -148,8 +153,8 @@ void c_function(ClientContext &context, TableFunctionInput &input, DataChunk &ou
     auto local_data = input.local_state->Cast<CTableLocalData>().ffi_data->DataPtr();
 
     duckdb_vx_error error_out = nullptr;
-    bind.info->vtab.function(ctx, bind_data, global_data, local_data, reinterpret_cast<duckdb_data_chunk>(&output),
-                             &error_out);
+    bind.info->vtab.function(ctx, bind_data, global_data, local_data,
+                             reinterpret_cast<duckdb_data_chunk>(&output), &error_out);
     if (error_out) {
         throw InvalidInputException(IntoErrString(error_out));
     }
@@ -255,8 +260,8 @@ virtual_column_map_t c_get_virtual_columns(ClientContext &context, optional_ptr<
 }
 
 extern "C" void duckdb_vx_tfunc_virtual_columns_push(duckdb_vx_tfunc_virtual_cols_result ffi_result,
-                                                  idx_t column_idx, const char *name_str, size_t name_len,
-                                                  duckdb_logical_type ffi_type) {
+                                                     idx_t column_idx, const char *name_str, size_t name_len,
+                                                     duckdb_logical_type ffi_type) {
     if (!ffi_result || !name_str || !ffi_type) {
         return;
     }
@@ -289,7 +294,7 @@ OperatorPartitionData c_get_partition_data(ClientContext &context, TableFunction
 InsertionOrderPreservingMap<string> c_to_string(TableFunctionToStringInput &input) {
     InsertionOrderPreservingMap<string> result;
     auto &bind = input.bind_data->Cast<CTableBindData>();
-    
+
     // Call the Rust side to get custom string representation if available
     if (bind.info->vtab.to_string) {
         auto map = bind.info->vtab.to_string(bind.ffi_data->DataPtr());
@@ -303,7 +308,7 @@ InsertionOrderPreservingMap<string> c_to_string(TableFunctionToStringInput &inpu
             duckdb_vx_string_map_free(map);
         }
     }
-    
+
     return result;
 }
 
diff --git a/vortex-duckdb/cpp/vector.cpp b/vortex-duckdb/cpp/vector.cpp
@@ -47,24 +47,26 @@ namespace vortex {
 // This is a complete hack to access the data buffer and pointer of a vector.
 class DataVector : public Vector {
 public:
-	inline void SetDataBuffer(buffer_ptr<VectorBuffer> new_buffer) {
-		buffer = std::move(new_buffer);
-	};
+    inline void SetDataBuffer(buffer_ptr<VectorBuffer> new_buffer) {
+        buffer = std::move(new_buffer);
+    };
 
-	inline void SetDataPtr(data_ptr_t ptr) {
-		data = ptr;
-	};
+    inline void SetDataPtr(data_ptr_t ptr) {
+        data = ptr;
+    };
 };
 
 } // namespace vortex
 
-extern "C" void duckdb_vx_string_vector_add_vector_data_buffer(duckdb_vector ffi_vector, duckdb_vx_vector_buffer buffer) {
+extern "C" void duckdb_vx_string_vector_add_vector_data_buffer(duckdb_vector ffi_vector,
+                                                               duckdb_vx_vector_buffer buffer) {
     auto vector = reinterpret_cast<Vector *>(ffi_vector);
     auto data = reinterpret_cast<shared_ptr<vortex::ExternalVectorBuffer> *>(buffer);
     StringVector::AddBuffer(*vector, *data);
 }
 
-extern "C" void duckdb_vx_vector_set_vector_data_buffer(duckdb_vector ffi_vector, duckdb_vx_vector_buffer buffer) {
+extern "C" void duckdb_vx_vector_set_vector_data_buffer(duckdb_vector ffi_vector,
+                                                        duckdb_vx_vector_buffer buffer) {
     auto vector = reinterpret_cast<Vector *>(ffi_vector);
     auto dvector = reinterpret_cast<vortex::DataVector *>(vector);
     auto data = reinterpret_cast<shared_ptr<vortex::ExternalVectorBuffer> *>(buffer);
diff --git a/vortex-duckdb/cpp/vector_buffer.cpp b/vortex-duckdb/cpp/vector_buffer.cpp
@@ -12,9 +12,8 @@ using namespace duckdb;
 
 extern "C" duckdb_vx_vector_buffer duckdb_vx_vector_buffer_create(duckdb_vx_data buffer) {
     auto data = reinterpret_cast<vortex::CData *>(buffer);
-    auto* shared_buffer = new duckdb::shared_ptr<vortex::ExternalVectorBuffer>(
-        duckdb::make_shared_ptr<vortex::ExternalVectorBuffer>(unique_ptr<vortex::CData>(data))
-    );
+    auto *shared_buffer = new duckdb::shared_ptr<vortex::ExternalVectorBuffer>(
+        duckdb::make_shared_ptr<vortex::ExternalVectorBuffer>(unique_ptr<vortex::CData>(data)));
     return reinterpret_cast<duckdb_vx_vector_buffer>(shared_buffer);
 }
 
diff --git a/vortex-duckdb/src/duckdb/table_function/mod.rs b/vortex-duckdb/src/duckdb/table_function/mod.rs
@@ -51,6 +51,10 @@ pub trait TableFunction: Sized + Debug {
     ///   - j does not need to leave the table function at all.
     const FILTER_PRUNE: bool = false;
 
+    /// Maximum number of threads the table function can use.
+    /// If not specified, DuckDB will use its default (GlobalTableFunctionState::MAX_THREADS).
+    const MAX_THREADS: u64 = u64::MAX;
+
     /// Returns the parameters of the table function.
     fn parameters() -> Vec<LogicalType> {
         // By default, we don't have any parameters.
@@ -180,6 +184,7 @@ impl Connection {
             filter_prune: T::FILTER_PRUNE,
             sampling_pushdown: false,
             late_materialization: false,
+            max_threads: T::MAX_THREADS,
         };
 
         duckdb_try!(
diff --git a/vortex-duckdb/src/scan.rs b/vortex-duckdb/src/scan.rs
@@ -40,6 +40,7 @@ pub struct VortexBindData {
     column_names: Vec<String>,
     column_types: Vec<LogicalType>,
     runtime: CurrentThreadRuntime,
+    max_threads: u64,
 }
 
 impl Clone for VortexBindData {
@@ -53,6 +54,7 @@ impl Clone for VortexBindData {
             column_names: self.column_names.clone(),
             column_types: self.column_types.clone(),
             runtime: self.runtime.clone(),
+            max_threads: self.max_threads,
         }
     }
 }
@@ -258,6 +260,7 @@ impl TableFunction for VortexTableFunction {
             column_names,
             column_types,
             runtime,
+            max_threads: u64::MAX,
         })
     }
 

Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,7 @@ pub struct VortexBindData {`
`40`	`40`	`column_names: Vec<String>,`
`41`	`41`	`column_types: Vec<LogicalType>,`
`42`	`42`	`runtime: CurrentThreadRuntime,`
	`43`	`+ max_threads: u64,`
`43`	`44`	`}`
`44`	`45`
`45`	`46`	`impl Clone for VortexBindData {`
`@@ -53,6 +54,7 @@ impl Clone for VortexBindData {`
`53`	`54`	`column_names: self.column_names.clone(),`
`54`	`55`	`column_types: self.column_types.clone(),`
`55`	`56`	`runtime: self.runtime.clone(),`
	`57`	`+ max_threads: self.max_threads,`
`56`	`58`	`}`
`57`	`59`	`}`
`58`	`60`	`}`
`@@ -258,6 +260,7 @@ impl TableFunction for VortexTableFunction {`
`258`	`260`	`column_names,`
`259`	`261`	`column_types,`
`260`	`262`	`runtime,`
	`263`	`+ max_threads: u64::MAX,`
`261`	`264`	`})`
`262`	`265`	`}`
`263`	`266`