/*
 * 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;

import java.util.List;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.Row;
import org.junit.Test;

/** Tests for BeamSQL aggregation analytics functionality. */
public class BeamAnalyticFunctionsTest extends BeamSqlDslBase {

  /**
   * Table schema and data taken from.
   * https://cloud.google.com/bigquery/docs/reference/standard-sql/analytic-function-concepts#produce_table
   */
  private PCollection<Row> inputData() {
    Schema schema =
        Schema.builder()
            .addStringField("item")
            .addStringField("category")
            .addInt32Field("purchases")
            .build();
    return pipeline
        .apply(
            Create.of(
                TestUtils.rowsBuilderOf(schema)
                    .addRows(
                        "kale",
                        "vegetable",
                        23,
                        "orange",
                        "fruit",
                        2,
                        "cabbage",
                        "vegetable",
                        9,
                        "apple",
                        "fruit",
                        8,
                        "leek",
                        "vegetable",
                        2,
                        "lettuce",
                        "vegetable",
                        10)
                    .getRows()))
        .setRowSchema(schema);
  }

  /**
   * Table schema and data taken from.
   * https://cloud.google.com/bigquery/docs/reference/standard-sql/analytic-function-concepts#numbering_function_concepts
   */
  private PCollection<Row> inputData2() {
    Schema schema = Schema.builder().addInt32Field("x").build();
    return pipeline
        .apply(Create.of(TestUtils.rowsBuilderOf(schema).addRows(1, 2, 2, 5, 8, 10, 10).getRows()))
        .setRowSchema(schema);
  }

  /**
   * Compute a cumulative sum query taken from.
   * https://cloud.google.com/bigquery/docs/reference/standard-sql/analytic-function-concepts#compute_a_cumulative_sum
   */
  @Test
  public void testOverCumulativeSum() throws Exception {
    pipeline.enableAbandonedNodeEnforcement(false);
    PCollection<Row> inputRows = inputData();
    String sql =
        "SELECT item, purchases, category, sum(purchases) over "
            + "("
            + "PARTITION BY category "
            + "ORDER BY purchases "
            + "ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW"
            + ")"
            + " as total_purchases  FROM PCOLLECTION";
    PCollection<Row> result = inputRows.apply("sql", SqlTransform.query(sql));

    Schema overResultSchema =
        Schema.builder()
            .addStringField("item")
            .addInt32Field("purchases")
            .addStringField("category")
            .addInt32Field("total_purchases")
            .build();

    List<Row> overResult =
        TestUtils.RowsBuilder.of(overResultSchema)
            .addRows(
                "orange",
                2,
                "fruit",
                2,
                "apple",
                8,
                "fruit",
                10,
                "leek",
                2,
                "vegetable",
                2,
                "cabbage",
                9,
                "vegetable",
                11,
                "lettuce",
                10,
                "vegetable",
                21,
                "kale",
                23,
                "vegetable",
                44)
            .getRows();

    PAssert.that(result).containsInAnyOrder(overResult);

    pipeline.run();
  }

  /**
   * Compute a bounded ROWS sum query taken from.
   * https://cloud.google.com/bigquery/docs/reference/standard-sql/analytic-function-concepts#compute_a_cumulative_sum
   */
  @Test
  public void testOverRowsBoundedSum() throws Exception {
    pipeline.enableAbandonedNodeEnforcement(false);
    PCollection<Row> inputRows = inputData();
    String sql =
        "SELECT item, purchases, category, sum(purchases) over "
            + "("
            + "PARTITION BY category "
            + "ORDER BY purchases "
            + "ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING"
            + ")"
            + " as total_purchases  FROM PCOLLECTION";
    PCollection<Row> result = inputRows.apply("sql", SqlTransform.query(sql));

    Schema overResultSchema =
        Schema.builder()
            .addStringField("item")
            .addInt32Field("purchases")
            .addStringField("category")
            .addInt32Field("total_purchases")
            .build();

    List<Row> overResult =
        TestUtils.RowsBuilder.of(overResultSchema)
            .addRows(
                "orange",
                2,
                "fruit",
                10,
                "apple",
                8,
                "fruit",
                10,
                "leek",
                2,
                "vegetable",
                11,
                "cabbage",
                9,
                "vegetable",
                21,
                "lettuce",
                10,
                "vegetable",
                42,
                "kale",
                23,
                "vegetable",
                33)
            .getRows();

    PAssert.that(result).containsInAnyOrder(overResult);

    pipeline.run();
  }

  /**
   * Compute a bounded RANGE sum query taken from.
   * https://cloud.google.com/bigquery/docs/reference/standard-sql/analytic-function-concepts#compute_a_cumulative_sum
   */
  @Test
  public void testOverRangeBoundedSum() throws Exception {
    pipeline.enableAbandonedNodeEnforcement(false);
    PCollection<Row> inputRows = inputData();
    String sql =
        "SELECT item, purchases, category, sum(purchases) over "
            + "("
            + "PARTITION BY category "
            + "ORDER BY purchases "
            + "RANGE BETWEEN 1 PRECEDING AND 1 FOLLOWING"
            + ")"
            + " as total_purchases  FROM PCOLLECTION";
    PCollection<Row> result = inputRows.apply("sql", SqlTransform.query(sql));

    Schema overResultSchema =
        Schema.builder()
            .addStringField("item")
            .addInt32Field("purchases")
            .addStringField("category")
            .addInt32Field("total_purchases")
            .build();

    List<Row> overResult =
        TestUtils.RowsBuilder.of(overResultSchema)
            .addRows(
                "orange",
                2,
                "fruit",
                2,
                "apple",
                8,
                "fruit",
                8,
                "leek",
                2,
                "vegetable",
                2,
                "cabbage",
                9,
                "vegetable",
                19,
                "lettuce",
                10,
                "vegetable",
                19,
                "kale",
                23,
                "vegetable",
                23)
            .getRows();

    PAssert.that(result).containsInAnyOrder(overResult);

    pipeline.run();
  }

  /**
   * Compute a cumulative sum (inverse order) query taken from.
   * https://cloud.google.com/bigquery/docs/reference/standard-sql/analytic-function-concepts#compute_a_cumulative_sum
   */
  @Test
  public void testOverCumulativeSumOrderByDesc() throws Exception {
    pipeline.enableAbandonedNodeEnforcement(false);
    PCollection<Row> inputRows = inputData();
    String sql =
        "SELECT item, purchases, category, sum(purchases) over "
            + "("
            + "PARTITION BY category "
            + "ORDER BY purchases DESC "
            + "ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW"
            + ")"
            + " as total_purchases  FROM PCOLLECTION";
    PCollection<Row> result = inputRows.apply("sql", SqlTransform.query(sql));

    Schema overResultSchema =
        Schema.builder()
            .addStringField("item")
            .addInt32Field("purchases")
            .addStringField("category")
            .addInt32Field("total_purchases")
            .build();

    List<Row> overResult =
        TestUtils.RowsBuilder.of(overResultSchema)
            .addRows(
                "orange",
                2,
                "fruit",
                10,
                "apple",
                8,
                "fruit",
                8,
                "leek",
                2,
                "vegetable",
                44,
                "cabbage",
                9,
                "vegetable",
                42,
                "lettuce",
                10,
                "vegetable",
                33,
                "kale",
                23,
                "vegetable",
                23)
            .getRows();

    PAssert.that(result).containsInAnyOrder(overResult);

    pipeline.run();
  }

  /**
   * Basic analytic function query taken from.
   * https://cloud.google.com/bigquery/docs/reference/standard-sql/analytic-function-concepts#compute_a_grand_total
   */
  @Test
  public void testSimpleOverFunction() throws Exception {
    pipeline.enableAbandonedNodeEnforcement(false);
    PCollection<Row> inputRows = inputData();
    String sql =
        "SELECT item, purchases, category, sum(purchases) over ()"
            + " as total_purchases  FROM PCOLLECTION";
    PCollection<Row> result = inputRows.apply("sql", SqlTransform.query(sql));

    Schema overResultSchema =
        Schema.builder()
            .addStringField("item")
            .addInt32Field("purchases")
            .addStringField("category")
            .addInt32Field("total_purchases")
            .build();

    List<Row> overResult =
        TestUtils.RowsBuilder.of(overResultSchema)
            .addRows(
                "orange",
                2,
                "fruit",
                54,
                "apple",
                8,
                "fruit",
                54,
                "leek",
                2,
                "vegetable",
                54,
                "cabbage",
                9,
                "vegetable",
                54,
                "lettuce",
                10,
                "vegetable",
                54,
                "kale",
                23,
                "vegetable",
                54)
            .getRows();

    PAssert.that(result).containsInAnyOrder(overResult);

    pipeline.run();
  }

  // Tests for new functions
  @Test
  public void testFirstValueFunction() throws Exception {
    pipeline.enableAbandonedNodeEnforcement(false);
    PCollection<Row> inputRows = inputData();
    String sql =
        "SELECT item, purchases, category, FIRST_VALUE(purchases) over "
            + "("
            + "PARTITION BY category "
            + "ORDER BY purchases "
            + "ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW"
            + ")"
            + " as total_purchases  FROM PCOLLECTION";
    PCollection<Row> result = inputRows.apply("sql", SqlTransform.query(sql));

    Schema overResultSchema =
        Schema.builder()
            .addStringField("item")
            .addInt32Field("purchases")
            .addStringField("category")
            .addInt32Field("total_purchases")
            .build();

    List<Row> overResult =
        TestUtils.RowsBuilder.of(overResultSchema)
            .addRows(
                "orange",
                2,
                "fruit",
                2,
                "apple",
                8,
                "fruit",
                2,
                "leek",
                2,
                "vegetable",
                2,
                "cabbage",
                9,
                "vegetable",
                2,
                "lettuce",
                10,
                "vegetable",
                2,
                "kale",
                23,
                "vegetable",
                2)
            .getRows();

    PAssert.that(result).containsInAnyOrder(overResult);

    pipeline.run();
  }

  @Test
  public void testLastValueFunction() throws Exception {
    pipeline.enableAbandonedNodeEnforcement(false);
    PCollection<Row> inputRows = inputData();
    String sql =
        "SELECT item, purchases, category, LAST_VALUE(purchases) over "
            + "("
            + "PARTITION BY category "
            + "ORDER BY purchases "
            + "ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW"
            + ")"
            + " as total_purchases  FROM PCOLLECTION";
    PCollection<Row> result = inputRows.apply("sql", SqlTransform.query(sql));

    Schema overResultSchema =
        Schema.builder()
            .addStringField("item")
            .addInt32Field("purchases")
            .addStringField("category")
            .addInt32Field("total_purchases")
            .build();

    List<Row> overResult =
        TestUtils.RowsBuilder.of(overResultSchema)
            .addRows(
                "orange",
                2,
                "fruit",
                2,
                "apple",
                8,
                "fruit",
                8,
                "leek",
                2,
                "vegetable",
                2,
                "cabbage",
                9,
                "vegetable",
                9,
                "lettuce",
                10,
                "vegetable",
                10,
                "kale",
                23,
                "vegetable",
                23)
            .getRows();

    PAssert.that(result).containsInAnyOrder(overResult);

    pipeline.run();
  }

  @Test
  public void testRowNumberFunction() throws Exception {
    pipeline.enableAbandonedNodeEnforcement(false);
    PCollection<Row> inputRows = inputData2();
    String sql = "SELECT x, ROW_NUMBER() over (ORDER BY x ) as agg  FROM PCOLLECTION";
    PCollection<Row> result = inputRows.apply("sql", SqlTransform.query(sql));

    Schema overResultSchema = Schema.builder().addInt32Field("x").addInt64Field("agg").build();

    List<Row> overResult =
        TestUtils.RowsBuilder.of(overResultSchema)
            .addRows(
                1, 1L,
                2, 2L,
                2, 3L,
                5, 4L,
                8, 5L,
                10, 6L,
                10, 7L)
            .getRows();

    PAssert.that(result).containsInAnyOrder(overResult);

    pipeline.run();
  }

  @Test
  public void testDenseRankFunction() throws Exception {
    pipeline.enableAbandonedNodeEnforcement(false);
    PCollection<Row> inputRows = inputData2();
    String sql = "SELECT x, DENSE_RANK() over (ORDER BY x ) as agg  FROM PCOLLECTION";
    PCollection<Row> result = inputRows.apply("sql", SqlTransform.query(sql));

    Schema overResultSchema = Schema.builder().addInt32Field("x").addInt64Field("agg").build();

    List<Row> overResult =
        TestUtils.RowsBuilder.of(overResultSchema)
            .addRows(
                1, 1L,
                2, 2L,
                2, 2L,
                5, 3L,
                8, 4L,
                10, 5L,
                10, 5L)
            .getRows();

    PAssert.that(result).containsInAnyOrder(overResult);

    pipeline.run();
  }

  @Test
  public void testRankFunction() throws Exception {
    pipeline.enableAbandonedNodeEnforcement(false);
    PCollection<Row> inputRows = inputData2();
    String sql = "SELECT x, RANK() over (ORDER BY x ) as agg  FROM PCOLLECTION";
    PCollection<Row> result = inputRows.apply("sql", SqlTransform.query(sql));

    Schema overResultSchema = Schema.builder().addInt32Field("x").addInt64Field("agg").build();

    List<Row> overResult =
        TestUtils.RowsBuilder.of(overResultSchema)
            .addRows(
                1, 1L,
                2, 2L,
                2, 2L,
                5, 4L,
                8, 5L,
                10, 6L,
                10, 6L)
            .getRows();

    PAssert.that(result).containsInAnyOrder(overResult);

    pipeline.run();
  }

  @Test
  public void testPercentRankFunction() throws Exception {
    pipeline.enableAbandonedNodeEnforcement(false);
    PCollection<Row> inputRows = inputData2();
    String sql = "SELECT x, PERCENT_RANK() over (ORDER BY x ) as agg  FROM PCOLLECTION";
    PCollection<Row> result = inputRows.apply("sql", SqlTransform.query(sql));

    Schema overResultSchema = Schema.builder().addInt32Field("x").addDoubleField("agg").build();

    List<Row> overResult =
        TestUtils.RowsBuilder.of(overResultSchema)
            .addRows(
                1, 0.0 / 6.0,
                2, 1.0 / 6.0,
                2, 1.0 / 6.0,
                5, 3.0 / 6.0,
                8, 4.0 / 6.0,
                10, 5.0 / 6.0,
                10, 5.0 / 6.0)
            .getRows();

    PAssert.that(result).containsInAnyOrder(overResult);

    pipeline.run();
  }
}
