Skip to content

Fix output bugs in FnApiDoFnRunner @OnTimer and @OnWindowExpiration #40024

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kennknowles wants to merge 1 commit into
base: master
Choose a base branch
from
+333 −20
Open

Fix output bugs in FnApiDoFnRunner @OnTimer and @OnWindowExpiration #40024

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 41 additions & 20 deletions sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1235,6 +1235,7 @@ private <K> void processTimer(
currentTimer = null;
currentTimeDomain = null;
currentWindow = null;
causedByDrain = null;
}
}

Expand Down Expand Up @@ -2295,10 +2296,10 @@ public ValueKind valueKind(DoFn<InputT, OutputT> doFn) {
* DoFn.OnWindowExpiration @OnWindowExpiration}.
*/
private class OnWindowExpirationContext<K> extends BaseArgumentProvider<InputT, OutputT> {
private class Context extends DoFn<InputT, OutputT>.OnWindowExpirationContext
private class WindowExpirationContext extends DoFn<InputT, OutputT>.OnWindowExpirationContext
implements OutputReceiver<OutputT> {

private Context() {
private WindowExpirationContext() {
doFn.super();
}

Expand Down Expand Up @@ -2369,7 +2370,7 @@ public <T> void output(TupleTag<T> tag, T output) {
null,
currentTimer.causedByDrain(),
null,
currentElement.getValueKind()));
ValueKind.INSERT));
}

@Override
Expand All @@ -2395,6 +2396,9 @@ public <T> void outputWindowedValue(
checkOnWindowExpirationTimestamp(timestamp);
FnDataReceiver<WindowedValue<T>> consumer =
(FnDataReceiver) localNameToConsumer.get(tag.getId());
if (consumer == null) {
throw new IllegalArgumentException(String.format("Unknown output tag %s", tag));
}
outputTo(consumer, WindowedValues.of(output, timestamp, windows, paneInfo));
}

Expand All @@ -2405,7 +2409,12 @@ public void outputWindowedValue(WindowedValue<OutputT> windowedValue) {

@Override
public <T> void outputWindowedValue(TupleTag<T> tag, WindowedValue<T> windowedValue) {
outputTo((FnDataReceiver) localNameToConsumer.get(tag.getId()), windowedValue);
FnDataReceiver<WindowedValue<T>> consumer =
(FnDataReceiver) localNameToConsumer.get(tag.getId());
if (consumer == null) {
throw new IllegalArgumentException(String.format("Unknown output tag %s", tag));
}
outputTo(consumer, windowedValue);
}

@SuppressWarnings(
Expand Down Expand Up @@ -2435,8 +2444,7 @@ private void checkOnWindowExpirationTimestamp(Instant timestamp) {
}
}

private final OnWindowExpirationContext.Context context =
new OnWindowExpirationContext.Context();
private final WindowExpirationContext context = new WindowExpirationContext();

@Override
public DoFn<InputT, OutputT>.OnWindowExpirationContext onWindowExpirationContext(
Expand All @@ -2459,11 +2467,6 @@ public Instant timestamp(DoFn<InputT, OutputT> doFn) {
return currentTimer.getHoldTimestamp();
}

@Override
public TimeDomain timeDomain(DoFn<InputT, OutputT> doFn) {
return currentTimeDomain;
}

@Override
public K key() {
return (K) currentTimer.getUserKey();
Expand Down Expand Up @@ -2622,9 +2625,9 @@ public String getErrorContext() {
/** Provides arguments for a {@link DoFnInvoker} for {@link DoFn.OnTimer @OnTimer}. */
private class OnTimerContext<K> extends BaseArgumentProvider<InputT, OutputT> {

private class Context extends DoFn<InputT, OutputT>.OnTimerContext
private class TimerContext extends DoFn<InputT, OutputT>.OnTimerContext
implements OutputReceiver<OutputT> {
private Context() {
private TimerContext() {
doFn.super();
}

Expand Down Expand Up @@ -2719,7 +2722,12 @@ public void outputWindowedValue(WindowedValue<OutputT> windowedValue) {

@Override
public <T> void outputWindowedValue(TupleTag<T> tag, WindowedValue<T> windowedValue) {
outputTo((FnDataReceiver) localNameToConsumer.get(tag.getId()), windowedValue);
FnDataReceiver<WindowedValue<T>> consumer =
(FnDataReceiver) localNameToConsumer.get(tag.getId());
if (consumer == null) {
throw new IllegalArgumentException(String.format("Unknown output tag %s", tag));
}
outputTo(consumer, windowedValue);
}

@Override
Expand All @@ -2728,7 +2736,15 @@ public <T> void outputWindowedValue(
T output,
Instant timestamp,
Collection<? extends BoundedWindow> windows,
PaneInfo paneInfo) {}
PaneInfo paneInfo) {
checkTimerTimestamp(timestamp);
FnDataReceiver<WindowedValue<T>> consumer =
(FnDataReceiver) localNameToConsumer.get(tag.getId());
if (consumer == null) {
throw new IllegalArgumentException(String.format("Unknown output tag %s", tag));
}
outputTo(consumer, WindowedValues.of(output, timestamp, windows, paneInfo));
}

@Override
public TimeDomain timeDomain() {
Expand Down Expand Up @@ -2772,7 +2788,7 @@ private void checkTimerTimestamp(Instant timestamp) {
}
}

private final OnTimerContext.Context context = new OnTimerContext.Context();
private final TimerContext context = new TimerContext();

@Override
public BoundedWindow window() {
Expand Down Expand Up @@ -2822,8 +2838,12 @@ public OutputReceiver<OutputT> outputReceiver(DoFn<InputT, OutputT> doFn) {

@Override
public OutputBuilder<Row> builder(Row value) {
return WindowedValues.builder(currentElement)
.withValue(value)
return WindowedValues.<Row>builder()
.setValue(value)
.setTimestamp(currentTimer.getHoldTimestamp())
.setWindow(currentWindow)
.setPaneInfo(currentTimer.getPaneInfo())
.setCausedByDrain(currentTimer.causedByDrain())
.setReceiver(
windowedValue ->
context.outputWindowedValue(
Expand Down Expand Up @@ -2860,7 +2880,7 @@ public OutputBuilder<T> builder(T value) {
.setWindow(currentWindow)
.setCausedByDrain(currentTimer.causedByDrain())
.setPaneInfo(currentTimer.getPaneInfo())
.setReceiver(windowedValue -> context.outputWindowedValue(windowedValue));
.setReceiver(windowedValue -> context.outputWindowedValue(tag, windowedValue));
}
};
}
Expand All @@ -2887,14 +2907,15 @@ private <T> OutputReceiver<Row> createTaggedRowReceiver(TupleTag<T> tag) {
@Override
public OutputBuilder<Row> builder(Row value) {
return WindowedValues.<Row>builder()
.withValue(value)
.setValue(value)
.setTimestamp(currentTimer.getHoldTimestamp())
.setWindow(currentWindow)
.setPaneInfo(currentTimer.getPaneInfo())
.setCausedByDrain(currentTimer.causedByDrain())
.setReceiver(
windowedValue ->
context.outputWindowedValue(
tag,
windowedValue.withValue(
fromRowFunction.apply(windowedValue.getValue()))));
}
Expand Down
Loading
Loading