Undo skipping qual addition once the query already includes a qual

src/backend/distributed/planner/multi_router_planner.c

@@ -78,12 +78,6 @@ typedef struct InstantiateQualWalker
 	Var *relationPartitionColumn;
 }InstantiateQualWalker;
 
-typedef struct EqualityCheckOnPartitionColumnWalker
-{
-	Var *relationPartitionColumn;
-	bool partitionColumnEqualityExists;
-}PartitionColumnEqualityCheckWalker;
-
 
 bool EnableRouterExecution = true;
 
@@ -141,7 +135,6 @@ static DeferredErrorMessage * InsertPartitionColumnMatchesSelect(Query *query,
 																 Oid *
 																 selectPartitionColumnTableId);
 static void AddUninstantiatedEqualityQual(Query *query, Var *targetPartitionColumnVar);
-static Node * PartitionColumnEqualityWalker(Node *originalNode, void *context);
 static DeferredErrorMessage * ErrorIfQueryHasModifyingCTE(Query *queryTree);
 
 
@@ -479,11 +472,17 @@ RouterModifyTaskForShardInterval(Query *originalQuery, ShardInterval *shardInter
 		}
 	}
 
-	/* this case indicate*/
+	/*
-	if (!ShardsColocated(anchorShardInterval, shardInterval))
+	 * We should not let the anchor shard interval and target shard interval
+	 * being not colocated. This case is mostly valid once original query already
+	 * includes a partition column
+	 */
+	if (!ShardsIntervalsEqual(anchorShardInterval, shardInterval))
 	{
 		ereport(DEBUG2, (errmsg("Skipping target shard interval %ld since "
-								"SELECT query for it pruned away", shardId)));
+								"it doesn't have the same shard range with the select "
+								"anchor shard interval %ld",
+								shardId, anchorShardInterval->shardId)));
 
 		return NULL;
 	}
@@ -492,6 +491,11 @@ RouterModifyTaskForShardInterval(Query *originalQuery, ShardInterval *shardInter
 	insertShardPlacementList = FinalizedShardPlacementList(shardId);
 	intersectedPlacementList = IntersectPlacementList(insertShardPlacementList,
 													  selectPlacementList);
+
+	/*
+	 * If insert target does not have exactly the same placements with the select,
+	 * we sholdn't run the query.
+	 */
 	if (list_length(insertShardPlacementList) != list_length(intersectedPlacementList))
 	{
 		ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
@@ -1188,22 +1192,6 @@ AddUninstantiatedEqualityQual(Query *query, Var *partitionColumn)
 	Oid equalsOperator = InvalidOid;
 	Oid greaterOperator = InvalidOid;
 	bool hashable = false;
-	PartitionColumnEqualityCheckWalker *walker =
-		palloc0(sizeof(PartitionColumnEqualityCheckWalker));
-
-	AssertArg(query->commandType == CMD_SELECT);
-
-	walker->partitionColumnEqualityExists = false;
-	walker->relationPartitionColumn = partitionColumn;
-
-	/* if the query already includes part_column = const, we don't need to add another. In fact,
-	 * adding another equality qual breaks lots of things.
-	 */
-	PartitionColumnEqualityWalker((Node *) query->jointree->quals, walker);
-	if (walker->partitionColumnEqualityExists)
-	{
-		return;
-	}
 
 	/* get the necessary equality operator */
 	get_sort_group_operators(partitionColumn->vartype, false, true, false,
@@ -1246,51 +1234,6 @@ AddUninstantiatedEqualityQual(Query *query, Var *partitionColumn)
 }
 
 
-/*
- *
- */
-static Node *
-PartitionColumnEqualityWalker(Node *originalNode, void *context)
-{
-	PartitionColumnEqualityCheckWalker *walker =
-		(PartitionColumnEqualityCheckWalker *) context;
-
-	if (originalNode == NULL)
-	{
-		return NULL;
-	}
-
-	if (IsA(originalNode, OpExpr) && list_length(((OpExpr *) originalNode)->args) == 2)
-	{
-		OpExpr *op = (OpExpr *) originalNode;
-		Node *leftop = get_leftop((Expr *) op);
-		Node *rightop = get_rightop((Expr *) op);
-
-
-		if (IsA(leftop, Var) && IsA(rightop, Const) &&
-			OperatorImplementsEquality(op->opno))
-		{
-			if (((Var *) leftop)->varattno == walker->relationPartitionColumn->varattno)
-			{
-				walker->partitionColumnEqualityExists = true;
-
-				return NULL;
-			}
-		}
-		else if (IsA(leftop, Const) && IsA(rightop, Var) &&
-				 OperatorImplementsEquality(op->opno))
-		{
-			walker->partitionColumnEqualityExists = true;
-			return NULL;
-		}
-	}
-
-
-	return expression_tree_mutator(originalNode, PartitionColumnEqualityWalker,
-								   (void *) context);
-}
-
-
 /*
  * ModifyQuerySupported returns NULL if the query only contains supported
  * features, otherwise it returns an error description.

src/backend/distributed/utils/colocation_utils.c

@@ -38,8 +38,6 @@
 /* local function forward declarations */
 static void MarkTablesColocated(Oid sourceRelationId, Oid targetRelationId);
 static void ErrorIfShardPlacementsNotColocated(Oid leftRelationId, Oid rightRelationId);
-static bool ShardsIntervalsEqual(ShardInterval *leftShardInterval,
-								 ShardInterval *rightShardInterval);
 static bool HashPartitionedShardIntervalsEqual(ShardInterval *leftShardInterval,
 											   ShardInterval *rightShardInterval);
 static int CompareShardPlacementsByNode(const void *leftElement,
@@ -296,7 +294,7 @@ ErrorIfShardPlacementsNotColocated(Oid leftRelationId, Oid rightRelationId)
  *       and shard min/max values). Thus, always return true for shards of reference
  *       tables.
  */
-static bool
+bool
 ShardsIntervalsEqual(ShardInterval *leftShardInterval, ShardInterval *rightShardInterval)
 {
 	char leftIntervalPartitionMethod = PartitionMethod(leftShardInterval->relationId);

src/include/distributed/colocation_utils.h

@@ -21,6 +21,8 @@ extern uint32 TableColocationId(Oid distributedTableId);
 extern bool TablesColocated(Oid leftDistributedTableId, Oid rightDistributedTableId);
 extern bool ShardsColocated(ShardInterval *leftShardInterval,
 							ShardInterval *rightShardInterval);
+extern bool ShardsIntervalsEqual(ShardInterval *leftShardInterval,
+								 ShardInterval *rightShardInterval);
 extern List * ColocatedTableList(Oid distributedTableId);
 extern List * ColocatedShardIntervalList(ShardInterval *shardInterval);
 extern Oid ColocatedTableId(Oid colocationId);

src/test/regress/expected/multi_insert_select.out

@@ -200,7 +200,7 @@ DEBUG:  name: unnamed; blockState:       DEFAULT; state: INPROGR, xid/subid/cid:
 DEBUG:  predicate pruning for shardId 13300000
 DEBUG:  predicate pruning for shardId 13300002
 DEBUG:  predicate pruning for shardId 13300003
-DEBUG:  Skipping target shard interval 13300004 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300004 since it doesn't have the same shard range with the select anchor shard interval 13300001
 DEBUG:  predicate pruning for shardId 13300000
 DEBUG:  predicate pruning for shardId 13300002
 DEBUG:  predicate pruning for shardId 13300003
@@ -208,11 +208,11 @@ DEBUG:  distributed statement: INSERT INTO public.raw_events_second_13300005 AS
 DEBUG:  predicate pruning for shardId 13300000
 DEBUG:  predicate pruning for shardId 13300002
 DEBUG:  predicate pruning for shardId 13300003
-DEBUG:  Skipping target shard interval 13300006 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300006 since it doesn't have the same shard range with the select anchor shard interval 13300001
 DEBUG:  predicate pruning for shardId 13300000
 DEBUG:  predicate pruning for shardId 13300002
 DEBUG:  predicate pruning for shardId 13300003
-DEBUG:  Skipping target shard interval 13300007 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300007 since it doesn't have the same shard range with the select anchor shard interval 13300001
 DEBUG:  ProcessQuery
 DEBUG:  Plan is router executable
 DEBUG:  CommitTransactionCommand
@@ -248,15 +248,15 @@ DEBUG:  distributed statement: INSERT INTO public.raw_events_second_13300004 AS
 DEBUG:  predicate pruning for shardId 13300001
 DEBUG:  predicate pruning for shardId 13300002
 DEBUG:  predicate pruning for shardId 13300003
-DEBUG:  Skipping target shard interval 13300005 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300005 since it doesn't have the same shard range with the select anchor shard interval 13300000
 DEBUG:  predicate pruning for shardId 13300001
 DEBUG:  predicate pruning for shardId 13300002
 DEBUG:  predicate pruning for shardId 13300003
-DEBUG:  Skipping target shard interval 13300006 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300006 since it doesn't have the same shard range with the select anchor shard interval 13300000
 DEBUG:  predicate pruning for shardId 13300001
 DEBUG:  predicate pruning for shardId 13300002
 DEBUG:  predicate pruning for shardId 13300003
-DEBUG:  Skipping target shard interval 13300007 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300007 since it doesn't have the same shard range with the select anchor shard interval 13300000
 DEBUG:  ProcessQuery
 DEBUG:  Plan is router executable
 DEBUG:  CommitTransactionCommand
@@ -365,9 +365,28 @@ DEBUG:  StartTransactionCommand
 DEBUG:  StartTransaction
 DEBUG:  name: unnamed; blockState:       DEFAULT; state: INPROGR, xid/subid/cid: 0/1/0, nestlvl: 1, children: 
 DEBUG:  predicate pruning for shardId 13300000
+DEBUG:  predicate pruning for shardId 13300001
 DEBUG:  predicate pruning for shardId 13300002
-ERROR:  cannot perform distributed planning for the given modification
+DEBUG:  predicate pruning for shardId 13300003
-DETAIL:  Select query cannot be pushed down to the worker.
+DEBUG:  Skipping target shard interval 13300004 since SELECT query for it pruned away
+DEBUG:  predicate pruning for shardId 13300000
+DEBUG:  predicate pruning for shardId 13300002
+DEBUG:  predicate pruning for shardId 13300003
+DEBUG:  distributed statement: INSERT INTO public.raw_events_second_13300005 AS citus_table_alias (user_id, value_1, value_3) SELECT user_id, value_1, value_3 FROM public.raw_events_first_13300001 raw_events_first WHERE (((user_id = 9) OR (user_id = 16)) AND ((hashint4(user_id) >= '-1073741824'::integer) AND (hashint4(user_id) <= '-1'::integer))) RETURNING citus_table_alias.user_id, citus_table_alias."time", citus_table_alias.value_1, citus_table_alias.value_2, citus_table_alias.value_3, citus_table_alias.value_4
+DEBUG:  predicate pruning for shardId 13300000
+DEBUG:  predicate pruning for shardId 13300001
+DEBUG:  predicate pruning for shardId 13300002
+DEBUG:  predicate pruning for shardId 13300003
+DEBUG:  Skipping target shard interval 13300006 since SELECT query for it pruned away
+DEBUG:  predicate pruning for shardId 13300000
+DEBUG:  predicate pruning for shardId 13300001
+DEBUG:  predicate pruning for shardId 13300002
+DEBUG:  distributed statement: INSERT INTO public.raw_events_second_13300007 AS citus_table_alias (user_id, value_1, value_3) SELECT user_id, value_1, value_3 FROM public.raw_events_first_13300003 raw_events_first WHERE (((user_id = 9) OR (user_id = 16)) AND ((hashint4(user_id) >= 1073741824) AND (hashint4(user_id) <= 2147483647))) RETURNING citus_table_alias.user_id, citus_table_alias."time", citus_table_alias.value_1, citus_table_alias.value_2, citus_table_alias.value_3, citus_table_alias.value_4
+DEBUG:  ProcessQuery
+DEBUG:  Plan is router executable
+ERROR:  duplicate key value violates unique constraint "raw_events_second_user_id_value_1_key_13300007"
+DETAIL:  Key (user_id, value_1)=(9, 90) already exists.
+CONTEXT:  while executing command on localhost:57638
 -- now do some aggregations
 INSERT INTO agg_events 
 SELECT
@@ -697,21 +716,21 @@ DEBUG:  predicate pruning for shardId 13300002
 DEBUG:  predicate pruning for shardId 13300004
 DEBUG:  predicate pruning for shardId 13300005
 DEBUG:  predicate pruning for shardId 13300006
-DEBUG:  Skipping target shard interval 13300004 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300004 since it doesn't have the same shard range with the select anchor shard interval 13300003
 DEBUG:  predicate pruning for shardId 13300000
 DEBUG:  predicate pruning for shardId 13300001
 DEBUG:  predicate pruning for shardId 13300002
 DEBUG:  predicate pruning for shardId 13300004
 DEBUG:  predicate pruning for shardId 13300005
 DEBUG:  predicate pruning for shardId 13300006
-DEBUG:  Skipping target shard interval 13300005 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300005 since it doesn't have the same shard range with the select anchor shard interval 13300003
 DEBUG:  predicate pruning for shardId 13300000
 DEBUG:  predicate pruning for shardId 13300001
 DEBUG:  predicate pruning for shardId 13300002
 DEBUG:  predicate pruning for shardId 13300004
 DEBUG:  predicate pruning for shardId 13300005
 DEBUG:  predicate pruning for shardId 13300006
-DEBUG:  Skipping target shard interval 13300006 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300006 since it doesn't have the same shard range with the select anchor shard interval 13300003
 DEBUG:  predicate pruning for shardId 13300000
 DEBUG:  predicate pruning for shardId 13300001
 DEBUG:  predicate pruning for shardId 13300002
@@ -1214,21 +1233,21 @@ DEBUG:  predicate pruning for shardId 13300003
 DEBUG:  predicate pruning for shardId 13300005
 DEBUG:  predicate pruning for shardId 13300006
 DEBUG:  predicate pruning for shardId 13300007
-DEBUG:  Skipping target shard interval 13300009 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300009 since it doesn't have the same shard range with the select anchor shard interval 13300000
 DEBUG:  predicate pruning for shardId 13300001
 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
-DEBUG:  Skipping target shard interval 13300010 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300010 since it doesn't have the same shard range with the select anchor shard interval 13300000
 DEBUG:  predicate pruning for shardId 13300001
 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
-DEBUG:  Skipping target shard interval 13300011 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300011 since it doesn't have the same shard range with the select anchor shard interval 13300000
 DEBUG:  ProcessQuery
 DEBUG:  Plan is router executable
 DEBUG:  CommitTransactionCommand
@@ -1256,21 +1275,21 @@ DEBUG:  predicate pruning for shardId 13300003
 DEBUG:  predicate pruning for shardId 13300005
 DEBUG:  predicate pruning for shardId 13300006
 DEBUG:  predicate pruning for shardId 13300007
-DEBUG:  Skipping target shard interval 13300009 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300009 since it doesn't have the same shard range with the select anchor shard interval 13300000
 DEBUG:  predicate pruning for shardId 13300001
 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
-DEBUG:  Skipping target shard interval 13300010 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300010 since it doesn't have the same shard range with the select anchor shard interval 13300000
 DEBUG:  predicate pruning for shardId 13300001
 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
-DEBUG:  Skipping target shard interval 13300011 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300011 since it doesn't have the same shard range with the select anchor shard interval 13300000
 DEBUG:  ProcessQuery
 DEBUG:  Plan is router executable
 DEBUG:  CommitTransactionCommand
@@ -1732,15 +1751,15 @@ DEBUG:  distributed statement: INSERT INTO public.raw_events_first_13300000 AS c
 DEBUG:  predicate pruning for shardId 13300005
 DEBUG:  predicate pruning for shardId 13300006
 DEBUG:  predicate pruning for shardId 13300007
-DEBUG:  Skipping target shard interval 13300001 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300001 since it doesn't have the same shard range with the select anchor shard interval 13300004
 DEBUG:  predicate pruning for shardId 13300005
 DEBUG:  predicate pruning for shardId 13300006
 DEBUG:  predicate pruning for shardId 13300007
-DEBUG:  Skipping target shard interval 13300002 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300002 since it doesn't have the same shard range with the select anchor shard interval 13300004
 DEBUG:  predicate pruning for shardId 13300005
 DEBUG:  predicate pruning for shardId 13300006
 DEBUG:  predicate pruning for shardId 13300007
-DEBUG:  Skipping target shard interval 13300003 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300003 since it doesn't have the same shard range with the select anchor shard interval 13300004
 DEBUG:  ProcessQuery
 DEBUG:  Plan is router executable
 DEBUG:  CommitTransactionCommand
@@ -1786,11 +1805,11 @@ DEBUG:  name: unnamed; blockState:       DEFAULT; state: INPROGR, xid/subid/cid:
 DEBUG:  predicate pruning for shardId 13300004
 DEBUG:  predicate pruning for shardId 13300005
 DEBUG:  predicate pruning for shardId 13300007
-DEBUG:  Skipping target shard interval 13300000 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300000 since it doesn't have the same shard range with the select anchor shard interval 13300006
 DEBUG:  predicate pruning for shardId 13300004
 DEBUG:  predicate pruning for shardId 13300005
 DEBUG:  predicate pruning for shardId 13300007
-DEBUG:  Skipping target shard interval 13300001 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300001 since it doesn't have the same shard range with the select anchor shard interval 13300006
 DEBUG:  predicate pruning for shardId 13300004
 DEBUG:  predicate pruning for shardId 13300005
 DEBUG:  predicate pruning for shardId 13300007
@@ -1798,7 +1817,7 @@ DEBUG:  distributed statement: INSERT INTO public.raw_events_first_13300002 AS c
 DEBUG:  predicate pruning for shardId 13300004
 DEBUG:  predicate pruning for shardId 13300005
 DEBUG:  predicate pruning for shardId 13300007
-DEBUG:  Skipping target shard interval 13300003 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300003 since it doesn't have the same shard range with the select anchor shard interval 13300006
 DEBUG:  ProcessQuery
 DEBUG:  Plan is router executable
 DEBUG:  CommitTransactionCommand
@@ -1866,15 +1885,15 @@ DEBUG:  distributed statement: INSERT INTO public.raw_events_first_13300000 AS c
 DEBUG:  predicate pruning for shardId 13300005
 DEBUG:  predicate pruning for shardId 13300006
 DEBUG:  predicate pruning for shardId 13300007
-DEBUG:  Skipping target shard interval 13300001 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300001 since it doesn't have the same shard range with the select anchor shard interval 13300004
 DEBUG:  predicate pruning for shardId 13300005
 DEBUG:  predicate pruning for shardId 13300006
 DEBUG:  predicate pruning for shardId 13300007
-DEBUG:  Skipping target shard interval 13300002 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300002 since it doesn't have the same shard range with the select anchor shard interval 13300004
 DEBUG:  predicate pruning for shardId 13300005
 DEBUG:  predicate pruning for shardId 13300006
 DEBUG:  predicate pruning for shardId 13300007
-DEBUG:  Skipping target shard interval 13300003 since SELECT query for it pruned away
+DEBUG:  Skipping target shard interval 13300003 since it doesn't have the same shard range with the select anchor shard interval 13300004
 DEBUG:  ProcessQuery
 DEBUG:  Plan is router executable
 DEBUG:  CommitTransactionCommand