Add more tests

src/backend/distributed/planner/local_distributed_join_planner.c

@@ -174,6 +174,10 @@ static bool AllRangeTableEntriesHaveUniqueIndex(List *rangeTableEntryDetailsList
 static bool FirstIsSuperSetOfSecond(List *firstIntList, List *secondIntList);
 static void ConvertRTEsToSubquery(List *rangeTableEntryDetailsList,
 								  RecursivePlanningContext *context);
+static int ResultRTEIdentity(Query *query);
+static List * RTEListToConvert(ConversionCandidates *conversionCandidates,
+							   ConversionChoice conversionChoice);
+
 
 /*
  * RecursivelyPlanLocalTableJoins gets a query and the planner
@@ -182,17 +186,13 @@ static void ConvertRTEsToSubquery(List *rangeTableEntryDetailsList,
  */
 void
 RecursivelyPlanLocalTableJoins(Query *query,
-							   RecursivePlanningContext *context, List *rangeTableList)
+							   RecursivePlanningContext *context)
 {
 	PlannerRestrictionContext *plannerRestrictionContext =
 		GetPlannerRestrictionContext(context);
 
-	int resultRTEIdentity = INVALID_RTE_IDENTITY;
+	List *rangeTableList = query->rtable;
-	if (IsModifyCommand(query))
+	int resultRTEIdentity = ResultRTEIdentity(query);
-	{
-		RangeTblEntry *resultRTE = ExtractResultRelationRTE(query);
-		resultRTEIdentity = GetRTEIdentity(resultRTE);
-	}
 	ConversionCandidates *conversionCandidates =
 		CreateConversionCandidates(plannerRestrictionContext,
 								   rangeTableList, resultRTEIdentity);
@@ -200,14 +200,48 @@ RecursivelyPlanLocalTableJoins(Query *query,
 	ConversionChoice conversionChoise =
 		GetConversionChoice(conversionCandidates, plannerRestrictionContext);
 
-	if (conversionChoise == CONVERT_LOCAL_TABLES)
+
+	List *rteListToConvert = RTEListToConvert(conversionCandidates, conversionChoise);
+	ConvertRTEsToSubquery(rteListToConvert, context);
+}
+
+
+/*
+ * ResultRTEIdentity returns the result RTE's identity if it exists,
+ * otherwise it returns INVALID_RTE_INDENTITY
+ */
+static int
+ResultRTEIdentity(Query *query)
+{
+	int resultRTEIdentity = INVALID_RTE_IDENTITY;
+	if (IsModifyCommand(query))
 	{
-		ConvertRTEsToSubquery(conversionCandidates->localTableList, context);
+		RangeTblEntry *resultRTE = ExtractResultRelationRTE(query);
+		resultRTEIdentity = GetRTEIdentity(resultRTE);
+	}
+	return resultRTEIdentity;
+}
+
+
+/*
+ * RTEListToConvert to converts returns a list of RTEs that should
+ * be converted to a subquery.
+ */
+static List *
+RTEListToConvert(ConversionCandidates *conversionCandidates, ConversionChoice
+				 conversionChoice)
+{
+	List *rtesToConvert = NIL;
+	if (conversionChoice == CONVERT_LOCAL_TABLES)
+	{
+		rtesToConvert = list_concat(rtesToConvert, conversionCandidates->localTableList);
 	}
 	else
 	{
-		ConvertRTEsToSubquery(conversionCandidates->distributedTableList, context);
+		rtesToConvert = list_concat(rtesToConvert,
+									conversionCandidates->distributedTableList);
 	}
+	return rtesToConvert;
 }
 
 
@@ -305,9 +339,7 @@ AllRangeTableEntriesHaveUniqueIndex(List *rangeTableEntryDetailsList)
  * convert local-dist table joins to subqueries.
  */
 bool
-ShouldConvertLocalTableJoinsToSubqueries(Query *query, List *rangeTableList,
+ShouldConvertLocalTableJoinsToSubqueries(List *rangeTableList)
-										 PlannerRestrictionContext *
-										 plannerRestrictionContext)
 {
 	if (LocalTableJoinPolicy == LOCAL_JOIN_POLICY_NEVER)
 	{
@@ -315,7 +347,7 @@ ShouldConvertLocalTableJoinsToSubqueries(Query *query, List *rangeTableList,
 		return false;
 	}
 
-	if (!ContainsTableToBeConvertedToSubquery(rangeTableList))
+	if (!ContainsLocalTableDistributedTableJoin(rangeTableList))
 	{
 		return false;
 	}

src/backend/distributed/planner/multi_logical_optimizer.c

@@ -292,6 +292,8 @@ static SortGroupClause * CreateSortGroupClause(Var *column);
 static const char * CountDistinctHashFunctionName(Oid argumentType);
 static int CountDistinctStorageSize(double approximationErrorRate);
 static Const * MakeIntegerConstInt64(int64 integerValue);
+static Const * MakeIntegerConst(int32 integerValue);
+
 
 /* Local functions forward declarations for aggregate expression checks */
 static bool HasNonDistributableAggregates(MultiNode *logicalPlanNode);
@@ -3789,7 +3791,7 @@ CountDistinctStorageSize(double approximationErrorRate)
 
 
 /* Makes an integer constant node from the given value, and returns that node. */
-Const *
+static Const *
 MakeIntegerConst(int32 integerValue)
 {
 	const int typeCollationId = get_typcollation(INT4OID);

src/backend/distributed/planner/multi_physical_planner.c

@@ -3639,7 +3639,7 @@ NodeIsRangeTblRefReferenceTable(Node *node, List *rangeTableList)
 
 /*
  * FetchEqualityAttrNumsForRTE fetches the attribute numbers from quals
- * which has equality operator
+ * which have an equality operator
  */
 List *
 FetchEqualityAttrNumsForRTE(Node *node)
@@ -3665,8 +3665,8 @@ FetchEqualityAttrNumsForRTE(Node *node)
 
 
 /*
- * FetchEqualityAttrNumsForList fetches the constant equality numbers
+ * FetchEqualityAttrNumsForList fetches the attribute numbers of expression
- * from the given node list.
+ * of the form "= constant" from the given node list.
  */
 static List *
 FetchEqualityAttrNumsForList(List *nodeList)
@@ -3696,8 +3696,8 @@ FetchEqualityAttrNumsForList(List *nodeList)
 
 
 /*
- * FetchEqualityAttrNumsForRTEOpExpr fetches the attribute numbers from opExpr
+ * FetchEqualityAttrNumsForRTEOpExpr fetches the attribute numbers of expression
- * which has equality operator.
+ * of the form "= constant" from the given opExpr.
  */
 static List *
 FetchEqualityAttrNumsForRTEOpExpr(OpExpr *opExpr)
@@ -3718,8 +3718,8 @@ FetchEqualityAttrNumsForRTEOpExpr(OpExpr *opExpr)
 
 
 /*
- * FetchEqualityAttrNumsForRTEBoolExpr fetches the attribute numbers from boolExpr
+ * FetchEqualityAttrNumsForRTEBoolExpr fetches the attribute numbers of expression
- * which has equality operator
+ * of the form "= constant" from the given boolExpr.
  */
 static List *
 FetchEqualityAttrNumsForRTEBoolExpr(BoolExpr *boolExpr)

src/backend/distributed/planner/recursive_planning.c

@@ -191,9 +191,7 @@ static Query * BuildReadIntermediateResultsQuery(List *targetEntryList,
 												 List *columnAliasList,
 												 Const *resultIdConst, Oid functionOid,
 												 bool useBinaryCopyFormat);
-static void UpdateVarNosInNode(Query *query, Index newVarNo);
+static void UpdateVarNosInNode(Node *node, Index newVarNo);
-static void GetRangeTableEntriesFromJoinTree(Node *joinNode, List *rangeTableList,
-											 List **joinRangeTableEntries);
 static Query * CreateOuterSubquery(RangeTblEntry *rangeTableEntry,
 								   List *outerSubqueryTargetList);
 static List * GenerateRequiredColNamesFromTargetList(List *targetList);
@@ -354,22 +352,14 @@ RecursivelyPlanSubqueriesAndCTEs(Query *query, RecursivePlanningContext *context
 	}
 
 
-	PlannerRestrictionContext *plannerRestrictionContext =
+	if (ShouldConvertLocalTableJoinsToSubqueries(query->rtable))
-		context->plannerRestrictionContext;
-
-	List *rangeTableList = NIL;
-	GetRangeTableEntriesFromJoinTree((Node *) query->jointree, query->rtable,
-									 &rangeTableList);
-
-	if (ShouldConvertLocalTableJoinsToSubqueries(query, rangeTableList,
-												 plannerRestrictionContext))
 	{
 		/*
 		 * Logical planner cannot handle "local_table" [OUTER] JOIN "dist_table", or
 		 * a query with local table/citus local table and subquery. We convert local/citus local
 		 * tables to a subquery until they can be planned.
 		 */
-		RecursivelyPlanLocalTableJoins(query, context, rangeTableList);
+		RecursivelyPlanLocalTableJoins(query, context);
 	}
 
 
@@ -388,50 +378,6 @@ GetPlannerRestrictionContext(RecursivePlanningContext *recursivePlanningContext)
 }
 
 
-/*
- * GetRangeTableEntriesFromJoinTree gets the range table entries that are
- * on the given join tree.
- */
-static void
-GetRangeTableEntriesFromJoinTree(Node *joinNode, List *rangeTableList,
-								 List **joinRangeTableEntries)
-{
-	if (joinNode == NULL)
-	{
-		return;
-	}
-	else if (IsA(joinNode, FromExpr))
-	{
-		FromExpr *fromExpr = (FromExpr *) joinNode;
-		Node *fromElement;
-
-		foreach_ptr(fromElement, fromExpr->fromlist)
-		{
-			GetRangeTableEntriesFromJoinTree(fromElement, rangeTableList,
-											 joinRangeTableEntries);
-		}
-	}
-	else if (IsA(joinNode, JoinExpr))
-	{
-		JoinExpr *joinExpr = (JoinExpr *) joinNode;
-		GetRangeTableEntriesFromJoinTree(joinExpr->larg, rangeTableList,
-										 joinRangeTableEntries);
-		GetRangeTableEntriesFromJoinTree(joinExpr->rarg, rangeTableList,
-										 joinRangeTableEntries);
-	}
-	else if (IsA(joinNode, RangeTblRef))
-	{
-		int rangeTableIndex = ((RangeTblRef *) joinNode)->rtindex;
-		RangeTblEntry *rte = rt_fetch(rangeTableIndex, rangeTableList);
-		*joinRangeTableEntries = lappend(*joinRangeTableEntries, rte);
-	}
-	else
-	{
-		pg_unreachable();
-	}
-}
-
-
 /*
  * ShouldRecursivelyPlanNonColocatedSubqueries returns true if the input query contains joins
  * that are not on the distribution key.
@@ -1476,7 +1422,8 @@ ReplaceRTERelationWithRteSubquery(RangeTblEntry *rangeTableEntry,
 	 * Originally: rtable: [rte1, current_rte, rte3...]
 	 * Now: rtable: [rte1, subquery[current_rte], rte3...] --subquery[current_rte] refers to its rtable.
 	 */
-	UpdateVarNosInNode(subquery, SINGLE_RTE_INDEX);
+	Node *quals = subquery->jointree->quals;
+	UpdateVarNosInNode(quals, SINGLE_RTE_INDEX);
 
 	/* replace the function with the constructed subquery */
 	rangeTableEntry->rtekind = RTE_SUBQUERY;
@@ -1590,12 +1537,12 @@ GenerateRequiredColNamesFromTargetList(List *targetList)
 
 /*
  * UpdateVarNosInNode iterates the Vars in the
- * given node's join tree quals and updates the varno's as the newVarNo.
+ * given node and updates the varno's as the newVarNo.
  */
 static void
-UpdateVarNosInNode(Query *query, Index newVarNo)
+UpdateVarNosInNode(Node *node, Index newVarNo)
 {
-	List *varList = pull_var_clause(query->jointree->quals, PVC_RECURSE_AGGREGATES |
+	List *varList = pull_var_clause(node, PVC_RECURSE_AGGREGATES |
 									PVC_RECURSE_PLACEHOLDERS);
 	Var *var = NULL;
 	foreach_ptr(var, varList)
@@ -1605,22 +1552,6 @@ UpdateVarNosInNode(Query *query, Index newVarNo)
 }
 
 
-/*
- * ContainsTableToBeConvertedToSubquery checks if the given range table list contains
- * any table that should be converted to a subquery, which otherwise is not plannable.
- */
-bool
-ContainsTableToBeConvertedToSubquery(List *rangeTableList)
-{
-	if (ContainsLocalTableDistributedTableJoin(rangeTableList))
-	{
-		return true;
-	}
-
-	return false;
-}
-
-
 /*
  * IsRecursivelyPlannableRelation returns true if the given range table entry
  * is a relation type that can be converted to a subquery.

src/include/distributed/local_distributed_join_planner.h

@@ -27,12 +27,8 @@ typedef enum
 
 extern int LocalTableJoinPolicy;
 
-extern bool ShouldConvertLocalTableJoinsToSubqueries(Query *query,
+extern bool ShouldConvertLocalTableJoinsToSubqueries(List *rangeTableList);
-													 List *rangeTableList,
-													 PlannerRestrictionContext *
-													 plannerRestrictionContext);
 extern void RecursivelyPlanLocalTableJoins(Query *query,
-										   RecursivePlanningContext *context,
+										   RecursivePlanningContext *context);
-										   List *rangeTableList);
 
 #endif /* LOCAL_DISTRIBUTED_JOIN_PLANNER_H */

src/include/distributed/multi_logical_optimizer.h

@@ -177,6 +177,5 @@ extern void FindReferencedTableColumn(Expr *columnExpression, List *parentQueryL
 extern char * WorkerColumnName(AttrNumber resno);
 extern bool IsGroupBySubsetOfDistinct(List *groupClauses, List *distinctClauses);
 extern bool TargetListHasAggregates(List *targetEntryList);
-extern Const * MakeIntegerConst(int32 integerValue);
 
 #endif   /* MULTI_LOGICAL_OPTIMIZER_H */

src/include/distributed/recursive_planning.h

@@ -47,7 +47,6 @@ extern bool ContainsLocalTableDistributedTableJoin(List *rangeTableList);
 extern void ReplaceRTERelationWithRteSubquery(RangeTblEntry *rangeTableEntry,
 											  List *requiredAttrNumbers,
 											  RecursivePlanningContext *context);
-extern bool ContainsTableToBeConvertedToSubquery(List *rangeTableList);
 extern bool IsRecursivelyPlannableRelation(RangeTblEntry *rangeTableEntry);
 extern bool IsRelationLocalTableOrMatView(Oid relationId);
 #endif /* RECURSIVE_PLANNING_H */

src/test/regress/expected/citus_local_dist_joins.out

@@ -458,6 +458,45 @@ DEBUG:  Plan XXX query after replacing subqueries and CTEs: SELECT count(*) AS c
      0
 (1 row)
 
+SELECT
+	COUNT(*)
+FROM
+	postgres_table p1
+JOIN
+	distributed_partitioned_table dp1
+USING (key)
+JOIN
+	distributed_table d1
+USING (key)
+JOIN
+	citus_local c1
+USING (key)	
+JOIN
+	postgres_table p2
+USING (key)
+JOIN
+	reference_table r1
+USING (key)
+JOIN
+	distributed_table d2
+USING (key)
+JOIN
+	citus_local c2
+USING (key);
+DEBUG:  Wrapping relation "postgres_table" "p1" to a subquery
+DEBUG:  generating subplan XXX_1 for subquery SELECT key FROM citus_local_dist_joins.postgres_table p1 WHERE true
+DEBUG:  Wrapping relation "citus_local" "c1" to a subquery
+DEBUG:  generating subplan XXX_2 for subquery SELECT key FROM citus_local_dist_joins.citus_local c1 WHERE true
+DEBUG:  Wrapping relation "postgres_table" "p2" to a subquery
+DEBUG:  generating subplan XXX_3 for subquery SELECT key FROM citus_local_dist_joins.postgres_table p2 WHERE true
+DEBUG:  Wrapping relation "citus_local" "c2" to a subquery
+DEBUG:  generating subplan XXX_4 for subquery SELECT key FROM citus_local_dist_joins.citus_local c2 WHERE true
+DEBUG:  Plan XXX query after replacing subqueries and CTEs: SELECT count(*) AS count FROM ((((((((SELECT p1_1.key, NULL::text AS value, NULL::jsonb AS value_2 FROM (SELECT intermediate_result.key FROM read_intermediate_result('XXX_1'::text, 'binary'::citus_copy_format) intermediate_result(key integer)) p1_1) p1 JOIN citus_local_dist_joins.distributed_partitioned_table dp1 USING (key)) JOIN citus_local_dist_joins.distributed_table d1 USING (key)) JOIN (SELECT c1_1.key, NULL::text AS value FROM (SELECT intermediate_result.key FROM read_intermediate_result('XXX_2'::text, 'binary'::citus_copy_format) intermediate_result(key integer)) c1_1) c1 USING (key)) JOIN (SELECT p2_1.key, NULL::text AS value, NULL::jsonb AS value_2 FROM (SELECT intermediate_result.key FROM read_intermediate_result('XXX_3'::text, 'binary'::citus_copy_format) intermediate_result(key integer)) p2_1) p2 USING (key)) JOIN citus_local_dist_joins.reference_table r1 USING (key)) JOIN citus_local_dist_joins.distributed_table d2 USING (key)) JOIN (SELECT c2_1.key, NULL::text AS value FROM (SELECT intermediate_result.key FROM read_intermediate_result('XXX_4'::text, 'binary'::citus_copy_format) intermediate_result(key integer)) c2_1) c2 USING (key))
+ count
+---------------------------------------------------------------------
+   100
+(1 row)
+
 SET client_min_messages to ERROR;
 DROP TABLE citus_local;
 SELECT master_remove_node('localhost', :master_port);

src/test/regress/expected/local_table_join.out

@@ -992,7 +992,15 @@ DEBUG:  Plan XXX query after replacing subqueries and CTEs: SELECT count(*) AS c
    100
 (1 row)
 
--- TODO:: we should support this?
+-- will error as we don't support complex joins
+SELECT COUNT(*) FROM postgres_table, distributed_table d1, distributed_table d2 WHERE d1.value = d2.value;
+DEBUG:  Wrapping relation "postgres_table" to a subquery
+DEBUG:  generating subplan XXX_1 for subquery SELECT NULL::integer AS key FROM local_table_join.postgres_table WHERE true
+DEBUG:  Plan XXX query after replacing subqueries and CTEs: SELECT count(*) AS count FROM (SELECT NULL::integer AS key, NULL::text AS value, NULL::jsonb AS value_2 FROM (SELECT intermediate_result.key FROM read_intermediate_result('XXX_1'::text, 'binary'::citus_copy_format) intermediate_result(key integer)) postgres_table_1) postgres_table, local_table_join.distributed_table d1, local_table_join.distributed_table d2 WHERE (d1.value OPERATOR(pg_catalog.=) d2.value)
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
+-- This will error because router planner will think that since reference tables have a single
+-- shard, it contains only a single task for modify. However, updating a reference tables
+-- will require multiple tasks. So requires some rewrite in router planner.
 UPDATE reference_table SET key = 1 FROM postgres_table WHERE postgres_table.key = 10;
 DEBUG:  Wrapping relation "postgres_table" to a subquery
 DEBUG:  generating subplan XXX_1 for subquery SELECT key FROM local_table_join.postgres_table WHERE (key OPERATOR(pg_catalog.=) 10)

src/test/regress/expected/multi_simple_queries.out

@@ -280,8 +280,7 @@ ORDER BY articles.id;
 -- subqueries are not supported in SELECT clause
 SELECT a.title AS name, (SELECT a2.id FROM articles_single_shard a2 WHERE a.id = a2.id  LIMIT 1)
 						 AS special_price FROM articles a;
-ERROR:  could not run distributed query with subquery outside the FROM, WHERE and HAVING clauses
+ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
-HINT:  Consider using an equality filter on the distributed table's partition column.
 -- joins are supported between local and distributed tables
 SELECT title, authors.name FROM authors, articles WHERE authors.id = articles.author_id;
  title | name

src/test/regress/expected/replicate_reference_tables_to_coordinator.out

@@ -281,23 +281,23 @@ SELECT test_reference_local_join_func();
 
 -- CTEs are allowed
 WITH ins AS (INSERT INTO numbers VALUES (1) RETURNING *)
-SELECT * FROM numbers, local_table;
+SELECT * FROM numbers, local_table ORDER BY 1,2;
 NOTICE:  executing the command locally: INSERT INTO replicate_ref_to_coordinator.numbers_8000001 (a) VALUES (1) RETURNING a
-NOTICE:  executing the command locally: SELECT numbers.a, local_table.a FROM replicate_ref_to_coordinator.numbers_8000001 numbers, (SELECT local_table_1.a FROM (SELECT intermediate_result.a FROM read_intermediate_result('XXX_2'::text, 'binary'::citus_copy_format) intermediate_result(a integer)) local_table_1) local_table
+NOTICE:  executing the command locally: SELECT numbers.a, local_table.a FROM replicate_ref_to_coordinator.numbers_8000001 numbers, (SELECT local_table_1.a FROM (SELECT intermediate_result.a FROM read_intermediate_result('XXX_2'::text, 'binary'::citus_copy_format) intermediate_result(a integer)) local_table_1) local_table ORDER BY numbers.a, local_table.a
  a  | a
 ---------------------------------------------------------------------
- 20 |  2
- 21 |  2
   2 |  2
- 20 |  4
- 21 |  4
   2 |  4
- 20 |  7
- 21 |  7
   2 |  7
- 20 | 20
- 21 | 20
   2 | 20
+ 20 |  2
+ 20 |  4
+ 20 |  7
+ 20 | 20
+ 21 |  2
+ 21 |  4
+ 21 |  7
+ 21 | 20
 (12 rows)
 
 WITH t AS (SELECT *, my_volatile_fn() x FROM numbers FOR UPDATE)

src/test/regress/sql/citus_local_dist_joins.sql

@@ -224,6 +224,32 @@ SELECT count(*) FROM postgres_table JOIN (SELECT * FROM (SELECT * FROM distribut
 SELECT count(*) FROM postgres_table JOIN (SELECT * FROM (SELECT * FROM distributed_table LIMIT 1) d1) d2 using (key) JOIN reference_table USING(key) JOIN citus_local USING (key) JOIN (SELECT * FROM citus_local) c1  USING (key) WHERE d2.key > 10 AND d2.key = 10;
 
 
+SELECT 
+	COUNT(*) 
+FROM 
+	postgres_table p1
+JOIN 
+	distributed_partitioned_table dp1
+USING (key)
+JOIN
+	distributed_table d1 
+USING (key)
+JOIN 
+	citus_local c1 
+USING (key)	
+JOIN 
+	postgres_table p2
+USING (key)
+JOIN
+	reference_table r1
+USING (key)
+JOIN 
+	distributed_table d2
+USING (key)
+JOIN
+	citus_local c2 
+USING (key);
+
 
 SET client_min_messages to ERROR;
 DROP TABLE citus_local;

src/test/regress/sql/local_table_join.sql

@@ -252,8 +252,12 @@ FROM
 JOIN postgres_table
 USING(key);
 
+-- will error as we don't support complex joins
+SELECT COUNT(*) FROM postgres_table, distributed_table d1, distributed_table d2 WHERE d1.value = d2.value;
 
--- TODO:: we should support this?
+-- This will error because router planner will think that since reference tables have a single
+-- shard, it contains only a single task for modify. However, updating a reference tables
+-- will require multiple tasks. So requires some rewrite in router planner.
 UPDATE reference_table SET key = 1 FROM postgres_table WHERE postgres_table.key = 10;
 UPDATE reference_table SET key = 1 FROM (SELECT * FROM postgres_table) l WHERE l.key = 10;
 

src/test/regress/sql/replicate_reference_tables_to_coordinator.sql

@@ -145,7 +145,7 @@ SELECT test_reference_local_join_func();
 
 -- CTEs are allowed
 WITH ins AS (INSERT INTO numbers VALUES (1) RETURNING *)
-SELECT * FROM numbers, local_table;
+SELECT * FROM numbers, local_table ORDER BY 1,2;
 
 WITH t AS (SELECT *, my_volatile_fn() x FROM numbers FOR UPDATE)
 SELECT * FROM numbers, local_table