src-d · ajnavarro · Jan 28, 2019 · Jan 28, 2019 · Jan 28, 2019 · kuba--
@@ -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
+}
@@ -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)