From 96c822e322652a29d0eab791565c54c9d0cd73af Mon Sep 17 00:00:00 2001 From: Richard Guo Date: Wed, 12 Feb 2025 17:55:35 +0900 Subject: [PATCH v2] Expand virtual generated columns in planner --- src/backend/optimizer/plan/planner.c | 6 + src/backend/optimizer/prep/prepjointree.c | 180 ++++++++++++++++++++++ src/backend/rewrite/rewriteHandler.c | 14 +- src/backend/rewrite/rewriteManip.c | 2 +- src/include/optimizer/prep.h | 1 + src/include/rewrite/rewriteManip.h | 2 + 6 files changed, 193 insertions(+), 12 deletions(-) diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c index 7b1a8a0a9f..c80108a06c 100644 --- a/src/backend/optimizer/plan/planner.c +++ b/src/backend/optimizer/plan/planner.c @@ -749,6 +749,12 @@ subquery_planner(PlannerGlobal *glob, Query *parse, PlannerInfo *parent_root, if (parse->setOperations) flatten_simple_union_all(root); + /* + * Expand virtual generated columns. + * XXX more comments + */ + parse = root->parse = expand_virtual_generated_columns(root); + /* * Survey the rangetable to see what kinds of entries are present. We can * skip some later processing if relevant SQL features are not used; for diff --git a/src/backend/optimizer/prep/prepjointree.c b/src/backend/optimizer/prep/prepjointree.c index 5d9225e990..99ef55f7ac 100644 --- a/src/backend/optimizer/prep/prepjointree.c +++ b/src/backend/optimizer/prep/prepjointree.c @@ -9,6 +9,7 @@ * preprocess_function_rtes * pull_up_subqueries * flatten_simple_union_all + * expand_virtual_generated_columns * do expression preprocessing (including flattening JOIN alias vars) * reduce_outer_joins * remove_useless_result_rtes @@ -25,6 +26,7 @@ */ #include "postgres.h" +#include "access/table.h" #include "catalog/pg_type.h" #include "funcapi.h" #include "miscadmin.h" @@ -39,7 +41,9 @@ #include "optimizer/tlist.h" #include "parser/parse_relation.h" #include "parser/parsetree.h" +#include "rewrite/rewriteHandler.h" #include "rewrite/rewriteManip.h" +#include "utils/rel.h" typedef struct nullingrel_info @@ -57,6 +61,7 @@ typedef struct pullup_replace_vars_context { PlannerInfo *root; List *targetlist; /* tlist of subquery being pulled up */ + int result_relation; /* the index of the result relation */ RangeTblEntry *target_rte; /* RTE of subquery */ Relids relids; /* relids within subquery, as numbered after * pullup (set only if target_rte->lateral) */ @@ -1273,6 +1278,7 @@ pull_up_simple_subquery(PlannerInfo *root, Node *jtnode, RangeTblEntry *rte, */ rvcontext.root = root; rvcontext.targetlist = subquery->targetList; + rvcontext.result_relation = 0; rvcontext.target_rte = rte; if (rte->lateral) { @@ -1833,6 +1839,7 @@ pull_up_simple_values(PlannerInfo *root, Node *jtnode, RangeTblEntry *rte) } rvcontext.root = root; rvcontext.targetlist = tlist; + rvcontext.result_relation = 0; rvcontext.target_rte = rte; rvcontext.relids = NULL; /* can't be any lateral references here */ rvcontext.nullinfo = NULL; @@ -1992,6 +1999,7 @@ pull_up_constant_function(PlannerInfo *root, Node *jtnode, 1, /* resno */ NULL, /* resname */ false)); /* resjunk */ + rvcontext.result_relation = 0; rvcontext.target_rte = rte; /* @@ -2499,6 +2507,10 @@ pullup_replace_vars_callback(Var *var, */ need_phv = (var->varnullingrels != NULL) || rcon->wrap_non_vars; + /* System columns are not replaced. */ + if (varattno < InvalidAttrNumber) + return (Node *) copyObject(var); + /* * If PlaceHolderVars are needed, we cache the modified expressions in * rcon->rv_cache[]. This is not in hopes of any material speed gain @@ -2559,6 +2571,22 @@ pullup_replace_vars_callback(Var *var, rowexpr->location = var->location; newnode = (Node *) rowexpr; + /* + * Handle any OLD/NEW RETURNING list Vars + * + * XXX comments about wrapping it into ReturningExpr + */ + if (var->varreturningtype != VAR_RETURNING_DEFAULT) + { + ReturningExpr *rexpr = makeNode(ReturningExpr); + + rexpr->retlevelsup = var->varlevelsup; + rexpr->retold = (var->varreturningtype == VAR_RETURNING_OLD); + rexpr->retexpr = (Expr *) rowexpr; + + newnode = (Node *) rexpr; + } + /* * Insert PlaceHolderVar if needed. Notice that we are wrapping one * PlaceHolderVar around the whole RowExpr, rather than putting one @@ -2588,6 +2616,37 @@ pullup_replace_vars_callback(Var *var, /* Make a copy of the tlist item to return */ newnode = (Node *) copyObject(tle->expr); + /* + * Handle any OLD/NEW RETURNING list Vars + * + * XXX comments about wrapping it into ReturningExpr + */ + if (var->varreturningtype != VAR_RETURNING_DEFAULT) + { + /* + * Copy varreturningtype onto any Vars in the tlist item that + * refer to result_relation (which had better be non-zero). + */ + if (rcon->result_relation == 0) + elog(ERROR, "variable returning old/new found outside RETURNING list"); + + SetVarReturningType((Node *) newnode, rcon->result_relation, + var->varlevelsup, var->varreturningtype); + + if (!IsA(newnode, Var) || + ((Var *) newnode)->varno != rcon->result_relation || + ((Var *) newnode)->varlevelsup != var->varlevelsup) + { + ReturningExpr *rexpr = makeNode(ReturningExpr); + + rexpr->retlevelsup = var->varlevelsup; + rexpr->retold = (var->varreturningtype == VAR_RETURNING_OLD); + rexpr->retexpr = (Expr *) newnode; + + newnode = (Node *) rexpr; + } + } + /* Insert PlaceHolderVar if needed */ if (need_phv) { @@ -2947,6 +3006,127 @@ flatten_simple_union_all(PlannerInfo *root) } +/* + * expand_virtual_generated_columns + * Expand all virtual generated column references in a query. + * + * XXX more comments + */ +Query * +expand_virtual_generated_columns(PlannerInfo *root) +{ + Query *parse = root->parse; + int rt_index; + ListCell *lc; + + rt_index = 0; + foreach(lc, parse->rtable) + { + RangeTblEntry *rte = (RangeTblEntry *) lfirst(lc); + Relation rel; + TupleDesc tupdesc; + + ++rt_index; + + /* + * Only normal relations can have virtual generated columns. + */ + if (rte->rtekind != RTE_RELATION) + continue; + + rel = table_open(rte->relid, NoLock); + + tupdesc = RelationGetDescr(rel); + if (tupdesc->constr && tupdesc->constr->has_generated_virtual) + { + List *tlist = NIL; + pullup_replace_vars_context rvcontext; + + for (int i = 0; i < tupdesc->natts; i++) + { + Form_pg_attribute attr = TupleDescAttr(tupdesc, i); + int attnum = i + 1; + TargetEntry *te; + + if (attr->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL) + { + Node *defexpr; + Oid attcollid; + + defexpr = build_column_default(rel, attnum); + if (defexpr == NULL) + elog(ERROR, "no generation expression found for column number %d of table \"%s\"", + attnum, RelationGetRelationName(rel)); + + /* + * If the column definition has a collation and it is + * different from the collation of the generation expression, + * put a COLLATE clause around the expression. + */ + attcollid = attr->attcollation; + if (attcollid && attcollid != exprCollation(defexpr)) + { + CollateExpr *ce = makeNode(CollateExpr); + + ce->arg = (Expr *) defexpr; + ce->collOid = attcollid; + ce->location = -1; + + defexpr = (Node *) ce; + } + + ChangeVarNodes(defexpr, 1, rt_index, 0); + + te = makeTargetEntry((Expr *) defexpr, attnum, 0, false); + tlist = lappend(tlist, te); + } + else + { + Var *var; + + var = makeVar(rt_index, + attnum, + attr->atttypid, + attr->atttypmod, + attr->attcollation, + 0); + te = makeTargetEntry((Expr *) var, attnum, 0, false); + tlist = lappend(tlist, te); + } + } + + Assert(list_length(tlist) > 0); + + rvcontext.root = root; + rvcontext.targetlist = tlist; + rvcontext.result_relation = parse->resultRelation; + rvcontext.target_rte = rte; + rvcontext.relids = NULL; + rvcontext.nullinfo = NULL; + rvcontext.outer_hasSubLinks = NULL; + rvcontext.varno = rt_index; + rvcontext.wrap_non_vars = false; + rvcontext.rv_cache = palloc0((list_length(tlist) + 1) * + sizeof(Node *)); + /* + * If the query uses grouping sets, we need a PlaceHolderVar for + * anything that's not a simple Var. This ensures that expressions + * retain their separate identity so that they will match grouping + * set columns when appropriate. + */ + if (parse->groupingSets) + rvcontext.wrap_non_vars = true; + + parse = (Query *) pullup_replace_vars((Node *) parse, &rvcontext); + } + + table_close(rel, NoLock); + } + + return parse; +} + + /* * reduce_outer_joins * Attempt to reduce outer joins to plain inner joins. diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c index e996bdc0d2..0b1fcb3758 100644 --- a/src/backend/rewrite/rewriteHandler.c +++ b/src/backend/rewrite/rewriteHandler.c @@ -2211,7 +2211,9 @@ fireRIRrules(Query *parsetree, List *activeRIRs) * Only normal relations can have RLS policies or virtual generated * columns. */ - if (rte->rtekind != RTE_RELATION) + if (rte->rtekind != RTE_RELATION || + (rte->relkind != RELKIND_RELATION && + rte->relkind != RELKIND_PARTITIONED_TABLE)) continue; rel = table_open(rte->relid, NoLock); @@ -2300,16 +2302,6 @@ fireRIRrules(Query *parsetree, List *activeRIRs) if (hasSubLinks) parsetree->hasSubLinks = true; - /* - * Expand any references to virtual generated columns of this table. - * Note that subqueries in virtual generated column expressions are - * not currently supported, so this cannot add any more sublinks. - */ - parsetree = (Query *) - expand_generated_columns_internal((Node *) parsetree, - rel, rt_index, rte, - parsetree->resultRelation); - table_close(rel, NoLock); } diff --git a/src/backend/rewrite/rewriteManip.c b/src/backend/rewrite/rewriteManip.c index 9433548d27..6994b8c542 100644 --- a/src/backend/rewrite/rewriteManip.c +++ b/src/backend/rewrite/rewriteManip.c @@ -1010,7 +1010,7 @@ SetVarReturningType_walker(Node *node, SetVarReturningType_context *context) return expression_tree_walker(node, SetVarReturningType_walker, context); } -static void +void SetVarReturningType(Node *node, int result_relation, int sublevels_up, VarReturningType returning_type) { diff --git a/src/include/optimizer/prep.h b/src/include/optimizer/prep.h index 0ae57ec24a..6bc75ac9f3 100644 --- a/src/include/optimizer/prep.h +++ b/src/include/optimizer/prep.h @@ -27,6 +27,7 @@ extern void pull_up_sublinks(PlannerInfo *root); extern void preprocess_function_rtes(PlannerInfo *root); extern void pull_up_subqueries(PlannerInfo *root); extern void flatten_simple_union_all(PlannerInfo *root); +extern Query *expand_virtual_generated_columns(PlannerInfo *root); extern void reduce_outer_joins(PlannerInfo *root); extern void remove_useless_result_rtes(PlannerInfo *root); extern Relids get_relids_in_jointree(Node *jtnode, bool include_outer_joins, diff --git a/src/include/rewrite/rewriteManip.h b/src/include/rewrite/rewriteManip.h index 5ec475c63e..2f09657ab2 100644 --- a/src/include/rewrite/rewriteManip.h +++ b/src/include/rewrite/rewriteManip.h @@ -97,5 +97,7 @@ extern Node *ReplaceVarsFromTargetList(Node *node, ReplaceVarsNoMatchOption nomatch_option, int nomatch_varno, bool *outer_hasSubLinks); +extern void SetVarReturningType(Node *node, int result_relation, int sublevels_up, + VarReturningType returning_type); #endif /* REWRITEMANIP_H */ -- 2.43.0