Merge pull request #1068 from mauidude/reduce-allocations

perf(col_str): Reduce heap allocations

Author: SpencerTorres

proto/buffer.go

@@ -19,7 +19,11 @@ func (b *Buffer) Reader() *Reader {
 
 // Ensure Buf length.
 func (b *Buffer) Ensure(n int) {
-	b.Buf = append(b.Buf[:0], make([]byte, n)...)
+	if cap(b.Buf) < n {
+		b.Buf = make([]byte, n)
+	} else {
+		b.Buf = b.Buf[:n] // Set length to n (zeros not guaranteed)
+	}
 }
 
 // Encoder implements encoding to Buffer.
@@ -67,8 +71,8 @@ func (b *Buffer) PutRaw(v []byte) {
 
 // PutUVarInt encodes x as uvarint.
 func (b *Buffer) PutUVarInt(x uint64) {
-	buf := make([]byte, binary.MaxVarintLen64)
-	n := binary.PutUvarint(buf, x)
+	buf := [binary.MaxVarintLen64]byte{}
+	n := binary.PutUvarint(buf[:], x)
 	b.Buf = append(b.Buf, buf[:n]...)
 }
 
@@ -98,27 +102,27 @@ func (b *Buffer) PutUInt8(x uint8) {
 }
 
 func (b *Buffer) PutUInt16(x uint16) {
-	buf := make([]byte, 16/8)
-	binary.LittleEndian.PutUint16(buf, x)
-	b.Buf = append(b.Buf, buf...)
+	buf := [16 / 8]byte{}
+	binary.LittleEndian.PutUint16(buf[:], x)
+	b.Buf = append(b.Buf, buf[:]...)
 }
 
 func (b *Buffer) PutUInt32(x uint32) {
-	buf := make([]byte, 32/8)
-	binary.LittleEndian.PutUint32(buf, x)
-	b.Buf = append(b.Buf, buf...)
+	buf := [32 / 8]byte{}
+	binary.LittleEndian.PutUint32(buf[:], x)
+	b.Buf = append(b.Buf, buf[:]...)
 }
 
 func (b *Buffer) PutUInt64(x uint64) {
-	buf := make([]byte, 64/8)
-	binary.LittleEndian.PutUint64(buf, x)
-	b.Buf = append(b.Buf, buf...)
+	buf := [64 / 8]byte{}
+	binary.LittleEndian.PutUint64(buf[:], x)
+	b.Buf = append(b.Buf, buf[:]...)
 }
 
 func (b *Buffer) PutUInt128(x UInt128) {
-	buf := make([]byte, 128/8)
-	binPutUInt128(buf, x)
-	b.Buf = append(b.Buf, buf...)
+	buf := [128 / 8]byte{}
+	binPutUInt128(buf[:], x)
+	b.Buf = append(b.Buf, buf[:]...)
 }
 
 func (b *Buffer) PutInt8(v int8) {

proto/col_str.go

@@ -68,22 +68,27 @@ func (c *ColStr) Reset() {
 
 // EncodeColumn encodes String rows to *Buffer.
 func (c ColStr) EncodeColumn(b *Buffer) {
-	buf := make([]byte, binary.MaxVarintLen64)
+	var buf [binary.MaxVarintLen64]byte
 	for _, p := range c.Pos {
-		n := binary.PutUvarint(buf, uint64(p.End-p.Start))
+		length := uint64(p.End - p.Start)
+
+		// Encode to temp buffer first
+		n := binary.PutUvarint(buf[:], length)
+
+		// Append only the bytes we need
 		b.Buf = append(b.Buf, buf[:n]...)
 		b.Buf = append(b.Buf, c.Buf[p.Start:p.End]...)
 	}
 }
 
 // WriteColumn writes String rows to *Writer.
 func (c ColStr) WriteColumn(w *Writer) {
-	buf := make([]byte, binary.MaxVarintLen64)
+	var buf [binary.MaxVarintLen64]byte
 	// Writing values from c.Buf directly might improve performance if [ColStr] contains a few rows of very long strings.
 	// However, most of the time it is quite opposite, so we copy data.
 	w.ChainBuffer(func(b *Buffer) {
 		for _, p := range c.Pos {
-			n := binary.PutUvarint(buf, uint64(p.End-p.Start))
+			n := binary.PutUvarint(buf[:], uint64(p.End-p.Start))
 			b.PutRaw(buf[:n])
 			b.PutRaw(c.Buf[p.Start:p.End])
 		}

proto/reader.go

@@ -135,12 +135,17 @@ func (r *Reader) StrBytes() ([]byte, error) {
 
 // Str decodes string.
 func (r *Reader) Str() (string, error) {
-	s, err := r.StrBytes()
+	defer r.b.Reset()
+
+	// call StrRaw instead of StrBytes and converting
+	// so we can avoid a second allocation and copy of
+	// the str/[]byte data
+	str, err := r.StrRaw()
 	if err != nil {
-		return "", errors.Wrap(err, "bytes")
+		return "", errors.Wrap(err, "raw")
 	}
 
-	return string(s), err
+	return string(str), err
 }
 
 // Int decodes uvarint as int.