Implement DistPartitionKeys

multi_column.txt

@@ -1,5 +1,12 @@
 - [ ] DistPartitionKey()
+   - [x] Add DistPartitionKeys
+   - [ ] Remove DistPartitionKey
+- [ ] DistPartitionKeyOrError()
+   - [ ] Add DistPartitionKeysOrError
+   - [ ] Remove DistPartitionKeyOrError
 - [ ] PartitionColumn()
+   - [ ] Add PartitionColumns
+   - [ ] Remove PartitionColumn
 - [ ] IsPartitionColumn()
 - [ ] PartitionColumnIndex(Var *targetVar, List *targetList)
 - [ ] PartitionColumnIndex(List *insertTargetList, Var *partitionColumn);

src/backend/distributed/planner/multi_join_order.c

@@ -1379,16 +1379,39 @@ PartitionColumn(Oid relationId, uint32 rangeTableId)
  */
 Var *
 DistPartitionKey(Oid relationId)
+{
+	List *partitionKeys = DistPartitionKeys(relationId);
+	if (list_length(partitionKeys) == 0)
+	{
+		return NULL;
+	}
+
+	return linitial(partitionKeys);
+}
+
+
+/*
+ * DistPartitionKeys returns the list of partition key columns for the given
+ * relation. Note that in the context of distributed join and query planning,
+ * the callers of this function *must* set the partition key column's range
+ * table reference (varno) to match the table's location in the query range
+ * table list.
+ *
+ * Note that reference tables do not have partition columns. Thus, this function
+ * returns NIL when called for reference tables.
+ */
+List *
+DistPartitionKeys(Oid relationId)
 {
 	CitusTableCacheEntry *partitionEntry = GetCitusTableCacheEntry(relationId);
 
 	/* non-distributed tables do not have partition column */
 	if (IsCitusTableTypeCacheEntry(partitionEntry, CITUS_TABLE_WITH_NO_DIST_KEY))
 	{
-		return NULL;
+		return NIL;
 	}
 
-	return copyObject(partitionEntry->partitionColumn);
+	return copyObject(partitionEntry->partitionColumns);
 }
 
 

src/include/distributed/multi_join_order.h

@@ -105,6 +105,7 @@ extern Var * LeftColumnOrNULL(OpExpr *joinClause);
 extern Var * RightColumnOrNULL(OpExpr *joinClause);
 extern Var * PartitionColumn(Oid relationId, uint32 rangeTableId);
 extern Var * DistPartitionKey(Oid relationId);
+extern List * DistPartitionKeys(Oid relationId);
 extern Var * DistPartitionKeyOrError(Oid relationId);
 extern char PartitionMethod(Oid relationId);
 extern char TableReplicationModel(Oid relationId);