examples/java/src/main/java/org/apache/beam/examples/complete/TfIdf.java

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

import java.io.File;
import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.HashSet;
import java.util.Set;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.StringDelegateCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.extensions.gcp.options.GcsOptions;
import org.apache.beam.sdk.io.TextIO;
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.Count;
import org.apache.beam.sdk.transforms.Distinct;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.Flatten;
import org.apache.beam.sdk.transforms.Keys;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.Values;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.WithKeys;
import org.apache.beam.sdk.transforms.join.CoGbkResult;
import org.apache.beam.sdk.transforms.join.CoGroupByKey;
import org.apache.beam.sdk.transforms.join.KeyedPCollectionTuple;
import org.apache.beam.sdk.util.GcsUtil;
import org.apache.beam.sdk.util.gcsfs.GcsPath;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PDone;
import org.apache.beam.sdk.values.TupleTag;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * An example that computes a basic TF-IDF search table for a directory or GCS prefix.
 *
 * <p>Concepts: joining data; side inputs; logging
 *
 * <p>To execute this pipeline locally, specify a local output file or output prefix on GCS:
 * <pre>{@code
 *   --output=[YOUR_LOCAL_FILE | gs://YOUR_OUTPUT_PREFIX]
 * }</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 default input is {@code gs://apache-beam-samples/shakespeare/} and can be overridden with
 * {@code --input}.
 */
public class TfIdf {
  /**
   * Options supported by {@link TfIdf}.
   *
   * <p>Inherits standard configuration options.
   */
  public interface Options extends PipelineOptions {
    @Description("Path to the directory or GCS prefix containing files to read from")
    @Default.String("gs://apache-beam-samples/shakespeare/")
    String getInput();
    void setInput(String value);

    @Description("Prefix of output URI to write to")
    @Validation.Required
    String getOutput();
    void setOutput(String value);
  }

  /**
   * Lists documents contained beneath the {@code options.input} prefix/directory.
   */
  public static Set<URI> listInputDocuments(Options options)
      throws URISyntaxException, IOException {
    URI baseUri = new URI(options.getInput());

    // List all documents in the directory or GCS prefix.
    URI absoluteUri;
    if (baseUri.getScheme() != null) {
      absoluteUri = baseUri;
    } else {
      absoluteUri = new URI(
          "file",
          baseUri.getAuthority(),
          baseUri.getPath(),
          baseUri.getQuery(),
          baseUri.getFragment());
    }

    Set<URI> uris = new HashSet<>();
    if (absoluteUri.getScheme().equals("file")) {
      File directory = new File(absoluteUri);
      for (String entry : directory.list()) {
        File path = new File(directory, entry);
        uris.add(path.toURI());
      }
    } else if (absoluteUri.getScheme().equals("gs")) {
      GcsUtil gcsUtil = options.as(GcsOptions.class).getGcsUtil();
      URI gcsUriGlob = new URI(
          absoluteUri.getScheme(),
          absoluteUri.getAuthority(),
          absoluteUri.getPath() + "*",
          absoluteUri.getQuery(),
          absoluteUri.getFragment());
      for (GcsPath entry : gcsUtil.expand(GcsPath.fromUri(gcsUriGlob))) {
        uris.add(entry.toUri());
      }
    }

    return uris;
  }

  /**
   * Reads the documents at the provided uris and returns all lines
   * from the documents tagged with which document they are from.
   */
  public static class ReadDocuments
      extends PTransform<PBegin, PCollection<KV<URI, String>>> {
    private Iterable<URI> uris;

    public ReadDocuments(Iterable<URI> uris) {
      this.uris = uris;
    }

    @Override
    public Coder<?> getDefaultOutputCoder() {
      return KvCoder.of(StringDelegateCoder.of(URI.class), StringUtf8Coder.of());
    }

    @Override
    public PCollection<KV<URI, String>> expand(PBegin input) {
      Pipeline pipeline = input.getPipeline();

      // Create one TextIO.Read transform for each document
      // and add its output to a PCollectionList
      PCollectionList<KV<URI, String>> urisToLines =
          PCollectionList.empty(pipeline);

      // TextIO.Read supports:
      //  - file: URIs and paths locally
      //  - gs: URIs on the service
      for (final URI uri : uris) {
        String uriString;
        if (uri.getScheme().equals("file")) {
          uriString = new File(uri).getPath();
        } else {
          uriString = uri.toString();
        }

        PCollection<KV<URI, String>> oneUriToLines = pipeline
            .apply("TextIO.Read(" + uriString + ")", TextIO.read().from(uriString))
            .apply("WithKeys(" + uriString + ")", WithKeys.<URI, String>of(uri));

        urisToLines = urisToLines.and(oneUriToLines);
      }

      return urisToLines.apply(Flatten.<KV<URI, String>>pCollections());
    }
  }

  /**
   * A transform containing a basic TF-IDF pipeline. The input consists of KV objects
   * where the key is the document's URI and the value is a piece
   * of the document's content. The output is mapping from terms to
   * scores for each document URI.
   */
  public static class ComputeTfIdf
      extends PTransform<PCollection<KV<URI, String>>, PCollection<KV<String, KV<URI, Double>>>> {
    public ComputeTfIdf() { }

    @Override
    public PCollection<KV<String, KV<URI, Double>>> expand(
      PCollection<KV<URI, String>> uriToContent) {

      // Compute the total number of documents, and
      // prepare this singleton PCollectionView for
      // use as a side input.
      final PCollectionView<Long> totalDocuments =
          uriToContent
          .apply("GetURIs", Keys.<URI>create())
          .apply("DistinctDocs", Distinct.<URI>create())
          .apply(Count.<URI>globally())
          .apply(View.<Long>asSingleton());

      // Create a collection of pairs mapping a URI to each
      // of the words in the document associated with that that URI.
      PCollection<KV<URI, String>> uriToWords = uriToContent
          .apply("SplitWords", ParDo.of(
              new DoFn<KV<URI, String>, KV<URI, String>>() {
                @ProcessElement
                public void processElement(ProcessContext c) {
                  URI uri = c.element().getKey();
                  String line = c.element().getValue();
                  for (String word : line.split("\\W+")) {
                    // Log INFO messages when the word “love” is found.
                    if (word.toLowerCase().equals("love")) {
                      LOG.info("Found {}", word.toLowerCase());
                    }

                    if (!word.isEmpty()) {
                      c.output(KV.of(uri, word.toLowerCase()));
                    }
                  }
                }
              }));

      // Compute a mapping from each word to the total
      // number of documents in which it appears.
      PCollection<KV<String, Long>> wordToDocCount = uriToWords
          .apply("DistinctWords", Distinct.<KV<URI, String>>create())
          .apply(Values.<String>create())
          .apply("CountDocs", Count.<String>perElement());

      // Compute a mapping from each URI to the total
      // number of words in the document associated with that URI.
      PCollection<KV<URI, Long>> uriToWordTotal = uriToWords
          .apply("GetURIs2", Keys.<URI>create())
          .apply("CountWords", Count.<URI>perElement());

      // Count, for each (URI, word) pair, the number of
      // occurrences of that word in the document associated
      // with the URI.
      PCollection<KV<KV<URI, String>, Long>> uriAndWordToCount = uriToWords
          .apply("CountWordDocPairs", Count.<KV<URI, String>>perElement());

      // Adjust the above collection to a mapping from
      // (URI, word) pairs to counts into an isomorphic mapping
      // from URI to (word, count) pairs, to prepare for a join
      // by the URI key.
      PCollection<KV<URI, KV<String, Long>>> uriToWordAndCount = uriAndWordToCount
          .apply("ShiftKeys", ParDo.of(
              new DoFn<KV<KV<URI, String>, Long>, KV<URI, KV<String, Long>>>() {
                @ProcessElement
                public void processElement(ProcessContext c) {
                  URI uri = c.element().getKey().getKey();
                  String word = c.element().getKey().getValue();
                  Long occurrences = c.element().getValue();
                  c.output(KV.of(uri, KV.of(word, occurrences)));
                }
              }));

      // Prepare to join the mapping of URI to (word, count) pairs with
      // the mapping of URI to total word counts, by associating
      // each of the input PCollection<KV<URI, ...>> with
      // a tuple tag. Each input must have the same key type, URI
      // in this case. The type parameter of the tuple tag matches
      // the types of the values for each collection.
      final TupleTag<Long> wordTotalsTag = new TupleTag<Long>();
      final TupleTag<KV<String, Long>> wordCountsTag = new TupleTag<KV<String, Long>>();
      KeyedPCollectionTuple<URI> coGbkInput = KeyedPCollectionTuple
          .of(wordTotalsTag, uriToWordTotal)
          .and(wordCountsTag, uriToWordAndCount);

      // Perform a CoGroupByKey (a sort of pre-join) on the prepared
      // inputs. This yields a mapping from URI to a CoGbkResult
      // (CoGroupByKey Result). The CoGbkResult is a mapping
      // from the above tuple tags to the values in each input
      // associated with a particular URI. In this case, each
      // KV<URI, CoGbkResult> group a URI with the total number of
      // words in that document as well as all the (word, count)
      // pairs for particular words.
      PCollection<KV<URI, CoGbkResult>> uriToWordAndCountAndTotal = coGbkInput
          .apply("CoGroupByUri", CoGroupByKey.<URI>create());

      // Compute a mapping from each word to a (URI, term frequency)
      // pair for each URI. A word's term frequency for a document
      // is simply the number of times that word occurs in the document
      // divided by the total number of words in the document.
      PCollection<KV<String, KV<URI, Double>>> wordToUriAndTf = uriToWordAndCountAndTotal
          .apply("ComputeTermFrequencies", ParDo.of(
              new DoFn<KV<URI, CoGbkResult>, KV<String, KV<URI, Double>>>() {
                @ProcessElement
                public void processElement(ProcessContext c) {
                  URI uri = c.element().getKey();
                  Long wordTotal = c.element().getValue().getOnly(wordTotalsTag);

                  for (KV<String, Long> wordAndCount
                           : c.element().getValue().getAll(wordCountsTag)) {
                    String word = wordAndCount.getKey();
                    Long wordCount = wordAndCount.getValue();
                    Double termFrequency = wordCount.doubleValue() / wordTotal.doubleValue();
                    c.output(KV.of(word, KV.of(uri, termFrequency)));
                  }
                }
              }));

      // Compute a mapping from each word to its document frequency.
      // A word's document frequency in a corpus is the number of
      // documents in which the word appears divided by the total
      // number of documents in the corpus. Note how the total number of
      // documents is passed as a side input; the same value is
      // presented to each invocation of the DoFn.
      PCollection<KV<String, Double>> wordToDf = wordToDocCount
          .apply("ComputeDocFrequencies", ParDo
              .of(new DoFn<KV<String, Long>, KV<String, Double>>() {
                @ProcessElement
                public void processElement(ProcessContext c) {
                  String word = c.element().getKey();
                  Long documentCount = c.element().getValue();
                  Long documentTotal = c.sideInput(totalDocuments);
                  Double documentFrequency = documentCount.doubleValue()
                      / documentTotal.doubleValue();

                  c.output(KV.of(word, documentFrequency));
                }
              }).withSideInputs(totalDocuments));

      // Join the term frequency and document frequency
      // collections, each keyed on the word.
      final TupleTag<KV<URI, Double>> tfTag = new TupleTag<KV<URI, Double>>();
      final TupleTag<Double> dfTag = new TupleTag<Double>();
      PCollection<KV<String, CoGbkResult>> wordToUriAndTfAndDf = KeyedPCollectionTuple
          .of(tfTag, wordToUriAndTf)
          .and(dfTag, wordToDf)
          .apply(CoGroupByKey.<String>create());

      // Compute a mapping from each word to a (URI, TF-IDF) score
      // for each URI. There are a variety of definitions of TF-IDF
      // ("term frequency - inverse document frequency") score;
      // here we use a basic version that is the term frequency
      // divided by the log of the document frequency.
      PCollection<KV<String, KV<URI, Double>>> wordToUriAndTfIdf = wordToUriAndTfAndDf
          .apply("ComputeTfIdf", ParDo.of(
              new DoFn<KV<String, CoGbkResult>, KV<String, KV<URI, Double>>>() {
                @ProcessElement
                public void processElement(ProcessContext c) {
                  String word = c.element().getKey();
                  Double df = c.element().getValue().getOnly(dfTag);

                  for (KV<URI, Double> uriAndTf : c.element().getValue().getAll(tfTag)) {
                    URI uri = uriAndTf.getKey();
                    Double tf = uriAndTf.getValue();
                    Double tfIdf = tf * Math.log(1 / df);
                    c.output(KV.of(word, KV.of(uri, tfIdf)));
                  }
                }
              }));

      return wordToUriAndTfIdf;
    }

    // Instantiate Logger.
    // It is suggested that the user specify the class name of the containing class
    // (in this case ComputeTfIdf).
    private static final Logger LOG = LoggerFactory.getLogger(ComputeTfIdf.class);
  }

  /**
   * A {@link PTransform} to write, in CSV format, a mapping from term and URI
   * to score.
   */
  public static class WriteTfIdf
      extends PTransform<PCollection<KV<String, KV<URI, Double>>>, PDone> {
    private String output;

    public WriteTfIdf(String output) {
      this.output = output;
    }

    @Override
    public PDone expand(PCollection<KV<String, KV<URI, Double>>> wordToUriAndTfIdf) {
      return wordToUriAndTfIdf
          .apply("Format", ParDo.of(new DoFn<KV<String, KV<URI, Double>>, String>() {
            @ProcessElement
            public void processElement(ProcessContext c) {
              c.output(String.format("%s,\t%s,\t%f",
                  c.element().getKey(),
                  c.element().getValue().getKey(),
                  c.element().getValue().getValue()));
            }
          }))
          .apply(TextIO.write()
              .to(output)
              .withSuffix(".csv"));
    }
  }

  public static void main(String[] args) throws Exception {
    Options options = PipelineOptionsFactory.fromArgs(args).withValidation().as(Options.class);
    Pipeline pipeline = Pipeline.create(options);
    pipeline.getCoderRegistry().registerCoderForClass(URI.class, StringDelegateCoder.of(URI.class));

    pipeline
        .apply(new ReadDocuments(listInputDocuments(options)))
        .apply(new ComputeTfIdf())
        .apply(new WriteTfIdf(options.getOutput()));

    pipeline.run().waitUntilFinish();
  }
}