Skip to content

External metadata for streaming runner v1 changes #36373

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Merged
liferoad merged 9 commits into from
Oct 13, 2025
Merged

External metadata for streaming runner v1 changes #36373

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
4d87d3c
encode empty metadata within WindowedValue
stankiewicz Sep 9, 2025
7cfd6f8
add todos for future otel metadata
stankiewicz Oct 3, 2025
5b7e0fc
add todos for future otel metadata
stankiewicz Oct 3, 2025
7654478
spotless
stankiewicz Oct 3, 2025
0274977
fix tests
stankiewicz Oct 3, 2025
bbb9474
fix
stankiewicz Oct 3, 2025
99edc76
fix
stankiewicz Oct 3, 2025
1e38bec
fix bugs
stankiewicz Oct 9, 2025
a4054cd
fix bugs
stankiewicz Oct 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions ...java/src/main/java/org/apache/beam/runners/dataflow/RedistributeByKeyOverrideFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,7 @@ public Duration getAllowedTimestampSkew() {
public void processElement(
@Element KV<K, ValueInSingleWindow<V>> kv,
OutputReceiver<KV<K, V>> outputReceiver) {
// todo #33176 specify additional metadata in the future

@kennknowles kennknowles Oct 3, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yea I think we'll want to find a way to update Redistribute to automatically carry along "all" metadata. We built a separate data type for Redistribute to avoid making WindowedValue public. But now that it is public and everything around it is public (or public enough) we should make it possible to be resilient to new metadata being added without having to edit Redistribute or the overrides.

@stankiewicz stankiewicz Oct 3, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestions? :)

outputReceiver
.builder(KV.of(kv.getKey(), kv.getValue().getValue()))
.setTimestamp(kv.getValue().getTimestamp())
Expand Down
5 changes: 5 additions & 0 deletions ...worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);
Expand Down
5 changes: 5 additions & 0 deletions ...worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,9 @@ protected WindowedValue<T> decodeMessage(Windmill.Message message) throws IOExce
Collection<? extends BoundedWindow> windows =
WindmillSink.decodeMetadataWindows(windowsCoder, message.getMetadata());
PaneInfo paneInfo = WindmillSink.decodeMetadataPane(message.getMetadata());
if (WindowedValues.WindowedValueCoder.isMetadataSupported()) {
WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata());
}
if (valueCoder instanceof KvCoder) {
KvCoder<?, ?> kvCoder = (KvCoder<?, ?>) valueCoder;
InputStream key = context.getSerializedKey().newInput();
Expand All @@ -125,9 +128,11 @@ protected WindowedValue<T> 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);
}
}
Expand Down
5 changes: 4 additions & 1 deletion ...a/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,9 +108,12 @@ public Iterable<WindowedValue<ElemT>> elementsIterable() {
Collection<? extends BoundedWindow> windows =
WindmillSink.decodeMetadataWindows(windowsCoder, message.getMetadata());
PaneInfo paneInfo = WindmillSink.decodeMetadataPane(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);
Expand Down
46 changes: 39 additions & 7 deletions ...aflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,9 +26,11 @@
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;
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;
Expand All @@ -40,6 +42,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;
Expand Down Expand Up @@ -75,11 +78,20 @@ private static ByteString encodeMetadata(
ByteStringOutputStream stream,
Coder<Collection<? extends BoundedWindow>> windowsCoder,
Collection<? extends BoundedWindow> windows,
PaneInfo paneInfo)
PaneInfo paneInfo,
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.withElementMetadata(true), 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();
Expand All @@ -90,23 +102,39 @@ private static ByteString encodeMetadata(
public static ByteString encodeMetadata(
Coder<Collection<? extends BoundedWindow>> windowsCoder,
Collection<? extends BoundedWindow> 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 {
InputStream inStream = metadata.newInput();
return PaneInfoCoder.INSTANCE.decode(inStream);
}

public static BeamFnApi.Elements.ElementMetadata decodeAdditionalMetadata(
Coder<Collection<? extends BoundedWindow>> windowsCoder, ByteString metadata)
throws IOException {
InputStream inStream = metadata.newInput();
PaneInfo paneInfo = PaneInfoCoder.INSTANCE.decode(inStream);
windowsCoder.decode(inStream);
if (paneInfo.isElementMetadata()) {
return BeamFnApi.Elements.ElementMetadata.parseFrom(
ByteArrayCoder.of().decode(inStream, Coder.Context.OUTER));
} else {
// empty
return BeamFnApi.Elements.ElementMetadata.newBuilder().build();
}
}

public static Collection<? extends BoundedWindow> decodeMetadataWindows(
Coder<Collection<? extends BoundedWindow>> 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. */
Expand Down Expand Up @@ -184,8 +212,12 @@ private <EncodeT> ByteString encode(Coder<EncodeT> coder, EncodeT object) throws
public long add(WindowedValue<T> data) throws IOException {
ByteString key, value;
ByteString id = ByteString.EMPTY;
// todo #33176 specify additional metadata in the future
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();
Expand Down
8 changes: 7 additions & 1 deletion .../test/java/org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowFnsTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -176,7 +177,12 @@ private <V> 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));
}
Expand Down
8 changes: 7 additions & 1 deletion ...org/apache/beam/runners/dataflow/worker/StreamingGroupAlsoByWindowsReshuffleDoFnTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -114,7 +115,12 @@ private <V> 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));
}
Expand Down
9 changes: 7 additions & 2 deletions ...rker/src/test/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItemTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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)))
Expand Down
1 change: 1 addition & 0 deletions sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,6 +179,7 @@ public Duration getAllowedTimestampSkew() {
public void processElement(
@Element KV<K, ValueInSingleWindow<V>> kv,
OutputReceiver<KV<K, V>> outputReceiver) {
// todo #33176 specify additional metadata in the future
outputReceiver
.builder(KV.of(kv.getKey(), kv.getValue().getValue()))
.setTimestamp(kv.getValue().getTimestamp())
Expand Down
1 change: 1 addition & 0 deletions sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reify.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,7 @@ public PCollection<KV<K, ValueInSingleWindow<V>>> expand(PCollection<KV<K, V>> i
KvCoder<K, V> coder = (KvCoder<K, V>) input.getCoder();
return input
.apply(
// todo #33176 specify additional metadata in the future
ParDo.of(
new DoFn<KV<K, V>, KV<K, ValueInSingleWindow<V>>>() {
@ProcessElement
Expand Down
1 change: 1 addition & 0 deletions sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reshuffle.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,7 @@ public Duration getAllowedTimestampSkew() {
public void processElement(
@Element KV<K, ValueInSingleWindow<V>> kv,
OutputReceiver<KV<K, V>> outputReceiver) {
// todo #33176 specify additional metadata in the future
outputReceiver
.builder(KV.of(kv.getKey(), kv.getValue().getValue()))
.setTimestamp(kv.getValue().getTimestamp())
Expand Down
Loading
Loading