Function NewMemBatchWithCapacity

pkg/col/coldata/batch.go:216–273

NewMemBatchWithCapacity allocates a new in-memory Batch with the given column size. Use for operators that have a precisely-sized output batch.

(typs []*types.T, capacity int, factory ColumnFactory)

// NewMemBatchWithCapacity allocates a new in-memory Batch with the given
// column size. Use for operators that have a precisely-sized output batch.
func NewMemBatchWithCapacity(typs []*types.T, capacity int, factory ColumnFactory) Batch {
	b := NewMemBatchNoCols(typs, capacity).(*MemBatch)
	vecs := make([]Vec, len(typs))
	// numForCanonicalType will track how many times a particular canonical type
	// family appears in typs.
	//
	//gcassert:noescape
	var numForCanonicalType [numCanonicalTypes]int
	for i, t := range typs {
		vecs[i].t = t
		vecs[i].canonicalTypeFamily = typeconv.TypeFamilyToCanonicalTypeFamily(t.Family())
		b.b[i] = &vecs[i]
		allocIdx := toAllocIdx(vecs[i].canonicalTypeFamily, vecs[i].t.Width())
		numForCanonicalType[allocIdx]++
	}
	var maxSame int
	for _, numColumns := range numForCanonicalType {
		maxSame = max(maxSame, numColumns)
	}
	// First, initialize columns that can be batch-allocated together for the
	// same canonical type.
	if maxSame > 1 {
		scratchColumns := make([]Column, maxSame)
		for allocIdx, numColumns := range numForCanonicalType {
			if numColumns > 1 {
				scratch := scratchColumns[:numColumns]
				t := toCanonicalType(allocIdx)
				factory.MakeColumns(scratch, t, capacity)
				for i := range vecs {
					if toAllocIdx(vecs[i].canonicalTypeFamily, vecs[i].t.Width()) == allocIdx {
						vecs[i].col = scratch[0]
						scratch = scratch[1:]
					}
				}
			}
		}
		if numForCanonicalType[datumAllocIdx] > 1 {
			// We need to set the correct type on each datum-backed vector.
			for i := range vecs {
				if vecs[i].canonicalTypeFamily == typeconv.DatumVecCanonicalTypeFamily {
					vecs[i].Datum().SetType(vecs[i].t)
				}
			}
		}
	}
	// Initialize all remaining columns while also initializing all Nulls
	// bitmaps.
	nullsCap := nullsStorageCap(capacity)
	nullsAlloc := make([]byte, len(typs)*nullsCap)
	for i := range vecs {
		if vecs[i].col == nil {
			vecs[i].col = factory.MakeColumn(vecs[i].t, capacity)
		}
		vecs[i].nulls = newNulls(nullsAlloc[:nullsCap:nullsCap])
		nullsAlloc = nullsAlloc[nullsCap:]
	}
	return b
}