0002-ExecPartitionCheck-pre-compute-partition-key-express.patch

From a7909c2b142ef255c034d7d319d09abdeede23eb Mon Sep 17 00:00:00 2001
From: amitlan <amitlangote09@gmail.com>
Date: Tue, 25 May 2021 22:55:12 +0900
Subject: [PATCH 2/2] ExecPartitionCheck: pre-compute partition key expression
 v2

---
 src/backend/executor/execMain.c      | 95 ++++++++++++++++++++++++++++
 src/backend/executor/execPartition.c | 52 +++++++++++++++
 src/include/nodes/execnodes.h        |  9 +++
 3 files changed, 156 insertions(+)

diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index b3ce4bae53..1fc2a9fe82 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -44,6 +44,7 @@
 #include "access/transam.h"
 #include "access/xact.h"
 #include "catalog/namespace.h"
+#include "catalog/partition.h"
 #include "catalog/pg_publication.h"
 #include "commands/matview.h"
 #include "commands/trigger.h"
@@ -52,6 +53,8 @@
 #include "foreign/fdwapi.h"
 #include "jit/jit.h"
 #include "mb/pg_wchar.h"
+#include "nodes/makefuncs.h"
+#include "nodes/nodeFuncs.h"
 #include "miscadmin.h"
 #include "parser/parsetree.h"
 #include "storage/bufmgr.h"
@@ -1686,6 +1689,32 @@ ExecRelCheck(ResultRelInfo *resultRelInfo,
 	return NULL;
 }
 
+/*
+ * Replaces the occurrence of cxt->matchexpr in the expression tree given by
+ * 'node' by an OUTER var with provided attribute number.
+ */
+typedef struct
+{
+	Expr	   *matchexpr;
+	AttrNumber	varattno;
+} replace_partexpr_with_dummy_var_context;
+
+static Node *
+replace_partexpr_with_dummy_var(Node *node,
+								replace_partexpr_with_dummy_var_context *cxt)
+{
+	if (node == NULL)
+		return NULL;
+
+	if (equal(node, cxt->matchexpr))
+		return (Node *) makeVar(OUTER_VAR, cxt->varattno,
+								exprType(node), exprTypmod(node),
+								exprCollation(node), 0);
+
+	return expression_tree_mutator(node, replace_partexpr_with_dummy_var,
+								   (void *) cxt);
+}
+
 /*
  * ExecPartitionCheck --- check that tuple meets the partition constraint.
  *
@@ -1716,6 +1745,45 @@ ExecPartitionCheck(ResultRelInfo *resultRelInfo, TupleTableSlot *slot,
 		MemoryContext oldcxt = MemoryContextSwitchTo(estate->es_query_cxt);
 		List	   *qual = RelationGetPartitionQual(resultRelInfo->ri_RelationDesc);
 
+		/*
+		 * Optimize the evaluation of partition key expressions.  The way we do
+		 * that is by replacing any occurrences of the individual expressions
+		 * in this relation's partition constraint by dummy Vars marked as
+		 * coming from the "OUTER" relation.  Then when actually executing such
+		 * modified partition constraint tree, we feed the actual partition
+		 * expression values via econtext->ecxt_outertuple; see below.
+		 */
+		if (resultRelInfo->ri_partConstrKeyExprs)
+		{
+			List	  *partexprs = resultRelInfo->ri_partConstrKeyExprs;
+			ListCell  *lc;
+			AttrNumber attrno = 1;
+			TupleDesc	partexprs_tupdesc;
+			replace_partexpr_with_dummy_var_context cxt;
+
+			partexprs_tupdesc = CreateTemplateTupleDesc(list_length(partexprs));
+			foreach(lc, partexprs)
+			{
+				Expr   *expr = lfirst(lc);
+
+				cxt.matchexpr = expr;
+				cxt.varattno = attrno;
+				qual = (List *) replace_partexpr_with_dummy_var((Node *) qual,
+																&cxt);
+
+				resultRelInfo->ri_partConstrKeyExprStates =
+					lappend(resultRelInfo->ri_partConstrKeyExprStates,
+							ExecPrepareExpr(expr, estate));
+				TupleDescInitEntry(partexprs_tupdesc, attrno, NULL,
+								   exprType((Node *) expr),
+								   exprTypmod((Node *) expr), 0);
+				attrno++;
+			}
+
+			resultRelInfo->ri_partConstrKeyExprsSlot =
+				ExecInitExtraTupleSlot(estate, partexprs_tupdesc, &TTSOpsVirtual);
+		}
+
 		resultRelInfo->ri_PartitionCheckExpr = ExecPrepareCheck(qual, estate);
 		MemoryContextSwitchTo(oldcxt);
 	}
@@ -1729,6 +1797,33 @@ ExecPartitionCheck(ResultRelInfo *resultRelInfo, TupleTableSlot *slot,
 	/* Arrange for econtext's scan tuple to be the tuple under test */
 	econtext->ecxt_scantuple = slot;
 
+	if (resultRelInfo->ri_partConstrKeyExprStates)
+	{
+		TupleTableSlot *partexprs_slot = resultRelInfo->ri_partConstrKeyExprsSlot;
+		Datum	*values;
+		bool	*isnull;
+		ListCell *lc;
+		AttrNumber attrno = 1;
+
+		Assert(partexprs_slot != NULL);
+		ExecClearTuple(partexprs_slot);
+
+		values = partexprs_slot->tts_values;
+		isnull = partexprs_slot->tts_isnull;
+
+		foreach(lc, resultRelInfo->ri_partConstrKeyExprStates)
+		{
+			ExprState   *partexpr = lfirst(lc);
+
+			values[attrno-1] = ExecEvalExprSwitchContext(partexpr, econtext,
+												  &isnull[attrno-1]);
+			attrno++;
+		}
+		ExecStoreVirtualTuple(partexprs_slot);
+
+		econtext->ecxt_outertuple = partexprs_slot;
+	}
+
 	/*
 	 * As in case of the catalogued constraints, we treat a NULL result as
 	 * success here, not a failure.
diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
index 1d0d8e63f6..0e689b9c56 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -143,6 +143,13 @@ struct PartitionTupleRouting
  *		If non-NULL, PartititionDispatch for the sub-partitioned partition
  *		that was most recently chosen as the routing target
  *
+ * partconstr_keyexprs
+ *		List of expressions present in the partition keys of all ancestors
+ *		of this table including itself, mapped to have the attribute
+ *		numbers of this table.  The field is so named because all of these
+ *		expressions appear in the partition constraint of each of this
+ *		table's partitions.
+ *
  * indexes
  *		Array of partdesc->nparts elements.  For leaf partitions the index
  *		corresponds to the partition's ResultRelInfo in the encapsulating
@@ -162,6 +169,7 @@ typedef struct PartitionDispatchData
 	AttrMap    *tupmap;
 	ResultRelInfo *savedPartInfo;
 	PartitionDispatch savedDispatchInfo;
+	List	   *partconstr_keyexprs;
 	int			indexes[FLEXIBLE_ARRAY_MEMBER];
 }			PartitionDispatchData;
 
@@ -1006,6 +1014,21 @@ InitRootToPartitionMap(ResultRelInfo *partRelInfo,
 		partRelInfo->ri_PartitionTupleSlot = NULL;
 }
 
+/*
+ * Save parent's partition key expressions in the partition ResultRelInfo
+ * after mapping them to have the partition's attribute numbers.
+ */
+static inline void
+InitPartitionConstraintKeyExprs(PartitionDispatch dispatch,
+								ResultRelInfo *partRelInfo)
+{
+	if (dispatch->partconstr_keyexprs)
+		partRelInfo->ri_partConstrKeyExprs =
+			map_partition_varattnos(dispatch->partconstr_keyexprs, 1,
+									partRelInfo->ri_RelationDesc,
+									dispatch->reldesc);
+}
+
 /*
  * ExecInitRoutingInfo
  *		Set up information needed for translating tuples between root
@@ -1057,6 +1080,8 @@ ExecInitRoutingInfo(ModifyTableState *mtstate,
 
 	partRelInfo->ri_CopyMultiInsertBuffer = NULL;
 
+	InitPartitionConstraintKeyExprs(dispatch, partRelInfo);
+
 	/*
 	 * Keep track of it in the PartitionTupleRouting->partitions array.
 	 */
@@ -1175,6 +1200,32 @@ ExecInitPartitionDispatchInfo(EState *estate,
 	pd->savedPartInfo = NULL;
 	pd->savedDispatchInfo = NULL;
 
+	if (pd->key->partexprs != NIL)
+	{
+		pd->partconstr_keyexprs = copyObject(pd->key->partexprs);
+		if (parent_pd)
+		{
+			List   *parent_keyexprs = parent_pd->partconstr_keyexprs;
+
+			if (parent_keyexprs && pd->tupmap)
+			{
+				bool	found_whole_row;
+
+				parent_keyexprs = (List *)
+					map_variable_attnos((Node *) parent_keyexprs, 1, 0,
+										pd->tupmap,
+										RelationGetForm(rel)->reltype,
+										&found_whole_row);
+			}
+			else if (parent_keyexprs)
+				parent_keyexprs = copyObject(parent_keyexprs);
+			pd->partconstr_keyexprs =
+				list_concat(pd->partconstr_keyexprs, parent_keyexprs);
+		}
+	}
+	else
+		pd->partconstr_keyexprs = NIL;
+
 	/*
 	 * Initialize with -1 to signify that the corresponding partition's
 	 * ResultRelInfo or PartitionDispatch has not been created yet.
@@ -1220,6 +1271,7 @@ ExecInitPartitionDispatchInfo(EState *estate,
 		InitResultRelInfo(rri, rel, 0, rootResultRelInfo, 0);
 		/* The map is needed in CanUseSavedPartitionForTuple(). */
 		InitRootToPartitionMap(rri, rootResultRelInfo, estate);
+		InitPartitionConstraintKeyExprs(pd, rri);
 		proute->nonleaf_partitions[dispatchidx] = rri;
 	}
 	else
diff --git a/src/include/nodes/execnodes.h b/src/include/nodes/execnodes.h
index 7795a69490..7f1ce732ea 100644
--- a/src/include/nodes/execnodes.h
+++ b/src/include/nodes/execnodes.h
@@ -496,6 +496,15 @@ typedef struct ResultRelInfo
 	/* partition check expression state (NULL if not set up yet) */
 	ExprState  *ri_PartitionCheckExpr;
 
+	/*
+	 * Information used by ExecPartitionCheck() to optimize some cases where
+	 * the partition's ancestors' partition keys contain arbitrary
+	 * expressions.
+	 */
+	List	   *ri_partConstrKeyExprs;
+	List	   *ri_partConstrKeyExprStates;
+	TupleTableSlot *ri_partConstrKeyExprsSlot;
+
 	/*
 	 * Information needed by tuple routing target relations
 	 *
-- 
2.24.1