Google Git
Sign in
apache / beam / 5a851b734f53206c20efe08d93d15760bbc15b0c / . / sdks / java / core / src / main / java / org / apache / beam / sdk / transforms / ParDo.java
blob: fb8524aae969ba9530c2ad32c2020fcd91219cc2 [file] [log] [blame]
/*
* 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.transforms;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
import java.io.Serializable;
import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Type;
import java.util.Arrays;
import java.util.Collections;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;
import javax.annotation.Nullable;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineRunner;
import org.apache.beam.sdk.annotations.Internal;
import org.apache.beam.sdk.coders.CannotProvideCoderException;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderRegistry;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.schemas.FieldAccessDescriptor;
import org.apache.beam.sdk.schemas.NoSuchSchemaException;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.SchemaCoder;
import org.apache.beam.sdk.schemas.SchemaRegistry;
import org.apache.beam.sdk.schemas.utils.ConvertHelpers;
import org.apache.beam.sdk.schemas.utils.SelectHelpers;
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.transforms.DoFn.WindowedContext;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.display.DisplayData.Builder;
import org.apache.beam.sdk.transforms.display.DisplayData.ItemSpec;
import org.apache.beam.sdk.transforms.display.HasDisplayData;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.FieldAccessDeclaration;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.MethodWithExtraParameters;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.OnTimerMethod;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.SchemaElementParameter;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.SideInputParameter;
import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.util.NameUtils;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PCollectionViews;
import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
/**
* {@link ParDo} is the core element-wise transform in Apache Beam, invoking a user-specified
* function on each of the elements of the input {@link PCollection} to produce zero or more output
* elements, all of which are collected into the output {@link PCollection}.
*
* <p>Elements are processed independently, and possibly in parallel across distributed cloud
* resources.
*
* <p>The {@link ParDo} processing style is similar to what happens inside the "Mapper" or "Reducer"
* class of a MapReduce-style algorithm.
*
* <h2>{@link DoFn DoFns}</h2>
*
* <p>The function to use to process each element is specified by a {@link DoFn DoFn&lt;InputT,
* OutputT&gt;}, primarily via its {@link DoFn.ProcessElement ProcessElement} method. The {@link
* DoFn} may also provide a {@link DoFn.StartBundle StartBundle} and {@link DoFn.FinishBundle
* finishBundle} method.
*
* <p>Conceptually, when a {@link ParDo} transform is executed, the elements of the input {@link
* PCollection} are first divided up into some number of "bundles". These are farmed off to
* distributed worker machines (or run locally, if using the {@code DirectRunner}). For each bundle
* of input elements processing proceeds as follows:
*
* <ol>
* <li>If required, a fresh instance of the argument {@link DoFn} is created on a worker, and the
* {@link DoFn.Setup} method is called on this instance. This may be through deserialization
* or other means. A {@link PipelineRunner} may reuse {@link DoFn} instances for multiple
* bundles. A {@link DoFn} that has terminated abnormally (by throwing an {@link Exception})
* will never be reused.
* <li>The {@link DoFn DoFn's} {@link DoFn.StartBundle} method, if provided, is called to
* initialize it.
* <li>The {@link DoFn DoFn's} {@link DoFn.ProcessElement} method is called on each of the input
* elements in the bundle.
* <li>The {@link DoFn DoFn's} {@link DoFn.FinishBundle} method, if provided, is called to
* complete its work. After {@link DoFn.FinishBundle} is called, the framework will not again
* invoke {@link DoFn.ProcessElement} or {@link DoFn.FinishBundle} until a new call to {@link
* DoFn.StartBundle} has occurred.
* <li>If any of {@link DoFn.Setup}, {@link DoFn.StartBundle}, {@link DoFn.ProcessElement} or
* {@link DoFn.FinishBundle} methods throw an exception, the {@link DoFn.Teardown} method, if
* provided, will be called on the {@link DoFn} instance.
* <li>If a runner will no longer use a {@link DoFn}, the {@link DoFn.Teardown} method, if
* provided, will be called on the discarded instance.
* <li>If a bundle requested bundle finalization by registering a {@link
* DoFn.BundleFinalizer.Callback bundle finalization callback}, the callback will be invoked
* after the runner has successfully committed the output of a successful bundle.
* </ol>
*
* <p>Note also that calls to {@link DoFn.Teardown} are best effort, and may not be called before a
* {@link DoFn} is discarded in the general case. As a result, use of the {@link DoFn.Teardown}
* method to perform side effects is not appropriate, because the elements that produced the side
* effect will not be replayed in case of failure, and those side effects are permanently lost.
*
* <p>Each of the calls to any of the {@link DoFn DoFn's} processing methods can produce zero or
* more output elements. All of the of output elements from all of the {@link DoFn} instances are
* included in an output {@link PCollection}.
*
* <p>For example:
*
* <pre>{@code PCollection<String> lines = ...;
* PCollection<String> words =
* lines.apply(ParDo.of(new DoFn<String, String>() }{
* {@code @ProcessElement
* public void processElement(@Element String line,
* OutputReceiver<String> r) {
* for (String word : line.split("[^a-zA-Z']+")) {
* r.output(word);
* }
* }}}));
* {@code PCollection<Integer> wordLengths =
* words.apply(ParDo.of(new DoFn<String, Integer>() }{
* {@code @ProcessElement
* public void processElement(@Element String word,
* OutputReceiver<Integer> r) {
* Integer length = word.length();
* r.output(length);
* }}}));
* </pre>
*
* <p>Each output element has the same timestamp and is in the same windows as its corresponding
* input element, and the output {@code PCollection} has the same {@link WindowFn} associated with
* it as the input.
*
* <h2>Naming {@link ParDo ParDo} transforms</h2>
*
* <p>The name of a transform is used to provide a name for any node in the {@link Pipeline} graph
* resulting from application of the transform. It is best practice to provide a name at the time of
* application, via {@link PCollection#apply(String, PTransform)}. Otherwise, a unique name - which
* may not be stable across pipeline revision - will be generated, based on the transform name.
*
* <p>For example:
*
* <pre>{@code PCollection<String> words =
* lines.apply("ExtractWords", ParDo.of(new DoFn<String, String>() { ... }));
* PCollection<Integer> wordLengths =
* words.apply("ComputeWordLengths", ParDo.of(new DoFn<String, Integer>() { ... }));
* }</pre>
*
* <h2>Side Inputs</h2>
*
* <p>While a {@link ParDo} processes elements from a single "main input" {@link PCollection}, it
* can take additional "side input" {@link PCollectionView PCollectionViews}. These side input
* {@link PCollectionView PCollectionViews} express styles of accessing {@link PCollection
* PCollections} computed by earlier pipeline operations, passed in to the {@link ParDo} transform
* using {@link SingleOutput#withSideInputs}, and their contents accessible to each of the {@link
* DoFn} operations via {@link DoFn.ProcessContext#sideInput sideInput}. For example:
*
* <pre>{@code PCollection<String> words = ...;
* PCollection<Integer> maxWordLengthCutOff = ...; // Singleton PCollection
* final PCollectionView<Integer> maxWordLengthCutOffView =
* maxWordLengthCutOff.apply(View.<Integer>asSingleton());
* PCollection<String> wordsBelowCutOff =
* words.apply(ParDo.of(new DoFn<String, String>() }{
* {@code @ProcessElement
* public void processElement(ProcessContext c) {
* String word = c.element();
* int lengthCutOff = c.sideInput(maxWordLengthCutOffView);
* if (word.length() <= lengthCutOff) {
* c.output(word);
* }
* }}}).withSideInputs(maxWordLengthCutOffView));
* </pre>
*
* <h2>Additional Outputs</h2>
*
* <p>Optionally, a {@link ParDo} transform can produce multiple output {@link PCollection
* PCollections}, both a "main output" {@code PCollection<OutputT>} plus any number of additional
* output {@link PCollection PCollections}, each keyed by a distinct {@link TupleTag}, and bundled
* in a {@link PCollectionTuple}. The {@link TupleTag TupleTags} to be used for the output {@link
* PCollectionTuple} are specified by invoking {@link SingleOutput#withOutputTags}. Unconsumed
* outputs do not necessarily need to be explicitly specified, even if the {@link DoFn} generates
* them. Within the {@link DoFn}, an element is added to the main output {@link PCollection} as
* normal, using {@link WindowedContext#output(Object)}, while an element is added to any additional
* output {@link PCollection} using {@link WindowedContext#output(TupleTag, Object)}. For example:
*
* <pre>{@code PCollection<String> words = ...;
* // Select words whose length is below a cut off,
* // plus the lengths of words that are above the cut off.
* // Also select words starting with "MARKER".
* final int wordLengthCutOff = 10;
* // Create tags to use for the main and additional outputs.
* final TupleTag<String> wordsBelowCutOffTag =
* new TupleTag<String>(){};
* final TupleTag<Integer> wordLengthsAboveCutOffTag =
* new TupleTag<Integer>(){};
* final TupleTag<String> markedWordsTag =
* new TupleTag<String>(){};
* PCollectionTuple results =
* words.apply(
* ParDo
* .of(new DoFn<String, String>() {
* // Create a tag for the unconsumed output.
* final TupleTag<String> specialWordsTag =
* new TupleTag<String>(){};}}
* {@code @ProcessElement
* public void processElement(@Element String word, MultiOutputReceiver r) {
* if (word.length() <= wordLengthCutOff) {
* // Emit this short word to the main output.
* r.get(wordsBelowCutOffTag).output(word);
* } else {
* // Emit this long word's length to a specified output.
* r.get(wordLengthsAboveCutOffTag).output(word.length());
* }
* if (word.startsWith("MARKER")) {
* // Emit this word to a different specified output.
* r.get(markedWordsTag).output(word);
* }
* if (word.startsWith("SPECIAL")) {
* // Emit this word to the unconsumed output.
* r.get(specialWordsTag).output(word);
* }
* }}})
* // Specify the main and consumed output tags of the
* // PCollectionTuple result:
* .withOutputTags(wordsBelowCutOffTag,
* TupleTagList.of(wordLengthsAboveCutOffTag)
* .and(markedWordsTag)));
* // Extract the PCollection results, by tag.{@code
* PCollection<String> wordsBelowCutOff =
* results.get(wordsBelowCutOffTag);
* PCollection<Integer> wordLengthsAboveCutOff =
* results.get(wordLengthsAboveCutOffTag);
* PCollection<String> markedWords =
* results.get(markedWordsTag);
* }</pre>
*
* <h2>Output Coders</h2>
*
* <p>By default, the {@link Coder Coder&lt;OutputT&gt;} for the elements of the main output {@link
* PCollection PCollection&lt;OutputT&gt;} is inferred from the concrete type of the {@link DoFn
* DoFn&lt;InputT, OutputT&gt;}.
*
* <p>By default, the {@link Coder Coder&lt;AdditionalOutputT&gt;} for the elements of an output
* {@link PCollection PCollection&lt;AdditionalOutputT&gt;} is inferred from the concrete type of
* the corresponding {@link TupleTag TupleTag&lt;AdditionalOutputT&gt;}. To be successful, the
* {@link TupleTag} should be created as an instance of a trivial anonymous subclass, with {@code
* {}} suffixed to the constructor call. Such uses block Java's generic type parameter inference, so
* the {@code <X>} argument must be provided explicitly. For example:
*
* <pre>{@code
* // A TupleTag to use for a side input can be written concisely:
* final TupleTag<Integer> sideInputag = new TupleTag<>();
* // A TupleTag to use for an output should be written with "{}",
* // and explicit generic parameter type:
* final TupleTag<String> additionalOutputTag = new TupleTag<String>(){};
* }</pre>
*
* This style of {@code TupleTag} instantiation is used in the example of {@link ParDo ParDos} that
* produce multiple outputs, above.
*
* <h2>Serializability of {@link DoFn DoFns}</h2>
*
* <p>A {@link DoFn} passed to a {@link ParDo} transform must be {@link Serializable}. This allows
* the {@link DoFn} instance created in this "main program" to be sent (in serialized form) to
* remote worker machines and reconstituted for bundles of elements of the input {@link PCollection}
* being processed. A {@link DoFn} can have instance variable state, and non-transient instance
* variable state will be serialized in the main program and then deserialized on remote worker
* machines for some number of bundles of elements to process.
*
* <p>{@link DoFn DoFns} expressed as anonymous inner classes can be convenient, but due to a quirk
* in Java's rules for serializability, non-static inner or nested classes (including anonymous
* inner classes) automatically capture their enclosing class's instance in their serialized state.
* This can lead to including much more than intended in the serialized state of a {@link DoFn}, or
* even things that aren't {@link Serializable}.
*
* <p>There are two ways to avoid unintended serialized state in a {@link DoFn}:
*
* <ul>
* <li>Define the {@link DoFn} as a named, static class.
* <li>Define the {@link DoFn} as an anonymous inner class inside of a static method.
* </ul>
*
* <p>Both of these approaches ensure that there is no implicit enclosing instance serialized along
* with the {@link DoFn} instance.
*
* <p>Prior to Java 8, any local variables of the enclosing method referenced from within an
* anonymous inner class need to be marked as {@code final}. If defining the {@link DoFn} as a named
* static class, such variables would be passed as explicit constructor arguments and stored in
* explicit instance variables.
*
* <p>There are three main ways to initialize the state of a {@link DoFn} instance processing a
* bundle:
*
* <ul>
* <li>Define instance variable state (including implicit instance variables holding final
* variables captured by an anonymous inner class), initialized by the {@link DoFn}'s
* constructor (which is implicit for an anonymous inner class). This state will be
* automatically serialized and then deserialized in the {@link DoFn} instances created for
* bundles. This method is good for state known when the original {@link DoFn} is created in
* the main program, if it's not overly large. This is not suitable for any state which must
* only be used for a single bundle, as {@link DoFn DoFn's} may be used to process multiple
* bundles.
* <li>Compute the state as a singleton {@link PCollection} and pass it in as a side input to the
* {@link DoFn}. This is good if the state needs to be computed by the pipeline, or if the
* state is very large and so is best read from file(s) rather than sent as part of the {@link
* DoFn DoFn's} serialized state.
* <li>Initialize the state in each {@link DoFn} instance, in a {@link DoFn.StartBundle} method.
* This is good if the initialization doesn't depend on any information known only by the main
* program or computed by earlier pipeline operations, but is the same for all instances of
* this {@link DoFn} for all program executions, say setting up empty caches or initializing
* constant data.
* </ul>
*
* <h2>No Global Shared State</h2>
*
* <p>{@link ParDo} operations are intended to be able to run in parallel across multiple worker
* machines. This precludes easy sharing and updating mutable state across those machines. There is
* no support in the Beam model for communicating and synchronizing updates to shared state across
* worker machines, so programs should not access any mutable static variable state in their {@link
* DoFn}, without understanding that the Java processes for the main program and workers will each
* have its own independent copy of such state, and there won't be any automatic copying of that
* state across Java processes. All information should be communicated to {@link DoFn} instances via
* main and side inputs and serialized state, and all output should be communicated from a {@link
* DoFn} instance via output {@link PCollection PCollections}, in the absence of external
* communication mechanisms written by user code.
*
* <h2>Fault Tolerance</h2>
*
* <p>In a distributed system, things can fail: machines can crash, machines can be unable to
* communicate across the network, etc. While individual failures are rare, the larger the job, the
* greater the chance that something, somewhere, will fail. Beam runners may strive to mask such
* failures by retrying failed {@link DoFn} bundle. This means that a {@link DoFn} instance might
* process a bundle partially, then crash for some reason, then be rerun (often in a new JVM) on
* that same bundle and on the same elements as before. Sometimes two or more {@link DoFn} instances
* will be running on the same bundle simultaneously, with the system taking the results of the
* first instance to complete successfully. Consequently, the code in a {@link DoFn} needs to be
* written such that these duplicate (sequential or concurrent) executions do not cause problems. If
* the outputs of a {@link DoFn} are a pure function of its inputs, then this requirement is
* satisfied. However, if a {@link DoFn DoFn's} execution has external side-effects, such as
* performing updates to external HTTP services, then the {@link DoFn DoFn's} code needs to take
* care to ensure that those updates are idempotent and that concurrent updates are acceptable. This
* property can be difficult to achieve, so it is advisable to strive to keep {@link DoFn DoFns} as
* pure functions as much as possible.
*
* <h2>Optimization</h2>
*
* <p>Beam runners may choose to apply optimizations to a pipeline before it is executed. A key
* optimization, <i>fusion</i>, relates to {@link ParDo} operations. If one {@link ParDo} operation
* produces a {@link PCollection} that is then consumed as the main input of another {@link ParDo}
* operation, the two {@link ParDo} operations will be <i>fused</i> together into a single ParDo
* operation and run in a single pass; this is "producer-consumer fusion". Similarly, if two or more
* ParDo operations have the same {@link PCollection} main input, they will be fused into a single
* {@link ParDo} that makes just one pass over the input {@link PCollection}; this is "sibling
* fusion".
*
* <p>If after fusion there are no more unfused references to a {@link PCollection} (e.g., one
* between a producer ParDo and a consumer {@link ParDo}), the {@link PCollection} itself is "fused
* away" and won't ever be written to disk, saving all the I/O and space expense of constructing it.
*
* <p>When Beam runners apply fusion optimization, it is essentially "free" to write {@link ParDo}
* operations in a very modular, composable style, each {@link ParDo} operation doing one clear
* task, and stringing together sequences of {@link ParDo} operations to get the desired overall
* effect. Such programs can be easier to understand, easier to unit-test, easier to extend and
* evolve, and easier to reuse in new programs. The predefined library of PTransforms that come with
* Beam makes heavy use of this modular, composable style, trusting to the runner to "flatten out"
* all the compositions into highly optimized stages.
*
* @see <a href= "https://beam.apache.org/documentation/programming-guide/#pardo"> the web
* documentation for ParDo</a>
*/
public class ParDo {
/**
* Creates a {@link ParDo} {@link PTransform} that will invoke the given {@link DoFn} function.
*
* <p>The resulting {@link PTransform PTransform} is ready to be applied, or further properties
* can be set on it first.
*/
public static <InputT, OutputT> SingleOutput<InputT, OutputT> of(DoFn<InputT, OutputT> fn) {
validate(fn);
return new SingleOutput<>(fn, Collections.emptyMap(), displayDataForFn(fn));
}
private static <T> DisplayData.ItemSpec<? extends Class<?>> displayDataForFn(T fn) {
return DisplayData.item("fn", fn.getClass()).withLabel("Transform Function");
}
private static void finishSpecifyingStateSpecs(
DoFn<?, ?> fn, CoderRegistry coderRegistry, Coder<?> inputCoder) {
DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());
Map<String, DoFnSignature.StateDeclaration> stateDeclarations = signature.stateDeclarations();
for (DoFnSignature.StateDeclaration stateDeclaration : stateDeclarations.values()) {
try {
StateSpec<?> stateSpec = (StateSpec<?>) stateDeclaration.field().get(fn);
stateSpec.offerCoders(codersForStateSpecTypes(stateDeclaration, coderRegistry, inputCoder));
stateSpec.finishSpecifying();
} catch (IllegalAccessException e) {
throw new RuntimeException(e);
}
}
}
private static void validateStateApplicableForInput(DoFn<?, ?> fn, PCollection<?> input) {
Coder<?> inputCoder = input.getCoder();
checkArgument(
inputCoder instanceof KvCoder,
"%s requires its input to use %s in order to use state and timers.",
ParDo.class.getSimpleName(),
KvCoder.class.getSimpleName());
KvCoder<?, ?> kvCoder = (KvCoder<?, ?>) inputCoder;
try {
kvCoder.getKeyCoder().verifyDeterministic();
} catch (Coder.NonDeterministicException exc) {
throw new IllegalArgumentException(
String.format(
"%s requires a deterministic key coder in order to use state and timers",
ParDo.class.getSimpleName()));
}
}
private static void validateSideInputTypes(
Map<String, PCollectionView<?>> sideInputs, DoFn<?, ?> fn) {
DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());
DoFnSignature.ProcessElementMethod processElementMethod = signature.processElement();
for (SideInputParameter sideInput : processElementMethod.getSideInputParameters()) {
PCollectionView<?> view = sideInputs.get(sideInput.sideInputId());
checkArgument(
view != null,
"the ProcessElement method expects a side input identified with the tag %s, but no such side input was"
+ " supplied. Use withSideInput(String, PCollectionView) to supply this side input.",
sideInput.sideInputId());
TypeDescriptor<?> viewType = view.getViewFn().getTypeDescriptor();
// Currently check that the types exactly match, even if the types are convertible.
checkArgument(
viewType.equals(sideInput.elementT()),
"Side Input with tag %s and type %s cannot be bound to ProcessElement parameter with type %s",
sideInput.sideInputId(),
viewType,
sideInput.elementT());
}
}
private static FieldAccessDescriptor getFieldAccessDescriptorFromParameter(
@Nullable String fieldAccessString,
Schema inputSchema,
Map<String, FieldAccessDeclaration> fieldAccessDeclarations,
DoFn<?, ?> fn) {
// Resolve the FieldAccessDescriptor against the Schema.
// This will be resolved anyway by the runner, however we want any resolution errors
// (i.e. caused by a FieldAccessDescriptor that references fields not in the schema) to
// be caught and presented to the user at graph-construction time. Therefore we resolve
// here as well to catch these errors.
FieldAccessDescriptor fieldAccessDescriptor = null;
if (fieldAccessString == null) {
// This is the case where no FieldId is defined. Default to all fields accessed.
fieldAccessDescriptor = FieldAccessDescriptor.withAllFields();
} else {
// If there is a FieldAccessDescriptor in the class with this id, use that.
FieldAccessDeclaration fieldAccessDeclaration =
fieldAccessDeclarations.get(fieldAccessString);
if (fieldAccessDeclaration != null) {
checkArgument(fieldAccessDeclaration.field().getType().equals(FieldAccessDescriptor.class));
try {
fieldAccessDescriptor = (FieldAccessDescriptor) fieldAccessDeclaration.field().get(fn);
} catch (IllegalAccessException e) {
throw new RuntimeException(e);
}
} else {
// Otherwise, interpret the string as a field-name expression.
fieldAccessDescriptor = FieldAccessDescriptor.withFieldNames(fieldAccessString);
}
}
return fieldAccessDescriptor.resolve(inputSchema);
}
/**
* Try to provide coders for as many of the type arguments of given {@link
* DoFnSignature.StateDeclaration} as possible.
*/
private static <InputT> Coder[] codersForStateSpecTypes(
DoFnSignature.StateDeclaration stateDeclaration,
CoderRegistry coderRegistry,
Coder<InputT> inputCoder) {
Type stateType = stateDeclaration.stateType().getType();
if (!(stateType instanceof ParameterizedType)) {
// No type arguments means no coders to infer.
return new Coder[0];
}
Type[] typeArguments = ((ParameterizedType) stateType).getActualTypeArguments();
Coder[] coders = new Coder[typeArguments.length];
for (int i = 0; i < typeArguments.length; i++) {
Type typeArgument = typeArguments[i];
TypeDescriptor<?> typeDescriptor = TypeDescriptor.of(typeArgument);
try {
coders[i] = coderRegistry.getCoder(typeDescriptor);
} catch (CannotProvideCoderException e) {
try {
coders[i] =
coderRegistry.getCoder(
typeDescriptor, inputCoder.getEncodedTypeDescriptor(), inputCoder);
} catch (CannotProvideCoderException ignored) {
// Since not all type arguments will have a registered coder we ignore this exception.
}
}
}
return coders;
}
/**
* Perform common validations of the {@link DoFn} against the input {@link PCollection}, for
* example ensuring that the window type expected by the {@link DoFn} matches the window type of
* the {@link PCollection}.
*/
private static <InputT, OutputT> void validateWindowType(
PCollection<? extends InputT> input, DoFn<InputT, OutputT> fn) {
DoFnSignature signature = DoFnSignatures.getSignature((Class) fn.getClass());
TypeDescriptor<? extends BoundedWindow> actualWindowT =
input.getWindowingStrategy().getWindowFn().getWindowTypeDescriptor();
validateWindowTypeForMethod(actualWindowT, signature.processElement());
for (OnTimerMethod method : signature.onTimerMethods().values()) {
validateWindowTypeForMethod(actualWindowT, method);
}
}
private static void validateWindowTypeForMethod(
TypeDescriptor<? extends BoundedWindow> actualWindowT,
MethodWithExtraParameters methodSignature) {
if (methodSignature.windowT() != null) {
checkArgument(
methodSignature.windowT().isSupertypeOf(actualWindowT),
"%s unable to provide window -- expected window type from parameter (%s) is not a "
+ "supertype of actual window type assigned by windowing (%s)",
methodSignature.targetMethod(),
methodSignature.windowT(),
actualWindowT);
}
}
/**
* Perform common validations of the {@link DoFn}, for example ensuring that state is used
* correctly and that its features can be supported.
*/
private static <InputT, OutputT> void validate(DoFn<InputT, OutputT> fn) {
DoFnSignature signature = DoFnSignatures.getSignature((Class) fn.getClass());
// State is semantically incompatible with splitting
if (!signature.stateDeclarations().isEmpty() && signature.processElement().isSplittable()) {
throw new UnsupportedOperationException(
String.format(
"%s is splittable and uses state, but these are not compatible",
fn.getClass().getName()));
}
// Timers are semantically incompatible with splitting
if (!signature.timerDeclarations().isEmpty() && signature.processElement().isSplittable()) {
throw new UnsupportedOperationException(
String.format(
"%s is splittable and uses timers, but these are not compatible",
fn.getClass().getName()));
}
}
/**
* Extract information on how the DoFn uses schemas. In particular, if the schema of an element
* parameter does not match the input PCollection's schema, convert.
*/
@Internal
public static DoFnSchemaInformation getDoFnSchemaInformation(
DoFn<?, ?> fn, PCollection<?> input) {
DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());
DoFnSignature.ProcessElementMethod processElementMethod = signature.processElement();
if (!processElementMethod.getSchemaElementParameters().isEmpty()) {
if (!input.hasSchema()) {
throw new IllegalArgumentException("Type of @Element must match the DoFn type" + input);
}
}
SchemaRegistry schemaRegistry = input.getPipeline().getSchemaRegistry();
DoFnSchemaInformation doFnSchemaInformation = DoFnSchemaInformation.create();
for (SchemaElementParameter parameter : processElementMethod.getSchemaElementParameters()) {
TypeDescriptor<?> elementT = parameter.elementT();
FieldAccessDescriptor accessDescriptor =
getFieldAccessDescriptorFromParameter(
parameter.fieldAccessString(),
input.getSchema(),
signature.fieldAccessDeclarations(),
fn);
Schema selectedSchema = SelectHelpers.getOutputSchema(input.getSchema(), accessDescriptor);
ConvertHelpers.ConvertedSchemaInformation converted =
ConvertHelpers.getConvertedSchemaInformation(selectedSchema, elementT, schemaRegistry);
if (converted.outputSchemaCoder != null) {
doFnSchemaInformation =
doFnSchemaInformation.withSelectFromSchemaParameter(
(SchemaCoder<?>) input.getCoder(),
accessDescriptor,
selectedSchema,
converted.outputSchemaCoder,
converted.unboxedType != null);
} else {
// If the selected schema is a Row containing a single primitive type (which is the output
// of Select when selecting a primitive), attempt to unbox it and match against the
// parameter.
checkArgument(converted.unboxedType != null);
doFnSchemaInformation =
doFnSchemaInformation.withUnboxPrimitiveParameter(
(SchemaCoder<?>) input.getCoder(), accessDescriptor, selectedSchema, elementT);
}
}
return doFnSchemaInformation;
}
/**
* A {@link PTransform} that, when applied to a {@code PCollection<InputT>}, invokes a
* user-specified {@code DoFn<InputT, OutputT>} on all its elements, with all its outputs
* collected into an output {@code PCollection<OutputT>}.
*
* <p>A multi-output form of this transform can be created with {@link
* SingleOutput#withOutputTags}.
*
* @param <InputT> the type of the (main) input {@link PCollection} elements
* @param <OutputT> the type of the (main) output {@link PCollection} elements
*/
public static class SingleOutput<InputT, OutputT>
extends PTransform<PCollection<? extends InputT>, PCollection<OutputT>> {
private static final String MAIN_OUTPUT_TAG = "output";
private final Map<String, PCollectionView<?>> sideInputs;
private final DoFn<InputT, OutputT> fn;
private final DisplayData.ItemSpec<? extends Class<?>> fnDisplayData;
SingleOutput(
DoFn<InputT, OutputT> fn,
Map<String, PCollectionView<?>> sideInputs,
DisplayData.ItemSpec<? extends Class<?>> fnDisplayData) {
this.fn = fn;
this.fnDisplayData = fnDisplayData;
this.sideInputs = sideInputs;
}
/**
* Returns a new {@link ParDo} {@link PTransform} that's like this {@link PTransform} but with
* the specified additional side inputs. Does not modify this {@link PTransform}.
*
* <p>See the discussion of Side Inputs above for more explanation.
*/
public SingleOutput<InputT, OutputT> withSideInputs(PCollectionView<?>... sideInputs) {
return withSideInputs(Arrays.asList(sideInputs));
}
/**
* Returns a new {@link ParDo} {@link PTransform} that's like this {@link PTransform} but with
* the specified additional side inputs. Does not modify this {@link PTransform}.
*
* <p>See the discussion of Side Inputs above for more explanation.
*/
public SingleOutput<InputT, OutputT> withSideInputs(
Iterable<? extends PCollectionView<?>> sideInputs) {
Map<String, PCollectionView<?>> mappedInputs =
StreamSupport.stream(sideInputs.spliterator(), false)
.collect(Collectors.toMap(v -> v.getTagInternal().getId(), v -> v));
return withSideInputs(mappedInputs);
}
/**
* Returns a new {@link ParDo} {@link PTransform} that's like this {@link PTransform} but with
* the specified additional side inputs. Does not modify this {@link PTransform}.
*
* <p>See the discussion of Side Inputs above for more explanation.
*/
public SingleOutput<InputT, OutputT> withSideInputs(
Map<String, PCollectionView<?>> sideInputs) {
return new SingleOutput<>(
fn,
ImmutableMap.<String, PCollectionView<?>>builder()
.putAll(this.sideInputs)
.putAll(sideInputs)
.build(),
fnDisplayData);
}
/**
* Returns a new {@link ParDo} {@link PTransform} that's like this {@link PTransform} but with
* the specified additional side inputs. Does not modify this {@link PTransform}.
*/
public SingleOutput<InputT, OutputT> withSideInput(
String tagId, PCollectionView<?> pCollectionView) {
return withSideInputs(Collections.singletonMap(tagId, pCollectionView));
}
/**
* Returns a new multi-output {@link ParDo} {@link PTransform} that's like this {@link
* PTransform} but with the specified output tags. Does not modify this {@link PTransform}.
*
* <p>See the discussion of Additional Outputs above for more explanation.
*/
public MultiOutput<InputT, OutputT> withOutputTags(
TupleTag<OutputT> mainOutputTag, TupleTagList additionalOutputTags) {
return new MultiOutput<>(fn, sideInputs, mainOutputTag, additionalOutputTags, fnDisplayData);
}
@Override
public PCollection<OutputT> expand(PCollection<? extends InputT> input) {
SchemaRegistry schemaRegistry = input.getPipeline().getSchemaRegistry();
CoderRegistry registry = input.getPipeline().getCoderRegistry();
finishSpecifyingStateSpecs(fn, registry, input.getCoder());
TupleTag<OutputT> mainOutput = new TupleTag<>(MAIN_OUTPUT_TAG);
PCollection<OutputT> res =
input.apply(withOutputTags(mainOutput, TupleTagList.empty())).get(mainOutput);
try {
res.setSchema(
schemaRegistry.getSchema(getFn().getOutputTypeDescriptor()),
schemaRegistry.getToRowFunction(getFn().getOutputTypeDescriptor()),
schemaRegistry.getFromRowFunction(getFn().getOutputTypeDescriptor()));
} catch (NoSuchSchemaException e) {
try {
res.setCoder(
registry.getCoder(
getFn().getOutputTypeDescriptor(),
getFn().getInputTypeDescriptor(),
((PCollection<InputT>) input).getCoder()));
} catch (CannotProvideCoderException e2) {
// Ignore and leave coder unset.
}
}
return res;
}
@Override
protected String getKindString() {
return String.format("ParDo(%s)", NameUtils.approximateSimpleName(getFn()));
}
/**
* {@inheritDoc}
*
* <p>{@link ParDo} registers its internal {@link DoFn} as a subcomponent for display data.
* {@link DoFn} implementations can register display data by overriding {@link
* DoFn#populateDisplayData}.
*/
@Override
public void populateDisplayData(Builder builder) {
super.populateDisplayData(builder);
ParDo.populateDisplayData(builder, (HasDisplayData) fn, fnDisplayData);
}
public DoFn<InputT, OutputT> getFn() {
return fn;
}
public Map<String, PCollectionView<?>> getSideInputs() {
return sideInputs;
}
/**
* Returns the side inputs of this {@link ParDo}, tagged with the tag of the {@link
* PCollectionView}. The values of the returned map will be equal to the result of {@link
* #getSideInputs()}.
*/
@Override
public Map<TupleTag<?>, PValue> getAdditionalInputs() {
return PCollectionViews.toAdditionalInputs(sideInputs.values());
}
@Override
public String toString() {
return fn.toString();
}
}
/**
* A {@link PTransform} that, when applied to a {@code PCollection<InputT>}, invokes a
* user-specified {@code DoFn<InputT, OutputT>} on all its elements, which can emit elements to
* any of the {@link PTransform}'s output {@code PCollection}s, which are bundled into a result
* {@code PCollectionTuple}.
*
* @param <InputT> the type of the (main) input {@code PCollection} elements
* @param <OutputT> the type of the main output {@code PCollection} elements
*/
public static class MultiOutput<InputT, OutputT>
extends PTransform<PCollection<? extends InputT>, PCollectionTuple> {
private final Map<String, PCollectionView<?>> sideInputs;
private final TupleTag<OutputT> mainOutputTag;
private final TupleTagList additionalOutputTags;
private final DisplayData.ItemSpec<? extends Class<?>> fnDisplayData;
private final DoFn<InputT, OutputT> fn;
MultiOutput(
DoFn<InputT, OutputT> fn,
Map<String, PCollectionView<?>> sideInputs,
TupleTag<OutputT> mainOutputTag,
TupleTagList additionalOutputTags,
ItemSpec<? extends Class<?>> fnDisplayData) {
this.sideInputs = sideInputs;
this.mainOutputTag = mainOutputTag;
this.additionalOutputTags = additionalOutputTags;
this.fn = fn;
this.fnDisplayData = fnDisplayData;
}
/**
* Returns a new multi-output {@link ParDo} {@link PTransform} that's like this {@link
* PTransform} but with the specified additional side inputs. Does not modify this {@link
* PTransform}.
*
* <p>See the discussion of Side Inputs above for more explanation.
*/
public MultiOutput<InputT, OutputT> withSideInputs(PCollectionView<?>... sideInputs) {
return withSideInputs(Arrays.asList(sideInputs));
}
public MultiOutput<InputT, OutputT> withSideInputs(
Iterable<? extends PCollectionView<?>> sideInputs) {
Map<String, PCollectionView<?>> mappedInputs =
StreamSupport.stream(sideInputs.spliterator(), false)
.collect(Collectors.toMap(v -> v.getTagInternal().getId(), v -> v));
return withSideInputs(mappedInputs);
}
/**
* Returns a new multi-output {@link ParDo} {@link PTransform} that's like this {@link
* PTransform} but with the specified additional side inputs. Does not modify this {@link
* PTransform}.
*
* <p>See the discussion of Side Inputs above for more explanation.
*/
public MultiOutput<InputT, OutputT> withSideInputs(Map<String, PCollectionView<?>> sideInputs) {
return new MultiOutput<>(
fn,
ImmutableMap.<String, PCollectionView<?>>builder()
.putAll(this.sideInputs)
.putAll(sideInputs)
.build(),
mainOutputTag,
additionalOutputTags,
fnDisplayData);
}
/**
* Returns a new multi-output {@link ParDo} {@link PTransform} that's like this {@link
* PTransform} but with the specified additional side inputs. Does not modify this {@link
* PTransform}.
*/
public MultiOutput<InputT, OutputT> withSideInput(
String tagId, PCollectionView<?> pCollectionView) {
return withSideInputs(Collections.singletonMap(tagId, pCollectionView));
}
@Override
public PCollectionTuple expand(PCollection<? extends InputT> input) {
// SplittableDoFn should be forbidden on the runner-side.
validateWindowType(input, fn);
// Use coder registry to determine coders for all StateSpec defined in the fn signature.
CoderRegistry registry = input.getPipeline().getCoderRegistry();
finishSpecifyingStateSpecs(fn, registry, input.getCoder());
DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());
if (signature.usesState() || signature.usesTimers()) {
validateStateApplicableForInput(fn, input);
}
validateSideInputTypes(sideInputs, fn);
// TODO: We should validate OutputReceiver<Row> only happens if the output PCollection
// as schema. However coder/schema inference may not have happened yet at this point.
// Need to figure out where to validate this.
PCollectionTuple outputs =
PCollectionTuple.ofPrimitiveOutputsInternal(
input.getPipeline(),
TupleTagList.of(mainOutputTag).and(additionalOutputTags.getAll()),
// TODO
Collections.emptyMap(),
input.getWindowingStrategy(),
input.isBounded().and(signature.isBoundedPerElement()));
@SuppressWarnings("unchecked")
Coder<InputT> inputCoder = ((PCollection<InputT>) input).getCoder();
for (PCollection<?> out : outputs.getAll().values()) {
try {
out.setCoder(
(Coder)
registry.getCoder(
out.getTypeDescriptor(), getFn().getInputTypeDescriptor(), inputCoder));
} catch (CannotProvideCoderException e) {
// Ignore and let coder inference happen later.
}
}
// The fn will likely be an instance of an anonymous subclass
// such as DoFn<Integer, String> { }, thus will have a high-fidelity
// TypeDescriptor for the output type.
outputs.get(mainOutputTag).setTypeDescriptor(getFn().getOutputTypeDescriptor());
return outputs;
}
@Override
protected String getKindString() {
return String.format("ParMultiDo(%s)", NameUtils.approximateSimpleName(getFn()));
}
@Override
public void populateDisplayData(Builder builder) {
super.populateDisplayData(builder);
ParDo.populateDisplayData(builder, fn, fnDisplayData);
}
public DoFn<InputT, OutputT> getFn() {
return fn;
}
public TupleTag<OutputT> getMainOutputTag() {
return mainOutputTag;
}
public TupleTagList getAdditionalOutputTags() {
return additionalOutputTags;
}
public Map<String, PCollectionView<?>> getSideInputs() {
return sideInputs;
}
/**
* Returns the side inputs of this {@link ParDo}, tagged with the tag of the {@link
* PCollectionView}. The values of the returned map will be equal to the result of {@link
* #getSideInputs()}.
*/
@Override
public Map<TupleTag<?>, PValue> getAdditionalInputs() {
return PCollectionViews.toAdditionalInputs(sideInputs.values());
}
@Override
public String toString() {
return fn.toString();
}
}
private static void populateDisplayData(
DisplayData.Builder builder,
HasDisplayData fn,
DisplayData.ItemSpec<? extends Class<?>> fnDisplayData) {
builder.include("fn", fn).add(fnDisplayData);
}
private static boolean isSplittable(DoFn<?, ?> fn) {
return DoFnSignatures.signatureForDoFn(fn).processElement().isSplittable();
}
}