Cache prepared statement plans for fast path queries on worker connections - #8825
Open
Colm (colm-mchugh) wants to merge 2 commits into
Open
Cache prepared statement plans for fast path queries on worker connections#8825Colm (colm-mchugh) wants to merge 2 commits into
Colm (colm-mchugh) wants to merge 2 commits into
Conversation
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #8825 +/- ##
==========================================
- Coverage 88.73% 88.70% -0.04%
==========================================
Files 289 290 +1
Lines 65012 65316 +304
Branches 8203 8257 +54
==========================================
+ Hits 57690 57937 +247
- Misses 4952 4993 +41
- Partials 2370 2386 +16 🚀 New features to boost your workflow:
|
Colm (colm-mchugh)
force-pushed
the
colm/prepared-stmts-poc
branch
from
September 4, 2026 15:42
f6ba6a2 to
598557f
Compare
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
It contains confirmed correctness issues (use-after-free in lazy deparse caching, UPSERT alias handling gaps, and modify-path caching semantics when coordinator evaluation is required).
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR introduces worker-connection prepared-statement caching for single-shard fast-path executions, gated by a new citus.enable_prepared_statement_caching GUC, aiming to avoid repeated parse/plan overhead on workers and repeated deparse work on the coordinator.
Changes:
- Add a per-
MultiConnectionprepared-statement cache keyed by(planId, shardId), plus executor send-path support usingPQprepare+PQsendQueryPrepared. - Add fast-path execution support that can build minimal per-execution
Tasks directly fromParamListInfofor eligible router queries. - Add regression coverage and scheduling for prepared-statement caching behavior.
File summaries
| File | Description |
|---|---|
| src/test/regress/sql/prepared_statement_caching.sql | New regression test SQL for prepared-statement caching scenarios |
| src/test/regress/expected/prepared_statement_caching.out | Expected output for the new regression test |
| src/test/regress/multi_schedule | Adds prepared_statement_caching to the regression schedule |
| src/test/regress/citus_tests/run_test.py | Adds missing dependency mapping for subquery_prepared_statements |
| src/include/distributed/remote_commands.h | Declares new remote prepare / prepared-query send wrappers |
| src/include/distributed/prepared_statement_cache.h | New public API for prepared-statement caching + fast-path hooks |
| src/include/distributed/multi_physical_planner.h | Adds Job/Task fields needed to carry cache keys/templates |
| src/include/distributed/distributed_planner.h | Extends fast-path context with distribution-key Param id |
| src/include/distributed/connection_management.h | Adds per-connection prepared-statement cache pointer |
| src/include/distributed/citus_custom_scan.h | Exposes EnsureAnchorShardsInJobExist for reuse |
| src/include/distributed/citus_clauses.h | Adds EVALUATE_FUNCTIONS mode + function-only evaluation API |
| src/backend/distributed/utils/citus_outfuncs.c | Extends node output for new/needed fields |
| src/backend/distributed/utils/citus_copyfuncs.c | Extends node copy logic for new/needed fields |
| src/backend/distributed/utils/citus_clauses.c | Implements function-only coordinator evaluation mode |
| src/backend/distributed/shared_library_init.c | Registers citus.enable_prepared_statement_caching GUC |
| src/backend/distributed/planner/multi_router_planner.c | Captures distribution-key Param id for router jobs (incl. INSERT path) |
| src/backend/distributed/planner/fast_path_router_planner.c | Records distribution-key Param id in fast-path restriction context |
| src/backend/distributed/planner/distributed_planner.c | Clears stale task pointers before copy/serialization checks |
| src/backend/distributed/planner/deparse_shard_query.c | Adds deparse-on-demand path for fast-path tasks (and caches result) |
| src/backend/distributed/metadata/XXnahQZC | New file present in PR metadata directory |
| src/backend/distributed/executor/prepared_statement_cache.c | New core implementation of per-connection prepared-statement caching |
| src/backend/distributed/executor/local_executor.c | Local planning support for fast-path tasks without query strings |
| src/backend/distributed/executor/citus_custom_scan.c | Integrates cache fast-path + task stamping in scan setup |
| src/backend/distributed/executor/adaptive_executor.c | Integrates cache dispatch path before parameterized-SQL fallback |
| src/backend/distributed/connection/remote_commands.c | Implements SendRemotePrepare / SendRemotePreparedQuery wrappers |
| src/backend/distributed/connection/connection_management.c | Ensures cache teardown when worker connections are closed |
Review details
Suppressed comments (1)
src/backend/distributed/executor/citus_custom_scan.c:499
- PreparedStatementCacheSaveTemplate() is called unconditionally in CitusBeginModifyScan(), and its returned template is attached to tasks even when ModifyJobNeedsEvaluation() runs (i.e., requiresCoordinatorEvaluation). In that case, the normal path evaluates coordinator-only functions/expressions in jobQuery, but the prepared-statement template still contains those function calls and would be PQprepare'd on the worker, changing semantics compared to the existing parameterized-SQL path. Until the template is produced with coordinator function evaluation while preserving Param nodes (e.g., using the new EVALUATE_FUNCTIONS mode), caching should be disabled for jobs that require coordinator evaluation.
Query *jobQuery = workerJob->jobQuery;
Query *savedJobQuery = PreparedStatementCacheSaveTemplate(originalDistributedPlan);
if (ModifyJobNeedsEvaluation(workerJob))
{
ExecuteCoordinatorEvaluableExpressions(jobQuery, planState);
- Files reviewed: 25/26 changed files
- Comments generated: 5
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Colm (colm-mchugh)
force-pushed
the
colm/prepared-stmts-poc
branch
from
September 9, 2026 16:36
598557f to
adcbd39
Compare
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Cached execution currently risks incorrect coordinator evaluation, shard routing, parameter handling, and placement selection.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (2)
src/backend/distributed/executor/prepared_statement_cache.c:406
- This hand-built SELECT task bypasses
GenerateSingleShardRouterTaskList()and therefore never appliesReorderTaskPlacementsByTaskAssignmentPolicy(). With replicated shards and round-robin assignment, cache-hit executions revert to metadata placement order (and can retain a coordinator placement that the normal path removes), so enabling the cache changes load distribution. Apply the same placement-policy processing as the normal regenerated-task path.
List *placementList =
CreateTaskPlacementListForShardIntervals(shardIntervalListList, shardsPresent,
true, false);
src/backend/distributed/executor/prepared_statement_cache.c:388
- The raw parameter Datum is passed to the partition function without checking that
param->ptypematches the distribution-column type. Single-row INSERT captures a Param after stripping implicit coercions, so e.g.PREPARE p(int) AS INSERT INTO numeric_dist VALUES ($1)reaches this code with anint4Datum for anumerickey, which can assert/crash or route to the wrong shard. Convert it withTransformPartitionRestrictionValue()before callingFindShardInterval(), asTargetShardIntervalForFastPathQuery()already does.
Oid relationId = linitial_oid(plan->relationIdList);
CitusTableCacheEntry *tableEntry = GetCitusTableCacheEntry(relationId);
return FindShardInterval(param->value, tableEntry);
- Files reviewed: 27/28 changed files
- Comments generated: 3
- Review effort level: Balanced
Colm (colm-mchugh)
force-pushed
the
colm/prepared-stmts-poc
branch
from
September 11, 2026 12:39
adcbd39 to
2cf3184
Compare
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Cache error handling, row-lock semantics, shard-split concurrency, and direct cache-path testing have unresolved issues.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 28/29 changed files
- Comments generated: 4
- Review effort level: Balanced
DESCRIPTION: Adds citus.enable_prepared_statement_caching to cache prepared
statement plans on worker connections for fast-path queries
Repeated executions of a prepared statement re-parse and re-plan the
shard query on the worker every time, because Citus sends the query via
PQsendQueryParams, which uses libpq's unnamed prepared statement. The
plan is built and discarded on each execution. The coordinator also
deparses the query tree to SQL for every execution.
This adds citus.enable_prepared_statement_caching (PGC_USERSET, default
off). When enabled, single-shard fast-path queries are sent using named
prepared statements: PQprepare on first use, then PQsendQueryPrepared for
subsequent executions. The worker keeps the CachedPlanSource/CachedPlan
for the lifetime of the connection, so later executions skip parse and
plan entirely and go straight to bind/execute.
This complements Citus's local plan caching, which already caches prepared
statement plans for local execution ((see `local_plan_cache.c`,
`local_executor.c`), by caching plans for remote execution so achieving
plan caching for all executions of a prepared statement for a fast path
query.
Cache
-----
Each MultiConnection gains a hash table keyed by (planId, shardId)
mapping to the statement name prepared on that connection. The cache
holds only names and parameter types (~264 bytes/entry); the plans
themselves live in the worker backend. It is capped at
`MAX_CACHED_STMTS_PER_CONNECTION` entries; beyond that, executions
fall back to plain parameterized SQL. Entries are dropped with the
connection, and a lost connection simply re-prepares on reconnect.
Design
------
Caching logic is encapsulated in prepared_statement_cache.c, which owns
both the hash table and the decisions about when caching applies. Call
sites invoke the API unconditionally and carry no policy of their own;
each entry point returns early when the GUC is off:
`PreparedStatementCacheTryFastPath()` build a Task straight from
ParamListInfo, skipping replan
`PreparedStatementCacheSaveTemplate()` save the parameterized job query
on the original plan
`PreparedStatementCacheAttachToTasks()` stamp planId and template onto
each task
`PreparedStatementCacheSendQuery()` lookup, prepare and dispatch on a
worker connection
`SendQuery()` reports its outcome as a PreparedStatementSendStatus value so
the executor can distinguish a dispatched query from a lost connection or a
cache-full fallback, for which the deparsed SQL is returned to the caller
to send down the existing path.
Executor
--------
`SendNextQuery()` consults the cache before the existing parameterized-SQL
path. Because parameters must survive to be bound remotely,
`MarkUnreferencedExternParams()` is skipped when caching is enabled.
Fast path
---------
`CitusBeginReadOnlyScan()` and `CitusBeginModifyScan()` try the cache-hit
fast path on second and subsequent executions of a generic plan. Instead
of deep-copying the `DistributedPlan`, evaluating coordinator expressions
and regenerating the task, the distribution key value is read directly
from `ParamListInfo`, the shard is found with `FindShardInterval()`, and a
minimal Task is built. The parameterized job query is saved once on the
original plan (Job->savedJobQueryForCaching) and reused, avoiding a
per-execution `copyObject()`. DML additionally acquires metadata locks and
assigns the first replica.
Applies to single-shard router queries with a parameterized distribution
key: SELECT, UPDATE, DELETE and single-row INSERT. Multi-row INSERT is
excluded: each shard's task carries only its own subset of VALUES rows,
so a statement deparsed once from the whole job query would send every
row to every shard.
Behaviour is unchanged when the GUC is off.
Colm (colm-mchugh)
force-pushed
the
colm/prepared-stmts-poc
branch
from
September 11, 2026 16:28
2cf3184 to
8b66706
Compare
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Cached execution currently regresses tenant attribution, partition-key statistics, and PostgreSQL 17 result chunking.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 30/32 changed files
- Comments generated: 3
- Review effort level: Balanced
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Custom plans pollute the cache, synchronous preparation undermines nonblocking execution, and reconnect coverage does not verify re-preparation.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 30/32 changed files
- Comments generated: 3
- Review effort level: Balanced
Colm (colm-mchugh)
force-pushed
the
colm/prepared-stmts-poc
branch
from
September 11, 2026 19:30
8b66706 to
d3ef235
Compare
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Fast-path placement and stale-plan-state handling contain correctness risks, and one regression test leaks session configuration.
Get a fresh assessment by requesting another Copilot review.
Review details
- Files reviewed: 30/32 changed files
- Comments generated: 3
- Review effort level: Balanced
Colm (colm-mchugh)
force-pushed
the
colm/prepared-stmts-poc
branch
2 times, most recently
from
September 11, 2026 21:52
e48ebcc to
ce41f4a
Compare
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The new function-only evaluator is ineffective, and remote DML cache engagement is not directly tested.
Get a fresh assessment by requesting another Copilot review.
Review details
- Files reviewed: 30/31 changed files
- Comments generated: 2
- Review effort level: Balanced
Colm (colm-mchugh)
force-pushed
the
colm/prepared-stmts-poc
branch
from
September 11, 2026 22:28
ce41f4a to
44d77e5
Compare
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
A connection failure during preparation currently aborts the query instead of entering the executor’s retry path.
Get a fresh assessment by requesting another Copilot review.
Review details
- Files reviewed: 28/29 changed files
- Comments generated: 1
- Review effort level: Balanced
Comment on lines
+611
to
+619
| if (PQresultStatus(result) != PGRES_COMMAND_OK) | ||
| { | ||
| /* | ||
| * A rejected statement is a query error, not a connection failure. | ||
| * Treating it as a lost connection sends the executor down the | ||
| * reconnect path while placements are still attached. | ||
| */ | ||
| ReportResultError(connection, result, ERROR); | ||
| } |
Review comments: - Drop the write-only paramTypes, paramCount and parameterizedQueryString from PreparedStatementCacheEntry; free queryString on the PREPARED_STMT_FAILED path; add an insert alias when deparsing shard-INSERTs. - Carry `relationRowLockList` on the fast-path task, so `SELECT ... FOR UPDATE` still opens a remote transaction and holds the row lock to the end of the enclosing one. - Publish the cache entry only after the worker accepts the statement; a rejected `PQprepare` was leaving an entry naming a statement that does not exist. - Recheck `ShardExists()` after `AcquireMetadataLocks()` and decline, so a concurrent split reroutes through normal planning instead of erroring. - Don't attach templates to custom plans: their planId is new per execution, so every execution prepared again and consumed a cache slot it could never hit. Gate on plan reuse, as `IsLocalPlanCachingSupported()` does. - Set `workerJob->partitionKeyValue` per execution; local execution and query statistics read it from the job. Save the planner's value first and put it back at end of scan rather than nulling the field: the fast path's Const dies with the portal, but clearing it made `ModifyJobNeedsEvaluation()` treat the distribution column value as unknown and forced every deferred-pruning INSERT through coordinator evaluation, with or without the cache. - `PQprepare()` blocks uninterruptibly, so a cancel against a slow worker was lost. Use `PQsendPrepare()` with `GetRemoteCommandResult()`. The wait is still synchronous for sibling placements; full executor integration is follow-up. - Share `SetRemoteRowMode()` with `SendNextQuery()`; caching was hard-coded to single-row mode and ignored `citus.executor_chunk_size`. - Decline caching while `citus.stat_tenants_track` is on. The tenant id travels in the query text, and a reused statement would freeze the first tenant rather than attribute per execution. Also address issues discovered with running multi, multi-1, and multi-mx test suites with `citus.enable_prepared_statement_caching` enabled: - Crash on shard key mismatch: `strip_implicit_coercions()` records the raw Param, so `PREPARE p(int)` on a numeric column hashed an int4 Datum as numeric. Instead, coerce via `TransformPartitionRestrictionValue()` and decline if not coercible. - Don't bypass coordinator evaluation and incorrectly ship nextval() to a worker. Correctly propagate `requiresCoordinatorEvaluation` on the deferred-pruning path. - Assert fail on unused params of a user-defined type; now mark against `jobQueryForPrepare` when that is present, instead of depending on the GUC value. - All three surfaced as the same assert: a rejected `PQprepare` was reported as MULTI_CONNECTION_LOST and sent the executor down the reconnect path with placements still attached. - `BuildFastPathTask()` should not ignore round-robin task assignment. - `parametersInQueryStringResolved` was set unconditionally, suppressing parameter types for a template that still carries $n. - EXPLAIN ANALYZE inherited `jobQueryForPrepare` through copyObject, and the fast path leaked $1 into the reported worker plan. - Deparse path: dangling pointer from freeing the task's own string, and a per-tuple memory context for a string the task retains. Drop the EVALUATE_FUNCTIONS evaluation mode and ExecuteCoordinatorEvaluableFunctions(). citus_evaluate_expr() bails out for every mode except EVALUATE_FUNCTIONS_PARAMS, so the helper never evaluated anything, and nothing called it; caching declines on requiresCoordinatorEvaluation instead. citus_clauses is untouched by this branch again. Tests asserted only results, which are identical whether or not the cache is engaged. Test 13 now asserts the wire protocol, Test 14 the remote transaction for FOR UPDATE, Test 15 the INSERT/UPDATE/DELETE dispatch, and Test 7 the re-prepare after a connection drop. DML reaches the cache through its own eligibility and task-building branches: with template saving disabled for non-SELECT, every diff came from Test 15 and no pre-existing test noticed. The suite sets `citus.stat_tenants_track = 'ALL'`, which the cache declines, so these tests set it to 'none'. Adds a CitusPreparedStatementCachingConfig arbitrary-configs class, which inherits the 'none' default and is the main CI coverage. Known limitation: caching is disabled while `citus.stat_tenants_track` is enabled. No regression test for: the SendRemotePrepare changes, the concurrent-split recheck, and the stale cache entry; all verified by hand.
Colm (colm-mchugh)
force-pushed
the
colm/prepared-stmts-poc
branch
from
September 11, 2026 22:58
44d77e5 to
d49e8af
Compare
Contributor
There was a problem hiding this comment.
🔵 Needs a closer look
The cross-layer planner, executor, memory-lifetime, and connection-state changes require final human validation despite extensive regression coverage.
Review details
- Files reviewed: 28/29 changed files
- Comments generated: 0 new
- Review effort level: Balanced
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
DESCRIPTION: Adds citus.enable_prepared_statement_caching to cache prepared statement plans on worker connections for fast-path queries
Repeated executions of a prepared statement re-parse and re-plan the shard query on the worker every time, because Citus sends the query via PQsendQueryParams, which uses libpq's unnamed prepared statement. The plan is built and discarded on each execution. The coordinator also deparses the query tree to SQL for every execution.
This adds citus.enable_prepared_statement_caching (PGC_USERSET, default off). When enabled, single-shard fast-path queries are sent using named prepared statements: PQprepare on first use, then PQsendQueryPrepared for subsequent executions. The worker keeps the CachedPlanSource/CachedPlan for the lifetime of the connection, so later executions skip parse and plan entirely and go straight to bind/execute.
This complements Citus's local plan caching, which already caches prepared statement plans for local execution ((see
local_plan_cache.c,local_executor.c), by caching plans for remote execution so achievingplan caching for all executions of a prepared statement for a fast path query.
Cache
Each MultiConnection gains a hash table keyed by (planId, shardId) mapping to the statement name prepared on that connection. The cache holds only names and parameter types (~264 bytes/entry); the plans themselves live in the worker backend. It is capped at
MAX_CACHED_STMTS_PER_CONNECTIONentries; beyond that, executions fall back to plain parameterized SQL. Entries are dropped with the connection, and a lost connection simply re-prepares on reconnect.Design
Caching logic is encapsulated in prepared_statement_cache.c, which owns both the hash table and the decisions about when caching applies. Call sites invoke the API unconditionally and carry no policy of their own; each entry point returns early when the GUC is off:
PreparedStatementCacheTryFastPath()build a Task straight fromParamListInfo, skipping replan
PreparedStatementCacheSaveTemplate()save the parameterized job queryon the original plan
PreparedStatementCacheAttachToTasks()stamp planId and template ontoeach task
PreparedStatementCacheSendQuery()lookup, prepare and dispatch on aworker connection
SendQuery()reports its outcome as a PreparedStatementSendStatus value so the executor can distinguish a dispatched query from a lost connection or a cache-full fallback, for which the deparsed SQL is returned to the caller to send down the existing path.Executor
SendNextQuery()consults the cache before the existing parameterized-SQL path. Because parameters must survive to be bound remotely,MarkUnreferencedExternParams()is skipped when caching is enabled.Fast path
CitusBeginReadOnlyScan()andCitusBeginModifyScan()try the cache-hit fast path on second and subsequent executions of a generic plan. Instead of deep-copying theDistributedPlan, evaluating coordinator expressions and regenerating the task, the distribution key value is read directly fromParamListInfo, the shard is found withFindShardInterval(), and a minimal Task is built. The parameterized job query is saved once on the original plan (Job->savedJobQueryForCaching) and reused, avoiding a per-executioncopyObject(). DML additionally acquires metadata locks and assigns the first replica.Applies to single-shard router queries with a parameterized distribution key: SELECT, UPDATE, DELETE and single-row INSERT. Multi-row INSERT is excluded: each shard's task carries only its own subset of VALUES rows, so a statement deparsed once from the whole job query would send every row to every shard.
Behaviour is unchanged when the GUC is off.