colserde: fix the edge case with nulls handling

When serializing the data of Bool, Bytes, Int, and Float types when they
don't have any nulls in the vector, we don't explicit specify the null
bitmap. Previously, when deserializing such vectors with no nulls we
would simply call `UnsetNulls` on the `coldata.Nulls` object that is
currently present. However, it is possible that already present nulls
object cannot support the desired batch length. This could lead to index
out of bounds accesses. Note that in the vast majority of cases this
likely doesn't happen in practice because we check `MaybeHasNulls`, and
that would return `false` making us omit the null checking code.

Release note (bug fix): Previously, CockroachDB could encounter an
internal error in rare circumstances when executing queries via the
vectorized engine that operate on columns of BOOL, BYTES, INT, and FLOAT
types that have a mix of NULL and non-NULL values.

pkg/col/coldata/nulls.go

@@ -311,6 +311,12 @@ func (n *Nulls) Slice(start int, end int) Nulls {
 	return s
 }
 
+// MaxNumElements returns the maximum number of elements that this Nulls can
+// accommodate.
+func (n *Nulls) MaxNumElements() int {
+	return len(n.nulls) * 8
+}
+
 // NullBitmap returns the null bitmap.
 func (n *Nulls) NullBitmap() []byte {
 	return n.nulls

pkg/col/coldata/vec_tmpl.go

@@ -267,6 +267,9 @@ func SetValueAt(v Vec, elem interface{}, rowIdx int) {
 // GetValueAt is an inefficient helper to get the value in a Vec when the type
 // is unknown.
 func GetValueAt(v Vec, rowIdx int) interface{} {
+	if v.Nulls().NullAt(rowIdx) {
+		return nil
+	}
 	t := v.Type()
 	switch v.CanonicalTypeFamily() {
 	// {{range .}}

pkg/col/coldatatestutils/random_testutils.go

@@ -380,18 +380,16 @@ func (o *RandomDataOp) Next(context.Context) coldata.Batch {
 		selProbability  float64
 		nullProbability float64
 	)
+	if o.selection {
+		selProbability = o.rng.Float64()
+	}
+	if o.nulls && o.rng.Float64() > 0.1 {
+		// Even if nulls are desired, in 10% of cases create a batch with no
+		// nulls at all.
+		nullProbability = o.rng.Float64()
+	}
 	for {
-		if o.selection {
-			selProbability = o.rng.Float64()
-		}
-		if o.nulls {
-			nullProbability = o.rng.Float64()
-		}
-
 		b := RandomBatchWithSel(o.allocator, o.rng, o.typs, o.batchSize, nullProbability, selProbability)
-		if !o.selection {
-			b.SetSelection(false)
-		}
 		if b.Length() == 0 {
 			// Don't return a zero-length batch until we return o.numBatches batches.
 			continue

pkg/col/colserde/arrowbatchconverter.go

@@ -483,6 +483,13 @@ func handleNulls(arr array.Interface, vec coldata.Vec, batchLength int) {
 		// For types with the canonical type family of Bool, Bytes, Int, or
 		// Float, when there are no nulls, we have a null bitmap with zero
 		// length.
-		vec.Nulls().UnsetNulls()
+		if vec.Nulls().MaxNumElements() < batchLength {
+			// The current null bitmap doesn't have enough space, so we need to
+			// allocate a new one.
+			nulls := coldata.NewNulls(batchLength)
+			vec.SetNulls(&nulls)
+		} else {
+			vec.Nulls().UnsetNulls()
+		}
 	}
 }

pkg/col/colserde/arrowbatchconverter_test.go

@@ -61,35 +61,27 @@ func TestArrowBatchConverterRandom(t *testing.T) {
 	coldata.AssertEquivalentBatches(t, expected, actual)
 }
 
-// roundTripBatch is a helper function that round trips a batch through the
-// ArrowBatchConverter and RecordBatchSerializer. Make sure to copy the input
-// batch before passing it to this function to assert equality.
+// roundTripBatch is a helper function that pushes the source batch through the
+// ArrowBatchConverter and RecordBatchSerializer. The result is written to dest.
 func roundTripBatch(
-	b coldata.Batch,
-	c *colserde.ArrowBatchConverter,
-	r *colserde.RecordBatchSerializer,
-	typs []*types.T,
-) (coldata.Batch, error) {
+	src, dest coldata.Batch, c *colserde.ArrowBatchConverter, r *colserde.RecordBatchSerializer,
+) error {
 	var buf bytes.Buffer
-	arrowDataIn, err := c.BatchToArrow(b)
+	arrowDataIn, err := c.BatchToArrow(src)
 	if err != nil {
-		return nil, err
+		return err
 	}
-	_, _, err = r.Serialize(&buf, arrowDataIn, b.Length())
+	_, _, err = r.Serialize(&buf, arrowDataIn, src.Length())
 	if err != nil {
-		return nil, err
+		return err
 	}
 
 	var arrowDataOut []*array.Data
 	batchLength, err := r.Deserialize(&arrowDataOut, buf.Bytes())
 	if err != nil {
-		return nil, err
+		return err
 	}
-	actual := testAllocator.NewMemBatchWithFixedCapacity(typs, batchLength)
-	if err := c.ArrowToBatch(arrowDataOut, batchLength, actual); err != nil {
-		return nil, err
-	}
-	return actual, nil
+	return c.ArrowToBatch(arrowDataOut, batchLength, dest)
 }
 
 func TestRecordBatchRoundtripThroughBytes(t *testing.T) {
@@ -98,26 +90,45 @@ func TestRecordBatchRoundtripThroughBytes(t *testing.T) {
 	rng, _ := randutil.NewPseudoRand()
 	for run := 0; run < 10; run++ {
 		var typs []*types.T
-		var b coldata.Batch
+		var src coldata.Batch
 		if rng.Float64() < 0.1 {
 			// In 10% of cases we'll use a zero length schema.
-			b = testAllocator.NewMemBatchWithFixedCapacity(typs, rng.Intn(coldata.BatchSize())+1)
-			b.SetLength(b.Capacity())
+			src = testAllocator.NewMemBatchWithFixedCapacity(typs, rng.Intn(coldata.BatchSize())+1)
+			src.SetLength(src.Capacity())
 		} else {
-			typs, b = randomBatch(testAllocator)
+			typs, src = randomBatch(testAllocator)
 		}
+		dest := testAllocator.NewMemBatchWithMaxCapacity(typs)
 		c, err := colserde.NewArrowBatchConverter(typs)
 		require.NoError(t, err)
 		r, err := colserde.NewRecordBatchSerializer(typs)
 		require.NoError(t, err)
 
-		// Make a copy of the original batch because the converter modifies and
-		// casts data without copying for performance reasons.
-		expected := coldatatestutils.CopyBatch(b, typs, testColumnFactory)
-		actual, err := roundTripBatch(b, c, r, typs)
-		require.NoError(t, err)
+		// Reuse the same destination batch as well as the ArrowBatchConverter
+		// and RecordBatchSerializer in order to simulate how these things are
+		// used in the production setting.
+		for i := 0; i < 10; i++ {
+			require.NoError(t, roundTripBatch(src, dest, c, r))
+
+			coldata.AssertEquivalentBatches(t, src, dest)
+			// Check that we can actually read each tuple from the destination
+			// batch.
+			for _, vec := range dest.ColVecs() {
+				for tupleIdx := 0; tupleIdx < dest.Length(); tupleIdx++ {
+					coldata.GetValueAt(vec, tupleIdx)
+				}
+			}
 
-		coldata.AssertEquivalentBatches(t, expected, actual)
+			// Generate the new source batch.
+			nullProbability := rng.Float64()
+			if rng.Float64() < 0.1 {
+				// In some cases, make sure that there are no nulls at all.
+				nullProbability = 0
+			}
+			capacity := rng.Intn(coldata.BatchSize()) + 1
+			length := rng.Intn(capacity)
+			src = coldatatestutils.RandomBatch(testAllocator, rng, typs, capacity, length, nullProbability)
+		}
 	}
 }
 

pkg/col/colserde/record_batch_test.go

@@ -302,8 +302,9 @@ func TestRecordBatchSerializerDeserializeMemoryEstimate(t *testing.T) {
 	rng, _ := randutil.NewPseudoRand()
 
 	typs := []*types.T{types.Bytes}
-	b := testAllocator.NewMemBatchWithFixedCapacity(typs, coldata.BatchSize())
-	bytesVec := b.ColVec(0).Bytes()
+	src := testAllocator.NewMemBatchWithMaxCapacity(typs)
+	dest := testAllocator.NewMemBatchWithMaxCapacity(typs)
+	bytesVec := src.ColVec(0).Bytes()
 	maxValueLen := coldata.BytesInitialAllocationFactor * 8
 	value := make([]byte, maxValueLen)
 	for i := 0; i < coldata.BatchSize(); i++ {
@@ -312,17 +313,16 @@ func TestRecordBatchSerializerDeserializeMemoryEstimate(t *testing.T) {
 		require.NoError(t, err)
 		bytesVec.Set(i, value)
 	}
-	b.SetLength(coldata.BatchSize())
-
-	originalMemorySize := colmem.GetBatchMemSize(b)
+	src.SetLength(coldata.BatchSize())
 
 	c, err := colserde.NewArrowBatchConverter(typs)
 	require.NoError(t, err)
 	r, err := colserde.NewRecordBatchSerializer(typs)
 	require.NoError(t, err)
-	b, err = roundTripBatch(b, c, r, typs)
-	require.NoError(t, err)
-	newMemorySize := colmem.GetBatchMemSize(b)
+	require.NoError(t, roundTripBatch(src, dest, c, r))
+
+	originalMemorySize := colmem.GetBatchMemSize(src)
+	newMemorySize := colmem.GetBatchMemSize(dest)
 
 	// We expect that the original and the new memory sizes are relatively close
 	// to each other (do not differ by more than a third). We cannot guarantee

pkg/sql/colcontainer/diskqueue_test.go

@@ -104,20 +104,22 @@ func TestDiskQueue(t *testing.T) {
 					require.NoError(t, err)
 					require.Equal(t, 1, len(directories))
 
-					// Run verification.
+					// Run verification. We reuse the same batch to dequeue into
+					// since that is the common pattern.
+					dest := coldata.NewMemBatch(typs, testColumnFactory)
 					for {
-						b := op.Next(ctx)
-						require.NoError(t, q.Enqueue(ctx, b))
-						if b.Length() == 0 {
+						src := op.Next(ctx)
+						require.NoError(t, q.Enqueue(ctx, src))
+						if src.Length() == 0 {
 							break
 						}
 						if rng.Float64() < dequeuedProbabilityBeforeAllEnqueuesAreDone {
-							if ok, err := q.Dequeue(ctx, b); !ok {
+							if ok, err := q.Dequeue(ctx, dest); !ok {
 								t.Fatal("queue incorrectly considered empty")
 							} else if err != nil {
 								t.Fatal(err)
 							}
-							coldata.AssertEquivalentBatches(t, batches[0], b)
+							coldata.AssertEquivalentBatches(t, batches[0], dest)
 							batches = batches[1:]
 						}
 					}
@@ -127,25 +129,24 @@ func TestDiskQueue(t *testing.T) {
 					}
 					for i := 0; i < numReadIterations; i++ {
 						batchIdx := 0
-						b := coldata.NewMemBatch(typs, testColumnFactory)
 						for batchIdx < len(batches) {
-							if ok, err := q.Dequeue(ctx, b); !ok {
+							if ok, err := q.Dequeue(ctx, dest); !ok {
 								t.Fatal("queue incorrectly considered empty")
 							} else if err != nil {
 								t.Fatal(err)
 							}
-							coldata.AssertEquivalentBatches(t, batches[batchIdx], b)
+							coldata.AssertEquivalentBatches(t, batches[batchIdx], dest)
 							batchIdx++
 						}
 
 						if testReuseCache {
 							// Trying to Enqueue after a Dequeue should return an error in these
 							// CacheModes.
-							require.Error(t, q.Enqueue(ctx, b))
+							require.Error(t, q.Enqueue(ctx, dest))
 						}
 
-						if ok, err := q.Dequeue(ctx, b); ok {
-							if b.Length() != 0 {
+						if ok, err := q.Dequeue(ctx, dest); ok {
+							if dest.Length() != 0 {
 								t.Fatal("queue should be empty")
 							}
 						} else if err != nil {