[#23461] YSQL: Fix memory corruption in UPDATE pushdown

Summary:
The mutated expressions should not be placed into the per tuple
memory context, it may be reset during the mutation.

The problem was introduced by #16255 where allocations were released
a bit too aggressively.

To fix the issue perform pushdown expression mutation outside of
the per tuple context.
Jira: DB-12381

Test Plan: ./yb_build.sh --java-test 'org.yb.pgsql.TestPgRegressDml#testPgRegressDml'

Reviewers: telgersma

Reviewed By: telgersma

Subscribers: yql

Tags: #jenkins-ready

Differential Revision: https://phorge.dev.yugabyte.com/D37251

src/postgres/src/backend/executor/ybcModifyTable.c

@@ -898,19 +898,21 @@ bool YBCExecuteUpdate(Relation rel,
 		 */
 		Assert(!bms_is_member(attnum - minattr, YBGetTablePrimaryKeyBms(rel)));
 
-		MemoryContext oldContext = MemoryContextSwitchTo(GetPerTupleMemoryContext(estate));
 		/* Assign this attr's value, handle expression pushdown if needed. */
 		if (pushdown_lc != NULL &&
 			((TargetEntry *) lfirst(pushdown_lc))->resno == attnum)
 		{
 			TargetEntry *tle = (TargetEntry *) lfirst(pushdown_lc);
 			Expr *expr = YbExprInstantiateParams(tle->expr, estate);
+			MemoryContext oldContext = MemoryContextSwitchTo(GetPerTupleMemoryContext(estate));
 			YBCPgExpr ybc_expr = YBCNewEvalExprCall(update_stmt, expr);
 			HandleYBStatus(YBCPgDmlAssignColumn(update_stmt, attnum, ybc_expr));
 			pushdown_lc = lnext(pushdown_lc);
+			MemoryContextSwitchTo(oldContext);
 		}
 		else
 		{
+			MemoryContext oldContext = MemoryContextSwitchTo(GetPerTupleMemoryContext(estate));
 			bool is_null = false;
 			Datum d = heap_getattr(tuple, attnum, inputTupleDesc, &is_null);
 			/*
@@ -924,8 +926,8 @@ bool YBCExecuteUpdate(Relation rel,
 			YBCPgExpr ybc_expr = YBCNewConstant(update_stmt, type_id, collation_id, d, is_null);
 
 			HandleYBStatus(YBCPgDmlAssignColumn(update_stmt, attnum, ybc_expr));
+			MemoryContextSwitchTo(oldContext);
 		}
-		MemoryContextSwitchTo(oldContext);
 	}
 
 	/*

src/postgres/src/test/regress/expected/yb_dml_pushdown.out

@@ -564,6 +564,23 @@ SELECT * FROM single_row;
  1 |  3 |  3
 (1 row)
 
+-- GHI 23461
+CREATE TABLE single_row_changes(k int, v1 int);
+INSERT INTO single_row_changes VALUES (1, 1), (1, 2);
+WITH deleted_rows AS (
+  DELETE FROM single_row_changes
+  WHERE k = 1
+  RETURNING v1
+)
+UPDATE single_row
+SET v1 = v1 + (SELECT COALESCE(SUM(v1), 0) FROM deleted_rows)
+WHERE k = 1;
+SELECT * FROM single_row;
+ k | v1 | v2
+---+----+----
+ 1 |  6 |  3
+(1 row)
+
 -- Updates on a table with a secondary index
 CREATE INDEX single_row_v2_idx ON single_row(v2);
 EXPLAIN (COSTS FALSE) WITH temp AS (UPDATE single_row SET v1 = v1 + 1 WHERE k = 1 RETURNING k, v1)

src/postgres/src/test/regress/sql/yb_dml_pushdown.sql

@@ -163,6 +163,21 @@ WITH temp AS (UPDATE single_row SET v1 = v1 + 1 WHERE k = 2 RETURNING k, v1)
 
 SELECT * FROM single_row;
 
+-- GHI 23461
+CREATE TABLE single_row_changes(k int, v1 int);
+INSERT INTO single_row_changes VALUES (1, 1), (1, 2);
+
+WITH deleted_rows AS (
+  DELETE FROM single_row_changes
+  WHERE k = 1
+  RETURNING v1
+)
+UPDATE single_row
+SET v1 = v1 + (SELECT COALESCE(SUM(v1), 0) FROM deleted_rows)
+WHERE k = 1;
+
+SELECT * FROM single_row;
+
 -- Updates on a table with a secondary index
 CREATE INDEX single_row_v2_idx ON single_row(v2);
 EXPLAIN (COSTS FALSE) WITH temp AS (UPDATE single_row SET v1 = v1 + 1 WHERE k = 1 RETURNING k, v1)