Merge #125236 #125272

125236: colexecdisk: fix incorrect context capture in partitionerToOperator r=yuzefovich a=yuzefovich This commit fixes recently exposed minor issue with how `partitionerToOperator` captures the context. In particular, these operators are reused multiple times by the external sort and the hash-based partitioner for each new partition, and previously these operator would capture the context that was provided on the first call to `Init`. In case of the external sort, that context comes from the ordered synchronizer which has a tracing span that is finished after each partition, so the captured context would have already finished span that can be problematic down the line. The fix in this commit is to make the operator resettable so that the new context is captured every time the operator is reused. Calling `Reset` is already in place in the hash-based partitioner (because other operators are reused there too), and this commit adds a call in the external sort. Fixes: #125177. Release note: None 125272: randgen: ignore virtual columns in foreign key mutator r=yuzefovich a=yuzefovich We currently don't support FK references to / from virtual columns, so we need to adjust the FK mutator after a recent change in fef5968. Fixes: #125145. Release note: None Co-authored-by: Yahor Yuzefovich <yahor@cockroachlabs.com>
cockroachdb · Jun 6, 2024 · 802554a · 802554a
3 parents 6bd7e35 + 62224d3 + cd4fc69
commit 802554a
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 1 deletion.
diff --git a/pkg/sql/colexec/colexecdisk/external_sort.go b/pkg/sql/colexec/colexecdisk/external_sort.go
@@ -668,6 +668,10 @@ func (s *externalSorter) Close(ctx context.Context) error {
 // to the last n current partitions to be merged.
 func (s *externalSorter) createPartitionerToOperators(n int) {
 	oldPartitioners := s.partitionerToOperators
+	for i := range oldPartitioners {
+		// Prepare the partitioner for reuse.
+		oldPartitioners[i].Reset(s.Ctx)
+	}
 	if len(oldPartitioners) < n {
 		s.partitionerToOperators = make([]*partitionerToOperator, n)
 		copy(s.partitionerToOperators, oldPartitioners)

diff --git a/pkg/sql/colexec/colexecdisk/utils.go b/pkg/sql/colexec/colexecdisk/utils.go
@@ -47,7 +47,16 @@ type partitionerToOperator struct {
 	batch        coldata.Batch
 }
 
-var _ colexecop.Operator = &partitionerToOperator{}
+var _ colexecop.ResettableOperator = &partitionerToOperator{}
+
+func (p *partitionerToOperator) Reset(ctx context.Context) {
+	// When resetting we simply want to update the context - this will allow us
+	// to make Init a no-op and reuse already allocated batch. If Init hasn't
+	// been called yet, then Reset is a no-op.
+	if p.Ctx != nil {
+		p.Ctx = ctx
+	}
+}
 
 func (p *partitionerToOperator) Init(ctx context.Context) {
 	if !p.InitHelper.Init(ctx) {

diff --git a/pkg/sql/randgen/mutator.go b/pkg/sql/randgen/mutator.go
@@ -463,6 +463,12 @@ func foreignKeyMutator(
 		for _, def := range table.Defs {
 			switch def := def.(type) {
 			case *tree.ColumnTableDef:
+				if def.Computed.Virtual {
+					// We currently don't support FK references to / from
+					// virtual columns.
+					// TODO(#59671): remove this skip.
+					continue
+				}
 				idx, ok := tableNameToColIdx(table.Table)
 				if !ok {
 					idx = len(cols)