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

Code before:
}

// Reserve allocates some memory for the column
func (c *Column) Reserve(nullBitmapExpectedCap int64, dataExpectedCap int64, offsetExpectedCap int64) {
	currentNullBitmapLen := int64(len(c.nullBitmap))
	nullBitmapExpectedCap += currentNullBitmapLen
	if int64(cap(c.nullBitmap)) < nullBitmapExpectedCap {
		c.nullBitmap = slices.Grow(c.nullBitmap, int(nullBitmapExpectedCap))
	}

	currentDataLen := int64(len(c.data))
	dataExpectedCap += currentDataLen
	if int64(cap(c.data)) < dataExpectedCap {
		c.data = slices.Grow(c.data, int(dataExpectedCap))
	}

	currentOffsetLen := int64(len(c.offsets))
	offsetExpectedCap += currentOffsetLen
	if int64(cap(c.offsets)) < offsetExpectedCap {
		c.offsets = slices.Grow(c.offsets, int(offsetExpectedCap))
	}
}

func (c *Column) calculateLenDeltaForAppendCellNTimesForNullBitMap(times int) int64 {
	nullBitMapLenDelta := int64(0)
	if times == 1 {
		if c.length>>3 >= len(c.nullBitmap) {
			nullBitMapLenDelta++
		}
	} else {
		nullBitMapLenDelta += int64(((c.length + times + 7) >> 3) - len(c.nullBitmap))
	}
	return nullBitMapLenDelta
}

// CalculateLenDeltaForAppendCellNTimesForFixedElem calculate the memory usage for `AppendCellNTimes` function
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

Review comment (suggestion):
if `calculateLenDeltaForAppendCellNTimesForNullBitMap` is calculated once for a chunk, i think we don't need to check `times == 1` here

Code after:
// AppendNullBitmap append a null/notnull value to the column's null map
func (c *Column) AppendNullBitmap(notNull bool) {
	c.appendNullBitmap(notNull)
}

func (c *Column) appendNullBitmap(notNull bool) {
	idx := c.length >> 3
	if idx >= len(c.nullBitmap) {
		c.nullBitmap = append(c.nullBitmap, 0)
	}
	if notNull {
		pos := uint(c.length) & 7
		c.nullBitmap[idx] |= byte(1 << pos)
	}
}

// Reserve allocates some memory for the column
func (c *Column) Reserve(moreBytesNumNullBitmapNeed int64, moreBytesNumDataNeed int64, moreBytesNumOffsetNeed int64) {
	currentNullBitmapLen := int64(len(c.nullBitmap))
	nullBitmapExpectedCap := currentNullBitmapLen + moreBytesNumNullBitmapNeed
	if int64(cap(c.nullBitmap)) < nullBitmapExpectedCap {
		c.nullBitmap = slices.Grow(c.nullBitmap, int(nullBitmapExpectedCap))
	}

	currentDataLen := int64(len(c.data))
	dataExpectedCap := currentDataLen + moreBytesNumDataNeed
	if int64(cap(c.data)) < dataExpectedCap {
		c.data = slices.Grow(c.data, int(dataExpectedCap))
	}

	currentOffsetLen := int64(len(c.offsets))
	offsetExpectedCap := currentOffsetLen + moreBytesNumOffsetNeed
	if int64(cap(c.offsets)) < offsetExpectedCap {
		c.offsets = slices.Grow(c.offsets, int(offsetExpectedCap))
	}
}

// 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)
}