Improve the logic some more

src/backend/distributed/planner/multi_router_planner.c

@@ -801,6 +801,7 @@ AllRelationRestrictionsContainUninstantiatedQual(
 		List *joinInfo = list_copy(restriction->relOptInfo->joininfo);
 		List *allRestrictions = list_concat(baseRestrictInfo, joinInfo);
 		ListCell *restrictionCell = NULL;
+		Var *relationPartitionKey = NULL;
 		bool relationHasRestriction = false;
 
 		if (ContainsFalseClause(extract_actual_clauses(allRestrictions, true)))
@@ -814,6 +815,8 @@ AllRelationRestrictionsContainUninstantiatedQual(
 			continue;
 		}
 
+		relationPartitionKey = PartitionKey(restriction->relationId);
+
 		foreach(restrictionCell, allRestrictions)
 		{
 			RestrictInfo *restrictInfo = (RestrictInfo *) lfirst(restrictionCell);
@@ -821,7 +824,7 @@ AllRelationRestrictionsContainUninstantiatedQual(
 			relationHasRestriction = relationHasRestriction ||
 									 HasUninstantiatedQualWalker(
 				(Node *) restrictInfo->clause,
-				NULL);
+				relationPartitionKey);
 
 			if (relationHasRestriction)
 			{
@@ -843,24 +846,75 @@ AllRelationRestrictionsContainUninstantiatedQual(
 static bool
 HasUninstantiatedQualWalker(Node *node, void *context)
 {
-	Param *param = NULL;
+	Var *relationPartitionColumn = (Var *) context;
 
 	if (node == NULL)
 	{
 		return false;
 	}
 
-	if (IsA(node, Param))
+	if (IsA(node, OpExpr) && list_length(((OpExpr *) node)->args) == 2)
 	{
-		param = (Param *) node;
+		OpExpr *op = (OpExpr *) node;
-	}
+		Node *leftop = get_leftop((Expr *) op);
+		Node *rightop = get_rightop((Expr *) op);
+		Param *param = NULL;
+		Var *currentColumn = NULL;
 
-	if (param && param->paramid == UNINSTANTIATED_PARAMETER_ID)
+		/* look for the Params */
-	{
+		if (IsA(leftop, Param))
+		{
+			param = (Param *) leftop;
+
+			/*
+			 * Before instantiating the qual, ensure that it is equal to
+			 * the partition key.
+			 */
+			if (IsA(rightop, Var))
+			{
+				currentColumn = (Var *) rightop;
+			}
+		}
+		else if (IsA(rightop, Param))
+		{
+			param = (Param *) rightop;
+
+			/*
+			 * Before instantiating the qual, ensure that it is equal to
+			 * the partition key.
+			 */
+			if (IsA(leftop, Var))
+			{
+				currentColumn = (Var *) leftop;
+			}
+		}
+		else
+		{
+			return expression_tree_walker(node, HasUninstantiatedQualWalker, context);
+		}
+
+		if (!(param && param->paramid == UNINSTANTIATED_PARAMETER_ID))
+		{
+			return false;
+		}
+
+		/* ensure that it is the relation's partition column */
+		if (relationPartitionColumn && currentColumn &&
+			currentColumn->varattno != relationPartitionColumn->varattno)
+		{
+			return false;
+		}
+
+		/*
+		 * We still return true here given that finding the parameter is the
+		 * actual goal of the walker. We only hit here once the query includes
+		 * (partitionColumn = Const) on the query and we artificially added
+		 * the uninstantiated parameter to the query.
+		 */
 		return true;
 	}
 
-	return expression_tree_walker(node, HasUninstantiatedQualWalker, NULL);
+	return expression_tree_walker(node, HasUninstantiatedQualWalker, context);
 }
 
 

src/test/regress/expected/multi_insert_select.out

@@ -1322,22 +1322,14 @@ SELECT
   raw_events_first.user_id
 FROM
   raw_events_first LEFT JOIN raw_events_second ON raw_events_first.user_id = raw_events_second.value_1;
-DEBUG:  predicate pruning for shardId 13300001
+ERROR:  cannot plan distributed query since all join conditions in the query need include two distribution keys using an equality operator
-DEBUG:  predicate pruning for shardId 13300002
-DEBUG:  predicate pruning for shardId 13300003
-ERROR:  cannot perform distributed planning for the given modification
-DETAIL:  Select query cannot be pushed down to the worker.
 -- same as the above with INNER JOIN
 INSERT INTO agg_events (user_id)
 SELECT
   raw_events_first.user_id
 FROM
   raw_events_first INNER JOIN raw_events_second ON raw_events_first.user_id = raw_events_second.value_1;
-DEBUG:  predicate pruning for shardId 13300001
+ERROR:  cannot plan distributed query since all join conditions in the query need include two distribution keys using an equality operator
-DEBUG:  predicate pruning for shardId 13300002
-DEBUG:  predicate pruning for shardId 13300003
-ERROR:  cannot perform distributed planning for the given modification
-DETAIL:  Select query cannot be pushed down to the worker.
 -- a not meaningful query
 INSERT INTO agg_events 
             (user_id) 
@@ -1359,11 +1351,7 @@ SELECT
   raw_events_first.user_id
 FROM
   raw_events_first INNER JOIN raw_events_second ON raw_events_first.user_id = raw_events_second.value_1;
-DEBUG:  predicate pruning for shardId 13300001
+ERROR:  cannot plan distributed query since all join conditions in the query need include two distribution keys using an equality operator
-DEBUG:  predicate pruning for shardId 13300002
-DEBUG:  predicate pruning for shardId 13300003
-ERROR:  cannot perform distributed planning for the given modification
-DETAIL:  Select query cannot be pushed down to the worker.
 -- although we do not support pushing down JOINs on non-partition
 -- columns here it is safe to push it down given that we're looking for
 -- a specific value (i.e., user_id = 10) on the joining column.
@@ -1453,11 +1441,7 @@ SELECT
 FROM
   raw_events_first INNER JOIN raw_events_second ON raw_events_first.user_id = raw_events_second.value_1
   WHERE raw_events_first.value_1 IN (10, 11,12) OR raw_events_second.user_id IN (1,2,3,4);
-DEBUG:  predicate pruning for shardId 13300001
+ERROR:  cannot plan distributed query since all join conditions in the query need include two distribution keys using an equality operator
-DEBUG:  predicate pruning for shardId 13300002
-DEBUG:  predicate pruning for shardId 13300003
-ERROR:  cannot perform distributed planning for the given modification
-DETAIL:  Select query cannot be pushed down to the worker.
 -- implicit join on non partition column should also not be pushed down
 INSERT INTO agg_events 
             (user_id) 
@@ -1507,14 +1491,7 @@ FROM
 ON (f.id = f2.id)) as outer_most
 GROUP BY
   outer_most.id;
-DEBUG:  predicate pruning for shardId 13300001
+ERROR:  cannot plan distributed query since all join conditions in the query need include two distribution keys using an equality operator
-DEBUG:  predicate pruning for shardId 13300002
-DEBUG:  predicate pruning for shardId 13300003
-DEBUG:  predicate pruning for shardId 13300005
-DEBUG:  predicate pruning for shardId 13300006
-DEBUG:  predicate pruning for shardId 13300007
-ERROR:  cannot perform distributed planning for the given modification
-DETAIL:  Select query cannot be pushed down to the worker.
 -- not equals on the partition column cannot be pushed down
 INSERT INTO agg_events
             (value_4_agg,
@@ -1762,11 +1739,7 @@ SELECT user_id
 FROM   raw_events_first 
 WHERE  user_id IN (SELECT value_2 
                    FROM   raw_events_second);
-DEBUG:  predicate pruning for shardId 13300001
+ERROR:  cannot plan distributed query since all join conditions in the query need include two distribution keys using an equality operator
-DEBUG:  predicate pruning for shardId 13300002
-DEBUG:  predicate pruning for shardId 13300003
-ERROR:  cannot perform distributed planning for the given modification
-DETAIL:  Select query cannot be pushed down to the worker.
 -- we currently not support grouping sets
 INSERT INTO agg_events
             (user_id,