Repository navigation
ParDoWithDlq: propagate output coder + add ParDo.MultiOutput.withDlq #168
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,169 @@ | ||
| package org.apache.beam.sdk.transforms; | ||
|
|
||
| import java.util.Collections; | ||
| import java.util.List; | ||
| import java.util.Map; | ||
| import javax.annotation.Nullable; | ||
| 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.options.PipelineOptions; | ||
| 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.vendor.guava.v26_0_jre.com.google.common.base.Preconditions; | ||
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList; | ||
|
|
||
| /** | ||
| * Package-private implementation of the multi-output DLQ-enabled ParDo transform. | ||
| * | ||
| * <p>Multi-output counterpart to {@link ParDoWithDlq}. Instances are created exclusively via | ||
| * {@link ParDo.MultiOutput#withDlq(DlqSink)} or | ||
| * {@link ParDo.MultiOutput#withDlq(DlqSink, SerializableFunction)}. Callers outside this package | ||
| * interact with this transform through the {@link ParDoWithDlqMultiOutputSpec} interface. | ||
| * | ||
| * <p>The Beam pipeline graph still exposes the same {@link PCollectionTuple} the wrapped | ||
| * {@code ParDo.MultiOutput} would produce — the DLQ side output is handled entirely within the | ||
| * Flink runner via the registered {@link DlqSink} and is not visible in the user-facing tuple. | ||
| * | ||
| * @param <InputT> the input element type | ||
| */ | ||
| @SuppressWarnings({ | ||
| "nullness", // matches the suppression on ParDo (TODO: apache/beam#20497) | ||
| "rawtypes" | ||
| }) | ||
| class ParDoWithDlqMultiOutput<InputT> | ||
| extends PTransform<PCollection<? extends InputT>, PCollectionTuple> | ||
| implements ParDoWithDlqMultiOutputSpec<InputT> { | ||
|
|
||
| private static final List<String> ALLOWED_RUNNERS = | ||
| ImmutableList.of( | ||
| "org.apache.beam.runners.flink.FlinkRunner", | ||
| "com.linkedin.beam.runners.LineageFlinkRunner"); | ||
|
|
||
| private final DoFn<InputT, ?> fn; | ||
| private final TupleTag<?> mainOutputTag; | ||
| private final TupleTagList additionalOutputTags; | ||
| private final Map<String, PCollectionView<?>> sideInputs; | ||
| private final DlqSink<InputT> dlqSink; | ||
| @Nullable private final SerializableFunction<Throwable, Boolean> dlqFilter; | ||
|
|
||
| ParDoWithDlqMultiOutput( | ||
| DoFn<InputT, ?> fn, | ||
| TupleTag<?> mainOutputTag, | ||
| TupleTagList additionalOutputTags, | ||
| Map<String, PCollectionView<?>> sideInputs, | ||
| DlqSink<InputT> dlqSink, | ||
| @Nullable SerializableFunction<Throwable, Boolean> dlqFilter) { | ||
| this.fn = fn; | ||
| this.mainOutputTag = mainOutputTag; | ||
| this.additionalOutputTags = additionalOutputTags; | ||
| this.sideInputs = sideInputs; | ||
| this.dlqSink = dlqSink; | ||
| this.dlqFilter = dlqFilter; | ||
| } | ||
|
|
||
| @Override | ||
| public DoFn<InputT, ?> getFn() { | ||
| return fn; | ||
| } | ||
|
|
||
| @Override | ||
| public TupleTag<?> getMainOutputTag() { | ||
| return mainOutputTag; | ||
| } | ||
|
|
||
| @Override | ||
| public TupleTagList getAdditionalOutputTags() { | ||
| return additionalOutputTags; | ||
| } | ||
|
|
||
| @Override | ||
| public Map<String, PCollectionView<?>> getSideInputs() { | ||
| return sideInputs; | ||
| } | ||
|
|
||
| /** | ||
| * Declares the side-input views this transform consumes, mirroring {@link | ||
| * ParDo.MultiOutput#getAdditionalInputs}. | ||
| * | ||
| * <p>Because {@link #expand} only emits primitive outputs and never applies an inner {@code | ||
| * ParDo}, the side-input views are not referenced anywhere in the visible graph. This override is | ||
| * how the runner's topological traversal learns the transform consumes those views and must | ||
| * translate the view-producing subgraph first; without it a multi-output {@code .withDlq(...)} | ||
| * with side inputs risks the consuming operator being translated before its side-input stream | ||
| * exists. | ||
| */ | ||
| @Override | ||
| public Map<TupleTag<?>, PValue> getAdditionalInputs() { | ||
| return PCollectionViews.toAdditionalInputs(sideInputs.values()); | ||
| } | ||
|
|
||
| @Override | ||
| public DlqSink<InputT> getDlqSink() { | ||
| return dlqSink; | ||
| } | ||
|
|
||
| @Override | ||
| @Nullable | ||
| public SerializableFunction<Throwable, Boolean> getDlqFilter() { | ||
| return dlqFilter; | ||
| } | ||
|
|
||
| @Override | ||
| public void validate(@Nullable PipelineOptions options) { | ||
| if (options == null) { | ||
| return; | ||
| } | ||
| Preconditions.checkState( | ||
| ALLOWED_RUNNERS.contains(options.getRunner().getName()), | ||
| "ParDoWithDlqMultiOutput is not supported in runner: %s", | ||
| options.getRunner().getName()); | ||
| } | ||
|
|
||
| /** | ||
| * Constructs a {@link PCollectionTuple} with one primitive output per tag (main + additional), | ||
| * matching {@link ParDo.MultiOutput#expand}. | ||
| * | ||
| * <p>Sets coders on each output via the {@code CoderRegistry} (same logic as | ||
| * {@code ParDo.MultiOutput.expand}) so that wrapping a {@code ParDo.MultiOutput} with | ||
| * {@code .withDlq(...)} does not erase coder inference — anonymous-class DoFns continue to | ||
| * have their output types resolved automatically. | ||
| */ | ||
| @Override | ||
| @SuppressWarnings({"unchecked", "rawtypes"}) | ||
| public PCollectionTuple expand(PCollection<? extends InputT> input) { | ||
| PCollectionTuple outputs = | ||
| PCollectionTuple.ofPrimitiveOutputsInternal( | ||
| input.getPipeline(), | ||
| TupleTagList.of(mainOutputTag).and(additionalOutputTags.getAll()), | ||
| Collections.emptyMap(), | ||
| input.getWindowingStrategy(), | ||
| input.isBounded()); | ||
|
There was a problem hiding this comment. upstream
Author
There was a problem hiding this comment. Intentional — leaving as input.isBounded(). The .and(signature.isBoundedPerElement()) only changes the result for a splittable @UnboundedPerElement DoFn over a bounded input, which .withDlq(...) realistically never wraps. More importantly, the single-output ParDoWithDlq.expand already uses the same simplification, so keeping multi-output identical preserves parity between the two paths; adding the signature lookup to only one would introduce exactly the kind of divergence flagged on the setTypeDescriptor thread. If we ever want strict upstream parity I'd add it to both expand() methods in a follow-up, but it's a no-op for supported .withDlq usage. |
||
| CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry(); | ||
| Coder<InputT> inputCoder = ((PCollection<InputT>) input).getCoder(); | ||
| for (PCollection<?> out : outputs.getAll().values()) { | ||
|
There was a problem hiding this comment. The main output never gets its type descriptor — this defeats the coder fix for the anonymous-DoFn case. This loop mirrors ParDo.MultiOutput.expand, but drops the line upstream runs right after the loop: The loop keys coder inference off each tag's own getTypeDescriptor(). For the common idiom of an anonymous DoFn with a raw main tag — Upstream recovers by setting the type descriptor from the DoFn's output type, which lets later inference succeed; without it we get "Unable to return a default Coder for .out" downstream — the exact failure mode the single-output half of this PR is fixing. Note ParDoWithDlq.expand (single-output) correctly uses fn.getOutputTypeDescriptor(), so the two paths are inconsistent. Please add the setTypeDescriptor call after the loop + a regression test.
Author
There was a problem hiding this comment. Good catch — fixed in 7abb7d5. Added after the per-tag coder loop, matching the final step of ParDo.MultiOutput.expand: This brings multi-output in line with the single-output ParDoWithDlq.expand, which already uses fn.getOutputTypeDescriptor(). Regression test testExpandResolvesMainTypeDescriptorAndAdditionalTagCoder wraps an anonymous DoFn<Integer, String> with a raw new TupleTag<>() main tag and asserts the main output's TypeDescriptor resolves to strings() (plus a type-captured side tag whose coder resolves to StringUtf8Coder via the loop). Verified it fails without the fix. |
||
| try { | ||
| out.setCoder( | ||
| (Coder) | ||
| coderRegistry.getCoder( | ||
| out.getTypeDescriptor(), fn.getInputTypeDescriptor(), inputCoder)); | ||
| } catch (CannotProvideCoderException e) { | ||
| // Leave coder unset for this tag — caller can .setCoder explicitly. Matches | ||
| // ParDo.MultiOutput's expand behavior. | ||
| } | ||
| } | ||
| // Mirror the final step of ParDo.MultiOutput.expand: the fn is typically an anonymous subclass | ||
| // (e.g. new DoFn<Integer, String>(){...}) carrying a high-fidelity output TypeDescriptor. | ||
| // Setting it on the main output recovers coder inference for the common idiom of a raw main tag | ||
| // (new TupleTag<>()), where out.getTypeDescriptor() above is null and the per-tag getCoder() | ||
| // could not resolve a coder — without this the main output is left with no coder and downstream | ||
| // applies throw "Unable to return a default Coder", the exact failure this PR fixes. | ||
| ((PCollection) outputs.get(mainOutputTag)).setTypeDescriptor(fn.getOutputTypeDescriptor()); | ||
| return outputs; | ||
| } | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,68 @@ | ||
| package org.apache.beam.sdk.transforms; | ||
|
|
||
| import java.util.Map; | ||
| import javax.annotation.Nullable; | ||
| import org.apache.beam.sdk.values.PCollectionView; | ||
| import org.apache.beam.sdk.values.TupleTag; | ||
| import org.apache.beam.sdk.values.TupleTagList; | ||
|
|
||
| /** | ||
| * Public accessor interface for the package-private {@link ParDoWithDlqMultiOutput} transform. | ||
| * | ||
| * <p>Multi-output counterpart to {@link ParDoWithDlqSpec}. Mirrors the relationship between | ||
| * {@link ParDo.SingleOutput} and {@link ParDo.MultiOutput}: a {@code ParDo.MultiOutput} can be | ||
| * wrapped with {@code .withDlq(DlqSink)} to route per-element user-code throws to the sink | ||
| * while still emitting successful elements to each of the configured output tags. | ||
| * | ||
| * <p>This interface allows the Flink runner (which lives in a different package) to inspect | ||
| * transform properties — {@link #getFn()}, {@link #getMainOutputTag()}, | ||
| * {@link #getAdditionalOutputTags()}, {@link #getSideInputs()}, {@link #getDlqSink()}, | ||
| * {@link #getDlqFilter()} — and reference the {@link #URN} constant, without needing to import | ||
| * the package-private implementation class. | ||
| * | ||
| * <p>Callers outside {@code org.apache.beam.sdk.transforms} should interact with this transform | ||
| * only through {@link ParDo.MultiOutput#withDlq(DlqSink)} and | ||
| * {@link ParDo.MultiOutput#withDlq(DlqSink, SerializableFunction)}. | ||
| * | ||
| * @param <InputT> the input element type | ||
| */ | ||
| public interface ParDoWithDlqMultiOutputSpec<InputT> { | ||
|
|
||
| /** URN used by the Flink runner to identify and translate this transform. */ | ||
| String URN = "beam:transform:li:pardo_with_dlq_multi_output:v1"; | ||
|
|
||
| /** Returns the wrapped {@link DoFn}. */ | ||
| DoFn<InputT, ?> getFn(); | ||
|
|
||
| /** Returns the main output tag from the wrapped {@code ParDo.MultiOutput}. */ | ||
| TupleTag<?> getMainOutputTag(); | ||
|
|
||
| /** Returns the additional output tags from the wrapped {@code ParDo.MultiOutput}. */ | ||
| TupleTagList getAdditionalOutputTags(); | ||
|
|
||
| /** | ||
| * Returns the side inputs threaded through from the wrapped {@code ParDo.MultiOutput} (an | ||
| * empty map when none were declared). | ||
| */ | ||
| Map<String, PCollectionView<?>> getSideInputs(); | ||
|
|
||
| /** Returns the {@link DlqSink} that receives failed elements. */ | ||
| DlqSink<InputT> getDlqSink(); | ||
|
|
||
| /** | ||
| * Returns the optional filter predicate, or {@code null} if all exceptions are routed to | ||
| * the DLQ. | ||
| */ | ||
| @Nullable | ||
| SerializableFunction<Throwable, Boolean> getDlqFilter(); | ||
|
|
||
| /** | ||
| * Returns the implementation {@link PTransform} class for use in payload translator | ||
| * registration. Accessible here because this interface is in the same package as the | ||
| * package-private {@link ParDoWithDlqMultiOutput}. | ||
| */ | ||
| @SuppressWarnings({"unchecked", "rawtypes"}) | ||
| static Class<? extends PTransform<?, ?>> implementationClass() { | ||
| return (Class<? extends PTransform<?, ?>>) (Class) ParDoWithDlqMultiOutput.class; | ||
| } | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Side inputs are threaded through but never declared via getAdditionalInputs().
We store sideInputs and expose getSideInputs(), but don't override getAdditionalInputs(), while real ParDo.MultiOutput does:
Because expand() only emits primitive outputs and never applies an inner ParDo, the side-input views are referenced nowhere in the visible graph. getAdditionalInputs() is how the runner's topological traversal learns this transform consumes those views and must translate the view-producing subgraph first. Without it, a multi-output ParDo with side inputs risks the consuming operator being translated before its side-input stream exists, plus broken pipeline-surgery accounting. Silent for side-input-free jobs, but a latent hole. Please override getAdditionalInputs() to match ParDo.MultiOutput + add a side-input test.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed in 7abb7d5. Overrode getAdditionalInputs() to mirror ParDo.MultiOutput:
@OverRide
public Map<TupleTag<?>, PValue> getAdditionalInputs() {
return PCollectionViews.toAdditionalInputs(sideInputs.values());
}
Since expand() only emits primitive outputs and never applies an inner ParDo, this is what lets the runner's topological traversal know the transform consumes those views and translate the view-producing subgraph first. Covered by testGetAdditionalInputsDeclaresSideInputs, which threads a real View.asSingleton() side input through .withDlq(...) and asserts it surfaces in getAdditionalInputs().