[YSQL] Batch updates that are pushed down to docdb (#5257)

Summary: As discussed, this is only enabling batching of updates for SQL procedures/functions.

Test Plan:
Running a slightly modified example from the github issue:
```
create table test(k int primary key, v int) split into 1 tablets;
insert into test values(3,3);

create procedure testproc() as $$
    insert into test values(1, 1);
    insert into test values(2, 2);
    update test set v=2 where k=3;
    insert into test values(4,4);
$$ LANGUAGE sql;

call testproc();
```
Only results in a single RPC to the tserver.

Reviewers: mihnea, dmitry, sudheer

Reviewed By: sudheer

Subscribers: yql

Differential Revision: https://phabricator.dev.yugabyte.com/D9307

src/postgres/src/backend/executor/execUtils.c

@@ -177,6 +177,7 @@ CreateExecutorState(void)
 	estate->yb_exec_params.limit_offset = 0;
 	estate->yb_exec_params.limit_use_default = true;
 	estate->yb_exec_params.rowmark = -1;
+	estate->yb_can_batch_updates = false;
 
 	return estate;
 }

src/postgres/src/backend/executor/functions.c

@@ -836,6 +836,12 @@ postquel_start(execution_state *es, SQLFunctionCachePtr fcache)
 		else
 			eflags = 0;			/* default run-to-completion flags */
 		ExecutorStart(es->qd, eflags);
+
+		/*
+		 * Since PGSQL functions don't require the row count from updates, we
+		 * can allow for batched updates.
+		 */
+		es->qd->estate->yb_can_batch_updates = true;
 	}
 
 	es->status = F_EXEC_RUN;

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

@@ -664,7 +664,13 @@ bool YBCExecuteUpdate(Relation rel,
 
 	/* Execute the statement. */
 	int rows_affected_count = 0;
-	YBCExecWriteStmt(update_stmt, rel, isSingleRow ? &rows_affected_count : NULL);
+
+	/* Currently only allows batching of single row updates for PGSQL procedures. */
+	bool can_batch_update = !isSingleRow ||
+							(YBCGetEnableUpdateBatching() && estate->yb_can_batch_updates);
+
+	/* If update batching is allowed, then ignore rows_affected_count. */
+	YBCExecWriteStmt(update_stmt, rel, can_batch_update ? NULL : &rows_affected_count);
 
 	/* Cleanup. */
 	update_stmt = NULL;

src/postgres/src/include/nodes/execnodes.h

@@ -142,7 +142,7 @@ typedef struct ExprState
  *		Am					Oid of index AM
  *		AmCache				private cache area for index AM
  *		Context				memory context holding this IndexInfo
- *		SplitOptions		Options to split index into tablets. 
+ *		SplitOptions		Options to split index into tablets.
  *
  * ii_Concurrent, ii_BrokenHotChain, and ii_ParallelWorkers are used only
  * during index build; they're conventionally zeroed otherwise.
@@ -590,6 +590,15 @@ typedef struct EState
 																		 * we cache the conflict tuple here when processing and
 																		 * then free the slot after the conflict is resolved. */
 	YBCPgExecParameters yb_exec_params;
+
+	/*
+	 * Whether we can batch updates - note that enabling this will cause batched
+	 * updates to not return a correct rows_affected_count, thus cannot be used
+	 * for plpgsql (which uses this value for GET DIAGNOSTICS...ROW_COUNT and
+	 * FOUND).
+	 * Currently only enabled for PGSQL functions / procedures.
+	 */
+	bool yb_can_batch_updates;
 } EState;
 
 

src/yb/yql/pggate/pggate_flags.cc

@@ -72,6 +72,10 @@ DEFINE_int32(ysql_output_buffer_size, 262144,
              "While fetched data resides within this buffer and hasn't been flushed to client yet, "
              "we're free to transparently restart operation in case of restart read error.");
 
+DEFINE_bool(ysql_enable_update_batching, true,
+            "Whether to enable batching of updates where possible. Currently update batching is "
+            "only supported for PGSQL procedures.");
+
 DEFINE_bool(ysql_suppress_unsupported_error, false,
             "Suppress ERROR on use of unsupported SQL statement and use WARNING instead");
 

src/yb/yql/pggate/pggate_flags.h

@@ -32,6 +32,7 @@ DECLARE_int32(ysql_max_read_restart_attempts);
 DECLARE_bool(TEST_ysql_disable_transparent_cache_refresh_retry);
 DECLARE_int32(ysql_output_buffer_size);
 DECLARE_int32(ysql_select_parallelism);
+DECLARE_bool(ysql_enable_update_batching);
 
 DECLARE_bool(ysql_suppress_unsupported_error);
 

src/yb/yql/pggate/ybc_pggate.cc

@@ -651,6 +651,10 @@ YBCStatus YBCPgExecUpdate(YBCPgStatement handle) {
   return ToYBCStatus(pgapi->ExecUpdate(handle));
 }
 
+bool YBCGetEnableUpdateBatching() {
+  return FLAGS_ysql_enable_update_batching;
+}
+
 // DELETE Operations -------------------------------------------------------------------------------
 YBCStatus YBCPgNewDelete(const YBCPgOid database_oid,
                          const YBCPgOid table_oid,

src/yb/yql/pggate/ybc_pggate.h

@@ -385,6 +385,9 @@ YBCStatus YBCPgNewUpdate(YBCPgOid database_oid,
 
 YBCStatus YBCPgExecUpdate(YBCPgStatement handle);
 
+// Retrieve value of ysql_enable_update_batching gflag.
+bool YBCGetEnableUpdateBatching();
+
 // DELETE ------------------------------------------------------------------------------------------
 YBCStatus YBCPgNewDelete(YBCPgOid database_oid,
                          YBCPgOid table_oid,