Refactor insert select planner to exclude identity columns from both insert and select target lists

m3hm3t/issue_7887
Mehmet Yilmaz 2025-02-28 15:10:38 +00:00
parent 67b390ee3e
commit f4835b7dbb
1 changed files with 50 additions and 32 deletions

View File

@ -1435,12 +1435,8 @@ CreateNonPushableInsertSelectPlan(uint64 planId, Query *parse, ParamListInfo bou
PrepareInsertSelectForCitusPlanner(insertSelectQuery); PrepareInsertSelectForCitusPlanner(insertSelectQuery);
/* /* get the SELECT query (may have changed after PrepareInsertSelectForCitusPlanner) */
* The insertTargetList are the columns we plan to insert into the target table. Query *selectQuery = selectRte->subquery;
* For partial inserts, it might incorrectly include the identity column if
* some rewriting logic added it. We'll fix that below.
*/
List *insertTargetList = insertSelectQuery->targetList;
/* /*
* 1) Open the target relation to inspect its attributes and detect identity columns. * 1) Open the target relation to inspect its attributes and detect identity columns.
@ -1448,19 +1444,40 @@ CreateNonPushableInsertSelectPlan(uint64 planId, Query *parse, ParamListInfo bou
Relation targetRel = RelationIdGetRelation(targetRelationId); Relation targetRel = RelationIdGetRelation(targetRelationId);
if (RelationIsValid(targetRel)) if (RelationIsValid(targetRel))
{ {
/* We'll build a new list of TLEs that excludes identity columns if user omitted them. */
List *newTargetList = NIL;
ListCell *lc = NULL;
foreach(lc, insertTargetList) ListCell *lcInsert = NULL;
// ListCell *lcSelect = list_head(selectQuery->targetList);
/* We'll build new lists for both sides */
List *newInsertTList = NIL;
List *newSelectTList = NIL;
int insertIndex = 0;
foreach(lcInsert, insertSelectQuery->targetList)
{ {
TargetEntry *tle = (TargetEntry *) lfirst(lc); TargetEntry *insertTle = (TargetEntry *) lfirst(lcInsert);
insertIndex++;
/* Get the corresponding TLE from the SELECT by position or resno */
TargetEntry *selectTle = NULL;
/* /*
* resno is 1-based attribute number: if we have 3 columns in table, they * If your plan is guaranteed to keep them in the same order, you can
* correspond to resno=1..3. Make sure attno is in range before we do anything. * do "selectTle = (TargetEntry *) list_nth(selectQuery->targetList, insertIndex - 1)".
*
* Alternatively, if you rely on resno alignment, you'd find the TLE with resno==insertTle->resno.
* For simplicity, let's assume same ordering:
*/ */
int attno = tle->resno; selectTle = (TargetEntry *) list_nth(selectQuery->targetList, insertIndex - 1);
/*
* Check if the insertTle is an identity column that the user didn't supply,
* e.g. by checking 'attr->attidentity == ATTRIBUTE_IDENTITY_ALWAYS' etc.
* If we skip it, also skip the SELECT TLE at the same position.
*/
int attno = insertTle->resno;
if (attno > 0 && attno <= targetRel->rd_att->natts) if (attno > 0 && attno <= targetRel->rd_att->natts)
{ {
Form_pg_attribute attr = TupleDescAttr(targetRel->rd_att, attno - 1); Form_pg_attribute attr = TupleDescAttr(targetRel->rd_att, attno - 1);
@ -1471,7 +1488,8 @@ CreateNonPushableInsertSelectPlan(uint64 planId, Query *parse, ParamListInfo bou
* indicated by something in the parse tree?), we remove or convert * indicated by something in the parse tree?), we remove or convert
* the TLE to a default. * the TLE to a default.
*/ */
bool userSpecifiedValue = CheckIfUserSpecifiedValue(tle, parse); // bool userSpecifiedValue = CheckIfUserSpecifiedValue(tle, parse);
bool userSpecifiedValue = false;
if ((attr->attidentity == ATTRIBUTE_IDENTITY_ALWAYS || if ((attr->attidentity == ATTRIBUTE_IDENTITY_ALWAYS ||
attr->attidentity == ATTRIBUTE_IDENTITY_BY_DEFAULT) && attr->attidentity == ATTRIBUTE_IDENTITY_BY_DEFAULT) &&
!userSpecifiedValue) !userSpecifiedValue)
@ -1481,18 +1499,18 @@ CreateNonPushableInsertSelectPlan(uint64 planId, Query *parse, ParamListInfo bou
} }
} }
/* If we get here, we keep the TLE. */ /* else keep both TLEs */
newTargetList = lappend(newTargetList, tle); newInsertTList = lappend(newInsertTList, insertTle);
newSelectTList = lappend(newSelectTList, selectTle);
} }
/* Update the plan's target list to the "cleaned" version */ /* Now we have 1:1 matching lists with the identity column removed from both sides */
insertSelectQuery->targetList = newTargetList; insertSelectQuery->targetList = newInsertTList;
selectQuery->targetList = newSelectTList;
RelationClose(targetRel); RelationClose(targetRel);
} }
/* get the SELECT query (may have changed after PrepareInsertSelectForCitusPlanner) */
Query *selectQuery = selectRte->subquery;
/* /*
* Later we might need to call WrapTaskListForProjection(), which requires * Later we might need to call WrapTaskListForProjection(), which requires