Add support for join filter pullups (#6272)

This commit adds functionality to the optimizer to pull up simple join
predicates. Currently this will only pullup simple predicate expressions
where keys are compared against constant values.

Author: mattnibs

compiler/optimizer/optimizer.go

@@ -144,6 +144,7 @@ func (o *Optimizer) Optimize(main *dag.Main) error {
 	seq = mergeFilters(seq)
 	seq = mergeValuesOps(seq)
 	inlineRecordExprSpreads(seq)
+	seq = joinFilterPullup(seq)
 	seq = removePassOps(seq)
 	seq = replaceSortAndHeadOrTailWithTop(seq)
 	o.optimizeParallels(seq)
@@ -506,6 +507,103 @@ func inlineRecordExprSpreads(v any) {
 	})
 }
 
+func joinFilterPullup(seq dag.Seq) dag.Seq {
+	seq = mergeFilters(seq)
+	for i := 0; i <= len(seq)-3; i++ {
+		fork, isfork := seq[i].(*dag.Fork)
+		leftAlias, rightAlias, isjoin := isJoin(seq[i+1])
+		filter, isfilter := seq[i+2].(*dag.Filter)
+		if !isfork || !isjoin || !isfilter {
+			continue
+		}
+		if len(fork.Paths) != 2 {
+			panic(seq[i])
+		}
+		var remaining []dag.Expr
+		for _, e := range splitPredicate(filter.Expr) {
+			if pullup, ok := pullupExpr(leftAlias, e); ok {
+				fork.Paths[0] = append(fork.Paths[0], dag.NewFilter(pullup))
+				continue
+			}
+			if pullup, ok := pullupExpr(rightAlias, e); ok {
+				fork.Paths[1] = append(fork.Paths[1], dag.NewFilter(pullup))
+				continue
+			}
+			remaining = append(remaining, e)
+		}
+		if len(remaining) == 0 {
+			// Filter has been fully pulled up and can be removed.
+			seq.Delete(i+2, i+3)
+		} else {
+			out := remaining[0]
+			for _, e := range remaining[1:] {
+				out = dag.NewBinaryExpr("and", e, out)
+			}
+			seq[i+2] = dag.NewFilter(out)
+		}
+		fork.Paths[0] = joinFilterPullup(fork.Paths[0])
+		fork.Paths[1] = joinFilterPullup(fork.Paths[1])
+	}
+	return seq
+}
+
+func isJoin(op dag.Op) (string, string, bool) {
+	switch op := op.(type) {
+	case *dag.HashJoin:
+		return op.LeftAlias, op.RightAlias, true
+	case *dag.Join:
+		return op.LeftAlias, op.RightAlias, true
+	default:
+		return "", "", false
+	}
+}
+
+func splitPredicate(e dag.Expr) []dag.Expr {
+	if b, ok := e.(*dag.BinaryExpr); ok && b.Op == "and" {
+		return append(splitPredicate(b.LHS), splitPredicate(b.RHS)...)
+	}
+	return []dag.Expr{e}
+}
+
+func pullupExpr(alias string, expr dag.Expr) (dag.Expr, bool) {
+	e, ok := expr.(*dag.BinaryExpr)
+	if !ok {
+		return nil, false
+	}
+	if e.Op == "and" {
+		lhs, lok := pullupExpr(alias, e.LHS)
+		rhs, rok := pullupExpr(alias, e.RHS)
+		if !lok || !rok {
+			return nil, false
+		}
+		return dag.NewBinaryExpr("and", lhs, rhs), true
+	}
+	if e.Op == "or" {
+		lhs, lok := pullupExpr(alias, e.LHS)
+		rhs, rok := pullupExpr(alias, e.RHS)
+		if !lok || !rok {
+			return nil, false
+		}
+		return dag.NewBinaryExpr("or", lhs, rhs), true
+
+	}
+	var literal *dag.Literal
+	var this *dag.This
+	for _, e := range []dag.Expr{e.RHS, e.LHS} {
+		if l, ok := e.(*dag.Literal); ok && literal == nil {
+			literal = l
+			continue
+		}
+		if t, ok := e.(*dag.This); ok && this == nil && len(t.Path) > 1 && t.Path[0] == alias {
+			this = t
+			continue
+		}
+		return nil, false
+	}
+	path := slices.Clone(this.Path[1:])
+	return dag.NewBinaryExpr(e.Op, dag.NewThis(path), literal), true
+}
+
 func liftFilterOps(seq dag.Seq) dag.Seq {
 	walkT(reflect.ValueOf(&seq), func(seq dag.Seq) dag.Seq {
 		for i := len(seq) - 2; i >= 0; i-- {

compiler/ztests/sql/join-filter-pullup.yaml

@@ -0,0 +1,52 @@
+script: |
+  super compile -C -O "
+    select *
+    from (values (1,'badgers')) as teams(id,team)
+    inner join (values (1,1,'blake','P')) as players(id,team_id,player,position) on teams.id = players.team_id
+    where players.position = 'P' and 'badgers' = team
+  "
+  echo // ===
+  super compile -C -O "
+    select *
+    from (values (1)) as t1(a1), (values (1)) as t2(a2), (values (1)) as t3(a3)
+    where a1 == 1 and a3 == 1 and (1 == a2 or a2 == 2)
+  "
+
+outputs:
+  - name: stdout
+    data: |
+      null
+      | fork
+        (
+          values {id:1,team:"badgers"}
+          | where team=="badgers"
+        )
+        (
+          values {id:1,team_id:1,player:"blake",position:"P"}
+          | where position=="P"
+        )
+      | inner hashjoin as {left,right} on id==team_id
+      | values {team:left.team,id:right.id,team_id:right.team_id,player:right.player,position:right.position}
+      | output main
+      // ===
+      null
+      | fork
+        (
+          values {a1:1}
+          | where a1==1
+        )
+        (
+          fork
+            (
+              values {a2:1}
+              | where a2==1 or a2==2
+            )
+            (
+              values {a3:1}
+              | where a3==1
+            )
+          | cross join as {left,right}
+        )
+      | cross join as {left,right}
+      | values {a1:left.a1,a2:right.left.a2,a3:right.right.a3}
+      | output main