Repository: pingcap/tidb
File: pkg/util/chunk/column.go (Go)

Code before:
func (c *Column) CalculateLenDeltaForAppendCellNTimesForFixedElem(src *Column, times int) (nullBitMapLenDelta int64, dataLenDelta int64) {
	nullBitMapLenDelta = c.calculateLenDeltaForAppendCellNTimesForNullBitMap(times)
	dataLenDelta = int64(len(src.elemBuf) * times)
	return
}

// CalculateLenDeltaForAppendCellNTimesForVarElem calculate the memory usage for `AppendCellNTimes` function
func (c *Column) CalculateLenDeltaForAppendCellNTimesForVarElem(src *Column, pos, times int) (nullBitMapLenDelta int64, dataLenDelta int64, offsetLenDelta int64) {
	nullBitMapLenDelta = c.calculateLenDeltaForAppendCellNTimesForNullBitMap(times)
	dataLenDelta = (src.offsets[pos+1] - src.offsets[pos]) * int64(times)
	offsetLenDelta = int64(times)
	return
}

// AppendCellNTimes append the pos-th Cell in source column to target column N times
func (c *Column) AppendCellNTimes(src *Column, pos, times int) {
	notNull := !src.IsNull(pos)
	if times == 1 {
		c.appendNullBitmap(notNull)
	} else {
		c.appendMultiSameNullBitmap(notNull, times)
	}
	if c.IsFixed() {
		elemLen := len(src.elemBuf)
		offset := pos * elemLen
		if intest.InTest {
			appendedByteNum := elemLen * times
			availableMem := cap(c.data) - len(c.data)
			if availableMem < appendedByteNum {
				panic("Can't pre-alloc enough memory")
			}
		}

		for range times {
			c.data = append(c.data, src.data[offset:offset+elemLen]...)
		}
	} else {
		start, end := src.offsets[pos], src.offsets[pos+1]

		if intest.InTest {
			appendedByteNum := int(end-start)*times
			availableMem := cap(c.data) - len(c.data)
			if availableMem < appendedByteNum {
				panic("Can't pre-alloc enough memory")
			}

			if cap(c.offsets)-len(c.offsets) < times {
				panic("Can't pre-alloc enough memory")
			}
		}
		

Review comment (suggestion):
I think we can check this once after all `AppendCellNTimes` is called in `appendProbeRowToChunkInternal`

Code after:

// CalculateLenDeltaForAppendCellNTimesForNullBitMap calculates the memory usage of nullBitmap for `AppendCellNTimes` function
func (c *Column) CalculateLenDeltaForAppendCellNTimesForNullBitMap(times int) int64 {
	return max(0, int64(((c.length+times+7)>>3)-len(c.nullBitmap)))
}

// CalculateLenDeltaForAppendCellNTimesForFixedElem calculates the memory usage of fixed elem for `AppendCellNTimes` function
func (*Column) CalculateLenDeltaForAppendCellNTimesForFixedElem(src *Column, times int) int64 {
	return int64(len(src.elemBuf) * times)
}

// CalculateLenDeltaForAppendCellNTimesForVarElem calculates the memory usage of variable length elem for `AppendCellNTimes` function
func (*Column) CalculateLenDeltaForAppendCellNTimesForVarElem(src *Column, pos, times int) int64 {
	return (src.offsets[pos+1] - src.offsets[pos]) * int64(times)
}

// AppendCellNTimes append the pos-th Cell in source column to target column N times
func (c *Column) AppendCellNTimes(src *Column, pos, times int) {
	notNull := !src.IsNull(pos)
	if times == 1 {
		c.appendNullBitmap(notNull)
	} else {
		c.appendMultiSameNullBitmap(notNull, times)
	}
	if c.IsFixed() {
		elemLen := len(src.elemBuf)
		offset := pos * elemLen

		for range times {
			c.data = append(c.data, src.data[offset:offset+elemLen]...)
		}
	} else {
		start, end := src.offsets[pos], src.offsets[pos+1]

		for range times {
			c.data = append(c.data, src.data[start:end]...)
			c.offsets = append(c.offsets, int64(len(c.data)))
		}
	}
	c.length += times
}

// appendMultiSameNullBitmap appends multiple same bit value to `nullBitMap`.
// notNull means not null.
// num means the number of bits that should be appended.
func (c *Column) appendMultiSameNullBitmap(notNull bool, num int) {
	numNewBytes := ((c.length + num + 7) >> 3) - len(c.nullBitmap)
	b := byte(0)
	if notNull {
		b = 0xff
	}