WIP

rohit-nayak-ps · Nov 21, 2023 · 2a3124e · 2a3124e
1 parent 786a532
commit 2a3124e
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 72 deletions.
diff --git a/go/vt/sqlparser/predicate_rewriting.go b/go/vt/sqlparser/predicate_rewriting.go
@@ -17,6 +17,8 @@ limitations under the License.
 package sqlparser
 
 import (
+	"fmt"
+
 	"vitess.io/vitess/go/vt/log"
 )
 
@@ -25,6 +27,17 @@ import (
 // Note: In order to re-plan, we need to empty the accumulated metadata in the AST,
 // so ColName.Metadata will be nil:ed out as part of this rewrite
 func RewritePredicate(ast SQLNode) SQLNode {
+	count := 0
+	_ = Walk(func(node SQLNode) (bool, error) {
+		if _, isExpr := node.(*OrExpr); isExpr {
+			count++
+		}
+
+		return true, nil
+	}, ast)
+
+	allowCNF := count < 4
+
 	for {
 		printExpr(ast)
 		exprChanged := false
@@ -37,7 +50,7 @@ func RewritePredicate(ast SQLNode) SQLNode {
 				return true
 			}
 
-			rewritten, state := simplifyExpression(e)
+			rewritten, state := simplifyExpression(e, allowCNF)
 			if ch, isChange := state.(changed); isChange {
 				printRule(ch.rule, ch.exprMatched)
 				exprChanged = true
@@ -56,12 +69,12 @@ func RewritePredicate(ast SQLNode) SQLNode {
 	}
 }
 
-func simplifyExpression(expr Expr) (Expr, rewriteState) {
+func simplifyExpression(expr Expr, allowCNF bool) (Expr, rewriteState) {
 	switch expr := expr.(type) {
 	case *NotExpr:
 		return simplifyNot(expr)
 	case *OrExpr:
-		return simplifyOr(expr)
+		return simplifyOr(expr, allowCNF)
 	case *XorExpr:
 		return simplifyXor(expr)
 	case *AndExpr:
@@ -117,55 +130,66 @@ func ExtractINFromOR(expr *OrExpr) []Expr {
 	return uniquefy(ins)
 }
 
-func simplifyOr(expr *OrExpr) (Expr, rewriteState) {
+func simplifyOr(expr *OrExpr, allowCNF bool) (Expr, rewriteState) {
 	or := expr
 
 	// first we search for ANDs and see how they can be simplified
 	land, lok := or.Left.(*AndExpr)
 	rand, rok := or.Right.(*AndExpr)
-	switch {
-	case lok && rok:
+
+	if lok && rok {
 		// (<> AND <>) OR (<> AND <>)
 		var a, b, c Expr
 		var change changed
 		switch {
 		case Equals.Expr(land.Left, rand.Left):
 			change = newChange("(A and B) or (A and C) => A AND (B OR C)", f(expr))
 			a, b, c = land.Left, land.Right, rand.Right
+			return &AndExpr{Left: a, Right: &OrExpr{Left: b, Right: c}}, change
 		case Equals.Expr(land.Left, rand.Right):
 			change = newChange("(A and B) or (C and A) => A AND (B OR C)", f(expr))
 			a, b, c = land.Left, land.Right, rand.Left
+			return &AndExpr{Left: a, Right: &OrExpr{Left: b, Right: c}}, change
 		case Equals.Expr(land.Right, rand.Left):
 			change = newChange("(B and A) or (A and C) => A AND (B OR C)", f(expr))
 			a, b, c = land.Right, land.Left, rand.Right
+			return &AndExpr{Left: a, Right: &OrExpr{Left: b, Right: c}}, change
 		case Equals.Expr(land.Right, rand.Right):
 			change = newChange("(B and A) or (C and A) => A AND (B OR C)", f(expr))
 			a, b, c = land.Right, land.Left, rand.Left
-		default:
-			return expr, noChange{}
+			return &AndExpr{Left: a, Right: &OrExpr{Left: b, Right: c}}, change
 		}
-		return &AndExpr{Left: a, Right: &OrExpr{Left: b, Right: c}}, change
-	case lok:
-		// (<> AND <>) OR <>
+	}
+
+	// (<> AND <>) OR <>
+	if lok {
 		// Simplification
 		if Equals.Expr(or.Right, land.Left) || Equals.Expr(or.Right, land.Right) {
 			return or.Right, newChange("(A AND B) OR A => A", f(expr))
 		}
-		// Distribution Law
-		return &AndExpr{Left: &OrExpr{Left: land.Left, Right: or.Right}, Right: &OrExpr{Left: land.Right, Right: or.Right}},
-			newChange("(A AND B) OR C => (A OR C) AND (B OR C)", f(expr))
-	case rok:
-		// <> OR (<> AND <>)
+
+		if allowCNF {
+			// Distribution Law
+			return &AndExpr{Left: &OrExpr{Left: land.Left, Right: or.Right}, Right: &OrExpr{Left: land.Right, Right: or.Right}},
+				newChange("(A AND B) OR C => (A OR C) AND (B OR C)", f(expr))
+		}
+	}
+
+	// <> OR (<> AND <>)
+	if rok {
 		// Simplification
 		if Equals.Expr(or.Left, rand.Left) || Equals.Expr(or.Left, rand.Right) {
 			return or.Left, newChange("A OR (A AND B) => A", f(expr))
 		}
-		// Distribution Law
-		return &AndExpr{
-				Left:  &OrExpr{Left: or.Left, Right: rand.Left},
-				Right: &OrExpr{Left: or.Left, Right: rand.Right},
-			},
-			newChange("C OR (A AND B) => (C OR A) AND (C OR B)", f(expr))
+
+		if allowCNF {
+			// Distribution Law
+			return &AndExpr{
+					Left:  &OrExpr{Left: or.Left, Right: rand.Left},
+					Right: &OrExpr{Left: or.Left, Right: rand.Right},
+				},
+				newChange("C OR (A AND B) => (C OR A) AND (C OR B)", f(expr))
+		}
 	}
 
 	// next, we want to try to turn multiple ORs into an IN when possible
@@ -431,18 +455,20 @@ func f(e Expr) func() Expr {
 }
 
 func printRule(rule string, expr func() Expr) {
-	if log.V(10) {
-		log.Infof("Rule: %s   ON   %s", rule, String(expr()))
-	}
+	//	if log.V(10) {
+	fmt.Printf("Rule: %s   ON   %s\n", rule, String(expr()))
+	// }
 }
 
 func printExpr(expr SQLNode) {
-	if log.V(10) {
-		log.Infof("Current: %s", String(expr))
-	}
+	//	if log.V(10) {
+	fmt.Printf("Current: %s\n", String(expr))
+	// }
 }
 
 func newChange(rule string, exprMatched func() Expr) changed {
+	fmt.Println(rule, String(exprMatched()))
+
 	return changed{
 		rule:        rule,
 		exprMatched: exprMatched,

diff --git a/go/vt/sqlparser/predicate_rewriting_test.go b/go/vt/sqlparser/predicate_rewriting_test.go
@@ -91,7 +91,7 @@ func TestSimplifyExpression(in *testing.T) {
 			expr, err := ParseExpr(tc.in)
 			require.NoError(t, err)
 
-			expr, didRewrite := simplifyExpression(expr)
+			expr, didRewrite := simplifyExpression(expr, false)
 			assert.True(t, didRewrite.changed())
 			assert.Equal(t, tc.expected, String(expr))
 		})
@@ -103,32 +103,15 @@ func TestRewritePredicate(in *testing.T) {
 		in       string
 		expected string
 	}{{
-		in:       "A xor B",
-		expected: "(A or B) and (not A or not B)",
-	}, {
-		in:       "(A and B) and (B and A) and (B and A) and (A and B)",
-		expected: "A and B",
-	}, {
-		in:       "((A and B) OR (A and C) OR (A and D)) and E and F",
-		expected: "A and (B or C or D) and E and F",
-	}, {
-		in:       "(A and B) OR (A and C)",
-		expected: "A and (B or C)",
-	}, {
-		in:       "(A and B) OR (C and A)",
-		expected: "A and (B or C)",
+		in:       "(a = 1 and b = 41) or (a = 2 and b = 42)",
+		expected: "a = 1 and b = 41 or a = 2 and b = 42",
 	}, {
-		in:       "(B and A) OR (A and C)",
-		expected: "A and (B or C)",
-	}, {
-		in:       "(A and B) or (A and C) or (A and D)",
-		expected: "A and (B or C or D)",
+		in:       "(a = 1 and b = 41) or (a = 2 and b = 42) or (a = 3 and b = 43)",
+		expected: "a = 1 and b = 41 or a = 2 and b = 42 or a = 3 and b = 43",
+		//		"((a = 1 and b = 41) or (a = 2 and b = 42) or a = 3) and ((a = 1 and b = 41) or (a = 2 and b = 42) or b = 43)"
 	}, {
-		in:       "(a=1 or a IN (1,2)) or (a = 2 or a = 3)",
-		expected: "a in (1, 2, 3)",
-	}, {
-		in:       "A and (B or A)",
-		expected: "A",
+		in:       "(a = 1 and b = 41) or (a = 2 and b = 42) or (a = 3 and b = 43) or (a = 4 and b = 44) or (a = 5 and b = 45)",
+		expected: "a = 1 and b = 41 or a = 2 and b = 42 or a = 3 and b = 43 or a = 4 and b = 44 or a = 5 and b = 45",
 	}}
 
 	for _, tc := range tests {
@@ -147,23 +130,8 @@ func TestExtractINFromOR(in *testing.T) {
 		in       string
 		expected string
 	}{{
-		in:       "(A and B) or (B and A)",
-		expected: "<nil>",
-	}, {
-		in:       "(a = 5 and B) or A",
-		expected: "<nil>",
-	}, {
-		in:       "a = 5 and B or a = 6 and C",
-		expected: "a in (5, 6)",
-	}, {
-		in:       "(a = 5 and b = 1 or b = 2 and a = 6)",
-		expected: "a in (5, 6) and b in (1, 2)",
-	}, {
-		in:       "(a in (1,5) and B or C and a = 6)",
-		expected: "a in (1, 5, 6)",
-	}, {
-		in:       "(a in (1, 5) and B or C and a in (5, 7))",
-		expected: "a in (1, 5, 7)",
+		in:       "(a = 1 and b = 41) or (a = 2 and b = 42) or (a = 3 and b = 43) or (a = 4 and b = 44)",
+		expected: "(a, b) in ((1, 41), (2, 42), (3, 43), (4, 44))",
 	}}
 
 	for _, tc := range tests {

diff --git a/go/vt/vtgate/planbuilder/testdata/onecase.json b/go/vt/vtgate/planbuilder/testdata/onecase.json
@@ -1,9 +1,9 @@
 [
   {
     "comment": "Add your test case here for debugging and run go test -run=One.",
-    "query": "",
+    "query": "SELECT COUNT(*) FROM music WHERE (music.user_id = 1 AND music.name = 'test' OR music.user_id = 2 AND music.name = 'test2' OR music.user_id = 3 AND music.name = 'test3' OR music.user_id = 4 AND music.name = 'test4' OR music.user_id = 5 AND music.name = 'test5') AND music.other = false GROUP BY user_id, name",
     "plan": {
 
     }
   }
-]
+]