wip

Signed-off-by: Florent Poinsard <florent.poinsard@outlook.fr>
planetscale · Jan 16, 2025 · 705e94d · 705e94d
1 parent ac8bbc6
commit 705e94d
Show file tree

Hide file tree

Showing 12 changed files with 980 additions and 8 deletions.
diff --git a/cpu.prof b/cpu.prof
diff --git a/go/vt/vtgate/planbuilder/operators/ast_to_op.go b/go/vt/vtgate/planbuilder/operators/ast_to_op.go
@@ -18,6 +18,7 @@ package operators
 
 import (
 	"fmt"
+	"log"
 
 	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vterrors"
@@ -109,14 +110,25 @@ func addWherePredsToSubQueryBuilder(ctx *plancontext.PlanningContext, in sqlpars
 }
 
 // cloneASTAndSemState clones the AST and the semantic state of the input node.
-func cloneASTAndSemState[T sqlparser.SQLNode](ctx *plancontext.PlanningContext, original T) T {
-	return sqlparser.CopyOnRewrite(original, nil, func(cursor *sqlparser.CopyOnWriteCursor) {
+func cloneASTAndSemState(ctx *plancontext.PlanningContext, original sqlparser.Expr, subq *sqlparser.Subquery) (sqlparser.Expr, *sqlparser.Subquery) {
+	var newSubQ *sqlparser.Subquery
+	newOriginal := sqlparser.CopyOnRewrite(original, nil, func(cursor *sqlparser.CopyOnWriteCursor) {
 		e, ok := cursor.Node().(sqlparser.Expr)
 		if !ok {
 			return
 		}
-		cursor.Replace(e) // We do this only to trigger the cloning of the AST
-	}, ctx.SemTable.CopySemanticInfo).(T)
+
+		if sq, ok := e.(*sqlparser.Subquery); ok {
+			log.Println("aa")
+			if subq == e {
+				newSubQ = sq
+			}
+		}
+
+		cursor.Replace(e) // We do this only to trigger the cloning of t
+	}, ctx.SemTable.CopySemanticInfo)
+
+	return newOriginal.(sqlparser.Expr), newSubQ
 }
 
 // findTablesContained returns the TableSet of all the contained

diff --git a/go/vt/vtgate/planbuilder/operators/subquery_builder.go b/go/vt/vtgate/planbuilder/operators/subquery_builder.go
@@ -162,8 +162,8 @@ func createSubquery(
 	isArg bool,
 ) *SubQuery {
 	topLevel := ctx.SemTable.EqualsExpr(original, parent)
-	original = cloneASTAndSemState(ctx, original)
-	originalSq := cloneASTAndSemState(ctx, subq)
+	original, originalSq := cloneASTAndSemState(ctx, original, subq)
+	// originalSq := cloneASTAndSemState(ctx, subq)
 	subqID := findTablesContained(ctx, subq.Select)
 	totalID := subqID.Merge(outerID)
 	sqc := &SubQueryBuilder{totalID: totalID, subqID: subqID, outerID: outerID}

diff --git a/go/vt/vtgate/planbuilder/plan_test.go b/go/vt/vtgate/planbuilder/plan_test.go
@@ -872,3 +872,41 @@ func benchmarkPlanner(b *testing.B, version plancontext.PlannerVersion, testCase
 		}
 	}
 }
+
+func (s *planTestSuite) TestMy() {
+	vschema, err := vschemawrapper.NewVschemaWrapper(
+		vtenv.NewTestEnv(),
+		loadSchema(s.T(), "vschemas/my_schema.json", true),
+		TestBuilder,
+	)
+	require.NoError(s.T(), err)
+	// vschema := &vschemawrapper.VSchemaWrapper{
+	// 	V:             loadSchema(s.T(), "vschemas/my_schema.json", true),
+	// 	SysVarEnabled: true,
+	// 	Version:       Gen4,
+	// 	Env:           vtenv.NewTestEnv(),
+	// }
+
+	s.testFile("my.json", vschema, false)
+}
+
+func BenchmarkMine(b *testing.B) {
+	vschema, err := vschemawrapper.NewVschemaWrapper(
+		vtenv.NewTestEnv(),
+		loadSchema(b, "vschemas/my_schema.json", true),
+		TestBuilder,
+	)
+	require.NoError(b, err)
+	testCases := readJSONTests("my.json")
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for range b.N {
+		for _, tcase := range testCases {
+			plan, _ := TestBuilder(tcase.Query, vschema, vschema.CurrentDb())
+			if plan == nil {
+				panic("")
+			}
+		}
+	}
+}
diff --git a/go/vt/vtgate/planbuilder/testdata/my.json b/go/vt/vtgate/planbuilder/testdata/my.json
diff --git a/go/vt/vtgate/planbuilder/testdata/onecase.json b/go/vt/vtgate/planbuilder/testdata/onecase.json
@@ -1,8 +1,32 @@
 [
   {
-    "comment": "Add your test case here for debugging and run go test -run=One.",
-    "query": "",
+    "comment": "sharded correlated subquery inside aggregation function on a sharded table on same vindex",
+    "query": "select max((select max(col2) from user u1 where u1.id = u2.id)) from user u2",
     "plan": {
+      "QueryType": "SELECT",
+      "Original": "select max((select max(col2) from user u1 where u1.id = u2.id)) from user u2",
+      "Instructions": {
+        "OperatorType": "Aggregate",
+        "Variant": "Scalar",
+        "Aggregates": "max(0|1) AS max((select max(col2) from `user` as u1 where u1.id = u2.id))",
+        "ResultColumns": 1,
+        "Inputs": [
+          {
+            "OperatorType": "Route",
+            "Variant": "Scatter",
+            "Keyspace": {
+              "Name": "user",
+              "Sharded": true
+            },
+            "FieldQuery": "select max((select max(col2) from `user` as u1 where 1 != 1)), weight_string(max((select max(col2) from `user` as u1 where 1 != 1))) from `user` as u2 where 1 != 1",
+            "Query": "select max((select max(col2) from `user` as u1 where u1.id = u2.id)), weight_string(max((select max(col2) from `user` as u1 where u1.id = u2.id))) from `user` as u2",
+            "Table": "`user`"
+          }
+        ]
+      },
+      "TablesUsed": [
+        "user.user"
+      ]
     }
   }
 ]
diff --git a/go/vt/vtgate/planbuilder/testdata/vschemas/my_schema.json b/go/vt/vtgate/planbuilder/testdata/vschemas/my_schema.json
@@ -0,0 +1,110 @@
+{
+  "keyspaces": {
+    "unsharded_ns": {
+      "sharded": false,
+      "tables": {
+        "my_table_seq": {
+          "type": "sequence",
+          "columns": [
+            {
+              "name": "next_id",
+              "type": "UINT64"
+            },
+            {
+              "name": "cache",
+              "type": "UINT32"
+            },
+            {
+              "name": "id",
+              "type": "UINT64"
+            }
+          ],
+          "column_list_authoritative": true
+        }
+      }
+    },
+    "main": {
+      "sharded": true,
+      "vindexes": {
+        "binary_vdx": {
+          "type": "binary"
+        }
+      },
+      "tables": {
+        "my_table": {
+          "column_vindexes": [
+            {
+              "columns": [
+                "sdKey"
+              ],
+              "name": "binary_vdx"
+            }
+          ],
+          "auto_increment": {
+            "column": "id",
+            "sequence": "my_table_seq"
+          },
+          "columns": [
+            {
+              "name": "col01",
+              "type": "VARCHAR"
+            },
+            {
+              "name": "col02",
+              "type": "INT64"
+            },
+            {
+              "name": "col03",
+              "type": "BLOB"
+            },
+            {
+              "name": "col04",
+              "type": "VARCHAR"
+            },
+            {
+              "name": "col05",
+              "type": "VARBINARY"
+            },
+            {
+              "name": "col06",
+              "type": "INT8"
+            },
+            {
+              "name": "col07",
+              "type": "INT64"
+            },
+            {
+              "name": "col08",
+              "type": "VARBINARY"
+            },
+            {
+              "name": "sdKey",
+              "type": "VARBINARY"
+            },
+            {
+              "name": "col09",
+              "type": "INT64"
+            },
+            {
+              "name": "col10",
+              "type": "INT64"
+            },
+            {
+              "name": "col11",
+              "type": "INT64"
+            },
+            {
+              "name": "col12",
+              "type": "INT64"
+            },
+            {
+              "name": "id",
+              "type": "UINT64"
+            }
+          ],
+          "column_list_authoritative": true
+        }
+      }
+    }
+  }
+}
diff --git a/memprofile.out b/memprofile.out
diff --git a/new.txt b/new.txt
@@ -0,0 +1,16 @@
+goos: darwin
+goarch: arm64
+pkg: vitess.io/vitess/go/vt/vtgate/planbuilder
+cpu: Apple M1 Pro
+BenchmarkMine-10    	     226	   5414522 ns/op	 3534348 B/op	   71107 allocs/op
+BenchmarkMine-10    	     242	   4954859 ns/op	 3534127 B/op	   71106 allocs/op
+BenchmarkMine-10    	     240	   4980136 ns/op	 3533776 B/op	   71105 allocs/op
+BenchmarkMine-10    	     235	   5310915 ns/op	 3533834 B/op	   71106 allocs/op
+BenchmarkMine-10    	     241	   5086943 ns/op	 3534022 B/op	   71107 allocs/op
+BenchmarkMine-10    	     242	   4990665 ns/op	 3533751 B/op	   71106 allocs/op
+BenchmarkMine-10    	     241	   5004904 ns/op	 3534083 B/op	   71107 allocs/op
+BenchmarkMine-10    	     242	   5009568 ns/op	 3534540 B/op	   71108 allocs/op
+BenchmarkMine-10    	     235	   5024306 ns/op	 3534362 B/op	   71108 allocs/op
+BenchmarkMine-10    	     237	   5163583 ns/op	 3534147 B/op	   71107 allocs/op
+PASS
+ok  	vitess.io/vitess/go/vt/vtgate/planbuilder	19.387s
diff --git a/old.txt b/old.txt
@@ -0,0 +1,16 @@
+goos: darwin
+goarch: arm64
+pkg: vitess.io/vitess/go/vt/vtgate/planbuilder
+cpu: Apple M1 Pro
+BenchmarkMine-10    	     206	   5468655 ns/op	 3683900 B/op	   74513 allocs/op
+BenchmarkMine-10    	     216	   5452287 ns/op	 3683589 B/op	   74511 allocs/op
+BenchmarkMine-10    	     214	   5510678 ns/op	 3684216 B/op	   74512 allocs/op
+BenchmarkMine-10    	     218	   5426798 ns/op	 3684354 B/op	   74513 allocs/op
+BenchmarkMine-10    	     222	   5427010 ns/op	 3684368 B/op	   74512 allocs/op
+BenchmarkMine-10    	     219	   5461373 ns/op	 3684293 B/op	   74513 allocs/op
+BenchmarkMine-10    	     218	   5544313 ns/op	 3683700 B/op	   74511 allocs/op
+BenchmarkMine-10    	     219	   5446030 ns/op	 3684065 B/op	   74512 allocs/op
+BenchmarkMine-10    	     218	   5513812 ns/op	 3683917 B/op	   74512 allocs/op
+BenchmarkMine-10    	     218	   5471942 ns/op	 3684090 B/op	   74513 allocs/op
+PASS
+ok  	vitess.io/vitess/go/vt/vtgate/planbuilder	18.180s
diff --git a/planbuilder.test b/planbuilder.test
diff --git a/with_cache b/with_cache
@@ -0,0 +1,16 @@
+goos: darwin
+goarch: arm64
+pkg: vitess.io/vitess/go/vt/vtgate/planbuilder
+cpu: Apple M1 Ultra
+BenchmarkMine-20    	     202	   5858150 ns/op	 3687144 B/op	   74512 allocs/op
+BenchmarkMine-20    	     205	   5884388 ns/op	 3687345 B/op	   74512 allocs/op
+BenchmarkMine-20    	     205	   5872149 ns/op	 3687565 B/op	   74513 allocs/op
+BenchmarkMine-20    	     206	   5813034 ns/op	 3687350 B/op	   74512 allocs/op
+BenchmarkMine-20    	     207	   5836727 ns/op	 3687054 B/op	   74512 allocs/op
+BenchmarkMine-20    	     207	   5826613 ns/op	 3687322 B/op	   74512 allocs/op
+BenchmarkMine-20    	     204	   5839730 ns/op	 3687249 B/op	   74512 allocs/op
+BenchmarkMine-20    	     205	   5814595 ns/op	 3687180 B/op	   74511 allocs/op
+BenchmarkMine-20    	     205	   5796803 ns/op	 3687704 B/op	   74513 allocs/op
+BenchmarkMine-20    	     207	   5861450 ns/op	 3687561 B/op	   74513 allocs/op
+PASS
+ok  	vitess.io/vitess/go/vt/vtgate/planbuilder	19.160s