content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to you 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 org.apache.calcite.piglet;
  
  import org.apache.calcite.adapter.enumerable.EnumerableConvention;
  import org.apache.calcite.adapter.enumerable.EnumerableRules;
  import org.apache.calcite.plan.RelOptPlanner;
  import org.apache.calcite.plan.RelOptRule;
  import org.apache.calcite.rel.RelCollation;
  import org.apache.calcite.rel.RelCollations;
  import org.apache.calcite.rel.RelNode;
  import org.apache.calcite.rel.core.Sort;
  import org.apache.calcite.rel.logical.ToLogicalConverter;
  import org.apache.calcite.rel.rel2sql.RelToSqlConverter;
  import org.apache.calcite.rel.rules.CoreRules;
  import org.apache.calcite.sql.SqlDialect;
  import org.apache.calcite.sql.SqlNode;
  import org.apache.calcite.sql.SqlWriter;
  import org.apache.calcite.sql.SqlWriterConfig;
  import org.apache.calcite.sql.pretty.SqlPrettyWriter;
  import org.apache.calcite.tools.FrameworkConfig;
  import org.apache.calcite.tools.Program;
  import org.apache.calcite.tools.Programs;
  import org.apache.calcite.tools.RuleSets;
  
  import org.apache.pig.ExecType;
  import org.apache.pig.PigServer;
  import org.apache.pig.impl.logicalLayer.FrontendException;
  import org.apache.pig.impl.util.PropertiesUtil;
  import org.apache.pig.newplan.logical.relational.LogicalPlan;
  
  import com.google.common.collect.ImmutableList;
  
  import java.io.IOException;
  import java.io.InputStream;
  import java.util.ArrayList;
  import java.util.List;
  import java.util.Map;
  import java.util.Properties;
  
  /**
   * Extension from PigServer to convert Pig scripts into logical relational
   * algebra plans and SQL statements.
   */
  public class PigConverter extends PigServer {
    // Basic transformation and implementation rules to optimize for Pig-translated logical plans
    private static final List<RelOptRule> PIG_RULES =
        ImmutableList.of(
            CoreRules.PROJECT_TO_LOGICAL_PROJECT_AND_WINDOW,
            PigToSqlAggregateRule.INSTANCE,
            EnumerableRules.ENUMERABLE_VALUES_RULE,
            EnumerableRules.ENUMERABLE_JOIN_RULE,
            EnumerableRules.ENUMERABLE_CORRELATE_RULE,
            EnumerableRules.ENUMERABLE_PROJECT_RULE,
            EnumerableRules.ENUMERABLE_FILTER_RULE,
            EnumerableRules.ENUMERABLE_AGGREGATE_RULE,
            EnumerableRules.ENUMERABLE_SORT_RULE,
            EnumerableRules.ENUMERABLE_LIMIT_RULE,
            EnumerableRules.ENUMERABLE_COLLECT_RULE,
            EnumerableRules.ENUMERABLE_UNCOLLECT_RULE,
            EnumerableRules.ENUMERABLE_UNION_RULE,
            EnumerableRules.ENUMERABLE_WINDOW_RULE,
            EnumerableRules.ENUMERABLE_TABLE_SCAN_RULE,
            EnumerableRules.TO_INTERPRETER);
  
    private static final List<RelOptRule> TRANSFORM_RULES =
        ImmutableList.of(
            CoreRules.PROJECT_WINDOW_TRANSPOSE,
            CoreRules.FILTER_MERGE,
            CoreRules.PROJECT_MERGE,
            CoreRules.FILTER_PROJECT_TRANSPOSE,
            EnumerableRules.ENUMERABLE_VALUES_RULE,
            EnumerableRules.ENUMERABLE_JOIN_RULE,
            EnumerableRules.ENUMERABLE_CORRELATE_RULE,
            EnumerableRules.ENUMERABLE_PROJECT_RULE,
            EnumerableRules.ENUMERABLE_FILTER_RULE,
            EnumerableRules.ENUMERABLE_AGGREGATE_RULE,
            EnumerableRules.ENUMERABLE_SORT_RULE,
            EnumerableRules.ENUMERABLE_LIMIT_RULE,
            EnumerableRules.ENUMERABLE_COLLECT_RULE,
            EnumerableRules.ENUMERABLE_UNCOLLECT_RULE,
            EnumerableRules.ENUMERABLE_UNION_RULE,
            EnumerableRules.ENUMERABLE_WINDOW_RULE,
            EnumerableRules.ENUMERABLE_TABLE_SCAN_RULE,
            EnumerableRules.TO_INTERPRETER);
  
    private final PigRelBuilder builder;
  
    /** Private constructor. */
    private PigConverter(FrameworkConfig config, ExecType execType,
        Properties properties) throws Exception {
      super(execType, properties);
      this.builder = PigRelBuilder.create(config);
    }
  
    /** Creates a PigConverter using the given property settings. */
    public static PigConverter create(FrameworkConfig config,
        Properties properties) throws Exception {
      return new PigConverter(config, ExecType.LOCAL, properties);
    }
  
    /** Creates a PigConverter using default property settings. */
    public static PigConverter create(FrameworkConfig config) throws Exception {
      return create(config, PropertiesUtil.loadDefaultProperties());
    }
  
    public PigRelBuilder getBuilder() {
      return builder;
    }
  
    /**
     * Parses a Pig script and converts it into relational algebra plans,
     * optimizing the result.
     *
     * <p>Equivalent to {@code pigQuery2Rel(pigQuery, true, true, true)}.
     *
     * @param pigQuery Pig script
     *
     * @return A list of root nodes of the translated relational plans. Each of
     * these root corresponds to a sink operator (normally a STORE command) in the
     * Pig plan
     *
     * @throws IOException Exception during parsing or translating Pig
     */
    public List<RelNode> pigQuery2Rel(String pigQuery) throws IOException {
      return pigQuery2Rel(pigQuery, true, true, true);
    }
  
    /**
     * Parses a Pig script and converts it into relational algebra plans.
     *
     * @param pigQuery Pig script
     * @param planRewrite Whether to rewrite the translated plan
     * @param validate Whether to validate the Pig logical plan before doing
     *   translation
     * @param usePigRules Whether to use Pig Rules (see PigRelPlanner} to rewrite
     *   translated rel plan
     *
     * @return A list of root nodes of the translated relational plans. Each of
     * these root corresponds to a sink operator (normally a STORE command) in the
     * Pig plan
     *
     * @throws IOException Exception during parsing or translating Pig
     */
    public List<RelNode> pigQuery2Rel(String pigQuery, boolean planRewrite,
        boolean validate, boolean usePigRules) throws IOException {
      setBatchOn();
      registerQuery(pigQuery);
      final LogicalPlan pigPlan = getCurrentDAG().getLogicalPlan();
      if (validate) {
        pigPlan.validate(getPigContext(), scope, false);
      }
      return pigPlan2Rel(pigPlan, planRewrite, usePigRules);
    }
  
    /**
     * Gets a Pig script string from a file after doing param substitution.
     *
     * @param in Pig script file
     * @param params Param sub map
     */
    public String getPigScript(InputStream in, Map<String, String> params)
        throws IOException {
      return getPigContext().doParamSubstitution(in, paramMapToList(params), null);
    }
  
    /**
     * Parses a Pig script and converts it into relational algebra plans.
     *
     * @param fileName File name
     * @param params Param substitution map
     * @param planRewrite Whether to rewrite the translated plan
     *
     * @return A list of root nodes of the translated relational plans. Each of
     * these root corresponds to a sink operator (normally a STORE command) in the
     * Pig plan
     *
     * @throws IOException Exception during parsing or translating Pig
     */
    public List<RelNode> pigScript2Rel(String fileName, Map<String, String> params,
        boolean planRewrite) throws IOException {
      setBatchOn();
      registerScript(fileName, params);
      final LogicalPlan pigPlan = getCurrentDAG().getLogicalPlan();
      pigPlan.validate(getPigContext(), scope, false);
  
      return pigPlan2Rel(pigPlan, planRewrite, true);
    }
  
    private List<RelNode> pigPlan2Rel(LogicalPlan pigPlan, boolean planRewrite,
        boolean usePigRules) throws FrontendException {
      final PigRelOpWalker walker = new PigRelOpWalker(pigPlan);
      List<RelNode> relNodes =
          new PigRelOpVisitor(pigPlan, walker, builder).translate();
      final List<RelNode> storeRels = builder.getRelsForStores();
      relNodes = storeRels != null ? storeRels : relNodes;
  
      if (usePigRules) {
        relNodes = optimizePlans(relNodes, PIG_RULES);
      }
      if (planRewrite) {
        relNodes = optimizePlans(relNodes, TRANSFORM_RULES);
      }
      return relNodes;
    }
  
    /**
     * Converts a Pig script to a list of SQL statements.
     *
     * @param pigQuery Pig script
     * @param sqlDialect Dialect of SQL language
     * @throws IOException Exception during parsing or translating Pig
     */
    public List<String> pigToSql(String pigQuery, SqlDialect sqlDialect)
        throws IOException {
      final SqlWriterConfig config = SqlPrettyWriter.config()
          .withQuoteAllIdentifiers(false)
          .withAlwaysUseParentheses(false)
          .withSelectListItemsOnSeparateLines(false)
          .withIndentation(2)
          .withDialect(sqlDialect);
      final SqlPrettyWriter writer = new SqlPrettyWriter(config);
      return pigToSql(pigQuery, writer);
    }
  
    /**
     * Converts a Pig script to a list of SQL statements.
     *
     * @param pigQuery Pig script
     * @param writer The SQL writer to decide dialect and format of SQL statements
     * @throws IOException Exception during parsing or translating Pig
     */
    private List<String> pigToSql(String pigQuery, SqlWriter writer)
        throws IOException {
      final RelToSqlConverter sqlConverter =
          new PigRelToSqlConverter(writer.getDialect());
      final List<RelNode> finalRels = pigQuery2Rel(pigQuery);
      final List<String> sqlStatements = new ArrayList<>();
      for (RelNode rel : finalRels) {
        final SqlNode sqlNode = sqlConverter.visitRoot(rel).asStatement();
        sqlNode.unparse(writer, 0, 0);
        sqlStatements.add(writer.toString());
        writer.reset();
      }
      return sqlStatements;
    }
  
    private List<RelNode> optimizePlans(List<RelNode> originalRels,
        List<RelOptRule> rules) {
      final RelOptPlanner planner = originalRels.get(0).getCluster().getPlanner();
      // Remember old rule set of the planner before resetting it with new rules
      final List<RelOptRule> oldRules = planner.getRules();
      resetPlannerRules(planner, rules);
      final Program program = Programs.of(RuleSets.ofList(planner.getRules()));
      final List<RelNode> optimizedPlans = new ArrayList<>();
      for (RelNode rel : originalRels) {
        final RelCollation collation = rel instanceof Sort
            ? ((Sort) rel).collation
            : RelCollations.EMPTY;
        // Apply the planner to obtain the physical plan
        final RelNode physicalPlan =
            program.run(planner, rel,
                rel.getTraitSet().replace(EnumerableConvention.INSTANCE)
                    .replace(collation).simplify(),
            ImmutableList.of(), ImmutableList.of());
  
        // Then convert the physical plan back to logical plan
        final RelNode logicalPlan = new ToLogicalConverter(builder).visit(physicalPlan);
        optimizedPlans.add(logicalPlan);
      }
      resetPlannerRules(planner, oldRules);
      return optimizedPlans;
    }
  
    private static void resetPlannerRules(RelOptPlanner planner,
        List<RelOptRule> rulesToSet) {
      planner.clear();
      for (RelOptRule rule : rulesToSet) {
        planner.addRule(rule);
      }
    }
  }
  