Support unlimited number of ORs in ExtractINFromOR

go/vt/sqlparser/predicate_rewriting.go

@@ -17,6 +17,8 @@ limitations under the License.
 package sqlparser
 
 import (
+	"slices"
+
 	"vitess.io/vitess/go/vt/log"
 )
 
@@ -96,36 +98,143 @@ func simplifyNot(expr *NotExpr) (Expr, rewriteState) {
 	return expr, noChange{}
 }
 
-// ExtractINFromOR will add additional predicated to an OR.
-// this rewriter should not be used in a fixed point way, since it returns the original expression with additions,
-// and it will therefor OOM before it stops rewriting
+// ExtractINFromOR rewrites the OR expression into an IN clause.
+// Each side of each ORs has to be an equality comparison expression and the column names have to
+// match for all sides of each comparison.
+// This rewriter takes a query that looks like this WHERE a = 1 and b = 11 or a = 2 and b = 12 or a = 3 and b = 13
+// And rewrite that to WHERE (a, b) IN ((1,11), (2,12), (3,13))
 func ExtractINFromOR(expr *OrExpr) []Expr {
-	// we check if we have two comparisons on either side of the OR
-	// that we can add as an ANDed comparison.
-	// WHERE (a = 5 and B) or (a = 6 AND C) =>
-	// WHERE (a = 5 AND B) OR (a = 6 AND C) AND a IN (5,6)
-	// This rewrite makes it possible to find a better route than Scatter if the `a` column has a helpful vindex
-	lftPredicates := SplitAndExpression(nil, expr.Left)
-	rgtPredicates := SplitAndExpression(nil, expr.Right)
-	var ins []Expr
-	for _, lft := range lftPredicates {
-		l, ok := lft.(*ComparisonExpr)
-		if !ok {
-			continue
+	var varNames []*ColName
+	var values []Exprs
+	orSlice := orToSlice(expr)
+	for _, expr := range orSlice {
+		andSlice := andToSlice(expr)
+		if len(andSlice) == 0 {
+			return nil
 		}
-		for _, rgt := range rgtPredicates {
-			r, ok := rgt.(*ComparisonExpr)
-			if !ok {
-				continue
+
+		var currentVarNames []*ColName
+		var currentValues []Expr
+		for _, comparisonExpr := range andSlice {
+			if comparisonExpr.Operator != EqualOp {
+				return nil
+			}
+
+			var colName *ColName
+			switch left := comparisonExpr.Left.(type) {
+			case *ColName:
+				colName = left
+				currentValues = append(currentValues, comparisonExpr.Right)
+			}
+
+			switch right := comparisonExpr.Right.(type) {
+			case *ColName:
+				if colName != nil {
+					return nil
+				}
+				colName = right
+				currentValues = append(currentValues, comparisonExpr.Left)
+			}
+
+			if colName == nil {
+				return nil
 			}
-			in, state := tryTurningOrIntoIn(l, r)
-			if state.changed() {
-				ins = append(ins, in)
+
+			currentVarNames = append(currentVarNames, colName)
+		}
+
+		if len(varNames) == 0 {
+			varNames = currentVarNames
+		} else if !slices.EqualFunc(varNames, currentVarNames, func(col1, col2 *ColName) bool { return col1.Equal(col2) }) {
+			return nil
+		}
+
+		values = append(values, currentValues)
+	}
+
+	if len(varNames) == 1 {
+		var valueTuple ValTuple
+		for _, value := range values {
+			valueTuple = append(valueTuple, value[0])
+		}
+
+		return []Expr{&ComparisonExpr{
+			Operator: InOp,
+			Left:     varNames[0],
+			Right:    valueTuple,
+		}}
+	}
+
+	var nameTuple ValTuple
+	for _, name := range varNames {
+		nameTuple = append(nameTuple, name)
+	}
+
+	var valueTuple ValTuple
+	for _, value := range values {
+		valueTuple = append(valueTuple, ValTuple(value))
+	}
+
+	return []Expr{&ComparisonExpr{
+		Operator: InOp,
+		Left:     nameTuple,
+		Right:    valueTuple,
+	}}
+}
+
+func orToSlice(expr *OrExpr) []Expr {
+	var exprs []Expr
+
+	handleOrSide := func(e Expr) {
+		switch e := e.(type) {
+		case *OrExpr:
+			exprs = append(exprs, orToSlice(e)...)
+		default:
+			exprs = append(exprs, e)
+		}
+	}
+
+	handleOrSide(expr.Left)
+	handleOrSide(expr.Right)
+	return exprs
+}
+
+func andToSlice(expr Expr) []*ComparisonExpr {
+	var andExpr *AndExpr
+	switch expr := expr.(type) {
+	case *AndExpr:
+		andExpr = expr
+	case *ComparisonExpr:
+		return []*ComparisonExpr{expr}
+	default:
+		return nil
+	}
+
+	var exprs []*ComparisonExpr
+	handleAndSide := func(e Expr) bool {
+		switch e := e.(type) {
+		case *AndExpr:
+			slice := andToSlice(e)
+			if slice == nil {
+				return false
 			}
+			exprs = append(exprs, slice...)
+		case *ComparisonExpr:
+			exprs = append(exprs, e)
+		default:
+			return false
 		}
+		return true
+	}
+
+	if !handleAndSide(andExpr.Left) {
+		return nil
+	}
+	if !handleAndSide(andExpr.Right) {
+		return nil
 	}
 
-	return uniquefy(ins)
+	return exprs
 }
 
 func simplifyOr(expr *OrExpr, allowCNF bool) (Expr, rewriteState) {

go/vt/sqlparser/predicate_rewriting_test.go

@@ -140,6 +140,18 @@ func TestRewritePredicate(in *testing.T) {
 		// this has too many OR expressions in it, so we don't even try the CNF rewriting
 		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 or a = 6 and b = 46",
 		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 or a = 6 and b = 46",
+	}, {
+		in:       "a = 5 and B or a = 6 and C",
+		expected: "a in (5, 6) and (a = 5 or C) and ((B or a = 6) and (B or C))",
+	}, {
+		in:       "(a = 5 and b = 1 or b = 2 and a = 6)",
+		expected: "(a = 5 or b = 2) and a in (5, 6) and (b in (1, 2) and (b = 1 or a = 6))",
+	}, {
+		in:       "(a in (1,5) and B or C and a = 6)",
+		expected: "(a in (1, 5) or C) and a in (1, 5, 6) and ((B or C) and (B or a = 6))",
+	}, {
+		in:       "(a in (1, 5) and B or C and a in (5, 7))",
+		expected: "(a in (1, 5) or C) and a in (1, 5, 7) and ((B or C) and (B or a in (5, 7)))",
 	}}
 
 	for _, tc := range tests {
@@ -158,26 +170,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 = 5 and b = 1 or b = 2 and a = 6 or b = 3 and a = 4)",
-		expected: "<nil>",
+		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 or a = 6 and b = 46",
+		expected: "(a, b) in ((1, 41), (2, 42), (3, 43), (4, 44), (5, 45), (6, 46))",
 	}}
 
 	for _, tc := range tests {