Merge #42845 #43409 #43887 #43980

42845: coldata: expose a modifier for batchSize variable r=yuzefovich a=yuzefovich

This commit exposes a setter for batchSize private variable to be
modified in tests and runs all tests in sql/colexec with a random
batch size in [3, 4096] range.

The tests in several places needed to be adjusted because their
assumptions might no longer hold when batch size is modified.

Fixes: #40791.

Release note: None

43409: execinfrapb: print input types one per line in EXPLAIN (DISTSQL, TYPES) r=yuzefovich a=yuzefovich

Previously, input types were printed in a single line within the input
synchronizer box on the flow diagram. However, when we have a processor
with two inputs and when each input has several columns, the boxes would
collide and make the types unreadable. This commit changes it to print
each type one per line, so there is no collision between different boxes
of input synchronizer. Separately, cockroachdb.github.io has been
adjusted so that the boxes for the input synchronizers could be
displayed well in such scenario.

Release note: None

43887: pkg/sql/opt: validate correct use of FOR UPDATE locking clauses r=nvanbenschoten a=nvanbenschoten

The second commit updates our SQL grammar to be identical to Postgres regarding select statement locking clauses. Concretely, this results in the following changes in behavior:
- We can now parse multiple locking clause items
- We can now parse the `FOR READ ONLY` syntax, which is a no-op
- We can now parse the `SELECT ... FOR UPDATE ... LIMIT` syntax

Ref: https://github.com/postgres/postgres/blob/1a4a0329650b0545a54afb3c317aa289fd817f8a/src/backend/parser/gram.y#L11834

The third commit introduces a number of validation checks into the optbuilder to ensure that FOR UPDATE locking clauses are only being used in places where they're allowed. Most of this was based off of Postgres, which disallows the same set of operations to be used with `FOR [KEY] UPDATE/SHARE` locking.

Ref: https://github.com/postgres/postgres/blob/1a4a0329650b0545a54afb3c317aa289fd817f8a/src/backend/parser/analyze.c#L2691

Release note (sql change): Invalid usages of FOR UPDATE locking clauses are now rejected by the SQL optimizer.

43980: storage: rename SSTSnapshotStorage{,Scratch,File} vars r=ajwerner a=irfansharif

Around snapshot receiving code, we have variables for SSTSnapshotStorage, SSTSnapshotStorageScratch, and SSTSnapshotStorageFile types
all abbreviated to some variation of sss{,s,f}. We also have "ssts" to
talk about SSTs. This is all terribly confusing, we could do with some
much needed clarity here.

Release note: None.

Co-authored-by: Yahor Yuzefovich <[email protected]>
Co-authored-by: Nathan VanBenschoten <[email protected]>
Co-authored-by: irfan sharif <[email protected]>

Author: 4 people

docs/generated/sql/bnf/select_stmt.bnf

@@ -1,3 +1,3 @@
 select_stmt ::=
-	( simple_select locking_clause | select_clause sort_clause locking_clause | select_clause ( sort_clause |  ) ( limit_clause offset_clause | offset_clause limit_clause | limit_clause | offset_clause ) locking_clause | ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) | 'WITH' 'RECURSIVE' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) select_clause locking_clause | ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) | 'WITH' 'RECURSIVE' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) select_clause sort_clause locking_clause | ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) | 'WITH' 'RECURSIVE' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) select_clause ( sort_clause |  ) ( limit_clause offset_clause | offset_clause limit_clause | limit_clause | offset_clause ) locking_clause )
+	( select_clause sort_clause | select_clause ( sort_clause |  ) for_locking_clause opt_select_limit | select_clause ( sort_clause |  ) ( limit_clause offset_clause | offset_clause limit_clause | limit_clause | offset_clause ) opt_for_locking_clause | ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) | 'WITH' 'RECURSIVE' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) select_clause | ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) | 'WITH' 'RECURSIVE' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) select_clause sort_clause | ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) | 'WITH' 'RECURSIVE' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) select_clause ( sort_clause |  ) for_locking_clause opt_select_limit | ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) | 'WITH' 'RECURSIVE' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) select_clause ( sort_clause |  ) ( limit_clause offset_clause | offset_clause limit_clause | limit_clause | offset_clause ) opt_for_locking_clause )
 

docs/generated/sql/bnf/stmt_block.bnf

@@ -439,12 +439,14 @@ scrub_database_stmt ::=
 	'EXPERIMENTAL' 'SCRUB' 'DATABASE' database_name opt_as_of_clause
 
 select_no_parens ::=
-	simple_select locking_clause
-	| select_clause sort_clause locking_clause
-	| select_clause opt_sort_clause select_limit locking_clause
-	| with_clause select_clause locking_clause
-	| with_clause select_clause sort_clause locking_clause
-	| with_clause select_clause opt_sort_clause select_limit locking_clause
+	simple_select
+	| select_clause sort_clause
+	| select_clause opt_sort_clause for_locking_clause opt_select_limit
+	| select_clause opt_sort_clause select_limit opt_for_locking_clause
+	| with_clause select_clause
+	| with_clause select_clause sort_clause
+	| with_clause select_clause opt_sort_clause for_locking_clause opt_select_limit
+	| with_clause select_clause opt_sort_clause select_limit opt_for_locking_clause
 
 select_with_parens ::=
 	'(' select_no_parens ')'
@@ -1172,19 +1174,28 @@ simple_select ::=
 	| table_clause
 	| set_operation
 
-locking_clause ::=
-	for_locking_strength opt_locked_rels opt_nowait_or_skip
-
 select_clause ::=
 	simple_select
 	| select_with_parens
 
+for_locking_clause ::=
+	for_locking_items
+	| 'FOR' 'READ' 'ONLY'
+
+opt_select_limit ::=
+	select_limit
+	| 
+
 select_limit ::=
 	limit_clause offset_clause
 	| offset_clause limit_clause
 	| limit_clause
 	| offset_clause
 
+opt_for_locking_clause ::=
+	for_locking_clause
+	| 
+
 set_rest_more ::=
 	generic_set
 
@@ -1578,18 +1589,8 @@ set_operation ::=
 	| select_clause 'INTERSECT' all_or_distinct select_clause
 	| select_clause 'EXCEPT' all_or_distinct select_clause
 
-for_locking_strength ::=
-	'FOR' 'UPDATE'
-	| 'FOR' 'NO' 'KEY' 'UPDATE'
-	| 'FOR' 'SHARE'
-	| 'FOR' 'KEY' 'SHARE'
-
-opt_locked_rels ::=
-	'OF' table_name_list
-
-opt_nowait_or_skip ::=
-	'SKIP' 'LOCKED'
-	| 'NOWAIT'
+for_locking_items ::=
+	( for_locking_item ) ( ( for_locking_item ) )*
 
 offset_clause ::=
 	'OFFSET' a_expr
@@ -1943,6 +1944,9 @@ all_or_distinct ::=
 	| 'DISTINCT'
 	| 
 
+for_locking_item ::=
+	for_locking_strength opt_locked_rels opt_nowait_or_skip
+
 var_list ::=
 	( var_value ) ( ( ',' var_value ) )*
 
@@ -2208,6 +2212,19 @@ interval_qualifier ::=
 window_definition_list ::=
 	( window_definition ) ( ( ',' window_definition ) )*
 
+for_locking_strength ::=
+	'FOR' 'UPDATE'
+	| 'FOR' 'NO' 'KEY' 'UPDATE'
+	| 'FOR' 'SHARE'
+	| 'FOR' 'KEY' 'SHARE'
+
+opt_locked_rels ::=
+	'OF' table_name_list
+
+opt_nowait_or_skip ::=
+	'SKIP' 'LOCKED'
+	| 'NOWAIT'
+
 opt_join_hint ::=
 	'HASH'
 	| 'MERGE'

pkg/col/coldata/batch.go

@@ -58,16 +58,39 @@ type Batch interface {
 
 var _ Batch = &MemBatch{}
 
-const maxBatchSize = 1024
+const (
+	// MinBatchSize is the minimum acceptable size of batches.
+	MinBatchSize = 3
+	// MaxBatchSize is the maximum acceptable size of batches.
+	MaxBatchSize = 4096
+)
 
+// TODO(jordan): tune.
 var batchSize = uint16(1024)
 
 // BatchSize is the maximum number of tuples that fit in a column batch.
-// TODO(jordan): tune
 func BatchSize() uint16 {
 	return batchSize
 }
 
+// SetBatchSizeForTests modifies batchSize variable. It should only be used in
+// tests.
+func SetBatchSizeForTests(newBatchSize uint16) {
+	if newBatchSize > MaxBatchSize {
+		panic(
+			fmt.Sprintf("requested batch size %d is greater than MaxBatchSize %d",
+				newBatchSize, MaxBatchSize),
+		)
+	}
+	if newBatchSize < MinBatchSize {
+		panic(
+			fmt.Sprintf("requested batch size %d is smaller than MinBatchSize %d",
+				newBatchSize, MinBatchSize),
+		)
+	}
+	batchSize = newBatchSize
+}
+
 // NewMemBatch allocates a new in-memory Batch. A coltypes.Unknown type
 // will create a placeholder Vec that may not be accessed.
 // TODO(jordan): pool these allocations.

pkg/col/coldata/nulls.go

@@ -10,13 +10,13 @@
 
 package coldata
 
-// zeroedNulls is a zeroed out slice representing a bitmap of size maxBatchSize.
+// zeroedNulls is a zeroed out slice representing a bitmap of size MaxBatchSize.
 // This is copied to efficiently set all nulls.
-var zeroedNulls [(maxBatchSize-1)/8 + 1]byte
+var zeroedNulls [(MaxBatchSize-1)/8 + 1]byte
 
-// filledNulls is a slice representing a bitmap of size maxBatchSize with every
+// filledNulls is a slice representing a bitmap of size MaxBatchSize with every
 // single bit set.
-var filledNulls [(maxBatchSize-1)/8 + 1]byte
+var filledNulls [(MaxBatchSize-1)/8 + 1]byte
 
 // bitMask[i] is a byte with a single bit set at i.
 var bitMask = [8]byte{0x1, 0x2, 0x4, 0x8, 0x10, 0x20, 0x40, 0x80}

pkg/sql/colexec/aggregator_test.go

@@ -222,7 +222,7 @@ func TestAggregatorOneFunc(t *testing.T) {
 				{7},
 				{8},
 			},
-			batchSize:       4,
+			batchSize:       3,
 			outputBatchSize: 1,
 			name:            "CarryBetweenInputAndOutputBatches",
 		},
@@ -476,6 +476,11 @@ func TestAggregatorRandom(t *testing.T) {
 	rng, _ := randutil.NewPseudoRand()
 	ctx := context.Background()
 	for _, groupSize := range []int{1, 2, int(coldata.BatchSize()) / 4, int(coldata.BatchSize()) / 2} {
+		if groupSize == 0 {
+			// We might be varying coldata.BatchSize() so that when it is divided by
+			// 4, groupSize is 0. We want to skip such configuration.
+			continue
+		}
 		for _, numInputBatches := range []int{1, 2, 64} {
 			for _, hasNulls := range []bool{true, false} {
 				for _, agg := range aggTypes {
@@ -495,7 +500,7 @@ func TestAggregatorRandom(t *testing.T) {
 							curGroup := -1
 							for i := range groups {
 								if i%groupSize == 0 {
-									expRowCounts = append(expRowCounts, int64(groupSize))
+									expRowCounts = append(expRowCounts, 0)
 									expCounts = append(expCounts, 0)
 									expSums = append(expSums, 0)
 									expMins = append(expMins, 2048)
@@ -508,8 +513,11 @@ func TestAggregatorRandom(t *testing.T) {
 								// slower.
 								aggCol[i] = 2048 * (rng.Float64() - 0.5)
 
+								// NULL values contribute to the row count, so we're updating
+								// the row counts outside of the if block.
+								expRowCounts[curGroup]++
 								if hasNulls && rng.Float64() < nullProbability {
-									aggColNulls.SetNull(uint16(i))
+									aggColNulls.SetNull64(uint64(i))
 								} else {
 									expNulls[curGroup] = false
 									expCounts[curGroup]++
@@ -704,12 +712,7 @@ func BenchmarkAggregator(b *testing.B) {
 											a.(resetter).reset()
 											source.reset()
 											// Exhaust aggregator until all batches have been read.
-											foundTuples := 0
 											for b := a.Next(ctx); b.Length() != 0; b = a.Next(ctx) {
-												foundTuples += int(b.Length())
-											}
-											if foundTuples != nTuples/groupSize {
-												b.Fatalf("Found %d tuples, expected %d", foundTuples, nTuples/groupSize)
 											}
 										}
 									},

pkg/sql/colexec/hashjoiner_test.go

@@ -912,6 +912,12 @@ func TestHashJoiner(t *testing.T) {
 	}
 
 	for _, outputBatchSize := range []uint16{1, 17, coldata.BatchSize()} {
+		if outputBatchSize > coldata.BatchSize() {
+			// It is possible for varied coldata.BatchSize() to be smaller than
+			// requested outputBatchSize. Such configuration is invalid, and we skip
+			// it.
+			continue
+		}
 		for _, tc := range tcs {
 			inputs := []tuples{tc.leftTuples, tc.rightTuples}
 			typs := [][]coltypes.T{tc.leftTypes, tc.rightTypes}

pkg/sql/colexec/main_test.go

@@ -12,9 +12,11 @@ package colexec
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"testing"
 
+	"github.com/cockroachdb/cockroach/pkg/col/coldata"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
 	"github.com/cockroachdb/cockroach/pkg/util/mon"
@@ -43,6 +45,13 @@ func TestMain(m *testing.M) {
 		testMemAcc = &memAcc
 		testAllocator = NewAllocator(ctx, testMemAcc)
 		defer testMemAcc.Close(ctx)
+		rng, _ := randutil.NewPseudoRand()
+		// Pick a random batch size in [coldata.MinBatchSize, coldata.MaxBatchSize]
+		// range.
+		randomBatchSize := uint16(coldata.MinBatchSize +
+			rng.Intn(coldata.MaxBatchSize-coldata.MinBatchSize))
+		fmt.Printf("coldata.BatchSize() is set to %d\n", randomBatchSize)
+		coldata.SetBatchSizeForTests(randomBatchSize)
 		return m.Run()
 	}())
 }

pkg/sql/colexec/mergejoiner_test.go

@@ -1814,10 +1814,25 @@ func TestMergeJoinerMultiBatchRuns(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	ctx := context.Background()
 	for _, groupSize := range []int{int(coldata.BatchSize()) / 8, int(coldata.BatchSize()) / 4, int(coldata.BatchSize()) / 2} {
+		if groupSize == 0 {
+			// We might be varying coldata.BatchSize() so that when it is divided by
+			// 4, groupSize is 0. We want to skip such configuration.
+			continue
+		}
 		for _, numInputBatches := range []int{1, 2, 16} {
 			t.Run(fmt.Sprintf("groupSize=%d/numInputBatches=%d", groupSize, numInputBatches),
 				func(t *testing.T) {
 					nTuples := int(coldata.BatchSize()) * numInputBatches
+					// There will be nTuples/groupSize "full" groups - i.e. groups of
+					// groupSize. Each of these "full" groups will produce groupSize^2
+					// tuples. The last group might be not full and will consist of
+					// nTuples % groupSize tuples. That group will produce
+					// lastGroupSize^2 tuples.
+					// Note that the math will still be correct in case when nTuples is
+					// divisible by groupSize - all the groups will be full and "last"
+					// group will be of size 0.
+					lastGroupSize := nTuples % groupSize
+					expCount := nTuples/groupSize*(groupSize*groupSize) + lastGroupSize*lastGroupSize
 					typs := []coltypes.T{coltypes.Int64, coltypes.Int64}
 					cols := []coldata.Vec{
 						testAllocator.NewMemColumn(typs[0], nTuples),
@@ -1870,9 +1885,9 @@ func TestMergeJoinerMultiBatchRuns(t *testing.T) {
 						i++
 					}
 
-					if count != groupSize*int(coldata.BatchSize())*numInputBatches {
+					if count != expCount {
 						t.Fatalf("found count %d, expected count %d",
-							count, groupSize*int(coldata.BatchSize())*numInputBatches)
+							count, expCount)
 					}
 				})
 		}

pkg/sql/colexec/orderedsynchronizer_test.go

@@ -157,6 +157,9 @@ func TestOrderedSyncRandomInput(t *testing.T) {
 	numInputs := 3
 	inputLen := 1024
 	batchSize := uint16(16)
+	if batchSize > coldata.BatchSize() {
+		batchSize = coldata.BatchSize()
+	}
 
 	// Generate a random slice of sorted ints.
 	randInts := make([]int, inputLen)

pkg/sql/colexec/routers.go

@@ -35,9 +35,14 @@ type routerOutput interface {
 	cancel()
 }
 
-// defaultRouterOutputBlockedThreshold is the number of unread values buffered
-// by the routerOutputOp after which the output is considered blocked.
-var defaultRouterOutputBlockedThreshold = int(coldata.BatchSize() * 2)
+// getDefaultRouterOutputBlockedThreshold returns the number of unread values
+// buffered by the routerOutputOp after which the output is considered blocked.
+// It is a function rather than a variable so that in tests we could modify
+// coldata.BatchSize() (if it were a variable, then its value would be
+// evaluated before we set the desired batch size).
+func getDefaultRouterOutputBlockedThreshold() int {
+	return int(coldata.BatchSize()) * 2
+}
 
 type routerOutputOp struct {
 	// input is a reference to our router.
@@ -90,7 +95,7 @@ func newRouterOutputOp(
 	allocator *Allocator, types []coltypes.T, unblockedEventsChan chan<- struct{},
 ) *routerOutputOp {
 	return newRouterOutputOpWithBlockedThresholdAndBatchSize(
-		allocator, types, unblockedEventsChan, defaultRouterOutputBlockedThreshold, int(coldata.BatchSize()),
+		allocator, types, unblockedEventsChan, getDefaultRouterOutputBlockedThreshold(), int(coldata.BatchSize()),
 	)
 }