internal/rule: fix squash joins when there are exchanges

Signed-off-by: Miguel Molina <[email protected]>

Author: erizocosmico

internal/rule/squashjoins.go

@@ -182,17 +182,21 @@ func joinTables(join *plan.InnerJoin) (*joinedTables, error) {
 func rearrange(join *plan.InnerJoin, squashedTable *joinedTables) sql.Node {
 	var projections []sql.Expression
 	var filters []sql.Expression
+	var parallelism int
 	plan.Inspect(join, func(node sql.Node) bool {
 		switch node := node.(type) {
 		case *plan.Project:
 			projections = append(projections, node.Projections...)
 		case *plan.Filter:
 			filters = append(filters, node.Expression)
+		case *plan.Exchange:
+			parallelism = node.Parallelism
 		}
 		return true
 	})
 
 	var node sql.Node = squashedTable
+
 	if len(filters) > 0 {
 		node = plan.NewFilter(expression.JoinAnd(filters...), node)
 	}
@@ -201,6 +205,10 @@ func rearrange(join *plan.InnerJoin, squashedTable *joinedTables) sql.Node {
 		node = plan.NewProject(projections, node)
 	}
 
+	if parallelism > 1 {
+		node = plan.NewExchange(parallelism, node)
+	}
+
 	return node
 }
 
@@ -865,7 +873,7 @@ func isJoinLeafSquashable(node sql.Node) bool {
 				return false
 			}
 			hasSquashableTables = true
-		case *plan.Project, *plan.Filter, *plan.TableAlias, nil:
+		case *plan.Project, *plan.Filter, *plan.TableAlias, *plan.Exchange, nil:
 		default:
 			hasUnsquashableNodes = true
 			return false

internal/rule/squashjoins_test.go

@@ -10,9 +10,39 @@ import (
 	"gopkg.in/src-d/go-mysql-server.v0/sql"
 	"gopkg.in/src-d/go-mysql-server.v0/sql/analyzer"
 	"gopkg.in/src-d/go-mysql-server.v0/sql/expression"
+	"gopkg.in/src-d/go-mysql-server.v0/sql/parse"
 	"gopkg.in/src-d/go-mysql-server.v0/sql/plan"
 )
 
+func TestAnalyzeSquashJoinsExchange(t *testing.T) {
+	require := require.New(t)
+
+	catalog := sql.NewCatalog()
+	catalog.AddDatabase(gitbase.NewDatabase("foo"))
+	a := analyzer.NewBuilder(catalog).
+		WithParallelism(2).
+		AddPostAnalyzeRule(SquashJoinsRule, SquashJoins).
+		Build()
+	a.CurrentDatabase = "foo"
+	ctx := sql.NewEmptyContext()
+
+	node, err := parse.Parse(ctx, `SELECT * FROM ref_commits NATURAL JOIN commits`)
+	require.NoError(err)
+
+	result, err := a.Analyze(ctx, node)
+	require.NoError(err)
+
+	project, ok := result.(*plan.Project)
+	require.True(ok)
+
+	exchange, ok := project.Child.(*plan.Exchange)
+	require.True(ok)
+	require.Equal(2, exchange.Parallelism)
+
+	_, ok = exchange.Child.(*gitbase.SquashedTable)
+	require.True(ok)
+}
+
 func TestSquashJoins(t *testing.T) {
 	require := require.New(t)