/*
 * 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.examples.cookbook;

import com.google.api.services.bigquery.model.TableFieldSchema;
import com.google.api.services.bigquery.model.TableRow;
import com.google.api.services.bigquery.model.TableSchema;
import java.util.ArrayList;
import java.util.List;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
import org.apache.beam.sdk.metrics.Counter;
import org.apache.beam.sdk.metrics.Metrics;
import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.Description;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.Validation;
import org.apache.beam.sdk.transforms.Combine;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;

/**
 * An example that reads the public 'Shakespeare' data, and for each word in the dataset that is
 * over a given length, generates a string containing the list of play names in which that word
 * appears, and saves this information to a bigquery table.
 *
 * <p>Note: Before running this example, you must create a BigQuery dataset to contain your output
 * table.
 *
 * <p>To execute this pipeline locally, specify the BigQuery table for the output:
 *
 * <pre>{@code
 * --output=YOUR_PROJECT_ID:DATASET_ID.TABLE_ID
 * }</pre>
 *
 * <p>To change the runner, specify:
 *
 * <pre>{@code
 * --runner=YOUR_SELECTED_RUNNER
 * }</pre>
 *
 * See examples/java/README.md for instructions about how to configure different runners.
 *
 * <p>The BigQuery input table defaults to {@code publicdata:samples.shakespeare} and can be
 * overridden with {@code --input}.
 */
public class CombinePerKeyExamples {
  // Use the shakespeare public BigQuery sample
  private static final String SHAKESPEARE_TABLE = "publicdata:samples.shakespeare";
  // We'll track words >= this word length across all plays in the table.
  private static final int MIN_WORD_LENGTH = 9;

  /**
   * Examines each row in the input table. If the word is greater than or equal to MIN_WORD_LENGTH,
   * outputs word, play_name.
   */
  static class ExtractLargeWordsFn extends DoFn<TableRow, KV<String, String>> {
    private final Counter smallerWords = Metrics.counter(ExtractLargeWordsFn.class, "smallerWords");

    @ProcessElement
    public void processElement(ProcessContext c) {
      TableRow row = c.element();
      String playName = (String) row.get("corpus");
      String word = (String) row.get("word");
      if (word.length() >= MIN_WORD_LENGTH) {
        c.output(KV.of(word, playName));
      } else {
        // Track how many smaller words we're not including. This information will be
        // visible in the Monitoring UI.
        smallerWords.inc();
      }
    }
  }

  /**
   * Prepares the data for writing to BigQuery by building a TableRow object containing a word with
   * a string listing the plays in which it appeared.
   */
  static class FormatShakespeareOutputFn extends DoFn<KV<String, String>, TableRow> {
    @ProcessElement
    public void processElement(ProcessContext c) {
      TableRow row =
          new TableRow().set("word", c.element().getKey()).set("all_plays", c.element().getValue());
      c.output(row);
    }
  }

  /**
   * Reads the public 'Shakespeare' data, and for each word in the dataset over a given length,
   * generates a string containing the list of play names in which that word appears. It does this
   * via the Combine.perKey transform, with the ConcatWords combine function.
   *
   * <p>Combine.perKey is similar to a GroupByKey followed by a ParDo, but has more restricted
   * semantics that allow it to be executed more efficiently. These records are then formatted as BQ
   * table rows.
   */
  static class PlaysForWord extends PTransform<PCollection<TableRow>, PCollection<TableRow>> {
    @Override
    public PCollection<TableRow> expand(PCollection<TableRow> rows) {

      // row... => <word, play_name> ...
      PCollection<KV<String, String>> words = rows.apply(ParDo.of(new ExtractLargeWordsFn()));

      // word, play_name => word, all_plays ...
      PCollection<KV<String, String>> wordAllPlays = words.apply(Combine.perKey(new ConcatWords()));

      // <word, all_plays>... => row...
      PCollection<TableRow> results = wordAllPlays.apply(ParDo.of(new FormatShakespeareOutputFn()));

      return results;
    }
  }

  /**
   * A 'combine function' used with the Combine.perKey transform. Builds a comma-separated string of
   * all input items. So, it will build a string containing all the different Shakespeare plays in
   * which the given input word has appeared.
   */
  public static class ConcatWords implements SerializableFunction<Iterable<String>, String> {
    @Override
    public String apply(Iterable<String> input) {
      StringBuilder all = new StringBuilder();
      for (String item : input) {
        if (!item.isEmpty()) {
          if (all.length() == 0) {
            all.append(item);
          } else {
            all.append(",");
            all.append(item);
          }
        }
      }
      return all.toString();
    }
  }

  /**
   * Options supported by {@link CombinePerKeyExamples}.
   *
   * <p>Inherits standard configuration options.
   */
  public interface Options extends PipelineOptions {
    @Description("Table to read from, specified as " + "<project_id>:<dataset_id>.<table_id>")
    @Default.String(SHAKESPEARE_TABLE)
    String getInput();

    void setInput(String value);

    @Description(
        "Table to write to, specified as "
            + "<project_id>:<dataset_id>.<table_id>. "
            + "The dataset_id must already exist")
    @Validation.Required
    String getOutput();

    void setOutput(String value);
  }

  public static void main(String[] args) throws Exception {

    Options options = PipelineOptionsFactory.fromArgs(args).withValidation().as(Options.class);
    runCombinePerKeyExamples(options);
  }

  static void runCombinePerKeyExamples(Options options) {
    Pipeline p = Pipeline.create(options);

    // Build the table schema for the output table.
    List<TableFieldSchema> fields = new ArrayList<>();
    fields.add(new TableFieldSchema().setName("word").setType("STRING"));
    fields.add(new TableFieldSchema().setName("all_plays").setType("STRING"));
    TableSchema schema = new TableSchema().setFields(fields);

    p.apply(BigQueryIO.readTableRows().from(options.getInput()))
        .apply(new PlaysForWord())
        .apply(
            BigQueryIO.writeTableRows()
                .to(options.getOutput())
                .withSchema(schema)
                .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED)
                .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_TRUNCATE));

    p.run().waitUntilFinish();
  }
}
