feat: fix count star creations from the planner

Signed-off-by: Manan Gupta <manan@planetscale.com>
slackhq · Jun 19, 2024 · 6de3492 · 6de3492
1 parent 748d9c2
commit 6de3492
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 16 deletions.
diff --git a/go/vt/sqlparser/cached_size.go b/go/vt/sqlparser/cached_size.go
diff --git a/go/vt/vtgate/planbuilder/select.go b/go/vt/vtgate/planbuilder/select.go
@@ -301,7 +301,7 @@ func buildSQLCalcFoundRowsPlan(
 	sel2.Limit = nil
 
 	countStartExpr := []sqlparser.SelectExpr{&sqlparser.AliasedExpr{
-		Expr: &sqlparser.CountStar{},
+		Expr: &sqlparser.CountStar{Name: "count"},
 	}}
 	if sel2.GroupBy == nil && sel2.Having == nil {
 		// if there is no grouping, we can use the same query and

diff --git a/go/vt/vtgate/planbuilder/testdata/onecase.json b/go/vt/vtgate/planbuilder/testdata/onecase.json
@@ -1,30 +1,67 @@
 [
   {
     "comment": "Add your test case here for debugging and run go test -run=One.",
-    "query": "select Count(*) from user",
+    "query": "select Count(*) from user join user_extra where user.intcol + user_extra.incol < 5 limit 5",
     "plan": {
       "QueryType": "SELECT",
-      "Original": "select Count(*) from user",
+      "Original": "select Count(*) from user join user_extra where user.intcol + user_extra.incol < 5 limit 5",
       "Instructions": {
-        "OperatorType": "Aggregate",
-        "Variant": "Scalar",
-        "Aggregates": "sum_count_star(0) AS Count(*)",
+        "OperatorType": "Limit",
+        "Count": "5",
         "Inputs": [
           {
-            "OperatorType": "Route",
-            "Variant": "Scatter",
-            "Keyspace": {
-              "Name": "user",
-              "Sharded": true
-            },
-            "FieldQuery": "select Count(*) from `user` where 1 != 1",
-            "Query": "select Count(*) from `user`",
-            "Table": "`user`"
+            "OperatorType": "Aggregate",
+            "Variant": "Scalar",
+            "Aggregates": "sum_count_star(0) AS Count(*)",
+            "Inputs": [
+              {
+                "OperatorType": "Projection",
+                "Expressions": [
+                  "count(*) * count(*) as Count(*)"
+                ],
+                "Inputs": [
+                  {
+                    "OperatorType": "Join",
+                    "Variant": "Join",
+                    "JoinColumnIndexes": "L:0,R:0",
+                    "JoinVars": {
+                      "user_intcol": 1
+                    },
+                    "TableName": "`user`_user_extra",
+                    "Inputs": [
+                      {
+                        "OperatorType": "Route",
+                        "Variant": "Scatter",
+                        "Keyspace": {
+                          "Name": "user",
+                          "Sharded": true
+                        },
+                        "FieldQuery": "select count(*), `user`.intcol from `user` where 1 != 1 group by `user`.intcol",
+                        "Query": "select count(*), `user`.intcol from `user` group by `user`.intcol",
+                        "Table": "`user`"
+                      },
+                      {
+                        "OperatorType": "Route",
+                        "Variant": "Scatter",
+                        "Keyspace": {
+                          "Name": "user",
+                          "Sharded": true
+                        },
+                        "FieldQuery": "select count(*) from user_extra where 1 != 1 group by .0",
+                        "Query": "select count(*) from user_extra where :user_intcol + user_extra.incol < 5 group by .0",
+                        "Table": "user_extra"
+                      }
+                    ]
+                  }
+                ]
+              }
+            ]
           }
         ]
       },
       "TablesUsed": [
-        "user.user"
+        "user.user",
+        "user.user_extra"
       ]
     }
   }