Repository: src-d/go-mysql-server
Pull request #605 (2019-01-28 09:20:49)

add experimental optional feature to use in memory joins
This exposes the env var `EXPERIMENTAL_IN_MEMORY_JOIN` to enable the in-memory join feature, which causes inner joins to be performed in memory, which is significantly faster.

This way, clients can decide if they should enable in memory joins or not depending on their environment. If they can trade memory usage for speed this is orders of magnitude faster (specially for computationally expensive inner join branches) than the regular inner joins.

These are the benchmarks:

```
goos: darwin
goarch: amd64
pkg: gopkg.in/src-d/go-mysql-server.v0/sql/plan
BenchmarkInnerJoin/inner_join-4         	   50000	     33170 ns/op	   10618 B/op	     144 allocs/op
BenchmarkInnerJoin/in_memory_inner_join-4         	  100000	     21972 ns/op	    7322 B/op	      84 allocs/op
BenchmarkInnerJoin/cross_join_with_filter-4       	   30000	     41750 ns/op	   12122 B/op	     181 allocs/op
PASS
ok  	gopkg.in/src-d/go-mysql-server.v0/sql/plan	6.112s
```


Patch:
diff --git a/sql/plan/innerjoin.go b/sql/plan/innerjoin.go
index ff61a57..5645ade 100644
--- a/sql/plan/innerjoin.go
+++ b/sql/plan/innerjoin.go
@@ -2,12 +2,18 @@ package plan
 
 import (
 	"io"
+	"os"
 	"reflect"
 
 	opentracing "github.com/opentracing/opentracing-go"
 	"gopkg.in/src-d/go-mysql-server.v0/sql"
 )
 
+const experimentalInMemoryJoinKey = "EXPERIMENTAL_IN_MEMORY_JOIN"
+const inMemoryJoinSessionVar = "inmemory_joins"
+
+var useInMemoryJoins = os.Getenv(experimentalInMemoryJoinKey) != ""
+
 // InnerJoin is an inner join between two tables.
 type InnerJoin struct {
 	BinaryNode
@@ -61,12 +67,36 @@ func (j *InnerJoin) RowIter(ctx *sql.Context) (sql.RowIter, error) {
 		return nil, err
 	}
 
-	return sql.NewSpanIter(span, &innerJoinIter{
-		l:    l,
-		rp:   j.Right,
-		ctx:  ctx,
-		cond: j.Cond,
-	}), nil
+	var inMemorySession bool
+	_, val := ctx.Get(inMemoryJoinSessionVar)
+	if val != nil {
+		inMemorySession = true
+	}
+
+	var iter sql.RowIter
+	if useInMemoryJoins || inMemorySession {
+		r, err := j.Right.RowIter(ctx)
+		if err != nil {
+			span.Finish()
+			return nil, err
+		}
+
+		iter = &innerJoinMemoryIter{
+			l:    l,
+			r:    r,
+			ctx:  ctx,
+			cond: j.Cond,
+		}
+	} else {
+		iter = &innerJoinIter{
+			l:    l,
+			rp:   j.Right,
+			ctx:  ctx,
+			cond: j.Cond,
+		}
+	}
+
+	return sql.NewSpanIter(span, iter), nil
 }
 
 // TransformUp implements the Transformable interface.
@@ -196,3 +226,78 @@ func (i *innerJoinIter) Close() error {
 
 	return nil
 }
+
+type innerJoinMemoryIter struct {
+	l       sql.RowIter
+	r       sql.RowIter
+	ctx     *sql.Context
+	cond    sql.Expression
+	pos     int
+	leftRow sql.Row
+	right   []sql.Row
+}
+
+func (i *innerJoinMemoryIter) Next() (sql.Row, error) {
+	for {
+		if i.leftRow == nil {
+			r, err := i.l.Next()
+			if err != nil {
+				return nil, err
+			}
+
+			i.leftRow = r
+		}
+
+		if i.r != nil {
+			for {
+				row, err := i.r.Next()
+				if err != nil {
+					if err == io.EOF {
+						break
+					}
+					return nil, err
+				}
+
+				i.right = append(i.right, row)
+			}
+			i.r = nil
+		}
+
+		if i.pos >= len(i.right) {
+			i.pos = 0
+			i.leftRow = nil
+			continue
+		}
+
+		rightRow := i.right[i.pos]
+		var row = make(sql.Row, len(i.leftRow)+len(rightRow))
+		copy(row, i.leftRow)
+		copy(row[len(i.leftRow):], rightRow)
+
+		i.pos++
+
+		v, err := i.cond.Eval(i.ctx, row)
+		if err != nil {
+			return nil, err
+		}
+
+		if v == true {
+			return row, nil
+		}
+	}
+}
+
+func (i *innerJoinMemoryIter) Close() error {
+	if err := i.l.Close(); err != nil {
+		if i.r != nil {
+			_ = i.r.Close()
+		}
+		return err
+	}
+
+	if i.r != nil {
+		return i.r.Close()
+	}
+
+	return nil
+}


Test patch:
diff --git a/sql/plan/innerjoin_test.go b/sql/plan/innerjoin_test.go
index 4ec1fbd..51ca6f4 100644
--- a/sql/plan/innerjoin_test.go
+++ b/sql/plan/innerjoin_test.go
@@ -38,6 +38,41 @@ func TestInnerJoin(t *testing.T) {
 	}, rows)
 }
 
+func TestInMemoryInnerJoin(t *testing.T) {
+	require := require.New(t)
+	finalSchema := append(lSchema, rSchema...)
+
+	ltable := mem.NewTable("left", lSchema)
+	rtable := mem.NewTable("right", rSchema)
+	insertData(t, ltable)
+	insertData(t, rtable)
+
+	j := NewInnerJoin(
+		NewResolvedTable(ltable),
+		NewResolvedTable(rtable),
+		expression.NewEquals(
+			expression.NewGetField(0, sql.Text, "lcol1", false),
+			expression.NewGetField(4, sql.Text, "rcol1", false),
+		))
+
+	require.Equal(finalSchema, j.Schema())
+
+	ctx := sql.NewEmptyContext()
+	ctx.Set(inMemoryJoinSessionVar, sql.Text, "true")
+
+	iter, err := j.RowIter(ctx)
+	require.NoError(err)
+
+	rows, err := sql.RowIterToRows(iter)
+	require.NoError(err)
+	require.Len(rows, 2)
+
+	require.Equal([]sql.Row{
+		{"col1_1", "col2_1", int32(1111), int64(2222), "col1_1", "col2_1", int32(1111), int64(2222)},
+		{"col1_2", "col2_2", int32(3333), int64(4444), "col1_2", "col2_2", int32(3333), int64(4444)},
+	}, rows)
+}
+
 func TestInnerJoinEmpty(t *testing.T) {
 	require := require.New(t)
 	ctx := sql.NewEmptyContext()
@@ -118,6 +153,23 @@ func BenchmarkInnerJoin(b *testing.B) {
 		}
 	})
 
+	b.Run("in memory inner join", func(b *testing.B) {
+		useInMemoryJoins = true
+		require := require.New(b)
+
+		for i := 0; i < b.N; i++ {
+			iter, err := n1.RowIter(ctx)
+			require.NoError(err)
+
+			rows, err := sql.RowIterToRows(iter)
+			require.NoError(err)
+
+			require.Equal(expected, rows)
+		}
+
+		useInMemoryJoins = false
+	})
+
 	b.Run("cross join with filter", func(b *testing.B) {
 		require := require.New(b)