From 4d87d3c803f23dd016946524fb30c42bfdf30dd0 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Tue, 9 Sep 2025 15:17:05 +0200 Subject: [PATCH 01/26] encode empty metadata within WindowedValue --- .../worker/StreamingDataflowWorker.java | 5 ++ .../worker/UngroupedWindmillReader.java | 1 + .../worker/WindmillKeyedWorkItem.java | 2 +- .../runners/dataflow/worker/WindmillSink.java | 21 ++++++- .../StreamingGroupAlsoByWindowFnsTest.java | 8 ++- ...ngGroupAlsoByWindowsReshuffleDoFnTest.java | 8 ++- .../worker/WindmillKeyedWorkItemTest.java | 9 ++- .../sdk/transforms/windowing/PaneInfo.java | 59 +++++++++++++++---- .../beam/sdk/values/WindowedValues.java | 23 +++++++- .../transforms/windowing/PaneInfoTest.java | 20 +++++++ .../beam/sdk/util/WindowedValueTest.java | 27 +++++++++ 11 files changed, 162 insertions(+), 21 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java index 2a4b111af225..59269f8ea55e 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java @@ -110,6 +110,7 @@ import org.apache.beam.sdk.io.gcp.bigquery.BigQuerySinkMetrics; import org.apache.beam.sdk.metrics.MetricsEnvironment; import org.apache.beam.sdk.util.construction.CoderTranslation; +import org.apache.beam.sdk.values.WindowedValues; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheStats; @@ -163,6 +164,7 @@ public final class StreamingDataflowWorker { private static final Random CLIENT_ID_GENERATOR = new Random(); private static final String CHANNELZ_PATH = "/channelz"; private static final String BEAM_FN_API_EXPERIMENT = "beam_fn_api"; + private static final String ELEMENT_METADATA_SUPPORTED_EXPERIMENT = "element_metadata_supported"; private static final String STREAMING_ENGINE_USE_JOB_SETTINGS_FOR_HEARTBEAT_POOL_EXPERIMENT = "streaming_engine_use_job_settings_for_heartbeat_pool"; // Experiment make the monitor within BoundedQueueExecutor fair @@ -815,6 +817,9 @@ public static void main(String[] args) throws Exception { validateWorkerOptions(options); CoderTranslation.verifyModelCodersRegistered(); + if (DataflowRunner.hasExperiment(options, ELEMENT_METADATA_SUPPORTED_EXPERIMENT)) { + WindowedValues.FullWindowedValueCoder.setMetadataSupported(); + } LOG.debug("Creating StreamingDataflowWorker from options: {}", options); StreamingDataflowWorker worker = StreamingDataflowWorker.fromOptions(options); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java index e031d1bb50eb..5df15e8d2109 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java @@ -117,6 +117,7 @@ protected WindowedValue decodeMessage(Windmill.Message message) throws IOExce Collection windows = WindmillSink.decodeMetadataWindows(windowsCoder, message.getMetadata()); PaneInfo paneInfo = WindmillSink.decodeMetadataPane(message.getMetadata()); + WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata()); if (valueCoder instanceof KvCoder) { KvCoder kvCoder = (KvCoder) valueCoder; InputStream key = context.getSerializedKey().newInput(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java index cee4894e3d68..bae41a156967 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java @@ -108,7 +108,7 @@ public Iterable> elementsIterable() { Collection windows = WindmillSink.decodeMetadataWindows(windowsCoder, message.getMetadata()); PaneInfo paneInfo = WindmillSink.decodeMetadataPane(message.getMetadata()); - + WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata()); InputStream inputStream = message.getData().newInput(); ElemT value = valueCoder.decode(inputStream, Coder.Context.OUTER); return WindowedValues.of(value, timestamp, windows, paneInfo); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java index 7cb6f2223472..4765eb74e7f5 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java @@ -26,6 +26,7 @@ import java.util.Collection; import java.util.HashMap; import java.util.Map; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.runners.dataflow.util.CloudObject; import org.apache.beam.runners.dataflow.worker.util.common.worker.Sink; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; @@ -40,6 +41,7 @@ import org.apache.beam.sdk.values.ValueWithRecordId; import org.apache.beam.sdk.values.ValueWithRecordId.ValueWithRecordIdCoder; import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; import org.apache.beam.sdk.values.WindowedValues.FullWindowedValueCoder; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; @@ -75,7 +77,8 @@ private static ByteString encodeMetadata( ByteStringOutputStream stream, Coder> windowsCoder, Collection windows, - PaneInfo paneInfo) + PaneInfo paneInfo, + BeamFnApi.Elements.ElementMetadata metadata) throws IOException { try { PaneInfoCoder.INSTANCE.encode(paneInfo, stream); @@ -101,12 +104,26 @@ public static PaneInfo decodeMetadataPane(ByteString metadata) throws IOExceptio return PaneInfoCoder.INSTANCE.decode(inStream); } + public static BeamFnApi.Elements.ElementMetadata decodeAdditionalMetadata( + Coder> windowsCoder, ByteString metadata) + throws IOException { + InputStream inStream = metadata.newInput(); + PaneInfo paneInfo = PaneInfoCoder.INSTANCE.decode(inStream); + windowsCoder.decode(inStream); + if (paneInfo.isElementMetadata() && WindowedValues.WindowedValueCoder.isMetadataSupported()) { + return BeamFnApi.Elements.ElementMetadata.parseDelimitedFrom(inStream); + } else { + // empty + return BeamFnApi.Elements.ElementMetadata.newBuilder().build(); + } + } + public static Collection decodeMetadataWindows( Coder> windowsCoder, ByteString metadata) throws IOException { InputStream inStream = metadata.newInput(); PaneInfoCoder.INSTANCE.decode(inStream); - return windowsCoder.decode(inStream, Coder.Context.OUTER); + return windowsCoder.decode(inStream); } /** A {@link SinkFactory.Registrar} for windmill sinks. */ diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java index c89a031b3728..1182a2c0b9e9 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java @@ -29,6 +29,7 @@ import java.util.Arrays; import java.util.Collection; import java.util.List; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.runners.core.DoFnRunner; import org.apache.beam.runners.core.DoFnRunners; import org.apache.beam.runners.core.InMemoryStateInternals; @@ -176,7 +177,12 @@ private void addElement( valueCoder.encode(value, dataOutput, Context.OUTER); messageBundle .addMessagesBuilder() - .setMetadata(WindmillSink.encodeMetadata(windowsCoder, windows, PaneInfo.NO_FIRING)) + .setMetadata( + WindmillSink.encodeMetadata( + windowsCoder, + windows, + PaneInfo.NO_FIRING, + BeamFnApi.Elements.ElementMetadata.newBuilder().build())) .setData(dataOutput.toByteString()) .setTimestamp(WindmillTimeUtils.harnessToWindmillTimestamp(timestamp)); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowsReshuffleDoFnTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowsReshuffleDoFnTest.java index c169c9b46a57..a348c0f00214 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowsReshuffleDoFnTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowsReshuffleDoFnTest.java @@ -24,6 +24,7 @@ import java.util.Arrays; import java.util.Collection; import java.util.List; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.runners.core.DoFnRunner; import org.apache.beam.runners.core.KeyedWorkItem; import org.apache.beam.runners.core.NullSideInputReader; @@ -114,7 +115,12 @@ private void addElement( valueCoder.encode(value, dataOutput, Context.OUTER); messageBundle .addMessagesBuilder() - .setMetadata(WindmillSink.encodeMetadata(windowsCoder, windows, PaneInfo.NO_FIRING)) + .setMetadata( + WindmillSink.encodeMetadata( + windowsCoder, + windows, + PaneInfo.NO_FIRING, + BeamFnApi.Elements.ElementMetadata.newBuilder().build())) .setData(dataOutput.toByteString()) .setTimestamp(WindmillTimeUtils.harnessToWindmillTimestamp(timestamp)); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java index ffe71176367a..53a36722e41c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java @@ -22,6 +22,7 @@ import java.io.IOException; import java.util.Collection; import java.util.Collections; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.runners.core.KeyedWorkItem; import org.apache.beam.runners.core.StateNamespace; import org.apache.beam.runners.core.StateNamespaces; @@ -107,10 +108,14 @@ private void addElement( long timestamp, String value, IntervalWindow window, - PaneInfo paneInfo) + PaneInfo pane) throws IOException { ByteString encodedMetadata = - WindmillSink.encodeMetadata(WINDOWS_CODER, Collections.singletonList(window), paneInfo); + WindmillSink.encodeMetadata( + WINDOWS_CODER, + Collections.singletonList(window), + pane, + BeamFnApi.Elements.ElementMetadata.newBuilder().build()); chunk .addMessagesBuilder() .setTimestamp(WindmillTimeUtils.harnessToWindmillTimestamp(new Instant(timestamp))) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PaneInfo.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PaneInfo.java index 6e4c694d48e3..80f931b8e6a6 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PaneInfo.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PaneInfo.java @@ -146,10 +146,10 @@ private static byte encodedByte(boolean isFirst, boolean isLast, Timing timing) ImmutableMap.Builder decodingBuilder = ImmutableMap.builder(); for (Timing timing : Timing.values()) { long onTimeIndex = timing == Timing.EARLY ? -1 : 0; - register(decodingBuilder, new PaneInfo(true, true, timing, 0, onTimeIndex)); - register(decodingBuilder, new PaneInfo(true, false, timing, 0, onTimeIndex)); - register(decodingBuilder, new PaneInfo(false, true, timing, -1, onTimeIndex)); - register(decodingBuilder, new PaneInfo(false, false, timing, -1, onTimeIndex)); + register(decodingBuilder, new PaneInfo(true, true, timing, 0, onTimeIndex, false)); + register(decodingBuilder, new PaneInfo(true, false, timing, 0, onTimeIndex, false)); + register(decodingBuilder, new PaneInfo(false, true, timing, -1, onTimeIndex, false)); + register(decodingBuilder, new PaneInfo(false, false, timing, -1, onTimeIndex, false)); } BYTE_TO_PANE_INFO = decodingBuilder.build(); } @@ -159,7 +159,7 @@ private static void register(ImmutableMap.Builder builder, PaneI } private final byte encodedByte; - + private final boolean containsElementMetadata; private final boolean isFirst; private final boolean isLast; private final Timing timing; @@ -177,13 +177,20 @@ private static void register(ImmutableMap.Builder builder, PaneI public static final PaneInfo ON_TIME_AND_ONLY_FIRING = PaneInfo.createPane(true, true, Timing.ON_TIME, 0, 0); - private PaneInfo(boolean isFirst, boolean isLast, Timing timing, long index, long onTimeIndex) { + private PaneInfo( + boolean isFirst, + boolean isLast, + Timing timing, + long index, + long onTimeIndex, + boolean containsElementMetadata) { this.encodedByte = encodedByte(isFirst, isLast, timing); this.isFirst = isFirst; this.isLast = isLast; this.timing = timing; this.index = index; this.nonSpeculativeIndex = onTimeIndex; + this.containsElementMetadata = containsElementMetadata; } public static PaneInfo createPane(boolean isFirst, boolean isLast, Timing timing) { @@ -191,13 +198,25 @@ public static PaneInfo createPane(boolean isFirst, boolean isLast, Timing timing return createPane(isFirst, isLast, timing, 0, timing == Timing.EARLY ? -1 : 0); } + /** Factory method to create a {@link PaneInfo} with the specified parameters. */ /** Factory method to create a {@link PaneInfo} with the specified parameters. */ public static PaneInfo createPane( boolean isFirst, boolean isLast, Timing timing, long index, long onTimeIndex) { + return createPane(isFirst, isLast, timing, index, onTimeIndex, false); + } + + /** Factory method to create a {@link PaneInfo} with the specified parameters. */ + public static PaneInfo createPane( + boolean isFirst, + boolean isLast, + Timing timing, + long index, + long onTimeIndex, + boolean containsElementMetadata) { if (isFirst || timing == Timing.UNKNOWN) { return checkNotNull(BYTE_TO_PANE_INFO.get(encodedByte(isFirst, isLast, timing))); } else { - return new PaneInfo(isFirst, isLast, timing, index, onTimeIndex); + return new PaneInfo(isFirst, isLast, timing, index, onTimeIndex, containsElementMetadata); } } @@ -219,6 +238,15 @@ public boolean isFirst() { return isFirst; } + public boolean isElementMetadata() { + return containsElementMetadata; + } + + public PaneInfo withElementMetadata(boolean elementMetadata) { + return new PaneInfo( + this.isFirst, this.isLast, this.timing, index, nonSpeculativeIndex, elementMetadata); + } + /** Return true if this is the last pane that will be produced in the associated window. */ public boolean isLast() { return isLast; @@ -295,6 +323,8 @@ public String toString() { /** A Coder for encoding PaneInfo instances. */ public static class PaneInfoCoder extends AtomicCoder { + private static final byte ELEMENT_METADATA_MASK = (byte) 0x80; + private enum Encoding { FIRST, ONE_INDEX, @@ -337,16 +367,17 @@ private PaneInfoCoder() {} public void encode(PaneInfo value, final OutputStream outStream) throws CoderException, IOException { Encoding encoding = chooseEncoding(value); + byte elementMetadata = value.containsElementMetadata ? ELEMENT_METADATA_MASK : 0x00; switch (chooseEncoding(value)) { case FIRST: - outStream.write(value.encodedByte); + outStream.write(value.encodedByte | elementMetadata); break; case ONE_INDEX: - outStream.write(value.encodedByte | encoding.tag); + outStream.write(value.encodedByte | encoding.tag | elementMetadata); VarInt.encode(value.index, outStream); break; case TWO_INDICES: - outStream.write(value.encodedByte | encoding.tag); + outStream.write(value.encodedByte | encoding.tag | elementMetadata); VarInt.encode(value.index, outStream); VarInt.encode(value.nonSpeculativeIndex, outStream); break; @@ -360,9 +391,10 @@ public PaneInfo decode(final InputStream inStream) throws CoderException, IOExce byte keyAndTag = (byte) inStream.read(); PaneInfo base = Preconditions.checkNotNull(BYTE_TO_PANE_INFO.get((byte) (keyAndTag & 0x0F))); long index, onTimeIndex; - switch (Encoding.fromTag(keyAndTag)) { + boolean elementMetadata = (keyAndTag & ELEMENT_METADATA_MASK) != 0; + switch (Encoding.fromTag((byte) (keyAndTag & ~ELEMENT_METADATA_MASK))) { case FIRST: - return base; + return base.withElementMetadata(elementMetadata); case ONE_INDEX: index = VarInt.decodeLong(inStream); onTimeIndex = base.timing == Timing.EARLY ? -1 : index; @@ -374,7 +406,8 @@ public PaneInfo decode(final InputStream inStream) throws CoderException, IOExce default: throw new CoderException("Unknown encoding " + (keyAndTag & 0xF0)); } - return new PaneInfo(base.isFirst, base.isLast, base.timing, index, onTimeIndex); + return new PaneInfo( + base.isFirst, base.isLast, base.timing, index, onTimeIndex, elementMetadata); } @Override diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java index 9b079b8699b9..26ba4948bc0a 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java @@ -36,6 +36,7 @@ import java.util.List; import java.util.Objects; import java.util.Set; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.sdk.annotations.Internal; import org.apache.beam.sdk.coders.ByteArrayCoder; import org.apache.beam.sdk.coders.Coder; @@ -763,6 +764,15 @@ public static ParamWindowedValueCoder getParamWindowedValueCoder(Coder /** Abstract class for {@code WindowedValue} coder. */ public abstract static class WindowedValueCoder extends StructuredCoder> { final Coder valueCoder; + private static boolean metadataSupported = false; + + public static void setMetadataSupported() { + metadataSupported = true; + } + + public static boolean isMetadataSupported() { + return metadataSupported; + } WindowedValueCoder(Coder valueCoder) { this.valueCoder = checkNotNull(valueCoder); @@ -829,7 +839,15 @@ public void encode(WindowedValue windowedElem, OutputStream outStream, Contex throws CoderException, IOException { InstantCoder.of().encode(windowedElem.getTimestamp(), outStream); windowsCoder.encode(windowedElem.getWindows(), outStream); - PaneInfoCoder.INSTANCE.encode(windowedElem.getPaneInfo(), outStream); + boolean metadataSupported = isMetadataSupported(); + PaneInfoCoder.INSTANCE.encode( + windowedElem.getPaneInfo().withElementMetadata(metadataSupported), outStream); + if (metadataSupported) { + BeamFnApi.Elements.ElementMetadata.Builder builder = + BeamFnApi.Elements.ElementMetadata.newBuilder(); + BeamFnApi.Elements.ElementMetadata em = builder.build(); + em.writeDelimitedTo(outStream); + } valueCoder.encode(windowedElem.getValue(), outStream, context); } @@ -844,6 +862,9 @@ public WindowedValue decode(InputStream inStream, Context context) Instant timestamp = InstantCoder.of().decode(inStream); Collection windows = windowsCoder.decode(inStream); PaneInfo paneInfo = PaneInfoCoder.INSTANCE.decode(inStream); + if (isMetadataSupported() && paneInfo.isElementMetadata()) { + BeamFnApi.Elements.ElementMetadata.parseDelimitedFrom(inStream); + } T value = valueCoder.decode(inStream, context); // Because there are some remaining (incorrect) uses of WindowedValue with no windows, diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/windowing/PaneInfoTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/windowing/PaneInfoTest.java index 946deba036db..cda8ee1ea55c 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/windowing/PaneInfoTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/windowing/PaneInfoTest.java @@ -52,6 +52,22 @@ public void testEncodingRoundTrip() throws Exception { } } + @Test + public void testEncodingRoundTripWithElementMetadata() throws Exception { + Coder coder = PaneInfo.PaneInfoCoder.INSTANCE; + for (Timing timing : Timing.values()) { + long onTimeIndex = timing == Timing.EARLY ? -1 : 37; + CoderProperties.coderDecodeEncodeEqual( + coder, PaneInfo.createPane(false, false, timing, 389, onTimeIndex, true)); + CoderProperties.coderDecodeEncodeEqual( + coder, PaneInfo.createPane(false, true, timing, 5077, onTimeIndex, true)); + CoderProperties.coderDecodeEncodeEqual( + coder, PaneInfo.createPane(true, false, timing, 0, 0, true)); + CoderProperties.coderDecodeEncodeEqual( + coder, PaneInfo.createPane(true, true, timing, 0, 0, true)); + } + } + @Test public void testEncodings() { assertEquals( @@ -82,5 +98,9 @@ public void testEncodings() { "PaneInfo encoding should remain the same.", 0xF, PaneInfo.createPane(true, true, Timing.UNKNOWN).getEncodedByte()); + assertEquals( + "PaneInfo encoding should remain the same.", + 0x1, + PaneInfo.createPane(true, false, Timing.EARLY, 1, -1, true).getEncodedByte()); } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/WindowedValueTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/WindowedValueTest.java index 18660c5e6c36..957a839eae94 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/WindowedValueTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/WindowedValueTest.java @@ -77,6 +77,33 @@ public void testWindowedValueCoder() throws CoderException { Assert.assertArrayEquals(value.getWindows().toArray(), decodedValue.getWindows().toArray()); } + @Test + public void testWindowedValueWithElementMetadataCoder() throws CoderException { + WindowedValues.WindowedValueCoder.setMetadataSupported(); + PaneInfo.PaneInfoCoder.setMetadataSupported(); + Instant timestamp = new Instant(1234); + WindowedValue value = + WindowedValues.of( + "abc", + new Instant(1234), + Arrays.asList( + new IntervalWindow(timestamp, timestamp.plus(Duration.millis(1000))), + new IntervalWindow( + timestamp.plus(Duration.millis(1000)), timestamp.plus(Duration.millis(2000)))), + PaneInfo.NO_FIRING); + + Coder> windowedValueCoder = + WindowedValues.getFullCoder(StringUtf8Coder.of(), IntervalWindow.getCoder()); + + byte[] encodedValue = CoderUtils.encodeToByteArray(windowedValueCoder, value); + WindowedValue decodedValue = + CoderUtils.decodeFromByteArray(windowedValueCoder, encodedValue); + + Assert.assertEquals(value.getValue(), decodedValue.getValue()); + Assert.assertEquals(value.getTimestamp(), decodedValue.getTimestamp()); + Assert.assertArrayEquals(value.getWindows().toArray(), decodedValue.getWindows().toArray()); + } + @Test public void testFullWindowedValueCoderIsSerializableWithWellKnownCoderType() { CoderProperties.coderSerializable( From 7cfd6f8c0539c2ca7fac4f87498686b67e949964 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 3 Oct 2025 11:39:38 +0200 Subject: [PATCH 02/26] add todos for future otel metadata --- .../worker/UngroupedWindmillReader.java | 6 ++++- .../worker/WindmillKeyedWorkItem.java | 5 +++- .../runners/dataflow/worker/WindmillSink.java | 25 ++++++++++++++----- .../sdk/transforms/windowing/PaneInfo.java | 1 - 4 files changed, 28 insertions(+), 9 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java index 5df15e8d2109..d2e3e84ee085 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java @@ -117,7 +117,9 @@ protected WindowedValue decodeMessage(Windmill.Message message) throws IOExce Collection windows = WindmillSink.decodeMetadataWindows(windowsCoder, message.getMetadata()); PaneInfo paneInfo = WindmillSink.decodeMetadataPane(message.getMetadata()); - WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata()); + if(WindowedValues.WindowedValueCoder.isMetadataSupported()) { + WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata()); + } if (valueCoder instanceof KvCoder) { KvCoder kvCoder = (KvCoder) valueCoder; InputStream key = context.getSerializedKey().newInput(); @@ -126,9 +128,11 @@ protected WindowedValue decodeMessage(Windmill.Message message) throws IOExce @SuppressWarnings("unchecked") T result = (T) KV.of(decode(kvCoder.getKeyCoder(), key), decode(kvCoder.getValueCoder(), data)); + // todo #33176 propagate metadata to windowed value return WindowedValues.of(result, timestampMillis, windows, paneInfo); } else { notifyElementRead(data.available() + metadata.available()); + // todo #33176 propagate metadata to windowed value return WindowedValues.of(decode(valueCoder, data), timestampMillis, windows, paneInfo); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java index bae41a156967..1b5a52ad71cf 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java @@ -108,9 +108,12 @@ public Iterable> elementsIterable() { Collection windows = WindmillSink.decodeMetadataWindows(windowsCoder, message.getMetadata()); PaneInfo paneInfo = WindmillSink.decodeMetadataPane(message.getMetadata()); - WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata()); + if(WindowedValues.WindowedValueCoder.isMetadataSupported()) { + WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata()); + } InputStream inputStream = message.getData().newInput(); ElemT value = valueCoder.decode(inputStream, Coder.Context.OUTER); + // todo #33176 specify additional metadata in the future return WindowedValues.of(value, timestamp, windows, paneInfo); } catch (IOException e) { throw new RuntimeException(e); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java index 4765eb74e7f5..b9f3e235379b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java @@ -30,6 +30,7 @@ import org.apache.beam.runners.dataflow.util.CloudObject; import org.apache.beam.runners.dataflow.worker.util.common.worker.Sink; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; +import org.apache.beam.sdk.coders.ByteArrayCoder; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.options.PipelineOptions; @@ -81,8 +82,16 @@ private static ByteString encodeMetadata( BeamFnApi.Elements.ElementMetadata metadata) throws IOException { try { - PaneInfoCoder.INSTANCE.encode(paneInfo, stream); - windowsCoder.encode(windows, stream, Coder.Context.OUTER); + // element metadata is behind the experiment + boolean elementMetadata = WindowedValues.WindowedValueCoder.isMetadataSupported(); + if (elementMetadata) { + PaneInfoCoder.INSTANCE.encode(paneInfo, stream); + windowsCoder.encode(windows, stream); + ByteArrayCoder.of().encode(metadata.toByteArray(), stream, Coder.Context.OUTER); + }else { + PaneInfoCoder.INSTANCE.encode(paneInfo, stream); + windowsCoder.encode(windows, stream, Coder.Context.OUTER); + } return stream.toByteStringAndReset(); } catch (Exception e) { stream.reset(); @@ -93,10 +102,11 @@ private static ByteString encodeMetadata( public static ByteString encodeMetadata( Coder> windowsCoder, Collection windows, - PaneInfo paneInfo) + PaneInfo paneInfo, + BeamFnApi.Elements.ElementMetadata metadata) throws IOException { ByteStringOutputStream stream = new ByteStringOutputStream(); - return encodeMetadata(stream, windowsCoder, windows, paneInfo); + return encodeMetadata(stream, windowsCoder, windows, paneInfo, metadata); } public static PaneInfo decodeMetadataPane(ByteString metadata) throws IOException { @@ -110,7 +120,7 @@ public static BeamFnApi.Elements.ElementMetadata decodeAdditionalMetadata( InputStream inStream = metadata.newInput(); PaneInfo paneInfo = PaneInfoCoder.INSTANCE.decode(inStream); windowsCoder.decode(inStream); - if (paneInfo.isElementMetadata() && WindowedValues.WindowedValueCoder.isMetadataSupported()) { + if (paneInfo.isElementMetadata()) { return BeamFnApi.Elements.ElementMetadata.parseDelimitedFrom(inStream); } else { // empty @@ -201,8 +211,11 @@ private ByteString encode(Coder coder, EncodeT object) throws public long add(WindowedValue data) throws IOException { ByteString key, value; ByteString id = ByteString.EMPTY; + // todo - add here all windowedValue metadata + BeamFnApi.Elements.ElementMetadata additionalMetadata = + BeamFnApi.Elements.ElementMetadata.newBuilder().build(); ByteString metadata = - encodeMetadata(stream, windowsCoder, data.getWindows(), data.getPaneInfo()); + encodeMetadata(stream, windowsCoder, data.getWindows(), data.getPaneInfo(), additionalMetadata); if (valueCoder instanceof KvCoder) { KvCoder kvCoder = (KvCoder) valueCoder; KV kv = (KV) data.getValue(); diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PaneInfo.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PaneInfo.java index 80f931b8e6a6..bc83687bae4e 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PaneInfo.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PaneInfo.java @@ -198,7 +198,6 @@ public static PaneInfo createPane(boolean isFirst, boolean isLast, Timing timing return createPane(isFirst, isLast, timing, 0, timing == Timing.EARLY ? -1 : 0); } - /** Factory method to create a {@link PaneInfo} with the specified parameters. */ /** Factory method to create a {@link PaneInfo} with the specified parameters. */ public static PaneInfo createPane( boolean isFirst, boolean isLast, Timing timing, long index, long onTimeIndex) { From 5b7e0fc21da9904fb3c52888a7c6d47d78bb1a49 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 3 Oct 2025 14:03:23 +0200 Subject: [PATCH 03/26] add todos for future otel metadata --- .../RedistributeByKeyOverrideFactory.java | 1 + .../runners/dataflow/worker/WindmillSink.java | 4 ++-- .../beam/sdk/transforms/Redistribute.java | 1 + .../org/apache/beam/sdk/transforms/Reify.java | 1 + .../apache/beam/sdk/transforms/Reshuffle.java | 1 + .../beam/sdk/values/ValueInSingleWindow.java | 19 ++++++++++++++++++- .../beam/sdk/values/WindowedValues.java | 5 +++-- 7 files changed, 27 insertions(+), 5 deletions(-) diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/RedistributeByKeyOverrideFactory.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/RedistributeByKeyOverrideFactory.java index 4375cc5adcfe..47ff5b764910 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/RedistributeByKeyOverrideFactory.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/RedistributeByKeyOverrideFactory.java @@ -135,6 +135,7 @@ public Duration getAllowedTimestampSkew() { public void processElement( @Element KV> kv, OutputReceiver> outputReceiver) { + // todo #33176 specify additional metadata in the future outputReceiver .builder(KV.of(kv.getKey(), kv.getValue().getValue())) .setTimestamp(kv.getValue().getTimestamp()) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java index b9f3e235379b..3fe6af954cb5 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java @@ -121,7 +121,7 @@ public static BeamFnApi.Elements.ElementMetadata decodeAdditionalMetadata( PaneInfo paneInfo = PaneInfoCoder.INSTANCE.decode(inStream); windowsCoder.decode(inStream); if (paneInfo.isElementMetadata()) { - return BeamFnApi.Elements.ElementMetadata.parseDelimitedFrom(inStream); + return BeamFnApi.Elements.ElementMetadata.parseFrom(ByteArrayCoder.of().decode(inStream)); } else { // empty return BeamFnApi.Elements.ElementMetadata.newBuilder().build(); @@ -211,7 +211,7 @@ private ByteString encode(Coder coder, EncodeT object) throws public long add(WindowedValue data) throws IOException { ByteString key, value; ByteString id = ByteString.EMPTY; - // todo - add here all windowedValue metadata + // todo #33176 specify additional metadata in the future BeamFnApi.Elements.ElementMetadata additionalMetadata = BeamFnApi.Elements.ElementMetadata.newBuilder().build(); ByteString metadata = diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java index a01b5f570a57..7382f29310c5 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java @@ -179,6 +179,7 @@ public Duration getAllowedTimestampSkew() { public void processElement( @Element KV> kv, OutputReceiver> outputReceiver) { + // todo #33176 specify additional metadata in the future outputReceiver .builder(KV.of(kv.getKey(), kv.getValue().getValue())) .setTimestamp(kv.getValue().getTimestamp()) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reify.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reify.java index 797af9538c53..905e94ccc4d0 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reify.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reify.java @@ -136,6 +136,7 @@ public PCollection>> expand(PCollection> i KvCoder coder = (KvCoder) input.getCoder(); return input .apply( + // todo #33176 specify additional metadata in the future ParDo.of( new DoFn, KV>>() { @ProcessElement diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reshuffle.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reshuffle.java index b2de48342d7c..0a8d058107b8 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reshuffle.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reshuffle.java @@ -184,6 +184,7 @@ public Duration getAllowedTimestampSkew() { public void processElement( @Element KV> kv, OutputReceiver> outputReceiver) { + // todo #33176 specify additional metadata in the future outputReceiver .builder(KV.of(kv.getKey(), kv.getValue().getValue())) .setTimestamp(kv.getValue().getTimestamp()) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java index 7dc5fef52ecb..56cb58ed6d15 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java @@ -22,7 +22,9 @@ import java.io.InputStream; import java.io.OutputStream; import java.util.List; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.coders.ByteArrayCoder; import org.apache.beam.sdk.coders.InstantCoder; import org.apache.beam.sdk.coders.StructuredCoder; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; @@ -64,6 +66,7 @@ public T getValue() { public abstract @Nullable Long getCurrentRecordOffset(); + // todo #33176 specify additional metadata in the future public static ValueInSingleWindow of( T value, Instant timestamp, @@ -110,7 +113,16 @@ public void encode(ValueInSingleWindow windowedElem, OutputStream outStream, throws IOException { InstantCoder.of().encode(windowedElem.getTimestamp(), outStream); windowCoder.encode(windowedElem.getWindow(), outStream); - PaneInfo.PaneInfoCoder.INSTANCE.encode(windowedElem.getPaneInfo(), outStream); + boolean metadataSupported = WindowedValues.WindowedValueCoder.isMetadataSupported(); + PaneInfo.PaneInfoCoder.INSTANCE.encode(windowedElem.getPaneInfo().withElementMetadata(metadataSupported), outStream); + if(metadataSupported){ + BeamFnApi.Elements.ElementMetadata.Builder builder = + BeamFnApi.Elements.ElementMetadata.newBuilder(); + // todo #33176 specify additional metadata in the future + BeamFnApi.Elements.ElementMetadata metadata = builder.build(); + ByteArrayCoder.of().encode(metadata.toByteArray(), outStream); + } + valueCoder.encode(windowedElem.getValue(), outStream, context); } @@ -124,7 +136,12 @@ public ValueInSingleWindow decode(InputStream inStream, Context context) thro Instant timestamp = InstantCoder.of().decode(inStream); BoundedWindow window = windowCoder.decode(inStream); PaneInfo paneInfo = PaneInfo.PaneInfoCoder.INSTANCE.decode(inStream); + if(WindowedValues.WindowedValueCoder.isMetadataSupported() && paneInfo.isElementMetadata()) { + BeamFnApi.Elements.ElementMetadata.parseFrom(ByteArrayCoder.of().decode(inStream)); + } + T value = valueCoder.decode(inStream, context); + // todo #33176 specify additional metadata in the future return new AutoValue_ValueInSingleWindow<>(value, timestamp, window, paneInfo, null, null); } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java index 26ba4948bc0a..91763002d294 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java @@ -846,7 +846,8 @@ public void encode(WindowedValue windowedElem, OutputStream outStream, Contex BeamFnApi.Elements.ElementMetadata.Builder builder = BeamFnApi.Elements.ElementMetadata.newBuilder(); BeamFnApi.Elements.ElementMetadata em = builder.build(); - em.writeDelimitedTo(outStream); + ByteArrayCoder.of().encode(em.toByteArray(), outStream); + } valueCoder.encode(windowedElem.getValue(), outStream, context); } @@ -863,7 +864,7 @@ public WindowedValue decode(InputStream inStream, Context context) Collection windows = windowsCoder.decode(inStream); PaneInfo paneInfo = PaneInfoCoder.INSTANCE.decode(inStream); if (isMetadataSupported() && paneInfo.isElementMetadata()) { - BeamFnApi.Elements.ElementMetadata.parseDelimitedFrom(inStream); + BeamFnApi.Elements.ElementMetadata.parseFrom(ByteArrayCoder.of().decode(inStream)); } T value = valueCoder.decode(inStream, context); From 7654478a5d2495382cdaa4e6ab57f65e36492e8a Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 3 Oct 2025 14:27:29 +0200 Subject: [PATCH 04/26] spotless --- .../runners/dataflow/worker/UngroupedWindmillReader.java | 2 +- .../runners/dataflow/worker/WindmillKeyedWorkItem.java | 2 +- .../beam/runners/dataflow/worker/WindmillSink.java | 7 ++++--- .../main/java/org/apache/beam/sdk/transforms/Reify.java | 2 +- .../org/apache/beam/sdk/values/ValueInSingleWindow.java | 9 +++++---- .../java/org/apache/beam/sdk/values/WindowedValues.java | 1 - 6 files changed, 12 insertions(+), 11 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java index d2e3e84ee085..a9a033c89ad7 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java @@ -117,7 +117,7 @@ protected WindowedValue decodeMessage(Windmill.Message message) throws IOExce Collection windows = WindmillSink.decodeMetadataWindows(windowsCoder, message.getMetadata()); PaneInfo paneInfo = WindmillSink.decodeMetadataPane(message.getMetadata()); - if(WindowedValues.WindowedValueCoder.isMetadataSupported()) { + if (WindowedValues.WindowedValueCoder.isMetadataSupported()) { WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata()); } if (valueCoder instanceof KvCoder) { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java index 1b5a52ad71cf..6690377d3de6 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java @@ -108,7 +108,7 @@ public Iterable> elementsIterable() { Collection windows = WindmillSink.decodeMetadataWindows(windowsCoder, message.getMetadata()); PaneInfo paneInfo = WindmillSink.decodeMetadataPane(message.getMetadata()); - if(WindowedValues.WindowedValueCoder.isMetadataSupported()) { + if (WindowedValues.WindowedValueCoder.isMetadataSupported()) { WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata()); } InputStream inputStream = message.getData().newInput(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java index 3fe6af954cb5..4c5841f2de97 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java @@ -88,7 +88,7 @@ private static ByteString encodeMetadata( PaneInfoCoder.INSTANCE.encode(paneInfo, stream); windowsCoder.encode(windows, stream); ByteArrayCoder.of().encode(metadata.toByteArray(), stream, Coder.Context.OUTER); - }else { + } else { PaneInfoCoder.INSTANCE.encode(paneInfo, stream); windowsCoder.encode(windows, stream, Coder.Context.OUTER); } @@ -213,9 +213,10 @@ public long add(WindowedValue data) throws IOException { ByteString id = ByteString.EMPTY; // todo #33176 specify additional metadata in the future BeamFnApi.Elements.ElementMetadata additionalMetadata = - BeamFnApi.Elements.ElementMetadata.newBuilder().build(); + BeamFnApi.Elements.ElementMetadata.newBuilder().build(); ByteString metadata = - encodeMetadata(stream, windowsCoder, data.getWindows(), data.getPaneInfo(), additionalMetadata); + encodeMetadata( + stream, windowsCoder, data.getWindows(), data.getPaneInfo(), additionalMetadata); if (valueCoder instanceof KvCoder) { KvCoder kvCoder = (KvCoder) valueCoder; KV kv = (KV) data.getValue(); diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reify.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reify.java index 905e94ccc4d0..af125d9e63e8 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reify.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reify.java @@ -136,7 +136,7 @@ public PCollection>> expand(PCollection> i KvCoder coder = (KvCoder) input.getCoder(); return input .apply( - // todo #33176 specify additional metadata in the future + // todo #33176 specify additional metadata in the future ParDo.of( new DoFn, KV>>() { @ProcessElement diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java index 56cb58ed6d15..2bbc28cac9fe 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java @@ -114,10 +114,11 @@ public void encode(ValueInSingleWindow windowedElem, OutputStream outStream, InstantCoder.of().encode(windowedElem.getTimestamp(), outStream); windowCoder.encode(windowedElem.getWindow(), outStream); boolean metadataSupported = WindowedValues.WindowedValueCoder.isMetadataSupported(); - PaneInfo.PaneInfoCoder.INSTANCE.encode(windowedElem.getPaneInfo().withElementMetadata(metadataSupported), outStream); - if(metadataSupported){ + PaneInfo.PaneInfoCoder.INSTANCE.encode( + windowedElem.getPaneInfo().withElementMetadata(metadataSupported), outStream); + if (metadataSupported) { BeamFnApi.Elements.ElementMetadata.Builder builder = - BeamFnApi.Elements.ElementMetadata.newBuilder(); + BeamFnApi.Elements.ElementMetadata.newBuilder(); // todo #33176 specify additional metadata in the future BeamFnApi.Elements.ElementMetadata metadata = builder.build(); ByteArrayCoder.of().encode(metadata.toByteArray(), outStream); @@ -136,7 +137,7 @@ public ValueInSingleWindow decode(InputStream inStream, Context context) thro Instant timestamp = InstantCoder.of().decode(inStream); BoundedWindow window = windowCoder.decode(inStream); PaneInfo paneInfo = PaneInfo.PaneInfoCoder.INSTANCE.decode(inStream); - if(WindowedValues.WindowedValueCoder.isMetadataSupported() && paneInfo.isElementMetadata()) { + if (WindowedValues.WindowedValueCoder.isMetadataSupported() && paneInfo.isElementMetadata()) { BeamFnApi.Elements.ElementMetadata.parseFrom(ByteArrayCoder.of().decode(inStream)); } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java index 91763002d294..d90637da9b1b 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java @@ -847,7 +847,6 @@ public void encode(WindowedValue windowedElem, OutputStream outStream, Contex BeamFnApi.Elements.ElementMetadata.newBuilder(); BeamFnApi.Elements.ElementMetadata em = builder.build(); ByteArrayCoder.of().encode(em.toByteArray(), outStream); - } valueCoder.encode(windowedElem.getValue(), outStream, context); } From 0274977bacf313b561559d1c31241bfdc608bc16 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 3 Oct 2025 14:56:47 +0200 Subject: [PATCH 05/26] fix tests --- .../test/java/org/apache/beam/sdk/util/WindowedValueTest.java | 1 - 1 file changed, 1 deletion(-) diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/WindowedValueTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/WindowedValueTest.java index 957a839eae94..3e3973e3720b 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/WindowedValueTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/WindowedValueTest.java @@ -80,7 +80,6 @@ public void testWindowedValueCoder() throws CoderException { @Test public void testWindowedValueWithElementMetadataCoder() throws CoderException { WindowedValues.WindowedValueCoder.setMetadataSupported(); - PaneInfo.PaneInfoCoder.setMetadataSupported(); Instant timestamp = new Instant(1234); WindowedValue value = WindowedValues.of( From bbb947471153f17d14417d22639aae1ad1fb27e3 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 3 Oct 2025 19:33:18 +0200 Subject: [PATCH 06/26] fix --- .../org/apache/beam/runners/dataflow/worker/WindmillSink.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java index 4c5841f2de97..18c023d62aa8 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java @@ -85,7 +85,7 @@ private static ByteString encodeMetadata( // element metadata is behind the experiment boolean elementMetadata = WindowedValues.WindowedValueCoder.isMetadataSupported(); if (elementMetadata) { - PaneInfoCoder.INSTANCE.encode(paneInfo, stream); + PaneInfoCoder.INSTANCE.encode(paneInfo.withElementMetadata(elementMetadata), stream); windowsCoder.encode(windows, stream); ByteArrayCoder.of().encode(metadata.toByteArray(), stream, Coder.Context.OUTER); } else { From 99edc7650e653ccf899be99f49e532f8479b5a23 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 3 Oct 2025 19:33:41 +0200 Subject: [PATCH 07/26] fix --- .../org/apache/beam/runners/dataflow/worker/WindmillSink.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java index 18c023d62aa8..6d3484f6642d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java @@ -85,7 +85,7 @@ private static ByteString encodeMetadata( // element metadata is behind the experiment boolean elementMetadata = WindowedValues.WindowedValueCoder.isMetadataSupported(); if (elementMetadata) { - PaneInfoCoder.INSTANCE.encode(paneInfo.withElementMetadata(elementMetadata), stream); + PaneInfoCoder.INSTANCE.encode(paneInfo.withElementMetadata(true), stream); windowsCoder.encode(windows, stream); ByteArrayCoder.of().encode(metadata.toByteArray(), stream, Coder.Context.OUTER); } else { From fb858321752c568f75c83209f1e9c11246740980 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Thu, 9 Oct 2025 13:48:11 +0200 Subject: [PATCH 08/26] proto change --- .../beam/model/fn_execution/v1/beam_fn_api.proto | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto b/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto index 9360522ab409..9b32048b4995 100644 --- a/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto +++ b/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto @@ -740,10 +740,18 @@ message Elements { bool is_last = 4; } + message DrainMode { + enum Enum { + UNSPECIFIED = 0; + NOT_DRAINING = 1; + DRAINING = 2; + } + } + // Element metadata passed as part of WindowedValue to make WindowedValue // extensible and backward compatible message ElementMetadata { - // empty message - add drain, kind, tracing metadata in the future + optional DrainMode.Enum drain = 1; } // Represent the encoded user timer for a given instruction, transform and From 373f159fd24bd4c2a3a9f674726ee8e2d6397b1f Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Thu, 9 Oct 2025 13:50:14 +0200 Subject: [PATCH 09/26] output builder change --- .../apache/beam/sdk/values/OutputBuilder.java | 2 + .../apache/beam/sdk/values/WindowedValue.java | 3 + .../beam/sdk/values/WindowedValues.java | 99 +++++++++++++------ 3 files changed, 76 insertions(+), 28 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/OutputBuilder.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/OutputBuilder.java index a7f8bc8e03b1..5762d32ae832 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/OutputBuilder.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/OutputBuilder.java @@ -48,5 +48,7 @@ public interface OutputBuilder extends WindowedValue { OutputBuilder setRecordOffset(@Nullable Long recordOffset); + OutputBuilder setDraining(@Nullable Boolean drain); + void output(); } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValue.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValue.java index ea6be129ecb4..762a602bc3f9 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValue.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValue.java @@ -52,6 +52,9 @@ public interface WindowedValue { @Nullable Long getRecordOffset(); + @Nullable + Boolean isDraining(); + /** * A representation of each of the actual values represented by this compressed {@link * WindowedValue}, one per window. diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java index d90637da9b1b..ed51be43ec1d 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java @@ -99,6 +99,7 @@ public static class Builder implements OutputBuilder { private @MonotonicNonNull Collection windows; private @Nullable String recordId; private @Nullable Long recordOffset; + private @Nullable Boolean draining; @Override public Builder setValue(T value) { @@ -142,6 +143,12 @@ public Builder setRecordOffset(@Nullable Long recordOffset) { return this; } + @Override + public Builder setDraining(@Nullable Boolean draining) { + this.draining = draining; + return this; + } + public Builder setReceiver(WindowedValueReceiver receiver) { this.receiver = receiver; return this; @@ -190,6 +197,11 @@ public PaneInfo getPaneInfo() { return recordOffset; } + @Override + public @Nullable Boolean isDraining() { + return draining; + } + @Override public Collection> explodeWindows() { throw new UnsupportedOperationException( @@ -235,7 +247,7 @@ public String toString() { public static WindowedValue of( T value, Instant timestamp, Collection windows, PaneInfo paneInfo) { - return of(value, timestamp, windows, paneInfo, null, null); + return of(value, timestamp, windows, paneInfo, null, null, null); } /** Returns a {@code WindowedValue} with the given value, timestamp, and windows. */ @@ -245,7 +257,8 @@ public static WindowedValue of( Collection windows, PaneInfo paneInfo, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { checkArgument(paneInfo != null, "WindowedValue requires PaneInfo, but it was null"); checkArgument(windows.size() > 0, "WindowedValue requires windows, but there were none"); @@ -253,7 +266,7 @@ public static WindowedValue of( return of(value, timestamp, windows.iterator().next(), paneInfo); } else { return new TimestampedValueInMultipleWindows<>( - value, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset); + value, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset, draining); } } @@ -265,7 +278,7 @@ static WindowedValue createWithoutValidation( return of(value, timestamp, windows.iterator().next(), paneInfo); } else { return new TimestampedValueInMultipleWindows<>( - value, timestamp, windows, paneInfo, null, null); + value, timestamp, windows, paneInfo, null, null, null); } } @@ -278,9 +291,10 @@ public static WindowedValue of( if (isGlobal && BoundedWindow.TIMESTAMP_MIN_VALUE.equals(timestamp)) { return valueInGlobalWindow(value, paneInfo); } else if (isGlobal) { - return new TimestampedValueInGlobalWindow<>(value, timestamp, paneInfo, null, null); + return new TimestampedValueInGlobalWindow<>(value, timestamp, paneInfo, null, null, null); } else { - return new TimestampedValueInSingleWindow<>(value, timestamp, window, paneInfo, null, null); + return new TimestampedValueInSingleWindow<>( + value, timestamp, window, paneInfo, null, null, null); } } @@ -289,7 +303,7 @@ public static WindowedValue of( * default timestamp and pane. */ public static WindowedValue valueInGlobalWindow(T value) { - return new ValueInGlobalWindow<>(value, PaneInfo.NO_FIRING, null, null); + return new ValueInGlobalWindow<>(value, PaneInfo.NO_FIRING, null, null, null); } /** @@ -297,7 +311,7 @@ public static WindowedValue valueInGlobalWindow(T value) { * default timestamp and the specified pane. */ public static WindowedValue valueInGlobalWindow(T value, PaneInfo paneInfo) { - return new ValueInGlobalWindow<>(value, paneInfo, null, null); + return new ValueInGlobalWindow<>(value, paneInfo, null, null, null); } /** @@ -308,7 +322,8 @@ public static WindowedValue timestampedValueInGlobalWindow(T value, Insta if (BoundedWindow.TIMESTAMP_MIN_VALUE.equals(timestamp)) { return valueInGlobalWindow(value); } else { - return new TimestampedValueInGlobalWindow<>(value, timestamp, PaneInfo.NO_FIRING, null, null); + return new TimestampedValueInGlobalWindow<>( + value, timestamp, PaneInfo.NO_FIRING, null, null, null); } } @@ -321,7 +336,7 @@ public static WindowedValue timestampedValueInGlobalWindow( if (paneInfo.equals(PaneInfo.NO_FIRING)) { return timestampedValueInGlobalWindow(value, timestamp); } else { - return new TimestampedValueInGlobalWindow<>(value, timestamp, paneInfo, null, null); + return new TimestampedValueInGlobalWindow<>(value, timestamp, paneInfo, null, null, null); } } @@ -337,7 +352,8 @@ public static WindowedValue withValue( windowedValue.getWindows(), windowedValue.getPaneInfo(), windowedValue.getRecordId(), - windowedValue.getRecordOffset()); + windowedValue.getRecordOffset(), + windowedValue.isDraining()); } public static boolean equals( @@ -388,6 +404,7 @@ private abstract static class SimpleWindowedValue implements WindowedValue private final PaneInfo paneInfo; private final @Nullable String currentRecordId; private final @Nullable Long currentRecordOffset; + private final @Nullable Boolean draining; @Override public @Nullable String getRecordId() { @@ -399,15 +416,22 @@ private abstract static class SimpleWindowedValue implements WindowedValue return currentRecordOffset; } + @Override + public @Nullable Boolean isDraining() { + return draining; + } + protected SimpleWindowedValue( T value, PaneInfo paneInfo, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { this.value = value; this.paneInfo = checkNotNull(paneInfo); this.currentRecordId = currentRecordId; this.currentRecordOffset = currentRecordOffset; + this.draining = draining; } @Override @@ -455,8 +479,9 @@ public MinTimestampWindowedValue( T value, PaneInfo pane, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { - super(value, pane, currentRecordId, currentRecordOffset); + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { + super(value, pane, currentRecordId, currentRecordOffset, draining); } @Override @@ -473,8 +498,9 @@ public ValueInGlobalWindow( T value, PaneInfo paneInfo, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { - super(value, paneInfo, currentRecordId, currentRecordOffset); + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { + super(value, paneInfo, currentRecordId, currentRecordOffset, draining); } @Override @@ -489,7 +515,8 @@ public BoundedWindow getWindow() { @Override public WindowedValue withValue(NewT newValue) { - return new ValueInGlobalWindow<>(newValue, getPaneInfo(), getRecordId(), getRecordOffset()); + return new ValueInGlobalWindow<>( + newValue, getPaneInfo(), getRecordId(), getRecordOffset(), isDraining()); } @Override @@ -526,8 +553,9 @@ public TimestampedWindowedValue( Instant timestamp, PaneInfo paneInfo, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { - super(value, paneInfo, currentRecordId, currentRecordOffset); + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { + super(value, paneInfo, currentRecordId, currentRecordOffset, draining); this.timestamp = checkNotNull(timestamp); } @@ -549,8 +577,9 @@ public TimestampedValueInGlobalWindow( Instant timestamp, PaneInfo paneInfo, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { - super(value, timestamp, paneInfo, currentRecordId, currentRecordOffset); + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { + super(value, timestamp, paneInfo, currentRecordId, currentRecordOffset, draining); } @Override @@ -566,7 +595,7 @@ public BoundedWindow getWindow() { @Override public WindowedValue withValue(NewT newValue) { return new TimestampedValueInGlobalWindow<>( - newValue, getTimestamp(), getPaneInfo(), getRecordId(), getRecordOffset()); + newValue, getTimestamp(), getPaneInfo(), getRecordId(), getRecordOffset(), isDraining()); } @Override @@ -615,15 +644,22 @@ public TimestampedValueInSingleWindow( BoundedWindow window, PaneInfo paneInfo, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { - super(value, timestamp, paneInfo, currentRecordId, currentRecordOffset); + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { + super(value, timestamp, paneInfo, currentRecordId, currentRecordOffset, draining); this.window = checkNotNull(window); } @Override public WindowedValue withValue(NewT newValue) { return new TimestampedValueInSingleWindow<>( - newValue, getTimestamp(), window, getPaneInfo(), getRecordId(), getRecordOffset()); + newValue, + getTimestamp(), + window, + getPaneInfo(), + getRecordId(), + getRecordOffset(), + isDraining()); } @Override @@ -679,8 +715,9 @@ public TimestampedValueInMultipleWindows( Collection windows, PaneInfo paneInfo, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { - super(value, timestamp, paneInfo, currentRecordId, currentRecordOffset); + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { + super(value, timestamp, paneInfo, currentRecordId, currentRecordOffset, draining); this.windows = checkNotNull(windows); } @@ -692,7 +729,13 @@ public Collection getWindows() { @Override public WindowedValue withValue(NewT newValue) { return new TimestampedValueInMultipleWindows<>( - newValue, getTimestamp(), getWindows(), getPaneInfo(), getRecordId(), getRecordOffset()); + newValue, + getTimestamp(), + getWindows(), + getPaneInfo(), + getRecordId(), + getRecordOffset(), + isDraining()); } @Override From aabe6d18393f41774ddd45c7987f37e145b6e574 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Thu, 9 Oct 2025 14:16:11 +0200 Subject: [PATCH 10/26] propagate drainMode and Element Metadata's drain enum from Windmill to WindowedValue and TimerData --- .../beam/runners/core/TimerInternals.java | 37 ++++++++++++++----- .../beam/runners/core/TimerInternalsTest.java | 32 +++++++++------- .../dataflow/worker/SimpleParDoFn.java | 3 +- .../worker/StreamingDataflowWorker.java | 2 + ...eamingGroupAlsoByWindowViaWindowSetFn.java | 1 + .../worker/StreamingModeExecutionContext.java | 14 ++++++- .../worker/UngroupedWindmillReader.java | 21 +++++++++-- .../worker/WindmillKeyedWorkItem.java | 24 ++++++++++-- .../worker/WindmillTimerInternals.java | 14 +++++-- .../worker/WindowingWindmillReader.java | 3 +- .../dataflow/worker/streaming/Work.java | 12 +++++- .../harness/SingleSourceWorkerHarness.java | 3 ++ ...tchGroupAlsoByWindowViaOutputBufferFn.java | 1 + .../grpc/GetWorkResponseChunkAssembler.java | 5 ++- .../client/grpc/GrpcDirectGetWorkStream.java | 1 + .../client/grpc/GrpcGetWorkStream.java | 1 + .../windmill/work/WorkItemReceiver.java | 1 + .../windmill/work/WorkItemScheduler.java | 3 ++ .../processing/StreamingWorkScheduler.java | 4 +- .../dataflow/worker/FakeWindmillServer.java | 1 + .../worker/StreamingDataflowWorkerTest.java | 2 + .../StreamingGroupAlsoByWindowFnsTest.java | 5 ++- ...ngGroupAlsoByWindowsReshuffleDoFnTest.java | 2 +- ...gKeyedWorkItemSideInputDoFnRunnerTest.java | 3 +- .../StreamingModeExecutionContextTest.java | 4 +- .../worker/StreamingSideInputFetcherTest.java | 3 +- .../worker/UserParDoFnFactoryTest.java | 9 +++-- .../worker/WindmillKeyedWorkItemTest.java | 9 +++-- .../worker/WindmillTimerInternalsTest.java | 33 ++++++++++++----- .../worker/WorkerCustomSourcesTest.java | 2 + .../worker/streaming/ActiveWorkStateTest.java | 2 + .../streaming/ComputationStateCacheTest.java | 1 + ...anOutStreamingEngineWorkerHarnessTest.java | 1 + .../harness/WindmillStreamSenderTest.java | 1 + .../worker/util/BoundedQueueExecutorTest.java | 1 + .../StreamingApplianceWorkCommitterTest.java | 1 + .../StreamingEngineWorkCommitterTest.java | 1 + .../grpc/GrpcDirectGetWorkStreamTest.java | 12 +++++- .../client/grpc/GrpcWindmillServerTest.java | 2 + .../failures/WorkFailureProcessorTest.java | 1 + .../work/refresh/ActiveWorkRefresherTest.java | 1 + .../control/TimerReceiverFactory.java | 3 +- 42 files changed, 215 insertions(+), 67 deletions(-) diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/TimerInternals.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/TimerInternals.java index 254e6f5fcf5b..dbf96e256970 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/TimerInternals.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/TimerInternals.java @@ -23,9 +23,11 @@ import java.io.OutputStream; import java.util.Arrays; import java.util.List; +import org.apache.beam.sdk.coders.BooleanCoder; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.CoderException; import org.apache.beam.sdk.coders.InstantCoder; +import org.apache.beam.sdk.coders.NullableCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.coders.StructuredCoder; import org.apache.beam.sdk.state.TimeDomain; @@ -188,6 +190,8 @@ abstract class TimerData implements Comparable { public abstract boolean getDeleted(); + public abstract @javax.annotation.Nullable Boolean getDraining(); + // When adding a new field, make sure to add it to the compareTo() method. /** Construct a {@link TimerData} for the given parameters. */ @@ -196,9 +200,10 @@ public static TimerData of( StateNamespace namespace, Instant timestamp, Instant outputTimestamp, - TimeDomain domain) { + TimeDomain domain, + @javax.annotation.Nullable Boolean draining) { return new AutoValue_TimerInternals_TimerData( - timerId, "", namespace, timestamp, outputTimestamp, domain, false); + timerId, "", namespace, timestamp, outputTimestamp, domain, false, draining); } /** @@ -211,9 +216,10 @@ public static TimerData of( StateNamespace namespace, Instant timestamp, Instant outputTimestamp, - TimeDomain domain) { + TimeDomain domain, + @Nullable Boolean draining) { return new AutoValue_TimerInternals_TimerData( - timerId, timerFamilyId, namespace, timestamp, outputTimestamp, domain, false); + timerId, timerFamilyId, namespace, timestamp, outputTimestamp, domain, false, draining); } /** @@ -221,9 +227,13 @@ public static TimerData of( * deterministically generated from the {@code timestamp} and {@code domain}. */ public static TimerData of( - StateNamespace namespace, Instant timestamp, Instant outputTimestamp, TimeDomain domain) { + StateNamespace namespace, + Instant timestamp, + Instant outputTimestamp, + TimeDomain domain, + @Nullable Boolean draining) { String timerId = String.valueOf(domain.ordinal()) + ':' + timestamp.getMillis(); - return of(timerId, namespace, timestamp, outputTimestamp, domain); + return of(timerId, namespace, timestamp, outputTimestamp, domain, draining); } public TimerData deleted() { @@ -234,7 +244,8 @@ public TimerData deleted() { getTimestamp(), getOutputTimestamp(), getDomain(), - true); + true, + getDraining()); } /** @@ -272,7 +283,9 @@ public String stringKey() { + "/" + getTimerFamilyId() + ":" - + getTimerId(); + + getTimerId() + + ":" + + getDraining(); } } @@ -280,6 +293,7 @@ public String stringKey() { class TimerDataCoderV2 extends StructuredCoder { private static final StringUtf8Coder STRING_CODER = StringUtf8Coder.of(); private static final InstantCoder INSTANT_CODER = InstantCoder.of(); + private static final NullableCoder BOOLEAN_CODER = NullableCoder.of(BooleanCoder.of()); private final Coder windowCoder; public static TimerDataCoderV2 of(Coder windowCoder) { @@ -298,6 +312,7 @@ public void encode(TimerData timer, OutputStream outStream) throws CoderExceptio INSTANT_CODER.encode(timer.getTimestamp(), outStream); INSTANT_CODER.encode(timer.getOutputTimestamp(), outStream); STRING_CODER.encode(timer.getDomain().name(), outStream); + // we don't need to serialize draining bit } @Override @@ -309,7 +324,9 @@ public TimerData decode(InputStream inStream) throws CoderException, IOException Instant timestamp = INSTANT_CODER.decode(inStream); Instant outputTimestamp = INSTANT_CODER.decode(inStream); TimeDomain domain = TimeDomain.valueOf(STRING_CODER.decode(inStream)); - return TimerData.of(timerId, timerFamilyId, namespace, timestamp, outputTimestamp, domain); + // we don't need to serialize draining bit + return TimerData.of( + timerId, timerFamilyId, namespace, timestamp, outputTimestamp, domain, false); } @Override @@ -355,7 +372,7 @@ public TimerData decode(InputStream inStream) throws CoderException, IOException StateNamespaces.fromString(STRING_CODER.decode(inStream), windowCoder); Instant timestamp = INSTANT_CODER.decode(inStream); TimeDomain domain = TimeDomain.valueOf(STRING_CODER.decode(inStream)); - return TimerData.of(timerId, namespace, timestamp, timestamp, domain); + return TimerData.of(timerId, namespace, timestamp, timestamp, domain, false); } @Override diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/TimerInternalsTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/TimerInternalsTest.java index 83fb30eb8120..640c2bb7128e 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/TimerInternalsTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/TimerInternalsTest.java @@ -46,7 +46,8 @@ public void testTimerDataCoder() throws Exception { StateNamespaces.global(), new Instant(0), new Instant(0), - TimeDomain.EVENT_TIME)); + TimeDomain.EVENT_TIME, + null)); Coder windowCoder = IntervalWindow.getCoder(); CoderProperties.coderDecodeEncodeEqual( @@ -57,7 +58,8 @@ public void testTimerDataCoder() throws Exception { windowCoder, new IntervalWindow(new Instant(0), new Instant(100))), new Instant(99), new Instant(99), - TimeDomain.PROCESSING_TIME)); + TimeDomain.PROCESSING_TIME, + null)); } @Test @@ -69,12 +71,13 @@ public void testCoderIsSerializableWithWellKnownCoderType() { public void testCompareEqual() { Instant timestamp = new Instant(100); StateNamespace namespace = StateNamespaces.global(); - TimerData timer = TimerData.of("id", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME); + TimerData timer = + TimerData.of("id", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME, null); assertThat( timer, comparesEqualTo( - TimerData.of("id", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME))); + TimerData.of("id", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME, null))); } @Test @@ -84,9 +87,9 @@ public void testCompareByTimestamp() { StateNamespace namespace = StateNamespaces.global(); TimerData firstTimer = - TimerData.of(namespace, firstTimestamp, firstTimestamp, TimeDomain.EVENT_TIME); + TimerData.of(namespace, firstTimestamp, firstTimestamp, TimeDomain.EVENT_TIME, null); TimerData secondTimer = - TimerData.of(namespace, secondTimestamp, secondTimestamp, TimeDomain.EVENT_TIME); + TimerData.of(namespace, secondTimestamp, secondTimestamp, TimeDomain.EVENT_TIME, null); assertThat(firstTimer, lessThan(secondTimer)); } @@ -96,10 +99,13 @@ public void testCompareByDomain() { Instant timestamp = new Instant(100); StateNamespace namespace = StateNamespaces.global(); - TimerData eventTimer = TimerData.of(namespace, timestamp, timestamp, TimeDomain.EVENT_TIME); - TimerData procTimer = TimerData.of(namespace, timestamp, timestamp, TimeDomain.PROCESSING_TIME); + TimerData eventTimer = + TimerData.of(namespace, timestamp, timestamp, TimeDomain.EVENT_TIME, null); + TimerData procTimer = + TimerData.of(namespace, timestamp, timestamp, TimeDomain.PROCESSING_TIME, null); TimerData synchronizedProcTimer = - TimerData.of(namespace, timestamp, timestamp, TimeDomain.SYNCHRONIZED_PROCESSING_TIME); + TimerData.of( + namespace, timestamp, timestamp, TimeDomain.SYNCHRONIZED_PROCESSING_TIME, null); assertThat(eventTimer, lessThan(procTimer)); assertThat(eventTimer, lessThan(synchronizedProcTimer)); @@ -117,9 +123,9 @@ public void testCompareByNamespace() { StateNamespace secondWindowNs = StateNamespaces.window(windowCoder, secondWindow); TimerData secondEventTime = - TimerData.of(firstWindowNs, timestamp, timestamp, TimeDomain.EVENT_TIME); + TimerData.of(firstWindowNs, timestamp, timestamp, TimeDomain.EVENT_TIME, null); TimerData thirdEventTime = - TimerData.of(secondWindowNs, timestamp, timestamp, TimeDomain.EVENT_TIME); + TimerData.of(secondWindowNs, timestamp, timestamp, TimeDomain.EVENT_TIME, null); assertThat(secondEventTime, lessThan(thirdEventTime)); } @@ -130,9 +136,9 @@ public void testCompareByTimerId() { StateNamespace namespace = StateNamespaces.global(); TimerData id0Timer = - TimerData.of("id0", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME); + TimerData.of("id0", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME, null); TimerData id1Timer = - TimerData.of("id1", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME); + TimerData.of("id1", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME, null); assertThat(id0Timer, lessThan(id1Timer)); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java index 00c8192b1e40..d3f280085155 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFn.java @@ -374,7 +374,8 @@ private void processUserTimer(TimerData timer) throws Exception { window, timer.getTimestamp(), timer.getOutputTimestamp(), - timer.getDomain()); + timer.getDomain(), + timer.getDraining()); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java index 59269f8ea55e..31729c3b71a9 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java @@ -263,6 +263,7 @@ private StreamingDataflowWorker( serializedWorkItemSize, watermarks, processingContext, + drainMode, getWorkStreamLatencies) -> computationStateCache .get(processingContext.computationId()) @@ -275,6 +276,7 @@ private StreamingDataflowWorker( serializedWorkItemSize, watermarks, processingContext, + drainMode, getWorkStreamLatencies); }), ChannelCachingRemoteStubFactory.create(options.getGcpCredential(), channelCache), diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowViaWindowSetFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowViaWindowSetFn.java index ec36644d1e68..e412f28a7583 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowViaWindowSetFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowViaWindowSetFn.java @@ -91,6 +91,7 @@ public void processElement( output, sideInputReader, reduceFn, + null, // elements and timers have relevant information about draining options); reduceFnRunner.processElements(keyedWorkItem.elementsIterable()); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java index b24ca561495c..82d18b0ab788 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java @@ -195,6 +195,10 @@ public boolean workIsFailed() { return work != null && work.isFailed(); } + public @javax.annotation.Nullable Boolean getDrainMode() { + return work != null ? work.getDrainMode() : null; + } + public boolean offsetBasedDeduplicationSupported() { return activeReader != null && activeReader.getCurrentSource().offsetBasedDeduplicationSupported(); @@ -816,7 +820,10 @@ public TimerData getNextFiredTimer(Coder windowCode .transform( timer -> WindmillTimerInternals.windmillTimerToTimerData( - WindmillNamespacePrefix.SYSTEM_NAMESPACE_PREFIX, timer, windowCoder)) + WindmillNamespacePrefix.SYSTEM_NAMESPACE_PREFIX, + timer, + windowCoder, + getDrainMode())) .iterator(); } @@ -876,7 +883,10 @@ public TimerData getNextFiredUserTimer(Coder window .transform( timer -> WindmillTimerInternals.windmillTimerToTimerData( - WindmillNamespacePrefix.USER_NAMESPACE_PREFIX, timer, windowCoder)) + WindmillNamespacePrefix.USER_NAMESPACE_PREFIX, + timer, + windowCoder, + getDrainMode())) .iterator()); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java index a9a033c89ad7..ff931b1bf6be 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java @@ -24,6 +24,7 @@ import java.io.InputStream; import java.util.Collection; import java.util.Map; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.runners.dataflow.util.CloudObject; import org.apache.beam.runners.dataflow.worker.util.common.worker.NativeReader; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; @@ -117,8 +118,14 @@ protected WindowedValue decodeMessage(Windmill.Message message) throws IOExce Collection windows = WindmillSink.decodeMetadataWindows(windowsCoder, message.getMetadata()); PaneInfo paneInfo = WindmillSink.decodeMetadataPane(message.getMetadata()); + Boolean drainingValueFromUpstream = null; if (WindowedValues.WindowedValueCoder.isMetadataSupported()) { - WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata()); + BeamFnApi.Elements.ElementMetadata elementMetadata = + WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata()); + drainingValueFromUpstream = + elementMetadata.hasDrain() + ? (elementMetadata.getDrain() == BeamFnApi.Elements.DrainMode.Enum.DRAINING) + : null; } if (valueCoder instanceof KvCoder) { KvCoder kvCoder = (KvCoder) valueCoder; @@ -129,11 +136,19 @@ protected WindowedValue decodeMessage(Windmill.Message message) throws IOExce T result = (T) KV.of(decode(kvCoder.getKeyCoder(), key), decode(kvCoder.getValueCoder(), data)); // todo #33176 propagate metadata to windowed value - return WindowedValues.of(result, timestampMillis, windows, paneInfo); + return WindowedValues.of( + result, timestampMillis, windows, paneInfo, null, null, drainingValueFromUpstream); } else { notifyElementRead(data.available() + metadata.available()); // todo #33176 propagate metadata to windowed value - return WindowedValues.of(decode(valueCoder, data), timestampMillis, windows, paneInfo); + return WindowedValues.of( + decode(valueCoder, data), + timestampMillis, + windows, + paneInfo, + null, + null, + drainingValueFromUpstream); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java index 6690377d3de6..4f67dfcef319 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java @@ -24,6 +24,7 @@ import java.util.Collection; import java.util.List; import java.util.Objects; +import org.apache.beam.model.fnexecution.v1.BeamFnApi; import org.apache.beam.runners.core.KeyedWorkItem; import org.apache.beam.runners.core.KeyedWorkItemCoder; import org.apache.beam.runners.core.TimerInternals.TimerData; @@ -60,6 +61,7 @@ public class WindmillKeyedWorkItem implements KeyedWorkItem private final Windmill.WorkItem workItem; private final K key; + private final @javax.annotation.Nullable Boolean drainMode; private final transient Coder windowCoder; private final transient Coder> windowsCoder; @@ -70,12 +72,14 @@ public WindmillKeyedWorkItem( Windmill.WorkItem workItem, Coder windowCoder, Coder> windowsCoder, - Coder valueCoder) { + Coder valueCoder, + @javax.annotation.Nullable Boolean drainMode) { this.key = key; this.workItem = workItem; this.windowCoder = windowCoder; this.windowsCoder = windowsCoder; this.valueCoder = valueCoder; + this.drainMode = drainMode; } @Override @@ -93,7 +97,10 @@ public Iterable timersIterable() { .transform( timer -> WindmillTimerInternals.windmillTimerToTimerData( - WindmillNamespacePrefix.SYSTEM_NAMESPACE_PREFIX, timer, windowCoder)); + WindmillNamespacePrefix.SYSTEM_NAMESPACE_PREFIX, + timer, + windowCoder, + drainMode)); } @Override @@ -108,13 +115,22 @@ public Iterable> elementsIterable() { Collection windows = WindmillSink.decodeMetadataWindows(windowsCoder, message.getMetadata()); PaneInfo paneInfo = WindmillSink.decodeMetadataPane(message.getMetadata()); + // Draining value is based on upstream data + Boolean drainingValueFromUpstream = null; if (WindowedValues.WindowedValueCoder.isMetadataSupported()) { - WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata()); + BeamFnApi.Elements.ElementMetadata elementMetadata = + WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata()); + drainingValueFromUpstream = + elementMetadata.hasDrain() + ? (elementMetadata.getDrain() + == BeamFnApi.Elements.DrainMode.Enum.DRAINING) + : null; } InputStream inputStream = message.getData().newInput(); ElemT value = valueCoder.decode(inputStream, Coder.Context.OUTER); // todo #33176 specify additional metadata in the future - return WindowedValues.of(value, timestamp, windows, paneInfo); + return WindowedValues.of( + value, timestamp, windows, paneInfo, null, null, drainingValueFromUpstream); } catch (IOException e) { throw new RuntimeException(e); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternals.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternals.java index 1dbc7b005345..9732e8633017 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternals.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternals.java @@ -111,7 +111,8 @@ public void setTimer( Instant outputTimestamp, TimeDomain timeDomain) { TimerData timer = - TimerData.of(timerId, timerFamilyId, namespace, timestamp, outputTimestamp, timeDomain); + TimerData.of( + timerId, timerFamilyId, namespace, timestamp, outputTimestamp, timeDomain, false); setTimer(timer); } @@ -148,7 +149,8 @@ public void deleteTimer( namespace, BoundedWindow.TIMESTAMP_MIN_VALUE, BoundedWindow.TIMESTAMP_MAX_VALUE, - timeDomain)); + timeDomain, + false)); } @Override @@ -298,7 +300,10 @@ static Timer timerDataToWindmillTimer( } public static TimerData windmillTimerToTimerData( - WindmillNamespacePrefix prefix, Timer timer, Coder windowCoder) { + WindmillNamespacePrefix prefix, + Timer timer, + Coder windowCoder, + @javax.annotation.Nullable Boolean draining) { // The tag is a path-structure string but cheaper to parse than a proper URI. It follows // this pattern, where no component but the ID can contain a slash @@ -391,7 +396,8 @@ public static TimerData windmillTimerToTimerData( namespace, timestamp, outputTimestamp, - timerTypeToTimeDomain(timer.getType())); + timerTypeToTimeDomain(timer.getType()), + draining); } private static boolean useNewTimerTagEncoding(TimerData timerData) { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java index d91a5412b917..f4a6eec61cbf 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindowingWindmillReader.java @@ -119,7 +119,8 @@ public NativeReaderIterator>> iterator() throw final K key = keyCoder.decode(context.getSerializedKey().newInput(), Coder.Context.OUTER); final WorkItem workItem = context.getWorkItem(); KeyedWorkItem keyedWorkItem = - new WindmillKeyedWorkItem<>(key, workItem, windowCoder, windowsCoder, valueCoder); + new WindmillKeyedWorkItem<>( + key, workItem, windowCoder, windowsCoder, valueCoder, context.getDrainMode()); final boolean isEmptyWorkItem = (Iterables.isEmpty(keyedWorkItem.timersIterable()) && Iterables.isEmpty(keyedWorkItem.elementsIterable())); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java index 8b41a2d13219..8a1a864ed4ea 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java @@ -28,6 +28,7 @@ import java.util.Optional; import java.util.function.Consumer; import java.util.function.Supplier; +import javax.annotation.Nullable; import javax.annotation.concurrent.NotThreadSafe; import org.apache.beam.repackaged.core.org.apache.commons.lang3.tuple.Pair; import org.apache.beam.runners.dataflow.worker.ActiveMessageMetadata; @@ -78,12 +79,14 @@ public final class Work implements RefreshableWork { private volatile TimedState currentState; private volatile boolean isFailed; private volatile String processingThreadName = ""; + private final @Nullable Boolean drainMode; private Work( WorkItem workItem, long serializedWorkItemSize, Watermarks watermarks, ProcessingContext processingContext, + @Nullable Boolean drainMode, Supplier clock) { this.shardedKey = ShardedKey.create(workItem.getKey(), workItem.getShardingKey()); this.workItem = workItem; @@ -91,6 +94,7 @@ private Work( this.processingContext = processingContext; this.watermarks = watermarks; this.clock = clock; + this.drainMode = drainMode; this.startTime = clock.get(); Preconditions.checkState(EMPTY_ENUM_MAP.isEmpty()); // Create by passing EMPTY_ENUM_MAP to avoid recreating @@ -110,8 +114,10 @@ public static Work create( long serializedWorkItemSize, Watermarks watermarks, ProcessingContext processingContext, + @Nullable Boolean drainMode, Supplier clock) { - return new Work(workItem, serializedWorkItemSize, watermarks, processingContext, clock); + return new Work( + workItem, serializedWorkItemSize, watermarks, processingContext, drainMode, clock); } public static ProcessingContext createProcessingContext( @@ -207,6 +213,10 @@ public State getState() { return currentState.state(); } + public @Nullable Boolean getDrainMode() { + return drainMode; + } + public void setState(State state) { Instant now = clock.get(); totalDurationPerState.compute( diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/SingleSourceWorkerHarness.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/SingleSourceWorkerHarness.java index 95023d117299..ce6c6a050861 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/SingleSourceWorkerHarness.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/SingleSourceWorkerHarness.java @@ -150,6 +150,7 @@ private void streamingEngineDispatchLoop( (computationId, inputDataWatermark, synchronizedProcessingTime, + drainMode, workItem, serializedWorkItemSize, getWorkStreamLatencies) -> @@ -173,6 +174,7 @@ private void streamingEngineDispatchLoop( getDataClient, workCommitter::commit, heartbeatSender), + drainMode, getWorkStreamLatencies); })); try { @@ -232,6 +234,7 @@ private void applianceDispatchLoop(Supplier getWorkFn) watermarks.setOutputDataWatermark(workItem.getOutputDataWatermark()).build(), Work.createProcessingContext( computationId, getDataClient, workCommitter::commit, heartbeatSender), + computationWork.getDrainMode(), /* getWorkStreamLatencies= */ ImmutableList.of()); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BatchGroupAlsoByWindowViaOutputBufferFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BatchGroupAlsoByWindowViaOutputBufferFn.java index 12fc97ca9015..b0251c7bb4a8 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BatchGroupAlsoByWindowViaOutputBufferFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BatchGroupAlsoByWindowViaOutputBufferFn.java @@ -88,6 +88,7 @@ public void processElement( output, NullSideInputReader.empty(), reduceFn, + null, options); // Process the elements. diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GetWorkResponseChunkAssembler.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GetWorkResponseChunkAssembler.java index 0ebb4726d3a1..017148e40e9f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GetWorkResponseChunkAssembler.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GetWorkResponseChunkAssembler.java @@ -124,7 +124,8 @@ private static ComputationMetadata fromProto( metadataProto.getComputationId(), WindmillTimeUtils.windmillToHarnessWatermark(metadataProto.getInputDataWatermark()), WindmillTimeUtils.windmillToHarnessWatermark( - metadataProto.getDependentRealtimeInputWatermark())); + metadataProto.getDependentRealtimeInputWatermark()), + metadataProto.getDrainMode()); } abstract String computationId(); @@ -132,6 +133,8 @@ private static ComputationMetadata fromProto( abstract @Nullable Instant inputDataWatermark(); abstract @Nullable Instant synchronizedProcessingTime(); + + abstract @Nullable Boolean drainMode(); } @AutoValue diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java index 2712bf1bd33d..8eb4c51a2b49 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java @@ -281,6 +281,7 @@ private void consumeAssembledWorkItem(AssembledWorkItem assembledWorkItem) { assembledWorkItem.bufferedSize(), createWatermarks(workItem, metadata), createProcessingContext(metadata.computationId()), + metadata.drainMode(), assembledWorkItem.latencyAttributions()); budgetTracker.recordBudgetReceived(assembledWorkItem.bufferedSize()); GetWorkBudget extension = budgetTracker.computeBudgetExtension(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetWorkStream.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetWorkStream.java index ae7ce85e13a8..58407ad8147f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetWorkStream.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetWorkStream.java @@ -203,6 +203,7 @@ private void consumeAssembledWorkItem(AssembledWorkItem assembledWorkItem) { assembledWorkItem.computationMetadata().computationId(), assembledWorkItem.computationMetadata().inputDataWatermark(), assembledWorkItem.computationMetadata().synchronizedProcessingTime(), + assembledWorkItem.computationMetadata().drainMode(), assembledWorkItem.workItem(), assembledWorkItem.bufferedSize(), assembledWorkItem.latencyAttributions()); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemReceiver.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemReceiver.java index e2f69585e48f..160bdaa07f89 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemReceiver.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemReceiver.java @@ -30,6 +30,7 @@ void receiveWork( String computation, @Nullable Instant inputDataWatermark, @Nullable Instant synchronizedProcessingTime, + @Nullable Boolean drainMode, Windmill.WorkItem workItem, long serializedWorkItemSize, ImmutableList getWorkStreamLatencies); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemScheduler.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemScheduler.java index b9d31fbe501d..4195dbca4bb5 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemScheduler.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/WorkItemScheduler.java @@ -18,6 +18,7 @@ package org.apache.beam.runners.dataflow.worker.windmill.work; import javax.annotation.CheckReturnValue; +import javax.annotation.Nullable; import org.apache.beam.runners.dataflow.worker.streaming.Watermarks; import org.apache.beam.runners.dataflow.worker.streaming.Work; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.LatencyAttribution; @@ -35,6 +36,7 @@ public interface WorkItemScheduler { * @param workItem {@link WorkItem} to be processed. * @param watermarks processing watermarks for the workItem. * @param processingContext for processing the workItem. + * @param drainMode is job is draining. * @param getWorkStreamLatencies Latencies per processing stage for the WorkItem for reporting * back to Streaming Engine backend. */ @@ -43,5 +45,6 @@ void scheduleWork( long serializedWorkItemSize, Watermarks watermarks, Work.ProcessingContext processingContext, + @Nullable Boolean drainMode, ImmutableList getWorkStreamLatencies); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java index a4cd5d6d8a6b..0f0b16a09167 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java @@ -210,10 +210,12 @@ public void scheduleWork( long serializedWorkItemSize, Watermarks watermarks, Work.ProcessingContext processingContext, + @javax.annotation.Nullable Boolean drainMode, ImmutableList getWorkStreamLatencies) { computationState.activateWork( ExecutableWork.create( - Work.create(workItem, serializedWorkItemSize, watermarks, processingContext, clock), + Work.create( + workItem, serializedWorkItemSize, watermarks, processingContext, drainMode, clock), work -> processWork(computationState, work, getWorkStreamLatencies))); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java index dd13d5b55930..0daf96d9400b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java @@ -269,6 +269,7 @@ public boolean awaitTermination(int time, TimeUnit unit) throws InterruptedExcep computationWork.getComputationId(), inputDataWatermark, Instant.now(), + null, workItem, workItem.getSerializedSize(), ImmutableList.of( diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java index a60535dfbd69..c9aaef4d2e2c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java @@ -361,6 +361,7 @@ private static ExecutableWork createMockWork( Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build(), Work.createProcessingContext( computationId, new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), + null, Instant::now), processWorkFn); } @@ -3541,6 +3542,7 @@ public void testLatencyAttributionProtobufsPopulated() { new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), + null, clock); clock.sleep(Duration.millis(10)); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java index 1182a2c0b9e9..654d9965f1da 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java @@ -156,7 +156,8 @@ private void addTimer( timestamp, type == Windmill.Timer.Type.WATERMARK ? TimeDomain.EVENT_TIME - : TimeDomain.PROCESSING_TIME))) + : TimeDomain.PROCESSING_TIME, + false))) .setTimestamp(WindmillTimeUtils.harnessToWindmillTimestamp(timestamp)) .setType(type) .setStateFamily(STATE_FAMILY); @@ -194,7 +195,7 @@ private WindowedValue> createValue( return new ValueInEmptyWindows<>( (KeyedWorkItem) new WindmillKeyedWorkItem<>( - KEY, workItem.build(), windowCoder, wildcardWindowsCoder, valueCoder)); + KEY, workItem.build(), windowCoder, wildcardWindowsCoder, valueCoder, false)); } @Test diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowsReshuffleDoFnTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowsReshuffleDoFnTest.java index a348c0f00214..52c9844add86 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowsReshuffleDoFnTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowsReshuffleDoFnTest.java @@ -132,7 +132,7 @@ private WindowedValue> createValue( return new ValueInEmptyWindows<>( (KeyedWorkItem) new WindmillKeyedWorkItem<>( - KEY, workItem.build(), windowCoder, wildcardWindowsCoder, valueCoder)); + KEY, workItem.build(), windowCoder, wildcardWindowsCoder, valueCoder, false)); } @Test diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunnerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunnerTest.java index de4ec8e5aa9e..1100d99b2a3d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunnerTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunnerTest.java @@ -166,7 +166,8 @@ private TimerData timerData(IntervalWindow window, Instant timestamp, Timer.Type StateNamespaces.window(IntervalWindow.getCoder(), window), timestamp, timestamp, - type == Windmill.Timer.Type.WATERMARK ? TimeDomain.EVENT_TIME : TimeDomain.PROCESSING_TIME); + type == Windmill.Timer.Type.WATERMARK ? TimeDomain.EVENT_TIME : TimeDomain.PROCESSING_TIME, + null); } private IntervalWindow window(long start, long end) { diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java index e216f912d77f..377acabca172 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java @@ -143,6 +143,7 @@ private static Work createMockWork(Windmill.WorkItem workItem, Watermarks waterm watermarks, Work.createProcessingContext( COMPUTATION_ID, new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), + null, Instant::now); } @@ -172,7 +173,8 @@ public void testTimerInternalsSetTimer() { new StateNamespaceForTest("key"), new Instant(5000), new Instant(5000), - TimeDomain.EVENT_TIME)); + TimeDomain.EVENT_TIME, + null)); executionContext.flushState(); Windmill.Timer timer = outputBuilder.buildPartial().getOutputTimers(0); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcherTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcherTest.java index cf616ee6ac0d..def611381ab4 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcherTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcherTest.java @@ -210,7 +210,8 @@ private TimerData createTimer(long timestamp) { StateNamespaces.window(IntervalWindow.getCoder(), createWindow(timestamp)), new Instant(timestamp), new Instant(timestamp), - TimeDomain.EVENT_TIME); + TimeDomain.EVENT_TIME, + null); } private IntervalWindow createWindow(long timestamp) { diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactoryTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactoryTest.java index 771ffe192baf..d2ac5e0e6cbd 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactoryTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactoryTest.java @@ -462,7 +462,8 @@ public void testCleanupTimerForGlobalWindowWithAllowedLateness() throws Exceptio globalWindowNamespace, BoundedWindow.TIMESTAMP_MAX_VALUE, BoundedWindow.TIMESTAMP_MAX_VALUE.minus(Duration.millis(1)), - TimeDomain.EVENT_TIME)) + TimeDomain.EVENT_TIME, + null)) .thenReturn(null); // Set up non-empty state. We don't mock + verify calls to clear() but instead @@ -537,7 +538,8 @@ public void testCleanupWorks() throws Exception { firstWindowNamespace, firstWindow.maxTimestamp().plus(Duration.millis(1L)), firstWindow.maxTimestamp().plus(Duration.millis(1L)), - TimeDomain.EVENT_TIME)) + TimeDomain.EVENT_TIME, + null)) .thenReturn(null); // This should fire the timer to clean up the first window @@ -553,7 +555,8 @@ public void testCleanupWorks() throws Exception { secondWindowNamespace, secondWindow.maxTimestamp().plus(Duration.millis(1L)), secondWindow.maxTimestamp().plus(Duration.millis(1L)), - TimeDomain.EVENT_TIME)) + TimeDomain.EVENT_TIME, + null)) .thenReturn(null); // And this should clean up the second window diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java index 53a36722e41c..8b0a29ad60aa 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java @@ -93,7 +93,7 @@ public void testElementIteration() throws Exception { KeyedWorkItem keyedWorkItem = new WindmillKeyedWorkItem<>( - KEY, workItem.build(), WINDOW_CODER, WINDOWS_CODER, VALUE_CODER); + KEY, workItem.build(), WINDOW_CODER, WINDOWS_CODER, VALUE_CODER, false); assertThat( keyedWorkItem.elementsIterable(), @@ -148,7 +148,7 @@ public void testTimerOrdering() throws Exception { .build(); KeyedWorkItem keyedWorkItem = - new WindmillKeyedWorkItem<>(KEY, workItem, WINDOW_CODER, WINDOWS_CODER, VALUE_CODER); + new WindmillKeyedWorkItem<>(KEY, workItem, WINDOW_CODER, WINDOWS_CODER, VALUE_CODER, false); assertThat( keyedWorkItem.timersIterable(), @@ -169,7 +169,8 @@ private static Windmill.Timer makeSerializedTimer( ns, new Instant(timestamp), new Instant(timestamp), - WindmillTimerInternals.timerTypeToTimeDomain(type)))) + WindmillTimerInternals.timerTypeToTimeDomain(type), + false))) .setTimestamp(WindmillTimeUtils.harnessToWindmillTimestamp(new Instant(timestamp))) .setType(type) .setStateFamily(STATE_FAMILY) @@ -177,7 +178,7 @@ private static Windmill.Timer makeSerializedTimer( } private static TimerData makeTimer(StateNamespace ns, long timestamp, TimeDomain domain) { - return TimerData.of(ns, new Instant(timestamp), new Instant(timestamp), domain); + return TimerData.of(ns, new Instant(timestamp), new Instant(timestamp), domain, false); } @Test diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternalsTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternalsTest.java index ec8672b6a75f..eaef5956d15b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternalsTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternalsTest.java @@ -84,9 +84,13 @@ public void testTimerDataToFromTimer() { for (Instant timestamp : TEST_TIMESTAMPS) { List anonymousTimers = ImmutableList.of( - TimerData.of(namespace, timestamp, timestamp, timeDomain), + TimerData.of(namespace, timestamp, timestamp, timeDomain, null), TimerData.of( - namespace, timestamp, timestamp.minus(Duration.millis(1)), timeDomain)); + namespace, + timestamp, + timestamp.minus(Duration.millis(1)), + timeDomain, + null)); for (TimerData timer : anonymousTimers) { Instant expectedTimestamp = timer.getOutputTimestamp().isBefore(BoundedWindow.TIMESTAMP_MIN_VALUE) @@ -96,12 +100,17 @@ public void testTimerDataToFromTimer() { WindmillTimerInternals.windmillTimerToTimerData( prefix, WindmillTimerInternals.timerDataToWindmillTimer(stateFamily, prefix, timer), - coder); + coder, + null); // The function itself bounds output, so we dont expect the original input as the // output, we expect it to be bounded TimerData expected = TimerData.of( - timer.getNamespace(), timestamp, expectedTimestamp, timer.getDomain()); + timer.getNamespace(), + timestamp, + expectedTimestamp, + timer.getDomain(), + null); assertThat(computed, equalTo(expected)); } @@ -109,22 +118,24 @@ public void testTimerDataToFromTimer() { for (String timerId : TEST_TIMER_IDS) { List timers = ImmutableList.of( - TimerData.of(timerId, namespace, timestamp, timestamp, timeDomain), + TimerData.of(timerId, namespace, timestamp, timestamp, timeDomain, null), TimerData.of( - timerId, "family", namespace, timestamp, timestamp, timeDomain), + timerId, "family", namespace, timestamp, timestamp, timeDomain, null), TimerData.of( timerId, namespace, timestamp, timestamp.minus(Duration.millis(1)), - timeDomain), + timeDomain, + null), TimerData.of( timerId, "family", namespace, timestamp, timestamp.minus(Duration.millis(1)), - timeDomain)); + timeDomain, + null)); for (TimerData timer : timers) { Instant expectedTimestamp = @@ -139,13 +150,15 @@ public void testTimerDataToFromTimer() { timer.getNamespace(), timer.getTimestamp(), expectedTimestamp, - timer.getDomain()); + timer.getDomain(), + null); assertThat( WindmillTimerInternals.windmillTimerToTimerData( prefix, WindmillTimerInternals.timerDataToWindmillTimer( stateFamily, prefix, timer), - coder), + coder, + null), equalTo(expected)); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java index df3b959c82c5..a3b42cc89c51 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java @@ -204,6 +204,7 @@ private static Work createMockWork(Windmill.WorkItem workItem, Watermarks waterm watermarks, Work.createProcessingContext( COMPUTATION_ID, new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), + null, Instant::now); } @@ -1014,6 +1015,7 @@ public void testFailedWorkItemsAbort() throws Exception { new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), + null, Instant::now); context.start( "key", diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java index c0cb8241d73e..c03aca8bb539 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java @@ -71,6 +71,7 @@ private static ExecutableWork createWork(Windmill.WorkItem workItem) { workItem.getSerializedSize(), Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build(), createWorkProcessingContext(), + null, Instant::now), ignored -> {}); } @@ -82,6 +83,7 @@ private static ExecutableWork expiredWork(Windmill.WorkItem workItem) { workItem.getSerializedSize(), Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build(), createWorkProcessingContext(), + null, () -> Instant.EPOCH), ignored -> {}); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java index 935b25acb6f2..e71a10276051 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java @@ -75,6 +75,7 @@ private static ExecutableWork createWork(ShardedKey shardedKey, long workToken, new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), + null, Instant::now), ignored -> {}); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java index 65e40f171b0c..94c8f4b75957 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java @@ -128,6 +128,7 @@ private static WorkItemScheduler noOpProcessWorkItemFn() { serializedWorkItemSize, watermarks, processingContext, + drainMode, getWorkStreamLatencies) -> {}; } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSenderTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSenderTest.java index b94270ad7bb7..3217c736adb1 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSenderTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSenderTest.java @@ -68,6 +68,7 @@ public class WindmillStreamSenderTest { serializedWorkItemSize, watermarks, processingContext, + drainMode, getWorkStreamLatencies) -> {}; @Rule public transient Timeout globalTimeout = Timeout.seconds(600); private ManagedChannel inProcessChannel; diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutorTest.java index a86e6060955c..a3ea64105e0c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutorTest.java @@ -83,6 +83,7 @@ private static ExecutableWork createWork(Consumer executeWorkFn) { new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), + null, Instant::now), executeWorkFn); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java index 477c764a70ef..fa1d79d100e2 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java @@ -74,6 +74,7 @@ private static Work createMockWork(long workToken) { throw new UnsupportedOperationException(); }, mock(HeartbeatSender.class)), + null, Instant::now); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java index 5748b128f971..2d319879c79f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java @@ -104,6 +104,7 @@ private static Work createMockWork(long workToken) { throw new UnsupportedOperationException(); }, mock(HeartbeatSender.class)), + null, Instant::now); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStreamTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStreamTest.java index 419000178381..76883bebdac0 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStreamTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStreamTest.java @@ -70,6 +70,7 @@ public class GrpcDirectGetWorkStreamTest { serializedWorkItemSize, watermarks, processingContext, + drainMode, getWorkStreamLatencies) -> {}; private static final Windmill.JobHeader TEST_JOB_HEADER = Windmill.JobHeader.newBuilder() @@ -283,6 +284,7 @@ public void testConsumedWorkItem_computesAndSendsCorrectExtension() throws Inter serializedWorkItemSize, watermarks, processingContext, + drainMode, getWorkStreamLatencies) -> { scheduledWorkItems.add(work); }); @@ -327,8 +329,12 @@ public void testConsumedWorkItem_doesNotSendExtensionIfOutstandingBudgetHigh() createGetWorkStream( testStub, initialBudget, - (work, serializedWorkItemSize, watermarks, processingContext, getWorkStreamLatencies) -> - scheduledWorkItems.add(work)); + (work, + serializedWorkItemSize, + watermarks, + processingContext, + drainMode, + getWorkStreamLatencies) -> scheduledWorkItems.add(work)); Windmill.WorkItem workItem = Windmill.WorkItem.newBuilder() .setKey(ByteString.copyFromUtf8("somewhat_long_key")) @@ -365,6 +371,7 @@ public void testConsumedWorkItems() throws InterruptedException { serializedWorkItemSize, watermarks, processingContext, + drainMode, getWorkStreamLatencies) -> { scheduledWorkItems.add(work); }); @@ -408,6 +415,7 @@ public void testConsumedWorkItems_itemsSplitAcrossResponses() throws Interrupted serializedWorkItemSize, watermarks, processingContext, + drainMode, getWorkStreamLatencies) -> { scheduledWorkItems.add(work); }); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServerTest.java index e52b6e8de4bf..6caffb485f0b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServerTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServerTest.java @@ -336,6 +336,7 @@ public void onCompleted() { (String computation, @Nullable Instant inputDataWatermark, Instant synchronizedProcessingTime, + Boolean drainMode, WorkItem workItem, long serializedWorkItemSize, ImmutableList getWorkStreamLatencies) -> { @@ -469,6 +470,7 @@ public void onCompleted() { (String computation, @Nullable Instant inputDataWatermark, Instant synchronizedProcessingTime, + Boolean drainMode, WorkItem workItem, long serializedWorkItemSize, ImmutableList getWorkStreamLatencies) -> { diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java index f55549f7e2d9..27f20c4cc0f1 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java @@ -95,6 +95,7 @@ private static ExecutableWork createWork(Supplier clock, Consumer new FakeGetDataClient(), ignored -> {}, mock(HeartbeatSender.class)), + null, clock), processWorkFn); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java index 115deccf6df4..413dd033416f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresherTest.java @@ -133,6 +133,7 @@ private ExecutableWork createOldWork( Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build(), Work.createProcessingContext( "computationId", new FakeGetDataClient(), ignored -> {}, heartbeatSender), + null, A_LONG_TIME_AGO), processWork); } diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/TimerReceiverFactory.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/TimerReceiverFactory.java index d99a733a4eb4..3c01bf244975 100644 --- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/TimerReceiverFactory.java +++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/TimerReceiverFactory.java @@ -94,7 +94,8 @@ public FnDataReceiver> create(String transformId, String timerFamil namespace, timer.getClearBit() ? BoundedWindow.TIMESTAMP_MAX_VALUE : timer.getFireTimestamp(), timer.getClearBit() ? BoundedWindow.TIMESTAMP_MAX_VALUE : timer.getHoldTimestamp(), - timerSpec.getTimerSpec().getTimeDomain()); + timerSpec.getTimerSpec().getTimeDomain(), + timer.getDraining()); timerDataConsumer.accept(timer, timerData); } }; From aacca04ccae453d0186264f6b5933e0a1b41b8cf Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Thu, 9 Oct 2025 14:47:06 +0200 Subject: [PATCH 11/26] extend onTimer signature to add draining flag --- .../apache/beam/runners/core/DoFnRunner.java | 3 ++- .../core/LateDataDroppingDoFnRunner.java | 6 +++-- .../beam/runners/core/SimpleDoFnRunner.java | 27 +++++++++++++------ .../beam/runners/core/StatefulDoFnRunner.java | 6 +++-- ...SimplePushbackSideInputDoFnRunnerTest.java | 6 +++-- .../metrics/DoFnRunnerWithMetricsUpdate.java | 6 +++-- .../ExecutableStageDoFnOperator.java | 7 ++--- .../stableinput/BufferedElements.java | 2 +- .../stableinput/BufferingDoFnRunner.java | 3 ++- .../wrappers/streaming/DoFnOperatorTest.java | 3 ++- .../worker/DataflowProcessFnRunner.java | 4 ++- .../worker/GroupAlsoByWindowFnRunner.java | 4 ++- ...amingKeyedWorkItemSideInputDoFnRunner.java | 4 ++- .../worker/StreamingSideInputDoFnRunner.java | 4 ++- .../samza/metrics/DoFnRunnerWithMetrics.java | 6 +++-- .../samza/runtime/AsyncDoFnRunner.java | 5 ++-- .../runtime/DoFnRunnerWithKeyedInternals.java | 6 +++-- .../samza/runtime/SamzaDoFnRunners.java | 6 +++-- .../translation/batch/DoFnRunnerFactory.java | 4 ++- .../batch/DoFnRunnerWithMetrics.java | 6 +++-- .../translation/DoFnRunnerWithMetrics.java | 6 +++-- .../SparkInputDataProcessorTest.java | 4 ++- .../beam/sdk/util/construction/Timer.java | 11 +++++--- 23 files changed, 94 insertions(+), 45 deletions(-) diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java index dce4fcaa68bd..02b68827d2e0 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java @@ -46,7 +46,8 @@ public interface DoFnRunner void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { - doFnRunner.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + TimeDomain timeDomain, + @Nullable Boolean draining) { + doFnRunner.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } @Override diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java index 3af90ea9a0a1..a3c41eb7af86 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java @@ -201,11 +201,12 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { Preconditions.checkNotNull(outputTimestamp, "outputTimestamp"); OnTimerArgumentProvider argumentProvider = - new OnTimerArgumentProvider<>(timerId, key, window, timestamp, outputTimestamp, timeDomain); + new OnTimerArgumentProvider<>(timerId, key, window, timestamp, outputTimestamp, timeDomain, draining); invoker.invokeOnTimer(timerId, timerFamilyId, argumentProvider); } @@ -363,7 +364,7 @@ public void output( Collections.singletonList(window), PaneInfo.NO_FIRING, currentRecordId, - currentRecordOffset)); + currentRecordOffset, null)); } } @@ -429,6 +430,9 @@ public InputT element() { return elem.getValue(); } + @Override + public Boolean draining() { return elem.isDraining(); } + @Override public T sideInput(PCollectionView view) { checkNotNull(view, "View passed to sideInput cannot be null"); @@ -524,7 +528,7 @@ public void outputWindowedValue( SimpleDoFnRunner.this.outputWindowedValue( tag, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset, elem.isDraining())); } @Override @@ -725,6 +729,7 @@ private class OnTimerArgumentProvider extends DoFn.OnTime private final TimeDomain timeDomain; private final String timerId; private final KeyT key; + private final Boolean draining; private final OutputBuilderSupplier builderSupplier; /** Lazily initialized; should only be accessed via {@link #getNamespace()}. */ @@ -750,7 +755,8 @@ private OnTimerArgumentProvider( BoundedWindow window, Instant fireTimestamp, Instant timestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + Boolean draining) { fn.super(); this.timerId = timerId; this.window = window; @@ -758,13 +764,15 @@ private OnTimerArgumentProvider( this.timestamp = timestamp; this.timeDomain = timeDomain; this.key = key; + this.draining = draining; this.builderSupplier = OutputBuilderSuppliers.supplierForElement( WindowedValues.builder() .setValue(null) .setTimestamp(timestamp) .setWindow(window) - .setPaneInfo(PaneInfo.NO_FIRING)); + .setPaneInfo(PaneInfo.NO_FIRING) + .setDraining(draining)); } @Override @@ -772,6 +780,9 @@ public Instant timestamp() { return timestamp; } + @Override + public Boolean draining() { return draining; } + @Override public Instant fireTimestamp() { return fireTimestamp; @@ -1026,7 +1037,7 @@ public void outputWindowedValue( SimpleDoFnRunner.this.outputWindowedValue( tag, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset, null)); } @Override @@ -1304,7 +1315,7 @@ public void outputWindowedValue( SimpleDoFnRunner.this.outputWindowedValue( tag, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset, null)); } @Override diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/StatefulDoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/StatefulDoFnRunner.java index 52bbb3a306a0..e54f45be6bef 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/StatefulDoFnRunner.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/StatefulDoFnRunner.java @@ -42,6 +42,7 @@ import org.apache.beam.sdk.values.WindowedValues; import org.apache.beam.sdk.values.WindowingStrategy; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; import org.joda.time.Instant; @@ -208,7 +209,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { if (timerId.equals(SORT_FLUSH_TIMER)) { onSortFlushTimer(window, stepContext.timerInternals().currentInputWatermarkTime()); @@ -232,7 +234,7 @@ public void onTimer( stepContext.timerInternals().currentInputWatermarkTime()); } else { doFnRunner.onTimer( - timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } } } diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java index b7dda599b8d2..f9dc647e842f 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java @@ -18,6 +18,7 @@ package org.apache.beam.runners.core; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; +import org.checkerframework.checker.nullness.qual.Nullable; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.contains; import static org.hamcrest.Matchers.containsInAnyOrder; @@ -351,7 +352,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { firedTimers.add( TimerData.of( timerId, @@ -359,7 +361,7 @@ public void onTimer( StateNamespaces.window(IntervalWindow.getCoder(), (IntervalWindow) window), timestamp, outputTimestamp, - timeDomain)); + timeDomain, draining)); } @Override diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java index f1ec36564a4b..6d8ee22fe9af 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java @@ -26,6 +26,7 @@ import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** @@ -73,10 +74,11 @@ public void onTimer( final BoundedWindow window, final Instant timestamp, final Instant outputTimestamp, - final TimeDomain timeDomain) { + final TimeDomain timeDomain, + @Nullable Boolean draining) { try (Closeable ignored = MetricsEnvironment.scopedMetricsContainer(container.getMetricsContainer(stepName))) { - delegate.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + delegate.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } catch (IOException e) { throw new RuntimeException(e); } diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java index d5a7ff035efc..fb1b11baac23 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java @@ -592,7 +592,7 @@ public void setTimer( Instant outputTimestamp, TimeDomain timeDomain) { setTimer( - TimerData.of(timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain)); + TimerData.of(timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain, false)); } @Override @@ -1001,7 +1001,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { Object timerKey = keyForTimer.get(); Preconditions.checkNotNull(timerKey, "Key for timer needs to be set before calling onTimer"); Preconditions.checkNotNull(remoteBundle, "Call to onTimer outside of a bundle"); @@ -1034,7 +1035,7 @@ public void onTimer( timestamp, outputTimestamp, // TODO: Support propagating the PaneInfo through. - PaneInfo.NO_FIRING); + PaneInfo.NO_FIRING, draining); try { timerReceiver.accept(timerValue); } catch (Exception e) { diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java index 0c2ba87c4baa..5cce6f9fab3b 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferedElements.java @@ -104,7 +104,7 @@ static final class Timer implements BufferedElement { @Override public void processWith(DoFnRunner doFnRunner) { doFnRunner.onTimer( - timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, null); } @Override diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java index 10a3182f90f4..2025e6c4e402 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java @@ -234,7 +234,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { minBufferedElementTimestamp = Math.min(outputTimestamp.getMillis(), minBufferedElementTimestamp); diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperatorTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperatorTest.java index c2556d7229ba..86e99bbb8017 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperatorTest.java +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperatorTest.java @@ -507,7 +507,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { if ("cleanup".equals(timerId)) { holdState.clear(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowProcessFnRunner.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowProcessFnRunner.java index 741716b82848..59250e4f15d3 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowProcessFnRunner.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowProcessFnRunner.java @@ -34,6 +34,7 @@ import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowedValues; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** @@ -117,7 +118,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { throw new UnsupportedOperationException("Unsupported for ProcessFn"); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowFnRunner.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowFnRunner.java index 1909a73dc8b9..b30525ee068a 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowFnRunner.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowFnRunner.java @@ -26,6 +26,7 @@ import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.util.WindowedValueReceiver; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** @@ -81,7 +82,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { throw new UnsupportedOperationException( String.format("Timers are not supported by %s", GroupAlsoByWindowFn.class.getSimpleName())); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunner.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunner.java index e5d8a18be76f..b35f0e7da3cb 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunner.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunner.java @@ -38,6 +38,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** @@ -142,7 +143,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { throw new UnsupportedOperationException( "Attempt to deliver a timer to a DoFn, but timers are not supported in Dataflow."); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java index b5b723adb2b5..6a556721cdd8 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputDoFnRunner.java @@ -24,6 +24,7 @@ import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** @@ -84,7 +85,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { throw new UnsupportedOperationException( "Attempt to deliver a timer to a DoFn, but timers are not supported in Dataflow."); } diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/DoFnRunnerWithMetrics.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/DoFnRunnerWithMetrics.java index b815649a7659..f7b6d3e8e68e 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/DoFnRunnerWithMetrics.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/DoFnRunnerWithMetrics.java @@ -22,6 +22,7 @@ import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** @@ -61,11 +62,12 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { withMetrics( () -> underlying.onTimer( - timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain), + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining), false); } diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/AsyncDoFnRunner.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/AsyncDoFnRunner.java index d07a9bda78c6..daac80b323a3 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/AsyncDoFnRunner.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/AsyncDoFnRunner.java @@ -154,8 +154,9 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { - underlying.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + TimeDomain timeDomain, + @Nullable Boolean draining) { + underlying.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } @Override diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnRunnerWithKeyedInternals.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnRunnerWithKeyedInternals.java index 34e3405660c2..8daefb3e3da2 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnRunnerWithKeyedInternals.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnRunnerWithKeyedInternals.java @@ -24,6 +24,7 @@ import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** This class wraps a DoFnRunner with keyed StateInternals and TimerInternals access. */ @@ -66,14 +67,15 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { // Note: wrap with KV.of(key, null) as a special use case of setKeyedInternals() to set key // directly. setKeyedInternals(KV.of(key, null)); try { underlying.onTimer( - timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } finally { clearKeyedInternals(); } diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java index 7129ba9145eb..1db27342dc8f 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java @@ -69,6 +69,7 @@ import org.apache.beam.sdk.values.WindowingStrategy; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.samza.context.Context; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** A factory for Samza runner translator to create underlying DoFnRunner used in {@link DoFnOp}. */ @@ -430,7 +431,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { final KV timerReceiverKey = TimerReceiverFactory.decodeTimerDataTimerId(timerFamilyId); final FnDataReceiver timerReceiver = @@ -443,7 +445,7 @@ public void onTimer( timestamp, outputTimestamp, // TODO: Support propagating the PaneInfo through. - PaneInfo.NO_FIRING); + PaneInfo.NO_FIRING, draining); try { timerReceiver.accept(timerValue); } catch (Exception e) { diff --git a/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerFactory.java b/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerFactory.java index 350f7daa56cd..3bbc11f84f3f 100644 --- a/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerFactory.java +++ b/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerFactory.java @@ -47,6 +47,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; import scala.Serializable; @@ -266,7 +267,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) { + TimeDomain timeDomain, + @Nullable Boolean draining) { throw new UnsupportedOperationException(); } diff --git a/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerWithMetrics.java b/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerWithMetrics.java index db4ee5be5786..f6fe810a2ae6 100644 --- a/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerWithMetrics.java +++ b/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerWithMetrics.java @@ -29,6 +29,7 @@ import org.apache.beam.sdk.transforms.reflect.DoFnInvokers; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** DoFnRunner decorator which registers {@link MetricsContainer}. */ @@ -77,9 +78,10 @@ public void onTimer( final BoundedWindow window, final Instant timestamp, final Instant outputTimestamp, - final TimeDomain timeDomain) { + final TimeDomain timeDomain, + @Nullable Boolean draining) { try (Closeable ignored = MetricsEnvironment.scopedMetricsContainer(metrics)) { - delegate.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + delegate.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } catch (IOException e) { throw new RuntimeException(e); } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/DoFnRunnerWithMetrics.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/DoFnRunnerWithMetrics.java index a6b1f65571dc..b80d451d64b0 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/DoFnRunnerWithMetrics.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/DoFnRunnerWithMetrics.java @@ -28,6 +28,7 @@ import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; /** DoFnRunner decorator which registers {@link MetricsContainerImpl}. */ @@ -76,9 +77,10 @@ public void onTimer( final BoundedWindow window, final Instant timestamp, final Instant outputTimestamp, - final TimeDomain timeDomain) { + final TimeDomain timeDomain, + @Nullable Boolean draining) { try (Closeable ignored = MetricsEnvironment.scopedMetricsContainer(metricsContainer())) { - delegate.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + delegate.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } catch (IOException e) { throw new RuntimeException(e); } diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkInputDataProcessorTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkInputDataProcessorTest.java index 2ff06b59f8eb..d3e929bee99c 100644 --- a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkInputDataProcessorTest.java +++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkInputDataProcessorTest.java @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.spark.translation; +import org.checkerframework.checker.nullness.qual.Nullable; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; @@ -252,7 +253,8 @@ public void onTimer( BoundedWindow window, Instant timestamp, Instant outputTimestamp, - TimeDomain timeDomain) {} + TimeDomain timeDomain, + @Nullable Boolean draining) {} @Override public void finishBundle() {} diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Timer.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Timer.java index d443f008b7d2..c00039750f4f 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Timer.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Timer.java @@ -65,9 +65,10 @@ public static Timer of( Collection windows, Instant fireTimestamp, Instant holdTimestamp, - PaneInfo paneInfo) { + PaneInfo paneInfo, + Boolean draining) { return new AutoValue_Timer( - userKey, dynamicTimerTag, windows, false, fireTimestamp, holdTimestamp, paneInfo); + userKey, dynamicTimerTag, windows, false, fireTimestamp, holdTimestamp, paneInfo, draining); } /** @@ -76,7 +77,7 @@ public static Timer of( */ public static Timer cleared( K userKey, String dynamicTimerTag, Collection windows) { - return new AutoValue_Timer(userKey, dynamicTimerTag, windows, true, null, null, null); + return new AutoValue_Timer(userKey, dynamicTimerTag, windows, true, null, null, null, false); } /** Returns the key that the timer is set on. */ @@ -116,6 +117,8 @@ public static Timer cleared( */ public abstract @Nullable PaneInfo getPaneInfo(); + public abstract @Nullable Boolean getDraining(); + @Override public final boolean equals(@Nullable Object other) { if (!(other instanceof Timer)) { @@ -201,7 +204,7 @@ public Timer decode(InputStream inStream) throws IOException { Instant fireTimestamp = InstantCoder.of().decode(inStream); Instant holdTimestamp = InstantCoder.of().decode(inStream); PaneInfo paneInfo = PaneInfoCoder.INSTANCE.decode(inStream); - return Timer.of(userKey, dynamicTimerTag, windows, fireTimestamp, holdTimestamp, paneInfo); + return Timer.of(userKey, dynamicTimerTag, windows, fireTimestamp, holdTimestamp, paneInfo, false); } @Override From 17ff5c43edcb47a84fbe6a5656801bbaa91d27d3 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 10 Oct 2025 11:19:46 +0200 Subject: [PATCH 12/26] expose draining to doFn --- .../GroupAlsoByWindowViaWindowSetNewDoFn.java | 1 + ...oundedSplittableProcessElementInvoker.java | 31 +++++++- .../apache/beam/runners/core/ReduceFn.java | 2 + .../runners/core/ReduceFnContextFactory.java | 46 ++++++++++-- .../beam/runners/core/ReduceFnRunner.java | 18 ++++- .../beam/runners/core/SimpleDoFnRunner.java | 24 ++++-- .../beam/runners/core/ReduceFnTester.java | 11 ++- .../GroupAlsoByWindowEvaluatorFactory.java | 1 + .../runners/dataflow/BatchViewOverrides.java | 5 ++ .../worker/util/ValueInEmptyWindows.java | 5 ++ .../org/apache/beam/sdk/transforms/DoFn.java | 15 ++++ .../beam/sdk/transforms/DoFnTester.java | 16 ++-- .../SplittableParDoNaiveBounded.java | 5 ++ .../beam/sdk/values/ValueInSingleWindow.java | 29 ++++++-- .../beam/fn/harness/FnApiDoFnRunner.java | 73 ++++++++++++++++--- 15 files changed, 243 insertions(+), 39 deletions(-) diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/GroupAlsoByWindowViaWindowSetNewDoFn.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/GroupAlsoByWindowViaWindowSetNewDoFn.java index f242c7d10003..38666b0ae901 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/GroupAlsoByWindowViaWindowSetNewDoFn.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/GroupAlsoByWindowViaWindowSetNewDoFn.java @@ -107,6 +107,7 @@ public void processElement(ProcessContext c) throws Exception { windowedValue -> outputManager.output(mainTag, windowedValue), sideInputReader, reduceFn, + c.draining(), c.getPipelineOptions()); reduceFnRunner.processElements(keyedWorkItem.elementsIterable()); diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java index 9bda4dd2cbca..5203c10c9deb 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java @@ -386,6 +386,11 @@ public PaneInfo pane() { return element.getPaneInfo(); } + @Override + public Boolean draining() { + return element.isDraining(); + } + @Override public String currentRecordId() { return element.getRecordId(); @@ -439,7 +444,13 @@ public void outputWindowedValue( outputReceiver.output( mainOutputTag, WindowedValues.of( - value, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + value, + timestamp, + windows, + paneInfo, + currentRecordId, + currentRecordOffset, + element.isDraining())); } @Override @@ -450,7 +461,15 @@ public void output(TupleTag tag, T value) { @Override public void outputWithTimestamp(TupleTag tag, T value, Instant timestamp) { outputReceiver.output( - tag, WindowedValues.of(value, timestamp, element.getWindows(), element.getPaneInfo())); + tag, + WindowedValues.of( + value, + timestamp, + element.getWindows(), + element.getPaneInfo(), + element.getRecordId(), + element.getRecordOffset(), + element.isDraining())); } @Override @@ -479,7 +498,13 @@ public void outputWindowedValue( outputReceiver.output( tag, WindowedValues.of( - value, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + value, + timestamp, + windows, + paneInfo, + currentRecordId, + currentRecordOffset, + element.isDraining())); } private void noteOutput() { diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFn.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFn.java index 0be0c6353783..1de6b3167350 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFn.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFn.java @@ -44,6 +44,8 @@ public abstract class Context { /** The window that is being processed. */ public abstract W window(); + public abstract Boolean draining(); + /** Access the current {@link WindowingStrategy}. */ public abstract WindowingStrategy windowingStrategy(); diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnContextFactory.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnContextFactory.java index 1ca927f16caa..48553f461de3 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnContextFactory.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnContextFactory.java @@ -49,6 +49,7 @@ public interface OnTriggerCallbacks { } private final K key; + private final @Nullable Boolean draining; private final ReduceFn reduceFn; private final WindowingStrategy windowingStrategy; private final StateInternals stateInternals; @@ -65,7 +66,9 @@ public interface OnTriggerCallbacks { ActiveWindowSet activeWindows, TimerInternals timerInternals, @Nullable SideInputReader sideInputReader, + @Nullable Boolean draining, @Nullable PipelineOptions options) { + this.draining = draining; this.key = key; this.reduceFn = reduceFn; this.windowingStrategy = windowingStrategy; @@ -98,8 +101,8 @@ public ReduceFn.Context base(W window, StateStyle style) } public ReduceFn.ProcessValueContext forValue( - W window, InputT value, Instant timestamp, StateStyle style) { - return new ProcessValueContextImpl(stateAccessor(window, style), value, timestamp); + W window, InputT value, Instant timestamp, StateStyle style, Boolean draining) { + return new ProcessValueContextImpl(stateAccessor(window, style), value, timestamp, draining); } public ReduceFn.OnTriggerContext forTrigger( @@ -135,17 +138,21 @@ public TimersImpl(StateNamespace namespace) { @Override public void setTimer(Instant timestamp, TimeDomain timeDomain) { - timerInternals.setTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain)); + // todo radoslaws@ should we pass draining bit to timerdata + timerInternals.setTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain, false)); } @Override public void setTimer(Instant timestamp, Instant outputTimestamp, TimeDomain timeDomain) { - timerInternals.setTimer(TimerData.of(namespace, timestamp, outputTimestamp, timeDomain)); + // todo radoslaws@ should we pass draining bit to timerdata + timerInternals.setTimer( + TimerData.of(namespace, timestamp, outputTimestamp, timeDomain, false)); } @Override public void deleteTimer(Instant timestamp, TimeDomain timeDomain) { - timerInternals.deleteTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain)); + // todo radoslaws@ should we pass draining bit to timerdata + timerInternals.deleteTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain, false)); } @Override @@ -327,6 +334,11 @@ public K key() { return key; } + @Override + public Boolean draining() { + return draining; + } + @Override public W window() { return state.window(); @@ -354,14 +366,16 @@ private class ProcessValueContextImpl private final Instant timestamp; private final StateAccessorImpl state; private final TimersImpl timers; + private final Boolean draining; private ProcessValueContextImpl( - StateAccessorImpl state, InputT value, Instant timestamp) { + StateAccessorImpl state, InputT value, Instant timestamp, Boolean draining) { reduceFn.super(); this.state = state; this.value = value; this.timestamp = timestamp; this.timers = new TimersImpl(state.namespace()); + this.draining = draining; } @Override @@ -369,6 +383,11 @@ public K key() { return key; } + @Override + public Boolean draining() { + return draining; + } + @Override public W window() { return state.window(); @@ -420,6 +439,11 @@ public K key() { return key; } + @Override + public Boolean draining() { + return draining; + } + @Override public W window() { return state.window(); @@ -466,6 +490,11 @@ public K key() { return key; } + @Override + public Boolean draining() { + return draining; + } + @Override public WindowingStrategy windowingStrategy() { return windowingStrategy; @@ -502,6 +531,11 @@ public K key() { return key; } + @Override + public Boolean draining() { + return draining; + } + @Override public WindowingStrategy windowingStrategy() { return windowingStrategy; diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java index b08bd42b0b22..1e479515620e 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java @@ -198,6 +198,8 @@ public class ReduceFnRunner { */ private final PaneInfoTracker paneInfoTracker; + private final @Nullable Boolean draining; + /** * Store whether we've seen any elements for a window since the last pane was emitted. * @@ -219,7 +221,9 @@ public ReduceFnRunner( WindowedValueReceiver> outputter, @Nullable SideInputReader sideInputReader, ReduceFn reduceFn, + @Nullable Boolean draining, @Nullable PipelineOptions options) { + this.draining = draining; this.key = key; this.timerInternals = timerInternals; this.paneInfoTracker = new PaneInfoTracker(timerInternals); @@ -248,6 +252,7 @@ public ReduceFnRunner( this.activeWindows, timerInternals, sideInputReader, + draining, options); this.watermarkHold = new WatermarkHold<>(timerInternals, windowingStrategy); @@ -583,7 +588,11 @@ private void processElement(Map windowToMergeResult, WindowedValue for (W window : windows) { ReduceFn.ProcessValueContext directContext = contextFactory.forValue( - window, value.getValue(), value.getTimestamp(), StateStyle.DIRECT); + window, + value.getValue(), + value.getTimestamp(), + StateStyle.DIRECT, + value.isDraining()); if (triggerRunner.isClosed(directContext.state())) { // This window has already been closed. droppedDueToClosedWindow.inc(); @@ -601,7 +610,11 @@ private void processElement(Map windowToMergeResult, WindowedValue activeWindows.ensureWindowIsActive(window); ReduceFn.ProcessValueContext renamedContext = contextFactory.forValue( - window, value.getValue(), value.getTimestamp(), StateStyle.RENAMED); + window, + value.getValue(), + value.getTimestamp(), + StateStyle.RENAMED, + value.isDraining()); nonEmptyPanes.recordContent(renamedContext.state()); scheduleGarbageCollectionTimer(directContext); @@ -722,6 +735,7 @@ public void onTimers(Iterable timers) throws Exception { continue; } + // pass draining bit ReduceFn.Context directContext = contextFactory.base(window, StateStyle.DIRECT); ReduceFn.Context renamedContext = diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java index a3c41eb7af86..f54687bfa629 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java @@ -206,7 +206,8 @@ public void onTimer( Preconditions.checkNotNull(outputTimestamp, "outputTimestamp"); OnTimerArgumentProvider argumentProvider = - new OnTimerArgumentProvider<>(timerId, key, window, timestamp, outputTimestamp, timeDomain, draining); + new OnTimerArgumentProvider<>( + timerId, key, window, timestamp, outputTimestamp, timeDomain, draining); invoker.invokeOnTimer(timerId, timerFamilyId, argumentProvider); } @@ -364,7 +365,8 @@ public void output( Collections.singletonList(window), PaneInfo.NO_FIRING, currentRecordId, - currentRecordOffset, null)); + currentRecordOffset, + null)); } } @@ -431,7 +433,9 @@ public InputT element() { } @Override - public Boolean draining() { return elem.isDraining(); } + public Boolean draining() { + return elem.isDraining(); + } @Override public T sideInput(PCollectionView view) { @@ -528,7 +532,13 @@ public void outputWindowedValue( SimpleDoFnRunner.this.outputWindowedValue( tag, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset, elem.isDraining())); + output, + timestamp, + windows, + paneInfo, + currentRecordId, + currentRecordOffset, + elem.isDraining())); } @Override @@ -772,7 +782,7 @@ private OnTimerArgumentProvider( .setTimestamp(timestamp) .setWindow(window) .setPaneInfo(PaneInfo.NO_FIRING) - .setDraining(draining)); + .setDraining(draining)); } @Override @@ -781,7 +791,9 @@ public Instant timestamp() { } @Override - public Boolean draining() { return draining; } + public Boolean draining() { + return draining; + } @Override public Instant fireTimestamp() { diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java index 55791d582717..d9582873477f 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java @@ -89,6 +89,7 @@ }) public class ReduceFnTester { private static final String KEY = "TEST_KEY"; + private static final Boolean DRAINING = null; private final TestInMemoryStateInternals stateInternals = new TestInMemoryStateInternals<>(KEY); @@ -302,6 +303,7 @@ ReduceFnRunner createRunner() { testOutputter, sideInputReader, reduceFn, + DRAINING, options); } @@ -574,7 +576,11 @@ public void fireTimer(W window, Instant timestamp, TimeDomain domain) throws Exc ArrayList timers = new ArrayList<>(1); timers.add( TimerData.of( - StateNamespaces.window(windowFn.windowCoder(), window), timestamp, timestamp, domain)); + StateNamespaces.window(windowFn.windowCoder(), window), + timestamp, + timestamp, + domain, + null)); runner.onTimers(timers); runner.persist(); } @@ -588,7 +594,8 @@ public void fireTimers(W window, TimestampedValue... timers) throws StateNamespaces.window(windowFn.windowCoder(), window), timer.getTimestamp(), timer.getTimestamp(), - timer.getValue())); + timer.getValue(), + null)); } runner.onTimers(timerData); runner.persist(); diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java index c6726fb3463f..77387bf497b3 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java @@ -179,6 +179,7 @@ public void processElement(WindowedValue> element) throws Ex new BundleWindowedValueReceiver<>(bundle), new UnsupportedSideInputReader(DirectGroupAlsoByWindow.class.getSimpleName()), reduceFn, + element.isDraining(), options); // Drop any elements within expired windows diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchViewOverrides.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchViewOverrides.java index 10b41bb5b5ba..7026396a42b9 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchViewOverrides.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchViewOverrides.java @@ -1378,6 +1378,11 @@ public T getValue() { return value; } + @Override + public @Nullable Boolean isDraining() { + return null; + } + @Override public Instant getTimestamp() { return BoundedWindow.TIMESTAMP_MIN_VALUE; diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ValueInEmptyWindows.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ValueInEmptyWindows.java index a51c9ed419e1..90b31b974e2f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ValueInEmptyWindows.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ValueInEmptyWindows.java @@ -59,6 +59,11 @@ public PaneInfo getPaneInfo() { return null; } + @Override + public @Nullable Boolean isDraining() { + return null; + } + @Override public Iterable> explodeWindows() { return Collections.emptyList(); diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java index d0714de60328..59bcde70edf1 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java @@ -361,6 +361,9 @@ public abstract class ProcessContext extends WindowedContext { @Pure public abstract Long currentRecordOffset(); + + @Pure + public abstract Boolean draining(); } /** Information accessible when running a {@link DoFn.OnTimer} method. */ @@ -377,6 +380,9 @@ public abstract class OnTimerContext extends WindowedContext { /** Returns the time domain of the current timer. */ public abstract TimeDomain timeDomain(); + + @Pure + public abstract Boolean draining(); } public abstract class OnWindowExpirationContext extends WindowedContext { @@ -858,6 +864,15 @@ public interface MultiOutputReceiver { @Target(ElementType.PARAMETER) public @interface Element {} + /** + * Parameter annotation for the drain information for {@link ProcessElement} and {@link OnTimer} + * methods. + */ + @Documented + @Retention(RetentionPolicy.RUNTIME) + @Target(ElementType.PARAMETER) + public @interface Draining {} + /** * Parameter annotation for the restriction for {@link GetSize}, {@link SplitRestriction}, {@link * GetInitialWatermarkEstimatorState}, {@link NewWatermarkEstimator}, and {@link NewTracker} diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java index c59d6b528c3f..2269f1fbd799 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java @@ -490,7 +490,7 @@ public void output(TupleTag tag, T output, Instant timestamp, BoundedWind getMutableOutput(tag) .add( ValueInSingleWindow.of( - output, timestamp, window, PaneInfo.NO_FIRING, null, null)); + output, timestamp, window, PaneInfo.NO_FIRING, null, null, null)); } @Override @@ -519,7 +519,8 @@ public void output( window, PaneInfo.NO_FIRING, currentRecordId, - currentRecordOffset)); + currentRecordOffset, + null)); } }; } @@ -608,6 +609,11 @@ public PaneInfo pane() { return element.getPaneInfo(); } + @Override + public Boolean draining() { + return element.isDraining(); + } + @Override public String currentRecordId() { return element.getCurrentRecordId(); @@ -670,7 +676,7 @@ public void outputWithTimestamp(TupleTag tag, T output, Instant timestamp getMutableOutput(tag) .add( ValueInSingleWindow.of( - output, timestamp, element.getWindow(), element.getPaneInfo(), null, null)); + output, timestamp, element.getWindow(), element.getPaneInfo(), null, null, null)); } @Override @@ -682,7 +688,7 @@ public void outputWindowedValue( PaneInfo paneInfo) { for (BoundedWindow w : windows) { getMutableOutput(tag) - .add(ValueInSingleWindow.of(output, timestamp, w, paneInfo, null, null)); + .add(ValueInSingleWindow.of(output, timestamp, w, paneInfo, null, null, null)); } } @@ -699,7 +705,7 @@ public void outputWindowedValue( getMutableOutput(tag) .add( ValueInSingleWindow.of( - output, timestamp, w, paneInfo, currentRecordId, currentRecordOffset)); + output, timestamp, w, paneInfo, currentRecordId, currentRecordOffset, null)); } } } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/SplittableParDoNaiveBounded.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/SplittableParDoNaiveBounded.java index e6394b8810a4..b4ce5594f382 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/SplittableParDoNaiveBounded.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/SplittableParDoNaiveBounded.java @@ -536,6 +536,11 @@ public Object key() { throw new UnsupportedOperationException(); } + @Override + public Boolean draining() { + return outerContext.draining(); + } + @Override public Object sideInput(String tagId) { PCollectionView view = sideInputMapping.get(tagId); diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java index 2bbc28cac9fe..74108328baf0 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java @@ -66,6 +66,8 @@ public T getValue() { public abstract @Nullable Long getCurrentRecordOffset(); + public abstract @Nullable Boolean isDraining(); + // todo #33176 specify additional metadata in the future public static ValueInSingleWindow of( T value, @@ -73,14 +75,15 @@ public static ValueInSingleWindow of( BoundedWindow window, PaneInfo paneInfo, @Nullable String currentRecordId, - @Nullable Long currentRecordOffset) { + @Nullable Long currentRecordOffset, + @Nullable Boolean draining) { return new AutoValue_ValueInSingleWindow<>( - value, timestamp, window, paneInfo, currentRecordId, currentRecordOffset); + value, timestamp, window, paneInfo, currentRecordId, currentRecordOffset, draining); } public static ValueInSingleWindow of( T value, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) { - return of(value, timestamp, window, paneInfo, null, null); + return of(value, timestamp, window, paneInfo, null, null, null); } /** A coder for {@link ValueInSingleWindow}. */ @@ -120,6 +123,15 @@ public void encode(ValueInSingleWindow windowedElem, OutputStream outStream, BeamFnApi.Elements.ElementMetadata.Builder builder = BeamFnApi.Elements.ElementMetadata.newBuilder(); // todo #33176 specify additional metadata in the future + Boolean draining = windowedElem.isDraining(); + if (draining != null) { + builder.setDrain( + draining + ? BeamFnApi.Elements.DrainMode.Enum.DRAINING + : BeamFnApi.Elements.DrainMode.Enum.NOT_DRAINING); + } else { + builder.setDrain(BeamFnApi.Elements.DrainMode.Enum.UNSPECIFIED); + } BeamFnApi.Elements.ElementMetadata metadata = builder.build(); ByteArrayCoder.of().encode(metadata.toByteArray(), outStream); } @@ -137,13 +149,20 @@ public ValueInSingleWindow decode(InputStream inStream, Context context) thro Instant timestamp = InstantCoder.of().decode(inStream); BoundedWindow window = windowCoder.decode(inStream); PaneInfo paneInfo = PaneInfo.PaneInfoCoder.INSTANCE.decode(inStream); + Boolean draining = null; if (WindowedValues.WindowedValueCoder.isMetadataSupported() && paneInfo.isElementMetadata()) { - BeamFnApi.Elements.ElementMetadata.parseFrom(ByteArrayCoder.of().decode(inStream)); + BeamFnApi.Elements.ElementMetadata elementMetadata = + BeamFnApi.Elements.ElementMetadata.parseFrom(ByteArrayCoder.of().decode(inStream)); + if (elementMetadata.hasDrain() + && elementMetadata.getDrain() != BeamFnApi.Elements.DrainMode.Enum.UNSPECIFIED) { + draining = elementMetadata.getDrain() == BeamFnApi.Elements.DrainMode.Enum.DRAINING; + } } T value = valueCoder.decode(inStream, context); // todo #33176 specify additional metadata in the future - return new AutoValue_ValueInSingleWindow<>(value, timestamp, window, paneInfo, null, null); + return new AutoValue_ValueInSingleWindow<>( + value, timestamp, window, paneInfo, null, null, draining); } @Override 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 0388d3c03f00..4d32bcf53e27 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 @@ -247,6 +247,7 @@ public final void addRunnerForPTransform(Context context) throws IOException { private WindowedValue currentElement; private Object currentKey; + private Boolean draining; /** * Only valid during {@link @@ -1200,6 +1201,8 @@ private void processTimer( checkNotNull(timerBundleTracker); try { currentKey = timer.getUserKey(); + draining = timer.getDraining(); + // add drain Iterator windowIterator = (Iterator) timer.getWindows().iterator(); while (windowIterator.hasNext()) { @@ -1531,7 +1534,8 @@ private Timer getTimerForTime(Instant scheduledTime) { Collections.singletonList(boundedWindow), scheduledTime, outputTimestamp, - paneInfo); + paneInfo, + draining); } } @@ -1684,7 +1688,8 @@ public void output( Collections.singletonList(window), PaneInfo.NO_FIRING, currentRecordId, - currentRecordOffset)); + currentRecordOffset, + currentElement != null ? currentElement.isDraining() : null)); } @Override @@ -1708,7 +1713,8 @@ public void output( Collections.singletonList(window), PaneInfo.NO_FIRING, currentRecordId, - currentRecordOffset)); + currentRecordOffset, + currentElement != null ? currentElement.isDraining() : null)); } } @@ -1821,7 +1827,7 @@ public void outputWindowedValue( outputTo( mainOutputConsumer, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset, null)); } @Override @@ -1872,7 +1878,7 @@ public void outputWindowedValue( outputTo( consumer, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset, null)); } @Override @@ -1990,7 +1996,13 @@ public void outputWindowedValue( outputTo( mainOutputConsumer, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, + timestamp, + windows, + paneInfo, + currentRecordId, + currentRecordOffset, + currentElement != null ? currentElement.isDraining() : null)); } @Override @@ -2041,7 +2053,13 @@ public void outputWindowedValue( outputTo( consumer, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, + timestamp, + windows, + paneInfo, + currentRecordId, + currentRecordOffset, + currentElement != null ? currentElement.isDraining() : null)); } } @@ -2333,6 +2351,11 @@ public PaneInfo pane() { return currentElement.getPaneInfo(); } + @Override + public Boolean draining() { + return draining; + } + @Override public Object watermarkEstimatorState() { return currentWatermarkEstimatorState; @@ -2392,6 +2415,7 @@ public OutputBuilder builder(OutputT value) { .setWindow(currentWindow) .setTimestamp(currentTimer.getHoldTimestamp()) .setPaneInfo(currentTimer.getPaneInfo()) + .setDraining(draining) .setReceiver( windowedValue -> { checkOnWindowExpirationTimestamp(windowedValue.getTimestamp()); @@ -2411,7 +2435,13 @@ public void outputWindowedValue( outputTo( mainOutputConsumer, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, + timestamp, + windows, + paneInfo, + currentRecordId, + currentRecordOffset, + currentElement != null ? currentElement.isDraining() : null)); } @Override @@ -2468,7 +2498,13 @@ public void outputWindowedValue( outputTo( consumer, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, + timestamp, + windows, + paneInfo, + currentRecordId, + currentRecordOffset, + currentElement != null ? currentElement.isDraining() : null)); } @SuppressWarnings( @@ -2539,6 +2575,7 @@ public OutputBuilder builder(Row value) { .setValue(value) .setTimestamp(currentTimer.getHoldTimestamp()) .setWindow(currentWindow) + .setDraining(draining) .setReceiver( windowedValue -> context.outputWindowedValue( @@ -2576,6 +2613,7 @@ public OutputBuilder builder(T value) { .setValue(value) .setTimestamp(currentTimer.getHoldTimestamp()) .setWindow(currentWindow) + .setDraining(draining) .setReceiver( windowedValue -> context.outputWindowedValue( @@ -2613,6 +2651,7 @@ public OutputBuilder builder(Row value) { .setValue(value) .setTimestamp(currentTimer.getHoldTimestamp()) .setWindow(currentWindow) + .setDraining(draining) .setReceiver( windowedValue -> context.outputWindowedValue( @@ -2690,6 +2729,11 @@ public BoundedWindow window() { return currentWindow; } + @Override + public Boolean draining() { + return draining; + } + @Override public OutputBuilder builder(OutputT value) { return WindowedValues.builder() @@ -2697,6 +2741,7 @@ public OutputBuilder builder(OutputT value) { .setTimestamp(currentTimer.getHoldTimestamp()) .setWindow(currentWindow) .setPaneInfo(currentTimer.getPaneInfo()) + .setDraining(draining) .setReceiver( windowedValue -> { checkTimerTimestamp(windowedValue.getTimestamp()); @@ -2735,7 +2780,13 @@ public void outputWindowedValue( outputTo( mainOutputConsumer, WindowedValues.of( - output, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset)); + output, + timestamp, + windows, + paneInfo, + currentRecordId, + currentRecordOffset, + currentElement != null ? currentElement.isDraining() : null)); } @Override @@ -2903,6 +2954,7 @@ public OutputBuilder builder(T value) { .setValue(value) .setTimestamp(currentTimer.getHoldTimestamp()) .setWindow(currentWindow) + .setDraining(draining) .setPaneInfo(currentTimer.getPaneInfo()) .setReceiver( windowedValue -> @@ -2941,6 +2993,7 @@ public OutputBuilder builder(Row value) { .setTimestamp(currentTimer.getHoldTimestamp()) .setWindow(currentWindow) .setPaneInfo(currentTimer.getPaneInfo()) + .setDraining(draining) .setReceiver( windowedValue -> context.outputWindowedValue( From 8c40543440653b5cf13df8f0c26cf981be1be653 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 10 Oct 2025 11:35:08 +0200 Subject: [PATCH 13/26] timer plumbing --- .../runners/core/InMemoryTimerInternals.java | 3 +- .../core/LateDataDroppingDoFnRunner.java | 3 +- .../SimplePushbackSideInputDoFnRunner.java | 3 +- .../SplittableParDoViaKeyedWorkItems.java | 3 +- .../core/InMemoryTimerInternalsTest.java | 28 ++++++++++--------- .../runners/core/KeyedWorkItemCoderTest.java | 3 +- .../runners/core/SimpleDoFnRunnerTest.java | 18 ++++++++---- ...SimplePushbackSideInputDoFnRunnerTest.java | 8 ++++-- .../runners/core/StatefulDoFnRunnerTest.java | 3 +- .../triggers/TriggerStateMachineTester.java | 7 +++-- .../runners/direct/DirectTimerInternals.java | 6 ++-- .../translation/PipelineTranslatorUtils.java | 3 +- .../beam/sdk/util/construction/Timer.java | 3 +- .../beam/sdk/util/construction/TimerTest.java | 27 ++++++++++++------ .../beam/fn/harness/FnApiDoFnRunnerTest.java | 3 +- ...ablePairWithRestrictionDoFnRunnerTest.java | 3 +- .../control/ProcessBundleHandlerTest.java | 9 ++++-- 17 files changed, 84 insertions(+), 49 deletions(-) diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/InMemoryTimerInternals.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/InMemoryTimerInternals.java index 7e84ccbbccfe..f9439a1542a4 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/InMemoryTimerInternals.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/InMemoryTimerInternals.java @@ -119,7 +119,8 @@ public void setTimer( Instant target, Instant outputTimestamp, TimeDomain timeDomain) { - setTimer(TimerData.of(timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain)); + setTimer( + TimerData.of(timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain, null)); } /** diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java index 13358e750962..aa387062e27a 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java @@ -92,7 +92,8 @@ public void onTimer( Instant outputTimestamp, TimeDomain timeDomain, @Nullable Boolean draining) { - doFnRunner.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); + doFnRunner.onTimer( + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } @Override diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunner.java index 2be8071f9835..fec32485b827 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunner.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunner.java @@ -116,7 +116,8 @@ public void onTimer( Instant timestamp, Instant outputTimestamp, TimeDomain timeDomain) { - underlying.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain); + underlying.onTimer( + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, null); } @Override diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java index 6af54da0a08b..9244e9158331 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java @@ -598,9 +598,10 @@ public String getErrorContext() { timerInternals.currentProcessingTime().plus(result.getContinuation().resumeDelay()); holdState.add(futureOutputWatermark); // Set a timer to continue processing this element. + // todo radoslws@ decide if draining should be set on timer timerInternals.setTimer( TimerInternals.TimerData.of( - stateNamespace, wakeupTime, wakeupTime, TimeDomain.PROCESSING_TIME)); + stateNamespace, wakeupTime, wakeupTime, TimeDomain.PROCESSING_TIME, false)); } private DoFnInvoker.ArgumentProvider wrapOptionsAsSetup( diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/InMemoryTimerInternalsTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/InMemoryTimerInternalsTest.java index abdc49d887ea..eca4fb4c39d1 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/InMemoryTimerInternalsTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/InMemoryTimerInternalsTest.java @@ -41,9 +41,9 @@ public class InMemoryTimerInternalsTest { public void testFiringEventTimers() throws Exception { InMemoryTimerInternals underTest = new InMemoryTimerInternals(); TimerData eventTimer1 = - TimerData.of(ID1, NS1, new Instant(19), new Instant(19), TimeDomain.EVENT_TIME); + TimerData.of(ID1, NS1, new Instant(19), new Instant(19), TimeDomain.EVENT_TIME, null); TimerData eventTimer2 = - TimerData.of(ID2, NS1, new Instant(29), new Instant(29), TimeDomain.EVENT_TIME); + TimerData.of(ID2, NS1, new Instant(29), new Instant(29), TimeDomain.EVENT_TIME, null); underTest.setTimer(eventTimer1); underTest.setTimer(eventTimer2); @@ -82,7 +82,9 @@ public void testResetById() throws Exception { underTest.advanceInputWatermark(laterTimestamp.plus(Duration.millis(1L))); assertThat( underTest.removeNextEventTimer(), - equalTo(TimerData.of(ID1, "", NS1, laterTimestamp, laterTimestamp, TimeDomain.EVENT_TIME))); + equalTo( + TimerData.of( + ID1, "", NS1, laterTimestamp, laterTimestamp, TimeDomain.EVENT_TIME, null))); } @Test @@ -111,9 +113,9 @@ public void testDeletionById() throws Exception { public void testFiringProcessingTimeTimers() throws Exception { InMemoryTimerInternals underTest = new InMemoryTimerInternals(); TimerData processingTime1 = - TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.PROCESSING_TIME); + TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.PROCESSING_TIME, null); TimerData processingTime2 = - TimerData.of(NS1, new Instant(29), new Instant(29), TimeDomain.PROCESSING_TIME); + TimerData.of(NS1, new Instant(29), new Instant(29), TimeDomain.PROCESSING_TIME, null); underTest.setTimer(processingTime1); underTest.setTimer(processingTime2); @@ -142,19 +144,19 @@ public void testFiringProcessingTimeTimers() throws Exception { public void testTimerOrdering() throws Exception { InMemoryTimerInternals underTest = new InMemoryTimerInternals(); TimerData eventTime1 = - TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.EVENT_TIME); + TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.EVENT_TIME, null); TimerData processingTime1 = - TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.PROCESSING_TIME); + TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.PROCESSING_TIME, null); TimerData synchronizedProcessingTime1 = TimerData.of( - NS1, new Instant(19), new Instant(19), TimeDomain.SYNCHRONIZED_PROCESSING_TIME); + NS1, new Instant(19), new Instant(19), TimeDomain.SYNCHRONIZED_PROCESSING_TIME, null); TimerData eventTime2 = - TimerData.of(NS1, new Instant(29), new Instant(29), TimeDomain.EVENT_TIME); + TimerData.of(NS1, new Instant(29), new Instant(29), TimeDomain.EVENT_TIME, null); TimerData processingTime2 = - TimerData.of(NS1, new Instant(29), new Instant(29), TimeDomain.PROCESSING_TIME); + TimerData.of(NS1, new Instant(29), new Instant(29), TimeDomain.PROCESSING_TIME, null); TimerData synchronizedProcessingTime2 = TimerData.of( - NS1, new Instant(29), new Instant(29), TimeDomain.SYNCHRONIZED_PROCESSING_TIME); + NS1, new Instant(29), new Instant(29), TimeDomain.SYNCHRONIZED_PROCESSING_TIME, null); underTest.setTimer(processingTime1); underTest.setTimer(eventTime1); @@ -188,9 +190,9 @@ public void testTimerOrdering() throws Exception { public void testDeduplicate() throws Exception { InMemoryTimerInternals underTest = new InMemoryTimerInternals(); TimerData eventTime = - TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.EVENT_TIME); + TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.EVENT_TIME, null); TimerData processingTime = - TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.PROCESSING_TIME); + TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.PROCESSING_TIME, null); underTest.setTimer(eventTime); underTest.setTimer(eventTime); underTest.setTimer(processingTime); diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/KeyedWorkItemCoderTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/KeyedWorkItemCoderTest.java index 6099f1fdc8e1..d52c797619ef 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/KeyedWorkItemCoderTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/KeyedWorkItemCoderTest.java @@ -48,7 +48,8 @@ public void testEncodeDecodeEqual() throws Exception { StateNamespaces.global(), new Instant(500L), new Instant(500L), - TimeDomain.EVENT_TIME)); + TimeDomain.EVENT_TIME, + null)); Iterable> elements = ImmutableList.of( WindowedValues.valueInGlobalWindow(1), diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java index 40bcb70e32c6..354730506e41 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java @@ -140,7 +140,8 @@ public void testOnTimerExceptionsWrappedAsUserCodeException() { GlobalWindow.INSTANCE, new Instant(0), new Instant(0), - TimeDomain.EVENT_TIME); + TimeDomain.EVENT_TIME, + null); } /** @@ -265,7 +266,8 @@ public void testOnTimerCalled() { GlobalWindow.INSTANCE, currentTime.plus(offset), currentTime.plus(offset), - TimeDomain.EVENT_TIME); + TimeDomain.EVENT_TIME, + null); assertThat( fn.onTimerInvocations, @@ -276,7 +278,8 @@ public void testOnTimerCalled() { StateNamespaces.window(windowFn.windowCoder(), GlobalWindow.INSTANCE), currentTime.plus(offset), currentTime.plus(offset), - TimeDomain.EVENT_TIME))); + TimeDomain.EVENT_TIME, + null))); } /** @@ -592,7 +595,8 @@ public void testOnTimerAllowedSkew() { GlobalWindow.INSTANCE, new Instant(0), new Instant(0), - TimeDomain.EVENT_TIME); + TimeDomain.EVENT_TIME, + null); } @Test @@ -624,7 +628,8 @@ public void testOnTimerNoSkew() { GlobalWindow.INSTANCE, new Instant(0), new Instant(0), - TimeDomain.EVENT_TIME); + TimeDomain.EVENT_TIME, + null); }); assertThat(exception.getCause(), isA(IllegalArgumentException.class)); @@ -701,7 +706,8 @@ public void onTimer(OnTimerContext context) { StateNamespaces.window(windowCoder, (W) context.window()), context.fireTimestamp(), context.timestamp(), - context.timeDomain())); + context.timeDomain(), + null)); } } diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java index f9dc647e842f..d4b1973af994 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java @@ -18,7 +18,6 @@ package org.apache.beam.runners.core; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; -import org.checkerframework.checker.nullness.qual.Nullable; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.contains; import static org.hamcrest.Matchers.containsInAnyOrder; @@ -65,6 +64,7 @@ import org.apache.beam.sdk.values.WindowingStrategy; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; import org.joda.time.Instant; import org.junit.Before; @@ -318,7 +318,8 @@ public void testOnTimerCalled() { StateNamespaces.window(IntervalWindow.getCoder(), window), timestamp, timestamp, - TimeDomain.EVENT_TIME))); + TimeDomain.EVENT_TIME, + null))); } private static class TestDoFnRunner implements DoFnRunner { @@ -361,7 +362,8 @@ public void onTimer( StateNamespaces.window(IntervalWindow.getCoder(), (IntervalWindow) window), timestamp, outputTimestamp, - timeDomain, draining)); + timeDomain, + draining)); } @Override diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/StatefulDoFnRunnerTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/StatefulDoFnRunnerTest.java index 4ed4f09fe62b..813d893c826b 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/StatefulDoFnRunnerTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/StatefulDoFnRunnerTest.java @@ -468,7 +468,8 @@ private static void advanceInputWatermark( window, timer.getTimestamp(), timer.getOutputTimestamp(), - timer.getDomain()); + timer.getDomain(), + timer.getDraining()); } } diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/TriggerStateMachineTester.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/TriggerStateMachineTester.java index 21dfddce7a6e..a5aed59ba223 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/TriggerStateMachineTester.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/TriggerStateMachineTester.java @@ -400,17 +400,18 @@ public TestTimers(StateNamespace namespace) { @Override public void setTimer(Instant timestamp, TimeDomain timeDomain) { - timerInternals.setTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain)); + timerInternals.setTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain, null)); } @Override public void setTimer(Instant timestamp, Instant outputTimestamp, TimeDomain timeDomain) { - timerInternals.setTimer(TimerData.of(namespace, timestamp, outputTimestamp, timeDomain)); + timerInternals.setTimer( + TimerData.of(namespace, timestamp, outputTimestamp, timeDomain, null)); } @Override public void deleteTimer(Instant timestamp, TimeDomain timeDomain) { - timerInternals.deleteTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain)); + timerInternals.deleteTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain, null)); } @Override diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectTimerInternals.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectTimerInternals.java index db1112f73885..d6ba428eb804 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectTimerInternals.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectTimerInternals.java @@ -67,7 +67,8 @@ public void setTimer( Instant target, Instant outputTimestamp, TimeDomain timeDomain) { - setTimer(TimerData.of(timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain)); + setTimer( + TimerData.of(timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain, null)); } /** @@ -92,7 +93,8 @@ public void deleteTimer( namespace, BoundedWindow.TIMESTAMP_MIN_VALUE, BoundedWindow.TIMESTAMP_MAX_VALUE, - timeDomain)); + timeDomain, + null)); } /** @deprecated use {@link #deleteTimer(StateNamespace, String, TimeDomain)}. */ diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/translation/PipelineTranslatorUtils.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/translation/PipelineTranslatorUtils.java index 5b7880ad56ee..3848bd99c127 100644 --- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/translation/PipelineTranslatorUtils.java +++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/translation/PipelineTranslatorUtils.java @@ -156,7 +156,8 @@ private static void fireTimer( Collections.singletonList(window), timestamp, outputTimestamp, - PaneInfo.NO_FIRING); + PaneInfo.NO_FIRING, + null); KV transformAndTimerId = TimerReceiverFactory.decodeTimerDataTimerId(timer.getTimerFamilyId()); FnDataReceiver fnTimerReceiver = timerReceivers.get(transformAndTimerId); diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Timer.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Timer.java index c00039750f4f..52cb87f96fa7 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Timer.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/Timer.java @@ -204,7 +204,8 @@ public Timer decode(InputStream inStream) throws IOException { Instant fireTimestamp = InstantCoder.of().decode(inStream); Instant holdTimestamp = InstantCoder.of().decode(inStream); PaneInfo paneInfo = PaneInfoCoder.INSTANCE.decode(inStream); - return Timer.of(userKey, dynamicTimerTag, windows, fireTimestamp, holdTimestamp, paneInfo, false); + return Timer.of( + userKey, dynamicTimerTag, windows, fireTimestamp, holdTimestamp, paneInfo, false); } @Override diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/TimerTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/TimerTest.java index e96bf27c6ed9..a5d52a2dc892 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/TimerTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/TimerTest.java @@ -57,7 +57,8 @@ public void testTimer() { Collections.singleton(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING); + PaneInfo.NO_FIRING, + null); assertEquals("key", timer.getUserKey()); assertEquals("tag", timer.getDynamicTimerTag()); assertEquals(FIRE_TIME, timer.getFireTimestamp()); @@ -79,7 +80,8 @@ public void testTimerCoderWithInconsistentWithEqualsComponentCoders() throws Exc Collections.singleton(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING)); + PaneInfo.NO_FIRING, + null)); CoderProperties.structuralValueDecodeEncodeEqual( coder, Timer.cleared("key", "tag", Collections.singleton(GlobalWindow.INSTANCE))); CoderProperties.structuralValueConsistentWithEquals( @@ -90,14 +92,16 @@ public void testTimerCoderWithInconsistentWithEqualsComponentCoders() throws Exc Collections.singleton(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING), + PaneInfo.NO_FIRING, + null), Timer.of( "key", "tag", Collections.singleton(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING)); + PaneInfo.NO_FIRING, + null)); CoderProperties.structuralValueConsistentWithEquals( coder, Timer.cleared("key", "tag", Collections.singleton(GlobalWindow.INSTANCE)), @@ -115,7 +119,8 @@ public void testTimerCoderWithConsistentWithEqualsComponentCoders() throws Excep Collections.singletonList(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING)); + PaneInfo.NO_FIRING, + null)); CoderProperties.coderDecodeEncodeEqual( coder, Timer.cleared("key", "tag", Collections.singletonList(GlobalWindow.INSTANCE))); CoderProperties.coderConsistentWithEquals( @@ -126,14 +131,16 @@ public void testTimerCoderWithConsistentWithEqualsComponentCoders() throws Excep Collections.singletonList(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING), + PaneInfo.NO_FIRING, + null), Timer.of( "key", "tag", Collections.singletonList(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING)); + PaneInfo.NO_FIRING, + null)); CoderProperties.coderConsistentWithEquals( coder, Timer.cleared("key", "tag", Collections.singletonList(GlobalWindow.INSTANCE)), @@ -146,14 +153,16 @@ public void testTimerCoderWithConsistentWithEqualsComponentCoders() throws Excep Collections.singletonList(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING), + PaneInfo.NO_FIRING, + null), Timer.of( "key", "tag", Collections.singletonList(GlobalWindow.INSTANCE), FIRE_TIME, HOLD_TIME, - PaneInfo.NO_FIRING)); + PaneInfo.NO_FIRING, + null)); CoderProperties.coderDeterministic( coder, Timer.cleared("key", "tag", Collections.singletonList(GlobalWindow.INSTANCE)), 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 ef19b7c18804..c68dc0ba8eac 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 @@ -1161,7 +1161,8 @@ private org.apache.beam.sdk.util.construction.Timer dynamicTimerInGlobalW Collections.singletonList(GlobalWindow.INSTANCE), fireTimestamp, holdTimestamp, - PaneInfo.NO_FIRING); + PaneInfo.NO_FIRING, + null); } private WindowedValue valueInWindows( diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/SplittablePairWithRestrictionDoFnRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/SplittablePairWithRestrictionDoFnRunnerTest.java index 13a0b105ec39..8101d6dcc83a 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/SplittablePairWithRestrictionDoFnRunnerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/SplittablePairWithRestrictionDoFnRunnerTest.java @@ -123,7 +123,8 @@ private org.apache.beam.sdk.util.construction.Timer dynamicTimerInGlobalW Collections.singletonList(GlobalWindow.INSTANCE), fireTimestamp, holdTimestamp, - PaneInfo.NO_FIRING); + PaneInfo.NO_FIRING, + null); } private WindowedValue valueInWindows( diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java index a7a62571e38e..11823886759c 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java @@ -1112,7 +1112,8 @@ public void testInstructionEmbeddedElementsAreProcessed() throws Exception { Collections.singletonList(GlobalWindow.INSTANCE), Instant.ofEpochMilli(1L), Instant.ofEpochMilli(1L), - PaneInfo.ON_TIME_AND_ONLY_FIRING), + PaneInfo.ON_TIME_AND_ONLY_FIRING, + null), encodedTimer); Elements elements = Elements.newBuilder() @@ -1231,7 +1232,8 @@ public void testInstructionEmbeddedElementsWithMalformedTimers() throws Exceptio Collections.singletonList(GlobalWindow.INSTANCE), Instant.ofEpochMilli(1L), Instant.ofEpochMilli(1L), - PaneInfo.ON_TIME_AND_ONLY_FIRING), + PaneInfo.ON_TIME_AND_ONLY_FIRING, + null), encodedTimer); assertThrows( @@ -1325,7 +1327,8 @@ public void testOutputEmbeddedElementsAreProcessed() throws Exception { Collections.singletonList(GlobalWindow.INSTANCE), Instant.ofEpochMilli(1L), Instant.ofEpochMilli(1L), - PaneInfo.ON_TIME_AND_ONLY_FIRING), + PaneInfo.ON_TIME_AND_ONLY_FIRING, + null), encodedTimer); InstructionResponse.Builder builder = From d345ac3c824cc908ee509cc19146b12b3b564879 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 10 Oct 2025 14:12:09 +0200 Subject: [PATCH 14/26] revert timerdata constructor --- .../runners/core/InMemoryTimerInternals.java | 3 +- .../runners/core/ReduceFnContextFactory.java | 7 ++-- .../SplittableParDoViaKeyedWorkItems.java | 2 +- .../beam/runners/core/TimerInternals.java | 35 +++++++++++-------- .../core/InMemoryTimerInternalsTest.java | 28 +++++++-------- .../runners/core/KeyedWorkItemCoderTest.java | 3 +- .../beam/runners/core/ReduceFnTester.java | 9 ++--- .../runners/core/SimpleDoFnRunnerTest.java | 6 ++-- ...SimplePushbackSideInputDoFnRunnerTest.java | 17 +++++---- .../beam/runners/core/TimerInternalsTest.java | 32 +++++++---------- .../runners/direct/DirectTimerInternals.java | 6 ++-- .../metrics/DoFnRunnerWithMetricsUpdate.java | 3 +- .../ExecutableStageDoFnOperator.java | 6 ++-- .../worker/WindmillTimerInternals.java | 20 +++++------ .../StreamingGroupAlsoByWindowFnsTest.java | 3 +- ...gKeyedWorkItemSideInputDoFnRunnerTest.java | 3 +- .../StreamingModeExecutionContextTest.java | 3 +- .../worker/StreamingSideInputFetcherTest.java | 3 +- .../worker/UserParDoFnFactoryTest.java | 9 ++--- .../worker/WindmillKeyedWorkItemTest.java | 9 +++-- .../worker/WindmillTimerInternalsTest.java | 33 ++++++----------- .../control/TimerReceiverFactory.java | 18 ++++++---- .../samza/metrics/DoFnRunnerWithMetrics.java | 9 ++++- .../samza/runtime/AsyncDoFnRunner.java | 3 +- .../samza/runtime/SamzaDoFnRunners.java | 3 +- .../batch/DoFnRunnerWithMetrics.java | 3 +- .../translation/DoFnRunnerWithMetrics.java | 3 +- .../SparkInputDataProcessorTest.java | 2 +- 28 files changed, 130 insertions(+), 151 deletions(-) diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/InMemoryTimerInternals.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/InMemoryTimerInternals.java index f9439a1542a4..7e84ccbbccfe 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/InMemoryTimerInternals.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/InMemoryTimerInternals.java @@ -119,8 +119,7 @@ public void setTimer( Instant target, Instant outputTimestamp, TimeDomain timeDomain) { - setTimer( - TimerData.of(timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain, null)); + setTimer(TimerData.of(timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain)); } /** diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnContextFactory.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnContextFactory.java index 48553f461de3..0511c6f3b8e0 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnContextFactory.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnContextFactory.java @@ -139,20 +139,19 @@ public TimersImpl(StateNamespace namespace) { @Override public void setTimer(Instant timestamp, TimeDomain timeDomain) { // todo radoslaws@ should we pass draining bit to timerdata - timerInternals.setTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain, false)); + timerInternals.setTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain)); } @Override public void setTimer(Instant timestamp, Instant outputTimestamp, TimeDomain timeDomain) { // todo radoslaws@ should we pass draining bit to timerdata - timerInternals.setTimer( - TimerData.of(namespace, timestamp, outputTimestamp, timeDomain, false)); + timerInternals.setTimer(TimerData.of(namespace, timestamp, outputTimestamp, timeDomain)); } @Override public void deleteTimer(Instant timestamp, TimeDomain timeDomain) { // todo radoslaws@ should we pass draining bit to timerdata - timerInternals.deleteTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain, false)); + timerInternals.deleteTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain)); } @Override diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java index 9244e9158331..7ddb146e0eb1 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java @@ -601,7 +601,7 @@ public String getErrorContext() { // todo radoslws@ decide if draining should be set on timer timerInternals.setTimer( TimerInternals.TimerData.of( - stateNamespace, wakeupTime, wakeupTime, TimeDomain.PROCESSING_TIME, false)); + stateNamespace, wakeupTime, wakeupTime, TimeDomain.PROCESSING_TIME)); } private DoFnInvoker.ArgumentProvider wrapOptionsAsSetup( diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/TimerInternals.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/TimerInternals.java index dbf96e256970..fabc755aae71 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/TimerInternals.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/TimerInternals.java @@ -200,10 +200,9 @@ public static TimerData of( StateNamespace namespace, Instant timestamp, Instant outputTimestamp, - TimeDomain domain, - @javax.annotation.Nullable Boolean draining) { + TimeDomain domain) { return new AutoValue_TimerInternals_TimerData( - timerId, "", namespace, timestamp, outputTimestamp, domain, false, draining); + timerId, "", namespace, timestamp, outputTimestamp, domain, false, null); } /** @@ -216,10 +215,9 @@ public static TimerData of( StateNamespace namespace, Instant timestamp, Instant outputTimestamp, - TimeDomain domain, - @Nullable Boolean draining) { + TimeDomain domain) { return new AutoValue_TimerInternals_TimerData( - timerId, timerFamilyId, namespace, timestamp, outputTimestamp, domain, false, draining); + timerId, timerFamilyId, namespace, timestamp, outputTimestamp, domain, false, null); } /** @@ -227,13 +225,9 @@ public static TimerData of( * deterministically generated from the {@code timestamp} and {@code domain}. */ public static TimerData of( - StateNamespace namespace, - Instant timestamp, - Instant outputTimestamp, - TimeDomain domain, - @Nullable Boolean draining) { + StateNamespace namespace, Instant timestamp, Instant outputTimestamp, TimeDomain domain) { String timerId = String.valueOf(domain.ordinal()) + ':' + timestamp.getMillis(); - return of(timerId, namespace, timestamp, outputTimestamp, domain, draining); + return of(timerId, namespace, timestamp, outputTimestamp, domain); } public TimerData deleted() { @@ -248,6 +242,18 @@ public TimerData deleted() { getDraining()); } + public TimerData draining(@Nullable Boolean draining) { + return new AutoValue_TimerInternals_TimerData( + getTimerId(), + getTimerFamilyId(), + getNamespace(), + getTimestamp(), + getOutputTimestamp(), + getDomain(), + getDeleted(), + draining); + } + /** * {@inheritDoc}. * @@ -325,8 +331,7 @@ public TimerData decode(InputStream inStream) throws CoderException, IOException Instant outputTimestamp = INSTANT_CODER.decode(inStream); TimeDomain domain = TimeDomain.valueOf(STRING_CODER.decode(inStream)); // we don't need to serialize draining bit - return TimerData.of( - timerId, timerFamilyId, namespace, timestamp, outputTimestamp, domain, false); + return TimerData.of(timerId, timerFamilyId, namespace, timestamp, outputTimestamp, domain); } @Override @@ -372,7 +377,7 @@ public TimerData decode(InputStream inStream) throws CoderException, IOException StateNamespaces.fromString(STRING_CODER.decode(inStream), windowCoder); Instant timestamp = INSTANT_CODER.decode(inStream); TimeDomain domain = TimeDomain.valueOf(STRING_CODER.decode(inStream)); - return TimerData.of(timerId, namespace, timestamp, timestamp, domain, false); + return TimerData.of(timerId, namespace, timestamp, timestamp, domain); } @Override diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/InMemoryTimerInternalsTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/InMemoryTimerInternalsTest.java index eca4fb4c39d1..abdc49d887ea 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/InMemoryTimerInternalsTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/InMemoryTimerInternalsTest.java @@ -41,9 +41,9 @@ public class InMemoryTimerInternalsTest { public void testFiringEventTimers() throws Exception { InMemoryTimerInternals underTest = new InMemoryTimerInternals(); TimerData eventTimer1 = - TimerData.of(ID1, NS1, new Instant(19), new Instant(19), TimeDomain.EVENT_TIME, null); + TimerData.of(ID1, NS1, new Instant(19), new Instant(19), TimeDomain.EVENT_TIME); TimerData eventTimer2 = - TimerData.of(ID2, NS1, new Instant(29), new Instant(29), TimeDomain.EVENT_TIME, null); + TimerData.of(ID2, NS1, new Instant(29), new Instant(29), TimeDomain.EVENT_TIME); underTest.setTimer(eventTimer1); underTest.setTimer(eventTimer2); @@ -82,9 +82,7 @@ public void testResetById() throws Exception { underTest.advanceInputWatermark(laterTimestamp.plus(Duration.millis(1L))); assertThat( underTest.removeNextEventTimer(), - equalTo( - TimerData.of( - ID1, "", NS1, laterTimestamp, laterTimestamp, TimeDomain.EVENT_TIME, null))); + equalTo(TimerData.of(ID1, "", NS1, laterTimestamp, laterTimestamp, TimeDomain.EVENT_TIME))); } @Test @@ -113,9 +111,9 @@ public void testDeletionById() throws Exception { public void testFiringProcessingTimeTimers() throws Exception { InMemoryTimerInternals underTest = new InMemoryTimerInternals(); TimerData processingTime1 = - TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.PROCESSING_TIME, null); + TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.PROCESSING_TIME); TimerData processingTime2 = - TimerData.of(NS1, new Instant(29), new Instant(29), TimeDomain.PROCESSING_TIME, null); + TimerData.of(NS1, new Instant(29), new Instant(29), TimeDomain.PROCESSING_TIME); underTest.setTimer(processingTime1); underTest.setTimer(processingTime2); @@ -144,19 +142,19 @@ public void testFiringProcessingTimeTimers() throws Exception { public void testTimerOrdering() throws Exception { InMemoryTimerInternals underTest = new InMemoryTimerInternals(); TimerData eventTime1 = - TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.EVENT_TIME, null); + TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.EVENT_TIME); TimerData processingTime1 = - TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.PROCESSING_TIME, null); + TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.PROCESSING_TIME); TimerData synchronizedProcessingTime1 = TimerData.of( - NS1, new Instant(19), new Instant(19), TimeDomain.SYNCHRONIZED_PROCESSING_TIME, null); + NS1, new Instant(19), new Instant(19), TimeDomain.SYNCHRONIZED_PROCESSING_TIME); TimerData eventTime2 = - TimerData.of(NS1, new Instant(29), new Instant(29), TimeDomain.EVENT_TIME, null); + TimerData.of(NS1, new Instant(29), new Instant(29), TimeDomain.EVENT_TIME); TimerData processingTime2 = - TimerData.of(NS1, new Instant(29), new Instant(29), TimeDomain.PROCESSING_TIME, null); + TimerData.of(NS1, new Instant(29), new Instant(29), TimeDomain.PROCESSING_TIME); TimerData synchronizedProcessingTime2 = TimerData.of( - NS1, new Instant(29), new Instant(29), TimeDomain.SYNCHRONIZED_PROCESSING_TIME, null); + NS1, new Instant(29), new Instant(29), TimeDomain.SYNCHRONIZED_PROCESSING_TIME); underTest.setTimer(processingTime1); underTest.setTimer(eventTime1); @@ -190,9 +188,9 @@ public void testTimerOrdering() throws Exception { public void testDeduplicate() throws Exception { InMemoryTimerInternals underTest = new InMemoryTimerInternals(); TimerData eventTime = - TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.EVENT_TIME, null); + TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.EVENT_TIME); TimerData processingTime = - TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.PROCESSING_TIME, null); + TimerData.of(NS1, new Instant(19), new Instant(19), TimeDomain.PROCESSING_TIME); underTest.setTimer(eventTime); underTest.setTimer(eventTime); underTest.setTimer(processingTime); diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/KeyedWorkItemCoderTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/KeyedWorkItemCoderTest.java index d52c797619ef..6099f1fdc8e1 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/KeyedWorkItemCoderTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/KeyedWorkItemCoderTest.java @@ -48,8 +48,7 @@ public void testEncodeDecodeEqual() throws Exception { StateNamespaces.global(), new Instant(500L), new Instant(500L), - TimeDomain.EVENT_TIME, - null)); + TimeDomain.EVENT_TIME)); Iterable> elements = ImmutableList.of( WindowedValues.valueInGlobalWindow(1), diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java index d9582873477f..b73d6a6e54b7 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java @@ -576,11 +576,7 @@ public void fireTimer(W window, Instant timestamp, TimeDomain domain) throws Exc ArrayList timers = new ArrayList<>(1); timers.add( TimerData.of( - StateNamespaces.window(windowFn.windowCoder(), window), - timestamp, - timestamp, - domain, - null)); + StateNamespaces.window(windowFn.windowCoder(), window), timestamp, timestamp, domain)); runner.onTimers(timers); runner.persist(); } @@ -594,8 +590,7 @@ public void fireTimers(W window, TimestampedValue... timers) throws StateNamespaces.window(windowFn.windowCoder(), window), timer.getTimestamp(), timer.getTimestamp(), - timer.getValue(), - null)); + timer.getValue())); } runner.onTimers(timerData); runner.persist(); diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java index 354730506e41..52e241b0e92c 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimpleDoFnRunnerTest.java @@ -278,8 +278,7 @@ public void testOnTimerCalled() { StateNamespaces.window(windowFn.windowCoder(), GlobalWindow.INSTANCE), currentTime.plus(offset), currentTime.plus(offset), - TimeDomain.EVENT_TIME, - null))); + TimeDomain.EVENT_TIME))); } /** @@ -706,8 +705,7 @@ public void onTimer(OnTimerContext context) { StateNamespaces.window(windowCoder, (W) context.window()), context.fireTimestamp(), context.timestamp(), - context.timeDomain(), - null)); + context.timeDomain())); } } diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java index d4b1973af994..5d0ba9d77c99 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java @@ -318,8 +318,7 @@ public void testOnTimerCalled() { StateNamespaces.window(IntervalWindow.getCoder(), window), timestamp, timestamp, - TimeDomain.EVENT_TIME, - null))); + TimeDomain.EVENT_TIME))); } private static class TestDoFnRunner implements DoFnRunner { @@ -357,13 +356,13 @@ public void onTimer( @Nullable Boolean draining) { firedTimers.add( TimerData.of( - timerId, - timerFamilyId, - StateNamespaces.window(IntervalWindow.getCoder(), (IntervalWindow) window), - timestamp, - outputTimestamp, - timeDomain, - draining)); + timerId, + timerFamilyId, + StateNamespaces.window(IntervalWindow.getCoder(), (IntervalWindow) window), + timestamp, + outputTimestamp, + timeDomain) + .draining(draining)); } @Override diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/TimerInternalsTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/TimerInternalsTest.java index 640c2bb7128e..83fb30eb8120 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/TimerInternalsTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/TimerInternalsTest.java @@ -46,8 +46,7 @@ public void testTimerDataCoder() throws Exception { StateNamespaces.global(), new Instant(0), new Instant(0), - TimeDomain.EVENT_TIME, - null)); + TimeDomain.EVENT_TIME)); Coder windowCoder = IntervalWindow.getCoder(); CoderProperties.coderDecodeEncodeEqual( @@ -58,8 +57,7 @@ public void testTimerDataCoder() throws Exception { windowCoder, new IntervalWindow(new Instant(0), new Instant(100))), new Instant(99), new Instant(99), - TimeDomain.PROCESSING_TIME, - null)); + TimeDomain.PROCESSING_TIME)); } @Test @@ -71,13 +69,12 @@ public void testCoderIsSerializableWithWellKnownCoderType() { public void testCompareEqual() { Instant timestamp = new Instant(100); StateNamespace namespace = StateNamespaces.global(); - TimerData timer = - TimerData.of("id", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME, null); + TimerData timer = TimerData.of("id", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME); assertThat( timer, comparesEqualTo( - TimerData.of("id", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME, null))); + TimerData.of("id", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME))); } @Test @@ -87,9 +84,9 @@ public void testCompareByTimestamp() { StateNamespace namespace = StateNamespaces.global(); TimerData firstTimer = - TimerData.of(namespace, firstTimestamp, firstTimestamp, TimeDomain.EVENT_TIME, null); + TimerData.of(namespace, firstTimestamp, firstTimestamp, TimeDomain.EVENT_TIME); TimerData secondTimer = - TimerData.of(namespace, secondTimestamp, secondTimestamp, TimeDomain.EVENT_TIME, null); + TimerData.of(namespace, secondTimestamp, secondTimestamp, TimeDomain.EVENT_TIME); assertThat(firstTimer, lessThan(secondTimer)); } @@ -99,13 +96,10 @@ public void testCompareByDomain() { Instant timestamp = new Instant(100); StateNamespace namespace = StateNamespaces.global(); - TimerData eventTimer = - TimerData.of(namespace, timestamp, timestamp, TimeDomain.EVENT_TIME, null); - TimerData procTimer = - TimerData.of(namespace, timestamp, timestamp, TimeDomain.PROCESSING_TIME, null); + TimerData eventTimer = TimerData.of(namespace, timestamp, timestamp, TimeDomain.EVENT_TIME); + TimerData procTimer = TimerData.of(namespace, timestamp, timestamp, TimeDomain.PROCESSING_TIME); TimerData synchronizedProcTimer = - TimerData.of( - namespace, timestamp, timestamp, TimeDomain.SYNCHRONIZED_PROCESSING_TIME, null); + TimerData.of(namespace, timestamp, timestamp, TimeDomain.SYNCHRONIZED_PROCESSING_TIME); assertThat(eventTimer, lessThan(procTimer)); assertThat(eventTimer, lessThan(synchronizedProcTimer)); @@ -123,9 +117,9 @@ public void testCompareByNamespace() { StateNamespace secondWindowNs = StateNamespaces.window(windowCoder, secondWindow); TimerData secondEventTime = - TimerData.of(firstWindowNs, timestamp, timestamp, TimeDomain.EVENT_TIME, null); + TimerData.of(firstWindowNs, timestamp, timestamp, TimeDomain.EVENT_TIME); TimerData thirdEventTime = - TimerData.of(secondWindowNs, timestamp, timestamp, TimeDomain.EVENT_TIME, null); + TimerData.of(secondWindowNs, timestamp, timestamp, TimeDomain.EVENT_TIME); assertThat(secondEventTime, lessThan(thirdEventTime)); } @@ -136,9 +130,9 @@ public void testCompareByTimerId() { StateNamespace namespace = StateNamespaces.global(); TimerData id0Timer = - TimerData.of("id0", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME, null); + TimerData.of("id0", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME); TimerData id1Timer = - TimerData.of("id1", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME, null); + TimerData.of("id1", namespace, timestamp, timestamp, TimeDomain.EVENT_TIME); assertThat(id0Timer, lessThan(id1Timer)); } diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectTimerInternals.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectTimerInternals.java index d6ba428eb804..db1112f73885 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectTimerInternals.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectTimerInternals.java @@ -67,8 +67,7 @@ public void setTimer( Instant target, Instant outputTimestamp, TimeDomain timeDomain) { - setTimer( - TimerData.of(timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain, null)); + setTimer(TimerData.of(timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain)); } /** @@ -93,8 +92,7 @@ public void deleteTimer( namespace, BoundedWindow.TIMESTAMP_MIN_VALUE, BoundedWindow.TIMESTAMP_MAX_VALUE, - timeDomain, - null)); + timeDomain)); } /** @deprecated use {@link #deleteTimer(StateNamespace, String, TimeDomain)}. */ diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java index 6d8ee22fe9af..cc0cf4e1a0b0 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java @@ -78,7 +78,8 @@ public void onTimer( @Nullable Boolean draining) { try (Closeable ignored = MetricsEnvironment.scopedMetricsContainer(container.getMetricsContainer(stepName))) { - delegate.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); + delegate.onTimer( + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } catch (IOException e) { throw new RuntimeException(e); } diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java index fb1b11baac23..76dad7d6f727 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java @@ -592,7 +592,8 @@ public void setTimer( Instant outputTimestamp, TimeDomain timeDomain) { setTimer( - TimerData.of(timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain, false)); + TimerData.of( + timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain)); } @Override @@ -1035,7 +1036,8 @@ public void onTimer( timestamp, outputTimestamp, // TODO: Support propagating the PaneInfo through. - PaneInfo.NO_FIRING, draining); + PaneInfo.NO_FIRING, + draining); try { timerReceiver.accept(timerValue); } catch (Exception e) { diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternals.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternals.java index 9732e8633017..a739ca7fcb1a 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternals.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternals.java @@ -111,8 +111,7 @@ public void setTimer( Instant outputTimestamp, TimeDomain timeDomain) { TimerData timer = - TimerData.of( - timerId, timerFamilyId, namespace, timestamp, outputTimestamp, timeDomain, false); + TimerData.of(timerId, timerFamilyId, namespace, timestamp, outputTimestamp, timeDomain); setTimer(timer); } @@ -149,8 +148,7 @@ public void deleteTimer( namespace, BoundedWindow.TIMESTAMP_MIN_VALUE, BoundedWindow.TIMESTAMP_MAX_VALUE, - timeDomain, - false)); + timeDomain)); } @Override @@ -391,13 +389,13 @@ public static TimerData windmillTimerToTimerData( StateNamespace namespace = StateNamespaces.fromString(namespaceString, windowCoder); return TimerData.of( - timerId, - timerFamily, - namespace, - timestamp, - outputTimestamp, - timerTypeToTimeDomain(timer.getType()), - draining); + timerId, + timerFamily, + namespace, + timestamp, + outputTimestamp, + timerTypeToTimeDomain(timer.getType())) + .draining(draining); } private static boolean useNewTimerTagEncoding(TimerData timerData) { diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java index 654d9965f1da..f7852ec1767d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java @@ -156,8 +156,7 @@ private void addTimer( timestamp, type == Windmill.Timer.Type.WATERMARK ? TimeDomain.EVENT_TIME - : TimeDomain.PROCESSING_TIME, - false))) + : TimeDomain.PROCESSING_TIME))) .setTimestamp(WindmillTimeUtils.harnessToWindmillTimestamp(timestamp)) .setType(type) .setStateFamily(STATE_FAMILY); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunnerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunnerTest.java index 1100d99b2a3d..de4ec8e5aa9e 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunnerTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingKeyedWorkItemSideInputDoFnRunnerTest.java @@ -166,8 +166,7 @@ private TimerData timerData(IntervalWindow window, Instant timestamp, Timer.Type StateNamespaces.window(IntervalWindow.getCoder(), window), timestamp, timestamp, - type == Windmill.Timer.Type.WATERMARK ? TimeDomain.EVENT_TIME : TimeDomain.PROCESSING_TIME, - null); + type == Windmill.Timer.Type.WATERMARK ? TimeDomain.EVENT_TIME : TimeDomain.PROCESSING_TIME); } private IntervalWindow window(long start, long end) { diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java index 377acabca172..dd5c36c9cf85 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java @@ -173,8 +173,7 @@ public void testTimerInternalsSetTimer() { new StateNamespaceForTest("key"), new Instant(5000), new Instant(5000), - TimeDomain.EVENT_TIME, - null)); + TimeDomain.EVENT_TIME)); executionContext.flushState(); Windmill.Timer timer = outputBuilder.buildPartial().getOutputTimers(0); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcherTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcherTest.java index def611381ab4..cf616ee6ac0d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcherTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingSideInputFetcherTest.java @@ -210,8 +210,7 @@ private TimerData createTimer(long timestamp) { StateNamespaces.window(IntervalWindow.getCoder(), createWindow(timestamp)), new Instant(timestamp), new Instant(timestamp), - TimeDomain.EVENT_TIME, - null); + TimeDomain.EVENT_TIME); } private IntervalWindow createWindow(long timestamp) { diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactoryTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactoryTest.java index d2ac5e0e6cbd..771ffe192baf 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactoryTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/UserParDoFnFactoryTest.java @@ -462,8 +462,7 @@ public void testCleanupTimerForGlobalWindowWithAllowedLateness() throws Exceptio globalWindowNamespace, BoundedWindow.TIMESTAMP_MAX_VALUE, BoundedWindow.TIMESTAMP_MAX_VALUE.minus(Duration.millis(1)), - TimeDomain.EVENT_TIME, - null)) + TimeDomain.EVENT_TIME)) .thenReturn(null); // Set up non-empty state. We don't mock + verify calls to clear() but instead @@ -538,8 +537,7 @@ public void testCleanupWorks() throws Exception { firstWindowNamespace, firstWindow.maxTimestamp().plus(Duration.millis(1L)), firstWindow.maxTimestamp().plus(Duration.millis(1L)), - TimeDomain.EVENT_TIME, - null)) + TimeDomain.EVENT_TIME)) .thenReturn(null); // This should fire the timer to clean up the first window @@ -555,8 +553,7 @@ public void testCleanupWorks() throws Exception { secondWindowNamespace, secondWindow.maxTimestamp().plus(Duration.millis(1L)), secondWindow.maxTimestamp().plus(Duration.millis(1L)), - TimeDomain.EVENT_TIME, - null)) + TimeDomain.EVENT_TIME)) .thenReturn(null); // And this should clean up the second window diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java index 8b0a29ad60aa..53a36722e41c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java @@ -93,7 +93,7 @@ public void testElementIteration() throws Exception { KeyedWorkItem keyedWorkItem = new WindmillKeyedWorkItem<>( - KEY, workItem.build(), WINDOW_CODER, WINDOWS_CODER, VALUE_CODER, false); + KEY, workItem.build(), WINDOW_CODER, WINDOWS_CODER, VALUE_CODER); assertThat( keyedWorkItem.elementsIterable(), @@ -148,7 +148,7 @@ public void testTimerOrdering() throws Exception { .build(); KeyedWorkItem keyedWorkItem = - new WindmillKeyedWorkItem<>(KEY, workItem, WINDOW_CODER, WINDOWS_CODER, VALUE_CODER, false); + new WindmillKeyedWorkItem<>(KEY, workItem, WINDOW_CODER, WINDOWS_CODER, VALUE_CODER); assertThat( keyedWorkItem.timersIterable(), @@ -169,8 +169,7 @@ private static Windmill.Timer makeSerializedTimer( ns, new Instant(timestamp), new Instant(timestamp), - WindmillTimerInternals.timerTypeToTimeDomain(type), - false))) + WindmillTimerInternals.timerTypeToTimeDomain(type)))) .setTimestamp(WindmillTimeUtils.harnessToWindmillTimestamp(new Instant(timestamp))) .setType(type) .setStateFamily(STATE_FAMILY) @@ -178,7 +177,7 @@ private static Windmill.Timer makeSerializedTimer( } private static TimerData makeTimer(StateNamespace ns, long timestamp, TimeDomain domain) { - return TimerData.of(ns, new Instant(timestamp), new Instant(timestamp), domain, false); + return TimerData.of(ns, new Instant(timestamp), new Instant(timestamp), domain); } @Test diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternalsTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternalsTest.java index eaef5956d15b..ec8672b6a75f 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternalsTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternalsTest.java @@ -84,13 +84,9 @@ public void testTimerDataToFromTimer() { for (Instant timestamp : TEST_TIMESTAMPS) { List anonymousTimers = ImmutableList.of( - TimerData.of(namespace, timestamp, timestamp, timeDomain, null), + TimerData.of(namespace, timestamp, timestamp, timeDomain), TimerData.of( - namespace, - timestamp, - timestamp.minus(Duration.millis(1)), - timeDomain, - null)); + namespace, timestamp, timestamp.minus(Duration.millis(1)), timeDomain)); for (TimerData timer : anonymousTimers) { Instant expectedTimestamp = timer.getOutputTimestamp().isBefore(BoundedWindow.TIMESTAMP_MIN_VALUE) @@ -100,17 +96,12 @@ public void testTimerDataToFromTimer() { WindmillTimerInternals.windmillTimerToTimerData( prefix, WindmillTimerInternals.timerDataToWindmillTimer(stateFamily, prefix, timer), - coder, - null); + coder); // The function itself bounds output, so we dont expect the original input as the // output, we expect it to be bounded TimerData expected = TimerData.of( - timer.getNamespace(), - timestamp, - expectedTimestamp, - timer.getDomain(), - null); + timer.getNamespace(), timestamp, expectedTimestamp, timer.getDomain()); assertThat(computed, equalTo(expected)); } @@ -118,24 +109,22 @@ public void testTimerDataToFromTimer() { for (String timerId : TEST_TIMER_IDS) { List timers = ImmutableList.of( - TimerData.of(timerId, namespace, timestamp, timestamp, timeDomain, null), + TimerData.of(timerId, namespace, timestamp, timestamp, timeDomain), TimerData.of( - timerId, "family", namespace, timestamp, timestamp, timeDomain, null), + timerId, "family", namespace, timestamp, timestamp, timeDomain), TimerData.of( timerId, namespace, timestamp, timestamp.minus(Duration.millis(1)), - timeDomain, - null), + timeDomain), TimerData.of( timerId, "family", namespace, timestamp, timestamp.minus(Duration.millis(1)), - timeDomain, - null)); + timeDomain)); for (TimerData timer : timers) { Instant expectedTimestamp = @@ -150,15 +139,13 @@ public void testTimerDataToFromTimer() { timer.getNamespace(), timer.getTimestamp(), expectedTimestamp, - timer.getDomain(), - null); + timer.getDomain()); assertThat( WindmillTimerInternals.windmillTimerToTimerData( prefix, WindmillTimerInternals.timerDataToWindmillTimer( stateFamily, prefix, timer), - coder, - null), + coder), equalTo(expected)); } } diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/TimerReceiverFactory.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/TimerReceiverFactory.java index 3c01bf244975..c2ad60d6280d 100644 --- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/TimerReceiverFactory.java +++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/TimerReceiverFactory.java @@ -89,13 +89,17 @@ public FnDataReceiver> create(String transformId, String timerFamil StateNamespace namespace = StateNamespaces.window(windowCoder, (BoundedWindow) window); TimerInternals.TimerData timerData = TimerInternals.TimerData.of( - timer.getDynamicTimerTag(), - encodeToTimerDataTimerId(timerSpec.transformId(), timerSpec.timerId()), - namespace, - timer.getClearBit() ? BoundedWindow.TIMESTAMP_MAX_VALUE : timer.getFireTimestamp(), - timer.getClearBit() ? BoundedWindow.TIMESTAMP_MAX_VALUE : timer.getHoldTimestamp(), - timerSpec.getTimerSpec().getTimeDomain(), - timer.getDraining()); + timer.getDynamicTimerTag(), + encodeToTimerDataTimerId(timerSpec.transformId(), timerSpec.timerId()), + namespace, + timer.getClearBit() + ? BoundedWindow.TIMESTAMP_MAX_VALUE + : timer.getFireTimestamp(), + timer.getClearBit() + ? BoundedWindow.TIMESTAMP_MAX_VALUE + : timer.getHoldTimestamp(), + timerSpec.getTimerSpec().getTimeDomain()) + .draining(timer.getDraining()); timerDataConsumer.accept(timer, timerData); } }; diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/DoFnRunnerWithMetrics.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/DoFnRunnerWithMetrics.java index f7b6d3e8e68e..237b9ba19a42 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/DoFnRunnerWithMetrics.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/metrics/DoFnRunnerWithMetrics.java @@ -67,7 +67,14 @@ public void onTimer( withMetrics( () -> underlying.onTimer( - timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining), + timerId, + timerFamilyId, + key, + window, + timestamp, + outputTimestamp, + timeDomain, + draining), false); } diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/AsyncDoFnRunner.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/AsyncDoFnRunner.java index daac80b323a3..08f4b367e9c2 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/AsyncDoFnRunner.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/AsyncDoFnRunner.java @@ -156,7 +156,8 @@ public void onTimer( Instant outputTimestamp, TimeDomain timeDomain, @Nullable Boolean draining) { - underlying.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); + underlying.onTimer( + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } @Override diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java index 1db27342dc8f..d012b33b428f 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/SamzaDoFnRunners.java @@ -445,7 +445,8 @@ public void onTimer( timestamp, outputTimestamp, // TODO: Support propagating the PaneInfo through. - PaneInfo.NO_FIRING, draining); + PaneInfo.NO_FIRING, + draining); try { timerReceiver.accept(timerValue); } catch (Exception e) { diff --git a/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerWithMetrics.java b/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerWithMetrics.java index f6fe810a2ae6..e9b42b169009 100644 --- a/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerWithMetrics.java +++ b/runners/spark/3/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/DoFnRunnerWithMetrics.java @@ -81,7 +81,8 @@ public void onTimer( final TimeDomain timeDomain, @Nullable Boolean draining) { try (Closeable ignored = MetricsEnvironment.scopedMetricsContainer(metrics)) { - delegate.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); + delegate.onTimer( + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } catch (IOException e) { throw new RuntimeException(e); } diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/DoFnRunnerWithMetrics.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/DoFnRunnerWithMetrics.java index b80d451d64b0..c26d9f75e98b 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/DoFnRunnerWithMetrics.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/DoFnRunnerWithMetrics.java @@ -80,7 +80,8 @@ public void onTimer( final TimeDomain timeDomain, @Nullable Boolean draining) { try (Closeable ignored = MetricsEnvironment.scopedMetricsContainer(metricsContainer())) { - delegate.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); + delegate.onTimer( + timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining); } catch (IOException e) { throw new RuntimeException(e); } diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkInputDataProcessorTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkInputDataProcessorTest.java index d3e929bee99c..346587993472 100644 --- a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkInputDataProcessorTest.java +++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/SparkInputDataProcessorTest.java @@ -17,7 +17,6 @@ */ package org.apache.beam.runners.spark.translation; -import org.checkerframework.checker.nullness.qual.Nullable; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; @@ -41,6 +40,7 @@ import org.apache.beam.sdk.values.WindowedValues; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterators; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; import org.junit.Rule; import org.junit.Test; From c24645b501792971a3e46dfbb9e617795c1b94ac Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 10 Oct 2025 14:43:21 +0200 Subject: [PATCH 15/26] revert timerdata constructor --- .../runners/core/triggers/TriggerStateMachineTester.java | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/TriggerStateMachineTester.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/TriggerStateMachineTester.java index a5aed59ba223..21dfddce7a6e 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/TriggerStateMachineTester.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/TriggerStateMachineTester.java @@ -400,18 +400,17 @@ public TestTimers(StateNamespace namespace) { @Override public void setTimer(Instant timestamp, TimeDomain timeDomain) { - timerInternals.setTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain, null)); + timerInternals.setTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain)); } @Override public void setTimer(Instant timestamp, Instant outputTimestamp, TimeDomain timeDomain) { - timerInternals.setTimer( - TimerData.of(namespace, timestamp, outputTimestamp, timeDomain, null)); + timerInternals.setTimer(TimerData.of(namespace, timestamp, outputTimestamp, timeDomain)); } @Override public void deleteTimer(Instant timestamp, TimeDomain timeDomain) { - timerInternals.deleteTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain, null)); + timerInternals.deleteTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain)); } @Override From 4daf59b67f83cf1dc6acec694127e4feabb02187 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 10 Oct 2025 14:43:55 +0200 Subject: [PATCH 16/26] revert timerdata constructor --- .../beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java index 53a36722e41c..ad9110152e74 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java @@ -93,7 +93,7 @@ public void testElementIteration() throws Exception { KeyedWorkItem keyedWorkItem = new WindmillKeyedWorkItem<>( - KEY, workItem.build(), WINDOW_CODER, WINDOWS_CODER, VALUE_CODER); + KEY, workItem.build(), WINDOW_CODER, WINDOWS_CODER, VALUE_CODER, null); assertThat( keyedWorkItem.elementsIterable(), From 19cb5e4efabd5150133b780935bf112787d1af7a Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 10 Oct 2025 14:44:18 +0200 Subject: [PATCH 17/26] revert timerdata constructor --- .../beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java index ad9110152e74..d66dc24a0c71 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java @@ -148,7 +148,7 @@ public void testTimerOrdering() throws Exception { .build(); KeyedWorkItem keyedWorkItem = - new WindmillKeyedWorkItem<>(KEY, workItem, WINDOW_CODER, WINDOWS_CODER, VALUE_CODER); + new WindmillKeyedWorkItem<>(KEY, workItem, WINDOW_CODER, WINDOWS_CODER, VALUE_CODER, null); assertThat( keyedWorkItem.timersIterable(), From 9fa4282b77b67988a4385f5b98f2ddc8827d9774 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 10 Oct 2025 14:46:09 +0200 Subject: [PATCH 18/26] revert timerdata constructor --- .../runners/dataflow/worker/WindmillTimerInternalsTest.java | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternalsTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternalsTest.java index ec8672b6a75f..853fe2ce4c62 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternalsTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillTimerInternalsTest.java @@ -96,7 +96,8 @@ public void testTimerDataToFromTimer() { WindmillTimerInternals.windmillTimerToTimerData( prefix, WindmillTimerInternals.timerDataToWindmillTimer(stateFamily, prefix, timer), - coder); + coder, + null); // The function itself bounds output, so we dont expect the original input as the // output, we expect it to be bounded TimerData expected = @@ -145,7 +146,8 @@ public void testTimerDataToFromTimer() { prefix, WindmillTimerInternals.timerDataToWindmillTimer( stateFamily, prefix, timer), - coder), + coder, + null), equalTo(expected)); } } From 7a33d72f9e28e33ac20b60d219ef0354d6a3a8b5 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 10 Oct 2025 17:50:27 +0200 Subject: [PATCH 19/26] reduceFnRunner changes --- .../GroupAlsoByWindowViaWindowSetNewDoFn.java | 1 - .../apache/beam/runners/core/ReduceFn.java | 4 +- .../runners/core/ReduceFnContextFactory.java | 44 +++++----- .../beam/runners/core/ReduceFnRunner.java | 43 +++++---- .../beam/runners/core/WindowMatchers.java | 52 +++++++++-- .../beam/runners/core/ReduceFnTester.java | 3 - .../GroupAlsoByWindowEvaluatorFactory.java | 1 - .../ExecutableStageDoFnOperator.java | 3 +- ...eamingGroupAlsoByWindowViaWindowSetFn.java | 1 - ...tchGroupAlsoByWindowViaOutputBufferFn.java | 1 - .../StreamingGroupAlsoByWindowFnsTest.java | 87 ++++++++++++++++++- .../beam/sdk/values/WindowedValues.java | 26 +++++- 12 files changed, 198 insertions(+), 68 deletions(-) diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/GroupAlsoByWindowViaWindowSetNewDoFn.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/GroupAlsoByWindowViaWindowSetNewDoFn.java index 38666b0ae901..f242c7d10003 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/GroupAlsoByWindowViaWindowSetNewDoFn.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/GroupAlsoByWindowViaWindowSetNewDoFn.java @@ -107,7 +107,6 @@ public void processElement(ProcessContext c) throws Exception { windowedValue -> outputManager.output(mainTag, windowedValue), sideInputReader, reduceFn, - c.draining(), c.getPipelineOptions()); reduceFnRunner.processElements(keyedWorkItem.elementsIterable()); diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFn.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFn.java index 1de6b3167350..146a60203f29 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFn.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFn.java @@ -44,8 +44,6 @@ public abstract class Context { /** The window that is being processed. */ public abstract W window(); - public abstract Boolean draining(); - /** Access the current {@link WindowingStrategy}. */ public abstract WindowingStrategy windowingStrategy(); @@ -61,6 +59,7 @@ public abstract class ProcessValueContext extends Context { /** Return the actual value being processed. */ public abstract InputT value(); + public abstract Boolean draining(); /** Return the timestamp associated with the value. */ public abstract Instant timestamp(); } @@ -77,6 +76,7 @@ public abstract class OnTriggerContext extends Context { /** Returns the {@link PaneInfo} for the trigger firing being processed. */ public abstract PaneInfo paneInfo(); + public abstract Boolean draining(); /** Output the given value in the current window. */ public abstract void output(OutputT value); } diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnContextFactory.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnContextFactory.java index 0511c6f3b8e0..ae175a5f0117 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnContextFactory.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnContextFactory.java @@ -49,7 +49,6 @@ public interface OnTriggerCallbacks { } private final K key; - private final @Nullable Boolean draining; private final ReduceFn reduceFn; private final WindowingStrategy windowingStrategy; private final StateInternals stateInternals; @@ -66,9 +65,7 @@ public interface OnTriggerCallbacks { ActiveWindowSet activeWindows, TimerInternals timerInternals, @Nullable SideInputReader sideInputReader, - @Nullable Boolean draining, @Nullable PipelineOptions options) { - this.draining = draining; this.key = key; this.reduceFn = reduceFn; this.windowingStrategy = windowingStrategy; @@ -97,7 +94,12 @@ private StateAccessorImpl stateAccessor(W window, StateStyle style) { } public ReduceFn.Context base(W window, StateStyle style) { - return new ContextImpl(stateAccessor(window, style)); + return new ContextImpl(stateAccessor(window, style), null); + } + + public ReduceFn.Context base( + W window, StateStyle style, Boolean draining) { + return new ContextImpl(stateAccessor(window, style), draining); } public ReduceFn.ProcessValueContext forValue( @@ -106,8 +108,12 @@ public ReduceFn.ProcessValueContext forValue( } public ReduceFn.OnTriggerContext forTrigger( - W window, PaneInfo paneInfo, StateStyle style, OnTriggerCallbacks callbacks) { - return new OnTriggerContextImpl(stateAccessor(window, style), paneInfo, callbacks); + W window, + PaneInfo paneInfo, + StateStyle style, + OnTriggerCallbacks callbacks, + Boolean draining) { + return new OnTriggerContextImpl(stateAccessor(window, style), paneInfo, callbacks, draining); } public ReduceFn.OnMergeContext forMerge( @@ -321,11 +327,13 @@ public Map accessInEachMergingWindow( private class ContextImpl extends ReduceFn.Context { private final StateAccessorImpl state; private final TimersImpl timers; + private final Boolean draining; - private ContextImpl(StateAccessorImpl state) { + private ContextImpl(StateAccessorImpl state, @Nullable Boolean draining) { reduceFn.super(); this.state = state; this.timers = new TimersImpl(state.namespace()); + this.draining = draining; } @Override @@ -333,11 +341,6 @@ public K key() { return key; } - @Override - public Boolean draining() { - return draining; - } - @Override public W window() { return state.window(); @@ -423,14 +426,19 @@ private class OnTriggerContextImpl extends ReduceFn.OnTri private final PaneInfo paneInfo; private final OnTriggerCallbacks callbacks; private final TimersImpl timers; + private final Boolean draining; private OnTriggerContextImpl( - StateAccessorImpl state, PaneInfo paneInfo, OnTriggerCallbacks callbacks) { + StateAccessorImpl state, + PaneInfo paneInfo, + OnTriggerCallbacks callbacks, + Boolean draining) { reduceFn.super(); this.state = state; this.paneInfo = paneInfo; this.callbacks = callbacks; this.timers = new TimersImpl(state.namespace()); + this.draining = draining; } @Override @@ -489,11 +497,6 @@ public K key() { return key; } - @Override - public Boolean draining() { - return draining; - } - @Override public WindowingStrategy windowingStrategy() { return windowingStrategy; @@ -530,11 +533,6 @@ public K key() { return key; } - @Override - public Boolean draining() { - return draining; - } - @Override public WindowingStrategy windowingStrategy() { return windowingStrategy; diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java index 1e479515620e..528a62f46c0a 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java @@ -198,8 +198,6 @@ public class ReduceFnRunner { */ private final PaneInfoTracker paneInfoTracker; - private final @Nullable Boolean draining; - /** * Store whether we've seen any elements for a window since the last pane was emitted. * @@ -221,9 +219,7 @@ public ReduceFnRunner( WindowedValueReceiver> outputter, @Nullable SideInputReader sideInputReader, ReduceFn reduceFn, - @Nullable Boolean draining, @Nullable PipelineOptions options) { - this.draining = draining; this.key = key; this.timerInternals = timerInternals; this.paneInfoTracker = new PaneInfoTracker(timerInternals); @@ -252,7 +248,6 @@ public ReduceFnRunner( this.activeWindows, timerInternals, sideInputReader, - draining, options); this.watermarkHold = new WatermarkHold<>(timerInternals, windowingStrategy); @@ -379,7 +374,8 @@ public void processElements(Iterable> values) throws Excep for (W window : windowsToFire) { emit( contextFactory.base(window, StateStyle.DIRECT), - contextFactory.base(window, StateStyle.RENAMED)); + contextFactory.base(window, StateStyle.RENAMED), + null); } // We're all done with merging and emitting elements so can compress the activeWindow state. @@ -652,12 +648,15 @@ private class WindowActivation { // garbage collect the window. We'll consider any timer at or after the // end-of-window time to be a signal to garbage collect. public final boolean isGarbageCollection; + public final Boolean draining; WindowActivation( ReduceFn.Context directContext, - ReduceFn.Context renamedContext) { + ReduceFn.Context renamedContext, + Boolean draining) { this.directContext = directContext; this.renamedContext = renamedContext; + this.draining = draining; W window = directContext.window(); // The output watermark is before the end of the window if it is either unknown @@ -714,12 +713,13 @@ public void onTimers(Iterable timers) throws Exception { WindowTracing.debug( "ReduceFnRunner: Received timer key:{}; window:{}; data:{} with " - + "inputWatermark:{}; outputWatermark:{}", + + "inputWatermark:{}; outputWatermark:{}; draining:{}", key, window, timer, timerInternals.currentInputWatermarkTime(), - timerInternals.currentOutputWatermarkTime()); + timerInternals.currentOutputWatermarkTime(), + timer.getDraining()); // Processing time timers for an expired window are ignored, just like elements // that show up too late. Window GC is management by an event time timer @@ -740,7 +740,8 @@ public void onTimers(Iterable timers) throws Exception { contextFactory.base(window, StateStyle.DIRECT); ReduceFn.Context renamedContext = contextFactory.base(window, StateStyle.RENAMED); - WindowActivation windowActivation = new WindowActivation(directContext, renamedContext); + WindowActivation windowActivation = + new WindowActivation(directContext, renamedContext, timer.getDraining()); windowActivations.put(window, windowActivation); // Perform prefetching of state to determine if the trigger should fire. @@ -771,11 +772,12 @@ public void onTimers(Iterable timers) throws Exception { if (windowActivation.isGarbageCollection) { WindowTracing.debug( - "ReduceFnRunner: Cleaning up for key:{}; window:{} with inputWatermark:{}; outputWatermark:{}", + "ReduceFnRunner: Cleaning up for key:{}; window:{} with inputWatermark:{}; outputWatermark:{}; draining:{}", key, directContext.window(), timerInternals.currentInputWatermarkTime(), - timerInternals.currentOutputWatermarkTime()); + timerInternals.currentOutputWatermarkTime(), + windowActivation.draining); boolean windowIsActiveAndOpen = windowActivation.windowIsActiveAndOpen(); if (windowIsActiveAndOpen) { @@ -788,7 +790,8 @@ public void onTimers(Iterable timers) throws Exception { directContext, renamedContext, true /* isFinished */, - windowActivation.isEndOfWindow); + windowActivation.isEndOfWindow, + windowActivation.draining); checkState(newHold == null, "Hold placed at %s despite isFinished being true.", newHold); } @@ -806,7 +809,7 @@ public void onTimers(Iterable timers) throws Exception { if (windowActivation.windowIsActiveAndOpen() && triggerRunner.shouldFire( directContext.window(), directContext.timers(), directContext.state())) { - emit(directContext, renamedContext); + emit(directContext, renamedContext, windowActivation.draining); } if (windowActivation.isEndOfWindow) { @@ -929,7 +932,8 @@ private void prefetchEmit( /** Emit if a trigger is ready to fire or timers require it, and cleanup state. */ private void emit( ReduceFn.Context directContext, - ReduceFn.Context renamedContext) + ReduceFn.Context renamedContext, + Boolean draining) throws Exception { checkState( triggerRunner.shouldFire( @@ -945,7 +949,7 @@ private void emit( // Run onTrigger to produce the actual pane contents. // As a side effect it will clear all element holds, but not necessarily any // end-of-window or garbage collection holds. - onTrigger(directContext, renamedContext, isFinished, false /*isEndOfWindow*/); + onTrigger(directContext, renamedContext, isFinished, false /*isEndOfWindow*/, draining); // Now that we've triggered, the pane is empty. nonEmptyPanes.clearPane(renamedContext.state()); @@ -1003,7 +1007,8 @@ private void prefetchOnTrigger( final ReduceFn.Context directContext, ReduceFn.Context renamedContext, final boolean isFinished, - boolean isEndOfWindow) + boolean isEndOfWindow, + Boolean draining) throws Exception { // Extract the window hold, and as a side effect clear it. final WatermarkHold.OldAndNewHolds pair = @@ -1075,10 +1080,12 @@ private void prefetchOnTrigger( .setValue(KV.of(key, toOutput)) .setTimestamp(outputTimestamp) .setWindows(windows) + .setDraining(draining) .setPaneInfo(paneInfo) .setReceiver(outputter) .output(); - }); + }, + draining); reduceFn.onTrigger(renamedTriggerContext); } diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/WindowMatchers.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/WindowMatchers.java index 463cb9320237..fc945791a788 100644 --- a/runners/core-java/src/main/java/org/apache/beam/runners/core/WindowMatchers.java +++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/WindowMatchers.java @@ -59,7 +59,7 @@ public static Matcher> isWindowedValue( Matcher> windowsMatcher, Matcher paneInfoMatcher) { return new WindowedValueMatcher<>( - valueMatcher, timestampMatcher, windowsMatcher, paneInfoMatcher); + valueMatcher, timestampMatcher, windowsMatcher, paneInfoMatcher, Matchers.anything()); } public static Matcher> isWindowedValue( @@ -67,19 +67,27 @@ public static Matcher> isWindowedValue( Matcher timestampMatcher, Matcher> windowsMatcher) { return new WindowedValueMatcher<>( - valueMatcher, timestampMatcher, windowsMatcher, Matchers.anything()); + valueMatcher, timestampMatcher, windowsMatcher, Matchers.anything(), Matchers.anything()); } public static Matcher> isWindowedValue( Matcher valueMatcher, Matcher timestampMatcher) { return new WindowedValueMatcher<>( - valueMatcher, timestampMatcher, Matchers.anything(), Matchers.anything()); + valueMatcher, + timestampMatcher, + Matchers.anything(), + Matchers.anything(), + Matchers.anything()); } public static Matcher> isWindowedValue( Matcher valueMatcher) { return new WindowedValueMatcher<>( - valueMatcher, Matchers.anything(), Matchers.anything(), Matchers.anything()); + valueMatcher, + Matchers.anything(), + Matchers.anything(), + Matchers.anything(), + Matchers.anything()); } public static Matcher> isSingleWindowedValue( @@ -143,7 +151,25 @@ public static Matcher> isSingleWindowedValue( Matcher timestampMatcher, Matcher windowMatcher) { return new WindowedValueMatcher<>( - valueMatcher, timestampMatcher, Matchers.contains(windowMatcher), Matchers.anything()); + valueMatcher, + timestampMatcher, + Matchers.contains(windowMatcher), + Matchers.anything(), + Matchers.anything()); + } + + public static Matcher> isSingleWindowedValue( + Matcher valueMatcher, + Matcher timestampMatcher, + Matcher windowMatcher, + Matcher paneInfoMatcher, + Matcher drainMatcher) { + return new WindowedValueMatcher<>( + valueMatcher, + timestampMatcher, + Matchers.contains(windowMatcher), + paneInfoMatcher, + drainMatcher); } public static Matcher> isSingleWindowedValue( @@ -152,7 +178,11 @@ public static Matcher> isSingleWindowedValue( Matcher windowMatcher, Matcher paneInfoMatcher) { return new WindowedValueMatcher<>( - valueMatcher, timestampMatcher, Matchers.contains(windowMatcher), paneInfoMatcher); + valueMatcher, + timestampMatcher, + Matchers.contains(windowMatcher), + paneInfoMatcher, + Matchers.anything()); } public static Matcher intervalWindow(long start, long end) { @@ -203,16 +233,19 @@ private static class WindowedValueMatcher extends TypeSafeMatcher timestampMatcher; private Matcher> windowsMatcher; private Matcher paneInfoMatcher; + private Matcher drainMatcher; private WindowedValueMatcher( Matcher valueMatcher, Matcher timestampMatcher, Matcher> windowsMatcher, - Matcher paneInfoMatcher) { + Matcher paneInfoMatcher, + Matcher drainMatcher) { this.valueMatcher = valueMatcher; this.timestampMatcher = timestampMatcher; this.windowsMatcher = windowsMatcher; this.paneInfoMatcher = paneInfoMatcher; + this.drainMatcher = drainMatcher; } @Override @@ -226,6 +259,8 @@ public void describeTo(Description description) { .appendValue(windowsMatcher) .appendText(", ") .appendValue(paneInfoMatcher) + .appendText(", ") + .appendValue(drainMatcher) .appendText(")"); } @@ -234,7 +269,8 @@ protected boolean matchesSafely(WindowedValue windowedValue) { return valueMatcher.matches(windowedValue.getValue()) && timestampMatcher.matches(windowedValue.getTimestamp()) && windowsMatcher.matches(windowedValue.getWindows()) - && paneInfoMatcher.matches(windowedValue.getPaneInfo()); + && paneInfoMatcher.matches(windowedValue.getPaneInfo()) + && drainMatcher.matches(windowedValue.isDraining()); } } } diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java index b73d6a6e54b7..51ba557749e2 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java @@ -89,8 +89,6 @@ }) public class ReduceFnTester { private static final String KEY = "TEST_KEY"; - private static final Boolean DRAINING = null; - private final TestInMemoryStateInternals stateInternals = new TestInMemoryStateInternals<>(KEY); private final InMemoryTimerInternals timerInternals = new InMemoryTimerInternals(); @@ -303,7 +301,6 @@ ReduceFnRunner createRunner() { testOutputter, sideInputReader, reduceFn, - DRAINING, options); } diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java index 77387bf497b3..c6726fb3463f 100644 --- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java +++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java @@ -179,7 +179,6 @@ public void processElement(WindowedValue> element) throws Ex new BundleWindowedValueReceiver<>(bundle), new UnsupportedSideInputReader(DirectGroupAlsoByWindow.class.getSimpleName()), reduceFn, - element.isDraining(), options); // Drop any elements within expired windows diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java index 76dad7d6f727..8b1b2d08cd49 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java @@ -592,8 +592,7 @@ public void setTimer( Instant outputTimestamp, TimeDomain timeDomain) { setTimer( - TimerData.of( - timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain)); + TimerData.of(timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain)); } @Override diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowViaWindowSetFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowViaWindowSetFn.java index e412f28a7583..ec36644d1e68 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowViaWindowSetFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowViaWindowSetFn.java @@ -91,7 +91,6 @@ public void processElement( output, sideInputReader, reduceFn, - null, // elements and timers have relevant information about draining options); reduceFnRunner.processElements(keyedWorkItem.elementsIterable()); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BatchGroupAlsoByWindowViaOutputBufferFn.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BatchGroupAlsoByWindowViaOutputBufferFn.java index b0251c7bb4a8..12fc97ca9015 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BatchGroupAlsoByWindowViaOutputBufferFn.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BatchGroupAlsoByWindowViaOutputBufferFn.java @@ -88,7 +88,6 @@ public void processElement( output, NullSideInputReader.empty(), reduceFn, - null, options); // Process the elements. diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java index f7852ec1767d..0cf2dc69111e 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java @@ -19,10 +19,7 @@ import static org.apache.beam.sdk.TestUtils.KvMatcher.isKv; import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.containsInAnyOrder; -import static org.hamcrest.Matchers.emptyIterable; -import static org.hamcrest.Matchers.equalTo; -import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.*; import static org.mockito.Mockito.when; import java.io.IOException; @@ -197,6 +194,16 @@ private WindowedValue> createValue( KEY, workItem.build(), windowCoder, wildcardWindowsCoder, valueCoder, false)); } + private WindowedValue> createDrainingValue( + WorkItem.Builder workItem, Coder valueCoder) { + @SuppressWarnings({"unchecked", "rawtypes"}) + Coder> wildcardWindowsCoder = (Coder) windowsCoder; + return new ValueInEmptyWindows<>( + (KeyedWorkItem) + new WindmillKeyedWorkItem<>( + KEY, workItem.build(), windowCoder, wildcardWindowsCoder, valueCoder, true)); + } + @Test public void testFixedWindows() throws Exception { TupleTag>> outputTag = new TupleTag<>(); @@ -255,6 +262,78 @@ public void testFixedWindows() throws Exception { equalTo(window(10, 20))))); } + @Test + public void testFixedWindowsWithDraining() throws Exception { + TupleTag>> outputTag = new TupleTag<>(); + ListOutputManager outputManager = new ListOutputManager(); + DoFnRunner, KV>> runner = + makeRunner( + output -> outputManager.output(outputTag, output), + WindowingStrategy.of(FixedWindows.of(Duration.millis(10)))); + + when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(new Instant(0)); + + runner.startBundle(); + + WorkItem.Builder workItem1 = WorkItem.newBuilder(); + workItem1.setKey(ByteString.copyFromUtf8(KEY)); + workItem1.setWorkToken(WORK_TOKEN); + InputMessageBundle.Builder messageBundle = workItem1.addMessageBundlesBuilder(); + messageBundle.setSourceComputationId(SOURCE_COMPUTATION_ID); + + Coder valueCoder = StringUtf8Coder.of(); + addElement(messageBundle, Arrays.asList(window(0, 10)), new Instant(1), valueCoder, "v1"); + addElement(messageBundle, Arrays.asList(window(0, 10)), new Instant(2), valueCoder, "v2"); + addElement(messageBundle, Arrays.asList(window(0, 10)), new Instant(0), valueCoder, "v0"); + addElement(messageBundle, Arrays.asList(window(10, 20)), new Instant(13), valueCoder, "v3"); + runner.processElement(createValue(workItem1, valueCoder)); + + runner.finishBundle(); + runner.startBundle(); + + WorkItem.Builder workItem2 = WorkItem.newBuilder(); + workItem2.setKey(ByteString.copyFromUtf8(KEY)); + workItem2.setWorkToken(WORK_TOKEN); + InputMessageBundle.Builder messageBundle2 = workItem1.addMessageBundlesBuilder(); + messageBundle2.setSourceComputationId(SOURCE_COMPUTATION_ID); + addTimer(workItem2, window(0, 10), new Instant(9), Timer.Type.WATERMARK); + when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(new Instant(10)); + + runner.processElement(createValue(workItem2, valueCoder)); + runner.finishBundle(); + + runner.startBundle(); + + WorkItem.Builder workItem3 = WorkItem.newBuilder(); + workItem3.setKey(ByteString.copyFromUtf8(KEY)); + workItem3.setWorkToken(WORK_TOKEN); + addTimer(workItem3, window(10, 20), new Instant(19), Timer.Type.WATERMARK); + + when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(new Instant(20)); + runner.processElement(createDrainingValue(workItem3, valueCoder)); + runner.finishBundle(); + + List>>> result = outputManager.getOutput(outputTag); + + assertThat(result.size(), equalTo(2)); + + assertThat( + result, + containsInAnyOrder( + WindowMatchers.isSingleWindowedValue( + isKv(equalTo(KEY), containsInAnyOrder("v0", "v1", "v2")), + equalTo(window(0, 10).maxTimestamp()), + equalTo(window(0, 10)), + anything(), + is(Boolean.FALSE)), + WindowMatchers.isSingleWindowedValue( + isKv(equalTo(KEY), containsInAnyOrder("v3")), + equalTo(window(10, 20).maxTimestamp()), + equalTo(window(10, 20)), + anything(), + is(Boolean.TRUE)))); + } + @Test public void testSlidingWindows() throws Exception { TupleTag>> outputTag = new TupleTag<>(); diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java index ed51be43ec1d..cbda32cf1453 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java @@ -230,7 +230,8 @@ public void output() { } public WindowedValue build() { - return WindowedValues.of(getValue(), getTimestamp(), getWindows(), getPaneInfo()); + return WindowedValues.of( + getValue(), getTimestamp(), getWindows(), getPaneInfo(), null, null, isDraining()); } @Override @@ -240,6 +241,7 @@ public String toString() { .add("timestamp", getTimestamp()) .add("windows", getWindows()) .add("paneInfo", getPaneInfo()) + .add("draining", isDraining()) .add("receiver", receiver) .toString(); } @@ -263,7 +265,7 @@ public static WindowedValue of( checkArgument(windows.size() > 0, "WindowedValue requires windows, but there were none"); if (windows.size() == 1) { - return of(value, timestamp, windows.iterator().next(), paneInfo); + return of(value, timestamp, windows.iterator().next(), paneInfo, draining); } else { return new TimestampedValueInMultipleWindows<>( value, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset, draining); @@ -287,14 +289,26 @@ public static WindowedValue of( T value, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) { checkArgument(paneInfo != null, "WindowedValue requires PaneInfo, but it was null"); + return of(value, timestamp, window, paneInfo, null); + } + + /** Returns a {@code WindowedValue} with the given value, timestamp, and window. */ + public static WindowedValue of( + T value, + Instant timestamp, + BoundedWindow window, + PaneInfo paneInfo, + @Nullable Boolean draining) { + checkArgument(paneInfo != null, "WindowedValue requires PaneInfo, but it was null"); + boolean isGlobal = GlobalWindow.INSTANCE.equals(window); if (isGlobal && BoundedWindow.TIMESTAMP_MIN_VALUE.equals(timestamp)) { return valueInGlobalWindow(value, paneInfo); } else if (isGlobal) { - return new TimestampedValueInGlobalWindow<>(value, timestamp, paneInfo, null, null, null); + return new TimestampedValueInGlobalWindow<>(value, timestamp, paneInfo, null, null, draining); } else { return new TimestampedValueInSingleWindow<>( - value, timestamp, window, paneInfo, null, null, null); + value, timestamp, window, paneInfo, null, null, draining); } } @@ -540,6 +554,7 @@ public String toString() { return MoreObjects.toStringHelper(getClass()) .add("value", getValue()) .add("paneInfo", getPaneInfo()) + .add("draining", isDraining()) .toString(); } } @@ -625,6 +640,7 @@ public String toString() { .add("value", getValue()) .add("timestamp", getTimestamp()) .add("paneInfo", getPaneInfo()) + .add("draining", isDraining()) .toString(); } } @@ -701,6 +717,7 @@ public String toString() { .add("timestamp", getTimestamp()) .add("window", window) .add("paneInfo", getPaneInfo()) + .add("draining", isDraining()) .toString(); } } @@ -773,6 +790,7 @@ public String toString() { .add("timestamp", getTimestamp()) .add("windows", windows) .add("paneInfo", getPaneInfo()) + .add("draining", isDraining()) .toString(); } From f3f61267d4a0e5b388f36f03c31ec24e4f0cd95f Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 10 Oct 2025 18:15:21 +0200 Subject: [PATCH 20/26] compile spark --- .../beam/runners/spark/translation/AbstractInOutIterator.java | 2 +- .../java/org/apache/beam/runners/spark/util/TimerUtils.java | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/AbstractInOutIterator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/AbstractInOutIterator.java index 9d73a605b3b5..d74a750487cb 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/AbstractInOutIterator.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/AbstractInOutIterator.java @@ -75,7 +75,7 @@ public void fireTimer(TimerInternals.TimerData timer) { window, timer.getTimestamp(), timer.getOutputTimestamp(), - timer.getDomain()); + timer.getDomain(), null); } finally { if (this.ctx.getTimerDataIterator() instanceof ParDoStateUpdateFn.SparkTimerInternalsIterator) { diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/TimerUtils.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/util/TimerUtils.java index 03735355de51..aa0b563ebfa1 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/TimerUtils.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/util/TimerUtils.java @@ -115,6 +115,9 @@ public PaneInfo getPaneInfo() { return null; } + @Override + public @Nullable Boolean isDraining() { return null; } + @Override public @Nullable Long getRecordOffset() { return null; From 424d0cca43a2173699efe65dd8efe43115f863c7 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 10 Oct 2025 18:22:20 +0200 Subject: [PATCH 21/26] compile errors --- .../translation/functions/FlinkStatefulDoFnFunction.java | 3 ++- .../wrappers/streaming/ExecutableStageDoFnOperatorTest.java | 6 ++++-- .../java/org/apache/beam/runners/samza/runtime/DoFnOp.java | 3 ++- .../apache/beam/runners/samza/runtime/PortableDoFnOp.java | 3 ++- 4 files changed, 10 insertions(+), 5 deletions(-) diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkStatefulDoFnFunction.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkStatefulDoFnFunction.java index 3aa5a0802b38..a387633f6ccd 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkStatefulDoFnFunction.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkStatefulDoFnFunction.java @@ -250,7 +250,8 @@ private void fireTimer( window, timer.getTimestamp(), timer.getOutputTimestamp(), - timer.getDomain()); + timer.getDomain(), + null); } @Override diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperatorTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperatorTest.java index 982871e59f14..8447fae34178 100644 --- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperatorTest.java +++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperatorTest.java @@ -534,7 +534,8 @@ public void testWatermarkHandling() throws Exception { windowedValue.getWindows(), timestamp, timestamp, - PaneInfo.NO_FIRING), + PaneInfo.NO_FIRING, + null), TimerInternals.TimerData.of( "", TimerReceiverFactory.encodeToTimerDataTimerId("transform", timerId), @@ -875,7 +876,8 @@ private void testEnsureDeferredStateCleanupTimerFiring(boolean withCheckpointing windowedValue.getWindows(), window.maxTimestamp(), window.maxTimestamp(), - PaneInfo.NO_FIRING), + PaneInfo.NO_FIRING, + null), userTimer2); assertThat(testHarness.numEventTimeTimers(), is(1)); diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java index 2f27e31d05dd..45c2ec4bde89 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/DoFnOp.java @@ -437,7 +437,8 @@ private void fireTimer(KeyedTimerData keyedTimerData) { window, timer.getTimestamp(), timer.getOutputTimestamp(), - timer.getDomain()); + timer.getDomain(), + null); } // todo: should this go through bundle manager to start and finish the bundle? diff --git a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/PortableDoFnOp.java b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/PortableDoFnOp.java index 743a42d14791..b7f6726f81e8 100644 --- a/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/PortableDoFnOp.java +++ b/runners/samza/src/main/java/org/apache/beam/runners/samza/runtime/PortableDoFnOp.java @@ -425,7 +425,8 @@ private void fireTimer(KeyedTimerData keyedTimerData) { window, timer.getTimestamp(), timer.getOutputTimestamp(), - timer.getDomain()); + timer.getDomain(), + null); } // todo: should this go through bundle manager to start and finish the bundle? From 652cff61b219275fface01683cbe1254e94d46d5 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 10 Oct 2025 18:36:24 +0200 Subject: [PATCH 22/26] compile errors --- .../runners/spark/translation/AbstractInOutIteratorTest.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/AbstractInOutIteratorTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/AbstractInOutIteratorTest.java index b1e0e3a39808..0a563f669d66 100644 --- a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/AbstractInOutIteratorTest.java +++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/AbstractInOutIteratorTest.java @@ -107,7 +107,7 @@ public void testFireTimer() { mockWindow, TEST_TIMESTAMP, TEST_OUTPUT_TIMESTAMP, - TEST_TIME_DOMAIN); + TEST_TIME_DOMAIN, null); // Verify that timer data iterator deletion was not called (no timer iterator was set in this // test) @@ -133,7 +133,7 @@ public void testFireTimerWithTimerDataIterator() { mockWindow, TEST_TIMESTAMP, TEST_OUTPUT_TIMESTAMP, - TEST_TIME_DOMAIN); + TEST_TIME_DOMAIN,null); // Verify that the timer data iterator's deleteTimer method was called verify(mockTimerDataIterator).deleteTimer(mockTimer); From b0f4a11d6b6a4c817427bb82d54e395349ceda26 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 10 Oct 2025 18:58:53 +0200 Subject: [PATCH 23/26] spotless --- .../runners/spark/translation/AbstractInOutIterator.java | 3 ++- .../java/org/apache/beam/runners/spark/util/TimerUtils.java | 4 +++- .../spark/translation/AbstractInOutIteratorTest.java | 6 ++++-- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/AbstractInOutIterator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/AbstractInOutIterator.java index d74a750487cb..6008ceb47bcc 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/AbstractInOutIterator.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/AbstractInOutIterator.java @@ -75,7 +75,8 @@ public void fireTimer(TimerInternals.TimerData timer) { window, timer.getTimestamp(), timer.getOutputTimestamp(), - timer.getDomain(), null); + timer.getDomain(), + null); } finally { if (this.ctx.getTimerDataIterator() instanceof ParDoStateUpdateFn.SparkTimerInternalsIterator) { diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/TimerUtils.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/util/TimerUtils.java index aa0b563ebfa1..0760c1aeb649 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/TimerUtils.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/util/TimerUtils.java @@ -116,7 +116,9 @@ public PaneInfo getPaneInfo() { } @Override - public @Nullable Boolean isDraining() { return null; } + public @Nullable Boolean isDraining() { + return null; + } @Override public @Nullable Long getRecordOffset() { diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/AbstractInOutIteratorTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/AbstractInOutIteratorTest.java index 0a563f669d66..e6da459e7c71 100644 --- a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/AbstractInOutIteratorTest.java +++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/AbstractInOutIteratorTest.java @@ -107,7 +107,8 @@ public void testFireTimer() { mockWindow, TEST_TIMESTAMP, TEST_OUTPUT_TIMESTAMP, - TEST_TIME_DOMAIN, null); + TEST_TIME_DOMAIN, + null); // Verify that timer data iterator deletion was not called (no timer iterator was set in this // test) @@ -133,7 +134,8 @@ public void testFireTimerWithTimerDataIterator() { mockWindow, TEST_TIMESTAMP, TEST_OUTPUT_TIMESTAMP, - TEST_TIME_DOMAIN,null); + TEST_TIME_DOMAIN, + null); // Verify that the timer data iterator's deleteTimer method was called verify(mockTimerDataIterator).deleteTimer(mockTimer); From 572e6ef8925dda555e1ebf6d00d37063d99bb29b Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 10 Oct 2025 19:12:23 +0200 Subject: [PATCH 24/26] spotless --- .../org/apache/beam/runners/jet/processors/StatefulParDoP.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/runners/jet/src/main/java/org/apache/beam/runners/jet/processors/StatefulParDoP.java b/runners/jet/src/main/java/org/apache/beam/runners/jet/processors/StatefulParDoP.java index f696873ec5c0..24f298c7e783 100644 --- a/runners/jet/src/main/java/org/apache/beam/runners/jet/processors/StatefulParDoP.java +++ b/runners/jet/src/main/java/org/apache/beam/runners/jet/processors/StatefulParDoP.java @@ -105,7 +105,8 @@ private static void fireTimer( window, timer.getTimestamp(), timer.getOutputTimestamp(), - timer.getDomain()); + timer.getDomain(), + null); } @Override From 9ede59ac5a34b9f753725cd0a465a543310443f9 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 10 Oct 2025 19:34:47 +0200 Subject: [PATCH 25/26] spotless --- .../beam/runners/fnexecution/control/RemoteExecutionTest.java | 3 ++- .../apache/beam/runners/fnexecution/wire/CommonCoderTest.java | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java index 556cc7993ae1..d6477095bf96 100644 --- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java +++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java @@ -2254,6 +2254,7 @@ private org.apache.beam.sdk.util.construction.Timer timerForTest( Collections.singletonList(GlobalWindow.INSTANCE), BoundedWindow.TIMESTAMP_MIN_VALUE.plus(Duration.millis(fireTimestamp)), BoundedWindow.TIMESTAMP_MIN_VALUE.plus(Duration.millis(holdTimestamp)), - PaneInfo.NO_FIRING); + PaneInfo.NO_FIRING, + null); } } diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/wire/CommonCoderTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/wire/CommonCoderTest.java index eccf1e66434e..839dc898d3c1 100644 --- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/wire/CommonCoderTest.java +++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/wire/CommonCoderTest.java @@ -326,7 +326,8 @@ private static Object convertValue(Object value, CommonCoder coderSpec, Coder co windows, new Instant(((Number) kvMap.get("fireTimestamp")).longValue()), new Instant(((Number) kvMap.get("holdTimestamp")).longValue()), - paneInfo); + paneInfo, + null); } else if (s.equals(getUrn(StandardCoders.Enum.INTERVAL_WINDOW))) { Map kvMap = (Map) value; Instant end = new Instant(((Number) kvMap.get("end")).longValue()); From eff6fe0cfe5e603ef8877488222e455d8b37a07d Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Fri, 10 Oct 2025 21:10:29 +0200 Subject: [PATCH 26/26] encode draining --- .../beam/sdk/values/WindowedValues.java | 32 +++++++++++++++---- 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java index cbda32cf1453..1f458a4a7729 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java @@ -275,12 +275,16 @@ public static WindowedValue of( /** @deprecated for use only in compatibility with old broken code */ @Deprecated static WindowedValue createWithoutValidation( - T value, Instant timestamp, Collection windows, PaneInfo paneInfo) { + T value, + Instant timestamp, + Collection windows, + PaneInfo paneInfo, + @Nullable Boolean draining) { if (windows.size() == 1) { - return of(value, timestamp, windows.iterator().next(), paneInfo); + return of(value, timestamp, windows.iterator().next(), paneInfo, draining); } else { return new TimestampedValueInMultipleWindows<>( - value, timestamp, windows, paneInfo, null, null, null); + value, timestamp, windows, paneInfo, null, null, draining); } } @@ -906,7 +910,16 @@ public void encode(WindowedValue windowedElem, OutputStream outStream, Contex if (metadataSupported) { BeamFnApi.Elements.ElementMetadata.Builder builder = BeamFnApi.Elements.ElementMetadata.newBuilder(); - BeamFnApi.Elements.ElementMetadata em = builder.build(); + BeamFnApi.Elements.ElementMetadata em = + builder + .setDrain( + windowedElem.isDraining() != null + ? (Boolean.TRUE.equals(windowedElem.isDraining()) + ? BeamFnApi.Elements.DrainMode.Enum.DRAINING + : BeamFnApi.Elements.DrainMode.Enum.NOT_DRAINING) + : BeamFnApi.Elements.DrainMode.Enum.UNSPECIFIED) + .build(); + ByteArrayCoder.of().encode(em.toByteArray(), outStream); } valueCoder.encode(windowedElem.getValue(), outStream, context); @@ -923,14 +936,21 @@ public WindowedValue decode(InputStream inStream, Context context) Instant timestamp = InstantCoder.of().decode(inStream); Collection windows = windowsCoder.decode(inStream); PaneInfo paneInfo = PaneInfoCoder.INSTANCE.decode(inStream); + Boolean draining = null; if (isMetadataSupported() && paneInfo.isElementMetadata()) { - BeamFnApi.Elements.ElementMetadata.parseFrom(ByteArrayCoder.of().decode(inStream)); + BeamFnApi.Elements.ElementMetadata elementMetadata = + BeamFnApi.Elements.ElementMetadata.parseFrom(ByteArrayCoder.of().decode(inStream)); + boolean b = elementMetadata.hasDrain(); + draining = + b + ? elementMetadata.getDrain().equals(BeamFnApi.Elements.DrainMode.Enum.DRAINING) + : null; } T value = valueCoder.decode(inStream, context); // Because there are some remaining (incorrect) uses of WindowedValue with no windows, // we call this deprecated no-validation path when decoding - return WindowedValues.createWithoutValidation(value, timestamp, windows, paneInfo); + return WindowedValues.createWithoutValidation(value, timestamp, windows, paneInfo, draining); } @Override