探究Presto SQL引擎(3)-代码生成

querySpecification    : SELECT selectItem (',' selectItem)*      (FROM relation (',' relation)*)?      (WHERE where=booleanExpression)?    ;... booleanExpression    : valueExpression predicate[$valueExpression.ctx]?             #predicated    | NOT booleanExpression                                        #logicalNot    | left=booleanExpression operator=AND right=booleanExpression  #logicalBinary    | left=booleanExpression operator=OR right=booleanExpression   #logicalBinary    ; predicate[ParserRuleContext value]    : comparisonOperator right=valueExpression                            #comparison    ;

antlr4 -package org.example.antlr -no-listener -visitor .\SqlBase.g4

@Overridepublic Node visitComparison(Select1Parser.ComparisonContext context){    return new ComparisonExpression(            getLocation(context.comparisonOperator()),            getComparisonOperator(((TerminalNode) context.comparisonOperator().getChild(0)).getSymbol()),            (Expression) visit(context.value),            (Expression) visit(context.right));}  @Overridepublic Node visitLogicalBinary(Select1Parser.LogicalBinaryContext context){    return new LogicalBinaryExpression(            getLocation(context.operator),            getLogicalBinaryOperator(context.operator),            (Expression) visit(context.left),            (Expression) visit(context.right));}

public abstract class Node{        /**     * Accessible for {@link AstVisitor}, use {@link AstVisitor#process(Node, Object)} instead.     */    protected <R, C> R accept(AstVisitor<R, C> visitor, C context)    {        return visitor.visitNode(this, context);    } } public abstract class AstVisitor<R, C>{        protected R visitStatement(Statement node, C context)    {        return visitNode(node, context);    }     protected R visitQuery(Query node, C context)    {        return visitStatement(node, context);    }     ....}

querySpecification    : SELECT setQuantifier? selectItem (',' selectItem)*      (FROM relation (',' relation)*)?      (WHERE where=booleanExpression)?      (GROUP BY groupBy)?      (HAVING having=booleanExpression)?    ;

@Overrideprotected Void visitQuerySpecification(QuerySpecification node, Integer indent){    process(node.getSelect(), indent);     if (node.getFrom().isPresent()) {        append(indent, "FROM");        builder.append('\n');        append(indent, "  ");        process(node.getFrom().get(), indent);    }     builder.append('\n');     if (node.getWhere().isPresent()) {        append(indent, "WHERE " + formatExpression(node.getWhere().get(), parameters))                .append('\n');    }     if (node.getGroupBy().isPresent()) {        append(indent, "GROUP BY " + (node.getGroupBy().get().isDistinct() ? " DISTINCT " : "") + formatGroupBy(node.getGroupBy().get().getGroupingElements())).append('\n');    }     if (node.getHaving().isPresent()) {        append(indent, "HAVING " + formatExpression(node.getHaving().get(), parameters))                .append('\n');    }     if (node.getOrderBy().isPresent()) {        process(node.getOrderBy().get(), indent);    }     if (node.getLimit().isPresent()) {        append(indent, "LIMIT " + node.getLimit().get())                .append('\n');    }    return null;}

/** * 处理比较表达式 * @param node * @param context * @return */@Overrideprotected Void visitComparisonExpression(ComparisonExpression node, Map<String,Long> context) {     Expression left = node.getLeft();         Expression right = node.getRight();         String leftKey  = ((Identifier) left).getValue();    Long rightKey = ((LongLiteral) right).getValue();    Long leftVal = context.get(leftKey);    if(leftVal == null){       stack.push(false);    }     ComparisonExpression.Operator op = node.getOperator();    switch (op){        case EQUAL:            stack.push(leftVal.equals(rightKey));break;        case LESS_THAN:            stack.push( leftVal < rightKey);;break;        case NOT_EQUAL:            stack.push( !leftVal.equals(rightKey));break;        case GREATER_THAN:            stack.push( leftVal>rightKey);break;        case LESS_THAN_OR_EQUAL:            stack.push( leftVal<=rightKey);break;        case GREATER_THAN_OR_EQUAL:            stack.push( leftVal>=rightKey);break;        case IS_DISTINCT_FROM:        default:            throw new UnsupportedOperationException("not supported");    }    return null;}

// antlr处理表达式语句，生成Expression对象SqlParser sqlParser = new SqlParser();Expression expression = sqlParser.createExpression("a>1 and b<2");// 基于AstVisitor实现WhereExpFilter rowFilter = new WhereExpFilter(expression);Random r = new Random();for(int i=;i<10;i++){    Map<String,Long> row = new HashMap<>();    row.put("a", (long) r.nextInt(10));    row.put("b", (long) r.nextInt(10));    System.out.println("exp: a>1 and b<2, param:"+row+", ret:"+rowFilter.filter(row));}// ====== 执行结果如下/**exp: a>1 and b<2, param:{a=9, b=8}, ret:falseexp: a>1 and b<2, param:{a=7, b=3}, ret:falseexp: a>1 and b<2, param:{a=, b=7}, ret:falseexp: a>1 and b<2, param:{a=6, b=}, ret:trueexp: a>1 and b<2, param:{a=2, b=}, ret:trueexp: a>1 and b<2, param:{a=9, b=}, ret:trueexp: a>1 and b<2, param:{a=3, b=6}, ret:falseexp: a>1 and b<2, param:{a=8, b=7}, ret:falseexp: a>1 and b<2, param:{a=6, b=1}, ret:trueexp: a>1 and b<2, param:{a=4, b=6}, ret:false*/

public class TestInline {     public int addOp(int a,int b){        return a+b;    }     @Benchmark    public int testAdd(){        int sum=;        for(int i=;i<100000;i++){            sum=addOp(sum,i);        }        return sum;    }         public static void main(String[] args) throws RunnerException {        Options options = new OptionsBuilder()                .warmupIterations(2).measurementIterations(2)                .forks(1).build();        new Runner(options).run();    }}// 执行结果如下：/**Benchmark            Mode  Cnt      Score   Error  UnitsTestInline.testAdd  thrpt    2  18588.318          ops/s(正常执行)TestInline.testAdd  thrpt    2  3131.466          ops/s(禁用内联)**/

private static final AtomicLong CLASS_ID = new AtomicLong();private static final DateTimeFormatter TIMESTAMP_FORMAT = DateTimeFormatter.ofPattern("YYYYMMdd_HHmmss"); private String clazzName;private ClassDefinition classDefinition; public ByteCodeGenDemo(String clazzName){    this.clazzName=clazzName;} public static ParameterizedType makeClassName(String baseName, Optional<String> suffix){    String className = baseName            + "_" + suffix.orElseGet(() -> Instant.now().atZone(UTC).format(TIMESTAMP_FORMAT))            + "_" + CLASS_ID.incrementAndGet();    return typeFromJavaClassName("org.shgy.demo.$gen." + toJavaIdentifierString(className));}public void buildClass(){    ClassDefinition classDefinition = new ClassDefinition(            a(PUBLIC, FINAL),            makeClassName(clazzName,Optional.empty()),            type(Object.class));    this.classDefinition=classDefinition;}

/** * 生成if分支代码 * if(a<){ *     System.out.println(a +" a<0"); * }else{ *     System.out.println(a +" a>=0"); * } * @param methodName */public void buildMethod1(String methodName){    Parameter argA = arg("a", int.class);    MethodDefinition method = classDefinition.declareMethod(            a(PUBLIC, STATIC),            methodName,            type(void.class),            ImmutableList.of(argA));         BytecodeExpression out = getStatic(System.class, "out");         IfStatement ifStatement = new IfStatement();    ifStatement.condition(lessThan(argA,constantInt()))            .ifTrue(new BytecodeBlock()                    .append(out.invoke("print", void.class, argA))                    .append(out.invoke("println", void.class, constantString(" a<0")))            )            .ifFalse(new BytecodeBlock()                    .append(out.invoke("print", void.class, argA))                    .append(out.invoke("println", void.class, constantString(" a>=0")))            );     method.getBody().append(ifStatement).ret();}

/** * 生成switch分支代码 *         switch (a){ *             case 1: *                 System.out.println("a=1"); *                 break; *             case 2: *                 System.out.println("a=2"); *                 break; *             default: *                 System.out.println("a=others"); *         } * @param methodName */public void buildMethod2(String methodName){    Parameter argA = arg("a", int.class);    MethodDefinition method = classDefinition.declareMethod(            a(PUBLIC, STATIC),            methodName,            type(void.class),            ImmutableList.of(argA));     SwitchStatement.SwitchBuilder switchBuilder = new SwitchStatement.SwitchBuilder().expression(argA);    switchBuilder.addCase(1, BytecodeExpressions.print(BytecodeExpressions.constantString("a=1")));    switchBuilder.addCase(2,BytecodeExpressions.print(BytecodeExpressions.constantString("a=2")));    switchBuilder.defaultCase(invokeStatic(ByteCodeGenDemo.class,"defaultCase", void.class));     method.getBody().append(switchBuilder.build()).ret();}public static void defaultCase(){    System.out.println("a=others");}

/**     * 生成循环逻辑代码     * int sum=;     * for(int i=s;i<=e;i++){     *     sum+=i;     *     System.out.println("i="+i+",sum="+sum);     * }     * @param methodName     */    public void buildMethodLoop(String methodName){        Parameter argS = arg("s", int.class);        Parameter argE = arg("e", int.class);         MethodDefinition method = classDefinition.declareMethod(                a(PUBLIC, STATIC),                methodName,                type(int.class),                ImmutableList.of(argS, argE));         Scope scope = method.getScope();        Variable i = scope.declareVariable(int.class,"i");        Variable sum = scope.declareVariable(int.class,"sum");         BytecodeExpression out = getStatic(System.class, "out");         ForLoop loop = new ForLoop()                .initialize(i.set(argS))                .condition(lessThanOrEqual(i, argE))                .update(incrementVariable(i,(byte)1))                .body(new BytecodeBlock()                        .append(sum.set(add(sum,i)))                        .append(out.invoke("print", void.class, constantString("i=")))                        .append(out.invoke("print", void.class, i))                        .append(out.invoke("print", void.class, constantString(",sum=")))                        .append(out.invoke("println", void.class,sum))                );         method.getBody().initializeVariable(i).putVariable(sum,).append(loop).append(sum).retInt();    }

public void executeLoop(String methodName) throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {    // invoke    Class<?> clazz = classGenerator(new DynamicClassLoader(this.getClass().getClassLoader())).defineClass(this.classDefinition,Object.class);     Method loopMethod = clazz.getMethod(methodName, int.class,int.class);    loopMethod.invoke(null,1,10);}

public void buildMapGetter(String methodName){       Parameter argRow = arg("row", Map.class);        MethodDefinition method = classDefinition.declareMethod(               a(PUBLIC, STATIC),               methodName,               type(void.class),               of(argRow));       BytecodeExpression out = getStatic(System.class, "out");       Scope scope = method.getScope();       Variable a = scope.declareVariable(int.class,"a");       // 从map中获取key=aa对应的值       method.getBody().append(out.invoke("print", void.class, argRow.invoke("get",Object.class,constantString("aa").cast(Object.class)))).ret();    }   // 代码执行   public void executeMapOp(String methodName) throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {       // invoke       Class<?> clazz = classGenerator(new DynamicClassLoader(this.getClass().getClassLoader())).defineClass(this.classDefinition,Object.class);        Method loopMethod = clazz.getMethod(methodName, Map.class);       Map<String,Integer> map = Maps.newHashMap();       map.put("aa",111);       loopMethod.invoke(null,map);   }

/** *   生成比较条件语句**/@Overrideprotected Void visitComparisonExpression(ComparisonExpression node, MethodDefinition context) {     ComparisonExpression.Operator op = node.getOperator();    Expression left = node.getLeft();    Expression right = node.getRight();     if(left instanceof Identifier && right instanceof LongLiteral){        String leftKey  = ((Identifier) left).getValue();        Long rightKey = ((LongLiteral) right).getValue();         Parameter argRow = context.getParameters().get();        Variable stack = context.getScope().getVariable("stack");        BytecodeBlock body = context.getBody();         BytecodeExpression leftVal = argRow.invoke("get", Object.class,constantString(leftKey).cast(Object.class)).cast(long.class);        BytecodeExpression cResult;        switch (op){            case EQUAL:                cResult = equal(leftVal,constantLong(rightKey));                break;            case LESS_THAN:                cResult = lessThan(leftVal,constantLong(rightKey));                break;            case GREATER_THAN:                cResult =greaterThan(leftVal,constantLong(rightKey));                break;            case NOT_EQUAL:                cResult = notEqual(leftVal,constantLong(rightKey));                break;            case LESS_THAN_OR_EQUAL:                cResult = lessThanOrEqual(leftVal,constantLong(rightKey));                break;            case GREATER_THAN_OR_EQUAL:                cResult = greaterThanOrEqual(leftVal,constantLong(rightKey));                break;            default:                throw new UnsupportedOperationException("not implemented");        }        body.append(stack.invoke("push",Object.class, cResult.cast(Object.class)));         return null;     }else{        throw new UnsupportedOperationException("not implemented");    }}

/**    *   生成比较and or语句   **/   @Override   protected Void visitLogicalBinaryExpression(LogicalBinaryExpression node, MethodDefinition context)   {       Expression left = node.getLeft();       Expression right = node.getRight();       // 处理左子树       if(left instanceof ComparisonExpression){           visitComparisonExpression((ComparisonExpression)left,context);       }else if(left instanceof LogicalBinaryExpression){           visitLogicalBinaryExpression((LogicalBinaryExpression) left,context);       }else{           throw new UnsupportedOperationException("not implemented");       }       // 处理右子树       if(right instanceof ComparisonExpression){           visitComparisonExpression((ComparisonExpression)right,context);       }else if(right instanceof LogicalBinaryExpression){           visitLogicalBinaryExpression((LogicalBinaryExpression) right,context);       }else{           throw new UnsupportedOperationException("not implemented");       }        // 根据运算符生成字节码       LogicalBinaryExpression.Operator op = node.getOperator();        Variable stack = context.getScope().getVariable("stack");       BytecodeExpression result;       switch (op){           case AND:               result = stack.invoke("push",Object.class,                       and(                               stack.invoke("pop",Object.class).cast(boolean.class),                               stack.invoke("pop",Object.class).cast(boolean.class)               ).cast(Object.class));               break;           case OR:               result = stack.invoke("push",Object.class,                       or(                               stack.invoke("pop",Object.class).cast(boolean.class),                               stack.invoke("pop",Object.class).cast(boolean.class)                       ).cast(Object.class));               break;           default:               throw  new UnsupportedOperationException("not implemented");       }       context.getBody().append(result);       return null;   }

@BenchmarkMode(Mode.Throughput)@Fork(1)@State(value = Scope.Benchmark)public class RowFilterBenchmark {     private RowFilter filter1;    private RowFilter filter2;     private List<Map<String,Long>> dataSet = Lists.newArrayListWithCapacity(100000);    @Setup    public void init(){        // antlr处理表达式语句，生成Expression对象        SqlParser sqlParser = new SqlParser();        Expression expression = sqlParser.createExpression("a>5 and b<5");        // 基于AstVisitor实现        this.filter1 = new WhereExpFilter(expression);        // 基于AstVisitor实现        ExpressionCodeCompiler compiler = new ExpressionCodeCompiler();        Class clazz = compiler.compile(expression);        try {            this.filter2 = (RowFilter) clazz.newInstance();        } catch (InstantiationException e) {            e.printStackTrace();        } catch (IllegalAccessException e) {            e.printStackTrace();        }        Random r = new Random();        for(int i=;i<100000;i++){            Map<String,Long> row = new HashMap<>();            row.put("a", (long) r.nextInt(10));            row.put("b", (long) r.nextInt(10));            dataSet.add(row);        }    }     @Benchmark    public int testAstDirect() {        int cnt =;        for(Map<String,Long> row:dataSet){            boolean ret = filter1.filter(row);            if(ret){                cnt++;            }        }        return cnt;    }     @Benchmark    public int testAstCompile() {        int cnt =;        for(Map<String,Long> row:dataSet){            boolean ret = filter2.filter(row);            if(ret){                cnt++;            }        }        return cnt;    }     public static void main(String[] args) throws RunnerException {        Options opt = new OptionsBuilder()                .include(RowFilterBenchmark.class.getSimpleName())                .build();         new Runner(opt).run();    }}

Benchmark                           Mode  Cnt    Score    Error  UnitsRowFilterBenchmark.testAstCompile  thrpt    5  211.298 ± 30.832  ops/sRowFilterBenchmark.testAstDirect   thrpt    5   62.254 ±  8.269  ops/s