Skip to content

Commit 0feb47b

Browse files
funbringerfunbringer
committed
implement state_tree_visitor()
1 parent dde5eb2 commit 0feb47b

File tree

5 files changed

+170
-51
lines changed

5 files changed

+170
-51
lines changed

‎src/hooks.c‎

Lines changed: 4 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -976,35 +976,10 @@ pathman_executor_hook(QueryDesc *queryDesc,
976976
#define EXECUTOR_RUN(q,d,c) standard_ExecutorRun((q),(d),(c))
977977
#endif
978978

979-
PlanState *state = (PlanState *) queryDesc->planstate;
980-
981-
/* FIXME: we should modify ALL ModifyTable nodes! They might be hidden deeper. */
982-
if (IsA(state, ModifyTableState))
983-
{
984-
ModifyTableState *mt_state = (ModifyTableState *) state;
985-
int i;
986-
987-
for (i = 0; i < mt_state->mt_nplans; i++)
988-
{
989-
CustomScanState *pr_state = (CustomScanState *) mt_state->mt_plans[i];
990-
991-
/* Check if this is a PartitionFilter + PartitionRouter combo */
992-
if (IsPartitionFilterState(pr_state) &&
993-
IsPartitionRouterState(linitial(pr_state->custom_ps)))
994-
{
995-
ResultRelInfo *rri = &mt_state->resultRelInfo[i];
996-
997-
/*
998-
* HACK: We unset junkfilter to disable
999-
* junk cleaning in ExecModifyTable.
1000-
*/
1001-
rri->ri_junkFilter = NULL;
1002-
1003-
/* HACK: change UPDATE operation to INSERT */
1004-
mt_state->operation = CMD_INSERT;
1005-
}
1006-
}
1007-
}
979+
/* Prepare ModifyTable nodes for PartitionRouter hackery */
980+
state_tree_visitor((PlanState *) queryDesc->planstate,
981+
prepare_modify_table_for_partition_router,
982+
NULL);
1008983

1009984
/* Call hooks set by other extensions if needed */
1010985
if (EXECUTOR_HOOK)

‎src/include/partition_router.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,8 @@ Plan *make_partition_router(Plan *subplan,
7272
int epq_param,
7373
List *returning_list);
7474

75+
void prepare_modify_table_for_partition_router(PlanState *state, void *context);
76+
7577

7678
Node *partition_router_create_scan_state(CustomScan *node);
7779

‎src/include/planner_tree_modification.h‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,14 @@ void assign_query_id(Query *query);
2525
void reset_query_id_generator(void);
2626

2727
/* Plan tree rewriting utility */
28-
void plan_tree_walker(Plan *plan,
29-
void (*visitor) (Plan *plan, void *context),
30-
void *context);
28+
void plan_tree_visitor(Plan *plan,
29+
void (*visitor) (Plan *plan, void *context),
30+
void *context);
31+
32+
/* PlanState tree rewriting utility */
33+
void state_tree_visitor(PlanState *state,
34+
void (*visitor) (PlanState *state, void *context),
35+
void *context);
3136

3237
/* Query tree rewriting utility */
3338
void pathman_transform_query(Query *parse, ParamListInfo params);

‎src/partition_router.c‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,38 @@ make_partition_router(Plan *subplan,
9595
return &cscan->scan.plan;
9696
}
9797

98+
void
99+
prepare_modify_table_for_partition_router(PlanState *state, void *context)
100+
{
101+
if (IsA(state, ModifyTableState))
102+
{
103+
ModifyTableState *mt_state = (ModifyTableState *) state;
104+
int i;
105+
106+
for (i = 0; i < mt_state->mt_nplans; i++)
107+
{
108+
CustomScanState *pr_state = (CustomScanState *) mt_state->mt_plans[i];
109+
110+
/* Check if this is a PartitionFilter + PartitionRouter combo */
111+
if (IsPartitionFilterState(pr_state) &&
112+
IsPartitionRouterState(linitial(pr_state->custom_ps)))
113+
{
114+
ResultRelInfo *rri = &mt_state->resultRelInfo[i];
115+
116+
/*
117+
* HACK: We unset junkfilter to disable
118+
* junk cleaning in ExecModifyTable.
119+
*/
120+
rri->ri_junkFilter = NULL;
121+
122+
/* HACK: change UPDATE operation to INSERT */
123+
mt_state->operation = CMD_INSERT;
124+
}
125+
}
126+
}
127+
}
128+
129+
98130
Node *
99131
partition_router_create_scan_state(CustomScan *node)
100132
{

‎src/planner_tree_modification.c‎

Lines changed: 124 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,9 @@ static void handle_modification_query(Query *parse, transform_query_cxt *context
113113
static void partition_filter_visitor(Plan *plan, void *context);
114114
static void partition_router_visitor(Plan *plan, void *context);
115115

116+
static void state_visit_subplans(List *plans, void (*visitor) (), void *context);
117+
static void state_visit_members(PlanState **planstates, int nplans, void (*visitor) (), void *context);
118+
116119
static Oid find_deepest_partition(Oid relid, Index rti, Expr *quals);
117120
static Node *eval_extern_params_mutator(Node *node, ParamListInfo params);
118121
static Node *adjust_appendrel_varnos(Node *node, adjust_appendrel_varnos_cxt *context);
@@ -152,9 +155,9 @@ reset_query_id_generator(void)
152155
* 'visitor' is applied right before return.
153156
*/
154157
void
155-
plan_tree_walker(Plan *plan,
156-
void (*visitor) (Plan *plan, void *context),
157-
void *context)
158+
plan_tree_visitor(Plan *plan,
159+
void (*visitor) (Plan *plan, void *context),
160+
void *context)
158161
{
159162
ListCell *l;
160163

@@ -167,50 +170,152 @@ plan_tree_walker(Plan *plan,
167170
switch (nodeTag(plan))
168171
{
169172
case T_SubqueryScan:
170-
plan_tree_walker(((SubqueryScan *) plan)->subplan, visitor, context);
173+
plan_tree_visitor(((SubqueryScan *) plan)->subplan, visitor, context);
171174
break;
172175

173176
case T_CustomScan:
174-
foreach(l, ((CustomScan *) plan)->custom_plans)
175-
plan_tree_walker((Plan *) lfirst(l), visitor, context);
177+
foreach (l, ((CustomScan *) plan)->custom_plans)
178+
plan_tree_visitor((Plan *) lfirst(l), visitor, context);
176179
break;
177180

178181
case T_ModifyTable:
179182
foreach (l, ((ModifyTable *) plan)->plans)
180-
plan_tree_walker((Plan *) lfirst(l), visitor, context);
183+
plan_tree_visitor((Plan *) lfirst(l), visitor, context);
181184
break;
182185

183186
case T_Append:
184-
foreach(l, ((Append *) plan)->appendplans)
185-
plan_tree_walker((Plan *) lfirst(l), visitor, context);
187+
foreach (l, ((Append *) plan)->appendplans)
188+
plan_tree_visitor((Plan *) lfirst(l), visitor, context);
186189
break;
187190

188191
case T_MergeAppend:
189-
foreach(l, ((MergeAppend *) plan)->mergeplans)
190-
plan_tree_walker((Plan *) lfirst(l), visitor, context);
192+
foreach (l, ((MergeAppend *) plan)->mergeplans)
193+
plan_tree_visitor((Plan *) lfirst(l), visitor, context);
191194
break;
192195

193196
case T_BitmapAnd:
194-
foreach(l, ((BitmapAnd *) plan)->bitmapplans)
195-
plan_tree_walker((Plan *) lfirst(l), visitor, context);
197+
foreach (l, ((BitmapAnd *) plan)->bitmapplans)
198+
plan_tree_visitor((Plan *) lfirst(l), visitor, context);
196199
break;
197200

198201
case T_BitmapOr:
199-
foreach(l, ((BitmapOr *) plan)->bitmapplans)
200-
plan_tree_walker((Plan *) lfirst(l), visitor, context);
202+
foreach (l, ((BitmapOr *) plan)->bitmapplans)
203+
plan_tree_visitor((Plan *) lfirst(l), visitor, context);
201204
break;
202205

203206
default:
204207
break;
205208
}
206209

207-
plan_tree_walker(plan->lefttree, visitor, context);
208-
plan_tree_walker(plan->righttree, visitor, context);
210+
plan_tree_visitor(plan->lefttree, visitor, context);
211+
plan_tree_visitor(plan->righttree, visitor, context);
209212

210213
/* Apply visitor to the current node */
211214
visitor(plan, context);
212215
}
213216

217+
void
218+
state_tree_visitor(PlanState *state,
219+
void (*visitor) (PlanState *plan, void *context),
220+
void *context)
221+
{
222+
Plan *plan;
223+
ListCell *lc;
224+
225+
if (state == NULL)
226+
return;
227+
228+
plan = state->plan;
229+
230+
check_stack_depth();
231+
232+
/* Plan-type-specific fixes */
233+
switch (nodeTag(plan))
234+
{
235+
case T_SubqueryScan:
236+
state_tree_visitor(((SubqueryScanState *) state)->subplan, visitor, context);
237+
break;
238+
239+
case T_CustomScan:
240+
foreach (lc, ((CustomScanState *) state)->custom_ps)
241+
state_tree_visitor((PlanState *) lfirst(lc),visitor, context);
242+
break;
243+
244+
case T_ModifyTable:
245+
state_visit_members(((ModifyTableState *) state)->mt_plans,
246+
((ModifyTableState *) state)->mt_nplans,
247+
visitor, context);
248+
break;
249+
250+
case T_Append:
251+
state_visit_members(((AppendState *) state)->appendplans,
252+
((AppendState *) state)->as_nplans,
253+
visitor, context);
254+
break;
255+
256+
case T_MergeAppend:
257+
state_visit_members(((MergeAppendState *) state)->mergeplans,
258+
((MergeAppendState *) state)->ms_nplans,
259+
visitor, context);
260+
break;
261+
262+
case T_BitmapAnd:
263+
state_visit_members(((BitmapAndState *) state)->bitmapplans,
264+
((BitmapAndState *) state)->nplans,
265+
visitor, context);
266+
break;
267+
268+
case T_BitmapOr:
269+
state_visit_members(((BitmapOrState *) state)->bitmapplans,
270+
((BitmapOrState *) state)->nplans,
271+
visitor, context);
272+
break;
273+
274+
default:
275+
break;
276+
}
277+
278+
state_visit_subplans(state->initPlan, visitor, context);
279+
state_visit_subplans(state->subPlan, visitor, context);
280+
281+
state_tree_visitor(state->lefttree, visitor, context);
282+
state_tree_visitor(state->righttree, visitor, context);
283+
284+
/* Apply visitor to the current node */
285+
visitor(state, context);
286+
}
287+
288+
/*
289+
* Walk a list of SubPlans (or initPlans, which also use SubPlan nodes).
290+
*/
291+
static void
292+
state_visit_subplans(List *plans,
293+
void (*visitor) (),
294+
void *context)
295+
{
296+
ListCell *lc;
297+
298+
foreach (lc, plans)
299+
{
300+
SubPlanState *sps = lfirst_node(SubPlanState, lc);
301+
visitor(sps->planstate, context);
302+
}
303+
}
304+
305+
/*
306+
* Walk the constituent plans of a ModifyTable, Append, MergeAppend,
307+
* BitmapAnd, or BitmapOr node.
308+
*/
309+
static void
310+
state_visit_members(PlanState **planstates, int nplans,
311+
void (*visitor) (), void *context)
312+
{
313+
int i;
314+
315+
for (i = 0; i < nplans; i++)
316+
visitor(planstates[i], context);
317+
}
318+
214319

215320
/*
216321
* -------------------------------
@@ -586,15 +691,15 @@ void
586691
add_partition_filters(List *rtable, Plan *plan)
587692
{
588693
if (pg_pathman_enable_partition_filter)
589-
plan_tree_walker(plan, partition_filter_visitor, rtable);
694+
plan_tree_visitor(plan, partition_filter_visitor, rtable);
590695
}
591696

592697
/* Add PartitionRouter nodes to the plan tree */
593698
void
594699
add_partition_routers(List *rtable, Plan *plan)
595700
{
596701
if (pg_pathman_enable_partition_router)
597-
plan_tree_walker(plan, partition_router_visitor, rtable);
702+
plan_tree_visitor(plan, partition_router_visitor, rtable);
598703
}
599704

600705
/*

0 commit comments

Comments
 (0)