Search all outer node levels for lateral join params.

src/backend/distributed/planner/distributed_planner.c

@@ -79,6 +79,7 @@ static void CheckNodeIsDumpable(Node *node);
 static Node * CheckNodeCopyAndSerialization(Node *node);
 static void AdjustReadIntermediateResultCost(RangeTblEntry *rangeTableEntry,
 											 RelOptInfo *relOptInfo);
+static List * OuterPlanParamsList(PlannerInfo *root);
 static List * CopyPlanParamList(List *originalPlanParamList);
 static PlannerRestrictionContext * CreateAndPushPlannerRestrictionContext(void);
 static PlannerRestrictionContext * CurrentPlannerRestrictionContext(void);
@@ -1299,15 +1300,10 @@ multi_relation_restriction_hook(PlannerInfo *root, RelOptInfo *relOptInfo, Index
 	relationRestriction->relOptInfo = relOptInfo;
 	relationRestriction->distributedRelation = distributedTable;
 	relationRestriction->plannerInfo = root;
-	relationRestriction->parentPlannerInfo = root->parent_root;
 	relationRestriction->prunedShardIntervalList = NIL;
 
 	/* see comments on GetVarFromAssignedParam() */
-	if (relationRestriction->parentPlannerInfo)
-	{
-		relationRestriction->parentPlannerParamList =
-			CopyPlanParamList(root->parent_root->plan_params);
-	}
+	relationRestriction->outerPlanParamsList = OuterPlanParamsList(root);
 
 	relationRestrictionContext = plannerRestrictionContext->relationRestrictionContext;
 	relationRestrictionContext->hasDistributedRelation |= distributedTable;
@@ -1462,6 +1458,36 @@ AdjustReadIntermediateResultCost(RangeTblEntry *rangeTableEntry, RelOptInfo *rel
 }
 
 
+/*
+ * OuterPlanParamsList creates a list of RootPlanParams for outer nodes of the
+ * given root. The first item in the list corresponds to parent_root, and the
+ * last item corresponds to the outer most node.
+ */
+static List *
+OuterPlanParamsList(PlannerInfo *root)
+{
+	List *planParamsList = NIL;
+	PlannerInfo *outerNodeRoot = NULL;
+
+	for (outerNodeRoot = root->parent_root; outerNodeRoot != NULL;
+		 outerNodeRoot = outerNodeRoot->parent_root)
+	{
+		RootPlanParams *rootPlanParams = palloc0(sizeof(RootPlanParams));
+		rootPlanParams->root = outerNodeRoot;
+
+		/*
+		 * TODO: In SearchPlannerParamList() we are only interested in Var plan
+		 * params, consider copying just them here.
+		 */
+		rootPlanParams->plan_params = CopyPlanParamList(outerNodeRoot->plan_params);
+
+		planParamsList = lappend(planParamsList, rootPlanParams);
+	}
+
+	return planParamsList;
+}
+
+
 /*
  * CopyPlanParamList deep copies the input PlannerParamItem list and returns the newly
  * allocated list.

src/backend/distributed/planner/relation_restriction_equivalence.c

@@ -98,8 +98,9 @@ static void AddRteRelationToAttributeEquivalenceClass(AttributeEquivalenceClass
 													  attrEquivalenceClass,
 													  RangeTblEntry *rangeTableEntry,
 													  Var *varToBeAdded);
-static Var * GetVarFromAssignedParam(List *parentPlannerParamList,
-									 Param *plannerParam);
+static Var * GetVarFromAssignedParam(List *outerPlanParamsList, Param *plannerParam,
+									 PlannerInfo **rootContainingVar);
+static Var * SearchPlannerParamList(List *plannerParamList, Param *plannerParam);
 static List * GenerateAttributeEquivalencesForJoinRestrictions(JoinRestrictionContext
 															   *joinRestrictionContext);
 static bool AttributeClassContainsAttributeClassMember(AttributeEquivalenceClassMember *
@@ -737,15 +738,15 @@ AttributeEquivalenceClassForEquivalenceClass(EquivalenceClass *plannerEqClass,
 
 		if (IsA(strippedEquivalenceExpr, Param))
 		{
-			List *parentParamList = relationRestriction->parentPlannerParamList;
+			PlannerInfo *outerNodeRoot = NULL;
 			Param *equivalenceParam = (Param *) strippedEquivalenceExpr;
 
-			expressionVar = GetVarFromAssignedParam(parentParamList,
-													equivalenceParam);
+			expressionVar =
+				GetVarFromAssignedParam(relationRestriction->outerPlanParamsList,
+										equivalenceParam, &outerNodeRoot);
 			if (expressionVar)
 			{
-				AddToAttributeEquivalenceClass(&attributeEquivalance,
-											   relationRestriction->parentPlannerInfo,
+				AddToAttributeEquivalenceClass(&attributeEquivalance, outerNodeRoot,
 											   expressionVar);
 			}
 		}
@@ -766,7 +767,7 @@ AttributeEquivalenceClassForEquivalenceClass(EquivalenceClass *plannerEqClass,
  * plannerParam if its kind is PARAM_EXEC.
  *
  * If the paramkind is not equal to PARAM_EXEC the function returns NULL. Similarly,
- * if there is no var that the given param is assigned to, the function returns NULL.
+ * if there is no Var corresponding to the given param is, the function returns NULL.
  *
  * Rationale behind this function:
  *
@@ -775,26 +776,26 @@ AttributeEquivalenceClassForEquivalenceClass(EquivalenceClass *plannerEqClass,
  *   the RTE_RELATIONs which actually belong to lateral vars from the other query
  *   levels.
  *
- *   We're also keeping track of the RTE_RELATION's parent_root's
- *   plan_param list which is expected to hold the parameters that are required
+ *   We're also keeping track of the RTE_RELATION's outer nodes'
+ *   plan_params lists which is expected to hold the parameters that are required
  *   for its lower level queries as it is documented:
  *
  *        plan_params contains the expressions that this query level needs to
  *        make available to a lower query level that is currently being planned.
  *
- *   This function is a helper function to iterate through the parent query's
+ *   This function is a helper function to iterate through the outer node's query's
  *   plan_params and looks for the param that the equivalence member has. The
  *   comparison is done via the "paramid" field. Finally, if the found parameter's
  *   item is a Var, we conclude that Postgres standard_planner replaced the Var
  *   with the Param on assign_param_for_var() function
- *   @src/backend/optimizer//plan/subselect.c.
- *
+ *   @src/backend/optimizer/plan/subselect.c.
  */
 static Var *
-GetVarFromAssignedParam(List *parentPlannerParamList, Param *plannerParam)
+GetVarFromAssignedParam(List *outerPlanParamsList, Param *plannerParam,
+						PlannerInfo **rootContainingVar)
 {
 	Var *assignedVar = NULL;
-	ListCell *plannerParameterCell = NULL;
+	ListCell *rootPlanParamsCell = NULL;
 
 	Assert(plannerParam != NULL);
 
@@ -804,7 +805,35 @@ GetVarFromAssignedParam(List *parentPlannerParamList, Param *plannerParam)
 		return NULL;
 	}
 
-	foreach(plannerParameterCell, parentPlannerParamList)
+	foreach(rootPlanParamsCell, outerPlanParamsList)
+	{
+		RootPlanParams *outerPlanParams = lfirst(rootPlanParamsCell);
+
+		assignedVar = SearchPlannerParamList(outerPlanParams->plan_params,
+											 plannerParam);
+		if (assignedVar != NULL)
+		{
+			*rootContainingVar = outerPlanParams->root;
+			break;
+		}
+	}
+
+	return assignedVar;
+}
+
+
+/*
+ * SearchPlannerParamList searches in plannerParamList and returns the Var that
+ * corresponds to the given plannerParam. If there is no Var corresponding to the
+ * given param is, the function returns NULL.
+ */
+static Var *
+SearchPlannerParamList(List *plannerParamList, Param *plannerParam)
+{
+	Var *assignedVar = NULL;
+	ListCell *plannerParameterCell = NULL;
+
+	foreach(plannerParameterCell, plannerParamList)
 	{
 		PlannerParamItem *plannerParamItem =
 			(PlannerParamItem *) lfirst(plannerParameterCell);
@@ -814,7 +843,7 @@ GetVarFromAssignedParam(List *parentPlannerParamList, Param *plannerParam)
 			continue;
 		}
 
-		/* TODO: Should we consider PlaceHolderVar?? */
+		/* TODO: Should we consider PlaceHolderVar? */
 		if (!IsA(plannerParamItem->item, Var))
 		{
 			continue;

src/include/distributed/distributed_planner.h

@@ -32,6 +32,18 @@ typedef struct RelationRestrictionContext
 	List *relationRestrictionList;
 } RelationRestrictionContext;
 
+
+typedef struct RootPlanParams
+{
+	PlannerInfo *root;
+
+	/*
+	 * Copy of root->plan_params. root->plan_params is not preserved in
+	 * relation_restriction_equivalence, so we need to create a copy.
+	 */
+	List *plan_params;
+} RootPlanParams;
+
 typedef struct RelationRestriction
 {
 	Index index;
@@ -40,9 +52,10 @@ typedef struct RelationRestriction
 	RangeTblEntry *rte;
 	RelOptInfo *relOptInfo;
 	PlannerInfo *plannerInfo;
-	PlannerInfo *parentPlannerInfo;
-	List *parentPlannerParamList;
 	List *prunedShardIntervalList;
+
+	/* list of RootPlanParams for all outer nodes */
+	List *outerPlanParamsList;
 } RelationRestriction;
 
 typedef struct JoinRestrictionContext

src/test/regress/expected/multi_subquery_behavioral_analytics.out

@@ -2249,6 +2249,37 @@ FROM
      6 |    
 (1 row)
 
+-- Test the case when a subquery has a lateral reference to two levels upper
+SELECT
+ b.user_id, b.value_2, b.cnt
+FROM (
+ SELECT
+   user_id,
+   value_2
+ FROM events_table
+ WHERE events_table.user_id BETWEEN 2 AND 5
+) a,
+LATERAL (
+ SELECT
+   user_id, value_2, count(*) as cnt
+ FROM (
+   SELECT
+     value_2, time, user_id
+   FROM events_table
+   WHERE user_id BETWEEN 2 AND 5
+     AND events_table.user_id = a.user_id
+     AND events_table.value_2 = a.value_2
+   ORDER BY time DESC
+ ) events
+ GROUP BY user_id, value_2
+) b
+ORDER BY user_id, value_2, cnt
+LIMIT 1;
+ user_id | value_2 | cnt 
+---------+---------+-----
+       2 |       0 |   1
+(1 row)
+
 DROP FUNCTION test_join_function_2(integer, integer);
 SELECT run_command_on_workers($f$
 

src/test/regress/expected/set_operations.out

@@ -671,12 +671,11 @@ SELECT * FROM test a WHERE x IN (SELECT x FROM test b UNION SELECT y FROM test c
 DEBUG:  generating subplan 144_1 for subquery SELECT x FROM recursive_union.test b
 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
 DEBUG:  Plan 144 query after replacing subqueries and CTEs: SELECT x, y FROM recursive_union.test a WHERE (x OPERATOR(pg_catalog.=) ANY (SELECT intermediate_result.x FROM read_intermediate_result('144_1'::text, 'binary'::citus_copy_format) intermediate_result(x integer) UNION SELECT c.y FROM recursive_union.test c WHERE (a.x OPERATOR(pg_catalog.=) c.x))) ORDER BY x, y
 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:  cannot push down this subquery
+DETAIL:  Complex subqueries and CTEs are not supported within a UNION
 -- force unions to be planned while subqueries are being planned
 SELECT * FROM ((SELECT * FROM test) UNION (SELECT * FROM test) ORDER BY 1,2 LIMIT 5) as foo ORDER BY 1 DESC LIMIT 3;
 DEBUG:  generating subplan 147_1 for subquery SELECT x, y FROM recursive_union.test

src/test/regress/sql/multi_subquery_behavioral_analytics.sql

@@ -1839,6 +1839,33 @@ FROM
   ) AS temp;
 
 
+-- Test the case when a subquery has a lateral reference to two levels upper
+SELECT
+ b.user_id, b.value_2, b.cnt
+FROM (
+ SELECT
+   user_id,
+   value_2
+ FROM events_table
+ WHERE events_table.user_id BETWEEN 2 AND 5
+) a,
+LATERAL (
+ SELECT
+   user_id, value_2, count(*) as cnt
+ FROM (
+   SELECT
+     value_2, time, user_id
+   FROM events_table
+   WHERE user_id BETWEEN 2 AND 5
+     AND events_table.user_id = a.user_id
+     AND events_table.value_2 = a.value_2
+   ORDER BY time DESC
+ ) events
+ GROUP BY user_id, value_2
+) b
+ORDER BY user_id, value_2, cnt
+LIMIT 1;
+
 
 DROP FUNCTION test_join_function_2(integer, integer);