arunpandianp commented on code in PR #39961:
URL: https://github.com/apache/beam/pull/39961#discussion_r4045652337
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java:
##########
@@ -860,23 +886,45 @@ private void startForNewKey(Work newWork) throws
CoderException {
}
}
- // Returns state bytes read during the bundle execution
- public long getStateBytesRead() {
- return stateBytesRead;
+ public void addBundleOutputMessages(Windmill.OutputMessageBundle
outputBundle) {
+ getOrCreateBundleOutputMessages().add(outputBundle);
+ }
+
+ public void addBundlePubsubMessages(Windmill.PubSubMessageBundle
pubsubBundle) {
+ getOrCreateBundlePubsubMessages().add(pubsubBundle);
+ }
+
+ private List<Windmill.OutputMessageBundle> getOrCreateBundleOutputMessages()
{
+ if (this.bundleOutputMessages == null) {
+ this.bundleOutputMessages = new ArrayList<>();
+ }
+ return this.bundleOutputMessages;
+ }
+
+ private List<Windmill.PubSubMessageBundle> getOrCreateBundlePubsubMessages()
{
+ if (this.bundlePubsubMessages == null) {
+ this.bundlePubsubMessages = new ArrayList<>();
+ }
+ return this.bundlePubsubMessages;
+ }
+
+ private List<Windmill.WorkItemCommitRequest> getOrCreateWorkItemCommits() {
Review Comment:
done.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java:
##########
@@ -350,26 +350,43 @@ public long add(WindowedValue<T> data) throws IOException
{
return (long) key.size() + value.size() + metadata.size() + id.size() +
offsetSize;
}
- @Override
- public void close() throws IOException {
+ private void flush(boolean bundleLevel) {
try {
outputBuilder.setDestinationStreamId(destinationName);
for (Windmill.KeyedMessageBundle.Builder keyedOutput :
productionMap.values()) {
outputBuilder.addBundles(keyedOutput.build());
}
if (outputBuilder.getBundlesCount() > 0) {
- context.getOutputBuilder().addOutputMessages(outputBuilder.build());
+ Windmill.OutputMessageBundle bundle = outputBuilder.build();
+ if (bundleLevel) {
+ context.addBundleOutputMessages(bundle);
+ } else {
+ context.getOutputBuilder().addOutputMessages(bundle);
Review Comment:
done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
