Fix incorrect union all pushdown issue

src/backend/distributed/planner/query_pushdown_planning.c

@@ -596,7 +596,7 @@ DeferErrorIfUnsupportedSubqueryPushdown(Query *originalQuery,
 	{
 		return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
 							 "complex joins are only supported when all distributed tables are "
-							 "joined on their distribution columns with equal operator",
+							 "co-located and joined on their distribution columns",
 							 NULL, NULL);
 	}
 

src/backend/distributed/planner/relation_restriction_equivalence.c

@@ -10,7 +10,9 @@
  */
 #include "postgres.h"
 
+#include "distributed/colocation_utils.h"
 #include "distributed/distributed_planner.h"
+#include "distributed/listutils.h"
 #include "distributed/metadata_cache.h"
 #include "distributed/multi_logical_planner.h"
 #include "distributed/multi_logical_optimizer.h"
@@ -133,6 +135,8 @@ static void ListConcatUniqueAttributeClassMemberLists(AttributeEquivalenceClass
 													  secondClass);
 static Index RelationRestrictionPartitionKeyIndex(RelationRestriction *
 												  relationRestriction);
+static bool AllRelationsInRestrictionContextColocated(RelationRestrictionContext *
+													  restrictionContext);
 static RelationRestrictionContext * FilterRelationRestrictionContext(
 	RelationRestrictionContext *relationRestrictionContext,
 	Relids
@@ -345,8 +349,20 @@ SafeToPushdownUnionSubquery(PlannerRestrictionContext *plannerRestrictionContext
 	allAttributeEquivalenceList = lappend(allAttributeEquivalenceList,
 										  attributeEquivalance);
 
-	return EquivalenceListContainsRelationsEquality(allAttributeEquivalenceList,
+	if (!EquivalenceListContainsRelationsEquality(allAttributeEquivalenceList,
-													restrictionContext);
+												  restrictionContext))
+	{
+		/* cannot confirm equality for all distribution colums */
+		return false;
+	}
+
+	if (!AllRelationsInRestrictionContextColocated(restrictionContext))
+	{
+		/* distribution columns are equal, but tables are not co-located */
+		return false;
+	}
+
+	return true;
 }
 
 
@@ -1650,6 +1666,42 @@ RelationRestrictionPartitionKeyIndex(RelationRestriction *relationRestriction)
 }
 
 
+/*
+ * AllRelationsInRestrictionContextColocated determines whether all of the relations in the
+ * given relation restrictions list are co-located.
+ */
+static bool
+AllRelationsInRestrictionContextColocated(RelationRestrictionContext *restrictionContext)
+{
+	RelationRestriction *relationRestriction = NULL;
+	int initialColocationId = INVALID_COLOCATION_ID;
+
+	/* check whether all relations exists in the main restriction list */
+	foreach_ptr(relationRestriction, restrictionContext->relationRestrictionList)
+	{
+		Oid relationId = relationRestriction->relationId;
+
+		if (PartitionMethod(relationId) == DISTRIBUTE_BY_NONE)
+		{
+			continue;
+		}
+
+		int colocationId = TableColocationId(relationId);
+
+		if (initialColocationId == INVALID_COLOCATION_ID)
+		{
+			initialColocationId = colocationId;
+		}
+		else if (colocationId != initialColocationId)
+		{
+			return false;
+		}
+	}
+
+	return true;
+}
+
+
 /*
  * RelationIdList returns list of unique relation ids in query tree.
  */

src/test/regress/expected/ch_bench_having.out

@@ -274,14 +274,14 @@ where   s_order_cnt > (select sum(s_order_cnt) * .005 as where_query from stock)
 group by s_i_id
 having   (select max(s_order_cnt) > 2 as having_query from stock where s_i_id = s.s_i_id)
 order by s_i_id;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- We don't support correlated subqueries in having
 select     s_i_id, sum(s_order_cnt) as ordercount
 from     stock s
 group by s_i_id
 having   (select max(s_order_cnt) > 2 as having_query from stock where s_i_id = s.s_i_id)
 order by s_i_id;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 DROP TABLE stock;
 CREATE TABLE stock (
   s_w_id int NOT NULL,

src/test/regress/expected/ch_bench_having_mx.out

@@ -279,14 +279,14 @@ where   s_order_cnt > (select sum(s_order_cnt) * .005 as where_query from stock)
 group by s_i_id
 having   (select max(s_order_cnt) > 2 as having_query from stock where s_i_id = s.s_i_id)
 order by s_i_id;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- We don't support correlated subqueries in having
 select     s_i_id, sum(s_order_cnt) as ordercount
 from     stock s
 group by s_i_id
 having   (select max(s_order_cnt) > 2 as having_query from stock where s_i_id = s.s_i_id)
 order by s_i_id;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 \c - - - :master_port
 SET citus.replication_model TO streaming;
 SET citus.shard_replication_factor to 1;

src/test/regress/expected/ch_bench_subquery_repartition.out

@@ -116,7 +116,7 @@ select  s_i_id
     where
         s_i_id in (select i_im_id from item)
         AND s_i_id = ol_i_id;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- Subquery + repartion is supported when it is a NOT IN query where the subquery
 -- returns unique results
 select  s_i_id
@@ -124,7 +124,7 @@ select  s_i_id
     where
         s_i_id not in (select i_id from item)
         AND s_i_id = ol_i_id;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- Subquery + repartion is not supported when it is a NOT IN where the subquery
 -- doesn't return unique results
 select  s_i_id
@@ -132,7 +132,7 @@ select  s_i_id
     where
         s_i_id not in (select i_im_id from item)
         AND s_i_id = ol_i_id;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- Actual CHbenCHmark query is supported
 select   su_name, su_address
 from     supplier, nation

src/test/regress/expected/dml_recursive.out

@@ -262,7 +262,7 @@ FROM
 	) as foo
 RETURNING *;
 DEBUG:  generating subplan 15_1 for subquery SELECT dept FROM recursive_dml_queries.second_distributed_table
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- again a corrolated subquery
 -- this time distribution key eq. exists
 -- however recursive planning is prevented due to correlated subqueries
@@ -292,7 +292,7 @@ FROM
 	) as baz
 	) as foo WHERE second_distributed_table.tenant_id = foo.tenant_id
 RETURNING *;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- we don't support subqueries/CTEs inside VALUES
 INSERT INTO
 	second_distributed_table (tenant_id, dept)

src/test/regress/expected/multi_dropped_column_aliases.out

@@ -59,4 +59,4 @@ FROM   (customer LEFT OUTER JOIN orders ON (c_custkey = o_custkey)) AS
        test(c_custkey, c_nationkey)
        INNER JOIN lineitem ON (test.c_custkey = l_orderkey)
 LIMIT 10;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns

src/test/regress/expected/multi_shard_modify.out

@@ -83,7 +83,7 @@ SELECT create_distributed_table('temp_nations', 'name', 'hash');
 SELECT master_modify_multiple_shards('DELETE FROM multi_shard_modify_test USING temp_nations WHERE multi_shard_modify_test.t_value = temp_nations.key AND temp_nations.name = ''foobar'' ');
 WARNING:  master_modify_multiple_shards is deprecated and will be removed in a future release.
 HINT:  Run the command directly
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- commands with a RETURNING clause are unsupported
 SELECT master_modify_multiple_shards('DELETE FROM multi_shard_modify_test WHERE t_key = 3 RETURNING *');
 WARNING:  master_modify_multiple_shards is deprecated and will be removed in a future release.

src/test/regress/expected/multi_shard_update_delete.out

@@ -656,7 +656,7 @@ WHERE  user_id IN (SELECT user_id
 UPDATE users_test_table
 SET    value_2 = (SELECT value_3
                   FROM   users_test_table);
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 UPDATE users_test_table
 SET value_2 = 2
 WHERE
@@ -671,7 +671,7 @@ WHERE
            GROUP BY
               user_id
           );
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 UPDATE users_test_table
 SET (value_1, value_2) = (2,1)
 WHERE  user_id IN

src/test/regress/expected/multi_subquery.out

@@ -426,7 +426,7 @@ FROM events_table t1
 LEFT JOIN users_table t2 ON t1.user_id > t2.user_id
 ORDER BY 1 DESC, 2 DESC, 3 DESC, 4 DESC
 LIMIT 5;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- outer joins on reference tables with expressions should work
 SELECT DISTINCT ON (t1.user_id) t1.user_id, t2.value_1, t2.value_2, t2.value_3
 FROM events_table t1
@@ -467,7 +467,7 @@ SELECT DISTINCT ON (t1.user_id) t1.user_id, t2.value_1, t2.value_2, t2.value_3
  LEFT JOIN users_reference_table t2 ON t1.user_id = trunc(t2.user_id)
  ORDER BY 1 DESC, 2 DESC, 3 DESC, 4 DESC
  LIMIT 5;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- outer joins as subqueries should work
 -- https://github.com/citusdata/citus/issues/2739
 SELECT user_id, value_1, event_type

src/test/regress/expected/multi_subquery_behavioral_analytics.out

@@ -2011,7 +2011,7 @@ FROM (
   GROUP BY user_id
 ) q
 ORDER BY 2 DESC, 1;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- note that the following query has both equi-joins on the partition keys
 -- and non-equi-joins on other columns. We now support query filters
 -- having non-equi-joins as long as they have equi-joins on partition keys.
@@ -2100,7 +2100,7 @@ FROM
     events_table.time = users_table.time AND
     events_table.value_2 IN (0, 4)
   ) as foo;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- we can even allow that on top level joins
 SELECT
   count(*)
@@ -2150,7 +2150,7 @@ FROM
     events_table.value_2 IN (1, 5)
   ) as bar
 WHERE foo.event_type = bar.event_type;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- DISTINCT in the outer query and DISTINCT in the subquery
 SELECT
     DISTINCT users_ids.user_id

src/test/regress/expected/multi_subquery_complex_queries.out

@@ -1616,7 +1616,7 @@ FROM
 ORDER BY
   user_id DESC, lastseen DESC
 LIMIT 10;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- not pushdownable since lower LATERAL JOIN is not on the partition key
 -- not recursively plannable due to LATERAL join where there is a reference
 -- from an outer query

src/test/regress/expected/multi_subquery_complex_reference_clause.out

@@ -1420,7 +1420,7 @@ WHERE
 GROUP BY 1
 ORDER BY 2 DESC, 1 DESC
 LIMIT 5;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 SELECT foo.user_id FROM
 (
   SELECT m.user_id, random() FROM users_table m JOIN events_reference_table r ON int4eq(m.user_id, r.user_id)

src/test/regress/expected/multi_subquery_in_where_clause.out

@@ -80,7 +80,7 @@ GROUP BY user_id
 HAVING count(*) > 1
 ORDER BY user_id
 LIMIT 5;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- subqueries in where with ALL operator
 SELECT
   user_id
@@ -465,7 +465,7 @@ SELECT user_id, value_2 FROM users_table WHERE
 		group by e1.user_id
 		HAVING sum(submit_card_info) > 0
 );
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- left leaf query does not return partition key
 SELECT
 	user_id
@@ -561,7 +561,7 @@ FROM (
   GROUP BY user_id
 ) q
 ORDER BY 2 DESC, 1;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- subquery in where clause doesn't have a relation, but is constant
 SELECT
   user_id
@@ -675,5 +675,5 @@ SELECT user_id, value_2 FROM users_table WHERE
   AND NOT EXISTS (SELECT user_id FROM events_table WHERE event_type=1 AND value_3 > 1 AND test_join_function(events_table.user_id, users_table.user_id))
 ORDER BY 1 DESC, 2 DESC
 LIMIT 3;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 DROP FUNCTION test_join_function(int,int);

src/test/regress/expected/non_colocated_subquery_joins.out

@@ -629,7 +629,7 @@ SELECT true AS valid FROM explain_json_2($$
 $$);
 DEBUG:  generating subplan 66_1 for subquery SELECT value_1, random() AS random FROM public.users_table
 DEBUG:  Plan 66 query after replacing subqueries and CTEs: SELECT count(*) AS count FROM ((public.users_table u1 JOIN public.users_table u2 USING (value_1)) a(value_1, user_id, "time", value_2, value_3, value_4, user_id_1, time_1, value_2_1, value_3_1, value_4_1) JOIN (SELECT intermediate_result.value_1, intermediate_result.random FROM read_intermediate_result('66_1'::text, 'binary'::citus_copy_format) intermediate_result(value_1 integer, random double precision)) u3 USING (value_1))
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- a very similar query to the above
 -- however, this time we users a subquery instead of join alias, and it works
 SELECT true AS valid FROM explain_json_2($$
@@ -926,7 +926,7 @@ DEBUG:  Router planner cannot handle multi-shard select queries
 DEBUG:  skipping recursive planning for the subquery since it contains references to outer queries
 DEBUG:  skipping recursive planning for the subquery since it contains references to outer queries
 DEBUG:  skipping recursive planning for the subquery since it contains references to outer queries
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- similar to the above, make sure that we skip recursive planning when
 -- the subquery doesn't have any tables
 SELECT true AS valid FROM explain_json_2($$
@@ -945,7 +945,7 @@ DEBUG:  Router planner cannot handle multi-shard select queries
 DEBUG:  skipping recursive planning for the subquery since it contains references to outer queries
 DEBUG:  skipping recursive planning for the subquery since it contains references to outer queries
 DEBUG:  skipping recursive planning for the subquery since it contains references to outer queries
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- similar to the above, make sure that we skip recursive planning when
 -- the subquery contains only intermediate results
 SELECT *
@@ -1061,14 +1061,14 @@ SELECT create_distributed_table('table1','tenant_id');
 -- all of the above queries are non-colocated subquery joins
 -- because the views are replaced with subqueries
 UPDATE table2 SET id=20 FROM table1_view WHERE table1_view.id=table2.id;
-DEBUG:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+DEBUG:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 DEBUG:  Router planner cannot handle multi-shard select queries
 DEBUG:  generating subplan 117_1 for subquery SELECT table1.id, table1.tenant_id FROM non_colocated_subquery.table1 WHERE (table1.id OPERATOR(pg_catalog.<) 100)
 DEBUG:  Plan 117 query after replacing subqueries and CTEs: UPDATE non_colocated_subquery.table2 SET id = 20 FROM (SELECT intermediate_result.id, intermediate_result.tenant_id FROM read_intermediate_result('117_1'::text, 'binary'::citus_copy_format) intermediate_result(id integer, tenant_id integer)) table1_view WHERE (table1_view.id OPERATOR(pg_catalog.=) table2.id)
 DEBUG:  Creating router plan
 DEBUG:  Plan is router executable
 UPDATE table2_p1 SET id=20 FROM table1_view WHERE table1_view.id=table2_p1.id;
-DEBUG:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+DEBUG:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 DEBUG:  Router planner cannot handle multi-shard select queries
 DEBUG:  generating subplan 119_1 for subquery SELECT table1.id, table1.tenant_id FROM non_colocated_subquery.table1 WHERE (table1.id OPERATOR(pg_catalog.<) 100)
 DEBUG:  Plan 119 query after replacing subqueries and CTEs: UPDATE non_colocated_subquery.table2_p1 SET id = 20 FROM (SELECT intermediate_result.id, intermediate_result.tenant_id FROM read_intermediate_result('119_1'::text, 'binary'::citus_copy_format) intermediate_result(id integer, tenant_id integer)) table1_view WHERE (table1_view.id OPERATOR(pg_catalog.=) table2_p1.id)

src/test/regress/expected/set_operations.out

@@ -14,6 +14,13 @@ SELECT create_reference_table('ref');
  
 (1 row)
 
+CREATE TABLE test_not_colocated (LIKE test);
+SELECT create_distributed_table('test_not_colocated', 'x', colocate_with := 'none');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
 INSERT INTO test VALUES (1,1), (2,2);
 INSERT INTO ref VALUES (2,2), (3,3);
 -- top-level set operations are supported through recursive planning
@@ -1045,11 +1052,49 @@ DEBUG:  Plan is router executable
 ---
 (0 rows)
 
+-- queries on non-colocated tables that would push down if they were not colocated are recursivelu planned
+SELECT * FROM (SELECT * FROM test UNION SELECT * FROM test_not_colocated) u ORDER BY 1,2;
+DEBUG:  Router planner cannot handle multi-shard select queries
+DEBUG:  Router planner cannot handle multi-shard select queries
+DEBUG:  generating subplan 188_1 for subquery SELECT x, y FROM recursive_union.test
+DEBUG:  Router planner cannot handle multi-shard select queries
+DEBUG:  generating subplan 188_2 for subquery SELECT x, y FROM recursive_union.test_not_colocated
+DEBUG:  Creating router plan
+DEBUG:  Plan is router executable
+DEBUG:  generating subplan 188_3 for subquery SELECT intermediate_result.x, intermediate_result.y FROM read_intermediate_result('188_1'::text, 'binary'::citus_copy_format) intermediate_result(x integer, y integer) UNION SELECT intermediate_result.x, intermediate_result.y FROM read_intermediate_result('188_2'::text, 'binary'::citus_copy_format) intermediate_result(x integer, y integer)
+DEBUG:  Plan 188 query after replacing subqueries and CTEs: SELECT x, y FROM (SELECT intermediate_result.x, intermediate_result.y FROM read_intermediate_result('188_3'::text, 'binary'::citus_copy_format) intermediate_result(x integer, y integer)) u ORDER BY x, y
+DEBUG:  Creating router plan
+DEBUG:  Plan is router executable
+ x | y 
+---+---
+ 1 | 1
+ 2 | 2
+(2 rows)
+
+SELECT * FROM (SELECT * FROM test UNION ALL SELECT * FROM test_not_colocated) u ORDER BY 1,2;
+DEBUG:  Router planner cannot handle multi-shard select queries
+DEBUG:  Router planner cannot handle multi-shard select queries
+DEBUG:  generating subplan 192_1 for subquery SELECT x, y FROM recursive_union.test
+DEBUG:  Router planner cannot handle multi-shard select queries
+DEBUG:  generating subplan 192_2 for subquery SELECT x, y FROM recursive_union.test_not_colocated
+DEBUG:  Creating router plan
+DEBUG:  Plan is router executable
+DEBUG:  generating subplan 192_3 for subquery SELECT intermediate_result.x, intermediate_result.y FROM read_intermediate_result('192_1'::text, 'binary'::citus_copy_format) intermediate_result(x integer, y integer) UNION ALL SELECT intermediate_result.x, intermediate_result.y FROM read_intermediate_result('192_2'::text, 'binary'::citus_copy_format) intermediate_result(x integer, y integer)
+DEBUG:  Plan 192 query after replacing subqueries and CTEs: SELECT x, y FROM (SELECT intermediate_result.x, intermediate_result.y FROM read_intermediate_result('192_3'::text, 'binary'::citus_copy_format) intermediate_result(x integer, y integer)) u ORDER BY x, y
+DEBUG:  Creating router plan
+DEBUG:  Plan is router executable
+ x | y 
+---+---
+ 1 | 1
+ 2 | 2
+(2 rows)
+
 RESET client_min_messages;
 DROP SCHEMA recursive_union CASCADE;
-NOTICE:  drop cascades to 5 other objects
+NOTICE:  drop cascades to 6 other objects
 DETAIL:  drop cascades to table test
 drop cascades to table ref
+drop cascades to table test_not_colocated
 drop cascades to view set_view_recursive
 drop cascades to view set_view_pushdown
 drop cascades to view set_view_recursive_second

src/test/regress/output/multi_outer_join.source

@@ -272,7 +272,7 @@ SELECT
 	count(*)
 FROM
 	multi_outer_join_left a LEFT JOIN multi_outer_join_right b ON (l_nationkey = r_nationkey);
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- Anti-join should return customers for which there is no row in the right table
 SELECT
 	min(l_custkey), max(l_custkey)

src/test/regress/output/multi_outer_join_reference.source

@@ -252,7 +252,7 @@ SELECT
 	count(*)
 FROM
 	multi_outer_join_left_hash a LEFT JOIN multi_outer_join_right_hash b ON (l_nationkey = r_nationkey);
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- Anti-join should return customers for which there is no row in the right table
 SELECT
 	min(l_custkey), max(l_custkey)
@@ -326,7 +326,7 @@ FROM
 	LEFT JOIN multi_outer_join_right_reference r1 ON (l1.l_custkey = r1.r_custkey)
 	LEFT JOIN multi_outer_join_right_reference r2 ON (l1.l_custkey  = r2.r_custkey)
 	RIGHT JOIN multi_outer_join_left_hash l2 ON (r2.r_custkey = l2.l_custkey);
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- add an anti-join, this should also error out
 SELECT
 	*
@@ -337,7 +337,7 @@ FROM
 	RIGHT JOIN multi_outer_join_left_hash l2 ON (r2.r_custkey = l2.l_custkey)
 WHERE
 	r1.r_custkey is NULL;
-ERROR:  complex joins are only supported when all distributed tables are joined on their distribution columns with equal operator
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
 -- Three way join 2-1-1 (broadcast + broadcast join) should work
 SELECT
 	l_custkey, r_custkey, t_custkey

src/test/regress/sql/set_operations.sql

@@ -7,6 +7,9 @@ SELECT create_distributed_table('test', 'x');
 CREATE TABLE recursive_union.ref (a int, b int);
 SELECT create_reference_table('ref');
 
+CREATE TABLE test_not_colocated (LIKE test);
+SELECT create_distributed_table('test_not_colocated', 'x', colocate_with := 'none');
+
 INSERT INTO test VALUES (1,1), (2,2);
 INSERT INTO ref VALUES (2,2), (3,3);
 
@@ -169,5 +172,9 @@ SELECT * FROM set_view_recursive_second ORDER BY 1,2;
 -- this should create lots of recursive calls since both views and set operations lead to recursive plans :)
 ((SELECT x FROM set_view_recursive_second) INTERSECT (SELECT * FROM set_view_recursive)) EXCEPT (SELECT * FROM set_view_pushdown);
 
+-- queries on non-colocated tables that would push down if they were not colocated are recursivelu planned
+SELECT * FROM (SELECT * FROM test UNION SELECT * FROM test_not_colocated) u ORDER BY 1,2;
+SELECT * FROM (SELECT * FROM test UNION ALL SELECT * FROM test_not_colocated) u ORDER BY 1,2;
+
 RESET client_min_messages;
 DROP SCHEMA recursive_union CASCADE;