chore(engine): Eliminate no-op Merge/SortMerge nodes during plan optimization (#19238)

Merge and SortMerge nodes with 1 or less children are effectively no-op nodes and can safely be removed from the plan during the optimization phase.

Signed-off-by: Christian Haudum <[email protected]>

Author: chaudum

pkg/engine/planner/physical/optimizer.go

@@ -35,6 +35,26 @@ func (r *removeNoopFilter) apply(node Node) bool {
 
 var _ rule = (*removeNoopFilter)(nil)
 
+// removeNoopMerge is a rule that removes merge/sortmerge nodes with only a single input
+type removeNoopMerge struct {
+	plan *Plan
+}
+
+// apply implements rule.
+func (r *removeNoopMerge) apply(node Node) bool {
+	changed := false
+	switch node := node.(type) {
+	case *Merge, *SortMerge:
+		if len(r.plan.Children(node)) <= 1 {
+			r.plan.eliminateNode(node)
+			changed = true
+		}
+	}
+	return changed
+}
+
+var _ rule = (*removeNoopMerge)(nil)
+
 // predicatePushdown is a rule that moves down filter predicates to the scan nodes.
 type predicatePushdown struct {
 	plan *Plan

pkg/engine/planner/physical/optimizer_test.go

@@ -1,6 +1,7 @@
 package physical
 
 import (
+	"fmt"
 	"sort"
 	"testing"
 	"time"
@@ -104,6 +105,7 @@ func dummyPlan() *Plan {
 }
 
 func TestOptimizer(t *testing.T) {
+
 	t.Run("noop", func(t *testing.T) {
 		plan := dummyPlan()
 		optimizations := []*optimization{
@@ -561,6 +563,43 @@ func TestOptimizer(t *testing.T) {
 		expected := PrintAsTree(expectedPlan)
 		require.Equal(t, expected, actual)
 	})
+
+	t.Run("cleanup no-op merge nodes", func(t *testing.T) {
+		plan := func() *Plan {
+			plan := &Plan{}
+			limit := plan.addNode(&Limit{id: "limit"})
+			merge := plan.addNode(&Merge{id: "merge"})
+			sortmerge := plan.addNode(&Merge{id: "sortmerge"})
+			scan := plan.addNode(&DataObjScan{id: "scan"})
+
+			_ = plan.addEdge(Edge{Parent: limit, Child: merge})
+			_ = plan.addEdge(Edge{Parent: merge, Child: sortmerge})
+			_ = plan.addEdge(Edge{Parent: sortmerge, Child: scan})
+			return plan
+		}()
+
+		optimizations := []*optimization{
+			newOptimization("cleanup", plan).withRules(
+				&removeNoopMerge{plan},
+			),
+		}
+
+		o := newOptimizer(plan, optimizations)
+		o.optimize(plan.Roots()[0])
+		actual := PrintAsTree(plan)
+
+		optimized := func() *Plan {
+			plan := &Plan{}
+			limit := plan.addNode(&Limit{id: "limit"})
+			scan := plan.addNode(&DataObjScan{id: "scan"})
+
+			_ = plan.addEdge(Edge{Parent: limit, Child: scan})
+			return plan
+		}()
+
+		expected := PrintAsTree(optimized)
+		require.Equal(t, expected, actual, fmt.Sprintf("Expected:\n%s\nActual:\n%s\n", expected, actual))
+	})
 }
 
 func TestProjectionPushdown_PushesRequestedKeysToParseNodes(t *testing.T) {

pkg/engine/planner/physical/planner.go

@@ -381,14 +381,17 @@ func (p *Planner) Optimize(plan *Plan) (*Plan, error) {
 		optimizations := []*optimization{
 			newOptimization("PredicatePushdown", plan).withRules(
 				&predicatePushdown{plan: plan},
-				&removeNoopFilter{plan: plan},
 			),
 			newOptimization("LimitPushdown", plan).withRules(
 				&limitPushdown{plan: plan},
 			),
 			newOptimization("ProjectionPushdown", plan).withRules(
 				&projectionPushdown{plan: plan},
 			),
+			newOptimization("Cleanup", plan).withRules(
+				&removeNoopFilter{plan: plan},
+				&removeNoopMerge{plan: plan},
+			),
 		}
 		optimizer := newOptimizer(plan, optimizations)
 		optimizer.optimize(root)