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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy