Small aggregation perf improvements (#1781)

* skip test on old

* Cheaper error and collation checks

* fixup test

Author: max-hoffman

sql/cache.go

@@ -21,7 +21,6 @@ import (
 	"github.com/cespare/xxhash"
 
 	lru "github.com/hashicorp/golang-lru"
-	errors "gopkg.in/src-d/go-errors.v1"
 )
 
 // HashOf returns a hash of the given value to be used as key in a cache.
@@ -37,7 +36,7 @@ func HashOf(v Row) (uint64, error) {
 }
 
 // ErrKeyNotFound is returned when the key could not be found in the cache.
-var ErrKeyNotFound = errors.NewKind("memory: key %d not found in cache")
+var ErrKeyNotFound = fmt.Errorf("memory: key not found in cache")
 
 type lruCache struct {
 	memory   Freeable
@@ -65,7 +64,7 @@ func (l *lruCache) Put(k uint64, v interface{}) error {
 func (l *lruCache) Get(k uint64) (interface{}, error) {
 	v, ok := l.cache.Get(k)
 	if !ok {
-		return nil, ErrKeyNotFound.New(k)
+		return nil, ErrKeyNotFound
 	}
 
 	return v, nil
@@ -133,7 +132,7 @@ func (m mapCache) Put(u uint64, i interface{}) error {
 func (m mapCache) Get(u uint64) (interface{}, error) {
 	v, ok := m.cache[u]
 	if !ok {
-		return nil, ErrKeyNotFound.New(u)
+		return nil, ErrKeyNotFound
 	}
 	return v, nil
 }
@@ -173,7 +172,7 @@ func (h *historyCache) Put(k uint64, v interface{}) error {
 func (h *historyCache) Get(k uint64) (interface{}, error) {
 	v, ok := h.cache[k]
 	if !ok {
-		return nil, ErrKeyNotFound.New(k)
+		return nil, ErrKeyNotFound
 	}
 	return v, nil
 }

sql/cache_test.go

@@ -15,6 +15,7 @@
 package sql
 
 import (
+	"errors"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -33,14 +34,14 @@ func TestLRUCache(t *testing.T) {
 
 		_, err = cache.Get(2)
 		require.Error(err)
-		require.True(ErrKeyNotFound.Is(err))
+		require.True(errors.Is(err, ErrKeyNotFound))
 
 		// Free the cache and check previous entry disappeared.
 		cache.Free()
 
 		_, err = cache.Get(1)
 		require.Error(err)
-		require.True(ErrKeyNotFound.Is(err))
+		require.True(errors.Is(err, ErrKeyNotFound))
 
 		cache.Dispose()
 		require.Panics(func() {
@@ -55,7 +56,7 @@ func TestLRUCache(t *testing.T) {
 		require.NoError(cache.Put(1, "foo"))
 		_, err := cache.Get(1)
 		require.Error(err)
-		require.True(ErrKeyNotFound.Is(err))
+		require.True(errors.Is(err, ErrKeyNotFound))
 	})
 
 	t.Run("free required to add entry", func(t *testing.T) {
@@ -94,7 +95,7 @@ func TestHistoryCache(t *testing.T) {
 
 		_, err = cache.Get(2)
 		require.Error(err)
-		require.True(ErrKeyNotFound.Is(err))
+		require.True(errors.Is(err, ErrKeyNotFound))
 
 		cache.Dispose()
 		require.Panics(func() {

sql/collations.go

@@ -18,6 +18,7 @@ import (
 	"fmt"
 	"io"
 	"strings"
+	"sync"
 	"unicode/utf8"
 
 	"github.com/cespare/xxhash"
@@ -851,6 +852,12 @@ func (c CollationID) Collation() Collation {
 	return collationArray[c]
 }
 
+var weightBuffers = sync.Pool{
+	New: func() interface{} {
+		return new([]byte)
+	},
+}
+
 // WriteWeightString writes the weights of each codepoint in the string into the given io.Writer.
 // Two strings with technically different contents may generate the same WeightString to the same value, as the collation
 // considers them the same string.
@@ -863,24 +870,30 @@ func (c CollationID) WriteWeightString(hash io.Writer, str string) error {
 		}
 	} else {
 		getRuneWeight := collationArray[c].Sorter
+		i := 0
+		buf := *weightBuffers.Get().(*[]byte)
+		if cap(buf) < len(str)*4 {
+			buf = make([]byte, len(str)*4)
+		}
 		for len(str) > 0 {
 			// All strings (should) have been decoded at this point, so we can rely on Go's internal string encoding
 			runeFromString, strRead := utf8.DecodeRuneInString(str)
 			if strRead == 0 || strRead == utf8.RuneError {
 				return ErrCollationMalformedString.New("hashing")
 			}
 			runeWeight := getRuneWeight(runeFromString)
-			_, err := hash.Write([]byte{
-				byte(runeWeight),
-				byte(runeWeight >> 8),
-				byte(runeWeight >> 16),
-				byte(runeWeight >> 24),
-			})
+			buf[i*4] = byte(runeWeight)
+			buf[i*4+1] = byte(runeWeight >> 8)
+			buf[i*4+2] = byte(runeWeight >> 16)
+			buf[i*4+3] = byte(runeWeight >> 24)
+			_, err := hash.Write(buf[i*4 : i*4+4])
 			if err != nil {
 				return err
 			}
 			str = str[strRead:]
+			i++
 		}
+		weightBuffers.Put(&buf)
 	}
 	return nil
 }

sql/rowexec/agg.go

@@ -15,6 +15,7 @@
 package rowexec
 
 import (
+	"errors"
 	"fmt"
 	"io"
 
@@ -162,7 +163,7 @@ func (i *groupByGroupingIter) compute(ctx *sql.Context) error {
 		}
 
 		b, err := i.get(key)
-		if sql.ErrKeyNotFound.Is(err) {
+		if errors.Is(err, sql.ErrKeyNotFound) {
 			b = make([]sql.AggregationBuffer, len(i.selectedExprs))
 			for j, a := range i.selectedExprs {
 				b[j], err = newAggregationBuffer(a)

sql/rowexec/update.go

@@ -15,6 +15,7 @@
 package rowexec
 
 import (
+	"errors"
 	"fmt"
 
 	"github.com/dolthub/go-mysql-server/sql"
@@ -209,7 +210,7 @@ func (u *updateJoinIter) Next(ctx *sql.Context) (sql.Row, error) {
 			}
 
 			_, err = cache.Get(hash)
-			if sql.ErrKeyNotFound.Is(err) {
+			if errors.Is(err, sql.ErrKeyNotFound) {
 				cache.Put(hash, struct{}{})
 				continue
 			} else if err != nil {