Merge #144635

144635: sql: disable buffered writes if ReturnRawMVCCValues needed r=rytaft a=rytaft

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

Co-authored-by: Rebecca Taft <[email protected]>

Author: craig[bot]

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(