Cache the partition column Var

src/backend/distributed/planner/multi_join_order.c

@@ -1350,8 +1350,6 @@ Var *
 DistPartitionKey(Oid relationId)
 {
 	DistTableCacheEntry *partitionEntry = DistributedTableCacheEntry(relationId);
-	Node *variableNode = NULL;
-	Var *partitionKey = NULL;
 
 	/* reference tables do not have partition column */
 	if (partitionEntry->partitionMethod == DISTRIBUTE_BY_NONE)
@@ -1359,13 +1357,7 @@ DistPartitionKey(Oid relationId)
 		return NULL;
 	}
 
-	/* now obtain partition key and build the var node */
+	return copyObject(partitionEntry->partitionColumn);
-	variableNode = stringToNode(partitionEntry->partitionKeyString);
-
-	partitionKey = (Var *) variableNode;
-	Assert(IsA(variableNode, Var));
-
-	return partitionKey;
 }
 
 

src/backend/distributed/utils/metadata_cache.c

@@ -843,8 +843,19 @@ BuildDistTableCacheEntry(DistTableCacheEntry *cacheEntry)
 	/* note that for reference tables partitionKeyisNull is true */
 	if (!partitionKeyIsNull)
 	{
+		Node *partitionNode = NULL;
+
 		oldContext = MemoryContextSwitchTo(CacheMemoryContext);
+
+		/* get the string representation of the partition column Var */
 		cacheEntry->partitionKeyString = TextDatumGetCString(partitionKeyDatum);
+
+		/* convert the string to a Node and ensure it is a Var */
+		partitionNode = stringToNode(cacheEntry->partitionKeyString);
+		Assert(IsA(partitionNode, Var));
+
+		cacheEntry->partitionColumn = (Var *) partitionNode;
+
 		MemoryContextSwitchTo(oldContext);
 	}
 	else
@@ -886,11 +897,9 @@ BuildDistTableCacheEntry(DistTableCacheEntry *cacheEntry)
 	if (cacheEntry->partitionMethod == DISTRIBUTE_BY_HASH)
 	{
 		TypeCacheEntry *typeEntry = NULL;
-		Node *partitionNode = stringToNode(cacheEntry->partitionKeyString);
-		Var *partitionColumn = (Var *) partitionNode;
 		FmgrInfo *hashFunction = NULL;
+		Var *partitionColumn = cacheEntry->partitionColumn;
 
-		Assert(IsA(partitionNode, Var));
 		typeEntry = lookup_type_cache(partitionColumn->vartype,
 									  TYPECACHE_HASH_PROC_FINFO);
 
@@ -2749,6 +2758,12 @@ ResetDistTableCacheEntry(DistTableCacheEntry *cacheEntry)
 		cacheEntry->hashFunction = NULL;
 	}
 
+	if (cacheEntry->partitionColumn != NULL)
+	{
+		pfree(cacheEntry->partitionColumn);
+		cacheEntry->partitionColumn = NULL;
+	}
+
 	if (cacheEntry->shardIntervalArrayLength == 0)
 	{
 		return;

src/include/distributed/metadata_cache.h

@@ -49,6 +49,7 @@ typedef struct
 
 	/* pg_dist_partition metadata for this table */
 	char *partitionKeyString;
+	Var *partitionColumn;
 	char partitionMethod;
 	uint32 colocationId;
 	char replicationModel;