Apache Calcite VolcanoPlanner源码学习

忘是亡心i 2022-11-18 13:58 138阅读 0赞

import org.apache.calcite.adapter.enumerable.EnumerableConvention;
    import org.apache.calcite.adapter.enumerable.EnumerableRules;
    import org.apache.calcite.config.CalciteConnectionConfig;
    import org.apache.calcite.config.CalciteConnectionConfigImpl;
    import org.apache.calcite.jdbc.CalciteSchema;
    import org.apache.calcite.plan.*;
    import org.apache.calcite.plan.volcano.VolcanoPlanner;
    import org.apache.calcite.prepare.CalciteCatalogReader;
    import org.apache.calcite.rel.RelDistributionTraitDef;
    import org.apache.calcite.rel.RelNode;
    import org.apache.calcite.rel.RelRoot;
    import org.apache.calcite.rel.rules.FilterJoinRule;
    import org.apache.calcite.rel.rules.PruneEmptyRules;
    import org.apache.calcite.rel.rules.ReduceExpressionsRule;
    import org.apache.calcite.rel.type.RelDataType;
    import org.apache.calcite.rel.type.RelDataTypeFactory;
    import org.apache.calcite.rel.type.RelDataTypeSystem;
    import org.apache.calcite.rel.type.RelDataTypeSystemImpl;
    import org.apache.calcite.rex.RexBuilder;
    import org.apache.calcite.schema.SchemaPlus;
    import org.apache.calcite.schema.impl.AbstractTable;
    import org.apache.calcite.sql.SqlNode;
    import org.apache.calcite.sql.fun.SqlStdOperatorTable;
    import org.apache.calcite.sql.parser.SqlParser;
    import org.apache.calcite.sql.type.BasicSqlType;
    import org.apache.calcite.sql.type.SqlTypeFactoryImpl;
    import org.apache.calcite.sql.type.SqlTypeName;
    import org.apache.calcite.sql.validate.SqlConformance;
    import org.apache.calcite.sql.validate.SqlConformanceEnum;
    import org.apache.calcite.sql.validate.SqlValidator;
    import org.apache.calcite.sql.validate.SqlValidatorUtil;
    import org.apache.calcite.sql2rel.RelDecorrelator;
    import org.apache.calcite.sql2rel.SqlToRelConverter;
    import org.apache.calcite.tools.FrameworkConfig;
    import org.apache.calcite.tools.Frameworks;
    import org.apache.calcite.tools.RelBuilder;
    import org.slf4j.Logger;
    import org.slf4j.LoggerFactory;
    
    import java.util.List;
    import java.util.Properties;
    
    
    /**
     * @Author: wpp
     * @Date: 2021/4/13 17:05
     */
    // SqlVolcanoTest.java
    public class SqlVolcanoTest {
        private static final Logger LOGGER = LoggerFactory.getLogger(SqlVolcanoTest.class);
        private static Object FilterJoinRule;
    
        public static void main(String[] args) {
            SchemaPlus rootSchema = CalciteUtils.registerRootSchema();
    
            final FrameworkConfig frameworkConfig = Frameworks.newConfigBuilder()
                    .parserConfig(SqlParser.Config.DEFAULT)
                    .defaultSchema(rootSchema)
                    .traitDefs(ConventionTraitDef.INSTANCE, RelDistributionTraitDef.INSTANCE)
                    .build();
    
            String sql
                    = "select u.id as user_id, u.name as user_name, j.company as user_company, u.age as user_age from users u"
                    + " join jobs j on u.id=j.id where u.age > 30 and j.id>10 order by user_id";
    
            // use HepPlanner
            VolcanoPlanner planner = new VolcanoPlanner();
            planner.addRelTraitDef(ConventionTraitDef.INSTANCE);
            planner.addRelTraitDef(RelDistributionTraitDef.INSTANCE);
            // add rules
    //        planner.addRule(FilterJoinRule.FilterIntoJoinRule.FILTER_ON_JOIN);
            planner.addRule(ReduceExpressionsRule.PROJECT_INSTANCE);
            planner.addRule(PruneEmptyRules.PROJECT_INSTANCE);
            // add ConverterRule
            planner.addRule(EnumerableRules.ENUMERABLE_MERGE_JOIN_RULE);
            planner.addRule(EnumerableRules.ENUMERABLE_SORT_RULE);
            planner.addRule(EnumerableRules.ENUMERABLE_VALUES_RULE);
            planner.addRule(EnumerableRules.ENUMERABLE_PROJECT_RULE);
            planner.addRule(EnumerableRules.ENUMERABLE_FILTER_RULE);
    
            try {
                SqlTypeFactoryImpl factory = new SqlTypeFactoryImpl(RelDataTypeSystem.DEFAULT);
                // sql parser
                SqlParser parser = SqlParser.create(sql, SqlParser.Config.DEFAULT);
                SqlNode parsed = parser.parseStmt();
                LOGGER.info("The SqlNode after parsed is:\n{}", parsed.toString());
    
                CalciteCatalogReader calciteCatalogReader = new CalciteCatalogReader(
                        CalciteSchema.from(rootSchema),
                        CalciteSchema.from(rootSchema).path(null),
                        factory,
                        new CalciteConnectionConfigImpl(new Properties()));
    
                // sql validate
                SqlValidator validator = SqlValidatorUtil.newValidator(SqlStdOperatorTable.instance(), calciteCatalogReader,
                        factory, CalciteUtils.conformance(frameworkConfig));
                SqlNode validated = validator.validate(parsed);
                LOGGER.info("The SqlNode after validated is:\n{}", validated.toString());
    
                final RexBuilder rexBuilder = CalciteUtils.createRexBuilder(factory);
                final RelOptCluster cluster = RelOptCluster.create(planner, rexBuilder);
    
                // init SqlToRelConverter config
                final SqlToRelConverter.Config config = SqlToRelConverter.configBuilder()
                        .withConfig(frameworkConfig.getSqlToRelConverterConfig())
                        .withTrimUnusedFields(false)
                        .withConvertTableAccess(false)
                        .build();
                // SqlNode toRelNode
                final SqlToRelConverter sqlToRelConverter = new SqlToRelConverter(new CalciteUtils.ViewExpanderImpl(),
                        validator, calciteCatalogReader, cluster, frameworkConfig.getConvertletTable(), config);
                RelRoot root = sqlToRelConverter.convertQuery(validated, false, true);
    
                root = root.withRel(sqlToRelConverter.flattenTypes(root.rel, true));
                final RelBuilder relBuilder = config.getRelBuilderFactory().create(cluster, null);
                root = root.withRel(RelDecorrelator.decorrelateQuery(root.rel, relBuilder));
                RelNode relNode = root.rel;
                LOGGER.info("The relational expression string before optimized is:\n{}", RelOptUtil.toString(relNode));
    
                RelTraitSet desiredTraits =
                        relNode.getCluster().traitSet().replace(EnumerableConvention.INSTANCE);
                relNode = planner.changeTraits(relNode, desiredTraits);
    
                planner.setRoot(relNode);
                relNode = planner.findBestExp();
                System.out.println("-----------------------------------------------------------");
                System.out.println("The Best relational expression string:");
                System.out.println(RelOptUtil.toString(relNode));
                System.out.println("-----------------------------------------------------------");
    
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }
    
    // CalciteUtils.java
     class CalciteUtils {
    
        public static SchemaPlus registerRootSchema() {
            SchemaPlus rootSchema = Frameworks.createRootSchema(true);
            rootSchema.add("USERS", new AbstractTable() { //note: add a table
                @Override
                public RelDataType getRowType(final RelDataTypeFactory typeFactory) {
                    RelDataTypeFactory.Builder builder = typeFactory.builder();
    
                    builder.add("ID", new BasicSqlType(new RelDataTypeSystemImpl() {}, SqlTypeName.INTEGER));
                    builder.add("NAME", new BasicSqlType(new RelDataTypeSystemImpl() {}, SqlTypeName.CHAR));
                    builder.add("AGE", new BasicSqlType(new RelDataTypeSystemImpl() {}, SqlTypeName.INTEGER));
                    return builder.build();
                }
            });
    
            rootSchema.add("JOBS", new AbstractTable() {
                @Override
                public RelDataType getRowType(final RelDataTypeFactory typeFactory) {
                    RelDataTypeFactory.Builder builder = typeFactory.builder();
    
                    builder.add("ID", new BasicSqlType(new RelDataTypeSystemImpl() {}, SqlTypeName.INTEGER));
                    builder.add("NAME", new BasicSqlType(new RelDataTypeSystemImpl() {}, SqlTypeName.CHAR));
                    builder.add("COMPANY", new BasicSqlType(new RelDataTypeSystemImpl() {}, SqlTypeName.CHAR));
                    return builder.build();
                }
            });
            return rootSchema;
        }
    
        public static SqlConformance conformance(FrameworkConfig config) {
            final Context context = config.getContext();
            if (context != null) {
                final CalciteConnectionConfig connectionConfig =
                        context.unwrap(CalciteConnectionConfig.class);
                if (connectionConfig != null) {
                    return connectionConfig.conformance();
                }
            }
            return SqlConformanceEnum.DEFAULT;
        }
    
        public static RexBuilder createRexBuilder(RelDataTypeFactory typeFactory) {
            return new RexBuilder(typeFactory);
        }
    
        public static class ViewExpanderImpl implements RelOptTable.ViewExpander {
            public ViewExpanderImpl() {
            }
    
            @Override
            public RelRoot expandView(RelDataType rowType, String queryString, List<String> schemaPath,
                                      List<String> viewPath) {
                return null;
            }
        }
    }