Refactor range table walkers (#3109)

src/backend/distributed/planner/distributed_planner.c

@@ -30,6 +30,7 @@
 #include "distributed/multi_physical_planner.h"
 #include "distributed/multi_master_planner.h"
 #include "distributed/multi_router_planner.h"
+#include "distributed/query_utils.h"
 #include "distributed/recursive_planning.h"
 #include "distributed/shardinterval_utils.h"
 #include "distributed/worker_shard_visibility.h"

src/backend/distributed/planner/multi_logical_planner.c

@@ -28,6 +28,7 @@
 #include "distributed/multi_physical_planner.h"
 #include "distributed/relation_restriction_equivalence.h"
 #include "distributed/query_pushdown_planning.h"
+#include "distributed/query_utils.h"
 #include "distributed/multi_router_planner.h"
 #include "distributed/worker_protocol.h"
 #include "distributed/version_compat.h"
@@ -1177,34 +1178,6 @@ HasComplexRangeTableType(Query *queryTree)
 }
 
 
-/*
- * ExtractRangeTableIndexWalker walks over a join tree, and finds all range
- * table indexes in that tree.
- */
-bool
-ExtractRangeTableIndexWalker(Node *node, List **rangeTableIndexList)
-{
-	bool walkerResult = false;
-	if (node == NULL)
-	{
-		return false;
-	}
-
-	if (IsA(node, RangeTblRef))
-	{
-		int rangeTableIndex = ((RangeTblRef *) node)->rtindex;
-		(*rangeTableIndexList) = lappend_int(*rangeTableIndexList, rangeTableIndex);
-	}
-	else
-	{
-		walkerResult = expression_tree_walker(node, ExtractRangeTableIndexWalker,
-											  rangeTableIndexList);
-	}
-
-	return walkerResult;
-}
-
-
 /*
  * WhereClauseList walks over the FROM expression in the query tree, and builds
  * a list of all clauses from the expression tree. The function checks for both
@@ -1967,105 +1940,6 @@ FindNodesOfType(MultiNode *node, int type)
 }
 
 
-/*
- * ExtractRangeTableRelationWalker gathers all range table relation entries
- * in a query.
- */
-bool
-ExtractRangeTableRelationWalker(Node *node, List **rangeTableRelationList)
-{
-	bool walkIsComplete = false;
-
-	if (node == NULL)
-	{
-		return false;
-	}
-
-	if (IsA(node, RangeTblEntry))
-	{
-		RangeTblEntry *rangeTable = (RangeTblEntry *) node;
-
-		if (rangeTable->rtekind == RTE_RELATION && rangeTable->relkind != RELKIND_VIEW)
-		{
-			(*rangeTableRelationList) = lappend(*rangeTableRelationList, rangeTable);
-
-			walkIsComplete = false;
-		}
-	}
-	else if (IsA(node, Query))
-	{
-		walkIsComplete = query_tree_walker((Query *) node,
-										   ExtractRangeTableRelationWalker,
-										   rangeTableRelationList,
-										   QTW_EXAMINE_RTES_BEFORE);
-	}
-	else
-	{
-		walkIsComplete = expression_tree_walker(node, ExtractRangeTableRelationWalker,
-												rangeTableRelationList);
-	}
-
-	return walkIsComplete;
-}
-
-
-/*
- * ExtractRangeTableEntryWalker walks over a query tree, and finds all range
- * table entries. For recursing into the query tree, this function uses the
- * query tree walker since the expression tree walker doesn't recurse into
- * sub-queries.
- */
-bool
-ExtractRangeTableEntryWalker(Node *node, List **rangeTableList)
-{
-	bool walkIsComplete = false;
-	if (node == NULL)
-	{
-		return false;
-	}
-
-	if (IsA(node, RangeTblEntry))
-	{
-		RangeTblEntry *rangeTable = (RangeTblEntry *) node;
-		(*rangeTableList) = lappend(*rangeTableList, rangeTable);
-	}
-	else if (IsA(node, Query))
-	{
-		Query *query = (Query *) node;
-
-		/*
-		 * Since we're only interested in range table entries, we only descend
-		 * into all parts of the query when it is necessary. Otherwise, it is
-		 * sufficient to descend into range table list since its the only part
-		 * of the query that could contain range table entries.
-		 */
-		if (query->hasSubLinks || query->cteList || query->setOperations)
-		{
-			/* descend into all parts of the query */
-			walkIsComplete = query_tree_walker((Query *) node,
-											   ExtractRangeTableEntryWalker,
-											   rangeTableList,
-											   QTW_EXAMINE_RTES_BEFORE);
-		}
-		else
-		{
-			/* descend only into RTEs */
-			walkIsComplete = range_table_walker(query->rtable,
-												ExtractRangeTableEntryWalker,
-												rangeTableList,
-												QTW_EXAMINE_RTES_BEFORE);
-		}
-	}
-	else
-	{
-		walkIsComplete = expression_tree_walker(node, ExtractRangeTableEntryWalker,
-												rangeTableList);
-	}
-
-	return walkIsComplete;
-}
-
-
 /*
  * pull_var_clause_default calls pull_var_clause with the most commonly used
  * arguments for distributed planning.

src/backend/distributed/planner/multi_router_planner.c

@@ -39,6 +39,7 @@
 #include "distributed/listutils.h"
 #include "distributed/citus_ruleutils.h"
 #include "distributed/query_pushdown_planning.h"
+#include "distributed/query_utils.h"
 #include "distributed/relation_restriction_equivalence.h"
 #include "distributed/relay_utility.h"
 #include "distributed/resource_lock.h"

src/backend/distributed/planner/query_pushdown_planning.c

@@ -29,6 +29,7 @@
 #include "distributed/multi_logical_planner.h"
 #include "distributed/multi_router_planner.h"
 #include "distributed/pg_dist_partition.h"
+#include "distributed/query_utils.h"
 #include "distributed/query_pushdown_planning.h"
 #include "distributed/relation_restriction_equivalence.h"
 #include "distributed/version_compat.h"

src/backend/distributed/planner/relation_restriction_equivalence.c

@@ -15,6 +15,7 @@
 #include "distributed/multi_logical_planner.h"
 #include "distributed/multi_logical_optimizer.h"
 #include "distributed/pg_dist_partition.h"
+#include "distributed/query_utils.h"
 #include "distributed/relation_restriction_equivalence.h"
 #include "nodes/nodeFuncs.h"
 #include "nodes/pg_list.h"

src/backend/distributed/test/distribution_metadata.c

@@ -27,6 +27,7 @@
 #include "distributed/multi_join_order.h"
 #include "distributed/multi_physical_planner.h"
 #include "distributed/pg_dist_shard.h"
+#include "distributed/query_utils.h"
 #include "distributed/resource_lock.h"
 #include "lib/stringinfo.h"
 #include "nodes/pg_list.h"

src/backend/distributed/utils/query_utils.c

@@ -0,0 +1,139 @@
+/*-------------------------------------------------------------------------
+ *
+ * query_utils.c
+ *
+ * Query-walker utility functions to be used to construct a logical plan
+ * tree from the given query tree structure.
+ *
+ * Copyright (c), Citus Data, Inc.
+ *
+ **----------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "catalog/pg_class.h"
+#include "distributed/query_utils.h"
+#include "distributed/version_compat.h"
+#include "nodes/nodeFuncs.h"
+
+/*
+ * ExtractRangeTableList walks over a tree to gather entries.
+ * Execution is parameterized by passing walkerMode flag via ExtractRangeTableWalkerContext
+ * as we cannot pass more than one parameter to query_tree_walker
+ */
+bool
+ExtractRangeTableList(Node *node, ExtractRangeTableWalkerContext *context)
+{
+	/* get parameters from context */
+	List **rangeTableRelationList = context->rangeTableList;
+	ExtractRangeTableMode walkerMode = context->walkerMode;
+
+	bool walkIsComplete = false;
+
+	if (node == NULL)
+	{
+		return false;
+	}
+
+	if (IsA(node, RangeTblEntry))
+	{
+		RangeTblEntry *rangeTable = (RangeTblEntry *) node;
+
+		if (walkerMode == EXTRACT_ALL_ENTRIES ||
+			(rangeTable->rtekind == RTE_RELATION && rangeTable->relkind != RELKIND_VIEW))
+		{
+			(*rangeTableRelationList) = lappend(*rangeTableRelationList, rangeTable);
+		}
+	}
+	else if (IsA(node, Query))
+	{
+		Query *query = (Query *) node;
+
+		if (query->hasSubLinks || query->cteList || query->setOperations)
+		{
+			/* descend into all parts of the query */
+			walkIsComplete = query_tree_walker(query,
+											   ExtractRangeTableList,
+											   context,
+											   QTW_EXAMINE_RTES_BEFORE);
+		}
+		else
+		{
+			/* descend only into RTEs */
+			walkIsComplete = range_table_walker(query->rtable,
+												ExtractRangeTableList,
+												context,
+												QTW_EXAMINE_RTES_BEFORE);
+		}
+	}
+	else
+	{
+		walkIsComplete = expression_tree_walker(node, ExtractRangeTableList,
+												context);
+	}
+
+	return walkIsComplete;
+}
+
+
+/*
+ * ExtractRangeTableRelationWalker gathers all range table relation entries
+ * in a query.
+ */
+bool
+ExtractRangeTableRelationWalker(Node *node, List **rangeTableRelationList)
+{
+	ExtractRangeTableWalkerContext context;
+
+	context.rangeTableList = rangeTableRelationList;
+	context.walkerMode = EXTRACT_RELATION_ENTRIES;
+
+	return ExtractRangeTableList(node, &context);
+}
+
+
+/*
+ * ExtractRangeTableEntryWalker walks over a query tree, and finds all range
+ * table entries. For recursing into the query tree, this function uses the
+ * query tree walker since the expression tree walker doesn't recurse into
+ * sub-queries.
+ */
+bool
+ExtractRangeTableEntryWalker(Node *node, List **rangeTableList)
+{
+	ExtractRangeTableWalkerContext context;
+
+	context.rangeTableList = rangeTableList;
+	context.walkerMode = EXTRACT_ALL_ENTRIES;
+
+	return ExtractRangeTableList(node, &context);
+}
+
+
+/*
+ * ExtractRangeTableIndexWalker walks over a join tree, and finds all range
+ * table indexes in that tree.
+ */
+bool
+ExtractRangeTableIndexWalker(Node *node, List **rangeTableIndexList)
+{
+	bool walkerResult = false;
+	if (node == NULL)
+	{
+		return false;
+	}
+
+	if (IsA(node, RangeTblRef))
+	{
+		int rangeTableIndex = ((RangeTblRef *) node)->rtindex;
+		(*rangeTableIndexList) = lappend_int(*rangeTableIndexList, rangeTableIndex);
+	}
+	else
+	{
+		walkerResult = expression_tree_walker(node, ExtractRangeTableIndexWalker,
+											  rangeTableIndexList);
+	}
+
+	return walkerResult;
+}

src/include/distributed/multi_logical_planner.h

@@ -206,13 +206,10 @@ extern List * JoinClauseList(List *whereClauseList);
 extern bool IsJoinClause(Node *clause);
 extern List * SubqueryEntryList(Query *queryTree);
 extern DeferredErrorMessage * DeferErrorIfQueryNotSupported(Query *queryTree);
-extern bool ExtractRangeTableIndexWalker(Node *node, List **rangeTableIndexList);
 extern List * WhereClauseList(FromExpr *fromExpr);
 extern List * QualifierList(FromExpr *fromExpr);
 extern List * TableEntryList(List *rangeTableList);
 extern List * UsedTableEntryList(Query *query);
-extern bool ExtractRangeTableRelationWalker(Node *node, List **rangeTableList);
-extern bool ExtractRangeTableEntryWalker(Node *node, List **rangeTableList);
 extern List * pull_var_clause_default(Node *node);
 extern bool OperatorImplementsEquality(Oid opno);
 extern bool FindNodeCheck(Node *node, bool (*check)(Node *));

src/include/distributed/query_utils.h

@@ -0,0 +1,40 @@
+/*-------------------------------------------------------------------------
+ *
+ * query_utils.h
+ *
+ * Declarations for query-walker utility functions and related types.
+ *
+ * Copyright (c), Citus Data, Inc.
+ *
+ **---------------------------------------------------------------------------
+ */
+#ifndef QUERY_UTILS_H
+#define QUERY_UTILS_H
+
+#include "postgres.h"
+#include "nodes/pg_list.h"
+
+/* Enum to define execution flow of ExtractRangeTableList */
+typedef enum ExtractRangeTableMode
+{
+	EXTRACT_RELATION_ENTRIES,
+	EXTRACT_ALL_ENTRIES
+} ExtractRangeTableMode;
+
+/* Struct to pass rtable list and execution flow flag to query_walker_tree */
+typedef struct ExtractRangeTableWalkerContext
+{
+	List **rangeTableList;
+	ExtractRangeTableMode walkerMode;
+} ExtractRangeTableWalkerContext;
+
+/* Function declarations for query-walker utility functions */
+extern bool ExtractRangeTableList(Node *node, ExtractRangeTableWalkerContext *context);
+
+/* Below two functions wrap ExtractRangeTableList function to determine the execution flow */
+extern bool ExtractRangeTableRelationWalker(Node *node, List **rangeTableList);
+extern bool ExtractRangeTableEntryWalker(Node *node, List **rangeTableList);
+
+extern bool ExtractRangeTableIndexWalker(Node *node, List **rangeTableIndexList);
+
+#endif /* QUERY_UTILS_H */