Gen4 init (#147)

* init gen4 plan
jd-opensource · Jan 5, 2024 · 3e1572d · 3e1572d
1 parent fe4d37b
commit 3e1572d
Show file tree

Hide file tree

Showing 213 changed files with 34,251 additions and 866 deletions.
diff --git a/.github/workflows/maven_test.yml b/.github/workflows/maven_test.yml
@@ -138,6 +138,27 @@ jobs:
       - name: Test charEncoding
         run: mvn clean test -Dtest=\
           com.jd.jdbc.engine.CharEncodingTest
+      - name: Test gen4
+        run: mvn clean test -Dtest=\
+          com.jd.jdbc.engine.gen4.ConcatenateGen4EngineTest,\
+          com.jd.jdbc.engine.gen4.DistinctGen4EngineTest,\
+          com.jd.jdbc.engine.gen4.FilterGen4EngineTest,\
+          com.jd.jdbc.engine.gen4.JoinGen4EngineTest,\
+          com.jd.jdbc.engine.gen4.LimitGen4EngineTest,\
+          com.jd.jdbc.engine.gen4.MemorySortGen4EngineTest,\
+          com.jd.jdbc.engine.gen4.OrderedAggregateGen4EngineTest,\
+          com.jd.jdbc.engine.gen4.RouteGen4EngineTest,\
+          com.jd.jdbc.engine.gen4.ScalarAggregateGen4EngineTest,\
+          com.jd.jdbc.engine.gen4.VitessCompareTest,\
+          com.jd.jdbc.planbuilder.Gen4PlanTest,\
+          com.jd.jdbc.planbuilder.gen4.AnalyzerTest,\
+          com.jd.jdbc.planbuilder.gen4.HorizonPlanningTest,\
+          com.jd.jdbc.planbuilder.gen4.OperatorTest,\
+          com.jd.jdbc.planbuilder.gen4.QueryProjectionTest,\
+          com.jd.jdbc.planbuilder.gen4.RewriterTest,\
+          com.jd.jdbc.planbuilder.gen4.operator.physical.RoutePlanningTest,\
+          com.jd.jdbc.planbuilder.semantics.EarlyRewriterTest,\
+          com.jd.jdbc.planbuilder.semantics.TableSetTest
       - name: install vtdriver
         run: mvn clean install -Dmaven.test.skip=true
 #      - name: Test starter

diff --git a/.gitignore b/.gitignore
@@ -170,3 +170,8 @@ hs_err_pid*
 
 /src/test/resources/vitess_env/build_vitess/
 dependency-reduced-pom.xml
+
+
+# vscode
+
+.vscode
diff --git a/src/main/java/com/jd/jdbc/common/Constant.java b/src/main/java/com/jd/jdbc/common/Constant.java
@@ -45,4 +45,5 @@ public class Constant {
 
     public static final String DEFAULT_SPLIT_TABLE_CONFIG_PATH = "vtdriver-split-table.yml";
 
+    public static final String GEN4_PLAN_ENABLE = "vtdriver.gen4plan.enable";
 }
diff --git a/src/main/java/com/jd/jdbc/context/PlanningContext.java b/src/main/java/com/jd/jdbc/context/PlanningContext.java
@@ -0,0 +1,61 @@
+/*
+Copyright 2023 JD Project Authors. Licensed under Apache-2.0.
+
+Copyright 2022 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package com.jd.jdbc.context;
+
+import com.jd.jdbc.planbuilder.semantics.SemTable;
+import com.jd.jdbc.planbuilder.semantics.VSchema;
+import com.jd.jdbc.sqlparser.ast.SQLExpr;
+import com.jd.jdbc.vindexes.VKeyspace;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import lombok.Getter;
+
+@Getter
+public class PlanningContext {
+
+    private Object reservedVars;
+
+    private SemTable semTable;
+
+    private VKeyspace keyspace;
+
+    private VSchema vschema;
+
+    // here we add all predicates that were created because of a join condition
+    // e.g. [FROM tblA JOIN tblB ON a.colA = b.colB] will be rewritten to [FROM tblB WHERE :a_colA = b.colB],
+    // if we assume that tblB is on the RHS of the join. This last predicate in the WHERE clause is added to the
+    // map below
+
+    private Map<SQLExpr, List<SQLExpr>> joinPredicates;
+
+    private Map<SQLExpr, Object> skipPredicates;
+
+    private boolean rewriteDerivedExpr;
+
+    public PlanningContext(Object reservedVars, SemTable semTable, VSchema vschema, VKeyspace keyspace) {
+        this.reservedVars = reservedVars;
+        this.semTable = semTable;
+        this.vschema = vschema;
+        this.keyspace = keyspace;
+        this.skipPredicates = new HashMap<>();
+        this.joinPredicates = new HashMap<>();
+        this.rewriteDerivedExpr = false;
+    }
+}
diff --git a/src/main/java/com/jd/jdbc/engine/Engine.java b/src/main/java/com/jd/jdbc/engine/Engine.java
@@ -20,6 +20,19 @@
 
 import com.google.common.collect.Lists;
 import com.jd.jdbc.IExecute;
+import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateAvg;
+import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateAvgCount;
+import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateAvgSum;
+import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateCount;
+import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateCountDistinct;
+import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateMax;
+import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateMin;
+import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateSum;
+import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateSumDistinct;
+import static com.jd.jdbc.engine.Engine.PulloutOpcode.PulloutExists;
+import static com.jd.jdbc.engine.Engine.PulloutOpcode.PulloutIn;
+import static com.jd.jdbc.engine.Engine.PulloutOpcode.PulloutNotIn;
+import static com.jd.jdbc.engine.Engine.PulloutOpcode.PulloutValue;
 import com.jd.jdbc.key.Destination;
 import com.jd.jdbc.sqlparser.ast.SQLObject;
 import com.jd.jdbc.sqlparser.dialect.mysql.visitor.VtRestoreVisitor;
@@ -43,20 +56,6 @@
 import java.util.stream.Collectors;
 import lombok.Getter;
 
-import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateAvg;
-import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateAvgCount;
-import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateAvgSum;
-import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateCount;
-import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateCountDistinct;
-import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateMax;
-import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateMin;
-import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateSum;
-import static com.jd.jdbc.engine.Engine.AggregateOpcode.AggregateSumDistinct;
-import static com.jd.jdbc.engine.Engine.PulloutOpcode.PulloutExists;
-import static com.jd.jdbc.engine.Engine.PulloutOpcode.PulloutIn;
-import static com.jd.jdbc.engine.Engine.PulloutOpcode.PulloutNotIn;
-import static com.jd.jdbc.engine.Engine.PulloutOpcode.PulloutValue;
-
 public class Engine {
 
     public static final String LIST_VAR_NAME = "__vals";
@@ -182,6 +181,20 @@ public static List<BoundQuery> getQueries(SQLObject query, List<Map<String, Bind
         return queries;
     }
 
+    public static List<BoundQuery> getQueriesGen4(SQLObject query, List<Map<String, BindVariable>> bindVariableMapList, String charEncoding) throws SQLException {
+        List<BoundQuery> queries = new ArrayList<>(bindVariableMapList.size());
+        for (Map<String, BindVariable> bindVariableMap : bindVariableMapList) {
+            StringBuilder realQueryOutput = new StringBuilder();
+            VtRestoreVisitor vtRestoreVisitor = new VtRestoreVisitor(realQueryOutput, bindVariableMap, charEncoding);
+            query.accept(vtRestoreVisitor);
+            if (vtRestoreVisitor.getException() != null) {
+                throw vtRestoreVisitor.getException();
+            }
+            queries.add(new BoundQuery(realQueryOutput.toString(), bindVariableMap));
+        }
+        return queries;
+    }
+
     /**
      * @param query
      * @param bindVariableMapList
@@ -289,10 +302,17 @@ public enum RouteOpcode {
          * SelectNone is used for queries that always return empty values
          */
         SelectNone(8),
+
+        /**
+         * ByDestination is to route explicitly to a given target destination.
+         * Is used when the query explicitly sets a target destination:
+         * in the clause e.g: UPDATE `keyspace[-]`.x1 SET foo=1
+         */
+        ByDestination(9),
         /**
          * NumRouteOpcodes is the number of opcodes
          */
-        NumRouteOpcodes(9);
+        NumRouteOpcodes(10);
 
         @Getter
         private final Integer value;
@@ -363,7 +383,8 @@ public enum AggregateOpcode {
         AggregateSumDistinct(5),
         AggregateAvg(6),
         AggregateAvgSum(7),
-        AggregateAvgCount(8);
+        AggregateAvgCount(8),
+        AggregateRandom(9);
 
         @Getter
         private final Integer value;
@@ -373,6 +394,31 @@ public enum AggregateOpcode {
         }
     }
 
+
+    /**
+     * These constants list the possible aggregate opcodes.
+     */
+    public enum AggregateOpcodeG4 {
+        /***/
+        AggregateUnassigned(0),
+        AggregateCount(1),
+        AggregateSum(2),
+        AggregateMin(3),
+        AggregateMax(4),
+        AggregateCountDistinct(5),
+        AggregateSumDistinct(6),
+        AggregateGtid(7),
+        AggregateRandom(8),
+        AggregateCountStar(9);
+
+        @Getter
+        private final Integer value;
+
+        AggregateOpcodeG4(Integer value) {
+            this.value = value;
+        }
+    }
+
     /**
      * This is the list of InsertOpcode values.
      */

diff --git a/src/main/java/com/jd/jdbc/engine/OrderedAggregateEngine.java b/src/main/java/com/jd/jdbc/engine/OrderedAggregateEngine.java
@@ -302,7 +302,6 @@ private IExecute.ExecuteMultiShardResponse orderedAggregateExecute(IContext ctx,
      * @throws SQLException
      */
     private rowProcessResponse merge(Query.Field[] fields, List<VtResultValue> row1, List<VtResultValue> row2, VtResultValue curDistinct) throws SQLException {
-        // result := sqltypes.CopyRow(row1)
         List<VtResultValue> result = new ArrayList<>(row1);
         for (AggregateParams aggr : this.aggregateParamsList) {
             if (aggr.isDistinct()) {

diff --git a/src/main/java/com/jd/jdbc/engine/ProjectionEngine.java b/src/main/java/com/jd/jdbc/engine/ProjectionEngine.java
@@ -21,6 +21,7 @@
 import com.jd.jdbc.IExecute;
 import com.jd.jdbc.context.IContext;
 import com.jd.jdbc.evalengine.EvalEngine;
+import com.jd.jdbc.evalengine.EvalResult;
 import com.jd.jdbc.sqlparser.utils.StringUtils;
 import com.jd.jdbc.sqltypes.VtResultSet;
 import com.jd.jdbc.sqltypes.VtResultValue;
@@ -122,7 +123,7 @@ public VtRowList fetch(boolean wantFields) throws SQLException {
                 for (List<VtResultValue> row : vtResultSet.getRows()) {
                     env.setRow(row);
                     for (EvalEngine.Expr expr : exprs) {
-                        EvalEngine.EvalResult res = expr.evaluate(env);
+                        EvalResult res = expr.evaluate(env);
                         row.add(convertToVtResultValue(res, vcursor.getCharEncoding()));
                     }
                     rows.add(row);
@@ -232,7 +233,7 @@ private IExecute.ExecuteMultiShardResponse getExecuteMultiShardResponse(VtResult
         for (List<VtResultValue> row : resultSet.getRows()) {
             env.setRow(row);
             for (EvalEngine.Expr expr : this.exprs) {
-                EvalEngine.EvalResult res = expr.evaluate(env);
+                EvalResult res = expr.evaluate(env);
                 row.add(convertToVtResultValue(res, charEncoding));
             }
             rows.add(row);
@@ -241,15 +242,15 @@ private IExecute.ExecuteMultiShardResponse getExecuteMultiShardResponse(VtResult
         return new IExecute.ExecuteMultiShardResponse(resultSet);
     }
 
-    private static VtResultValue convertToVtResultValue(EvalEngine.EvalResult res, String charEncoding) throws SQLException {
+    private static VtResultValue convertToVtResultValue(EvalResult res, String charEncoding) throws SQLException {
         VtResultValue resultValue;
         switch (res.getType()) {
             case FLOAT64:
-                EvalEngine.EvalResult evalResult1 = new EvalEngine.EvalResult(BigDecimal.valueOf(res.getFval()).setScale(4, RoundingMode.HALF_UP), Query.Type.DECIMAL);
+                EvalResult evalResult1 = new EvalResult(BigDecimal.valueOf(res.getFval()).setScale(4, RoundingMode.HALF_UP), Query.Type.DECIMAL);
                 resultValue = evalResult1.resultValue();
                 break;
             case VARBINARY:
-                EvalEngine.EvalResult evalResult2 = new EvalEngine.EvalResult(res.getBytes(), Query.Type.VARBINARY);
+                EvalResult evalResult2 = new EvalResult(res.getBytes(), Query.Type.VARBINARY);
                 Charset cs = StringUtils.isEmpty(charEncoding) ? Charset.defaultCharset() : Charset.forName(charEncoding);
                 resultValue = VtResultValue.newVtResultValue(Query.Type.VARBINARY, new String(evalResult2.getBytes(), cs));
                 break;

diff --git a/src/main/java/com/jd/jdbc/engine/RouteEngine.java b/src/main/java/com/jd/jdbc/engine/RouteEngine.java
@@ -22,6 +22,7 @@
 import com.jd.jdbc.IExecute.ExecuteMultiShardResponse;
 import com.jd.jdbc.context.IContext;
 import com.jd.jdbc.evalengine.EvalEngine;
+import com.jd.jdbc.evalengine.EvalResult;
 import com.jd.jdbc.key.Destination;
 import com.jd.jdbc.key.DestinationAllShard;
 import com.jd.jdbc.key.DestinationAnyShard;
@@ -336,7 +337,7 @@ public ParamsResponse paramsSystemQuery(Vcursor vcursor, Map<String, BindVariabl
         boolean schemaExists = false;
         EvalEngine.ExpressionEnv env = new EvalEngine.ExpressionEnv(bindVariableMap);
         for (EvalEngine.Expr expr : this.sysTableKeyspaceExpr) {
-            EvalEngine.EvalResult evalResult = expr.evaluate(env);
+            EvalResult evalResult = expr.evaluate(env);
             String other = evalResult.value().toString();
             if (StringUtil.isNullOrEmpty(keyspace)) {
                 keyspace = other;