internal/rule: fix squashjoins rule to squash projections properly too

Signed-off-by: Manuel Carmona <[email protected]>

Author: mcarmonaa

integration_test.go

@@ -132,6 +132,17 @@ func TestIntegration(t *testing.T) {
 				{int32(4), "b029517f6300c2da0f4b651b8642506cd6aaf45d"},
 			},
 		},
+		{
+			`SELECT count(1), refs.repository_id
+				FROM refs
+				NATURAL JOIN commits
+				NATURAL JOIN commit_blobs
+				WHERE refs.ref_name = 'HEAD'
+				GROUP BY refs.repository_id`,
+			[]sql.Row{
+				{int32(9), "worktree"},
+			},
+		},
 	}
 
 	runTests := func(t *testing.T) {

internal/rule/squashjoins.go

@@ -31,6 +31,9 @@ func SquashJoins(
 	defer span.Finish()
 
 	a.Log("squashing joins, node of type %T", n)
+
+	projectSquashes := countProjectSquashes(n)
+
 	n, err := n.TransformUp(func(n sql.Node) (sql.Node, error) {
 		join, ok := n.(*plan.InnerJoin)
 		if !ok {
@@ -39,18 +42,103 @@ func SquashJoins(
 
 		return squashJoin(join)
 	})
+
 	if err != nil {
 		return nil, err
 	}
 
-	return n.TransformUp(func(n sql.Node) (sql.Node, error) {
+	n, err = n.TransformUp(func(n sql.Node) (sql.Node, error) {
 		t, ok := n.(*joinedTables)
 		if !ok {
 			return n, nil
 		}
 
 		return buildSquashedTable(t.tables, t.filters, t.columns, t.indexes)
 	})
+
+	if err != nil {
+		return nil, err
+	}
+
+	return n.TransformUp(func(n sql.Node) (sql.Node, error) {
+		if projectSquashes <= 0 {
+			return n, nil
+		}
+
+		project, ok := n.(*plan.Project)
+		if !ok {
+			return n, nil
+		}
+
+		child, ok := project.Child.(*plan.Project)
+		if !ok {
+			return n, nil
+		}
+
+		squashedProject, err := squashProjects(project, child)
+		if err != nil {
+			return nil, err
+		}
+
+		projectSquashes--
+		return squashedProject, nil
+	})
+}
+
+func countProjectSquashes(n sql.Node) int {
+	var squashableProjects int
+	plan.Inspect(n, func(node sql.Node) bool {
+		if project, ok := node.(*plan.Project); ok {
+			if _, ok := project.Child.(*plan.InnerJoin); ok {
+				squashableProjects++
+			}
+		}
+
+		return true
+	})
+
+	return squashableProjects - 1
+}
+
+// ErrWrongProjection is raised if a plan.Project node contains a wrong expression.
+var ErrWrongProjection = errors.NewKind("wrong expression found in project node %s")
+
+func squashProjects(parent, child *plan.Project) (sql.Node, error) {
+	projections := []sql.Expression{}
+	for _, expr := range parent.Expressions() {
+		parentField, ok := expr.(*expression.GetField)
+		if !ok {
+			return nil, ErrWrongProjection.New(parent.String())
+		}
+
+		index := parentField.Index()
+		for _, e := range child.Expressions() {
+			childField, ok := e.(*expression.GetField)
+			if !ok {
+				return nil, ErrWrongProjection.New(child.String())
+			}
+
+			if referenceSameColumn(parentField, childField) {
+				index = childField.Index()
+			}
+		}
+
+		projection := expression.NewGetFieldWithTable(
+			index,
+			parentField.Type(),
+			parentField.Table(),
+			parentField.Name(),
+			parentField.IsNullable(),
+		)
+
+		projections = append(projections, projection)
+	}
+
+	return plan.NewProject(projections, child.Child), nil
+}
+
+func referenceSameColumn(parent, child *expression.GetField) bool {
+	return parent.Name() == child.Name() && parent.Table() == child.Table()
 }
 
 func squashJoin(join *plan.InnerJoin) (sql.Node, error) {