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

Code before:
	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(moreBytesNumNullBitmapNeed))
	}

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

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

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

Review comment (suggestion):
I think you don't need to manually check `int64(cap(c.data)) < dataExpectedCap`, just call `c.data = slices.Grow(c.data, int(moreBytesNumDataNeed))` is enough.

Code after:
	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) {
	c.nullBitmap = slices.Grow(c.nullBitmap, int(moreBytesNumNullBitmapNeed))
	c.data = slices.Grow(c.data, int(moreBytesNumDataNeed))
	c.offsets = slices.Grow(c.offsets, int(moreBytesNumOffsetNeed))
}

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