Fix HAVING rewriting made in #11306 (#11515)

* don't rewrite HAVING predicates that use table columns Signed-off-by: Andres Taylor <[email protected]> * Revert the changes made in #11306 Signed-off-by: Florent Poinsard <[email protected]> Signed-off-by: Andres Taylor <[email protected]> * Fix early rewriter test Signed-off-by: Florent Poinsard <[email protected]> Signed-off-by: Andres Taylor <[email protected]> Signed-off-by: Florent Poinsard <[email protected]> Co-authored-by: Andres Taylor <[email protected]>
vitessio · Oct 18, 2022 · 83543b8 · 83543b8
1 parent 8846bab
commit 83543b8
Show file tree

Hide file tree

Showing 4 changed files with 118 additions and 21 deletions.
diff --git a/go/vt/vtgate/planbuilder/rewrite.go b/go/vt/vtgate/planbuilder/rewrite.go
@@ -174,29 +174,40 @@ func rewriteHavingClause(node *sqlparser.Select) {
 	exprs := sqlparser.SplitAndExpression(nil, node.Having.Expr)
 	node.Having = nil
 	for _, expr := range exprs {
-		var hasAggr bool
-		sqlparser.Rewrite(expr, func(cursor *sqlparser.Cursor) bool {
-			switch x := cursor.Node().(type) {
-			case *sqlparser.ColName:
-				if !x.Qualifier.IsEmpty() {
-					return false
-				}
-				originalExpr, isInMap := selectExprMap[x.Name.Lowered()]
-				if isInMap && sqlparser.ContainsAggregation(originalExpr) {
-					hasAggr = true
-				}
-				return false
-			default:
-				_, isAggregate := x.(sqlparser.AggrFunc)
-				hasAggr = hasAggr || isAggregate
-			}
-			return true
-		}, nil)
-
+		hasAggr := sqlparser.ContainsAggregation(expr)
+		if !hasAggr {
+			sqlparser.Rewrite(expr, func(cursor *sqlparser.Cursor) bool {
+				visitColName(cursor.Node(), selectExprMap, func(original sqlparser.Expr) {
+					if sqlparser.ContainsAggregation(original) {
+						hasAggr = true
+					}
+				})
+				return true
+			}, nil)
+		}
 		if hasAggr {
 			node.AddHaving(expr)
 		} else {
+			sqlparser.Rewrite(expr, func(cursor *sqlparser.Cursor) bool {
+				visitColName(cursor.Node(), selectExprMap, func(original sqlparser.Expr) {
+					cursor.Replace(original)
+				})
+				return true
+			}, nil)
 			node.AddWhere(expr)
 		}
 	}
 }
+func visitColName(cursor sqlparser.SQLNode, selectExprMap map[string]sqlparser.Expr, f func(original sqlparser.Expr)) {
+	switch x := cursor.(type) {
+	case *sqlparser.ColName:
+		if !x.Qualifier.IsEmpty() {
+			return
+		}
+		originalExpr, isInMap := selectExprMap[x.Name.Lowered()]
+		if isInMap {
+			f(originalExpr)
+		}
+		return
+	}
+}
diff --git a/go/vt/vtgate/planbuilder/testdata/filter_cases.json b/go/vt/vtgate/planbuilder/testdata/filter_cases.json
@@ -6421,5 +6421,79 @@
         "user.user"
       ]
     }
+  },
+  {
+    "comment": "HAVING predicates that use table columns are safe to rewrite if we can move them to the WHERE clause",
+    "query": "select user.col + 2 as a from user having a = 42",
+    "v3-plan": {
+      "QueryType": "SELECT",
+      "Original": "select user.col + 2 as a from user having a = 42",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "Scatter",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select `user`.col + 2 as a from `user` where 1 != 1",
+        "Query": "select `user`.col + 2 as a from `user` having a = 42",
+        "Table": "`user`"
+      }
+    },
+    "gen4-plan": {
+      "QueryType": "SELECT",
+      "Original": "select user.col + 2 as a from user having a = 42",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "Scatter",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select `user`.col + 2 as a from `user` where 1 != 1",
+        "Query": "select `user`.col + 2 as a from `user` where `user`.col + 2 = 42",
+        "Table": "`user`"
+      },
+      "TablesUsed": [
+        "user.user"
+      ]
+    }
+  },
+  {
+    "comment": "HAVING predicates that use table columns should not get rewritten on unsharded keyspaces",
+    "query": "select col + 2 as a from unsharded having a = 42",
+    "v3-plan": {
+      "QueryType": "SELECT",
+      "Original": "select col + 2 as a from unsharded having a = 42",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "Unsharded",
+        "Keyspace": {
+          "Name": "main",
+          "Sharded": false
+        },
+        "FieldQuery": "select col + 2 as a from unsharded where 1 != 1",
+        "Query": "select col + 2 as a from unsharded having a = 42",
+        "Table": "unsharded"
+      }
+    },
+    "gen4-plan": {
+      "QueryType": "SELECT",
+      "Original": "select col + 2 as a from unsharded having a = 42",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "Unsharded",
+        "Keyspace": {
+          "Name": "main",
+          "Sharded": false
+        },
+        "FieldQuery": "select col + 2 as a from unsharded where 1 != 1",
+        "Query": "select col + 2 as a from unsharded having a = 42",
+        "Table": "unsharded"
+      },
+      "TablesUsed": [
+        "main.unsharded"
+      ]
+    }
   }
 ]
diff --git a/go/vt/vtgate/semantics/early_rewriter.go b/go/vt/vtgate/semantics/early_rewriter.go
@@ -146,7 +146,19 @@ func rewriteHavingAndOrderBy(cursor *sqlparser.Cursor, node sqlparser.SQLNode) {
 					continue
 				}
 				if ae.As.Equal(col.Name) {
-					inner.Replace(ae.Expr)
+					safeToRewrite := true
+					sqlparser.Rewrite(ae.Expr, func(cursor *sqlparser.Cursor) bool {
+						switch cursor.Node().(type) {
+						case *sqlparser.ColName:
+							safeToRewrite = false
+						case sqlparser.AggrFunc:
+							return false
+						}
+						return true
+					}, nil)
+					if safeToRewrite {
+						inner.Replace(ae.Expr)
+					}
 				}
 			}
 		}

diff --git a/go/vt/vtgate/semantics/early_rewriter_test.go b/go/vt/vtgate/semantics/early_rewriter_test.go
@@ -157,7 +157,7 @@ func TestExpandStar(t *testing.T) {
 		expSQL: "select t1.b as b, t1.a as a, t1.c as c, t5.a as a from t1 join t5 where t1.b = t5.b",
 	}, {
 		sql:    "select * from t1 join t5 using (b) having b = 12",
-		expSQL: "select t1.b as b, t1.a as a, t1.c as c, t5.a as a from t1 join t5 where t1.b = t5.b having t1.b = 12",
+		expSQL: "select t1.b as b, t1.a as a, t1.c as c, t5.a as a from t1 join t5 where t1.b = t5.b having b = 12",
 	}, {
 		sql:    "select 1 from t1 join t5 using (b) having b = 12",
 		expSQL: "select 1 from t1 join t5 where t1.b = t5.b having t1.b = 12",