sql/plan: speed up InnerJoin iterator

erizocosmico · erizocosmico · commit b19c55a959ae · 2018-12-21T15:18:48.000+01:00
pkg: gopkg.in/src-d/go-mysql-server.v0/sql/plan
BenchmarkInnerJoin/inner_join-4         	   50000	     34188 ns/op	   11415 B/op	     169 allocs/op
BenchmarkInnerJoin/cross_join_with_filter-4         	   30000	     41608 ns/op	   12119 B/op	     181 allocs/op
PASS
ok  	gopkg.in/src-d/go-mysql-server.v0/sql/plan	3.795s

Signed-off-by: Miguel Molina &lt;miguel@erizocosmi.co&gt;
diff --git a/sql/plan/innerjoin.go b/sql/plan/innerjoin.go
@@ -1,6 +1,7 @@
 package plan
 
 import (
+	"io"
 	"reflect"
 
 	opentracing "github.com/opentracing/opentracing-go"
@@ -60,15 +61,12 @@ func (j *InnerJoin) RowIter(ctx *sql.Context) (sql.RowIter, error) {
 		return nil, err
 	}
 
-	return sql.NewSpanIter(span, NewFilterIter(
-		ctx,
-		j.Cond,
-		&crossJoinIterator{
-			l:  l,
-			rp: j.Right,
-			s:  ctx,
-		},
-	)), nil
+	return sql.NewSpanIter(span, &innerJoinIter{
+		l:    l,
+		rp:   j.Right,
+		ctx:  ctx,
+		cond: j.Cond,
+	}), nil
 }
 
 // TransformUp implements the Transformable interface.
@@ -127,3 +125,74 @@ func (j *InnerJoin) TransformExpressions(f sql.TransformExprFunc) (sql.Node, err
 
 	return NewInnerJoin(j.Left, j.Right, cond), nil
 }
+
+type innerJoinIter struct {
+	l    sql.RowIter
+	rp   rowIterProvider
+	r    sql.RowIter
+	ctx  *sql.Context
+	cond sql.Expression
+
+	leftRow sql.Row
+}
+
+func (i *innerJoinIter) 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 {
+			iter, err := i.rp.RowIter(i.ctx)
+			if err != nil {
+				return nil, err
+			}
+
+			i.r = iter
+		}
+
+		rightRow, err := i.r.Next()
+		if err == io.EOF {
+			i.r = nil
+			i.leftRow = nil
+			continue
+		}
+
+		if err != nil {
+			return nil, err
+		}
+
+		var row sql.Row
+		row = append(row, i.leftRow...)
+		row = append(row, rightRow...)
+
+		v, err := i.cond.Eval(i.ctx, row)
+		if err != nil {
+			return nil, err
+		}
+
+		if v == true {
+			return row, nil
+		}
+	}
+}
+
+func (i *innerJoinIter) 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
+}
diff --git a/sql/plan/innerjoin_test.go b/sql/plan/innerjoin_test.go
@@ -1,6 +1,7 @@
 package plan
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -57,3 +58,77 @@ func TestInnerJoinEmpty(t *testing.T) {
 
 	assertRows(t, iter, 0)
 }
+
+func BenchmarkInnerJoin(b *testing.B) {
+	t1 := mem.NewTable("foo", sql.Schema{
+		{Name: "a", Source: "foo", Type: sql.Int64},
+		{Name: "b", Source: "foo", Type: sql.Text},
+	})
+
+	t2 := mem.NewTable("bar", sql.Schema{
+		{Name: "a", Source: "bar", Type: sql.Int64},
+		{Name: "b", Source: "bar", Type: sql.Text},
+	})
+
+	for i := 0; i < 5; i++ {
+		t1.Insert(sql.NewEmptyContext(), sql.NewRow(int64(i), fmt.Sprintf("t1_%d", i)))
+		t2.Insert(sql.NewEmptyContext(), sql.NewRow(int64(i), fmt.Sprintf("t2_%d", i)))
+	}
+
+	n1 := NewInnerJoin(
+		NewResolvedTable(t1),
+		NewResolvedTable(t2),
+		expression.NewEquals(
+			expression.NewGetField(0, sql.Int64, "a", false),
+			expression.NewGetField(2, sql.Int64, "a", false),
+		),
+	)
+
+	n2 := NewFilter(
+		expression.NewEquals(
+			expression.NewGetField(0, sql.Int64, "a", false),
+			expression.NewGetField(2, sql.Int64, "a", false),
+		),
+		NewCrossJoin(
+			NewResolvedTable(t1),
+			NewResolvedTable(t2),
+		),
+	)
+
+	expected := []sql.Row{
+		{int64(0), "t1_0", int64(0), "t2_0"},
+		{int64(1), "t1_1", int64(1), "t2_1"},
+		{int64(2), "t1_2", int64(2), "t2_2"},
+		{int64(3), "t1_3", int64(3), "t2_3"},
+		{int64(4), "t1_4", int64(4), "t2_4"},
+	}
+
+	ctx := sql.NewEmptyContext()
+	b.Run("inner join", func(b *testing.B) {
+		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)
+		}
+	})
+
+	b.Run("cross join with filter", func(b *testing.B) {
+		require := require.New(b)
+
+		for i := 0; i < b.N; i++ {
+			iter, err := n2.RowIter(ctx)
+			require.NoError(err)
+
+			rows, err := sql.RowIterToRows(iter)
+			require.NoError(err)
+
+			require.Equal(expected, rows)
+		}
+	})
+}
