From d092ef5e0235a3812d2494287e551cc75db93a67 Mon Sep 17 00:00:00 2001
From: Nils Dijk <nils@citusdata.com>
Date: Wed, 15 Dec 2021 17:59:37 +0100
Subject: [PATCH] compare_path_hook

---
 .../distributed/planner/path_based_planner.c  | 118 ++++++++++++++++++
 src/backend/distributed/shared_library_init.c |   1 +
 src/include/distributed/path_based_planner.h  |   2 +
 3 files changed, 121 insertions(+)

diff --git a/src/backend/distributed/planner/path_based_planner.c b/src/backend/distributed/planner/path_based_planner.c
index d3ccefdf9..8a5def730 100644
--- a/src/backend/distributed/planner/path_based_planner.c
+++ b/src/backend/distributed/planner/path_based_planner.c
@@ -2453,3 +2453,121 @@ RepartitionAggPath(PlannerInfo *root, Path *originalPath)
 	return NIL;
 }
 
+
+static bool
+PathTreeWalker(Node *node, bool (*walker)(), void *context)
+{
+	switch(nodeTag(node))
+	{
+		case T_List:
+		{
+			Path *path = NULL;
+			foreach_ptr(path, castNode(List, node))
+			{
+				if (walker((Node *) path, context))
+				{
+					return true;
+				}
+			}
+			return false;
+		}
+
+		case T_CustomPath:
+		{
+			CustomPath *path = castNode(CustomPath, node);
+			if (walker((Node *)path->custom_paths, context))
+			{
+				return true;
+			}
+			return false;
+		}
+
+		case T_NestPath:
+		case T_MergePath:
+		case T_HashPath:
+		{
+			JoinPath *path = (JoinPath *) node;
+			if (walker((Node *) path->outerjoinpath, context))
+				return true;
+			if (walker((Node *) path->innerjoinpath, context))
+				return true;
+			return false;
+		}
+
+		case T_AggPath:
+		{
+			AggPath *path = (AggPath *) node;
+			if (walker(path->subpath, context))
+				return true;
+			return false;
+		}
+
+		/* paths not having nesting */
+		case T_Path:
+		case T_IndexPath:
+		case T_BitmapHeapPath:
+		{
+			return false;
+		}
+
+		default:
+		{
+			elog(ERROR, "unrecognized node type: %d", (int) nodeTag(node));
+			return false;
+		}
+	}
+}
+
+typedef struct ColocationGroups
+{
+	Bitmapset *colocationIds;
+} ColocationGroups;
+
+static bool ColocationGroupsForPathWalker(Node *node, ColocationGroups *context);
+static bool
+ColocationGroupsForPathWalker(Node *node, ColocationGroups *context)
+{
+	if (node == NULL)
+		return false;
+
+	DistributedUnionPath *collect = NULL;
+	if (IsDistributedUnion((Path *) node, false, &collect))
+	{
+		context->colocationIds = bms_add_member(context->colocationIds,
+												(int) collect->colocationId);
+		return false;
+	}
+
+	return PathTreeWalker(node, ColocationGroupsForPathWalker, context);
+}
+
+
+static Bitmapset *
+ColocationGroupsForPath(Path *node)
+{
+	ColocationGroups context = { 0 };
+	ColocationGroupsForPathWalker((Node *)node, &context);
+	return context.colocationIds;
+}
+
+PathComparison
+PathBasedPlannerComparePath(Path *new_path, Path *old_path)
+{
+	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;
+}
diff --git a/src/backend/distributed/shared_library_init.c b/src/backend/distributed/shared_library_init.c
index 539b26f7d..cfdc979a1 100644
--- a/src/backend/distributed/shared_library_init.c
+++ b/src/backend/distributed/shared_library_init.c
@@ -299,6 +299,7 @@ _PG_init(void)
 	set_rel_pathlist_hook = multi_relation_restriction_hook;
 	set_join_pathlist_hook = multi_join_restriction_hook;
 	create_upper_paths_hook = PathBasedPlannedUpperPathHook;
+	compare_path_hook = PathBasedPlannerComparePath;
 
 	ExecutorStart_hook = CitusExecutorStart;
 	ExecutorRun_hook = CitusExecutorRun;
diff --git a/src/include/distributed/path_based_planner.h b/src/include/distributed/path_based_planner.h
index 34f6db686..104a97b73 100644
--- a/src/include/distributed/path_based_planner.h
+++ b/src/include/distributed/path_based_planner.h
@@ -46,4 +46,6 @@ extern void PathBasedPlannedUpperPathHook(PlannerInfo *root,
 										  RelOptInfo *output_rel,
 										  void *extra);
 
+extern PathComparison PathBasedPlannerComparePath(Path *new_path, Path *old_path);
+
 #endif //CITUS_PATH_BASED_PLANNER_H