Cache prepared statement plans for fast path queries on worker connections

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);

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

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