Open
Conversation
…ection integration Add per-connection prepared statement cache infrastructure: - New GUC citus.enable_prepared_statement_caching (default off, PGC_USERSET) - PreparedStatementCacheEntry/Key types with HTAB-based cache keyed by (planId, shardId), with MAX_CACHED_STMTS_PER_CONNECTION=1000 hard limit - Create/Lookup/Insert/Destroy cache lifecycle functions in prepared_statement_cache.c (TopMemoryContext for connection-scoped lifetime) - HTAB *preparedStatementCache field added to MultiConnection struct - Cache cleanup integrated into ShutdownConnection() and CitusPQFinish() to ensure no leaks on connection close
Core integration (prepared statement caching in executor):
1. Query template preservation (citus_clauses.c/h):
- Add EVALUATE_FUNCTIONS mode to CoordinatorEvaluationMode enum
- New ExecuteCoordinatorEvaluableFunctions() evaluates functions
while preserving Param nodes for parameterized queries
2. Task metadata (multi_physical_planner.h, citus_custom_scan.c):
- Add preparedStatementPlanId and jobQueryForPrepare fields to Task
- CitusBeginReadOnlyScan and CitusBeginModifyScan save pre-evaluation
job query and annotate tasks with plan ID for cache lookup
- INSERT commands excluded (query trees have special RTEs that
pg_get_query_def cannot deparse)
3. SendNextQuery prepared statement path (adaptive_executor.c):
- New if-block at top of SendNextQuery() checks
EnablePreparedStatementCaching && task->jobQueryForPrepare
- Cache miss: UpdateRelationToShardNames + pg_get_query_def to
construct parameterized shard query, PQprepare via SendRemotePrepare,
cache the entry, then PQsendQueryPrepared via SendRemotePreparedQuery
- Cache hit: reuse cached statement name with PQsendQueryPrepared
- Fallback: goto plain_sql when cache is full
- Skip MarkUnreferencedExternParams when caching enabled to preserve
original parameter types
4. libpq wrappers (remote_commands.c/h):
- SendRemotePrepare: synchronous PQprepare wrapper
- SendRemotePreparedQuery: async PQsendQueryPrepared wrapper
5. Cache key fix (prepared_statement_cache.c/h):
- Changed planId from uint32 to uint64 to match DistributedPlan.planId
Tests:
- Regression test suite for prepared statement caching (GUC toggle,
single-shard SELECT, multi-shard routing, INSERT/UPDATE/DELETE,
function evaluation, concurrent statements, connection loss re-prepare)
On cache hits (2nd+ execution with caching enabled), skip the expensive CopyDistributedPlanWithoutCache, ExecuteCoordinatorEvaluableExpressions, and RegenerateTaskForFasthPathQuery. Instead: 1. Preserve distribution key param index at plan time so the executor can extract the value directly from ParamListInfo. 2. In CitusBeginReadOnlyScan, when caching is enabled and the plan has been executed before, extract the distribution key from params, look up the shard interval via FindShardInterval (two hash lookups), and build a minimal Task - no deep copy, no tree walks, no deparse. 3. In SendNextQuery, defer TaskQueryStringAtIndex until the plain_sql label (cache hits never need it). Handle cache-full fallback by constructing the query string inline from the saved job query. Safety: always reset workerJob->taskList = NIL at the top of CitusBeginReadOnlyScan to prevent dangling pointers if the GUC is toggled between executions.
…ogging Add DML cache-hit fast path in CitusBeginModifyScan() for UPDATE/DELETE queries with deferred pruning and no coordinator-evaluable functions. Mirrors the existing SELECT fast path in CitusBeginReadOnlyScan(), skipping CopyDistributedPlanWithoutCache, expression evaluation, and RegenerateTaskForFasthPathQuery by extracting the distribution key directly from ParamListInfo and building a minimal MODIFY_TASK. Bug fixes: - Add missing COPY_SCALAR_FIELD(distributionKeyParamId) in copyJobInfo() and COPY_SCALAR_FIELD(preparedStatementPlanId) in CopyNodeTask() so these fields survive copyObject() in the PG plan cache. - Add missing WRITE_INT_FIELD(distributionKeyParamId) in OutJobFields() and WRITE_UINT64_FIELD(preparedStatementPlanId) in OutTask(). - Fix shardIntervalListList nesting: use list_make1(list_make1(...)) instead of flat list in the SELECT cache-hit fast path. - Fix use-after-free: move stale taskList/parametersInJobQueryResolved cleanup from AdaptiveExecutorCreateScan to CitusEndScan, gated on deferredPruning, so the plan is always clean for the next execution. - Fix DML returning 0 rows: gate stale-state cleanup on deferredPruning to avoid clearing real taskList on non-deferred plans (INSERTs, etc.). - Fix DML savedJobQueryForCaching: save template from original plan's jobQuery (pre-copy, pre-eval) since PartiallyEvaluateExpression with EVALUATE_FUNCTIONS silently switches to EVALUATE_FUNCTIONS_PARAMS for non-SELECT queries, resolving Params prematurely. - Skip expensive two-phase evaluation path for INSERT statements since they are excluded from worker-side prepared statement caching. Add DEBUG2 logging on four code paths: SELECT cache-hit fast path, DML cache-hit fast path, DML two-phase evaluation, and SendNextQuery cache hit/miss in adaptive_executor.c.
Single-row INSERT with a parameterized distribution key and no coordinator evaluation now uses the same cache-hit fast path as SELECT/UPDATE/DELETE. Changes: - RouterInsertJob: capture Param index for distribution key column so the executor can extract the value from ParamListInfo - CitusBeginModifyScan: remove CMD_INSERT exclusion from three guards (fast path entry, saved query template, metadata attachment), set anchorDistributedTableId on fast-path tasks for INSERT - SendNextQuery: branch on CMD_INSERT to use deparse_shard_query() instead of UpdateRelationToShardNames + pg_get_query_def for both cache-miss and cache-full fallback paths - Update test comments to reflect INSERT is now cached Also fixes INSERT ON CONFLICT deparse with shard name in rtable_names When deparse_shard_query deparses an INSERT for a specific shard, qualified column references in ON CONFLICT DO UPDATE (e.g. 'tablename.col') used the original base table name instead of the shard name. This caused 'missing FROM-clause entry' errors on workers and connection hangs. The fix: in get_query_def_extended, after set_deparse_for_query builds the rtable_names list, replace the INSERT target relation's name with the shard name when (a) a valid distrelid/shardid is set, (b) the RTE matches the distributed relation, and (c) the RTE has no user-defined alias (preserving the citus_table_alias used by INSERT...SELECT).
The fast-path optimisation (Phase 4+) creates tasks with TASK_QUERY_NULL type to skip expensive deparsing. Multiple code paths call TaskQueryString(task) which crashes on TASK_QUERY_NULL tasks: 1. Local executor — ParseQueryString(TaskQueryString(task), ...) when a fast-path task routes to local execution (shards on coordinator). 2. Adaptive executor plain_sql fallback — TaskQueryStringAtIndex() when the prepared statement cache is full. 3. CheckNodeCopyAndSerialization (cassert only) — serialises stale task pointers from a previous execution's freed memory. Fixes: - TaskQueryString(): handle TASK_QUERY_NULL by deparsing from jobQueryForPrepare template instead of crashing. Caches the result so subsequent calls don't re-deparse. - Local executor: add TASK_QUERY_NULL handler that deparses from the template and plans locally. - SELECT fast path: add CacheLocalPlanForShardQuery so local executor can use cached plan (DML fast path already had this). - SELECT fast path: set anchorDistributedTableId (needed for INSERT deparsing via deparse_shard_query). - GetDistributedPlan: clear stale taskList before CheckNodeCopyAndSerialization to prevent use-after-free in cassert. - CopyNodeTask: add missing fields (jobQueryForPrepare, partitionKeyValue, colocationId). - copyJobInfo: add missing fields (colocationId, savedJobQueryForCaching). - OutTask/OutJobFields: add colocationId for serialisation.
|
@SarthakDalmia1 please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
|
@microsoft-github-policy-service agree |
colm-mchugh
reviewed
Apr 5, 2026
Contributor
There was a problem hiding this comment.
@SarthakDalmia1 PG16 support is now in the dev branch.
The feature is not yet in main, so this PR is not needed.
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: PR description that will go into the change log, up to 78 characters