From 6b0fd95bd5941bc0885cdad10ab7df23b1b1f9db Mon Sep 17 00:00:00 2001 From: otegami Date: Sat, 30 May 2026 11:53:12 +0800 Subject: [PATCH] feat: per-worker proxy for table function execute callback Wire the table execute path to per-worker proxy threads on DuckDB >= 1.5.0. A local_init callback registered via duckdb_table_function_set_local_init runs once per worker thread, creates a proxy (allocating its Ruby thread under the GVL through the global executor, since local_init runs on a non-Ruby thread), and stores it as thread-local init data via duckdb_init_set_init_data. The execute callback retrieves that proxy with duckdb_function_get_local_init_data and dispatches through it via rbduckdb_function_executor_dispatch_via_proxy, so callbacks from different workers run concurrently instead of serializing on the single global executor. bind and init stay on the global executor. DuckDB frees each proxy through rbduckdb_worker_proxy_destroy. The proxy-creating wrapper runs rbduckdb_worker_proxy_create under rb_protect, implementing the raise contract documented on that function: the executor runs callbacks unprotected, so an uncaught raise would longjmp past its done-signaling and block the waiting DuckDB worker forever. On failure the proxy stays NULL and the execute callback falls back to the global executor. On DuckDB < 1.5.0 the local_init hook is absent and the execute callback keeps using the global executor unchanged. The added test records which Ruby threads run the execute callback and asserts more than two distinct threads, which the old implementation can never produce (calling thread plus the single global executor), in addition to result correctness. Verified to fail against a build without this change. Simultaneity assertions are avoided as scheduler-dependent. sample/issue1136.rb gains a table UDF section: a GVL-releasing emitter with both planner hints (set_cardinality, max_threads) that demonstrates the throughput win (locally about 4x at SET threads=4, execute callbacks on 4 distinct Ruby threads instead of 1). --- ext/duckdb/table_function.c | 68 ++++++++++++++++++- sample/issue1136.rb | 88 +++++++++++++++++++++---- test/duckdb_test/table_function_test.rb | 65 ++++++++++++++++++ 3 files changed, 208 insertions(+), 13 deletions(-) diff --git a/ext/duckdb/table_function.c b/ext/duckdb/table_function.c index c5273a15..0f23c494 100644 --- a/ext/duckdb/table_function.c +++ b/ext/duckdb/table_function.c @@ -21,6 +21,11 @@ static VALUE rbduckdb_table_function_set_init(VALUE self); static void table_function_init_callback(duckdb_init_info info); static VALUE rbduckdb_table_function_set_execute(VALUE self); static void table_function_execute_callback(duckdb_function_info info, duckdb_data_chunk output); +#ifdef HAVE_DUCKDB_H_GE_V1_5_0 +/* Thread detection (declared in function_executor.c); used to skip the proxy on Ruby threads. */ +extern int ruby_native_thread_p(void); +static void table_function_local_init_callback(duckdb_init_info info); +#endif static const rb_data_type_t table_function_data_type = { "DuckDB/TableFunction", @@ -358,6 +363,10 @@ static VALUE rbduckdb_table_function_set_execute(VALUE self) { ctx->execute_proc = rb_block_proc(); duckdb_table_function_set_function(ctx->table_function, table_function_execute_callback); +#ifdef HAVE_DUCKDB_H_GE_V1_5_0 + /* Per-worker proxy threads for the execute path (DuckDB >= 1.5.0). */ + duckdb_table_function_set_local_init(ctx->table_function, table_function_local_init_callback); +#endif rbduckdb_function_executor_ensure_started(); @@ -405,6 +414,7 @@ static void execute_execute_callback_protected(void *user_data) { static void table_function_execute_callback(duckdb_function_info info, duckdb_data_chunk output) { rubyDuckDBTableFunction *ctx; struct execute_dispatch_arg darg; + struct worker_proxy *proxy = NULL; ctx = (rubyDuckDBTableFunction *)duckdb_function_get_extra_info(info); if (!ctx || ctx->execute_proc == Qnil) return; @@ -413,8 +423,64 @@ static void table_function_execute_callback(duckdb_function_info info, duckdb_da darg.info = info; darg.output = output; - rbduckdb_function_executor_dispatch(execute_execute_callback_protected, &darg); +#ifdef HAVE_DUCKDB_H_GE_V1_5_0 + /* On DuckDB >= 1.5.0 each worker thread carries its own proxy (see local_init). */ + proxy = (struct worker_proxy *)duckdb_function_get_local_init_data(info); +#endif + rbduckdb_function_executor_dispatch_via_proxy(execute_execute_callback_protected, &darg, proxy); +} + +#ifdef HAVE_DUCKDB_H_GE_V1_5_0 +/* + * Per-worker init for the execute path (DuckDB >= 1.5.0). + * + * DuckDB calls this once on each worker thread that will run the execute + * callback. We create a per-worker proxy (allocating its Ruby thread under the + * GVL via the global executor, since this runs on a non-Ruby thread) and store + * it as thread-local init data. The execute callback then dispatches through it + * instead of the shared global executor, so workers run callbacks concurrently. + * DuckDB invokes rbduckdb_worker_proxy_destroy when the local state is freed. + */ +struct create_proxy_callback_arg { + struct worker_proxy *proxy; +}; + +static VALUE create_proxy_callback(VALUE varg) { + struct create_proxy_callback_arg *arg = (struct create_proxy_callback_arg *)varg; + arg->proxy = rbduckdb_worker_proxy_create(); + return Qnil; +} + +/* + * rbduckdb_worker_proxy_create may raise (NoMemError, Thread.new failure), + * and the executor runs callbacks unprotected — a raise would longjmp past + * its done-signaling and block the waiting DuckDB worker forever. Swallow + * the exception instead: the proxy stays NULL, local_init sets no state, and + * the execute callback falls back to the global executor. + */ +static void create_proxy_callback_protected(void *user_data) { + int exception_state; + + rb_protect(create_proxy_callback, (VALUE)user_data, &exception_state); + if (exception_state) { + rb_set_errinfo(Qnil); + } +} + +static void table_function_local_init_callback(duckdb_init_info info) { + struct create_proxy_callback_arg arg; + + /* A Ruby calling thread runs the callback inline (Case 1/2); no proxy needed. */ + if (ruby_native_thread_p()) return; + + arg.proxy = NULL; + rbduckdb_function_executor_dispatch(create_proxy_callback_protected, &arg); + + if (arg.proxy != NULL) { + duckdb_init_set_init_data(info, arg.proxy, rbduckdb_worker_proxy_destroy); + } } +#endif rubyDuckDBTableFunction *get_struct_table_function(VALUE self) { rubyDuckDBTableFunction *ctx; diff --git a/sample/issue1136.rb b/sample/issue1136.rb index 481ef8e4..39bec375 100644 --- a/sample/issue1136.rb +++ b/sample/issue1136.rb @@ -1,6 +1,7 @@ # frozen_string_literal: true -# GH-1136: per-worker proxy threads for scalar UDF callbacks (DuckDB >= 1.5.0). +# GH-1136: per-worker proxy threads for scalar and table UDF callbacks +# (DuckDB >= 1.5.0). # # Each DuckDB worker thread gets its own dedicated Ruby proxy thread, so UDF # callbacks from different workers run concurrently instead of serializing on @@ -11,10 +12,14 @@ # # Note: pure-CPU Ruby callbacks stay effectively serialized by the GVL; the # throughput win is specific to callbacks that release it (e.g. on I/O, -# simulated here with sleep). A large base table is used so the morsel-driven -# scan actually distributes work across workers. +# simulated here with sleep). The scalar section uses a large base table so +# the morsel-driven scan actually distributes work across workers; the table +# section sets both planner hints (set_cardinality + max_threads) for the +# same reason. require 'duckdb' +# --- scalar UDF ------------------------------------------------------------- + ROWS = 500_000 SLEEP_EVERY = 1_000 # simulate I/O on every Nth value SLEEP_SEC = 0.002 @@ -32,29 +37,88 @@ def register_slow_triple(con, threads_seen) con.register_scalar_function(sf) end -def timed_sum(con) +def timed_sum(con, query) started = Process.clock_gettime(Process::CLOCK_MONOTONIC) - sum = con.execute('SELECT SUM(slow_triple(value)) FROM t').first.first + sum = con.execute(query).first.first [Process.clock_gettime(Process::CLOCK_MONOTONIC) - started, sum] end -def measure(threads) +def measure_scalar(threads) db = DuckDB::Database.open con = db.connect con.execute("SET threads=#{threads}") con.execute("CREATE TABLE t AS SELECT range::INTEGER AS value FROM range(#{ROWS})") threads_seen = {} register_slow_triple(con, threads_seen) - elapsed, sum = timed_sum(con) + elapsed, sum = timed_sum(con, 'SELECT SUM(slow_triple(value)) FROM t') + con.close + db.close + [elapsed, threads_seen.size, sum] +end + +# --- table UDF -------------------------------------------------------------- + +CHUNKS = 200 +ROWS_PER_CHUNK = 50 +CHUNK_SLEEP_SEC = 0.005 # simulate I/O per emitted chunk + +def emit_chunk(output) + ROWS_PER_CHUNK.times { |i| output.set_value(0, i, 1) } + output.size = ROWS_PER_CHUNK + sleep(CHUNK_SLEEP_SEC) +end + +def emitter_bind_block + proc do |bind_info| + bind_info.add_result_column('v', DuckDB::LogicalType::BIGINT) + # Tell the planner there is real work so it distributes across workers. + bind_info.set_cardinality(CHUNKS * ROWS_PER_CHUNK, false) + end +end + +def emitter_execute_block(threads_seen) + remaining = CHUNKS + mutex = Mutex.new + proc do |_info, output| + threads_seen[Thread.current] = true + has_work = mutex.synchronize { (remaining -= 1) >= 0 } + has_work ? emit_chunk(output) : output.size = 0 + end +end + +def register_slow_emitter(con, threads_seen) + tf = DuckDB::TableFunction.new + tf.name = 'slow_emitter' + tf.bind(&emitter_bind_block) + # Without max_threads DuckDB assigns a single worker and the proxy never fires. + tf.init { |init_info| init_info.max_threads = 4 } + tf.execute(&emitter_execute_block(threads_seen)) + con.register_table_function(tf) +end + +def measure_table(threads) + db = DuckDB::Database.open + con = db.connect + con.execute("SET threads=#{threads}") + threads_seen = {} + register_slow_emitter(con, threads_seen) + elapsed, sum = timed_sum(con, 'SELECT SUM(v) FROM slow_emitter()') con.close db.close [elapsed, threads_seen.size, sum] end -expected = (ROWS - 1) * ROWS / 2 * 3 -[1, 4].each do |threads| - elapsed, distinct, sum = measure(threads) - raise "wrong result: #{sum} (expected #{expected})" unless sum == expected +# --- run both --------------------------------------------------------------- - puts "SET threads=#{threads}: #{elapsed.round(3)}s, callbacks ran on #{distinct} distinct Ruby thread(s)" +def report(label, expected) + puts "#{label}:" + [1, 4].each do |threads| + elapsed, distinct, sum = yield(threads) + raise "wrong result: #{sum} (expected #{expected})" unless sum == expected + + puts " SET threads=#{threads}: #{elapsed.round(3)}s, callbacks ran on #{distinct} distinct Ruby thread(s)" + end end + +report('scalar UDF (slow_triple)', (ROWS - 1) * ROWS / 2 * 3) { |threads| measure_scalar(threads) } +report('table UDF (slow_emitter)', CHUNKS * ROWS_PER_CHUNK) { |threads| measure_table(threads) } diff --git a/test/duckdb_test/table_function_test.rb b/test/duckdb_test/table_function_test.rb index 102b542d..f1352774 100644 --- a/test/duckdb_test/table_function_test.rb +++ b/test/duckdb_test/table_function_test.rb @@ -199,6 +199,71 @@ def test_symbol_columns db.close end + # Per-worker proxy: exercises the local_init -> proxy -> destroy lifecycle + # under real multi-threaded execution (SET threads=4) and asserts the + # proxy path actually fired: execute callbacks must run on more than two + # distinct Ruby threads. Without per-worker proxies that count can never + # exceed two (the calling thread plus the single global executor), so this + # fails on the old implementation. Simultaneity assertions (max + # concurrency) are avoided as scheduler-dependent; sample/issue1136.rb + # demonstrates the throughput win for the scalar twin of this mechanism. + # Requires DuckDB >= 1.5.0 (duckdb_table_function_set_local_init). + def test_execute_runs_on_per_worker_proxy_threads + if ::DuckDBTest.duckdb_library_version < Gem::Version.new('1.5.0') + skip 'per-worker proxy requires DuckDB >= 1.5.0' + end + + chunks = 64 + rows_per_chunk = 100 + remaining = chunks + mutex = Mutex.new + threads_seen = {} + + db = DuckDB::Database.open + conn = db.connect + conn.execute('SET threads=4') + + tf = DuckDB::TableFunction.new + tf.name = 'parallel_emitter' + tf.bind do |bind_info| + bind_info.add_result_column('v', DuckDB::LogicalType::BIGINT) + # Tell the planner there is real work so it distributes across workers. + bind_info.set_cardinality(chunks * rows_per_chunk, false) + end + tf.init do |init_info| + # Without this DuckDB assigns a single worker and the proxy never fires. + init_info.max_threads = 4 + end + tf.execute do |_info, output| + threads_seen[Thread.current] = true + has_work = mutex.synchronize do + next false if remaining.zero? + + remaining -= 1 + true + end + + unless has_work + output.size = 0 + next + end + + rows_per_chunk.times { |i| output.set_value(0, i, 1) } + output.size = rows_per_chunk + sleep 0.001 # release the GVL so workers can overlap + end + + conn.register_table_function(tf) + result = conn.query('SELECT COUNT(*), SUM(v) FROM parallel_emitter()').each.to_a + + assert_equal [chunks * rows_per_chunk, chunks * rows_per_chunk], result.first + assert_operator threads_seen.size, :>, 2, + 'expected callbacks on per-worker proxy threads, not just caller + global executor' + + conn.disconnect + db.close + end + private def setup_incomplete_function