From 2577544cc8b93d9a08a872353e3d260c5999ccef Mon Sep 17 00:00:00 2001
From: Matheus Alcantara <mths.dev@pm.me>
Date: Wed, 18 Dec 2024 16:54:01 -0300
Subject: [PATCH v1] JIT compilation per plan node

Introduce a new jit field on Planner struct to make it possible to
decide whether a plan node should be compiled or not.

Previously the cost of the entire query was used to decide whether it
should be compiled or not. This approach has a problem that it can exist
some expressions on plan tree that is not executed frequently, so it is
not worth to compile.

This commit change the consideration of JIT from plan level to
per-plan-node. The plan total cost and the number of times that the node
will be executed is used to decide if the plan node should be compiled
or not.
---
 src/backend/jit/jit.c                   |  8 ++++++++
 src/backend/optimizer/plan/createplan.c | 27 +++++++++++++++++++++++++
 src/backend/optimizer/plan/planner.c    |  3 +--
 src/include/nodes/plannodes.h           |  2 ++
 4 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/src/backend/jit/jit.c b/src/backend/jit/jit.c
index 815b58f33c..dfe010392f 100644
--- a/src/backend/jit/jit.c
+++ b/src/backend/jit/jit.c
@@ -170,6 +170,14 @@ jit_compile_expr(struct ExprState *state)
 	if (!(state->parent->state->es_jit_flags & PGJIT_EXPR))
 		return false;
 
+	/* don't jit if the plan node is missing */
+	if (state->parent->plan == NULL)
+		return false;
+
+	/* don't jit if it's not enabled for this plan node */
+	if (!state->parent->plan->jit)
+		return false;
+
 	/* this also takes !jit_enabled into account */
 	if (provider_init())
 		return provider.compile_expr(state);
diff --git a/src/backend/optimizer/plan/createplan.c b/src/backend/optimizer/plan/createplan.c
index 178c572b02..1e787c3437 100644
--- a/src/backend/optimizer/plan/createplan.c
+++ b/src/backend/optimizer/plan/createplan.c
@@ -25,6 +25,7 @@
 #include "nodes/extensible.h"
 #include "nodes/makefuncs.h"
 #include "nodes/nodeFuncs.h"
+#include "jit/jit.h"
 #include "optimizer/clauses.h"
 #include "optimizer/cost.h"
 #include "optimizer/optimizer.h"
@@ -320,6 +321,7 @@ static ModifyTable *make_modifytable(PlannerInfo *root, Plan *subplan,
 static GatherMerge *create_gather_merge_plan(PlannerInfo *root,
 											 GatherMergePath *best_path);
 
+static void plan_consider_jit(Plan *plan);
 
 /*
  * create_plan
@@ -551,9 +553,34 @@ create_plan_recurse(PlannerInfo *root, Path *best_path, int flags)
 			break;
 	}
 
+	plan_consider_jit(plan);
 	return plan;
 }
 
+static void
+plan_consider_jit(Plan *plan)
+{
+	plan->jit = false;
+
+	 /* Determine which JIT options to enable for this plan node */
+	if (jit_enabled)
+	{
+		Cost	total_cost;
+
+		/*
+		 * Take into account the number of times that we expect to rescan a
+		 * given plan node.  For example, subplans being invoked under the
+		 * inside of a Nested Loop may be rescanned many times.  JITing these
+		 * may be more worthwhile.
+		 */
+		total_cost = plan->total_cost * plan->plan_rows;
+
+		if (total_cost > jit_above_cost)
+			plan->jit = true;
+	}
+}
+
+
 /*
  * create_scan_plan
  *	 Create a scan plan for the parent relation of 'best_path'.
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index f3856c519f..e0b0e2f1ae 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -571,8 +571,7 @@ standard_planner(Query *parse, const char *query_string, int cursorOptions,
 	result->stmt_len = parse->stmt_len;
 
 	result->jitFlags = PGJIT_NONE;
-	if (jit_enabled && jit_above_cost >= 0 &&
-		top_plan->total_cost > jit_above_cost)
+	if (jit_enabled && jit_above_cost >= 0 && top_plan->jit)
 	{
 		result->jitFlags |= PGJIT_PERFORM;
 
diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h
index 52f29bcdb6..c4c6cd057f 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -135,6 +135,8 @@ typedef struct Plan
 	Cardinality plan_rows;		/* number of rows plan is expected to emit */
 	int			plan_width;		/* average row width in bytes */
 
+	bool jit; /* true if the plan node is worth to compile */
+
 	/*
 	 * information needed for parallel query
 	 */
-- 
2.39.3 (Apple Git-146)

