diff --git a/CHANGELOG.md b/CHANGELOG.md index 29027545..34350697 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,8 @@ All notable changes to this project will be documented in this file. # Unreleased +- add `DuckDB::TableFunction::BindInfo#set_bind_data` to store an arbitrary Ruby object as a custom table function's bind data. +- add `DuckDB::TableFunction::FunctionInfo#get_bind_data` to retrieve, during execution, the object stored by `BindInfo#set_bind_data`. # 1.5.5.0 - 2026-07-27 diff --git a/ext/duckdb/function_executor.c b/ext/duckdb/function_executor.c index b41ad76a..0153eb8c 100644 --- a/ext/duckdb/function_executor.c +++ b/ext/duckdb/function_executor.c @@ -565,3 +565,39 @@ void rbduckdb_function_executor_dispatch_via_proxy(rbduckdb_function_callback_t void rbduckdb_function_executor_dispatch(rbduckdb_function_callback_t cb, void *user_data) { rbduckdb_function_executor_dispatch_via_proxy(cb, user_data, NULL); } + +static VALUE g_function_data_registry = Qnil; +static unsigned long long g_function_data_next_id = 0; + +static void function_data_registry_ensure(void) { + if (g_function_data_registry == Qnil) { + g_function_data_registry = rb_hash_new(); + rb_global_variable(&g_function_data_registry); + } +} + +void *rbduckdb_function_data_register(VALUE value) { + unsigned long long id; + + function_data_registry_ensure(); + id = ++g_function_data_next_id; + rb_hash_aset(g_function_data_registry, ULL2NUM(id), value); + return (void *)(uintptr_t)id; +} + +VALUE rbduckdb_function_data_lookup(void *data) { + if (data == NULL || g_function_data_registry == Qnil) { + return Qnil; + } + return rb_hash_aref(g_function_data_registry, ULL2NUM((unsigned long long)(uintptr_t)data)); +} + +void rbduckdb_function_data_release(void *data) { + if (g_function_data_registry == Qnil) return; + rb_hash_delete(g_function_data_registry, ULL2NUM((unsigned long long)(uintptr_t)data)); +} + +void rbduckdb_function_data_destroy(void *data) { + if (data == NULL) return; + rbduckdb_function_executor_dispatch(rbduckdb_function_data_release, data); +} diff --git a/ext/duckdb/function_executor.h b/ext/duckdb/function_executor.h index ba71c04e..fded625f 100644 --- a/ext/duckdb/function_executor.h +++ b/ext/duckdb/function_executor.h @@ -87,4 +87,9 @@ void rbduckdb_worker_proxy_destroy(void *proxy); */ void rbduckdb_function_executor_dispatch_via_proxy(rbduckdb_function_callback_t cb, void *user_data, struct worker_proxy *proxy); +void *rbduckdb_function_data_register(VALUE value); +VALUE rbduckdb_function_data_lookup(void *data); +void rbduckdb_function_data_release(void *data); +void rbduckdb_function_data_destroy(void *data); + #endif diff --git a/ext/duckdb/table_function_bind_info.c b/ext/duckdb/table_function_bind_info.c index e25aa8b2..bd1ff269 100644 --- a/ext/duckdb/table_function_bind_info.c +++ b/ext/duckdb/table_function_bind_info.c @@ -10,6 +10,7 @@ static VALUE table_function_bind_info_get_parameter(VALUE self, VALUE index); static VALUE table_function_bind_info_get_named_parameter(VALUE self, VALUE name); static VALUE table_function_bind_info__add_result_column(VALUE self, VALUE column_name, VALUE logical_type); static VALUE table_function_bind_info_set_cardinality(VALUE self, VALUE cardinality, VALUE is_exact); +static VALUE table_function_bind_info_set_bind_data(VALUE self, VALUE data); static VALUE table_function_bind_info_set_error(VALUE self, VALUE error); static const rb_data_type_t bind_info_data_type = { @@ -167,6 +168,32 @@ static VALUE table_function_bind_info_set_cardinality(VALUE self, VALUE cardinal return self; } +/* + * call-seq: + * bind_info.set_bind_data(data) -> self + * + * Stores an arbitrary Ruby object as the table function's bind data. The same + * object can be retrieved during init and execution, and is kept alive until + * DuckDB frees the bind data. + * + * bind_info.set_bind_data({ rows: 100 }) + */ +static VALUE table_function_bind_info_set_bind_data(VALUE self, VALUE data) { + rubyDuckDBBindInfo *ctx; + void *handle; + + TypedData_Get_Struct(self, rubyDuckDBBindInfo, &bind_info_data_type, ctx); + + if (ctx->bind_data_handle != NULL) { + rbduckdb_function_data_release(ctx->bind_data_handle); + } + handle = rbduckdb_function_data_register(data); + ctx->bind_data_handle = handle; + duckdb_bind_set_bind_data(ctx->bind_info, handle, rbduckdb_function_data_destroy); + + return self; +} + /* * call-seq: * bind_info.set_error(error_message) -> self @@ -198,6 +225,7 @@ void rbduckdb_init_table_function_bind_info(void) { rb_define_method(cDuckDBTableFunctionBindInfo, "get_parameter", table_function_bind_info_get_parameter, 1); rb_define_method(cDuckDBTableFunctionBindInfo, "get_named_parameter", table_function_bind_info_get_named_parameter, 1); rb_define_method(cDuckDBTableFunctionBindInfo, "set_cardinality", table_function_bind_info_set_cardinality, 2); + rb_define_method(cDuckDBTableFunctionBindInfo, "set_bind_data", table_function_bind_info_set_bind_data, 1); rb_define_method(cDuckDBTableFunctionBindInfo, "set_error", table_function_bind_info_set_error, 1); rb_define_private_method(cDuckDBTableFunctionBindInfo, "_add_result_column", table_function_bind_info__add_result_column, 2); diff --git a/ext/duckdb/table_function_bind_info.h b/ext/duckdb/table_function_bind_info.h index 47f18eb7..e362823e 100644 --- a/ext/duckdb/table_function_bind_info.h +++ b/ext/duckdb/table_function_bind_info.h @@ -3,6 +3,7 @@ struct _rubyDuckDBBindInfo { duckdb_bind_info bind_info; + void *bind_data_handle; }; typedef struct _rubyDuckDBBindInfo rubyDuckDBBindInfo; diff --git a/ext/duckdb/table_function_function_info.c b/ext/duckdb/table_function_function_info.c index ef4a5113..4ebcd416 100644 --- a/ext/duckdb/table_function_function_info.c +++ b/ext/duckdb/table_function_function_info.c @@ -5,6 +5,7 @@ VALUE cDuckDBTableFunctionFunctionInfo; static void deallocate(void *ctx); static VALUE allocate(VALUE klass); static size_t memsize(const void *p); +static VALUE table_function_function_info_get_bind_data(VALUE self); static VALUE table_function_function_info_set_error(VALUE self, VALUE error); static const rb_data_type_t function_info_data_type = { @@ -33,6 +34,26 @@ rubyDuckDBFunctionInfo *rbduckdb_get_struct_function_info(VALUE obj) { return ctx; } +/* + * call-seq: + * function_info.get_bind_data -> object or nil + * + * Returns the object stored during the bind phase with + * DuckDB::TableFunction::BindInfo#set_bind_data, or nil if none was set. + * + * data = function_info.get_bind_data + */ +static VALUE table_function_function_info_get_bind_data(VALUE self) { + rubyDuckDBFunctionInfo *ctx; + void *bind_data; + + TypedData_Get_Struct(self, rubyDuckDBFunctionInfo, &function_info_data_type, ctx); + + bind_data = duckdb_function_get_bind_data(ctx->info); + + return rbduckdb_function_data_lookup(bind_data); +} + /* * call-seq: * function_info.set_error(error_message) -> self @@ -61,5 +82,6 @@ void rbduckdb_init_table_function_function_info(void) { cDuckDBTableFunctionFunctionInfo = rb_define_class_under(cDuckDBTableFunction, "FunctionInfo", rb_cObject); rb_define_alloc_func(cDuckDBTableFunctionFunctionInfo, allocate); + rb_define_method(cDuckDBTableFunctionFunctionInfo, "get_bind_data", table_function_function_info_get_bind_data, 0); rb_define_method(cDuckDBTableFunctionFunctionInfo, "set_error", table_function_function_info_set_error, 1); } diff --git a/test/duckdb_test/table_function/function_info_test.rb b/test/duckdb_test/table_function/function_info_test.rb index 4d49d554..73f3aac5 100644 --- a/test/duckdb_test/table_function/function_info_test.rb +++ b/test/duckdb_test/table_function/function_info_test.rb @@ -43,6 +43,104 @@ def test_execute_callback assert_equal table_function, result2 end + def test_bind_data_round_trip + skip 'GC.compact hangs on Windows in parallel test execution' if Gem.win_platform? + + table_function = DuckDB::TableFunction.new + table_function.name = 'test_bind_data' + + set_result = nil + table_function.bind do |bind_info| + bind_info.add_result_column('value', DuckDB::LogicalType::BIGINT) + set_result = bind_info.set_bind_data({ token: 'round-trip', n: 7 }) + end + + table_function.init { |_init_info| GC.compact } + + observed_bind_data = nil + done = false + table_function.execute do |func_info, output| + if done + output.size = 0 + else + observed_bind_data = func_info.get_bind_data + output.set_value(0, 0, 1) + output.size = 1 + done = true + end + end + + @connection.register_table_function(table_function) + @connection.query('SELECT * FROM test_bind_data()').each.to_a + + assert_instance_of DuckDB::TableFunction::BindInfo, set_result + assert_equal({ token: 'round-trip', n: 7 }, observed_bind_data) + end + + def test_bind_data_last_set_wins + skip 'GC.compact hangs on Windows in parallel test execution' if Gem.win_platform? + + table_function = DuckDB::TableFunction.new + table_function.name = 'test_bind_data_overwrite' + + table_function.bind do |bind_info| + bind_info.add_result_column('value', DuckDB::LogicalType::BIGINT) + bind_info.set_bind_data({ which: 'first' }) + bind_info.set_bind_data({ which: 'second' }) + end + + table_function.init { |_init_info| GC.compact } + + observed_bind_data = nil + done = false + table_function.execute do |func_info, output| + if done + output.size = 0 + else + observed_bind_data = func_info.get_bind_data + output.set_value(0, 0, 1) + output.size = 1 + done = true + end + end + + @connection.register_table_function(table_function) + @connection.query('SELECT * FROM test_bind_data_overwrite()').each.to_a + + assert_equal({ which: 'second' }, observed_bind_data) + end + + def test_bind_data_nil_when_unset + skip 'GC.compact hangs on Windows in parallel test execution' if Gem.win_platform? + + table_function = DuckDB::TableFunction.new + table_function.name = 'test_bind_data_unset' + + table_function.bind do |bind_info| + bind_info.add_result_column('value', DuckDB::LogicalType::BIGINT) + end + + table_function.init { |_init_info| GC.compact } + + observed_bind_data = :unset + done = false + table_function.execute do |func_info, output| + if done + output.size = 0 + else + observed_bind_data = func_info.get_bind_data + output.set_value(0, 0, 1) + output.size = 1 + done = true + end + end + + @connection.register_table_function(table_function) + @connection.query('SELECT * FROM test_bind_data_unset()').each.to_a + + assert_nil observed_bind_data + end + def test_execute_without_block table_function = DuckDB::TableFunction.new