From a2737115000abe73fbd52027cf759d27f05d7f66 Mon Sep 17 00:00:00 2001
From: Onder Kalaci <onder@citusdata.com>
Date: Tue, 28 Nov 2017 20:54:07 +0200
Subject: [PATCH] The common attribute equivalance class always includes the
 input relations

We added the ability to filter out the planner restriction information
for specific parts of the query. This might lead to situations where
the common restriction includes some other relations that we're searching
for. The reason is that while filtering for join restrictions, we add the
restriction as soon as we find the relation.

With this commit we make sure that the common attribute
equivalance class always includes the input relations.
---
 .../relation_restriction_equivalence.c        | 66 +++++++++++++++++--
 1 file changed, 59 insertions(+), 7 deletions(-)

diff --git a/src/backend/distributed/planner/relation_restriction_equivalence.c b/src/backend/distributed/planner/relation_restriction_equivalence.c
index 68e0c7351..13bbb4b3f 100644
--- a/src/backend/distributed/planner/relation_restriction_equivalence.c
+++ b/src/backend/distributed/planner/relation_restriction_equivalence.c
@@ -113,7 +113,13 @@ static bool AttributeEquivalancesAreEqual(AttributeEquivalenceClass *
 										  AttributeEquivalenceClass *
 										  secondAttributeEquivalance);
 static AttributeEquivalenceClass * GenerateCommonEquivalence(List *
-															 attributeEquivalenceList);
+															 attributeEquivalenceList,
+															 RelationRestrictionContext *
+															 relationRestrictionContext);
+static AttributeEquivalenceClass * GenerateEquivalanceClassForRelationRestriction(
+	RelationRestrictionContext
+	*
+	relationRestrictionContext);
 static void ListConcatUniqueAttributeClassMemberLists(AttributeEquivalenceClass **
 													  firstClass,
 													  AttributeEquivalenceClass *
@@ -478,7 +484,8 @@ EquivalenceListContainsRelationsEquality(List *attributeEquivalenceList,
 	 * common equivalence class. The main goal is to test whether this main class
 	 * contains all partition keys of the existing relations.
 	 */
-	commonEquivalenceClass = GenerateCommonEquivalence(attributeEquivalenceList);
+	commonEquivalenceClass = GenerateCommonEquivalence(attributeEquivalenceList,
+													   restrictionContext);
 
 	/* add the rte indexes of relations to a bitmap */
 	foreach(commonEqClassCell, commonEquivalenceClass->equivalentAttributes)
@@ -707,7 +714,8 @@ GetVarFromAssignedParam(List *parentPlannerParamList, Param *plannerParam)
  *      - Finally, return the common equivalence class.
  */
 static AttributeEquivalenceClass *
-GenerateCommonEquivalence(List *attributeEquivalenceList)
+GenerateCommonEquivalence(List *attributeEquivalenceList,
+						  RelationRestrictionContext *relationRestrictionContext)
 {
 	AttributeEquivalenceClass *commonEquivalenceClass = NULL;
 	AttributeEquivalenceClass *firstEquivalenceClass = NULL;
@@ -718,14 +726,20 @@ GenerateCommonEquivalence(List *attributeEquivalenceList)
 	commonEquivalenceClass = palloc0(sizeof(AttributeEquivalenceClass));
 	commonEquivalenceClass->equivalenceId = 0;
 
-	/* think more on this. */
-	if (equivalenceListSize < 1)
+	/*
+	 * We seed the common equivalence class with a the first distributed
+	 * table since we always want the input distributed relations to be
+	 * on the common class.
+	 */
+	firstEquivalenceClass =
+		GenerateEquivalanceClassForRelationRestriction(relationRestrictionContext);
+
+	/* we skip the calculation if there are not enough information */
+	if (equivalenceListSize < 1 || firstEquivalenceClass == NULL)
 	{
 		return commonEquivalenceClass;
 	}
 
-	/* setup the initial state of the main equivalence class */
-	firstEquivalenceClass = linitial(attributeEquivalenceList);
 	commonEquivalenceClass->equivalentAttributes =
 		firstEquivalenceClass->equivalentAttributes;
 	addedEquivalenceIds = bms_add_member(addedEquivalenceIds,
@@ -778,6 +792,44 @@ GenerateCommonEquivalence(List *attributeEquivalenceList)
 }
 
 
+/*
+ * GenerateEquivalanceClassForRelationRestriction generates an AttributeEquivalenceClass
+ * with a single AttributeEquivalenceClassMember.
+ */
+static AttributeEquivalenceClass *
+GenerateEquivalanceClassForRelationRestriction(
+	RelationRestrictionContext *relationRestrictionContext)
+{
+	ListCell *relationRestrictionCell = NULL;
+	AttributeEquivalenceClassMember *eqMember = NULL;
+	AttributeEquivalenceClass *eqClassForRelation = NULL;
+
+	foreach(relationRestrictionCell, relationRestrictionContext->relationRestrictionList)
+	{
+		RelationRestriction *relationRestriction =
+			(RelationRestriction *) lfirst(relationRestrictionCell);
+		Var *relationPartitionKey = DistPartitionKey(relationRestriction->relationId);
+
+		if (relationPartitionKey)
+		{
+			eqClassForRelation = palloc0(sizeof(AttributeEquivalenceClass));
+			eqMember = palloc0(sizeof(AttributeEquivalenceClassMember));
+			eqMember->relationId = relationRestriction->relationId;
+			eqMember->rteIdentity = GetRTEIdentity(relationRestriction->rte);
+			eqMember->varno = relationRestriction->index;
+			eqMember->varattno = relationPartitionKey->varattno;
+
+			eqClassForRelation->equivalentAttributes =
+				lappend(eqClassForRelation->equivalentAttributes, eqMember);
+
+			break;
+		}
+	}
+
+	return eqClassForRelation;
+}
+
+
 /*
  * ListConcatUniqueAttributeClassMemberLists gets two attribute equivalence classes. It
  * basically concatenates attribute equivalence member lists uniquely and updates the