colbuilder: don't use direct scans in an invalid case

We have an experimental feature to build `coldata.Batch`es directly on
the KV server side when handling Scan and ReverseScan requests.
It assumes that each request corresponds to any number of full SQL rows,
but previously we didn't disable usage of this feature when a single SQL
row is populated via multiple requests. This is the case when we have
multiple column families, and since we don't support Gets in the direct
scans machinery right now, we need at least 5 column families - so that
we get Scan for 0-1 and 3-4 CFs. This commit fixes this oversight to
ensure that if we have two consecutive key spans that come from the same
SQL row, we don't use the direct scans even if they are enabled.

There is no release note since this feature is disabled by default.

Release note: None

Author: yuzefovich

pkg/ccl/logictestccl/tests/3node-tenant/generated_test.go

@@ -12,6 +12,7 @@ go_library(
         "//pkg/col/coldata",
         "//pkg/col/coldataext",
         "//pkg/col/typeconv",
+        "//pkg/keys",
         "//pkg/kv/kvserver/concurrency/lock",
         "//pkg/settings",
         "//pkg/sql/catalog/descpb",

pkg/ccl/logictestccl/tests/local-read-committed/generated_test.go

@@ -6,13 +6,15 @@
 package colbuilder
 
 import (
+	"bytes"
 	"context"
 	"reflect"
 	"strings"
 
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
 	"github.com/cockroachdb/cockroach/pkg/col/coldataext"
 	"github.com/cockroachdb/cockroach/pkg/col/typeconv"
+	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
@@ -876,17 +878,38 @@ func NewColOperator(
 						core.TableReader.LockingWaitPolicy == descpb.ScanLockingWaitPolicy_SKIP_LOCKED {
 						return false
 					}
-					// At the moment, the ColBatchDirectScan cannot handle Gets
-					// (it's not clear whether it is worth to handle them via
-					// the same path as for Scans and ReverseScans (which could
-					// have too large of an overhead) or by teaching the
-					// operator to also decode a single KV (similar to what
-					// regular ColBatchScan does)).
-					// TODO(yuzefovich, 23.1): explore supporting Gets somehow.
-					for i := range core.TableReader.Spans {
-						if len(core.TableReader.Spans[i].EndKey) == 0 {
+					var prevRowPrefix []byte
+					for i, sp := range core.TableReader.Spans {
+						if len(sp.EndKey) == 0 {
+							// At the moment, the ColBatchDirectScan cannot
+							// handle Gets (it's not clear whether it is worth
+							// to handle them via the same path as for Scans and
+							// ReverseScans (which could have too large of an
+							// overhead) or by teaching the operator to also
+							// decode a single KV (similar to what regular
+							// ColBatchScan does)).
+							// TODO(yuzefovich, 23.1): explore supporting Gets
+							// somehow.
 							return false
 						}
+						l, err := keys.GetRowPrefixLength(sp.Key)
+						if err != nil {
+							// This should rarely happen since an error here
+							// indicates that we're dealing with a non-SQL key,
+							// so we'll be conservative and simply disable
+							// direct columnar scans. (One example where we can
+							// get an error if we had manually split the range.)
+							return false
+						}
+						curRowPrefix := sp.Key[:l]
+						if i > 0 && bytes.Equal(prevRowPrefix, curRowPrefix) {
+							// Two consecutive requests are part of the
+							// same SQL row in which case we cannot use direct
+							// columnar scans since we'd create a separate
+							// coldata.Batch for each.
+							return false
+						}
+						prevRowPrefix = curRowPrefix
 					}
 					fetchSpec := core.TableReader.FetchSpec
 					// Handling user-defined types requires type hydration which

pkg/ccl/logictestccl/tests/local-repeatable-read/generated_test.go

@@ -0,0 +1,26 @@
+statement ok
+SET direct_columnar_scans_enabled = true
+
+statement ok
+CREATE TABLE t145232 (
+      k INT PRIMARY KEY,
+      a INT NOT NULL,
+      b INT NOT NULL,
+      c INT NOT NULL,
+      v INT NOT NULL DEFAULT 5,
+      FAMILY (c),
+      FAMILY (v),
+      FAMILY (k),
+      FAMILY (a),
+      FAMILY (b)
+);
+
+statement ok
+INSERT INTO t145232 VALUES (2,2,2,2);
+
+# Regression test for #145232 where we tried to use direct columnar scans even
+# though a single SQL row is handled by multiple Scan requests.
+query IIIII
+SELECT * FROM t145232 WHERE k = 2
+----
+2  2  2  2  5

pkg/sql/colexec/colbuilder/BUILD.bazel

@@ -88,11 +88,11 @@ SELECT t1.x from t1 INNER LOOKUP JOIN t2 ON t1.x=t2.x WHERE t2.y='hello'
 ├ Node 2
 │ └ *colrpc.Outbox
 │   └ *rowexec.joinReader
-│     └ *colfetcher.ColBatchDirectScan
+│     └ *colfetcher.ColBatchScan
 └ Node 3
   └ *colrpc.Outbox
     └ *rowexec.joinReader
-      └ *colfetcher.ColBatchDirectScan
+      └ *colfetcher.ColBatchScan
 
 query I
 SELECT t1.x from t1 INNER LOOKUP JOIN t2 ON t1.x=t2.x WHERE t2.y='hello'