Java tutorial
/* * 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 static com.google.datastore.v1.client.DatastoreHelper.makeKey; import static com.google.datastore.v1.client.DatastoreHelper.makeValue; import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument; import com.google.api.services.bigquery.model.TableFieldSchema; import com.google.api.services.bigquery.model.TableReference; import com.google.api.services.bigquery.model.TableRow; import com.google.api.services.bigquery.model.TableSchema; import com.google.datastore.v1.Entity; import com.google.datastore.v1.Key; import com.google.datastore.v1.Value; import java.io.IOException; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.regex.Matcher; import java.util.regex.Pattern; import org.apache.beam.examples.common.ExampleBigQueryTableOptions; import org.apache.beam.examples.common.ExampleOptions; import org.apache.beam.examples.common.ExampleUtils; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.coders.AvroCoder; import org.apache.beam.sdk.coders.DefaultCoder; import org.apache.beam.sdk.io.TextIO; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO; import org.apache.beam.sdk.io.gcp.datastore.DatastoreIO; import org.apache.beam.sdk.options.Default; import org.apache.beam.sdk.options.Description; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.options.StreamingOptions; import org.apache.beam.sdk.options.Validation; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.transforms.Count; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.Filter; import org.apache.beam.sdk.transforms.Flatten; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.Partition; import org.apache.beam.sdk.transforms.Partition.PartitionFn; import org.apache.beam.sdk.transforms.SerializableFunction; import org.apache.beam.sdk.transforms.Sum; import org.apache.beam.sdk.transforms.Top; import org.apache.beam.sdk.transforms.windowing.GlobalWindows; import org.apache.beam.sdk.transforms.windowing.SlidingWindows; import org.apache.beam.sdk.transforms.windowing.Window; import org.apache.beam.sdk.transforms.windowing.WindowFn; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionList; import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.MoreObjects; import org.joda.time.Duration; /** * An example that computes the most popular hash tags for every prefix, which can be used for * auto-completion. * * <p>Concepts: Using the same pipeline in both streaming and batch, combiners, composite * transforms. * * <p>To execute this pipeline in streaming mode, specify: * * <pre>{@code * --streaming * }</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>This will update the Cloud Datastore every 10 seconds based on the last 30 minutes of data * received. */ public class AutoComplete { /** * A PTransform that takes as input a list of tokens and returns the most common tokens per * prefix. */ public static class ComputeTopCompletions extends PTransform<PCollection<String>, PCollection<KV<String, List<CompletionCandidate>>>> { private final int candidatesPerPrefix; private final boolean recursive; protected ComputeTopCompletions(int candidatesPerPrefix, boolean recursive) { this.candidatesPerPrefix = candidatesPerPrefix; this.recursive = recursive; } public static ComputeTopCompletions top(int candidatesPerPrefix, boolean recursive) { return new ComputeTopCompletions(candidatesPerPrefix, recursive); } @Override public PCollection<KV<String, List<CompletionCandidate>>> expand(PCollection<String> input) { PCollection<CompletionCandidate> candidates = input // First count how often each token appears. .apply(Count.perElement()) // Map the KV outputs of Count into our own CompletionCandiate class. .apply("CreateCompletionCandidates", ParDo.of(new DoFn<KV<String, Long>, CompletionCandidate>() { @ProcessElement public void processElement(ProcessContext c) { c.output(new CompletionCandidate(c.element().getKey(), c.element().getValue())); } })); // Compute the top via either a flat or recursive algorithm. if (recursive) { return candidates.apply(new ComputeTopRecursive(candidatesPerPrefix, 1)) .apply(Flatten.pCollections()); } else { return candidates.apply(new ComputeTopFlat(candidatesPerPrefix, 1)); } } } /** Lower latency, but more expensive. */ private static class ComputeTopFlat extends PTransform<PCollection<CompletionCandidate>, PCollection<KV<String, List<CompletionCandidate>>>> { private final int candidatesPerPrefix; private final int minPrefix; public ComputeTopFlat(int candidatesPerPrefix, int minPrefix) { this.candidatesPerPrefix = candidatesPerPrefix; this.minPrefix = minPrefix; } @Override public PCollection<KV<String, List<CompletionCandidate>>> expand(PCollection<CompletionCandidate> input) { return input // For each completion candidate, map it to all prefixes. .apply(ParDo.of(new AllPrefixes(minPrefix))) // Find and return the top candiates for each prefix. .apply(Top.<String, CompletionCandidate>largestPerKey(candidatesPerPrefix) .withHotKeyFanout(new HotKeyFanout())); } private static class HotKeyFanout implements SerializableFunction<String, Integer> { @Override public Integer apply(String input) { return (int) Math.pow(4, 5 - input.length()); } } } /** * Cheaper but higher latency. * * <p>Returns two PCollections, the first is top prefixes of size greater than minPrefix, and the * second is top prefixes of size exactly minPrefix. */ private static class ComputeTopRecursive extends PTransform<PCollection<CompletionCandidate>, PCollectionList<KV<String, List<CompletionCandidate>>>> { private final int candidatesPerPrefix; private final int minPrefix; public ComputeTopRecursive(int candidatesPerPrefix, int minPrefix) { this.candidatesPerPrefix = candidatesPerPrefix; this.minPrefix = minPrefix; } private class KeySizePartitionFn implements PartitionFn<KV<String, List<CompletionCandidate>>> { @Override public int partitionFor(KV<String, List<CompletionCandidate>> elem, int numPartitions) { return elem.getKey().length() > minPrefix ? 0 : 1; } } private static class FlattenTops extends DoFn<KV<String, List<CompletionCandidate>>, CompletionCandidate> { @ProcessElement public void processElement(ProcessContext c) { for (CompletionCandidate cc : c.element().getValue()) { c.output(cc); } } } @Override public PCollectionList<KV<String, List<CompletionCandidate>>> expand( PCollection<CompletionCandidate> input) { if (minPrefix > 10) { // Base case, partitioning to return the output in the expected format. return input.apply(new ComputeTopFlat(candidatesPerPrefix, minPrefix)) .apply(Partition.of(2, new KeySizePartitionFn())); } else { // If a candidate is in the top N for prefix a...b, it must also be in the top // N for a...bX for every X, which is typlically a much smaller set to consider. // First, compute the top candidate for prefixes of size at least minPrefix + 1. PCollectionList<KV<String, List<CompletionCandidate>>> larger = input .apply(new ComputeTopRecursive(candidatesPerPrefix, minPrefix + 1)); // Consider the top candidates for each prefix of length minPrefix + 1... PCollection<KV<String, List<CompletionCandidate>>> small = PCollectionList .of(larger.get(1).apply(ParDo.of(new FlattenTops()))) // ...together with those (previously excluded) candidates of length // exactly minPrefix... .and(input.apply(Filter.by(c -> c.getValue().length() == minPrefix))) .apply("FlattenSmall", Flatten.pCollections()) // ...set the key to be the minPrefix-length prefix... .apply(ParDo.of(new AllPrefixes(minPrefix, minPrefix))) // ...and (re)apply the Top operator to all of them together. .apply(Top.largestPerKey(candidatesPerPrefix)); PCollection<KV<String, List<CompletionCandidate>>> flattenLarger = larger.apply("FlattenLarge", Flatten.pCollections()); return PCollectionList.of(flattenLarger).and(small); } } } /** A DoFn that keys each candidate by all its prefixes. */ private static class AllPrefixes extends DoFn<CompletionCandidate, KV<String, CompletionCandidate>> { private final int minPrefix; private final int maxPrefix; public AllPrefixes(int minPrefix) { this(minPrefix, Integer.MAX_VALUE); } public AllPrefixes(int minPrefix, int maxPrefix) { this.minPrefix = minPrefix; this.maxPrefix = maxPrefix; } @ProcessElement public void processElement(ProcessContext c) { String word = c.element().value; for (int i = minPrefix; i <= Math.min(word.length(), maxPrefix); i++) { c.output(KV.of(word.substring(0, i), c.element())); } } } /** Class used to store tag-count pairs. */ @DefaultCoder(AvroCoder.class) static class CompletionCandidate implements Comparable<CompletionCandidate> { private long count; private String value; public CompletionCandidate(String value, long count) { this.value = value; this.count = count; } public long getCount() { return count; } public String getValue() { return value; } // Empty constructor required for Avro decoding. public CompletionCandidate() { } @Override public int compareTo(CompletionCandidate o) { if (this.count < o.count) { return -1; } else if (this.count == o.count) { return this.value.compareTo(o.value); } else { return 1; } } @Override public boolean equals(Object other) { if (other instanceof CompletionCandidate) { CompletionCandidate that = (CompletionCandidate) other; return this.count == that.count && this.value.equals(that.value); } else { return false; } } @Override public int hashCode() { return Long.valueOf(count).hashCode() ^ value.hashCode(); } @Override public String toString() { return "CompletionCandidate[" + value + ", " + count + "]"; } } /** Takes as input a set of strings, and emits each #hashtag found therein. */ static class ExtractHashtags extends DoFn<String, String> { @ProcessElement public void processElement(ProcessContext c) { Matcher m = Pattern.compile("#\\S+").matcher(c.element()); while (m.find()) { c.output(m.group().substring(1)); } } } static class FormatForBigquery extends DoFn<KV<String, List<CompletionCandidate>>, TableRow> { @ProcessElement public void processElement(ProcessContext c) { List<TableRow> completions = new ArrayList<>(); for (CompletionCandidate cc : c.element().getValue()) { completions.add(new TableRow().set("count", cc.getCount()).set("tag", cc.getValue())); } TableRow row = new TableRow().set("prefix", c.element().getKey()).set("tags", completions); c.output(row); } /** Defines the BigQuery schema used for the output. */ static TableSchema getSchema() { List<TableFieldSchema> tagFields = new ArrayList<>(); tagFields.add(new TableFieldSchema().setName("count").setType("INTEGER")); tagFields.add(new TableFieldSchema().setName("tag").setType("STRING")); List<TableFieldSchema> fields = new ArrayList<>(); fields.add(new TableFieldSchema().setName("prefix").setType("STRING")); fields.add(new TableFieldSchema().setName("tags").setType("RECORD").setMode("REPEATED") .setFields(tagFields)); return new TableSchema().setFields(fields); } } /** * Takes as input a the top candidates per prefix, and emits an entity suitable for writing to * Cloud Datastore. * * <p>Note: We use ancestor keys for strong consistency. See the Cloud Datastore documentation on * <a href="https://cloud.google.com/datastore/docs/concepts/structuring_for_strong_consistency"> * Structuring Data for Strong Consistency</a> */ static class FormatForDatastore extends DoFn<KV<String, List<CompletionCandidate>>, Entity> { private String kind; private String ancestorKey; public FormatForDatastore(String kind, String ancestorKey) { this.kind = kind; this.ancestorKey = ancestorKey; } @ProcessElement public void processElement(ProcessContext c) { Entity.Builder entityBuilder = Entity.newBuilder(); Key key = makeKey(makeKey(kind, ancestorKey).build(), kind, c.element().getKey()).build(); entityBuilder.setKey(key); List<Value> candidates = new ArrayList<>(); Map<String, Value> properties = new HashMap<>(); for (CompletionCandidate tag : c.element().getValue()) { Entity.Builder tagEntity = Entity.newBuilder(); properties.put("tag", makeValue(tag.value).build()); properties.put("count", makeValue(tag.count).build()); candidates.add(makeValue(tagEntity).build()); } properties.put("candidates", makeValue(candidates).build()); entityBuilder.putAllProperties(properties); c.output(entityBuilder.build()); } } /** * Options supported by this class. * * <p>Inherits standard Beam example configuration options. */ public interface Options extends ExampleOptions, ExampleBigQueryTableOptions, StreamingOptions { @Description("Input text file") @Validation.Required String getInputFile(); void setInputFile(String value); @Description("Whether to use the recursive algorithm") @Default.Boolean(true) Boolean getRecursive(); void setRecursive(Boolean value); @Description("Cloud Datastore entity kind") @Default.String("autocomplete-demo") String getKind(); void setKind(String value); @Description("Whether output to BigQuery") @Default.Boolean(true) Boolean getOutputToBigQuery(); void setOutputToBigQuery(Boolean value); @Description("Whether to send output to checksum Transform.") @Default.Boolean(true) Boolean getOutputToChecksum(); void setOutputToChecksum(Boolean value); @Description("Expected result of the checksum transform.") Long getExpectedChecksum(); void setExpectedChecksum(Long value); @Description("Whether output to Cloud Datastore") @Default.Boolean(false) Boolean getOutputToDatastore(); void setOutputToDatastore(Boolean value); @Description("Cloud Datastore ancestor key") @Default.String("root") String getDatastoreAncestorKey(); void setDatastoreAncestorKey(String value); @Description("Cloud Datastore output project ID, defaults to project ID") String getOutputProject(); void setOutputProject(String value); } public static void runAutocompletePipeline(Options options) throws IOException { options.setBigQuerySchema(FormatForBigquery.getSchema()); ExampleUtils exampleUtils = new ExampleUtils(options); // We support running the same pipeline in either // batch or windowed streaming mode. WindowFn<Object, ?> windowFn; if (options.isStreaming()) { checkArgument(!options.getOutputToDatastore(), "DatastoreIO is not supported in streaming."); windowFn = SlidingWindows.of(Duration.standardMinutes(30)).every(Duration.standardSeconds(5)); } else { windowFn = new GlobalWindows(); } // Create the pipeline. Pipeline p = Pipeline.create(options); PCollection<KV<String, List<CompletionCandidate>>> toWrite = p .apply(TextIO.read().from(options.getInputFile())).apply(ParDo.of(new ExtractHashtags())) .apply(Window.into(windowFn)).apply(ComputeTopCompletions.top(10, options.getRecursive())); if (options.getOutputToDatastore()) { toWrite.apply("FormatForDatastore", ParDo.of(new FormatForDatastore(options.getKind(), options.getDatastoreAncestorKey()))) .apply(DatastoreIO.v1().write().withProjectId( MoreObjects.firstNonNull(options.getOutputProject(), options.getProject()))); } if (options.getOutputToBigQuery()) { exampleUtils.setupBigQueryTable(); TableReference tableRef = new TableReference(); tableRef.setProjectId(options.getProject()); tableRef.setDatasetId(options.getBigQueryDataset()); tableRef.setTableId(options.getBigQueryTable()); toWrite.apply(ParDo.of(new FormatForBigquery())) .apply(BigQueryIO.writeTableRows().to(tableRef).withSchema(FormatForBigquery.getSchema()) .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED) .withWriteDisposition( options.isStreaming() ? BigQueryIO.Write.WriteDisposition.WRITE_APPEND : BigQueryIO.Write.WriteDisposition.WRITE_TRUNCATE)); } if (options.getOutputToChecksum()) { PCollection<Long> checksum = toWrite .apply(ParDo.of(new DoFn<KV<String, List<CompletionCandidate>>, Long>() { @ProcessElement public void process(ProcessContext c) { KV<String, List<CompletionCandidate>> elm = c.element(); Long listHash = c.element().getValue().stream().mapToLong(cc -> cc.hashCode()).sum(); c.output(Long.valueOf(elm.getKey().hashCode()) + listHash); } })).apply(Sum.longsGlobally()); PAssert.that(checksum).containsInAnyOrder(options.getExpectedChecksum()); } // Run the pipeline. PipelineResult result = p.run(); // ExampleUtils will try to cancel the pipeline and the injector before the program exists. exampleUtils.waitToFinish(result); } public static void main(String[] args) throws IOException { Options options = PipelineOptionsFactory.fromArgs(args).withValidation().as(Options.class); runAutocompletePipeline(options); } }