/*
 * 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.beam.sdk.extensions.sql.impl.rule;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;
import org.apache.beam.sdk.extensions.sql.impl.BeamSqlEnv;
import org.apache.beam.sdk.extensions.sql.impl.planner.BeamRuleSets;
import org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode;
import org.apache.beam.sdk.extensions.sql.meta.provider.test.TestTableProvider;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.vendor.calcite.v1_20_0.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.calcite.v1_20_0.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.DataContext;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.adapter.enumerable.EnumerableConvention;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.adapter.enumerable.EnumerableRules;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.linq4j.Enumerable;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.linq4j.Linq4j;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.plan.ConventionTraitDef;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.plan.RelOptRule;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.plan.RelTraitSet;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.RelCollationTraitDef;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.RelCollations;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.RelFieldCollation;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.RelNode;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.RelRoot;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.core.Join;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.core.TableScan;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.rules.JoinCommuteRule;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.rules.SortProjectTransposeRule;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.type.RelDataType;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.type.RelDataTypeFactory;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.schema.ScannableTable;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.schema.SchemaPlus;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.schema.Statistic;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.schema.Statistics;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.schema.Table;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.schema.impl.AbstractSchema;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.schema.impl.AbstractTable;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlNode;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.parser.SqlParser;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.tools.FrameworkConfig;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.tools.Frameworks;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.tools.Planner;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.tools.Programs;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.tools.RuleSet;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.tools.RuleSets;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.util.ImmutableBitSet;
import org.junit.Assert;
import org.junit.Test;

/**
 * This test ensures that we are reordering joins and get a plan similar to Join(large,Join(small,
 * medium)) instead of Join(small, Join(medium,large).
 */
public class JoinReorderingTest {
  private final PipelineOptions defaultPipelineOptions = PipelineOptionsFactory.create();

  @Test
  public void testTableSizes() {
    TestTableProvider tableProvider = new TestTableProvider();
    createThreeTables(tableProvider);

    Assert.assertEquals(
        1d,
        tableProvider
            .buildBeamSqlTable(tableProvider.getTable("small_table"))
            .getTableStatistics(null)
            .getRowCount(),
        0.01);

    Assert.assertEquals(
        3d,
        tableProvider
            .buildBeamSqlTable(tableProvider.getTable("medium_table"))
            .getTableStatistics(null)
            .getRowCount(),
        0.01);

    Assert.assertEquals(
        100d,
        tableProvider
            .buildBeamSqlTable(tableProvider.getTable("large_table"))
            .getTableStatistics(null)
            .getRowCount(),
        0.01);
  }

  @Test
  public void testBeamJoinAssociationRule() throws Exception {
    RuleSet prepareRules =
        RuleSets.ofList(
            SortProjectTransposeRule.INSTANCE,
            EnumerableRules.ENUMERABLE_JOIN_RULE,
            EnumerableRules.ENUMERABLE_PROJECT_RULE,
            EnumerableRules.ENUMERABLE_SORT_RULE,
            EnumerableRules.ENUMERABLE_TABLE_SCAN_RULE);

    String sqlQuery =
        "select * from \"tt\".\"large_table\" as large_table "
            + " JOIN \"tt\".\"medium_table\" as medium_table on large_table.\"medium_key\" = medium_table.\"large_key\" "
            + " JOIN \"tt\".\"small_table\" as small_table on medium_table.\"small_key\" = small_table.\"medium_key\" ";

    RelNode originalPlan = transform(sqlQuery, prepareRules);
    RelNode optimizedPlan =
        transform(
            sqlQuery,
            RuleSets.ofList(
                ImmutableList.<RelOptRule>builder()
                    .addAll(prepareRules)
                    .add(BeamJoinAssociateRule.INSTANCE)
                    .build()));

    assertTopTableInJoins(originalPlan, "small_table");
    assertTopTableInJoins(optimizedPlan, "large_table");
  }

  @Test
  public void testBeamJoinPushThroughJoinRuleLeft() throws Exception {
    RuleSet prepareRules =
        RuleSets.ofList(
            SortProjectTransposeRule.INSTANCE,
            EnumerableRules.ENUMERABLE_JOIN_RULE,
            EnumerableRules.ENUMERABLE_PROJECT_RULE,
            EnumerableRules.ENUMERABLE_SORT_RULE,
            EnumerableRules.ENUMERABLE_TABLE_SCAN_RULE);

    String sqlQuery =
        "select * from \"tt\".\"large_table\" as large_table "
            + " JOIN \"tt\".\"medium_table\" as medium_table on large_table.\"medium_key\" = medium_table.\"large_key\" "
            + " JOIN \"tt\".\"small_table\" as small_table on medium_table.\"small_key\" = small_table.\"medium_key\" ";

    RelNode originalPlan = transform(sqlQuery, prepareRules);
    RelNode optimizedPlan =
        transform(
            sqlQuery,
            RuleSets.ofList(
                ImmutableList.<RelOptRule>builder()
                    .addAll(prepareRules)
                    .add(BeamJoinPushThroughJoinRule.LEFT)
                    .build()));

    assertTopTableInJoins(originalPlan, "small_table");
    assertTopTableInJoins(optimizedPlan, "large_table");
  }

  @Test
  public void testBeamJoinPushThroughJoinRuleRight() throws Exception {
    RuleSet prepareRules =
        RuleSets.ofList(
            SortProjectTransposeRule.INSTANCE,
            EnumerableRules.ENUMERABLE_JOIN_RULE,
            EnumerableRules.ENUMERABLE_PROJECT_RULE,
            EnumerableRules.ENUMERABLE_SORT_RULE,
            EnumerableRules.ENUMERABLE_TABLE_SCAN_RULE);

    String sqlQuery =
        "select * from \"tt\".\"medium_table\" as medium_table "
            + " JOIN \"tt\".\"large_table\" as large_table on large_table.\"medium_key\" = medium_table.\"large_key\" "
            + " JOIN \"tt\".\"small_table\" as small_table on medium_table.\"small_key\" = small_table.\"medium_key\" ";

    RelNode originalPlan = transform(sqlQuery, prepareRules);
    RelNode optimizedPlan =
        transform(
            sqlQuery,
            RuleSets.ofList(
                ImmutableList.<RelOptRule>builder()
                    .addAll(prepareRules)
                    .add(BeamJoinPushThroughJoinRule.RIGHT)
                    .build()));

    assertTopTableInJoins(originalPlan, "small_table");
    assertTopTableInJoins(optimizedPlan, "large_table");
  }

  @Test
  public void testSystemReorderingLargeMediumSmall() {
    TestTableProvider tableProvider = new TestTableProvider();
    createThreeTables(tableProvider);
    BeamSqlEnv env = BeamSqlEnv.withTableProvider(tableProvider);

    // This is Join(Join(large, medium), small) which should be converted to a join that large table
    // is on the top.
    BeamRelNode parsedQuery =
        env.parseQuery(
            "select * from large_table "
                + " JOIN medium_table on large_table.medium_key = medium_table.large_key "
                + " JOIN small_table on medium_table.small_key = small_table.medium_key ");
    assertTopTableInJoins(parsedQuery, "large_table");
  }

  @Test
  public void testSystemReorderingMediumLargeSmall() {
    TestTableProvider tableProvider = new TestTableProvider();
    createThreeTables(tableProvider);
    BeamSqlEnv env = BeamSqlEnv.withTableProvider(tableProvider);

    // This is Join(Join(medium, large), small) which should be converted to a join that large table
    // is on the top.
    BeamRelNode parsedQuery =
        env.parseQuery(
            "select * from medium_table "
                + " JOIN large_table on large_table.medium_key = medium_table.large_key "
                + " JOIN small_table on medium_table.small_key = small_table.medium_key ");
    assertTopTableInJoins(parsedQuery, "large_table");
  }

  @Test
  public void testSystemNotReorderingWithoutRules() {
    TestTableProvider tableProvider = new TestTableProvider();
    createThreeTables(tableProvider);
    List<RelOptRule> ruleSet =
        Arrays.stream(BeamRuleSets.getRuleSets())
            .flatMap(rules -> StreamSupport.stream(rules.spliterator(), false))
            .filter(rule -> !(rule instanceof BeamJoinPushThroughJoinRule))
            .filter(rule -> !(rule instanceof BeamJoinAssociateRule))
            .filter(rule -> !(rule instanceof JoinCommuteRule))
            .collect(Collectors.toList());

    BeamSqlEnv env =
        BeamSqlEnv.builder(tableProvider)
            .setPipelineOptions(PipelineOptionsFactory.create())
            .setRuleSets(new RuleSet[] {RuleSets.ofList(ruleSet)})
            .build();

    // This is Join(Join(medium, large), small) which should be converted to a join that large table
    // is on the top.
    BeamRelNode parsedQuery =
        env.parseQuery(
            "select * from medium_table "
                + " JOIN large_table on large_table.medium_key = medium_table.large_key "
                + " JOIN small_table on medium_table.small_key = small_table.medium_key ");
    assertTopTableInJoins(parsedQuery, "small_table");
  }

  @Test
  public void testSystemNotReorderingMediumSmallLarge() {
    TestTableProvider tableProvider = new TestTableProvider();
    createThreeTables(tableProvider);
    BeamSqlEnv env = BeamSqlEnv.withTableProvider(tableProvider);

    // This is a correct ordered join because large table is on the top. It should not change that.
    BeamRelNode parsedQuery =
        env.parseQuery(
            "select * from medium_table "
                + " JOIN small_table on medium_table.small_key = small_table.medium_key "
                + " JOIN large_table on large_table.medium_key = medium_table.large_key ");
    assertTopTableInJoins(parsedQuery, "large_table");
  }

  @Test
  public void testSystemNotReorderingSmallMediumLarge() {
    TestTableProvider tableProvider = new TestTableProvider();
    createThreeTables(tableProvider);
    BeamSqlEnv env = BeamSqlEnv.withTableProvider(tableProvider);

    // This is a correct ordered join because large table is on the top. It should not change that.
    BeamRelNode parsedQuery =
        env.parseQuery(
            "select * from small_table "
                + " JOIN medium_table on medium_table.small_key = small_table.medium_key "
                + " JOIN large_table on large_table.medium_key = medium_table.large_key ");
    assertTopTableInJoins(parsedQuery, "large_table");
  }

  private RelNode transform(String sql, RuleSet prepareRules) throws Exception {
    final SchemaPlus rootSchema = Frameworks.createRootSchema(true);
    final SchemaPlus defSchema = rootSchema.add("tt", new ThreeTablesSchema());
    final FrameworkConfig config =
        Frameworks.newConfigBuilder()
            .parserConfig(SqlParser.Config.DEFAULT)
            .defaultSchema(defSchema)
            .traitDefs(ConventionTraitDef.INSTANCE, RelCollationTraitDef.INSTANCE)
            .programs(Programs.of(prepareRules))
            .build();
    Planner planner = Frameworks.getPlanner(config);
    SqlNode parse = planner.parse(sql);
    SqlNode validate = planner.validate(parse);
    RelRoot planRoot = planner.rel(validate);
    RelNode planBefore = planRoot.rel;
    RelTraitSet desiredTraits = planBefore.getTraitSet().replace(EnumerableConvention.INSTANCE);
    return planner.transform(0, desiredTraits, planBefore);
  }

  private void assertTopTableInJoins(RelNode parsedQuery, String expectedTableName) {
    RelNode firstJoin = parsedQuery;
    while (!(firstJoin instanceof Join)) {
      firstJoin = firstJoin.getInput(0);
    }

    RelNode topRight = ((Join) firstJoin).getRight();
    while (!(topRight instanceof Join) && !(topRight instanceof TableScan)) {
      topRight = topRight.getInput(0);
    }

    if (topRight instanceof TableScan) {
      Assert.assertTrue(topRight.getDescription().contains(expectedTableName));
    } else {
      RelNode topLeft = ((Join) firstJoin).getLeft();
      while (!(topLeft instanceof TableScan)) {
        topLeft = topLeft.getInput(0);
      }

      Assert.assertTrue(topLeft.getDescription().contains(expectedTableName));
    }
  }

  private void createThreeTables(TestTableProvider tableProvider) {
    BeamSqlEnv env = BeamSqlEnv.withTableProvider(tableProvider);
    env.executeDdl("CREATE EXTERNAL TABLE small_table (id INTEGER, medium_key INTEGER) TYPE text");

    env.executeDdl(
        "CREATE EXTERNAL TABLE medium_table ("
            + "id INTEGER,"
            + "small_key INTEGER,"
            + "large_key INTEGER"
            + ") TYPE text");

    env.executeDdl(
        "CREATE EXTERNAL TABLE large_table ("
            + "id INTEGER,"
            + "medium_key INTEGER"
            + ") TYPE text");

    Row row =
        Row.withSchema(tableProvider.getTable("small_table").getSchema()).addValues(1, 1).build();
    tableProvider.addRows("small_table", row);

    for (int i = 0; i < 3; i++) {
      row =
          Row.withSchema(tableProvider.getTable("medium_table").getSchema())
              .addValues(i, 1, 2)
              .build();
      tableProvider.addRows("medium_table", row);
    }

    for (int i = 0; i < 100; i++) {
      row =
          Row.withSchema(tableProvider.getTable("large_table").getSchema()).addValues(i, 2).build();
      tableProvider.addRows("large_table", row);
    }
  }
}

final class ThreeTablesSchema extends AbstractSchema {

  private final ImmutableMap<String, Table> tables;

  public ThreeTablesSchema() {
    super();
    ArrayList<Object[]> mediumData = new ArrayList<>();
    for (int i = 0; i < 3; i++) {
      mediumData.add(new Object[] {i, 1, 2});
    }

    ArrayList<Object[]> largeData = new ArrayList<>();
    for (int i = 0; i < 100; i++) {
      largeData.add(new Object[] {i, 2});
    }

    tables =
        ImmutableMap.<String, Table>builder()
            .put(
                "small_table",
                new PkClusteredTable(
                    factory ->
                        new RelDataTypeFactory.Builder(factory)
                            .add("id", factory.createJavaType(int.class))
                            .add("medium_key", factory.createJavaType(int.class))
                            .build(),
                    ImmutableBitSet.of(0),
                    Arrays.asList(new Object[] {1, 1}, new Object[] {2, 1})))
            .put(
                "medium_table",
                new PkClusteredTable(
                    factory ->
                        new RelDataTypeFactory.Builder(factory)
                            .add("id", factory.createJavaType(int.class))
                            .add("small_key", factory.createJavaType(int.class))
                            .add("large_key", factory.createJavaType(int.class))
                            .build(),
                    ImmutableBitSet.of(0),
                    mediumData))
            .put(
                "large_table",
                new PkClusteredTable(
                    factory ->
                        new RelDataTypeFactory.Builder(factory)
                            .add("id", factory.createJavaType(int.class))
                            .add("medium_key", factory.createJavaType(int.class))
                            .build(),
                    ImmutableBitSet.of(0),
                    largeData))
            .build();
  }

  @Override
  protected Map<String, org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.schema.Table>
      getTableMap() {
    return tables;
  }

  /** A table sorted (ascending direction and nulls last) on the primary key. */
  private static class PkClusteredTable extends AbstractTable implements ScannableTable {
    private final ImmutableBitSet pkColumns;
    private final List<Object[]> data;
    private final java.util.function.Function<RelDataTypeFactory, RelDataType> typeBuilder;

    PkClusteredTable(
        Function<RelDataTypeFactory, RelDataType> dataTypeBuilder,
        ImmutableBitSet pkColumns,
        List<Object[]> data) {
      this.data = data;
      this.typeBuilder = dataTypeBuilder;
      this.pkColumns = pkColumns;
    }

    @Override
    public Statistic getStatistic() {
      List<RelFieldCollation> collationFields = new ArrayList<>();
      for (Integer key : pkColumns) {
        collationFields.add(
            new RelFieldCollation(
                key, RelFieldCollation.Direction.ASCENDING, RelFieldCollation.NullDirection.LAST));
      }
      return Statistics.of(
          data.size(),
          ImmutableList.of(pkColumns),
          ImmutableList.of(RelCollations.of(collationFields)));
    }

    @Override
    public RelDataType getRowType(final RelDataTypeFactory typeFactory) {
      return typeBuilder.apply(typeFactory);
    }

    @Override
    public Enumerable<Object[]> scan(final DataContext root) {
      return Linq4j.asEnumerable(data);
    }
  }
}