Repository navigation
ParDoWithDlq: propagate output coder + add ParDo.MultiOutput.withDlq - #168
Conversation
Two related changes to the .withDlq DLQ API in OSS Beam (LinkedIn fork).
1) Fix ParDoWithDlq.expand() coder propagation
--------------------------------------------
Today, the wrapping PTransform created by .withDlq returns a primitive
output PCollection with `null` coder:
return PCollection.createPrimitiveOutputInternal(
input.getPipeline(), input.getWindowingStrategy(), input.isBounded(), null);
For wraps over anonymous-class DoFns (e.g. `ParDo.of(new DoFn<X, Y>() {...})`),
Beam's plain ParDo.SingleOutput.expand normally reads the type parameters via
reflection on the anonymous class's bytecode and sets the output coder /
schema. But because .withDlq wraps the original ParDo in a different
PTransform whose own type signature is `PCollection<? extends InputT> ->
PCollection<OutputT>`, the coder inference walk lands on the wrapping
transform's signature instead of the inner DoFn's, and the runtime throws
IllegalStateException: Unable to return a default Coder for <stage>.out
downstream. Named-class DoFns happened to work because their call sites
already had an explicit .setCoder(...) chained after .withDlq for other
reasons.
This commit mirrors ParDo.SingleOutput.expand's schema/coder inference
in ParDoWithDlq.expand: try the schema registry first, fall back to the
coder registry with the DoFn's input/output TypeDescriptors and the
input PCollection's coder, and leave the coder unset only if neither
resolves. Anonymous-class .withDlq sites now behave identically to plain
anonymous-class ParDo.of.
Caught downstream by cgp-nearline apache#2449's CI: test failures in
TestFanoutPipeline, TestRepartitionerPipeline, and TestFilteringPipeline
all stemmed from this missing coder propagation around the 4 anonymous
gate / filter DoFns wrapped with .withDlq.
2) Add ParDo.MultiOutput.withDlq(...)
-----------------------------------
OSS Beam only exposes .withDlq on ParDo.SingleOutput today. Multi-output
ParDos (the ones that chain .withOutputTags(mainTag, additionalTags)) had
no first-class DLQ wrap, forcing consumers either to refactor multi-output
logic into multiple single-output stages or to write a runner-only wrapper.
This commit adds the multi-output counterpart, mirroring the relationship
between ParDo.SingleOutput and ParDo.MultiOutput:
PCollectionTuple outputs = input.apply(
ParDo.of(new MyMultiOutputDoFn())
.withOutputTags(mainTag, TupleTagList.of(sideTagA).and(sideTagB))
.withDlq(myDlqSink));
New package-private class ParDoWithDlqMultiOutput<InputT> +
public-accessor interface ParDoWithDlqMultiOutputSpec<InputT> follow the
same structure as ParDoWithDlq / ParDoWithDlqSpec. URN
"beam:transform:li:pardo_with_dlq_multi_output:v1" reserved for the
runner-side translator (to be added in a beam-runner follow-up; the
existing in-flight beam-runner PR apache#2335 implementing
LiParDoWithDlqMultiOutput can be closed once that translator lands).
expand() builds the same PCollectionTuple ParDo.MultiOutput.expand would
build and runs the same per-tag coder registration logic, so multi-output
.withDlq sites get coder inference for free as well.
…qMultiOutput Matches the same suppression already used on ParDo (line 390). Required because the Checker Framework's nullness analysis on the per-tag coderRegistry.getCoder(out.getTypeDescriptor(), ...) call needs the same suppression that lets the identical pattern compile inside ParDo.MultiOutput.expand. Verified locally: ./gradlew :sdks:java:core:compileJava -> BUILD SUCCESSFUL.
12 tests over the public surface introduced in commit 7582568e9
(ParDo.MultiOutput.withDlq + ParDoWithDlqMultiOutput + spec):
- testWithDlqReturnedFromMultiOutput
- testGetFnReturnsSameFn
- testGetMainOutputTagReturnsSameTag
- testGetAdditionalOutputTagsReturnsSameTags
- testGetSideInputsIsEmptyByDefault
- testGetDlqSinkReturnsSameSink
- testDlqFilterIsNullByDefault
- testWithDlqFilterSetsFilter
- testWithDlqFilterPreservesFnTagsAndSink
- testUrnConstant
- testWithDlqRejectsNullSink
- testWithDlqFilterRejectsNullFilter
Mirrors ParDoWithDlqTest's structure; the DLQ-side FailedRecord and
DlqSink contracts are exercised there and not duplicated here.
Pipeline execution (translation + runner) is covered in the runner-side
beam-runner translator follow-up, not in this OSS module.
Verified locally:
./gradlew :sdks:java:core:test --tests
'org.apache.beam.sdk.transforms.ParDoWithDlqMultiOutputTest'
-> 12 tests, 0 failures, 0 errors, 0.111s.
| input.isBounded()); | ||
| 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:
outputs.get(mainOutputTag).setTypeDescriptor(fn.getOutputTypeDescriptor());
The loop keys coder inference off each tag's own getTypeDescriptor(). For the common idiom of an anonymous DoFn with a raw main tag — ParDo.of(new DoFn<Integer, String>(){...}).withOutputTags(new TupleTag<>(), ...).withDlq(sink) — the main tag carries no type info, so getCoder(...) throws CannotProvideCoderException and the main output is left with no coder.
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.
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:
((PCollection) outputs.get(mainOutputTag)).setTypeDescriptor(fn.getOutputTypeDescriptor());
(raw cast because here fn is DoFn<InputT, ?> and mainOutputTag is TupleTag<?>, so the captures don't unify like they do upstream — covered by the class-level @SuppressWarnings.)
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.
| public Map<String, PCollectionView<?>> getSideInputs() { | ||
| return sideInputs; | ||
| } |
There was a problem hiding this comment.
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:
@Override
public Map<TupleTag<?>, PValue> getAdditionalInputs() {
return PCollectionViews.toAdditionalInputs(sideInputs.values());
}
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.
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().
| TupleTagList.of(mainOutputTag).and(additionalOutputTags.getAll()), | ||
| Collections.emptyMap(), | ||
| input.getWindowingStrategy(), | ||
| input.isBounded()); |
There was a problem hiding this comment.
upstream ParDo.MultiOutput.expand uses input.isBounded().and(signature.isBoundedPerElement()) here; this uses input.isBounded() alone. Only diverges for a splittable DoFn that's unbounded-per-element — unlikely under .withDlq, and single-output ParDoWithDlq.expand already has the same simplification — so fine to leave, just flagging.
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.
| * {@link ParDoWithDlqTest}; this file focuses on the multi-output-specific surface. | ||
| */ | ||
| @RunWith(JUnit4.class) | ||
| public class ParDoWithDlqMultiOutputTest implements Serializable { |
There was a problem hiding this comment.
These tests only cover accessor plumbing, the URN string, and null-arg rejection — none exercise expand(), where the actual coder logic (the PR's primary purpose) lives. Two missing cases would have caught the issues flagged in ParDoWithDlqMultiOutput.java:
- Wrap an anonymous multi-output DoFn with a raw main tag, call expand, assert outputs.get(mainTag).getCoder() is non-null (catches the missing setTypeDescriptor).
- Wrap a ParDo with a real side input and assert it appears in getAdditionalInputs().
A validate()-rejects-non-Flink-runner test would also be worth adding. Given these are the core of the change, I'd treat the coverage as part of this PR rather than a follow-up.
There was a problem hiding this comment.
Agreed — treated as part of this PR. Added in 7abb7d5, covering all three cases you called out:
- testExpandResolvesMainTypeDescriptorAndAdditionalTagCoder — anonymous multi-output DoFn with a raw main tag; asserts the main output's TypeDescriptor is recovered (catches the missing setTypeDescriptor) and a type-captured additional tag's coder resolves via the loop.
- testGetAdditionalInputsDeclaresSideInputs — a real side input appears in getAdditionalInputs().
- testValidateRejectsNonFlinkRunner — validate() throws for a non-Flink runner (CrashingRunner).
The two fix-related tests were verified to fail without the corresponding source changes. Full class is 15/15 green.
…and() tests Addresses @yananhao12's review on ParDoWithDlqMultiOutput: - expand(): set the main output's TypeDescriptor from fn.getOutputTypeDescriptor() after the per-tag coder loop, matching the final step of ParDo.MultiOutput.expand. For the common idiom of an anonymous DoFn with a raw main tag (new TupleTag<>()), the per-tag getCoder() can't resolve a coder; this fallback recovers the type so downstream inference no longer fails with "Unable to return a default Coder" — the same failure the single-output ParDoWithDlq.expand already guards against. - Override getAdditionalInputs() to declare threaded side inputs via PCollectionViews.toAdditionalInputs(sideInputs.values()), mirroring ParDo.MultiOutput. expand() only emits primitive outputs, so without this the runner's topological traversal wouldn't know to translate the view-producing subgraph first. - Add expand()/getAdditionalInputs()/validate() tests (the prior tests only covered accessors/URN/null-args): raw-main-tag type-descriptor recovery + typed-tag coder resolution, side inputs surfacing in getAdditionalInputs(), and validate() rejecting a non-Flink runner. The first two were verified to fail without the fixes above. input.isBounded() (flagged at line 128) left as-is for consistency with the single-output ParDoWithDlq.expand, per the reviewer's note that it's fine to leave. ParDoWithDlqMultiOutputTest: 15/15 pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Please add a meaningful description for your change here
Problem & Solution Overview
Two fixes to the
.withDlq(...)DLQ API in the LinkedIn fork of OSS Beam.1) Fix
ParDoWithDlq.expand()coder propagationToday, the wrapping
PTransformcreated by.withDlqon aParDo.SingleOutputreturns a primitive outputPCollectionwith anullcoder: