scwhittle commented on code in PR #39961:
URL: https://github.com/apache/beam/pull/39961#discussion_r4013574396
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/IsmSink.java:
##########
@@ -295,6 +296,9 @@ private void finish() throws IOException {
.encode(Footer.of(startOfIndex, startOfBloomFilter,
numberOfKeysWritten), out);
}
+ @Override
+ public void finishKey(@Nullable Object key) throws IOException {}
Review Comment:
should we add a default no-op impl to the Sink base class? Since it's not
called in batch anyway, forcing a lot of these batch sinks to implement it
seems like unneeded clutter.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubDynamicSink.java:
##########
@@ -127,32 +134,52 @@ public long add(WindowedValue<PubsubMessage> data) throws
IOException {
!dataTopic.isEmpty(), "No topic set for message when using dynamic
topics.");
ByteString byteString = getDataFromMessage(data.getValue(), stream);
Windmill.PubSubMessageBundle.Builder builder =
- outputBuilders.computeIfAbsent(
- dataTopic,
- topic ->
- context
- .getOutputBuilder()
- .addPubsubMessagesBuilder()
- .setTopic(topic)
- .setTimestampLabel(timestampLabel)
- .setIdLabel(idLabel)
- .setWithAttributes(true));
+ outputBuilders.computeIfAbsent(dataTopic, this::createOutputBuilder);
builder.addMessages(
Windmill.Message.newBuilder()
.setData(byteString)
.setTimestamp(WindmillTimeUtils.harnessToWindmillTimestamp(data.getTimestamp()))
.build());
+ bufferedBytes += byteString.size();
+
return byteString.size();
}
+ private void flush(boolean bundleLevel) {
+ try {
+ for (Windmill.PubSubMessageBundle.Builder builder :
outputBuilders.values()) {
+ if (builder.getMessagesCount() > 0) {
+ Windmill.PubSubMessageBundle pubsubMessages = builder.build();
+ if (bundleLevel) {
+ context.addBundlePubsubMessages(pubsubMessages);
+ } else {
+ context.getOutputBuilder().addPubsubMessages(pubsubMessages);
+ }
+ }
+ }
+ } finally {
+ outputBuilders.clear();
+ bufferedBytes = 0;
+ }
+ }
+
+ @Override
+ public void finishKey(@Nullable Object key) throws IOException {
+ if (context.multiKeyBundleEnabled() && bufferedBytes >=
MAX_PUBSUB_BUNDLE_BYTES) {
Review Comment:
ditto about ordered keys
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java:
##########
@@ -321,6 +323,8 @@ public void reset() {
// these lists and maps are returned to callers after processing
Review Comment:
see other comment, maybe we can avoid setting to emptyList and then
reallocating if unused
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubDynamicSink.java:
##########
@@ -127,32 +134,52 @@ public long add(WindowedValue<PubsubMessage> data) throws
IOException {
!dataTopic.isEmpty(), "No topic set for message when using dynamic
topics.");
ByteString byteString = getDataFromMessage(data.getValue(), stream);
Windmill.PubSubMessageBundle.Builder builder =
- outputBuilders.computeIfAbsent(
- dataTopic,
- topic ->
- context
- .getOutputBuilder()
- .addPubsubMessagesBuilder()
- .setTopic(topic)
- .setTimestampLabel(timestampLabel)
- .setIdLabel(idLabel)
- .setWithAttributes(true));
+ outputBuilders.computeIfAbsent(dataTopic, this::createOutputBuilder);
builder.addMessages(
Windmill.Message.newBuilder()
.setData(byteString)
.setTimestamp(WindmillTimeUtils.harnessToWindmillTimestamp(data.getTimestamp()))
.build());
+ bufferedBytes += byteString.size();
+
return byteString.size();
}
+ private void flush(boolean bundleLevel) {
+ try {
+ for (Windmill.PubSubMessageBundle.Builder builder :
outputBuilders.values()) {
+ if (builder.getMessagesCount() > 0) {
+ Windmill.PubSubMessageBundle pubsubMessages = builder.build();
+ if (bundleLevel) {
+ context.addBundlePubsubMessages(pubsubMessages);
+ } else {
+ context.getOutputBuilder().addPubsubMessages(pubsubMessages);
+ }
+ }
+ }
+ } finally {
+ outputBuilders.clear();
+ bufferedBytes = 0;
+ }
+ }
+
+ @Override
+ public void finishKey(@Nullable Object key) throws IOException {
+ if (context.multiKeyBundleEnabled() && bufferedBytes >=
MAX_PUBSUB_BUNDLE_BYTES) {
+ flush(/* bundleLevel= */ false);
Review Comment:
why not flush at bundle level? It seems odd to attach output made from
processing key A to key B later if the buffered size was not exceeded on A but
later was on B.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java:
##########
Review Comment:
what about changing context.flushState() to return ExecuteWorkResult and
reset as needed?
Then we can get rid of the accessors just used for the flushing and we can
be smarter about how we reset to avoid allocations.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubSink.java:
##########
@@ -183,22 +190,44 @@ public long add(WindowedValue<T> data) throws IOException
{
.setData(byteString)
.setTimestamp(WindmillTimeUtils.harnessToWindmillTimestamp(data.getTimestamp()))
.build());
+ bufferedBytes += byteString.size();
return byteString.size();
}
+ private void flush(boolean bundleLevel) {
+ try {
+ Windmill.PubSubMessageBundle pubsubMessages = outputBuilder.build();
+ if (pubsubMessages.getMessagesCount() > 0) {
+ if (bundleLevel) {
+ context.addBundlePubsubMessages(pubsubMessages);
+ } else {
+ context.getOutputBuilder().addPubsubMessages(pubsubMessages);
+ }
+ }
+ } finally {
+ outputBuilder = createOutputBuilder();
+ bufferedBytes = 0;
+ }
+ }
+
@Override
- public void close() throws IOException {
- Windmill.PubSubMessageBundle pubsubMessages = outputBuilder.build();
- if (pubsubMessages.getMessagesCount() > 0) {
- context.getOutputBuilder().addPubsubMessages(pubsubMessages);
+ public void finishKey(@Nullable Object key) throws IOException {
+ if (context.multiKeyBundleEnabled() && bufferedBytes >=
MAX_PUBSUB_BUNDLE_BYTES) {
Review Comment:
this will need to be per-key if we change windmill pubsub to support
ordering keys. Maybe you could put a comment here and in the `add` method so
we don't forget?
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubSink.java:
##########
@@ -187,18 +190,41 @@ public long add(WindowedValue<T> data) throws IOException
{
return byteString.size();
}
+ private void flush(boolean bundleLevel) {
+ try {
+ Windmill.PubSubMessageBundle pubsubMessages = outputBuilder.build();
+ if (pubsubMessages.getMessagesCount() > 0) {
+ if (bundleLevel) {
+ context.addBundlePubsubMessages(pubsubMessages);
+ } else {
+ context.getOutputBuilder().addPubsubMessages(pubsubMessages);
+ }
+ }
+ } finally {
+ outputBuilder = createOutputBuilder();
+ }
+ }
+
+ @Override
+ public void finishKey(@Nullable Object key) throws IOException {
+ if (context.multiKeyBundleEnabled()) {
+ flush(/* bundleLevel= */ false);
Review Comment:
why not flush at bundle level? seems odd to attach outputs made from Key A
to key B if limit is only exceeded when reaching key B.
--
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]
