more add_path hook experiments

src/backend/distributed/planner/distributed_planner.c

@@ -74,6 +74,7 @@ static uint64 NextPlanId = 1;
 /* keep track of planner call stack levels */
 int PlannerLevel = 0;
 bool UseCustomPath = false;
+bool PlanAllPaths = false;
 bool OnlyGeoPartitioning = false;
 bool UseGeoPartitioning = true;
 bool EnableGeoPartitioningGrouping = true;

src/backend/distributed/planner/path_based_planner.c

@@ -2502,6 +2502,30 @@ PathTreeWalker(Node *node, bool (*walker)(), void *context)
 			return false;
 		}
 
+		case T_MaterialPath:
+		{
+			MaterialPath *path = castNode(MaterialPath, node);
+			if (walker(path->subpath, context))
+				return true;
+			return false;
+		}
+
+		case T_SortPath:
+		{
+			SortPath *path = castNode(SortPath, node);
+			if (walker(path->subpath, context))
+				return true;
+			return false;
+		}
+
+		case T_ProjectionPath:
+		{
+			ProjectionPath *path = castNode(ProjectionPath, node);
+			if (walker(path->subpath, context))
+				return true;
+			return false;
+		}
+
 		/* paths not having nesting */
 		case T_Path:
 		case T_IndexPath:
@@ -2533,8 +2557,13 @@ ColocationGroupsForPathWalker(Node *node, ColocationGroups *context)
 	DistributedUnionPath *collect = NULL;
 	if (IsDistributedUnion((Path *) node, false, &collect))
 	{
+		int colocationId = (int) collect->colocationId;
+		if (colocationId == -1)
+		{
+			colocationId = 0;
+		}
 		context->colocationIds = bms_add_member(context->colocationIds,
-												(int) collect->colocationId);
+												colocationId);
 		return false;
 	}
 
@@ -2550,24 +2579,77 @@ ColocationGroupsForPath(Path *node)
 	return context.colocationIds;
 }
 
+
+typedef struct ColocationGroupsList
+{
+	/* configuration */
+	bool unique;
+	bool sorted;
+
+	/* output */
+	List *list;
+} ColocationGroupsList;
+
+static bool ColocationGroupsForPathListWalker(Node *node, ColocationGroupsList *context);
+static bool
+ColocationGroupsForPathListWalker(Node *node, ColocationGroupsList *context)
+{
+	if (node == NULL)
+		return false;
+
+	DistributedUnionPath *collect = NULL;
+	if (IsDistributedUnion((Path *) node, false, &collect))
+	{
+		int colocationId = (int) collect->colocationId;
+		if (!context->unique || !list_member_int(context->list, colocationId))
+			context->list = lappend_int(context->list, colocationId);
+
+		return false;
+	}
+
+	return PathTreeWalker(node, ColocationGroupsForPathListWalker, context);
+}
+
+
+
+static List *
+ColocationGroupsForPathList(Path *node)
+{
+	ColocationGroupsList context = { .unique =  true, .sorted =  true };
+	ColocationGroupsForPathListWalker((Node *)node, &context);
+
+	if (context.sorted)
+		list_sort(context.list, list_int_cmp);
+
+	return context.list;
+}
+
 PathComparison
 PathBasedPlannerComparePath(Path *new_path, Path *old_path)
 {
+	if (PlanAllPaths)
+	{
+		return PATH_DIFFERENT;
+	}
+
 	if (!UseCustomPath)
 	{
 		return PATH_EQUAL;
 	}
 
-	Bitmapset *newColocationIds = ColocationGroupsForPath(new_path);
-	Bitmapset *oldColocationIds = ColocationGroupsForPath(old_path);
-
-	BMS_Comparison cmp = bms_subset_compare(newColocationIds, oldColocationIds);
-
-	bms_free(newColocationIds);
-	bms_free(oldColocationIds);
-
-	if (cmp == BMS_EQUAL)
-		return PATH_EQUAL;
-
-	return PATH_DIFFERENT;
+	return PATH_EQUAL;
+}
+
+add_path_merit_list_hook_type prev_add_path_merit_list_hook = NULL;
+List *
+PathBasedMeritListHook(Path *path)
+{
+	List *merits = NIL;
+	if (prev_add_path_merit_list_hook)
+	{
+		merits = prev_add_path_merit_list_hook(path);
+	}
+
+	List *colocationIds = ColocationGroupsForPathList(path);
+	return lappend(merits, colocationIds);
 }

src/backend/distributed/shared_library_init.c

@@ -300,6 +300,8 @@ _PG_init(void)
 	set_join_pathlist_hook = multi_join_restriction_hook;
 	create_upper_paths_hook = PathBasedPlannedUpperPathHook;
 	compare_path_hook = PathBasedPlannerComparePath;
+	prev_add_path_merit_list_hook = add_path_merit_list_hook;
+	add_path_merit_list_hook = PathBasedMeritListHook;
 
 	ExecutorStart_hook = CitusExecutorStart;
 	ExecutorRun_hook = CitusExecutorRun;
@@ -1148,6 +1150,16 @@ RegisterCitusConfigVariables(void)
 		GUC_STANDARD,
 		NULL, NULL, NULL);
 
+	DefineCustomBoolVariable(
+		"citus.plan_all_paths",
+		gettext_noop("Disable any pruning of paths in add_path for debugging purposes"),
+		NULL,
+		&PlanAllPaths,
+		false,
+		PGC_USERSET,
+		GUC_STANDARD,
+		NULL, NULL, NULL);
+
 	DefineCustomStringVariable(
 		"citus.local_hostname",
 		gettext_noop("Sets the hostname when connecting back to itself."),

src/include/distributed/distributed_planner.h

@@ -33,6 +33,7 @@
 /* level of planner calls */
 extern int PlannerLevel;
 extern bool UseCustomPath;
+extern bool PlanAllPaths;
 extern bool OnlyGeoPartitioning;
 extern bool UseGeoPartitioning;
 extern bool EnableGeoPartitioningGrouping;

src/include/distributed/path_based_planner.h

@@ -7,6 +7,7 @@
 
 #include "nodes/parsenodes.h"
 #include "nodes/pathnodes.h"
+#include "optimizer/paths.h"
 
 typedef List *(*optimizeFn)(PlannerInfo *root, Path *originalPath);
 
@@ -47,5 +48,8 @@ extern void PathBasedPlannedUpperPathHook(PlannerInfo *root,
 										  void *extra);
 
 extern PathComparison PathBasedPlannerComparePath(Path *new_path, Path *old_path);
+extern List * PathBasedMeritListHook(Path *path);
+
+extern add_path_merit_list_hook_type prev_add_path_merit_list_hook;
 
 #endif //CITUS_PATH_BASED_PLANNER_H