sql: disable buffered writes if ReturnRawMVCCValues needed

Buffered writes don't yet support the ReturnRawMVCCValues option of
Gets, Scans, and ReverseScans, which is needed to produce certain
system columns. Therefore, this commit disables buffered writes if
those system columns are requested when scanning data.

Fixes #144273

Release note: None

Author: rytaft

pkg/sql/colfetcher/cfetcher.go

@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/apd/v3"
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
 	"github.com/cockroachdb/cockroach/pkg/keys"
+	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catenumpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catpb"
@@ -202,6 +203,8 @@ type cFetcherArgs struct {
 	// the last one returned on the NextBatch calls if the caller wishes to keep
 	// multiple batches at the same time.
 	alwaysReallocate bool
+	// Txn is the txn for the fetch. It might be nil.
+	txn *kv.Txn
 }
 
 // noOutputColumn is a sentinel value to denote that a system column is not
@@ -440,6 +443,14 @@ func (cf *cFetcher) Init(
 		}
 	}
 
+	// Disable buffered writes if any system columns are needed that require
+	// MVCC decoding.
+	if cf.mvccDecodeStrategy == storage.MVCCDecodingRequired {
+		if cf.txn != nil && cf.txn.BufferedWritesEnabled() {
+			cf.txn.SetBufferedWritesEnabled(false /* enabled */)
+		}
+	}
+
 	fullColumns := table.spec.KeyFullColumns()
 
 	nIndexCols := len(fullColumns)

pkg/sql/colfetcher/cfetcher_wrapper.go

@@ -238,6 +238,7 @@ func newCFetcherWrapper(
 		true,  /* singleUse */
 		collectStats,
 		alwaysReallocate,
+		nil, /* txn */
 	}
 
 	// This memory monitor is not connected to the memory accounting system

pkg/sql/colfetcher/colbatch_scan.go

@@ -347,6 +347,7 @@ func NewColBatchScan(
 		true, /* singleUse */
 		shouldCollectStats,
 		false, /* alwaysReallocate */
+		flowCtx.Txn,
 	}
 	if err = fetcher.Init(fetcherAllocator, kvFetcher, tableArgs); err != nil {
 		fetcher.Release()

pkg/sql/colfetcher/index_join.go

@@ -602,6 +602,7 @@ func NewColIndexJoin(
 		false, /* singleUse */
 		shouldCollectStats,
 		false, /* alwaysReallocate */
+		txn,
 	}
 	if err = fetcher.Init(
 		fetcherAllocator, kvFetcher, tableArgs,

pkg/sql/logictest/testdata/logic_test/buffered_writes

@@ -252,9 +252,70 @@ SELECT * FROM t4
 2  200
 3  300
 
+subtest regression
+
 # Regression test for #144274.
 statement ok
 EXPLAIN CREATE DATABASE foo
 
 statement ok
 EXPLAIN ANALYZE CREATE DATABASE foo
+
+# Regression test for #144273.
+statement ok
+CREATE TABLE t144273 (
+  k INT PRIMARY KEY,
+  a INT,
+  b INT,
+  INDEX (a),
+  INDEX (b),
+  FAMILY (k, a, b)
+)
+
+statement ok
+PREPARE p AS UPDATE t144273 t1 SET a = t2.a + 1, b = t2.b + 1 FROM t144273 t2 WHERE t1.k = t2.a AND t1.k = $1
+
+statement ok
+SET vectorize = on
+
+statement ok
+INSERT INTO t144273 VALUES (1, 1, 1);
+
+# The table IDs in the kv trace below are different with the legacy schema
+# changer, so disable that configuration.
+skipif config local-legacy-schema-changer
+query T kvtrace
+EXECUTE p(1)
+----
+Scan /Table/114/2/{1-2}
+Scan /Table/114/1/1/0
+Scan /Table/114/1/1/0
+Put (locking) /Table/114/1/1/0 -> /TUPLE/2:2:Int/2/1:3:Int/2
+Del /Table/114/2/1/1/0
+Put /Table/114/2/2/1/0 -> /BYTES/
+Del /Table/114/3/1/1/0
+Put /Table/114/3/2/1/0 -> /BYTES/
+
+statement ok
+SET vectorize = off
+
+statement ok
+INSERT INTO t144273 VALUES (2, 2, 2);
+
+# The table IDs in the kv trace below are different with the legacy schema
+# changer, so disable that configuration.
+skipif config local-legacy-schema-changer
+query T kvtrace
+EXECUTE p(2)
+----
+Scan /Table/114/2/{2-3}
+Scan /Table/114/1/1/0, /Table/114/1/2/0
+Scan /Table/114/1/2/0
+Put (locking) /Table/114/1/2/0 -> /TUPLE/2:2:Int/3/1:3:Int/3
+Del /Table/114/2/2/2/0
+Put /Table/114/2/3/2/0 -> /BYTES/
+Del /Table/114/3/2/2/0
+Put /Table/114/3/3/2/0 -> /BYTES/
+
+statement ok
+RESET vectorize

pkg/sql/row/fetcher.go

@@ -409,6 +409,14 @@ func (rf *Fetcher) Init(ctx context.Context, args FetcherInitArgs) error {
 		}
 	}
 
+	// Disable buffered writes if any system columns are needed that require
+	// MVCC decoding.
+	if rf.mvccDecodeStrategy == storage.MVCCDecodingRequired {
+		if rf.args.Txn != nil && rf.args.Txn.BufferedWritesEnabled() {
+			rf.args.Txn.SetBufferedWritesEnabled(false /* enabled */)
+		}
+	}
+
 	if len(args.Spec.FetchedColumns) > 0 {
 		table.neededValueColsByIdx.AddRange(0, len(args.Spec.FetchedColumns)-1)
 	}
@@ -528,6 +536,14 @@ func (rf *Fetcher) SetTxn(txn *kv.Txn) error {
 // setTxnAndSendFn peeks inside of the KVFetcher to update the underlying
 // txnKVFetcher with the new txn and sendFn.
 func (rf *Fetcher) setTxnAndSendFn(txn *kv.Txn, sendFn sendFunc) error {
+	// Disable buffered writes if any system columns are needed that require
+	// MVCC decoding.
+	if rf.mvccDecodeStrategy == storage.MVCCDecodingRequired {
+		if txn != nil && txn.BufferedWritesEnabled() {
+			txn.SetBufferedWritesEnabled(false /* enabled */)
+		}
+	}
+
 	f, ok := rf.kvFetcher.KVBatchFetcher.(*txnKVFetcher)
 	if !ok {
 		return errors.AssertionFailedf(