diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java index a5914a25f799..d391807d798f 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java @@ -1235,6 +1235,7 @@ private void processTimer( currentTimer = null; currentTimeDomain = null; currentWindow = null; + causedByDrain = null; } } @@ -2295,10 +2296,10 @@ public ValueKind valueKind(DoFn doFn) { * DoFn.OnWindowExpiration @OnWindowExpiration}. */ private class OnWindowExpirationContext extends BaseArgumentProvider { - private class Context extends DoFn.OnWindowExpirationContext + private class WindowExpirationContext extends DoFn.OnWindowExpirationContext implements OutputReceiver { - private Context() { + private WindowExpirationContext() { doFn.super(); } @@ -2369,7 +2370,7 @@ public void output(TupleTag tag, T output) { null, currentTimer.causedByDrain(), null, - currentElement.getValueKind())); + ValueKind.INSERT)); } @Override @@ -2395,6 +2396,9 @@ public void outputWindowedValue( checkOnWindowExpirationTimestamp(timestamp); FnDataReceiver> 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)); } @@ -2405,7 +2409,12 @@ public void outputWindowedValue(WindowedValue windowedValue) { @Override public void outputWindowedValue(TupleTag tag, WindowedValue windowedValue) { - outputTo((FnDataReceiver) localNameToConsumer.get(tag.getId()), windowedValue); + FnDataReceiver> consumer = + (FnDataReceiver) localNameToConsumer.get(tag.getId()); + if (consumer == null) { + throw new IllegalArgumentException(String.format("Unknown output tag %s", tag)); + } + outputTo(consumer, windowedValue); } @SuppressWarnings( @@ -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.OnWindowExpirationContext onWindowExpirationContext( @@ -2459,11 +2467,6 @@ public Instant timestamp(DoFn doFn) { return currentTimer.getHoldTimestamp(); } - @Override - public TimeDomain timeDomain(DoFn doFn) { - return currentTimeDomain; - } - @Override public K key() { return (K) currentTimer.getUserKey(); @@ -2622,9 +2625,9 @@ public String getErrorContext() { /** Provides arguments for a {@link DoFnInvoker} for {@link DoFn.OnTimer @OnTimer}. */ private class OnTimerContext extends BaseArgumentProvider { - private class Context extends DoFn.OnTimerContext + private class TimerContext extends DoFn.OnTimerContext implements OutputReceiver { - private Context() { + private TimerContext() { doFn.super(); } @@ -2719,7 +2722,12 @@ public void outputWindowedValue(WindowedValue windowedValue) { @Override public void outputWindowedValue(TupleTag tag, WindowedValue windowedValue) { - outputTo((FnDataReceiver) localNameToConsumer.get(tag.getId()), windowedValue); + FnDataReceiver> consumer = + (FnDataReceiver) localNameToConsumer.get(tag.getId()); + if (consumer == null) { + throw new IllegalArgumentException(String.format("Unknown output tag %s", tag)); + } + outputTo(consumer, windowedValue); } @Override @@ -2728,7 +2736,15 @@ public void outputWindowedValue( T output, Instant timestamp, Collection windows, - PaneInfo paneInfo) {} + PaneInfo paneInfo) { + checkTimerTimestamp(timestamp); + FnDataReceiver> 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() { @@ -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() { @@ -2822,8 +2838,12 @@ public OutputReceiver outputReceiver(DoFn doFn) { @Override public OutputBuilder builder(Row value) { - return WindowedValues.builder(currentElement) - .withValue(value) + return WindowedValues.builder() + .setValue(value) + .setTimestamp(currentTimer.getHoldTimestamp()) + .setWindow(currentWindow) + .setPaneInfo(currentTimer.getPaneInfo()) + .setCausedByDrain(currentTimer.causedByDrain()) .setReceiver( windowedValue -> context.outputWindowedValue( @@ -2860,7 +2880,7 @@ public OutputBuilder builder(T value) { .setWindow(currentWindow) .setCausedByDrain(currentTimer.causedByDrain()) .setPaneInfo(currentTimer.getPaneInfo()) - .setReceiver(windowedValue -> context.outputWindowedValue(windowedValue)); + .setReceiver(windowedValue -> context.outputWindowedValue(tag, windowedValue)); } }; } @@ -2887,7 +2907,7 @@ private OutputReceiver createTaggedRowReceiver(TupleTag tag) { @Override public OutputBuilder builder(Row value) { return WindowedValues.builder() - .withValue(value) + .setValue(value) .setTimestamp(currentTimer.getHoldTimestamp()) .setWindow(currentWindow) .setPaneInfo(currentTimer.getPaneInfo()) @@ -2895,6 +2915,7 @@ public OutputBuilder builder(Row value) { .setReceiver( windowedValue -> context.outputWindowedValue( + tag, windowedValue.withValue( fromRowFunction.apply(windowedValue.getValue())))); } diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java index d24ab39c0471..4166b2a48dcc 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java @@ -87,6 +87,7 @@ import org.apache.beam.sdk.metrics.MetricsEnvironment; import org.apache.beam.sdk.options.ExperimentalOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.state.BagState; import org.apache.beam.sdk.state.CombiningState; import org.apache.beam.sdk.state.StateSpec; @@ -128,6 +129,7 @@ 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.Row; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.TupleTagList; import org.apache.beam.sdk.values.WindowedValue; @@ -1140,6 +1142,296 @@ public void testTimers() throws Exception { assertThat(result, containsInAnyOrder(expected.toArray())); } + private static class TestTimerTaggedOutputDoFn extends DoFn, String> { + @TimerId("event") + private final TimerSpec eventTimerSpec = TimerSpecs.timer(TimeDomain.EVENT_TIME); + + private final TupleTag additionalOutput; + + private TestTimerTaggedOutputDoFn(TupleTag additionalOutput) { + this.additionalOutput = additionalOutput; + } + + @ProcessElement + public void processElement(ProcessContext context, @TimerId("event") Timer eventTimeTimer) { + eventTimeTimer.withOutputTimestamp(context.timestamp()).set(context.timestamp()); + } + + @OnTimer("event") + public void eventTimer( + OnTimerContext context, @Key String key, MultiOutputReceiver receiver) { + context.output("main:" + key); + context.output(additionalOutput, "output:" + key); + context.outputWindowedValue( + additionalOutput, + "outputWindowedValue:" + key, + context.timestamp(), + Collections.singletonList(GlobalWindow.INSTANCE), + PaneInfo.NO_FIRING); + receiver.get(additionalOutput).output("receiver:" + key); + } + } + + @Test + public void testTimerTaggedOutputs() throws Exception { + Pipeline p = Pipeline.create(); + PCollection> valuePCollection = + p.apply(Create.of(KV.of("unused", "unused"))); + TupleTag mainOutput = new TupleTag("main") {}; + TupleTag additionalOutput = new TupleTag("additional") {}; + PCollectionTuple outputPCollection = + valuePCollection.apply( + TEST_TRANSFORM_ID, + ParDo.of(new TestTimerTaggedOutputDoFn(additionalOutput)) + .withOutputTags(mainOutput, TupleTagList.of(additionalOutput))); + + SdkComponents sdkComponents = SdkComponents.create(); + sdkComponents.registerEnvironment(Environment.getDefaultInstance()); + RunnerApi.Pipeline pProto = PipelineTranslation.toProto(p, sdkComponents); + String outputPCollectionId = + sdkComponents.registerPCollection(outputPCollection.get(mainOutput)); + String additionalPCollectionId = + sdkComponents.registerPCollection(outputPCollection.get(additionalOutput)); + RunnerApi.PTransform pTransform = + pProto.getComponents().getTransformsOrThrow(TEST_TRANSFORM_ID); + + List> mainOutputValues = new ArrayList<>(); + List> additionalOutputValues = new ArrayList<>(); + PTransformRunnerFactoryTestContext context = + PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform) + .beamFnStateClient(new FakeBeamFnStateClient(StringUtf8Coder.of(), ImmutableMap.of())) + .processBundleInstructionId("57L") + .components( + RunnerApi.Components.newBuilder() + .putAllCoders(pProto.getComponents().getCodersMap()) + .putAllEnvironments(Collections.emptyMap()) + .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap()) + .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap()) + .build()) + .outboundAggregators( + ImmutableMap.of( + ApiServiceDescriptor.getDefaultInstance(), + new TestBeamFnDataOutboundAggregator(() -> "57L"))) + .timerApiServiceDescriptor(ApiServiceDescriptor.getDefaultInstance()) + .build(); + context.addPCollectionConsumer( + outputPCollectionId, + (FnDataReceiver) (FnDataReceiver>) mainOutputValues::add); + context.addPCollectionConsumer( + additionalPCollectionId, + (FnDataReceiver) (FnDataReceiver>) additionalOutputValues::add); + + new FnApiDoFnRunner.Factory<>().addRunnerForPTransform(context); + Iterables.getOnlyElement(context.getStartBundleFunctions()).run(); + + context + .getIncomingTimerEndpoint("ts-event") + .getReceiver() + .accept(timerInGlobalWindow("A", new Instant(1400L), new Instant(2400L))); + + assertThat(mainOutputValues, contains(isValueInGlobalWindow("main:A", new Instant(1400L)))); + assertThat( + additionalOutputValues, + contains( + isValueInGlobalWindow("output:A", new Instant(1400L)), + isValueInGlobalWindow("outputWindowedValue:A", new Instant(1400L)), + isValueInGlobalWindow("receiver:A", new Instant(1400L)))); + + Iterables.getOnlyElement(context.getFinishBundleFunctions()).run(); + Iterables.getOnlyElement(context.getTearDownFunctions()).run(); + } + + private static final Schema ROW_SCHEMA = + Schema.of(Schema.Field.of("field", Schema.FieldType.STRING)); + + private static Row row(String value) { + return Row.withSchema(ROW_SCHEMA).addValue(value).build(); + } + + private static class TestTimerRowOutputDoFn extends DoFn, Row> { + @TimerId("event") + private final TimerSpec eventTimerSpec = TimerSpecs.timer(TimeDomain.EVENT_TIME); + + private final TupleTag mainOutput; + private final TupleTag additionalOutput; + + private TestTimerRowOutputDoFn(TupleTag mainOutput, TupleTag additionalOutput) { + this.mainOutput = mainOutput; + this.additionalOutput = additionalOutput; + } + + @ProcessElement + public void processElement(ProcessContext context, @TimerId("event") Timer eventTimeTimer) { + eventTimeTimer.withOutputTimestamp(context.timestamp()).set(context.timestamp()); + } + + @OnTimer("event") + public void eventTimer(@Key String key, MultiOutputReceiver receiver) { + receiver.getRowReceiver(mainOutput).output(row("mainRow:" + key)); + receiver.getRowReceiver(additionalOutput).output(row("taggedRow:" + key)); + } + } + + @Test + public void testTimerRowOutputReceivers() throws Exception { + Pipeline p = Pipeline.create(); + PCollection> valuePCollection = + p.apply(Create.of(KV.of("unused", "unused"))); + TupleTag mainOutput = new TupleTag("main") {}; + TupleTag additionalOutput = new TupleTag("additional") {}; + PCollectionTuple outputPCollection = + valuePCollection.apply( + TEST_TRANSFORM_ID, + ParDo.of(new TestTimerRowOutputDoFn(mainOutput, additionalOutput)) + .withOutputTags(mainOutput, TupleTagList.of(additionalOutput))); + outputPCollection.get(mainOutput).setRowSchema(ROW_SCHEMA); + outputPCollection.get(additionalOutput).setRowSchema(ROW_SCHEMA); + + SdkComponents sdkComponents = SdkComponents.create(); + sdkComponents.registerEnvironment(Environment.getDefaultInstance()); + RunnerApi.Pipeline pProto = PipelineTranslation.toProto(p, sdkComponents); + String outputPCollectionId = + sdkComponents.registerPCollection(outputPCollection.get(mainOutput)); + String additionalPCollectionId = + sdkComponents.registerPCollection(outputPCollection.get(additionalOutput)); + RunnerApi.PTransform pTransform = + pProto.getComponents().getTransformsOrThrow(TEST_TRANSFORM_ID); + + List> mainOutputValues = new ArrayList<>(); + List> additionalOutputValues = new ArrayList<>(); + PTransformRunnerFactoryTestContext context = + PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform) + .beamFnStateClient(new FakeBeamFnStateClient(StringUtf8Coder.of(), ImmutableMap.of())) + .processBundleInstructionId("57L") + .components( + RunnerApi.Components.newBuilder() + .putAllCoders(pProto.getComponents().getCodersMap()) + .putAllEnvironments(Collections.emptyMap()) + .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap()) + .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap()) + .build()) + .outboundAggregators( + ImmutableMap.of( + ApiServiceDescriptor.getDefaultInstance(), + new TestBeamFnDataOutboundAggregator(() -> "57L"))) + .timerApiServiceDescriptor(ApiServiceDescriptor.getDefaultInstance()) + .build(); + context.addPCollectionConsumer( + outputPCollectionId, + (FnDataReceiver) (FnDataReceiver>) mainOutputValues::add); + context.addPCollectionConsumer( + additionalPCollectionId, + (FnDataReceiver) (FnDataReceiver>) additionalOutputValues::add); + + new FnApiDoFnRunner.Factory<>().addRunnerForPTransform(context); + Iterables.getOnlyElement(context.getStartBundleFunctions()).run(); + + context + .getIncomingTimerEndpoint("ts-event") + .getReceiver() + .accept(timerInGlobalWindow("A", new Instant(1400L), new Instant(2400L))); + + assertThat( + mainOutputValues, contains(isValueInGlobalWindow(row("mainRow:A"), new Instant(1400L)))); + assertThat( + additionalOutputValues, + contains(isValueInGlobalWindow(row("taggedRow:A"), new Instant(1400L)))); + + Iterables.getOnlyElement(context.getFinishBundleFunctions()).run(); + Iterables.getOnlyElement(context.getTearDownFunctions()).run(); + } + + private static class TestWindowExpirationDoFn extends DoFn, String> { + @StateId("bag") + private final StateSpec> bagStateSpec = StateSpecs.bag(StringUtf8Coder.of()); + + private final TupleTag additionalOutput; + + private TestWindowExpirationDoFn(TupleTag additionalOutput) { + this.additionalOutput = additionalOutput; + } + + @ProcessElement + public void processElement( + ProcessContext context, @StateId("bag") BagState bagState) { + bagState.add(context.element().getValue()); + } + + @OnWindowExpiration + public void onWindowExpiration(OnWindowExpirationContext context, @Key String key) { + context.output("main:" + key); + context.output(additionalOutput, "output:" + key); + } + } + + @Test + public void testOnWindowExpirationTaggedOutputs() throws Exception { + Pipeline p = Pipeline.create(); + PCollection> valuePCollection = + p.apply(Create.of(KV.of("unused", "unused"))); + TupleTag mainOutput = new TupleTag("main") {}; + TupleTag additionalOutput = new TupleTag("additional") {}; + PCollectionTuple outputPCollection = + valuePCollection.apply( + TEST_TRANSFORM_ID, + ParDo.of(new TestWindowExpirationDoFn(additionalOutput)) + .withOutputTags(mainOutput, TupleTagList.of(additionalOutput))); + + SdkComponents sdkComponents = SdkComponents.create(); + sdkComponents.registerEnvironment(Environment.getDefaultInstance()); + RunnerApi.Pipeline pProto = PipelineTranslation.toProto(p, sdkComponents); + String outputPCollectionId = + sdkComponents.registerPCollection(outputPCollection.get(mainOutput)); + String additionalPCollectionId = + sdkComponents.registerPCollection(outputPCollection.get(additionalOutput)); + RunnerApi.PTransform pTransform = + pProto.getComponents().getTransformsOrThrow(TEST_TRANSFORM_ID); + String onWindowExpirationFamilyId = + RunnerApi.ParDoPayload.parseFrom(pTransform.getSpec().getPayload()) + .getOnWindowExpirationTimerFamilySpec(); + + List> mainOutputValues = new ArrayList<>(); + List> additionalOutputValues = new ArrayList<>(); + PTransformRunnerFactoryTestContext context = + PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID, pTransform) + .beamFnStateClient(new FakeBeamFnStateClient(StringUtf8Coder.of(), ImmutableMap.of())) + .processBundleInstructionId("57L") + .components( + RunnerApi.Components.newBuilder() + .putAllCoders(pProto.getComponents().getCodersMap()) + .putAllEnvironments(Collections.emptyMap()) + .putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap()) + .putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap()) + .build()) + .outboundAggregators( + ImmutableMap.of( + ApiServiceDescriptor.getDefaultInstance(), + new TestBeamFnDataOutboundAggregator(() -> "57L"))) + .timerApiServiceDescriptor(ApiServiceDescriptor.getDefaultInstance()) + .build(); + context.addPCollectionConsumer( + outputPCollectionId, + (FnDataReceiver) (FnDataReceiver>) mainOutputValues::add); + context.addPCollectionConsumer( + additionalPCollectionId, + (FnDataReceiver) (FnDataReceiver>) additionalOutputValues::add); + + new FnApiDoFnRunner.Factory<>().addRunnerForPTransform(context); + Iterables.getOnlyElement(context.getStartBundleFunctions()).run(); + + context + .getIncomingTimerEndpoint(onWindowExpirationFamilyId) + .getReceiver() + .accept(timerInGlobalWindow("A", new Instant(1400L), new Instant(2400L))); + + assertThat(mainOutputValues, contains(isValueInGlobalWindow("main:A", new Instant(1400L)))); + assertThat( + additionalOutputValues, contains(isValueInGlobalWindow("output:A", new Instant(1400L)))); + + Iterables.getOnlyElement(context.getFinishBundleFunctions()).run(); + Iterables.getOnlyElement(context.getTearDownFunctions()).run(); + } + private org.apache.beam.sdk.util.construction.Timer timerInGlobalWindow( K userKey, Instant holdTimestamp, Instant fireTimestamp) { return dynamicTimerInGlobalWindow(userKey, "", holdTimestamp, fireTimestamp);