feat: per-worker proxy for scalar function execute callback - #1366
Conversation
Wire the scalar execute path to per-worker proxy threads on DuckDB >= 1.5.0. An init callback registered via duckdb_scalar_function_set_init runs once per worker thread, creates a proxy (allocating its Ruby thread under the GVL through the global executor, since init runs on a non-Ruby thread), and stores it as per-worker state via duckdb_scalar_function_init_set_state. The execute callback retrieves that proxy with duckdb_scalar_function_get_state 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. 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 init hook is absent and the execute callback keeps using the global executor unchanged. The added test records which Ruby threads run the 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. Simultaneity assertions are avoided as scheduler-dependent; sample/issue1136.rb demonstrates the throughput win with a GVL-releasing callback (about 3.8x at SET threads=4 locally).
📝 WalkthroughWalkthroughThis PR introduces per-worker proxy initialization and dispatch for DuckDB scalar function callbacks in version 1.5.0+. The C extension captures per-worker proxies during query execution and routes Ruby UDF callbacks through proxy-based dispatchers, with fallback to a shared global executor. Automated tests and a sample script validate multi-threaded callback execution. ChangesPer-Worker Proxy Threading for Scalar UDFs
Sequence DiagramsequenceDiagram
participant DuckDB as DuckDB Query
participant InitCB as init_callback
participant Proxy as worker_proxy
participant ScalarCB as scalar_callback
participant Executor as ruby_executor
Note over DuckDB: Parallel execution on multiple workers
DuckDB->>InitCB: runs on each worker thread
InitCB->>InitCB: detect thread type
alt Non-Ruby worker thread
InitCB->>Proxy: create worker_proxy (rb_protect)
Proxy-->>InitCB: proxy registered
else Ruby calling thread
InitCB->>InitCB: skip proxy setup
end
DuckDB->>ScalarCB: invokes for each row
alt Proxy available
ScalarCB->>Proxy: dispatch_via_proxy
Proxy->>Executor: execute Ruby callback
Executor-->>Proxy: return result
else No proxy (fallback)
ScalarCB->>Executor: dispatch global executor
Executor-->>ScalarCB: return result
end
Proxy-->>ScalarCB: result
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 Infer (1.2.0)ext/duckdb/scalar_function.cIn file included from ext/duckdb/scalar_function.c:1: ... [truncated 746 characters] ... -linux-x86_64-v1.2.0/lib/infer/facebook-clang-plugins/clang/install/lib/clang/18/include" Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
I added the sample code to see whether these changes work.
If we don't need this file, I will remove it and add it as a comment on this PR.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
ext/duckdb/scalar_function.c (1)
26-27: ⚡ Quick winPrefix newly added helper symbols with
rbduckdb_to match extension naming rules.New helpers introduced in this patch use unprefixed C symbol names. Please rename them (and their call sites) to
rbduckdb_*for consistency with the extension’s namespace policy.Proposed rename patch
-static void scalar_function_init_callback(duckdb_init_info info); +static void rbduckdb_scalar_function_init_callback(duckdb_init_info info); -struct create_proxy_callback_arg { +struct rbduckdb_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; +static VALUE rbduckdb_create_proxy_callback(VALUE varg) { + struct rbduckdb_create_proxy_callback_arg *arg = (struct rbduckdb_create_proxy_callback_arg *)varg; arg->proxy = rbduckdb_worker_proxy_create(); return Qnil; } -static void create_proxy_callback_protected(void *user_data) { +static void rbduckdb_create_proxy_callback_protected(void *user_data) { int exception_state; - rb_protect(create_proxy_callback, (VALUE)user_data, &exception_state); + rb_protect(rbduckdb_create_proxy_callback, (VALUE)user_data, &exception_state); if (exception_state) { rb_set_errinfo(Qnil); } } -static void scalar_function_init_callback(duckdb_init_info info) { - struct create_proxy_callback_arg arg; +static void rbduckdb_scalar_function_init_callback(duckdb_init_info info) { + struct rbduckdb_create_proxy_callback_arg arg; if (ruby_native_thread_p()) return; arg.proxy = NULL; - rbduckdb_function_executor_dispatch(create_proxy_callback_protected, &arg); + rbduckdb_function_executor_dispatch(rbduckdb_create_proxy_callback_protected, &arg); if (arg.proxy != NULL) { duckdb_scalar_function_init_set_state(info, arg.proxy, rbduckdb_worker_proxy_destroy); } } - duckdb_scalar_function_set_init(p->scalar_function, scalar_function_init_callback); + duckdb_scalar_function_set_init(p->scalar_function, rbduckdb_scalar_function_init_callback);As per coding guidelines,
ext/duckdb/**/*.c: “C symbols must be prefixed withrbduckdb_to avoid namespace conflicts.”Also applies to: 325-329, 338-359, 382-385
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@ext/duckdb/scalar_function.c` around lines 26 - 27, Rename all newly added helper C symbols in this file to use the rbduckdb_ prefix and update their declarations and call sites accordingly; for example, change scalar_function_init_callback to rbduckdb_scalar_function_init_callback (and any other helpers introduced in this patch) and update every reference, prototype, and registration that uses those names so the file consistently uses rbduckdb_* symbols.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@sample/issue1136.rb`:
- Around line 41-51: The measure method currently opens DuckDB::Database and a
connection (db and con) but only closes them on the success path; wrap the work
in a begin...ensure block so con.close and db.close are always called (guarding
with con && con.closed? checks or con && con.close and similarly for db) to
guarantee resources are released if register_slow_triple or timed_sum raises,
and re-raise any exception after ensure so behavior doesn't change; locate this
logic in the measure method.
---
Nitpick comments:
In `@ext/duckdb/scalar_function.c`:
- Around line 26-27: Rename all newly added helper C symbols in this file to use
the rbduckdb_ prefix and update their declarations and call sites accordingly;
for example, change scalar_function_init_callback to
rbduckdb_scalar_function_init_callback (and any other helpers introduced in this
patch) and update every reference, prototype, and registration that uses those
names so the file consistently uses rbduckdb_* symbols.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: fcc6c17e-a573-402e-86fc-91a42331a995
📒 Files selected for processing (3)
ext/duckdb/scalar_function.csample/issue1136.rbtest/duckdb_test/scalar_function_test.rb
| def measure(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) | ||
| con.close | ||
| db.close | ||
| [elapsed, threads_seen.size, sum] |
There was a problem hiding this comment.
Ensure database resources are closed on exceptions in measure.
On Line 41, connection and database are only closed on success. If register_slow_triple or timed_sum raises, handles remain open until process exit.
Proposed fix
def measure(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)
- con.close
- db.close
- [elapsed, threads_seen.size, sum]
+ begin
+ 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, threads_seen.size, sum]
+ ensure
+ con&.close
+ db&.close
+ end
end📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def measure(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) | |
| con.close | |
| db.close | |
| [elapsed, threads_seen.size, sum] | |
| def measure(threads) | |
| db = DuckDB::Database.open | |
| con = db.connect | |
| begin | |
| 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, threads_seen.size, sum] | |
| ensure | |
| con&.close | |
| db&.close | |
| end | |
| end |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@sample/issue1136.rb` around lines 41 - 51, The measure method currently opens
DuckDB::Database and a connection (db and con) but only closes them on the
success path; wrap the work in a begin...ensure block so con.close and db.close
are always called (guarding with con && con.closed? checks or con && con.close
and similarly for db) to guarantee resources are released if
register_slow_triple or timed_sum raises, and re-raise any exception after
ensure so behavior doesn't change; locate this logic in the measure method.
GitHub: GH-1136
Step 3 of the per-worker proxy work and the first consumer of the proxy: on
DuckDB >= 1.5.0, scalar execute callbacks now run on one proxy thread per
DuckDB worker instead of all funneling through the single global executor.
The table function integration follows in the last PR.
Full picture: otegami#7
Why
With one global executor, callbacks from different workers can never overlap —
even when they release the GVL (e.g. on I/O). Per-worker proxies lift exactly
that ceiling. Measured with
sample/issue1136.rb(GVL-releasing callback overa 500k-row scan):
The before run caps at 2 threads (calling thread + global executor) no matter
how many workers DuckDB spawns. Pure-CPU callbacks stay bounded by the GVL,
so the win is specific to GVL-releasing UDFs.
Notes for review:
duckdb_scalar_function_set_initis a1.5.0 API, so everything is behind
HAVE_DUCKDB_H_GE_V1_5_0.The old path structurally caps at two (calling thread + global executor),
so the test fails without this change. Timing/simultaneity assertions were
deliberately avoided as scheduler-dependent.
Summary by CodeRabbit
New Features
Tests