diff --git a/src/backend/distributed/connection/connection_management.c b/src/backend/distributed/connection/connection_management.c index 0cb2b7c30c1..1491110d513 100644 --- a/src/backend/distributed/connection/connection_management.c +++ b/src/backend/distributed/connection/connection_management.c @@ -36,6 +36,7 @@ #include "distributed/memutils.h" #include "distributed/metadata_cache.h" #include "distributed/placement_connection.h" +#include "distributed/prepared_statement_cache.h" #include "distributed/remote_commands.h" #include "distributed/run_from_same_connection.h" #include "distributed/shared_connection_stats.h" @@ -794,6 +795,9 @@ ShutdownConnection(MultiConnection *connection) { SendCancelationRequest(connection); } + + PreparedStatementCacheDestroy(&connection->preparedStatementCache); + CitusPQFinish(connection); } @@ -1232,6 +1236,8 @@ CloseNotReadyMultiConnectionStates(List *connectionStates) static void CitusPQFinish(MultiConnection *connection) { + PreparedStatementCacheDestroy(&connection->preparedStatementCache); + if (connection->pgConn != NULL) { PQfinish(connection->pgConn); diff --git a/src/backend/distributed/connection/remote_commands.c b/src/backend/distributed/connection/remote_commands.c index cd0c1951999..82f8f3777b3 100644 --- a/src/backend/distributed/connection/remote_commands.c +++ b/src/backend/distributed/connection/remote_commands.c @@ -576,6 +576,86 @@ SendRemoteCommand(MultiConnection *connection, const char *command) } +/* + * SendRemotePrepare prepares a statement on a worker connection and waits for + * the worker to confirm it. The wait is interruptible, so a cancel is honoured + * even when the worker is slow to answer. Returns 1 on success, 0 on failure. + */ +int +SendRemotePrepare(MultiConnection *connection, const char *stmtName, + const char *query, int nParams, const Oid *paramTypes) +{ + PGconn *pgConn = connection->pgConn; + + LogRemoteCommand(connection, query); + + if (!pgConn || PQstatus(pgConn) != CONNECTION_OK) + { + return 0; + } + + Assert(PQisnonblocking(pgConn)); + + if (PQsendPrepare(pgConn, stmtName, query, nParams, paramTypes) == 0) + { + return 0; + } + + bool raiseInterrupts = true; + PGresult *result = GetRemoteCommandResult(connection, raiseInterrupts); + if (result == NULL) + { + return 0; + } + + if (PQresultStatus(result) != PGRES_COMMAND_OK) + { + /* + * Raise for a lost connection as well, which GetRemoteCommandResult() + * reports as a synthetic PGRES_FATAL_ERROR. Returning failure instead + * marks the connection lost while placements are still attached, and + * RestartConnection() asserts that they are not. + */ + ReportResultError(connection, result, ERROR); + } + + PQclear(result); + ForgetResults(connection); + + return 1; +} + + +/* + * SendRemotePreparedQuery wraps PQsendQueryPrepared() to asynchronously + * execute a previously prepared statement on a worker connection. Follows + * the same pattern as SendRemoteCommandParams(). Returns the result of + * PQsendQueryPrepared (1 on success, 0 on failure). + */ +int +SendRemotePreparedQuery(MultiConnection *connection, const char *stmtName, + int nParams, const char *const *paramValues, + bool binaryResults) +{ + PGconn *pgConn = connection->pgConn; + + /* log the statement name for debugging */ + LogRemoteCommand(connection, stmtName); + + if (!pgConn || PQstatus(pgConn) != CONNECTION_OK) + { + return 0; + } + + Assert(PQisnonblocking(pgConn)); + + int rc = PQsendQueryPrepared(pgConn, stmtName, nParams, paramValues, + NULL, NULL, binaryResults ? 1 : 0); + + return rc; +} + + /* * ExecuteRemoteCommandAndCheckResult executes the given command in the remote node and * checks if the result is equal to the expected result. If the result is equal to the diff --git a/src/backend/distributed/executor/adaptive_executor.c b/src/backend/distributed/executor/adaptive_executor.c index 443e9ab2ac5..7845bad38a6 100644 --- a/src/backend/distributed/executor/adaptive_executor.c +++ b/src/backend/distributed/executor/adaptive_executor.c @@ -149,6 +149,7 @@ #include "distributed/backend_data.h" #include "distributed/cancel_utils.h" #include "distributed/citus_custom_scan.h" +#include "distributed/citus_ruleutils.h" #include "distributed/citus_safe_lib.h" #include "distributed/commands/multi_copy.h" #include "distributed/connection_management.h" @@ -166,6 +167,7 @@ #include "distributed/param_utils.h" #include "distributed/placement_access.h" #include "distributed/placement_connection.h" +#include "distributed/prepared_statement_cache.h" #include "distributed/relation_access_tracking.h" #include "distributed/remote_commands.h" #include "distributed/repartition_join_execution.h" @@ -739,6 +741,7 @@ static bool StartPlacementExecutionOnSession(TaskPlacementExecution *placementEx WorkerSession *session); static bool SendNextQuery(TaskPlacementExecution *placementExecution, WorkerSession *session); +static bool SetRemoteRowMode(MultiConnection *connection); static void ConnectionStateMachine(WorkerSession *session); static bool HasUnfinishedTaskForSession(WorkerSession *session); static void HandleMultiConnectionSuccess(WorkerSession *session, bool newConnection); @@ -927,8 +930,23 @@ AdaptiveExecutorStart(CitusScanState *scanState) */ if (paramListInfo && !paramListInfo->paramFetch) { + /* + * A cached statement is deparsed from the saved template, which still + * carries the Params. job->jobQuery has them resolved to constants, so + * marking against it would strip types the worker still needs. + */ + Node *queryForParams = (Node *) job->jobQuery; + if (taskList != NIL) + { + Task *firstTask = (Task *) linitial(taskList); + if (firstTask->jobQueryForPrepare != NULL) + { + queryForParams = (Node *) firstTask->jobQueryForPrepare; + } + } + paramListInfo = copyParamList(paramListInfo); - MarkUnreferencedExternParams((Node *) job->jobQuery, paramListInfo); + MarkUnreferencedExternParams(queryForParams, paramListInfo); } DistributedExecution *execution = CreateDistributedExecution( @@ -4256,7 +4274,29 @@ SendNextQuery(TaskPlacementExecution *placementExecution, uint32 queryIndex = placementExecution->queryIndex; Assert(queryIndex < task->queryCount); - char *queryString = TaskQueryStringAtIndex(task, queryIndex); + char *queryString = NULL; + + /* + * Must be attempted before the paramListInfo path below: a fast-path task + * reports its parameters as resolved, so that path would send plain SQL. + */ + PreparedStatementSendStatus cacheStatus = + PreparedStatementCacheSendQuery(connection, task, paramListInfo, + binaryResults, &queryString); + if (cacheStatus == PREPARED_STMT_SENT) + { + return SetRemoteRowMode(connection); + } + else if (cacheStatus == PREPARED_STMT_FAILED) + { + return false; + } + + /* resolve queryString if not already set (e.g. from cache-full fallback) */ + if (queryString == NULL) + { + queryString = TaskQueryStringAtIndex(task, queryIndex); + } if (paramListInfo != NULL && !task->parametersInQueryStringResolved) { @@ -4304,6 +4344,17 @@ SendNextQuery(TaskPlacementExecution *placementExecution, return false; } + return SetRemoteRowMode(connection); +} + + +/* + * SetRemoteRowMode puts the connection into the incremental result mode the + * executor expects, chunked where libpq supports it. + */ +static bool +SetRemoteRowMode(MultiConnection *connection) +{ #ifdef LIBPQ_HAS_CHUNK_MODE int rowMode = PQsetChunkedRowsMode(connection->pgConn, ExecutorChunkSize); #else diff --git a/src/backend/distributed/executor/citus_custom_scan.c b/src/backend/distributed/executor/citus_custom_scan.c index 87fca7422b1..ebac9263fcc 100644 --- a/src/backend/distributed/executor/citus_custom_scan.c +++ b/src/backend/distributed/executor/citus_custom_scan.c @@ -43,10 +43,13 @@ #include "distributed/local_plan_cache.h" #include "distributed/merge_executor.h" #include "distributed/merge_planner.h" +#include "distributed/metadata_cache.h" #include "distributed/multi_executor.h" #include "distributed/multi_router_planner.h" #include "distributed/multi_server_executor.h" +#include "distributed/prepared_statement_cache.h" #include "distributed/shard_utils.h" +#include "distributed/shardinterval_utils.h" #include "distributed/sorted_merge.h" #include "distributed/stats/query_stats.h" #include "distributed/stats/stat_counters.h" @@ -82,7 +85,7 @@ static void SortedMergeReScan(CustomScanState *node); static void CitusEndScanCommon(CitusScanState *scanState); static void CitusReScanCommon(CustomScanState *node); static void EnsureForceDelegationDistributionKey(Job *job); -static void EnsureAnchorShardsInJobExist(Job *job); + static bool AnchorShardsInTaskListExist(List *taskList); static void TryToRerouteFastPathModifyQuery(Job *job); static void CheckQueryDeparseSafety(Query *query); @@ -355,6 +358,28 @@ CitusBeginReadOnlyScan(CustomScanState *node, EState *estate, int eflags) return; } + Job *workerJob = originalDistributedPlan->workerJob; + + /* + * Always clear stale task pointers from the previous fast-path execution. + * For deferred-pruning plans, workerJob->taskList is NIL until populated + * per-execution by RegenerateTaskForFasthPathQuery (or the fast path below). + * Without this reset, if a previous execution took the fast path and then the + * GUC is disabled, CopyDistributedPlanWithoutCache would deep-copy a stale + * taskList pointing into freed memory. + */ + workerJob->taskList = NIL; + workerJob->parametersInJobQueryResolved = false; + + /* + * A cached plan can build its task straight from the bound parameters, + * skipping the plan copy, coordinator evaluation and task regeneration. + */ + if (PreparedStatementCacheTryFastPath(scanState, estate, false)) + { + return; + } + /* * Create a copy of the generic plan for the current execution, but make a shallow * copy of the plan cache. That means we'll be able to access the plan cache via @@ -365,8 +390,8 @@ CitusBeginReadOnlyScan(CustomScanState *node, EState *estate, int eflags) CopyDistributedPlanWithoutCache(originalDistributedPlan); scanState->distributedPlan = currentPlan; - Job *workerJob = currentPlan->workerJob; - Query *jobQuery = workerJob->jobQuery; + Job *currentJob = currentPlan->workerJob; + Query *jobQuery = currentJob->jobQuery; PlanState *planState = &(scanState->customScanState.ss.ps); /* @@ -375,6 +400,8 @@ CitusBeginReadOnlyScan(CustomScanState *node, EState *estate, int eflags) */ Assert(currentPlan->fastPathRouterPlan || !EnableFastPathRouterPlanner); + Query *savedJobQuery = PreparedStatementCacheSaveTemplate(originalDistributedPlan); + /* * Evaluate parameters, because the parameters are only available on the * coordinator and are required for pruning. @@ -388,14 +415,16 @@ CitusBeginReadOnlyScan(CustomScanState *node, EState *estate, int eflags) ExecuteCoordinatorEvaluableExpressions(jobQuery, planState); /* job query no longer has parameters, so we should not send any */ - workerJob->parametersInJobQueryResolved = true; + currentJob->parametersInJobQueryResolved = true; /* parameters are filled in, so we can generate a task for this execution */ - RegenerateTaskForFasthPathQuery(workerJob); + RegenerateTaskForFasthPathQuery(currentJob); - if (IsLocalPlanCachingSupported(workerJob, originalDistributedPlan)) + PreparedStatementCacheAttachToTasks(currentPlan, currentJob, savedJobQuery); + + if (IsLocalPlanCachingSupported(currentJob, originalDistributedPlan)) { - Task *task = linitial(workerJob->taskList); + Task *task = linitial(currentJob->taskList); /* * We are going to execute this task locally. If it's not already in @@ -426,6 +455,28 @@ CitusBeginModifyScan(CustomScanState *node, EState *estate, int eflags) CitusScanState *scanState = (CitusScanState *) node; PlanState *planState = &(scanState->customScanState.ss.ps); DistributedPlan *originalDistributedPlan = scanState->distributedPlan; + Job *origWorkerJob = originalDistributedPlan->workerJob; + + /* + * Safety: clear stale task pointers from a previous fast-path execution. + * For deferred-pruning plans, workerJob->taskList is NIL until populated + * per-execution. Without this reset, CopyDistributedPlanWithoutCache + * would deep-copy a stale taskList pointing into freed memory. + */ + if (origWorkerJob->deferredPruning) + { + origWorkerJob->taskList = NIL; + origWorkerJob->parametersInJobQueryResolved = false; + } + + /* + * A cached plan can build its task straight from the bound parameters, + * skipping the plan copy, coordinator evaluation and task regeneration. + */ + if (PreparedStatementCacheTryFastPath(scanState, estate, true)) + { + return; + } MemoryContext localContext = AllocSetContextCreate(CurrentMemoryContext, "CitusBeginModifyScan", @@ -440,6 +491,8 @@ CitusBeginModifyScan(CustomScanState *node, EState *estate, int eflags) Query *jobQuery = workerJob->jobQuery; + Query *savedJobQuery = PreparedStatementCacheSaveTemplate(originalDistributedPlan); + if (ModifyJobNeedsEvaluation(workerJob)) { ExecuteCoordinatorEvaluableExpressions(jobQuery, planState); @@ -472,6 +525,8 @@ CitusBeginModifyScan(CustomScanState *node, EState *estate, int eflags) { RegenerateTaskForFasthPathQuery(workerJob); } + + PreparedStatementCacheAttachToTasks(currentPlan, workerJob, savedJobQuery); } else if (workerJob->requiresCoordinatorEvaluation) { @@ -574,7 +629,7 @@ TryToRerouteFastPathModifyQuery(Job *job) * EnsureAnchorShardsInJobExist ensures all shards are valid in job. * If it finds a non-existent shard in given job, it throws an error. */ -static void +void EnsureAnchorShardsInJobExist(Job *job) { if (!AnchorShardsInTaskListExist(job->taskList)) @@ -648,13 +703,17 @@ CopyDistributedPlanWithoutCache(DistributedPlan *originalDistributedPlan) { List *localPlannedStatements = originalDistributedPlan->workerJob->localPlannedStatements; + Query *savedJobQueryForCaching = + originalDistributedPlan->workerJob->savedJobQueryForCaching; originalDistributedPlan->workerJob->localPlannedStatements = NIL; + originalDistributedPlan->workerJob->savedJobQueryForCaching = NULL; DistributedPlan *distributedPlan = copyObject(originalDistributedPlan); - /* set back the immutable field */ + /* set back the immutable/cached fields */ originalDistributedPlan->workerJob->localPlannedStatements = localPlannedStatements; distributedPlan->workerJob->localPlannedStatements = localPlannedStatements; + originalDistributedPlan->workerJob->savedJobQueryForCaching = savedJobQueryForCaching; return distributedPlan; } @@ -915,6 +974,27 @@ CitusEndScanCommon(CitusScanState *scanState) /* queries without partition key are also recorded */ CitusQueryStatsExecutorsEntry(queryId, executorType, partitionKeyString); } + + /* + * Clear mutable per-execution state so the cached plan is clean for + * the next execution. The cache-hit fast paths in CitusBeginReadOnlyScan() + * and CitusBeginModifyScan() store a Task list and a partition key Const + * directly on the original plan's workerJob; both live in the per-execution + * memory context and become dangling after EndScan. In assert-checking + * builds the next execution's GetDistributedPlan() → copyObject() would + * traverse freed memory without this reset. + * + * partitionKeyValue goes back to what the planner set rather than to NULL: + * ModifyJobNeedsEvaluation() reads it to decide whether the distribution + * column value is already known, and nulling it forces every later + * execution through needless coordinator evaluation. + */ + if (workerJob != NULL && workerJob->savedJobQueryForCaching != NULL) + { + workerJob->taskList = NIL; + workerJob->parametersInJobQueryResolved = false; + workerJob->partitionKeyValue = workerJob->plannerPartitionKeyValue; + } } diff --git a/src/backend/distributed/executor/local_executor.c b/src/backend/distributed/executor/local_executor.c index 5480a1d142d..d19c48977bf 100644 --- a/src/backend/distributed/executor/local_executor.c +++ b/src/backend/distributed/executor/local_executor.c @@ -392,7 +392,42 @@ ExecuteLocalTaskListExtended(List *taskList, continue; } - if (taskType != TASK_QUERY_LOCAL_PLAN) + if (taskType == TASK_QUERY_NULL && task->jobQueryForPrepare != NULL) + { + /* + * Fast-path task from prepared statement caching: the task + * was built without a query string to avoid expensive + * deparsing. Deparse from the saved job query template now + * so we can plan locally. + */ + Query *queryForDeparse = copyObject(task->jobQueryForPrepare); + StringInfoData buf; + initStringInfo(&buf); + + if (queryForDeparse->commandType == CMD_INSERT) + { + /* upserts reference the target by name, which becomes the shard name */ + AddInsertAliasIfNeeded(queryForDeparse); + + deparse_shard_query(queryForDeparse, + task->anchorDistributedTableId, + task->anchorShardId, &buf); + } + else + { + UpdateRelationToShardNames((Node *) queryForDeparse, + task->relationShardList); + pg_get_query_def(queryForDeparse, &buf); + } + + Query *shardQuery = ParseQueryString(buf.data, + taskParameterTypes, + taskNumParams); + localPlan = planner(shardQuery, NULL, CURSOR_OPT_PARALLEL_OK, + paramListInfo); + pfree(buf.data); + } + else if (taskType != TASK_QUERY_LOCAL_PLAN) { Query *shardQuery = ParseQueryString(TaskQueryString(task), taskParameterTypes, diff --git a/src/backend/distributed/executor/prepared_statement_cache.c b/src/backend/distributed/executor/prepared_statement_cache.c new file mode 100644 index 00000000000..2e23806447a --- /dev/null +++ b/src/backend/distributed/executor/prepared_statement_cache.c @@ -0,0 +1,620 @@ +/*------------------------------------------------------------------------- + * + * prepared_statement_cache.c + * Per-connection cache for prepared statements on worker connections. + * + * When citus.enable_prepared_statement_caching is ON, the coordinator + * uses PQprepare/PQsendQueryPrepared on worker connections instead + * of PQsendQuery for fast-path prepared statement executions (generic + * plan, execution 6+). This module manages the per-connection hash + * table that tracks which statements have already been prepared on + * each connection. + * + * It also owns the two integration points that make the cache usable: + * the planner-side fast path that builds a Task straight from + * ParamListInfo (skipping replanning), and the executor-side dispatch + * that prepares or reuses the named statement on a connection. + * + * Copyright (c) Citus Data, Inc. + * + *------------------------------------------------------------------------- + */ + +#include "postgres.h" + +#include "access/htup_details.h" +#include "nodes/makefuncs.h" +#include "utils/hsearch.h" +#include "utils/lsyscache.h" +#include "utils/memutils.h" + +#include "distributed/citus_custom_scan.h" +#include "distributed/citus_ruleutils.h" +#include "distributed/citus_safe_lib.h" +#include "distributed/deparse_shard_query.h" +#include "distributed/distributed_execution_locks.h" +#include "distributed/executor_util.h" +#include "distributed/listutils.h" +#include "distributed/local_plan_cache.h" +#include "distributed/metadata_cache.h" +#include "distributed/metadata_utility.h" +#include "distributed/multi_executor.h" +#include "distributed/multi_explain.h" +#include "distributed/multi_physical_planner.h" +#include "distributed/multi_router_planner.h" +#include "distributed/prepared_statement_cache.h" +#include "distributed/remote_commands.h" +#include "distributed/shard_cleaner.h" +#include "distributed/shard_pruning.h" +#include "distributed/shardinterval_utils.h" +#include "distributed/stats/stat_tenants.h" + + +/* GUC: citus.enable_prepared_statement_caching */ +bool EnablePreparedStatementCaching = false; + + +/* + * PreparedStatementCachingUsable returns whether the cache may be used at all + * for this execution. + * + * Tenant statistics reach the worker as a comment on the query text. A reused + * prepared statement cannot carry that per execution -- it would freeze the + * tenant of whichever execution prepared it -- so decline while tracking is on + * rather than mis-attribute. + */ +static bool +PreparedStatementCachingUsable(void) +{ + return EnablePreparedStatementCaching && + StatTenantsTrack == STAT_TENANTS_TRACK_NONE; +} + + +/* + * PreparedStatementCacheCreate allocates a new hash table for caching + * prepared statement entries on a single worker connection. The hash + * table is allocated in TopMemoryContext so it survives across + * transactions (matching the connection lifetime). + */ +HTAB * +PreparedStatementCacheCreate(void) +{ + HASHCTL info; + + memset(&info, 0, sizeof(info)); + info.keysize = sizeof(PreparedStatementCacheKey); + info.entrysize = sizeof(PreparedStatementCacheEntry); + info.hcxt = TopMemoryContext; + + HTAB *cache = hash_create("Prepared Statement Cache", + 32, /* initial size */ + &info, + HASH_ELEM | HASH_BLOBS | HASH_CONTEXT); + + return cache; +} + + +/* + * PreparedStatementCacheLookup looks up a cache entry by (planId, shardId). + * Returns the entry if found, NULL otherwise. + */ +PreparedStatementCacheEntry * +PreparedStatementCacheLookup(HTAB *cache, uint64 planId, uint64 shardId) +{ + PreparedStatementCacheKey key; + + memset(&key, 0, sizeof(key)); + key.planId = planId; + key.shardId = shardId; + + PreparedStatementCacheEntry *entry = + (PreparedStatementCacheEntry *) hash_search(cache, &key, + HASH_FIND, NULL); + + return entry; +} + + +/* + * PreparedStatementCacheInsert records that stmtName has been prepared on this + * connection for (planId, shardId). + * + * Only call this once the worker has accepted the statement: an entry for a + * statement that does not exist makes every later execution on this connection + * fail with "prepared statement does not exist". + */ +PreparedStatementCacheEntry * +PreparedStatementCacheInsert(HTAB *cache, uint64 planId, uint64 shardId, + const char *stmtName) +{ + PreparedStatementCacheKey key; + + memset(&key, 0, sizeof(key)); + key.planId = planId; + key.shardId = shardId; + + bool found = false; + PreparedStatementCacheEntry *entry = + (PreparedStatementCacheEntry *) hash_search(cache, &key, + HASH_ENTER, &found); + + if (!found) + { + strlcpy(entry->stmtName, stmtName, MAX_STMT_NAME_LENGTH); + } + + return entry; +} + + +/* + * PreparedStatementCacheDestroy frees all memory used by the cache. + * Safe to call if the pointed to cache is NULL (NOP) + * Otherwise, frees all memory used by the cache and + * sets it to NULL. + */ +void +PreparedStatementCacheDestroy(HTAB **cache_ptr) +{ + Assert(cache_ptr != NULL); + HTAB *cache = *cache_ptr; + + if (cache == NULL) + { + return; + } + + hash_destroy(cache); + *cache_ptr = NULL; +} + + +/* + * DeparseTaskTemplate deparses the task's parameterized job query into SQL + * targeting the task's shard, leaving Param nodes ($1, ...) intact. + * + * The template is shared across executions, so a working copy is deparsed. + */ +static char * +DeparseTaskTemplate(Task *task) +{ + Query *queryCopy = copyObject(task->jobQueryForPrepare); + StringInfoData buf; + initStringInfo(&buf); + + if (queryCopy->commandType == CMD_INSERT) + { + /* upserts reference the target by name, which becomes the shard name */ + AddInsertAliasIfNeeded(queryCopy); + deparse_shard_query(queryCopy, task->anchorDistributedTableId, + task->anchorShardId, &buf); + } + else + { + UpdateRelationToShardNames((Node *) queryCopy, task->relationShardList); + pg_get_query_def(queryCopy, &buf); + } + + return buf.data; +} + + +/* + * PreparedStatementCacheSendQuery dispatches a task on a worker connection + * using a named prepared statement, preparing it first if this connection has + * not seen it before. + * + * Returns PREPARED_STMT_NOT_APPLICABLE when the task carries no template, in + * which case the caller must use its normal send path. When the connection's + * cache is full it returns PREPARED_STMT_FALLBACK and sets fallbackQueryString + * to parameterized SQL for the caller to send instead. + */ +PreparedStatementSendStatus +PreparedStatementCacheSendQuery(MultiConnection *connection, Task *task, + ParamListInfo paramListInfo, bool binaryResults, + char **fallbackQueryString) +{ + if (!PreparedStatementCachingUsable() || task->jobQueryForPrepare == NULL || + paramListInfo == NULL) + { + return PREPARED_STMT_NOT_APPLICABLE; + } + + Oid *parameterTypes = NULL; + const char **parameterValues = NULL; + + /* force evaluation of bound params */ + paramListInfo = copyParamList(paramListInfo); + int parameterCount = paramListInfo->numParams; + + ExtractParametersForRemoteExecution(paramListInfo, ¶meterTypes, + ¶meterValues); + + if (connection->preparedStatementCache == NULL) + { + connection->preparedStatementCache = PreparedStatementCacheCreate(); + } + + PreparedStatementCacheEntry *cacheEntry = + PreparedStatementCacheLookup(connection->preparedStatementCache, + task->preparedStatementPlanId, + task->anchorShardId); + if (cacheEntry == NULL) + { + elog(DEBUG2, "prepared statement cache miss: plan " UINT64_FORMAT + " shard " UINT64_FORMAT, + task->preparedStatementPlanId, task->anchorShardId); + + HTAB *cache = connection->preparedStatementCache; + + if (hash_get_num_entries(cache) >= MAX_CACHED_STMTS_PER_CONNECTION) + { + /* + * Cache full. The fast-path task has no query string of its own, so + * hand the caller freshly deparsed SQL. Clearing the resolved flag + * routes it through the parameterized send path. + */ + *fallbackQueryString = DeparseTaskTemplate(task); + task->parametersInQueryStringResolved = false; + return PREPARED_STMT_FALLBACK; + } + + char stmtName[MAX_STMT_NAME_LENGTH]; + SafeSnprintf(stmtName, MAX_STMT_NAME_LENGTH, "__citus_stmt_%ld", + (long) hash_get_num_entries(cache) + 1); + + char *queryString = DeparseTaskTemplate(task); + + if (SendRemotePrepare(connection, stmtName, queryString, + parameterCount, parameterTypes) == 0) + { + pfree(queryString); + connection->connectionState = MULTI_CONNECTION_LOST; + return PREPARED_STMT_FAILED; + } + + pfree(queryString); + + /* + * Publish only now: SendRemotePrepare raises on a rejected statement, so + * an earlier insert would leave an entry naming a statement the worker + * never created. + */ + cacheEntry = PreparedStatementCacheInsert(cache, + task->preparedStatementPlanId, + task->anchorShardId, stmtName); + } + else + { + elog(DEBUG2, "prepared statement cache hit: plan " UINT64_FORMAT + " shard " UINT64_FORMAT " stmt %s", + task->preparedStatementPlanId, task->anchorShardId, + cacheEntry->stmtName); + } + + if (SendRemotePreparedQuery(connection, cacheEntry->stmtName, parameterCount, + parameterValues, binaryResults) == 0) + { + connection->connectionState = MULTI_CONNECTION_LOST; + return PREPARED_STMT_FAILED; + } + + return PREPARED_STMT_SENT; +} + + +/* + * PreparedStatementCacheSaveTemplate returns the job query template used to + * build parameterized statements, saving it on the original plan on first use. + * + * The template must be captured before coordinator evaluation resolves Param + * nodes, and is kept on the original (cached) plan so later executions reuse + * it instead of paying for a copyObject each time. + * + * Returns NULL when the plan is not eligible for caching. + */ +Query * +PreparedStatementCacheSaveTemplate(DistributedPlan *originalPlan) +{ + Job *originalJob = originalPlan->workerJob; + Query *jobQuery = originalJob->jobQuery; + + if (!PreparedStatementCachingUsable()) + { + return NULL; + } + + if (jobQuery->commandType != CMD_SELECT) + { + /* + * nextval(), now() and other coordinator-evaluated expressions are + * resolved into the per-execution copy of the query, never into this + * template, so a cached statement would evaluate them on the worker. + */ + if (originalJob->requiresCoordinatorEvaluation) + { + return NULL; + } + + /* + * Multi-row INSERT can't be cached: each shard's task carries only its + * own subset of VALUES rows, but the statement is deparsed once from + * the whole job query, so every shard would receive every row. + */ + if (!originalJob->deferredPruning || + ExtractDistributedInsertValuesRTE(jobQuery) != NULL) + { + return NULL; + } + } + + if (originalJob->savedJobQueryForCaching == NULL) + { + MemoryContext oldContext = + MemoryContextSwitchTo(GetMemoryChunkContext(originalPlan)); + originalJob->savedJobQueryForCaching = copyObject(jobQuery); + MemoryContextSwitchTo(oldContext); + + /* the fast path has not run yet, so this is still the planner's value */ + originalJob->plannerPartitionKeyValue = originalJob->partitionKeyValue; + } + + return originalJob->savedJobQueryForCaching; +} + + +/* + * PreparedStatementCacheAttachToTasks stamps the cache key and query template + * onto each task so SendNextQuery can look them up. Does nothing when the plan + * is not eligible for caching. + */ +void +PreparedStatementCacheAttachToTasks(DistributedPlan *currentPlan, Job *workerJob, + Query *savedJobQuery) +{ + if (!PreparedStatementCachingUsable() || savedJobQuery == NULL) + { + return; + } + + /* + * Only attach to a plan that is being reused. A custom plan is rebuilt with a + * fresh planId per execution, so its entries could never be hit and would fill + * the connection's cache. Same reuse test as IsLocalPlanCachingSupported(). + */ + if (currentPlan->numberOfTimesExecuted < 1) + { + return; + } + + bool isInsert = workerJob->jobQuery->commandType == CMD_INSERT; + + Task *task = NULL; + foreach_declared_ptr(task, workerJob->taskList) + { + task->preparedStatementPlanId = currentPlan->planId; + task->jobQueryForPrepare = savedJobQuery; + + /* deparse_shard_query needs the target relation for INSERT */ + if (isInsert) + { + task->anchorDistributedTableId = linitial_oid(currentPlan->relationIdList); + } + } +} + + +/* + * FastPathShardInterval returns the shard the distribution key parameter routes + * to, or NULL if the fast path cannot be used for this execution. + * + * The value is returned via partitionKeyValue, coerced to the distribution + * column's type. Single-row INSERT records the Param with implicit coercions + * stripped, so the parameter type can differ from the column type. + */ +static ShardInterval * +FastPathShardInterval(DistributedPlan *plan, Job *workerJob, EState *estate, + Const **partitionKeyValue) +{ + int paramId = workerJob->distributionKeyParamId; + ParamListInfo paramListInfo = estate->es_param_list_info; + + if (paramId < 1 || paramListInfo == NULL || paramId > paramListInfo->numParams) + { + return NULL; + } + + ParamExternData *param = ¶mListInfo->params[paramId - 1]; + if (!OidIsValid(param->ptype) || param->isnull) + { + return NULL; + } + + Oid relationId = linitial_oid(plan->relationIdList); + CitusTableCacheEntry *tableEntry = GetCitusTableCacheEntry(relationId); + Var *partitionColumn = tableEntry->partitionColumn; + + if (partitionColumn == NULL) + { + return NULL; + } + + int16 typeLength; + bool typeByValue; + get_typlenbyval(param->ptype, &typeLength, &typeByValue); + Const *valueConst = makeConst(param->ptype, -1, InvalidOid, (int) typeLength, + param->value, false, typeByValue); + + if (param->ptype != partitionColumn->vartype) + { + bool missingOk = true; + valueConst = TransformPartitionRestrictionValue(partitionColumn, valueConst, + missingOk); + if (valueConst == NULL || valueConst->constisnull) + { + return NULL; + } + } + + *partitionKeyValue = valueConst; + + return FindShardInterval(valueConst->constvalue, tableEntry); +} + + +/* + * BuildFastPathTask builds the minimal Task for a single-shard execution, + * bypassing plan copying, coordinator evaluation and task regeneration. + * + * Returns NULL when the shard has no placement to run on, so the caller can + * fall back to normal planning and raise its "found no worker with all shard + * placements" error. + */ +static Task * +BuildFastPathTask(DistributedPlan *plan, Job *workerJob, ShardInterval *shardInterval, + Const *partitionKeyValue, bool isModify) +{ + List *shardIntervalListList = list_make1(list_make1(shardInterval)); + bool shardsPresent = false; + List *relationShardList = + RelationShardListForShardIntervalList(shardIntervalListList, &shardsPresent); + List *placementList = + CreateTaskPlacementListForShardIntervals(shardIntervalListList, shardsPresent, + true, false); + + if (placementList == NIL) + { + return NULL; + } + + /* + * Modifications run on every placement and are assigned first-replica by the + * caller, so only reads honour citus.task_assignment_policy. + */ + if (!isModify && TaskAssignmentPolicy == TASK_ASSIGNMENT_ROUND_ROBIN) + { + placementList = RemoveCoordinatorPlacementIfNotSingleNode(placementList); + placementList = RoundRobinReorder(placementList); + } + + Task *task = CitusMakeNode(Task); + task->taskType = isModify ? MODIFY_TASK : READ_TASK; + task->anchorShardId = shardInterval->shardId; + task->anchorDistributedTableId = linitial_oid(plan->relationIdList); + task->taskPlacementList = placementList; + task->queryCount = 1; + task->preparedStatementPlanId = plan->planId; + task->jobQueryForPrepare = workerJob->savedJobQueryForCaching; + task->relationShardList = relationShardList; + task->relationRowLockList = + RelationRowLockListForQuery(workerJob->savedJobQueryForCaching); + task->colocationId = workerJob->colocationId; + task->partitionKeyValue = partitionKeyValue; + + return task; +} + + +/* + * PreparedStatementCacheTryFastPath builds the task for this execution directly + * from the bound parameters, reusing the plan and query template saved on the + * first execution. + * + * Returns false when the fast path does not apply, in which case the caller + * must fall through to normal planning. + */ +bool +PreparedStatementCacheTryFastPath(struct CitusScanState *scanStateArg, EState *estate, + bool isModify) +{ + CitusScanState *scanState = (CitusScanState *) scanStateArg; + DistributedPlan *originalPlan = scanState->distributedPlan; + Job *workerJob = originalPlan->workerJob; + + /* the first execution populates the template the fast path depends on */ + if (!PreparedStatementCachingUsable() || + originalPlan->numberOfTimesExecuted == 0 || + workerJob->savedJobQueryForCaching == NULL) + { + return false; + } + + /* + * EXPLAIN ANALYZE wraps the task's query string, which would then carry the + * template's Params into the reported worker plan. + */ + if (RequestedForExplainAnalyze(scanState)) + { + return false; + } + + /* volatile functions must be re-evaluated per execution on the coordinator */ + if (isModify && + (!workerJob->deferredPruning || workerJob->requiresCoordinatorEvaluation)) + { + return false; + } + + Const *partitionKeyValue = NULL; + ShardInterval *shardInterval = FastPathShardInterval(originalPlan, workerJob, + estate, &partitionKeyValue); + if (shardInterval == NULL) + { + return false; + } + + Task *task = BuildFastPathTask(originalPlan, workerJob, shardInterval, + partitionKeyValue, isModify); + if (task == NULL) + { + return false; + } + + workerJob->taskList = list_make1(task); + workerJob->parametersInJobQueryResolved = true; + + /* local execution and query stats read the key from the job, not the task */ + workerJob->partitionKeyValue = partitionKeyValue; + + if (isModify) + { + AcquireMetadataLocks(workerJob->taskList); + + /* + * A concurrent split may have dropped the shard between pruning and + * locking. Normal planning reroutes here, but it does so by rewriting + * jobQuery in place, which would pin the cached plan to this execution's + * shards. Decline instead, so normal planning reroutes on its own copy. + */ + if (!ShardExists(shardInterval->shardId)) + { + workerJob->taskList = NIL; + workerJob->parametersInJobQueryResolved = false; + workerJob->partitionKeyValue = workerJob->plannerPartitionKeyValue; + return false; + } + + workerJob->taskList = FirstReplicaAssignTaskList(workerJob->taskList); + } + + elog(DEBUG2, "prepared statement cache-hit fast path%s: plan " UINT64_FORMAT + " shard " UINT64_FORMAT, + isModify ? " (DML)" : "", originalPlan->planId, shardInterval->shardId); + + /* the executor reads the plan back from the scan state */ + scanState->distributedPlan = originalPlan; + + /* + * A fast-path task has no query string, so local execution needs a cached + * local plan rather than deparsed SQL. + */ + if (IsLocalPlanCachingSupported(workerJob, originalPlan)) + { + CacheLocalPlanForShardQuery(linitial(workerJob->taskList), originalPlan, + estate->es_param_list_info); + } + + return true; +} diff --git a/src/backend/distributed/planner/deparse_shard_query.c b/src/backend/distributed/planner/deparse_shard_query.c index d38491baef2..9724d692bc0 100644 --- a/src/backend/distributed/planner/deparse_shard_query.c +++ b/src/backend/distributed/planner/deparse_shard_query.c @@ -790,6 +790,46 @@ TaskQueryString(Task *task) int taskQueryType = GetTaskQueryType(task); if (taskQueryType == TASK_QUERY_NULL) { + if (task->jobQueryForPrepare != NULL) + { + /* + * Fast-path task from prepared statement caching: deparse the + * query from the saved template and cache the result on the task + * so subsequent calls don't re-deparse. + */ + Query *queryForDeparse = copyObject(task->jobQueryForPrepare); + + /* + * The task retains the string, so allocate it in the task's own + * context rather than CurrentMemoryContext, which may be a + * short-lived per-tuple context. + */ + MemoryContext previousContext = + MemoryContextSwitchTo(GetMemoryChunkContext(task)); + StringInfoData buf; + initStringInfo(&buf); + MemoryContextSwitchTo(previousContext); + + if (queryForDeparse->commandType == CMD_INSERT) + { + /* upserts reference the target by name, which becomes the shard name */ + AddInsertAliasIfNeeded(queryForDeparse); + + deparse_shard_query(queryForDeparse, + task->anchorDistributedTableId, + task->anchorShardId, &buf); + } + else + { + UpdateRelationToShardNames((Node *) queryForDeparse, + task->relationShardList); + pg_get_query_def(queryForDeparse, &buf); + } + + SetTaskQueryString(task, buf.data); + return task->taskQuery.data.queryStringLazy; + } + /* if task query type is TASK_QUERY_NULL then the data will be NULL, * this is unexpected state */ ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), diff --git a/src/backend/distributed/planner/distributed_planner.c b/src/backend/distributed/planner/distributed_planner.c index 949bfbd7519..8633d1a5dab 100644 --- a/src/backend/distributed/planner/distributed_planner.c +++ b/src/backend/distributed/planner/distributed_planner.c @@ -1448,11 +1448,24 @@ GetDistributedPlan(CustomScan *customScan) Node *node = (Node *) linitial(customScan->custom_private); Assert(CitusIsA(node, DistributedPlan)); - CheckNodeCopyAndSerialization(node); + /* + * Undo what the previous execution's fast path left on the original plan. + * Its Task and partition key Const were allocated in a portal context that + * is now gone, so serializing the plan would follow dangling pointers. + * CitusEndScanCommon() does this too, but an execution that errors out + * never reaches it. + */ + DistributedPlan *plan = (DistributedPlan *) node; + if (plan->workerJob != NULL && plan->workerJob->savedJobQueryForCaching != NULL) + { + plan->workerJob->taskList = NIL; + plan->workerJob->parametersInJobQueryResolved = false; + plan->workerJob->partitionKeyValue = plan->workerJob->plannerPartitionKeyValue; + } - DistributedPlan *distributedPlan = (DistributedPlan *) node; + CheckNodeCopyAndSerialization(node); - return distributedPlan; + return plan; } @@ -2677,6 +2690,8 @@ CreateAndPushPlannerRestrictionContext(DistributedPlanningContext *planContext, fastPathRestrictionContext->distributionKeyValue; plannersFastPathCtx->distributionKeyHasParam = fastPathRestrictionContext->distributionKeyHasParam; + plannersFastPathCtx->distributionKeyParamId = + fastPathRestrictionContext->distributionKeyParamId; plannersFastPathCtx->delayFastPathPlanning = fastPathRestrictionContext->delayFastPathPlanning; } diff --git a/src/backend/distributed/planner/fast_path_router_planner.c b/src/backend/distributed/planner/fast_path_router_planner.c index 63c68f03af5..01493398092 100644 --- a/src/backend/distributed/planner/fast_path_router_planner.c +++ b/src/backend/distributed/planner/fast_path_router_planner.c @@ -180,7 +180,9 @@ InitializeFastPathContext(FastPathRestrictionContext *fastPathContext, } else if (IsA(distributionKeyValue, Param)) { + Param *distributionKeyParam = (Param *) distributionKeyValue; fastPathContext->distributionKeyHasParam = true; + fastPathContext->distributionKeyParamId = distributionKeyParam->paramid; } /* diff --git a/src/backend/distributed/planner/multi_explain.c b/src/backend/distributed/planner/multi_explain.c index 9c00f34c2fa..84a868e0cb1 100644 --- a/src/backend/distributed/planner/multi_explain.c +++ b/src/backend/distributed/planner/multi_explain.c @@ -1938,6 +1938,9 @@ ExplainAnalyzeTaskList(List *originalTaskList, SetTaskQueryStringList(explainAnalyzeTask, list_make2(wrappedQuery, fetchQuery)); + /* the wrapper replaces the query, so the saved template no longer applies */ + explainAnalyzeTask->jobQueryForPrepare = NULL; + TupleDestination *originalTaskDest = originalTask->tupleDest ? originalTask->tupleDest : defaultTupleDest; diff --git a/src/backend/distributed/planner/multi_router_planner.c b/src/backend/distributed/planner/multi_router_planner.c index 4bc038b3d4d..4c95e290088 100644 --- a/src/backend/distributed/planner/multi_router_planner.c +++ b/src/backend/distributed/planner/multi_router_planner.c @@ -1768,6 +1768,33 @@ RouterInsertJob(Query *originalQuery) job->deferredPruning = true; job->partitionKeyValue = ExtractInsertPartitionKeyValue(originalQuery); + /* + * For single-row INSERTs with a parameterized distribution key, + * capture the Param index so the executor fast path can extract + * the value from ParamListInfo at execution time. + */ + if (!isMultiRowInsert && job->partitionKeyValue == NULL && + !job->requiresCoordinatorEvaluation) + { + Oid distributedTableId = ExtractFirstCitusTableId(originalQuery); + if (HasDistributionKey(distributedTableId)) + { + Var *partitionColumn = PartitionColumn(distributedTableId, 1); + TargetEntry *targetEntry = get_tle_by_resno(originalQuery->targetList, + partitionColumn->varattno); + if (targetEntry != NULL) + { + Node *targetExpression = strip_implicit_coercions( + (Node *) targetEntry->expr); + if (IsA(targetExpression, Param)) + { + job->distributionKeyParamId = + ((Param *) targetExpression)->paramid; + } + } + } + } + return job; } @@ -1955,6 +1982,8 @@ RouterJob(Query *originalQuery, PlannerRestrictionContext *plannerRestrictionCon { Job *job = CreateJob(originalQuery); job->deferredPruning = true; + job->distributionKeyParamId = fastPathRestrictionContext->distributionKeyParamId; + job->requiresCoordinatorEvaluation = requiresCoordinatorEvaluation; ereport(DEBUG2, (errmsg("Deferred pruning for a fast-path router " "query"))); @@ -2721,6 +2750,21 @@ RowLocksOnRelations(Node *node, List **relationRowLockList) } +/* + * RelationRowLockListForQuery returns the FOR UPDATE/SHARE row locks the query + * takes on Citus tables. + */ +List * +RelationRowLockListForQuery(Query *query) +{ + List *relationRowLockList = NIL; + + RowLocksOnRelations((Node *) query, &relationRowLockList); + + return relationRowLockList; +} + + /* * SelectsFromDistributedTable checks if there is a select on a distributed * table by looking into range table entries. diff --git a/src/backend/distributed/shared_library_init.c b/src/backend/distributed/shared_library_init.c index 603f20f9384..bb7051d5b44 100644 --- a/src/backend/distributed/shared_library_init.c +++ b/src/backend/distributed/shared_library_init.c @@ -91,6 +91,7 @@ #include "distributed/multi_server_executor.h" #include "distributed/pg_dist_partition.h" #include "distributed/placement_connection.h" +#include "distributed/prepared_statement_cache.h" #include "distributed/priority.h" #include "distributed/procedure_body_analysis.h" #include "distributed/query_pushdown_planning.h" @@ -1602,6 +1603,18 @@ RegisterCitusConfigVariables(void) GUC_NO_SHOW_ALL | GUC_NOT_IN_SAMPLE, WarnIfDeprecatedPseudoconstantQualsGucIsSet, NULL, NULL); + DefineCustomBoolVariable( + "citus.enable_prepared_statement_caching", + gettext_noop("Enables caching prepared statement plans on worker connections."), + gettext_noop("When enabled, Citus uses PQprepare/PQsendQueryPrepared on " + "worker connections for fast-path prepared statement executions, " + "eliminating redundant parse/plan cycles on workers."), + &EnablePreparedStatementCaching, + false, + PGC_USERSET, + GUC_STANDARD, + NULL, NULL, NULL); + DefineCustomBoolVariable( "citus.enable_procedure_transaction_skip", gettext_noop("Skip coordinated transactions for single-statement, " diff --git a/src/backend/distributed/utils/citus_copyfuncs.c b/src/backend/distributed/utils/citus_copyfuncs.c index 58ef1697564..c8cb29ff3d2 100644 --- a/src/backend/distributed/utils/citus_copyfuncs.c +++ b/src/backend/distributed/utils/citus_copyfuncs.c @@ -99,9 +99,13 @@ copyJobInfo(Job *newnode, Job *from) COPY_SCALAR_FIELD(subqueryPushdown); COPY_SCALAR_FIELD(requiresCoordinatorEvaluation); COPY_SCALAR_FIELD(deferredPruning); + COPY_SCALAR_FIELD(distributionKeyParamId); COPY_NODE_FIELD(partitionKeyValue); COPY_NODE_FIELD(localPlannedStatements); COPY_SCALAR_FIELD(parametersInJobQueryResolved); + COPY_SCALAR_FIELD(colocationId); + COPY_NODE_FIELD(savedJobQueryForCaching); + COPY_NODE_FIELD(plannerPartitionKeyValue); } @@ -359,6 +363,8 @@ CopyNodeTask(COPYFUNC_ARGS) COPY_NODE_FIELD(rowValuesLists); COPY_SCALAR_FIELD(partiallyLocalOrRemote); COPY_SCALAR_FIELD(parametersInQueryStringResolved); + COPY_SCALAR_FIELD(preparedStatementPlanId); + COPY_NODE_FIELD(jobQueryForPrepare); COPY_SCALAR_FIELD(tupleDest); COPY_SCALAR_FIELD(queryCount); COPY_SCALAR_FIELD(totalReceivedTupleData); @@ -367,6 +373,8 @@ CopyNodeTask(COPYFUNC_ARGS) COPY_SCALAR_FIELD(fetchedExplainAnalyzeExecutionDuration); COPY_SCALAR_FIELD(isLocalTableModification); COPY_SCALAR_FIELD(cannotBeExecutedInTransaction); + COPY_NODE_FIELD(partitionKeyValue); + COPY_SCALAR_FIELD(colocationId); } diff --git a/src/backend/distributed/utils/citus_outfuncs.c b/src/backend/distributed/utils/citus_outfuncs.c index 3ddabbbca2e..d6b6051e4cd 100644 --- a/src/backend/distributed/utils/citus_outfuncs.c +++ b/src/backend/distributed/utils/citus_outfuncs.c @@ -393,9 +393,11 @@ OutJobFields(StringInfo str, const Job *node) WRITE_BOOL_FIELD(subqueryPushdown); WRITE_BOOL_FIELD(requiresCoordinatorEvaluation); WRITE_BOOL_FIELD(deferredPruning); + WRITE_INT_FIELD(distributionKeyParamId); WRITE_NODE_FIELD(partitionKeyValue); WRITE_NODE_FIELD(localPlannedStatements); WRITE_BOOL_FIELD(parametersInJobQueryResolved); + WRITE_UINT_FIELD(colocationId); } @@ -576,6 +578,7 @@ OutTask(OUTFUNC_ARGS) WRITE_NODE_FIELD(rowValuesLists); WRITE_BOOL_FIELD(partiallyLocalOrRemote); WRITE_BOOL_FIELD(parametersInQueryStringResolved); + WRITE_UINT64_FIELD(preparedStatementPlanId); WRITE_INT_FIELD(queryCount); WRITE_UINT64_FIELD(totalReceivedTupleData); WRITE_INT_FIELD(fetchedExplainAnalyzePlacementIndex); @@ -583,6 +586,7 @@ OutTask(OUTFUNC_ARGS) WRITE_FLOAT_FIELD(fetchedExplainAnalyzeExecutionDuration, "%.2f"); WRITE_BOOL_FIELD(isLocalTableModification); WRITE_BOOL_FIELD(cannotBeExecutedInTransaction); + WRITE_INT_FIELD(colocationId); } diff --git a/src/include/distributed/citus_custom_scan.h b/src/include/distributed/citus_custom_scan.h index 1c74da08756..5d73733ea3f 100644 --- a/src/include/distributed/citus_custom_scan.h +++ b/src/include/distributed/citus_custom_scan.h @@ -59,5 +59,6 @@ extern bool IsCitusPlan(Plan *plan); extern bool IsCitusCustomScan(Plan *plan); extern void SetJobColocationId(Job *job); +extern void EnsureAnchorShardsInJobExist(Job *job); #endif /* CITUS_CUSTOM_SCAN_H */ diff --git a/src/include/distributed/connection_management.h b/src/include/distributed/connection_management.h index 8032a0ea37b..aa8178f9d26 100644 --- a/src/include/distributed/connection_management.h +++ b/src/include/distributed/connection_management.h @@ -232,6 +232,9 @@ typedef struct MultiConnection bool requiresReplication; MultiConnectionStructInitializationState initializationState; + + /* per-connection cache of prepared statements on this worker connection */ + HTAB *preparedStatementCache; } MultiConnection; diff --git a/src/include/distributed/distributed_planner.h b/src/include/distributed/distributed_planner.h index f07ac78c2e6..92d915b5c29 100644 --- a/src/include/distributed/distributed_planner.h +++ b/src/include/distributed/distributed_planner.h @@ -112,6 +112,14 @@ typedef struct FastPathRestrictionContext */ bool distributionKeyHasParam; + /* + * When distributionKeyHasParam is true, stores the Param's paramid + * so the executor can extract the distribution key value directly + * from ParamListInfo without walking the query tree. 0 when not set + * (paramid is 1-based in PostgreSQL). + */ + int distributionKeyParamId; + /* * Indicates to hold off calling the fast path planner until its * known if the shard is local or not. diff --git a/src/include/distributed/multi_physical_planner.h b/src/include/distributed/multi_physical_planner.h index 977b51825cd..f899e0b16bd 100644 --- a/src/include/distributed/multi_physical_planner.h +++ b/src/include/distributed/multi_physical_planner.h @@ -141,6 +141,7 @@ typedef struct Job bool subqueryPushdown; bool requiresCoordinatorEvaluation; /* only applies to modify jobs */ bool deferredPruning; + int distributionKeyParamId; Const *partitionKeyValue; /* for local shard queries, we may save the local plan here */ @@ -153,6 +154,23 @@ typedef struct Job */ bool parametersInJobQueryResolved; uint32 colocationId; /* common colocation group ID of the relations */ + + /* + * Cached copy of jobQuery with Param nodes ($1, $2, ...) still intact, + * used for prepared statement caching. Populated lazily on first + * execution and reused across subsequent executions to avoid a + * per-execution copyObject(jobQuery). Lives on originalDistributedPlan + * only; the per-execution copy (currentPlan) gets a NULL here. + */ + Query *savedJobQueryForCaching; + + /* + * partitionKeyValue as the planner left it, captured alongside + * savedJobQueryForCaching. The cache fast path overwrites + * partitionKeyValue with a per-execution Const, and this is what it gets + * put back to once that Const is about to become dangling. + */ + Const *plannerPartitionKeyValue; } Job; @@ -311,6 +329,12 @@ typedef struct Task */ bool parametersInQueryStringResolved; + /* prepared statement cache key: plan ID from DistributedPlan */ + uint64 preparedStatementPlanId; + + /* pre-evaluation job query with Param nodes intact, for PQprepare on cache miss */ + Query *jobQueryForPrepare; + /* * Destination of tuples generated as a result of executing this task. Can be * NULL, in which case executor might use a default destination. diff --git a/src/include/distributed/multi_router_planner.h b/src/include/distributed/multi_router_planner.h index 434fd966fe1..9441e0e61a0 100644 --- a/src/include/distributed/multi_router_planner.h +++ b/src/include/distributed/multi_router_planner.h @@ -56,6 +56,7 @@ extern List * CreateTaskPlacementListForShardIntervals(List *shardIntervalList, bool hasLocalRelation); extern List * RouterInsertTaskList(Query *query, bool parametersInQueryResolved, DeferredErrorMessage **planningError); +extern List * RelationRowLockListForQuery(Query *query); extern Const * ExtractInsertPartitionKeyValue(Query *query); extern List * TargetShardIntervalsForRestrictInfo(RelationRestrictionContext * restrictionContext, diff --git a/src/include/distributed/prepared_statement_cache.h b/src/include/distributed/prepared_statement_cache.h new file mode 100644 index 00000000000..789ecc97bb7 --- /dev/null +++ b/src/include/distributed/prepared_statement_cache.h @@ -0,0 +1,100 @@ +/*------------------------------------------------------------------------- + * prepared_statement_cache.h + * + * Declarations for per-connection prepared statement caching on worker + * connections. + * + * Copyright (c) Citus Data, Inc. + * + *------------------------------------------------------------------------- + */ + +#ifndef PREPARED_STATEMENT_CACHE_H +#define PREPARED_STATEMENT_CACHE_H + +#include "postgres.h" + +#include "nodes/execnodes.h" +#include "nodes/parsenodes.h" +#include "utils/hsearch.h" + +#include "distributed/connection_management.h" +#include "distributed/multi_physical_planner.h" + +struct CitusScanState; + +/* compile-time limit for per-connection cached prepared statements */ +#define MAX_CACHED_STMTS_PER_CONNECTION 1000 + +/* maximum length for generated statement names ("__citus_stmt_NNNN") */ +#define MAX_STMT_NAME_LENGTH 64 + + +/* + * PreparedStatementCacheKey uniquely identifies a prepared statement on a + * given worker connection: planId identifies the cached generic plan, + * shardId identifies the target shard. + */ +typedef struct PreparedStatementCacheKey +{ + uint64 planId; + uint64 shardId; +} PreparedStatementCacheKey; + + +/* + * PreparedStatementCacheEntry names the statement prepared on a connection + * for a given (planId, shardId). + */ +typedef struct PreparedStatementCacheEntry +{ + PreparedStatementCacheKey key; + + char stmtName[MAX_STMT_NAME_LENGTH]; +} PreparedStatementCacheEntry; + + +/* + * Outcome of an attempt to dispatch a task through the connection's + * prepared statement cache. + */ +typedef enum PreparedStatementSendStatus +{ + PREPARED_STMT_NOT_APPLICABLE, /* caching off or task carries no template */ + PREPARED_STMT_SENT, /* dispatched on the connection */ + PREPARED_STMT_FAILED, /* connection lost */ + PREPARED_STMT_FALLBACK /* cache full; caller sends the returned SQL */ +} PreparedStatementSendStatus; + + +/* GUC variable */ +extern bool EnablePreparedStatementCaching; + +/* cache lifecycle */ +extern HTAB * PreparedStatementCacheCreate(void); +extern void PreparedStatementCacheDestroy(HTAB **cache_ptr); + +/* cache operations */ +extern PreparedStatementCacheEntry * PreparedStatementCacheLookup(HTAB *cache, uint64 + planId, uint64 shardId); +extern PreparedStatementCacheEntry * PreparedStatementCacheInsert(HTAB *cache, uint64 + planId, uint64 shardId, + const char *stmtName); + +/* planner-side integration (see citus_custom_scan.c) */ +extern bool PreparedStatementCacheTryFastPath(struct CitusScanState *scanState, + EState *estate, bool isModify); +extern Query * PreparedStatementCacheSaveTemplate(DistributedPlan *originalPlan); +extern void PreparedStatementCacheAttachToTasks(DistributedPlan *currentPlan, + Job *workerJob, Query *savedJobQuery); + +/* executor-side integration (see adaptive_executor.c) */ +extern PreparedStatementSendStatus PreparedStatementCacheSendQuery(MultiConnection * + connection, Task *task, + ParamListInfo + paramListInfo, + bool binaryResults, + char ** + fallbackQueryString); + +#endif /* PREPARED_STATEMENT_CACHE_H */ diff --git a/src/include/distributed/remote_commands.h b/src/include/distributed/remote_commands.h index 71cb9dad27f..88662503767 100644 --- a/src/include/distributed/remote_commands.h +++ b/src/include/distributed/remote_commands.h @@ -68,6 +68,11 @@ extern bool PutRemoteCopyEnd(MultiConnection *connection, const char *errormsg); /* waiting for multiple command results */ extern void WaitForAllConnections(List *connectionList, bool raiseInterrupts); +extern int SendRemotePrepare(MultiConnection *connection, const char *stmtName, + const char *query, int nParams, const Oid *paramTypes); +extern int SendRemotePreparedQuery(MultiConnection *connection, const char *stmtName, + int nParams, const char *const *paramValues, + bool binaryResults); extern bool SendCancelationRequest(MultiConnection *connection); extern bool EvaluateSingleQueryResult(MultiConnection *connection, PGresult *queryResult, diff --git a/src/test/regress/bin/normalize.sed b/src/test/regress/bin/normalize.sed index 34762fe1538..90b19035320 100644 --- a/src/test/regress/bin/normalize.sed +++ b/src/test/regress/bin/normalize.sed @@ -88,6 +88,9 @@ s/(NOTICE: [a-z]+ cascades to table ".*)_[0-9]{5,}"/\1_xxxxx"/g # connection id s/connectionId: [0-9]+/connectionId: xxxxxxx/g +# cached prepared statement names are numbered per connection +s/issuing __citus_stmt_[0-9]+/issuing __citus_stmt_xxx/g + # Remove trailing whitespace s/ *$//g diff --git a/src/test/regress/citus_tests/config.py b/src/test/regress/citus_tests/config.py index 1d3d122be4f..dd1684f0382 100644 --- a/src/test/regress/citus_tests/config.py +++ b/src/test/regress/citus_tests/config.py @@ -296,6 +296,14 @@ def __init__(self, arguments): } +class CitusPreparedStatementCachingConfig(CitusDefaultClusterConfig): + def __init__(self, arguments): + super().__init__(arguments) + self.new_settings = { + "citus.enable_prepared_statement_caching": True, + } + + class CitusUnusualExecutorConfig(CitusDefaultClusterConfig): def __init__(self, arguments): super().__init__(arguments) diff --git a/src/test/regress/citus_tests/run_test.py b/src/test/regress/citus_tests/run_test.py index 77c13287557..bf64c16eb9b 100755 --- a/src/test/regress/citus_tests/run_test.py +++ b/src/test/regress/citus_tests/run_test.py @@ -337,6 +337,9 @@ def extra_tests(self): "multi_subquery_window_functions": TestDeps( "minimal_schedule", ["multi_behavioral_analytics_create_table"] ), + "subquery_prepared_statements": TestDeps( + "minimal_schedule", ["multi_behavioral_analytics_create_table"] + ), } diff --git a/src/test/regress/expected/prepared_statement_caching.out b/src/test/regress/expected/prepared_statement_caching.out new file mode 100644 index 00000000000..65d21e51cb2 --- /dev/null +++ b/src/test/regress/expected/prepared_statement_caching.out @@ -0,0 +1,1293 @@ +-- +-- PREPARED_STATEMENT_CACHING +-- +-- Tests for citus.enable_prepared_statement_caching, which enables +-- worker-side prepared statement plan caching for fast-path queries. +-- +CREATE SCHEMA prepared_stmt_caching; +SET search_path TO prepared_stmt_caching; +-- Test 8 prints shard names in EXPLAIN output, so pin the shard ids +SET citus.next_shard_id TO 105000; +SET citus.shard_count TO 4; +SET citus.shard_replication_factor TO 1; +-- the regression suite runs with citus.stat_tenants_track = 'ALL', which the +-- cache declines in order to keep per-execution tenant attribution +SET citus.stat_tenants_track TO 'none'; +-- Create test tables +CREATE TABLE dist_table ( + key int PRIMARY KEY, + value int, + label text +); +SELECT create_distributed_table('dist_table', 'key'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +-- Insert base data +INSERT INTO dist_table SELECT i, i * 10, 'label-' || i FROM generate_series(1, 20) i; +CREATE TABLE dist_table_ts ( + key int PRIMARY KEY, + value int, + created_at timestamptz DEFAULT now() +); +SELECT create_distributed_table('dist_table_ts', 'key'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +-- ============================================================ +-- Test 1: GUC toggle — verify default is OFF, SET to ON succeeds +-- ============================================================ +SHOW citus.enable_prepared_statement_caching; + citus.enable_prepared_statement_caching +--------------------------------------------------------------------- + off +(1 row) + +SET citus.enable_prepared_statement_caching = on; +SHOW citus.enable_prepared_statement_caching; + citus.enable_prepared_statement_caching +--------------------------------------------------------------------- + on +(1 row) + +SET citus.enable_prepared_statement_caching = off; +SHOW citus.enable_prepared_statement_caching; + citus.enable_prepared_statement_caching +--------------------------------------------------------------------- + off +(1 row) + +-- ============================================================ +-- Test 2: Basic caching — PREPARE a single-shard SELECT, EXECUTE +-- 10 times with GUC ON, verify correct results +-- ============================================================ +SET citus.enable_prepared_statement_caching = on; +set citus.max_cached_connection_lifetime to '60min'; +SET search_path TO prepared_stmt_caching; +PREPARE cached_select(int) AS + SELECT key, value FROM dist_table WHERE key = $1; +-- Execute 10 times to ensure generic plan path and cache hit path +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +DEALLOCATE cached_select; +-- ============================================================ +-- Test 3: Multi-shard-value — EXECUTE with different partition +-- key values routing to different shards +-- ============================================================ +PREPARE cached_multi_shard(int) AS + SELECT key, value FROM dist_table WHERE key = $1; +-- Different keys likely route to different shards +EXECUTE cached_multi_shard(1); + key | value +--------------------------------------------------------------------- + 1 | 10 +(1 row) + +EXECUTE cached_multi_shard(2); + key | value +--------------------------------------------------------------------- + 2 | 20 +(1 row) + +EXECUTE cached_multi_shard(3); + key | value +--------------------------------------------------------------------- + 3 | 30 +(1 row) + +EXECUTE cached_multi_shard(4); + key | value +--------------------------------------------------------------------- + 4 | 40 +(1 row) + +EXECUTE cached_multi_shard(5); + key | value +--------------------------------------------------------------------- + 5 | 50 +(1 row) + +EXECUTE cached_multi_shard(6); + key | value +--------------------------------------------------------------------- + 6 | 60 +(1 row) + +EXECUTE cached_multi_shard(7); + key | value +--------------------------------------------------------------------- + 7 | 70 +(1 row) + +EXECUTE cached_multi_shard(8); + key | value +--------------------------------------------------------------------- + 8 | 80 +(1 row) + +EXECUTE cached_multi_shard(9); + key | value +--------------------------------------------------------------------- + 9 | 90 +(1 row) + +EXECUTE cached_multi_shard(10); + key | value +--------------------------------------------------------------------- + 10 | 100 +(1 row) + +DEALLOCATE cached_multi_shard; +-- ============================================================ +-- Test 4: INSERT/UPDATE/DELETE with caching ON +-- Include now() to verify coordinator-side function +-- evaluation still works. +-- Single-row INSERT, UPDATE, and DELETE all use +-- the cached prepared statement path. +-- ============================================================ +-- INSERT (cached via deparse_shard_query path) +PREPARE cached_insert(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2); +EXECUTE cached_insert(100, 1000); +EXECUTE cached_insert(101, 1010); +EXECUTE cached_insert(102, 1020); +EXECUTE cached_insert(103, 1030); +EXECUTE cached_insert(104, 1040); +EXECUTE cached_insert(105, 1050); +EXECUTE cached_insert(106, 1060); +EXECUTE cached_insert(107, 1070); +EXECUTE cached_insert(108, 1080); +EXECUTE cached_insert(109, 1090); +EXECUTE cached_insert(110, 1100); +-- Verify inserts +SELECT key, value FROM dist_table_ts WHERE key >= 100 ORDER BY key; + key | value +--------------------------------------------------------------------- + 100 | 1000 + 101 | 1010 + 102 | 1020 + 103 | 1030 + 104 | 1040 + 105 | 1050 + 106 | 1060 + 107 | 1070 + 108 | 1080 + 109 | 1090 + 110 | 1100 +(11 rows) + +DEALLOCATE cached_insert; +-- UPDATE +PREPARE cached_update(int, int) AS + UPDATE dist_table SET value = $2 WHERE key = $1; +EXECUTE cached_update(1, 100); +EXECUTE cached_update(2, 200); +EXECUTE cached_update(3, 300); +EXECUTE cached_update(4, 400); +EXECUTE cached_update(5, 500); +EXECUTE cached_update(6, 600); +EXECUTE cached_update(7, 700); +SELECT key, value FROM dist_table WHERE key <= 7 ORDER BY key; + key | value +--------------------------------------------------------------------- + 1 | 100 + 2 | 200 + 3 | 300 + 4 | 400 + 5 | 500 + 6 | 600 + 7 | 700 +(7 rows) + +DEALLOCATE cached_update; +-- DELETE +PREPARE cached_delete(int) AS + DELETE FROM dist_table WHERE key = $1; +EXECUTE cached_delete(18); +EXECUTE cached_delete(19); +EXECUTE cached_delete(20); +EXECUTE cached_delete(18); +EXECUTE cached_delete(19); +EXECUTE cached_delete(20); +EXECUTE cached_delete(18); +-- Verify deletes +SELECT count(*) FROM dist_table WHERE key >= 18; + count +--------------------------------------------------------------------- + 0 +(1 row) + +DEALLOCATE cached_delete; +-- INSERT with now() function evaluation +PREPARE cached_insert_ts(int) AS + INSERT INTO dist_table_ts (key, value, created_at) VALUES ($1, $1 * 10, now()); +EXECUTE cached_insert_ts(200); +EXECUTE cached_insert_ts(201); +EXECUTE cached_insert_ts(202); +EXECUTE cached_insert_ts(203); +EXECUTE cached_insert_ts(204); +EXECUTE cached_insert_ts(205); +EXECUTE cached_insert_ts(206); +-- Verify that each row has created_at populated (functions were evaluated) +SELECT key, value, created_at IS NOT NULL AS has_ts FROM dist_table_ts + WHERE key >= 200 ORDER BY key; + key | value | has_ts +--------------------------------------------------------------------- + 200 | 2000 | t + 201 | 2010 | t + 202 | 2020 | t + 203 | 2030 | t + 204 | 2040 | t + 205 | 2050 | t + 206 | 2060 | t +(7 rows) + +DEALLOCATE cached_insert_ts; +-- INSERT ... ON CONFLICT DO UPDATE (upsert). The qualified reference to the +-- target table in DO UPDATE must resolve to the shard alias on the worker. +PREPARE cached_upsert(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2) + ON CONFLICT (key) DO UPDATE SET value = dist_table_ts.value + EXCLUDED.value; +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(501, 5); +EXECUTE cached_upsert(501, 5); +-- key 500 accumulated 6 increments, key 501 accumulated 2 +SELECT key, value FROM dist_table_ts WHERE key >= 500 AND key < 600 ORDER BY key; + key | value +--------------------------------------------------------------------- + 500 | 6 + 501 | 10 +(2 rows) + +DEALLOCATE cached_upsert; +-- Multi-row INSERT is not cacheable: each shard's task carries only its own +-- subset of VALUES rows, so the rows must not be duplicated across shards. +PREPARE cached_multirow(int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, 1), ($1 + 1, 2); +EXECUTE cached_multirow(600); +EXECUTE cached_multirow(602); +EXECUTE cached_multirow(604); +EXECUTE cached_multirow(606); +EXECUTE cached_multirow(608); +EXECUTE cached_multirow(610); +-- exactly 12 rows, one per key, no duplicates +SELECT count(*) AS row_count, count(DISTINCT key) AS distinct_keys + FROM dist_table_ts WHERE key >= 600; + row_count | distinct_keys +--------------------------------------------------------------------- + 12 | 12 +(1 row) + +DEALLOCATE cached_multirow; +-- ============================================================ +-- Test 5: GUC OFF baseline — same queries produce identical results +-- ============================================================ +SET citus.enable_prepared_statement_caching = off; +PREPARE uncached_select(int) AS + SELECT key, value FROM dist_table WHERE key = $1; +EXECUTE uncached_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE uncached_select(2); + key | value +--------------------------------------------------------------------- + 2 | 200 +(1 row) + +EXECUTE uncached_select(3); + key | value +--------------------------------------------------------------------- + 3 | 300 +(1 row) + +EXECUTE uncached_select(4); + key | value +--------------------------------------------------------------------- + 4 | 400 +(1 row) + +EXECUTE uncached_select(5); + key | value +--------------------------------------------------------------------- + 5 | 500 +(1 row) + +EXECUTE uncached_select(6); + key | value +--------------------------------------------------------------------- + 6 | 600 +(1 row) + +EXECUTE uncached_select(7); + key | value +--------------------------------------------------------------------- + 7 | 700 +(1 row) + +EXECUTE uncached_select(8); + key | value +--------------------------------------------------------------------- + 8 | 80 +(1 row) + +EXECUTE uncached_select(9); + key | value +--------------------------------------------------------------------- + 9 | 90 +(1 row) + +EXECUTE uncached_select(10); + key | value +--------------------------------------------------------------------- + 10 | 100 +(1 row) + +DEALLOCATE uncached_select; +-- ============================================================ +-- Test 6: Multiple prepared statements in same session +-- ============================================================ +SET citus.enable_prepared_statement_caching = on; +PREPARE stmt_a(int) AS SELECT key, value FROM dist_table WHERE key = $1; +PREPARE stmt_b(int) AS SELECT key, label FROM dist_table WHERE key = $1; +PREPARE stmt_c(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2); +-- Interleave executions to verify independent caching +EXECUTE stmt_a(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE stmt_b(1); + key | label +--------------------------------------------------------------------- + 1 | label-1 +(1 row) + +EXECUTE stmt_a(2); + key | value +--------------------------------------------------------------------- + 2 | 200 +(1 row) + +EXECUTE stmt_b(2); + key | label +--------------------------------------------------------------------- + 2 | label-2 +(1 row) + +EXECUTE stmt_a(3); + key | value +--------------------------------------------------------------------- + 3 | 300 +(1 row) + +EXECUTE stmt_b(3); + key | label +--------------------------------------------------------------------- + 3 | label-3 +(1 row) + +EXECUTE stmt_a(4); + key | value +--------------------------------------------------------------------- + 4 | 400 +(1 row) + +EXECUTE stmt_b(4); + key | label +--------------------------------------------------------------------- + 4 | label-4 +(1 row) + +EXECUTE stmt_a(5); + key | value +--------------------------------------------------------------------- + 5 | 500 +(1 row) + +EXECUTE stmt_b(5); + key | label +--------------------------------------------------------------------- + 5 | label-5 +(1 row) + +EXECUTE stmt_a(6); + key | value +--------------------------------------------------------------------- + 6 | 600 +(1 row) + +EXECUTE stmt_b(6); + key | label +--------------------------------------------------------------------- + 6 | label-6 +(1 row) + +EXECUTE stmt_a(7); + key | value +--------------------------------------------------------------------- + 7 | 700 +(1 row) + +EXECUTE stmt_b(7); + key | label +--------------------------------------------------------------------- + 7 | label-7 +(1 row) + +EXECUTE stmt_c(300, 3000); +EXECUTE stmt_c(301, 3010); +EXECUTE stmt_c(302, 3020); +EXECUTE stmt_c(303, 3030); +EXECUTE stmt_c(304, 3040); +EXECUTE stmt_c(305, 3050); +EXECUTE stmt_c(306, 3060); +-- Verify inserts from stmt_c +SELECT key, value FROM dist_table_ts WHERE key >= 300 AND key < 400 ORDER BY key; + key | value +--------------------------------------------------------------------- + 300 | 3000 + 301 | 3010 + 302 | 3020 + 303 | 3030 + 304 | 3040 + 305 | 3050 + 306 | 3060 +(7 rows) + +DEALLOCATE stmt_a; +DEALLOCATE stmt_b; +DEALLOCATE stmt_c; +-- ============================================================ +-- Test 7: Connection loss re-prepare — force worker connection +-- close, verify the statement is re-prepared on the new +-- connection. Results alone cannot show this (see Test 13), +-- so assert the wire protocol: the new connection must be +-- sent the parameterized SQL before a statement name. +-- ============================================================ +SET citus.enable_prepared_statement_caching = on; +PREPARE reconnect_test(int) AS + SELECT key, value FROM dist_table WHERE key = $1; +-- Execute enough times to get into generic plan + cache hit +EXECUTE reconnect_test(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE reconnect_test(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE reconnect_test(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE reconnect_test(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE reconnect_test(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE reconnect_test(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE reconnect_test(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +-- Force all cached connections to be dropped by setting lifetime to 0 +SET citus.max_cached_connection_lifetime TO '0s'; +-- The lifetime is only applied when the connection is released at end of +-- transaction, so the first execution below still reuses the prepared +-- connection and the second lands on a fresh one. Both use the same key, +-- so a re-prepare can only be caused by the new connection. +SET citus.log_remote_commands TO on; +EXECUTE reconnect_test(1); +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE reconnect_test(1); +NOTICE: issuing SELECT key, value FROM prepared_stmt_caching.dist_table_105000 dist_table WHERE (key OPERATOR(pg_catalog.=) $1) +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +SET citus.log_remote_commands TO off; +EXECUTE reconnect_test(2); + key | value +--------------------------------------------------------------------- + 2 | 200 +(1 row) + +EXECUTE reconnect_test(3); + key | value +--------------------------------------------------------------------- + 3 | 300 +(1 row) + +-- Restore default +RESET citus.max_cached_connection_lifetime; +DEALLOCATE reconnect_test; +-- ============================================================ +-- Test 8: EXPLAIN a cached fast-path statement. EXPLAIN builds the +-- shard query from the saved template rather than from the +-- executor's cache path, and the result is retained on the +-- task, so it must outlive the deparse. +-- ============================================================ +SET citus.enable_prepared_statement_caching = on; +PREPARE explain_select(int) AS + SELECT key, value FROM dist_table WHERE key = $1; +-- reach the generic plan and the cache-hit fast path first +EXECUTE explain_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE explain_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE explain_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE explain_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE explain_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE explain_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE explain_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXPLAIN (COSTS OFF) EXECUTE explain_select(1); + QUERY PLAN +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Index Scan using dist_table_pkey_105000 on dist_table_105000 dist_table + Index Cond: (key = 1) +(7 rows) + +-- the statement must still execute correctly afterwards +EXECUTE explain_select(1); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +DEALLOCATE explain_select; +-- INSERT takes the other deparse branch +PREPARE explain_insert(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2); +EXECUTE explain_insert(700, 7000); +EXECUTE explain_insert(701, 7010); +EXECUTE explain_insert(702, 7020); +EXECUTE explain_insert(703, 7030); +EXECUTE explain_insert(704, 7040); +EXECUTE explain_insert(705, 7050); +EXECUTE explain_insert(706, 7060); +EXPLAIN (COSTS OFF) EXECUTE explain_insert(707, 7070); + QUERY PLAN +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on dist_table_ts_105005 + -> Result +(7 rows) + +-- EXPLAIN without ANALYZE must not have inserted key 707 +SELECT count(*) AS inserted_rows FROM dist_table_ts WHERE key >= 700 AND key < 800; + inserted_rows +--------------------------------------------------------------------- + 7 +(1 row) + +DEALLOCATE explain_insert; +-- An upsert is what actually needs the insert alias: the DO UPDATE reference to +-- the target table has to resolve against the shard alias, not the shard name. +PREPARE explain_upsert(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2) + ON CONFLICT (key) DO UPDATE SET value = dist_table_ts.value + EXCLUDED.value; +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXPLAIN (COSTS OFF) EXECUTE explain_upsert(800, 1); + QUERY PLAN +--------------------------------------------------------------------- + Custom Scan (Citus Adaptive) + Task Count: 1 + Tasks Shown: All + -> Task + Node: host=localhost port=xxxxx dbname=regression + -> Insert on dist_table_ts_105005 citus_table_alias + Conflict Resolution: UPDATE + Conflict Arbiter Indexes: dist_table_ts_pkey_105005 + -> Result +(9 rows) + +-- seven increments, and EXPLAIN must not have added an eighth +SELECT key, value FROM dist_table_ts WHERE key = 800; + key | value +--------------------------------------------------------------------- + 800 | 7 +(1 row) + +DEALLOCATE explain_upsert; +-- ============================================================ +-- Test 9: citus.task_assignment_policy must still apply on the +-- cache-hit fast path, which builds its task directly +-- instead of going through GenerateSingleShardRouterTaskList(). +-- ============================================================ +-- returns 'shardId@port' for the placement the task was assigned to +CREATE OR REPLACE FUNCTION parse_explain_output(in qry text, in table_name text, out r text) +RETURNS SETOF TEXT AS $$ +DECLARE + portOfTheTask text; + shardOfTheTask text; +begin + for r in execute qry loop + IF r LIKE '%port%' THEN + portOfTheTask = substring(r, '([0-9]{1,10})'); + END IF; + + IF r LIKE '%' || table_name || '%' THEN + shardOfTheTask = substring(r, '([0-9]{5,10})'); + END IF; + + end loop; + return QUERY SELECT shardOfTheTask || '@' || portOfTheTask; +end; $$ language plpgsql; +-- round-robin only has something to choose between when shards are replicated +SET citus.shard_replication_factor TO 2; +CREATE TABLE replicated_table (key int PRIMARY KEY, value int); +SELECT create_distributed_table('replicated_table', 'key'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +SET citus.shard_replication_factor TO 1; +INSERT INTO replicated_table SELECT i, i * 10 FROM generate_series(1, 20) i; +SET citus.task_assignment_policy TO 'round-robin'; +SET citus.explain_distributed_queries TO on; +PREPARE round_robin_select(int) AS + SELECT value FROM replicated_table WHERE key = $1; +-- reach the generic plan so that later executions take the fast path +EXECUTE round_robin_select(1); + value +--------------------------------------------------------------------- + 10 +(1 row) + +EXECUTE round_robin_select(1); + value +--------------------------------------------------------------------- + 10 +(1 row) + +EXECUTE round_robin_select(1); + value +--------------------------------------------------------------------- + 10 +(1 row) + +EXECUTE round_robin_select(1); + value +--------------------------------------------------------------------- + 10 +(1 row) + +EXECUTE round_robin_select(1); + value +--------------------------------------------------------------------- + 10 +(1 row) + +EXECUTE round_robin_select(1); + value +--------------------------------------------------------------------- + 10 +(1 row) + +EXECUTE round_robin_select(1); + value +--------------------------------------------------------------------- + 10 +(1 row) + +CREATE TEMPORARY TABLE explain_outputs (value text); +INSERT INTO explain_outputs + SELECT parse_explain_output('EXPLAIN EXECUTE round_robin_select(1)', 'replicated_table'); +INSERT INTO explain_outputs + SELECT parse_explain_output('EXPLAIN EXECUTE round_robin_select(1)', 'replicated_table'); +-- outside a transaction round-robin must alternate placements, so the fast +-- path has to reach both nodes rather than pinning to the first placement +SELECT count(DISTINCT value) FROM explain_outputs; + count +--------------------------------------------------------------------- + 2 +(1 row) + +DROP TABLE explain_outputs; +DEALLOCATE round_robin_select; +RESET citus.task_assignment_policy; +RESET citus.explain_distributed_queries; +-- ============================================================ +-- Test 10: the distribution key parameter need not have the same type +-- as the distribution column. A single-row INSERT records the +-- Param with implicit coercions stripped, so $1 here is int4 +-- while the column is numeric. +-- ============================================================ +CREATE TABLE numeric_dist (key numeric PRIMARY KEY, value int); +SELECT create_distributed_table('numeric_dist', 'key'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +PREPARE numeric_insert(int) AS + INSERT INTO numeric_dist (key, value) VALUES ($1, $1 * 10); +EXECUTE numeric_insert(1); +EXECUTE numeric_insert(2); +EXECUTE numeric_insert(3); +EXECUTE numeric_insert(4); +EXECUTE numeric_insert(5); +EXECUTE numeric_insert(6); +EXECUTE numeric_insert(7); +EXECUTE numeric_insert(8); +EXECUTE numeric_insert(9); +EXECUTE numeric_insert(10); +SELECT count(*) AS total FROM numeric_dist; + total +--------------------------------------------------------------------- + 10 +(1 row) + +-- rows written on the fast path must be reachable by a router lookup, which +-- only holds if the value was hashed as numeric rather than as int4 +SELECT key, value FROM numeric_dist WHERE key = 8; + key | value +--------------------------------------------------------------------- + 8 | 80 +(1 row) + +SELECT key, value FROM numeric_dist WHERE key = 9; + key | value +--------------------------------------------------------------------- + 9 | 90 +(1 row) + +SELECT key, value FROM numeric_dist WHERE key = 10; + key | value +--------------------------------------------------------------------- + 10 | 100 +(1 row) + +DEALLOCATE numeric_insert; +-- ============================================================ +-- Test 11: DML whose expressions must be evaluated on the coordinator +-- cannot be cached. The template is copied before evaluation, +-- so a cached statement would evaluate nextval() on the worker, +-- where the sequence does not exist. +-- ============================================================ +CREATE SEQUENCE coord_eval_seq; +CREATE TABLE coord_eval (key int PRIMARY KEY, seq_value bigint); +SELECT create_distributed_table('coord_eval', 'key'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +PREPARE coord_eval_insert(int) AS + INSERT INTO coord_eval (key, seq_value) VALUES ($1, nextval('coord_eval_seq')); +EXECUTE coord_eval_insert(1); +EXECUTE coord_eval_insert(2); +EXECUTE coord_eval_insert(3); +EXECUTE coord_eval_insert(4); +EXECUTE coord_eval_insert(5); +EXECUTE coord_eval_insert(6); +EXECUTE coord_eval_insert(7); +EXECUTE coord_eval_insert(8); +-- values come from the coordinator's sequence, in execution order +SELECT key, seq_value FROM coord_eval ORDER BY key; + key | seq_value +--------------------------------------------------------------------- + 1 | 1 + 2 | 2 + 3 | 3 + 4 | 4 + 5 | 5 + 6 | 6 + 7 | 7 + 8 | 8 +(8 rows) + +DEALLOCATE coord_eval_insert; +-- UPDATE reaches the deferred-pruning router path, which must carry the same +-- coordinator-evaluation flag +ALTER TABLE coord_eval ADD COLUMN ts timestamptz; +PREPARE coord_eval_update(int) AS + UPDATE coord_eval SET ts = now() WHERE key = $1; +BEGIN; +EXECUTE coord_eval_update(1); +EXECUTE coord_eval_update(2); +EXECUTE coord_eval_update(3); +EXECUTE coord_eval_update(4); +EXECUTE coord_eval_update(5); +EXECUTE coord_eval_update(6); +EXECUTE coord_eval_update(7); +EXECUTE coord_eval_update(8); +-- a coordinator-evaluated now() is one value for the whole transaction; +-- evaluated on the workers it would be one value per node +SELECT count(DISTINCT ts) AS distinct_timestamps FROM coord_eval; + distinct_timestamps +--------------------------------------------------------------------- + 1 +(1 row) + +COMMIT; +DEALLOCATE coord_eval_update; +-- ============================================================ +-- Test 12: a prepared statement may carry parameters the query never +-- uses. Their types still have to be normalized, or the worker +-- cannot infer a type for a parameter absent from the SQL. +-- ============================================================ +CREATE TYPE unused_param_type AS (a int); +PREPARE unused_param(int, unused_param_type) AS + SELECT key, value FROM dist_table WHERE key = $1; +EXECUTE unused_param(1, '(1)'); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE unused_param(1, '(1)'); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE unused_param(1, '(1)'); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE unused_param(1, '(1)'); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE unused_param(1, '(1)'); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE unused_param(1, '(1)'); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +EXECUTE unused_param(1, '(1)'); + key | value +--------------------------------------------------------------------- + 1 | 100 +(1 row) + +DEALLOCATE unused_param; +-- ============================================================ +-- Test 13: assert the cache is actually engaged. Every other test +-- here checks only query results, which are identical +-- whether or not the feature is active, so observe the +-- wire protocol instead: PQprepare logs the parameterized +-- SQL, PQsendQueryPrepared logs the statement name. +-- ============================================================ +CREATE TABLE observe_cache (key int PRIMARY KEY, value int); +SELECT create_distributed_table('observe_cache', 'key'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +INSERT INTO observe_cache SELECT i, i FROM generate_series(1, 20) i; +PREPARE observe_select(int) AS + SELECT value FROM observe_cache WHERE key = $1; +-- reach the generic plan before enabling logging; the cache is populated on +-- the second use of the generic plan, so this needs one more than the +-- five executions it takes to get there +EXECUTE observe_select(1); + value +--------------------------------------------------------------------- + 1 +(1 row) + +EXECUTE observe_select(1); + value +--------------------------------------------------------------------- + 1 +(1 row) + +EXECUTE observe_select(1); + value +--------------------------------------------------------------------- + 1 +(1 row) + +EXECUTE observe_select(1); + value +--------------------------------------------------------------------- + 1 +(1 row) + +EXECUTE observe_select(1); + value +--------------------------------------------------------------------- + 1 +(1 row) + +EXECUTE observe_select(1); + value +--------------------------------------------------------------------- + 1 +(1 row) + +EXECUTE observe_select(1); + value +--------------------------------------------------------------------- + 1 +(1 row) + +SET citus.log_remote_commands TO on; +-- key 1's shard is already prepared on this connection, so this reuses +-- the statement by name rather than sending SQL +EXECUTE observe_select(1); +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx + value +--------------------------------------------------------------------- + 1 +(1 row) + +-- an unprepared shard prepares first, then executes by name +EXECUTE observe_select(3); +NOTICE: issuing SELECT value FROM prepared_stmt_caching.observe_cache_105021 observe_cache WHERE (key OPERATOR(pg_catalog.=) $1) +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx + value +--------------------------------------------------------------------- + 3 +(1 row) + +EXECUTE observe_select(3); +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx + value +--------------------------------------------------------------------- + 3 +(1 row) + +SET citus.log_remote_commands TO off; +DEALLOCATE observe_select; +-- ============================================================ +-- Test 14: SELECT ... FOR UPDATE is fast-path eligible, so the +-- cached task must carry its row locks. Without them the +-- executor treats it as an ordinary read and, with +-- select_opens_transaction_block off, opens no remote +-- transaction — releasing the row lock at statement end +-- rather than at the end of the enclosing transaction. +-- ============================================================ +SET citus.select_opens_transaction_block TO off; +PREPARE lock_row(int) AS + SELECT value FROM observe_cache WHERE key = $1 FOR UPDATE; +-- reach the generic plan before enabling logging +EXECUTE lock_row(5); + value +--------------------------------------------------------------------- + 5 +(1 row) + +EXECUTE lock_row(5); + value +--------------------------------------------------------------------- + 5 +(1 row) + +EXECUTE lock_row(5); + value +--------------------------------------------------------------------- + 5 +(1 row) + +EXECUTE lock_row(5); + value +--------------------------------------------------------------------- + 5 +(1 row) + +EXECUTE lock_row(5); + value +--------------------------------------------------------------------- + 5 +(1 row) + +EXECUTE lock_row(5); + value +--------------------------------------------------------------------- + 5 +(1 row) + +EXECUTE lock_row(5); + value +--------------------------------------------------------------------- + 5 +(1 row) + +-- the worker must still be given a transaction to hold the row lock in +SET citus.log_remote_commands TO on; +BEGIN; +EXECUTE lock_row(5); +NOTICE: issuing BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;SELECT assign_distributed_transaction_id(xx, xx, 'xxxxxxx'); +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx + value +--------------------------------------------------------------------- + 5 +(1 row) + +COMMIT; +NOTICE: issuing COMMIT +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +SET citus.log_remote_commands TO off; +RESET citus.select_opens_transaction_block; +DEALLOCATE lock_row; +-- ============================================================ +-- Test 15: the same wire-protocol assertion for the remote DML +-- paths. INSERT and UPDATE/DELETE reach the cache through +-- their own eligibility and task-building branches, and a +-- silent fallback in either is invisible in query results. +-- ============================================================ +-- no primary key, so the same shard can be hit repeatedly by INSERT +CREATE TABLE observe_dml (key int, value int); +SELECT create_distributed_table('observe_dml', 'key'); + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +PREPARE observe_insert(int, int) AS + INSERT INTO observe_dml (key, value) VALUES ($1, $2); +-- reach the generic plan and populate the cache before enabling logging +EXECUTE observe_insert(1, 1); +EXECUTE observe_insert(1, 2); +EXECUTE observe_insert(1, 3); +EXECUTE observe_insert(1, 4); +EXECUTE observe_insert(1, 5); +EXECUTE observe_insert(1, 6); +EXECUTE observe_insert(1, 7); +SET citus.log_remote_commands TO on; +-- key 1's shard is already prepared on this connection, so this reuses +-- the statement by name rather than sending SQL +EXECUTE observe_insert(1, 8); +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +-- an unprepared shard prepares first, then executes by name +EXECUTE observe_insert(3, 9); +NOTICE: issuing INSERT INTO prepared_stmt_caching.observe_dml_105025 (key, value) VALUES ($1, $2) +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +EXECUTE observe_insert(3, 10); +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +SET citus.log_remote_commands TO off; +PREPARE observe_update(int, int) AS + UPDATE observe_dml SET value = $2 WHERE key = $1; +EXECUTE observe_update(1, 11); +EXECUTE observe_update(1, 12); +EXECUTE observe_update(1, 13); +EXECUTE observe_update(1, 14); +EXECUTE observe_update(1, 15); +EXECUTE observe_update(1, 16); +EXECUTE observe_update(1, 17); +SET citus.log_remote_commands TO on; +EXECUTE observe_update(1, 18); +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +SET citus.log_remote_commands TO off; +PREPARE observe_delete(int) AS + DELETE FROM observe_dml WHERE key = $1; +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +SET citus.log_remote_commands TO on; +EXECUTE observe_delete(3); +NOTICE: issuing __citus_stmt_xxx +DETAIL: on server postgres@localhost:xxxxx connectionId: xxxxxxx +SET citus.log_remote_commands TO off; +SELECT key, value FROM observe_dml ORDER BY key, value; + key | value +--------------------------------------------------------------------- + 1 | 18 + 1 | 18 + 1 | 18 + 1 | 18 + 1 | 18 + 1 | 18 + 1 | 18 + 1 | 18 +(8 rows) + +DEALLOCATE observe_insert; +DEALLOCATE observe_update; +DEALLOCATE observe_delete; +-- ============================================================ +-- Cleanup +-- ============================================================ +SET citus.enable_prepared_statement_caching = off; +DROP SCHEMA prepared_stmt_caching CASCADE; +NOTICE: drop cascades to 10 other objects +DETAIL: drop cascades to table dist_table +drop cascades to table dist_table_ts +drop cascades to function parse_explain_output(text,text) +drop cascades to table replicated_table +drop cascades to table numeric_dist +drop cascades to sequence coord_eval_seq +drop cascades to table coord_eval +drop cascades to type unused_param_type +drop cascades to table observe_cache +drop cascades to table observe_dml diff --git a/src/test/regress/expected/single_node.out b/src/test/regress/expected/single_node.out index bdd89459875..e0038152510 100644 --- a/src/test/regress/expected/single_node.out +++ b/src/test/regress/expected/single_node.out @@ -1103,7 +1103,7 @@ ROLLBACK; -- explain analyze should work on a single node EXPLAIN (COSTS FALSE, ANALYZE TRUE, TIMING FALSE, SUMMARY FALSE, BUFFERS OFF) SELECT * FROM test; - QUERY PLAN + QUERY PLAN --------------------------------------------------------------------- Custom Scan (Citus Adaptive) (actual rows=5 loops=1) Task Count: 4 @@ -2629,6 +2629,84 @@ NOTICE: executing the command locally: UPDATE single_node.another_schema_table_ (1 row) ROLLBACK; +-- A cached fast-path upsert executed locally is deparsed for local planning, +-- and needs the insert alias so the DO UPDATE reference to the target table +-- resolves against the shard alias. +SET citus.shard_replication_factor TO 1; +SET citus.shard_count TO 4; +SET citus.enable_prepared_statement_caching TO on; +-- the suite runs with citus.stat_tenants_track = 'ALL', which the cache declines +SET citus.stat_tenants_track TO 'none'; +CREATE TABLE single_node.upsert_cache_test (key int PRIMARY KEY, value int); +SELECT create_distributed_table('single_node.upsert_cache_test', 'key'); +NOTICE: executing the command locally: SELECT worker_apply_shard_ddl_command (90830508, 'single_node', 'CREATE TABLE single_node.upsert_cache_test (key integer NOT NULL, value integer) USING heap');SELECT worker_apply_shard_ddl_command (90830508, 'single_node', 'ALTER TABLE single_node.upsert_cache_test OWNER TO postgres');SELECT worker_apply_shard_ddl_command (90830508, 'single_node', 'ALTER TABLE single_node.upsert_cache_test ADD CONSTRAINT upsert_cache_test_pkey PRIMARY KEY (key)') +NOTICE: executing the command locally: SELECT worker_apply_shard_ddl_command (90830509, 'single_node', 'CREATE TABLE single_node.upsert_cache_test (key integer NOT NULL, value integer) USING heap');SELECT worker_apply_shard_ddl_command (90830509, 'single_node', 'ALTER TABLE single_node.upsert_cache_test OWNER TO postgres');SELECT worker_apply_shard_ddl_command (90830509, 'single_node', 'ALTER TABLE single_node.upsert_cache_test ADD CONSTRAINT upsert_cache_test_pkey PRIMARY KEY (key)') +NOTICE: executing the command locally: SELECT worker_apply_shard_ddl_command (90830510, 'single_node', 'CREATE TABLE single_node.upsert_cache_test (key integer NOT NULL, value integer) USING heap');SELECT worker_apply_shard_ddl_command (90830510, 'single_node', 'ALTER TABLE single_node.upsert_cache_test OWNER TO postgres');SELECT worker_apply_shard_ddl_command (90830510, 'single_node', 'ALTER TABLE single_node.upsert_cache_test ADD CONSTRAINT upsert_cache_test_pkey PRIMARY KEY (key)') +NOTICE: executing the command locally: SELECT worker_apply_shard_ddl_command (90830511, 'single_node', 'CREATE TABLE single_node.upsert_cache_test (key integer NOT NULL, value integer) USING heap');SELECT worker_apply_shard_ddl_command (90830511, 'single_node', 'ALTER TABLE single_node.upsert_cache_test OWNER TO postgres');SELECT worker_apply_shard_ddl_command (90830511, 'single_node', 'ALTER TABLE single_node.upsert_cache_test ADD CONSTRAINT upsert_cache_test_pkey PRIMARY KEY (key)') + create_distributed_table +--------------------------------------------------------------------- + +(1 row) + +PREPARE local_cached_upsert(int, int) AS + INSERT INTO single_node.upsert_cache_test (key, value) VALUES ($1, $2) + ON CONFLICT (key) DO UPDATE + SET value = upsert_cache_test.value + EXCLUDED.value; +-- with local command logging on, LogLocalCommand() deparses the task first, so +-- the notice below shows the shard SQL that the alias has to appear in +EXECUTE local_cached_upsert(1, 10); +NOTICE: executing the command locally: INSERT INTO single_node.upsert_cache_test_90830508 AS citus_table_alias (key, value) VALUES (1, 10) ON CONFLICT(key) DO UPDATE SET value = (citus_table_alias.value OPERATOR(pg_catalog.+) excluded.value) +EXECUTE local_cached_upsert(1, 10); +NOTICE: executing the command locally: INSERT INTO single_node.upsert_cache_test_90830508 AS citus_table_alias (key, value) VALUES (1, 10) ON CONFLICT(key) DO UPDATE SET value = (citus_table_alias.value OPERATOR(pg_catalog.+) excluded.value) +EXECUTE local_cached_upsert(1, 10); +NOTICE: executing the command locally: INSERT INTO single_node.upsert_cache_test_90830508 AS citus_table_alias (key, value) VALUES (1, 10) ON CONFLICT(key) DO UPDATE SET value = (citus_table_alias.value OPERATOR(pg_catalog.+) excluded.value) +EXECUTE local_cached_upsert(1, 10); +NOTICE: executing the command locally: INSERT INTO single_node.upsert_cache_test_90830508 AS citus_table_alias (key, value) VALUES (1, 10) ON CONFLICT(key) DO UPDATE SET value = (citus_table_alias.value OPERATOR(pg_catalog.+) excluded.value) +EXECUTE local_cached_upsert(1, 10); +NOTICE: executing the command locally: INSERT INTO single_node.upsert_cache_test_90830508 AS citus_table_alias (key, value) VALUES (1, 10) ON CONFLICT(key) DO UPDATE SET value = (citus_table_alias.value OPERATOR(pg_catalog.+) excluded.value) +EXECUTE local_cached_upsert(1, 10); +NOTICE: executing the command locally: INSERT INTO single_node.upsert_cache_test_90830508 AS citus_table_alias (key, value) VALUES (1, 10) ON CONFLICT(key) DO UPDATE SET value = (citus_table_alias.value OPERATOR(pg_catalog.+) excluded.value) +EXECUTE local_cached_upsert(1, 10); +NOTICE: executing the command locally: INSERT INTO single_node.upsert_cache_test_90830508 AS citus_table_alias (key, value) VALUES ($1, $2) ON CONFLICT(key) DO UPDATE SET value = (citus_table_alias.value OPERATOR(pg_catalog.+) excluded.value) +EXECUTE local_cached_upsert(2, 5); +NOTICE: executing the command locally: INSERT INTO single_node.upsert_cache_test_90830511 AS citus_table_alias (key, value) VALUES ($1, $2) ON CONFLICT(key) DO UPDATE SET value = (citus_table_alias.value OPERATOR(pg_catalog.+) excluded.value) +-- key 1 accumulated seven increments, key 2 was inserted once +SELECT key, value FROM single_node.upsert_cache_test ORDER BY key; +NOTICE: executing the command locally: SELECT key, value FROM single_node.upsert_cache_test_90830508 upsert_cache_test WHERE true ORDER BY key +NOTICE: executing the command locally: SELECT key, value FROM single_node.upsert_cache_test_90830509 upsert_cache_test WHERE true ORDER BY key +NOTICE: executing the command locally: SELECT key, value FROM single_node.upsert_cache_test_90830510 upsert_cache_test WHERE true ORDER BY key +NOTICE: executing the command locally: SELECT key, value FROM single_node.upsert_cache_test_90830511 upsert_cache_test WHERE true ORDER BY key + key | value +--------------------------------------------------------------------- + 1 | 70 + 2 | 5 +(2 rows) + +-- with logging off nothing deparses the task up front, so the local executor +-- reaches its own fast-path deparse instead +SET citus.log_local_commands TO off; +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(4, 7); +SELECT key, value FROM single_node.upsert_cache_test WHERE key >= 3 ORDER BY key; + key | value +--------------------------------------------------------------------- + 3 | 700 + 4 | 7 +(2 rows) + +RESET citus.log_local_commands; +DEALLOCATE local_cached_upsert; +DROP TABLE single_node.upsert_cache_test; +RESET citus.enable_prepared_statement_caching; +RESET citus.stat_tenants_track; +RESET citus.shard_count; +RESET citus.shard_replication_factor; -- if the local execution is disabled, we cannot failover to -- local execution and the queries would fail SET citus.enable_local_execution TO false; diff --git a/src/test/regress/multi_schedule b/src/test/regress/multi_schedule index 0960831df83..24fbaa169db 100644 --- a/src/test/regress/multi_schedule +++ b/src/test/regress/multi_schedule @@ -61,6 +61,7 @@ test: set_operation_and_local_tables test: subqueries_deep subquery_view subquery_partitioning subqueries_not_supported test: subquery_in_targetlist subquery_in_where subquery_complex_target_list subquery_append test: subquery_prepared_statements +test: prepared_statement_caching test: non_colocated_leaf_subquery_joins non_colocated_subquery_joins test: cte_inline recursive_view_local_table values sequences_with_different_types multi_level_recursive_queries test: pg13 pg12 diff --git a/src/test/regress/sql/prepared_statement_caching.sql b/src/test/regress/sql/prepared_statement_caching.sql new file mode 100644 index 00000000000..85bc1cd14d7 --- /dev/null +++ b/src/test/regress/sql/prepared_statement_caching.sql @@ -0,0 +1,709 @@ +-- +-- PREPARED_STATEMENT_CACHING +-- +-- Tests for citus.enable_prepared_statement_caching, which enables +-- worker-side prepared statement plan caching for fast-path queries. +-- + +CREATE SCHEMA prepared_stmt_caching; +SET search_path TO prepared_stmt_caching; + +-- Test 8 prints shard names in EXPLAIN output, so pin the shard ids +SET citus.next_shard_id TO 105000; +SET citus.shard_count TO 4; +SET citus.shard_replication_factor TO 1; + +-- the regression suite runs with citus.stat_tenants_track = 'ALL', which the +-- cache declines in order to keep per-execution tenant attribution +SET citus.stat_tenants_track TO 'none'; + +-- Create test tables +CREATE TABLE dist_table ( + key int PRIMARY KEY, + value int, + label text +); +SELECT create_distributed_table('dist_table', 'key'); + +-- Insert base data +INSERT INTO dist_table SELECT i, i * 10, 'label-' || i FROM generate_series(1, 20) i; + +CREATE TABLE dist_table_ts ( + key int PRIMARY KEY, + value int, + created_at timestamptz DEFAULT now() +); +SELECT create_distributed_table('dist_table_ts', 'key'); + +-- ============================================================ +-- Test 1: GUC toggle — verify default is OFF, SET to ON succeeds +-- ============================================================ + +SHOW citus.enable_prepared_statement_caching; +SET citus.enable_prepared_statement_caching = on; +SHOW citus.enable_prepared_statement_caching; +SET citus.enable_prepared_statement_caching = off; +SHOW citus.enable_prepared_statement_caching; + +-- ============================================================ +-- Test 2: Basic caching — PREPARE a single-shard SELECT, EXECUTE +-- 10 times with GUC ON, verify correct results +-- ============================================================ + +SET citus.enable_prepared_statement_caching = on; +set citus.max_cached_connection_lifetime to '60min'; +SET search_path TO prepared_stmt_caching; + +PREPARE cached_select(int) AS + SELECT key, value FROM dist_table WHERE key = $1; + +-- Execute 10 times to ensure generic plan path and cache hit path +EXECUTE cached_select(1); +EXECUTE cached_select(1); +EXECUTE cached_select(1); +EXECUTE cached_select(1); +EXECUTE cached_select(1); +EXECUTE cached_select(1); +EXECUTE cached_select(1); +EXECUTE cached_select(1); +EXECUTE cached_select(1); +EXECUTE cached_select(1); + +DEALLOCATE cached_select; + +-- ============================================================ +-- Test 3: Multi-shard-value — EXECUTE with different partition +-- key values routing to different shards +-- ============================================================ + +PREPARE cached_multi_shard(int) AS + SELECT key, value FROM dist_table WHERE key = $1; + +-- Different keys likely route to different shards +EXECUTE cached_multi_shard(1); +EXECUTE cached_multi_shard(2); +EXECUTE cached_multi_shard(3); +EXECUTE cached_multi_shard(4); +EXECUTE cached_multi_shard(5); +EXECUTE cached_multi_shard(6); +EXECUTE cached_multi_shard(7); +EXECUTE cached_multi_shard(8); +EXECUTE cached_multi_shard(9); +EXECUTE cached_multi_shard(10); + +DEALLOCATE cached_multi_shard; + +-- ============================================================ +-- Test 4: INSERT/UPDATE/DELETE with caching ON +-- Include now() to verify coordinator-side function +-- evaluation still works. +-- Single-row INSERT, UPDATE, and DELETE all use +-- the cached prepared statement path. +-- ============================================================ + +-- INSERT (cached via deparse_shard_query path) +PREPARE cached_insert(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2); + +EXECUTE cached_insert(100, 1000); +EXECUTE cached_insert(101, 1010); +EXECUTE cached_insert(102, 1020); +EXECUTE cached_insert(103, 1030); +EXECUTE cached_insert(104, 1040); +EXECUTE cached_insert(105, 1050); +EXECUTE cached_insert(106, 1060); + +EXECUTE cached_insert(107, 1070); +EXECUTE cached_insert(108, 1080); +EXECUTE cached_insert(109, 1090); +EXECUTE cached_insert(110, 1100); + +-- Verify inserts +SELECT key, value FROM dist_table_ts WHERE key >= 100 ORDER BY key; + +DEALLOCATE cached_insert; + +-- UPDATE +PREPARE cached_update(int, int) AS + UPDATE dist_table SET value = $2 WHERE key = $1; + +EXECUTE cached_update(1, 100); +EXECUTE cached_update(2, 200); +EXECUTE cached_update(3, 300); +EXECUTE cached_update(4, 400); +EXECUTE cached_update(5, 500); +EXECUTE cached_update(6, 600); +EXECUTE cached_update(7, 700); + +SELECT key, value FROM dist_table WHERE key <= 7 ORDER BY key; + +DEALLOCATE cached_update; + +-- DELETE +PREPARE cached_delete(int) AS + DELETE FROM dist_table WHERE key = $1; + +EXECUTE cached_delete(18); +EXECUTE cached_delete(19); +EXECUTE cached_delete(20); +EXECUTE cached_delete(18); +EXECUTE cached_delete(19); +EXECUTE cached_delete(20); +EXECUTE cached_delete(18); + +-- Verify deletes +SELECT count(*) FROM dist_table WHERE key >= 18; + +DEALLOCATE cached_delete; + +-- INSERT with now() function evaluation +PREPARE cached_insert_ts(int) AS + INSERT INTO dist_table_ts (key, value, created_at) VALUES ($1, $1 * 10, now()); + +EXECUTE cached_insert_ts(200); +EXECUTE cached_insert_ts(201); +EXECUTE cached_insert_ts(202); +EXECUTE cached_insert_ts(203); +EXECUTE cached_insert_ts(204); +EXECUTE cached_insert_ts(205); +EXECUTE cached_insert_ts(206); + +-- Verify that each row has created_at populated (functions were evaluated) +SELECT key, value, created_at IS NOT NULL AS has_ts FROM dist_table_ts + WHERE key >= 200 ORDER BY key; + +DEALLOCATE cached_insert_ts; + +-- INSERT ... ON CONFLICT DO UPDATE (upsert). The qualified reference to the +-- target table in DO UPDATE must resolve to the shard alias on the worker. +PREPARE cached_upsert(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2) + ON CONFLICT (key) DO UPDATE SET value = dist_table_ts.value + EXCLUDED.value; + +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(500, 1); +EXECUTE cached_upsert(501, 5); +EXECUTE cached_upsert(501, 5); + +-- key 500 accumulated 6 increments, key 501 accumulated 2 +SELECT key, value FROM dist_table_ts WHERE key >= 500 AND key < 600 ORDER BY key; + +DEALLOCATE cached_upsert; + +-- Multi-row INSERT is not cacheable: each shard's task carries only its own +-- subset of VALUES rows, so the rows must not be duplicated across shards. +PREPARE cached_multirow(int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, 1), ($1 + 1, 2); + +EXECUTE cached_multirow(600); +EXECUTE cached_multirow(602); +EXECUTE cached_multirow(604); +EXECUTE cached_multirow(606); +EXECUTE cached_multirow(608); +EXECUTE cached_multirow(610); + +-- exactly 12 rows, one per key, no duplicates +SELECT count(*) AS row_count, count(DISTINCT key) AS distinct_keys + FROM dist_table_ts WHERE key >= 600; + +DEALLOCATE cached_multirow; + +-- ============================================================ +-- Test 5: GUC OFF baseline — same queries produce identical results +-- ============================================================ + +SET citus.enable_prepared_statement_caching = off; + +PREPARE uncached_select(int) AS + SELECT key, value FROM dist_table WHERE key = $1; + +EXECUTE uncached_select(1); +EXECUTE uncached_select(2); +EXECUTE uncached_select(3); +EXECUTE uncached_select(4); +EXECUTE uncached_select(5); +EXECUTE uncached_select(6); +EXECUTE uncached_select(7); +EXECUTE uncached_select(8); +EXECUTE uncached_select(9); +EXECUTE uncached_select(10); + +DEALLOCATE uncached_select; + +-- ============================================================ +-- Test 6: Multiple prepared statements in same session +-- ============================================================ + +SET citus.enable_prepared_statement_caching = on; + +PREPARE stmt_a(int) AS SELECT key, value FROM dist_table WHERE key = $1; +PREPARE stmt_b(int) AS SELECT key, label FROM dist_table WHERE key = $1; +PREPARE stmt_c(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2); + +-- Interleave executions to verify independent caching +EXECUTE stmt_a(1); +EXECUTE stmt_b(1); +EXECUTE stmt_a(2); +EXECUTE stmt_b(2); +EXECUTE stmt_a(3); +EXECUTE stmt_b(3); +EXECUTE stmt_a(4); +EXECUTE stmt_b(4); +EXECUTE stmt_a(5); +EXECUTE stmt_b(5); +EXECUTE stmt_a(6); +EXECUTE stmt_b(6); +EXECUTE stmt_a(7); +EXECUTE stmt_b(7); +EXECUTE stmt_c(300, 3000); +EXECUTE stmt_c(301, 3010); +EXECUTE stmt_c(302, 3020); +EXECUTE stmt_c(303, 3030); +EXECUTE stmt_c(304, 3040); +EXECUTE stmt_c(305, 3050); +EXECUTE stmt_c(306, 3060); + +-- Verify inserts from stmt_c +SELECT key, value FROM dist_table_ts WHERE key >= 300 AND key < 400 ORDER BY key; + +DEALLOCATE stmt_a; +DEALLOCATE stmt_b; +DEALLOCATE stmt_c; + +-- ============================================================ +-- Test 7: Connection loss re-prepare — force worker connection +-- close, verify the statement is re-prepared on the new +-- connection. Results alone cannot show this (see Test 13), +-- so assert the wire protocol: the new connection must be +-- sent the parameterized SQL before a statement name. +-- ============================================================ + +SET citus.enable_prepared_statement_caching = on; + +PREPARE reconnect_test(int) AS + SELECT key, value FROM dist_table WHERE key = $1; + +-- Execute enough times to get into generic plan + cache hit +EXECUTE reconnect_test(1); +EXECUTE reconnect_test(1); +EXECUTE reconnect_test(1); +EXECUTE reconnect_test(1); +EXECUTE reconnect_test(1); +EXECUTE reconnect_test(1); +EXECUTE reconnect_test(1); + +-- Force all cached connections to be dropped by setting lifetime to 0 +SET citus.max_cached_connection_lifetime TO '0s'; + +-- The lifetime is only applied when the connection is released at end of +-- transaction, so the first execution below still reuses the prepared +-- connection and the second lands on a fresh one. Both use the same key, +-- so a re-prepare can only be caused by the new connection. +SET citus.log_remote_commands TO on; +EXECUTE reconnect_test(1); +EXECUTE reconnect_test(1); +SET citus.log_remote_commands TO off; + +EXECUTE reconnect_test(2); +EXECUTE reconnect_test(3); + +-- Restore default +RESET citus.max_cached_connection_lifetime; + +DEALLOCATE reconnect_test; + +-- ============================================================ +-- Test 8: EXPLAIN a cached fast-path statement. EXPLAIN builds the +-- shard query from the saved template rather than from the +-- executor's cache path, and the result is retained on the +-- task, so it must outlive the deparse. +-- ============================================================ + +SET citus.enable_prepared_statement_caching = on; + +PREPARE explain_select(int) AS + SELECT key, value FROM dist_table WHERE key = $1; + +-- reach the generic plan and the cache-hit fast path first +EXECUTE explain_select(1); +EXECUTE explain_select(1); +EXECUTE explain_select(1); +EXECUTE explain_select(1); +EXECUTE explain_select(1); +EXECUTE explain_select(1); +EXECUTE explain_select(1); + +EXPLAIN (COSTS OFF) EXECUTE explain_select(1); + +-- the statement must still execute correctly afterwards +EXECUTE explain_select(1); + +DEALLOCATE explain_select; + +-- INSERT takes the other deparse branch +PREPARE explain_insert(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2); + +EXECUTE explain_insert(700, 7000); +EXECUTE explain_insert(701, 7010); +EXECUTE explain_insert(702, 7020); +EXECUTE explain_insert(703, 7030); +EXECUTE explain_insert(704, 7040); +EXECUTE explain_insert(705, 7050); +EXECUTE explain_insert(706, 7060); + +EXPLAIN (COSTS OFF) EXECUTE explain_insert(707, 7070); + +-- EXPLAIN without ANALYZE must not have inserted key 707 +SELECT count(*) AS inserted_rows FROM dist_table_ts WHERE key >= 700 AND key < 800; + +DEALLOCATE explain_insert; + +-- An upsert is what actually needs the insert alias: the DO UPDATE reference to +-- the target table has to resolve against the shard alias, not the shard name. +PREPARE explain_upsert(int, int) AS + INSERT INTO dist_table_ts (key, value) VALUES ($1, $2) + ON CONFLICT (key) DO UPDATE SET value = dist_table_ts.value + EXCLUDED.value; + +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); +EXECUTE explain_upsert(800, 1); + +EXPLAIN (COSTS OFF) EXECUTE explain_upsert(800, 1); + +-- seven increments, and EXPLAIN must not have added an eighth +SELECT key, value FROM dist_table_ts WHERE key = 800; + +DEALLOCATE explain_upsert; + +-- ============================================================ +-- Test 9: citus.task_assignment_policy must still apply on the +-- cache-hit fast path, which builds its task directly +-- instead of going through GenerateSingleShardRouterTaskList(). +-- ============================================================ + +-- returns 'shardId@port' for the placement the task was assigned to +CREATE OR REPLACE FUNCTION parse_explain_output(in qry text, in table_name text, out r text) +RETURNS SETOF TEXT AS $$ +DECLARE + portOfTheTask text; + shardOfTheTask text; +begin + for r in execute qry loop + IF r LIKE '%port%' THEN + portOfTheTask = substring(r, '([0-9]{1,10})'); + END IF; + + IF r LIKE '%' || table_name || '%' THEN + shardOfTheTask = substring(r, '([0-9]{5,10})'); + END IF; + + end loop; + return QUERY SELECT shardOfTheTask || '@' || portOfTheTask; +end; $$ language plpgsql; + +-- round-robin only has something to choose between when shards are replicated +SET citus.shard_replication_factor TO 2; +CREATE TABLE replicated_table (key int PRIMARY KEY, value int); +SELECT create_distributed_table('replicated_table', 'key'); +SET citus.shard_replication_factor TO 1; + +INSERT INTO replicated_table SELECT i, i * 10 FROM generate_series(1, 20) i; + +SET citus.task_assignment_policy TO 'round-robin'; +SET citus.explain_distributed_queries TO on; + +PREPARE round_robin_select(int) AS + SELECT value FROM replicated_table WHERE key = $1; + +-- reach the generic plan so that later executions take the fast path +EXECUTE round_robin_select(1); +EXECUTE round_robin_select(1); +EXECUTE round_robin_select(1); +EXECUTE round_robin_select(1); +EXECUTE round_robin_select(1); +EXECUTE round_robin_select(1); +EXECUTE round_robin_select(1); + +CREATE TEMPORARY TABLE explain_outputs (value text); + +INSERT INTO explain_outputs + SELECT parse_explain_output('EXPLAIN EXECUTE round_robin_select(1)', 'replicated_table'); +INSERT INTO explain_outputs + SELECT parse_explain_output('EXPLAIN EXECUTE round_robin_select(1)', 'replicated_table'); + +-- outside a transaction round-robin must alternate placements, so the fast +-- path has to reach both nodes rather than pinning to the first placement +SELECT count(DISTINCT value) FROM explain_outputs; + +DROP TABLE explain_outputs; +DEALLOCATE round_robin_select; +RESET citus.task_assignment_policy; +RESET citus.explain_distributed_queries; + +-- ============================================================ +-- Test 10: the distribution key parameter need not have the same type +-- as the distribution column. A single-row INSERT records the +-- Param with implicit coercions stripped, so $1 here is int4 +-- while the column is numeric. +-- ============================================================ + +CREATE TABLE numeric_dist (key numeric PRIMARY KEY, value int); +SELECT create_distributed_table('numeric_dist', 'key'); + +PREPARE numeric_insert(int) AS + INSERT INTO numeric_dist (key, value) VALUES ($1, $1 * 10); + +EXECUTE numeric_insert(1); +EXECUTE numeric_insert(2); +EXECUTE numeric_insert(3); +EXECUTE numeric_insert(4); +EXECUTE numeric_insert(5); +EXECUTE numeric_insert(6); +EXECUTE numeric_insert(7); +EXECUTE numeric_insert(8); +EXECUTE numeric_insert(9); +EXECUTE numeric_insert(10); + +SELECT count(*) AS total FROM numeric_dist; + +-- rows written on the fast path must be reachable by a router lookup, which +-- only holds if the value was hashed as numeric rather than as int4 +SELECT key, value FROM numeric_dist WHERE key = 8; +SELECT key, value FROM numeric_dist WHERE key = 9; +SELECT key, value FROM numeric_dist WHERE key = 10; + +DEALLOCATE numeric_insert; + +-- ============================================================ +-- Test 11: DML whose expressions must be evaluated on the coordinator +-- cannot be cached. The template is copied before evaluation, +-- so a cached statement would evaluate nextval() on the worker, +-- where the sequence does not exist. +-- ============================================================ + +CREATE SEQUENCE coord_eval_seq; +CREATE TABLE coord_eval (key int PRIMARY KEY, seq_value bigint); +SELECT create_distributed_table('coord_eval', 'key'); + +PREPARE coord_eval_insert(int) AS + INSERT INTO coord_eval (key, seq_value) VALUES ($1, nextval('coord_eval_seq')); + +EXECUTE coord_eval_insert(1); +EXECUTE coord_eval_insert(2); +EXECUTE coord_eval_insert(3); +EXECUTE coord_eval_insert(4); +EXECUTE coord_eval_insert(5); +EXECUTE coord_eval_insert(6); +EXECUTE coord_eval_insert(7); +EXECUTE coord_eval_insert(8); + +-- values come from the coordinator's sequence, in execution order +SELECT key, seq_value FROM coord_eval ORDER BY key; + +DEALLOCATE coord_eval_insert; + +-- UPDATE reaches the deferred-pruning router path, which must carry the same +-- coordinator-evaluation flag +ALTER TABLE coord_eval ADD COLUMN ts timestamptz; + +PREPARE coord_eval_update(int) AS + UPDATE coord_eval SET ts = now() WHERE key = $1; + +BEGIN; +EXECUTE coord_eval_update(1); +EXECUTE coord_eval_update(2); +EXECUTE coord_eval_update(3); +EXECUTE coord_eval_update(4); +EXECUTE coord_eval_update(5); +EXECUTE coord_eval_update(6); +EXECUTE coord_eval_update(7); +EXECUTE coord_eval_update(8); + +-- a coordinator-evaluated now() is one value for the whole transaction; +-- evaluated on the workers it would be one value per node +SELECT count(DISTINCT ts) AS distinct_timestamps FROM coord_eval; +COMMIT; + +DEALLOCATE coord_eval_update; + +-- ============================================================ +-- Test 12: a prepared statement may carry parameters the query never +-- uses. Their types still have to be normalized, or the worker +-- cannot infer a type for a parameter absent from the SQL. +-- ============================================================ + +CREATE TYPE unused_param_type AS (a int); + +PREPARE unused_param(int, unused_param_type) AS + SELECT key, value FROM dist_table WHERE key = $1; + +EXECUTE unused_param(1, '(1)'); +EXECUTE unused_param(1, '(1)'); +EXECUTE unused_param(1, '(1)'); +EXECUTE unused_param(1, '(1)'); +EXECUTE unused_param(1, '(1)'); +EXECUTE unused_param(1, '(1)'); +EXECUTE unused_param(1, '(1)'); + +DEALLOCATE unused_param; + +-- ============================================================ +-- Test 13: assert the cache is actually engaged. Every other test +-- here checks only query results, which are identical +-- whether or not the feature is active, so observe the +-- wire protocol instead: PQprepare logs the parameterized +-- SQL, PQsendQueryPrepared logs the statement name. +-- ============================================================ + +CREATE TABLE observe_cache (key int PRIMARY KEY, value int); +SELECT create_distributed_table('observe_cache', 'key'); +INSERT INTO observe_cache SELECT i, i FROM generate_series(1, 20) i; + +PREPARE observe_select(int) AS + SELECT value FROM observe_cache WHERE key = $1; + +-- reach the generic plan before enabling logging; the cache is populated on +-- the second use of the generic plan, so this needs one more than the +-- five executions it takes to get there +EXECUTE observe_select(1); +EXECUTE observe_select(1); +EXECUTE observe_select(1); +EXECUTE observe_select(1); +EXECUTE observe_select(1); +EXECUTE observe_select(1); +EXECUTE observe_select(1); + +SET citus.log_remote_commands TO on; + +-- key 1's shard is already prepared on this connection, so this reuses +-- the statement by name rather than sending SQL +EXECUTE observe_select(1); + +-- an unprepared shard prepares first, then executes by name +EXECUTE observe_select(3); +EXECUTE observe_select(3); + +SET citus.log_remote_commands TO off; + +DEALLOCATE observe_select; + +-- ============================================================ +-- Test 14: SELECT ... FOR UPDATE is fast-path eligible, so the +-- cached task must carry its row locks. Without them the +-- executor treats it as an ordinary read and, with +-- select_opens_transaction_block off, opens no remote +-- transaction — releasing the row lock at statement end +-- rather than at the end of the enclosing transaction. +-- ============================================================ + +SET citus.select_opens_transaction_block TO off; + +PREPARE lock_row(int) AS + SELECT value FROM observe_cache WHERE key = $1 FOR UPDATE; + +-- reach the generic plan before enabling logging +EXECUTE lock_row(5); +EXECUTE lock_row(5); +EXECUTE lock_row(5); +EXECUTE lock_row(5); +EXECUTE lock_row(5); +EXECUTE lock_row(5); +EXECUTE lock_row(5); + +-- the worker must still be given a transaction to hold the row lock in +SET citus.log_remote_commands TO on; +BEGIN; +EXECUTE lock_row(5); +COMMIT; +SET citus.log_remote_commands TO off; + +RESET citus.select_opens_transaction_block; +DEALLOCATE lock_row; + +-- ============================================================ +-- Test 15: the same wire-protocol assertion for the remote DML +-- paths. INSERT and UPDATE/DELETE reach the cache through +-- their own eligibility and task-building branches, and a +-- silent fallback in either is invisible in query results. +-- ============================================================ + +-- no primary key, so the same shard can be hit repeatedly by INSERT +CREATE TABLE observe_dml (key int, value int); +SELECT create_distributed_table('observe_dml', 'key'); + +PREPARE observe_insert(int, int) AS + INSERT INTO observe_dml (key, value) VALUES ($1, $2); + +-- reach the generic plan and populate the cache before enabling logging +EXECUTE observe_insert(1, 1); +EXECUTE observe_insert(1, 2); +EXECUTE observe_insert(1, 3); +EXECUTE observe_insert(1, 4); +EXECUTE observe_insert(1, 5); +EXECUTE observe_insert(1, 6); +EXECUTE observe_insert(1, 7); + +SET citus.log_remote_commands TO on; + +-- key 1's shard is already prepared on this connection, so this reuses +-- the statement by name rather than sending SQL +EXECUTE observe_insert(1, 8); + +-- an unprepared shard prepares first, then executes by name +EXECUTE observe_insert(3, 9); +EXECUTE observe_insert(3, 10); + +SET citus.log_remote_commands TO off; + +PREPARE observe_update(int, int) AS + UPDATE observe_dml SET value = $2 WHERE key = $1; + +EXECUTE observe_update(1, 11); +EXECUTE observe_update(1, 12); +EXECUTE observe_update(1, 13); +EXECUTE observe_update(1, 14); +EXECUTE observe_update(1, 15); +EXECUTE observe_update(1, 16); +EXECUTE observe_update(1, 17); + +SET citus.log_remote_commands TO on; +EXECUTE observe_update(1, 18); +SET citus.log_remote_commands TO off; + +PREPARE observe_delete(int) AS + DELETE FROM observe_dml WHERE key = $1; + +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); +EXECUTE observe_delete(3); + +SET citus.log_remote_commands TO on; +EXECUTE observe_delete(3); +SET citus.log_remote_commands TO off; + +SELECT key, value FROM observe_dml ORDER BY key, value; + +DEALLOCATE observe_insert; +DEALLOCATE observe_update; +DEALLOCATE observe_delete; + +-- ============================================================ +-- Cleanup +-- ============================================================ + +SET citus.enable_prepared_statement_caching = off; +DROP SCHEMA prepared_stmt_caching CASCADE; diff --git a/src/test/regress/sql/single_node.sql b/src/test/regress/sql/single_node.sql index 079e59a93d8..93a24857fa3 100644 --- a/src/test/regress/sql/single_node.sql +++ b/src/test/regress/sql/single_node.sql @@ -1333,6 +1333,61 @@ BEGIN; SELECT coordinated_transaction_should_use_2PC(); ROLLBACK; +-- A cached fast-path upsert executed locally is deparsed for local planning, +-- and needs the insert alias so the DO UPDATE reference to the target table +-- resolves against the shard alias. +SET citus.shard_replication_factor TO 1; +SET citus.shard_count TO 4; +SET citus.enable_prepared_statement_caching TO on; + +-- the suite runs with citus.stat_tenants_track = 'ALL', which the cache declines +SET citus.stat_tenants_track TO 'none'; + +CREATE TABLE single_node.upsert_cache_test (key int PRIMARY KEY, value int); +SELECT create_distributed_table('single_node.upsert_cache_test', 'key'); + +PREPARE local_cached_upsert(int, int) AS + INSERT INTO single_node.upsert_cache_test (key, value) VALUES ($1, $2) + ON CONFLICT (key) DO UPDATE + SET value = upsert_cache_test.value + EXCLUDED.value; + +-- with local command logging on, LogLocalCommand() deparses the task first, so +-- the notice below shows the shard SQL that the alias has to appear in +EXECUTE local_cached_upsert(1, 10); +EXECUTE local_cached_upsert(1, 10); +EXECUTE local_cached_upsert(1, 10); +EXECUTE local_cached_upsert(1, 10); +EXECUTE local_cached_upsert(1, 10); +EXECUTE local_cached_upsert(1, 10); +EXECUTE local_cached_upsert(1, 10); +EXECUTE local_cached_upsert(2, 5); + +-- key 1 accumulated seven increments, key 2 was inserted once +SELECT key, value FROM single_node.upsert_cache_test ORDER BY key; + +-- with logging off nothing deparses the task up front, so the local executor +-- reaches its own fast-path deparse instead +SET citus.log_local_commands TO off; + +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(3, 100); +EXECUTE local_cached_upsert(4, 7); + +SELECT key, value FROM single_node.upsert_cache_test WHERE key >= 3 ORDER BY key; + +RESET citus.log_local_commands; +DEALLOCATE local_cached_upsert; +DROP TABLE single_node.upsert_cache_test; +RESET citus.enable_prepared_statement_caching; +RESET citus.stat_tenants_track; +RESET citus.shard_count; +RESET citus.shard_replication_factor; + -- if the local execution is disabled, we cannot failover to -- local execution and the queries would fail SET citus.enable_local_execution TO false;