Skip to content

Commit 4289ea0

Browse files
kennknowleskennknowles
committed
Add OutputBuilder to the Java SDK and use in runners
1 parent b26ecf1 commit 4289ea0

38 files changed

Lines changed: 979 additions & 570 deletions

File tree

‎buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1190,6 +1190,7 @@ class BeamModulePlugin implements Plugin<Project> {
11901190

11911191
List<String> skipDefRegexes = []
11921192
skipDefRegexes << "AutoValue_.*"
1193+
skipDefRegexes << "AutoBuilder_.*"
11931194
skipDefRegexes << "AutoOneOf_.*"
11941195
skipDefRegexes << ".*\\.jmh_generated\\..*"
11951196
skipDefRegexes += configuration.generatedClassPatterns
@@ -1283,7 +1284,8 @@ class BeamModulePlugin implements Plugin<Project> {
12831284
'**/org/apache/beam/gradle/**',
12841285
'**/org/apache/beam/model/**',
12851286
'**/org/apache/beam/runners/dataflow/worker/windmill/**',
1286-
'**/AutoValue_*'
1287+
'**/AutoValue_*',
1288+
'**/AutoBuilder_*',
12871289
]
12881290

12891291
def jacocoEnabled = project.hasProperty('enableJacocoReport')

‎runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataUtils.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,9 @@ public static <K, V> Iterable<WindowedValue<V>> dropExpiredWindows(
8181
if (input == null) {
8282
return null;
8383
}
84-
return input.explodeWindows();
84+
// The generics in this chain of calls line up best if we drop the covariance
85+
// in the return value of explodeWindows()
86+
return (Iterable<WindowedValue<V>>) input.explodeWindows();
8587
})
8688
.filter(
8789
input -> {

‎runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1057,8 +1057,13 @@ private void prefetchOnTrigger(
10571057
}
10581058

10591059
// Output the actual value.
1060-
outputter.output(
1061-
WindowedValues.of(KV.of(key, toOutput), outputTimestamp, windows, paneInfo));
1060+
WindowedValues.<KV<K, OutputT>>builder()
1061+
.setValue(KV.of(key, toOutput))
1062+
.setTimestamp(outputTimestamp)
1063+
.setWindows(windows)
1064+
.setPaneInfo(paneInfo)
1065+
.setReceiver(outputter)
1066+
.output();
10621067
});
10631068

10641069
reduceFn.onTrigger(renamedTriggerContext);

‎runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java‎

Lines changed: 36 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -447,8 +447,13 @@ public <T> void outputWindowedValue(
447447
Instant timestamp,
448448
Collection<? extends BoundedWindow> windows,
449449
PaneInfo paneInfo) {
450-
SimpleDoFnRunner.this.outputWindowedValue(
451-
tag, WindowedValues.of(output, timestamp, windows, paneInfo));
450+
WindowedValues.builder(elem)
451+
.withValue(output)
452+
.setTimestamp(timestamp)
453+
.setWindows(windows)
454+
.setPaneInfo(paneInfo)
455+
.setReceiver(wv -> SimpleDoFnRunner.this.outputWindowedValue(tag, wv))
456+
.output();
452457
}
453458

454459
@Override
@@ -638,6 +643,7 @@ private class OnTimerArgumentProvider<KeyT> extends DoFn<InputT, OutputT>.OnTime
638643
private final TimeDomain timeDomain;
639644
private final String timerId;
640645
private final KeyT key;
646+
private final WindowedValues.Builder<Void> outputTemplate;
641647

642648
/** Lazily initialized; should only be accessed via {@link #getNamespace()}. */
643649
private @Nullable StateNamespace namespace;
@@ -670,6 +676,12 @@ private OnTimerArgumentProvider(
670676
this.timestamp = timestamp;
671677
this.timeDomain = timeDomain;
672678
this.key = key;
679+
this.outputTemplate =
680+
WindowedValues.<Void>builder()
681+
.setValue(null)
682+
.setWindow(window)
683+
.setTimestamp(timestamp)
684+
.setPaneInfo(PaneInfo.NO_FIRING);
673685
}
674686

675687
@Override
@@ -888,8 +900,13 @@ public <T> void outputWindowedValue(
888900
Collection<? extends BoundedWindow> windows,
889901
PaneInfo paneInfo) {
890902
checkTimestamp(timestamp(), timestamp);
891-
SimpleDoFnRunner.this.outputWindowedValue(
892-
tag, WindowedValues.of(output, timestamp, windows, paneInfo));
903+
WindowedValues.builder(outputTemplate)
904+
.withValue(output)
905+
.setTimestamp(timestamp)
906+
.setWindows(windows)
907+
.setPaneInfo(paneInfo)
908+
.setReceiver(wv -> SimpleDoFnRunner.this.outputWindowedValue(tag, wv))
909+
.output();
893910
}
894911

895912
@Override
@@ -909,6 +926,8 @@ private class OnWindowExpirationArgumentProvider<KeyT>
909926
private final BoundedWindow window;
910927
private final Instant timestamp;
911928
private final KeyT key;
929+
private final WindowedValues.Builder<Void> outputTemplate;
930+
912931
/** Lazily initialized; should only be accessed via {@link #getNamespace()}. */
913932
private @Nullable StateNamespace namespace;
914933

@@ -931,6 +950,12 @@ private OnWindowExpirationArgumentProvider(BoundedWindow window, Instant timesta
931950
this.window = window;
932951
this.timestamp = timestamp;
933952
this.key = key;
953+
this.outputTemplate =
954+
WindowedValues.<Void>builder()
955+
.setValue(null)
956+
.setWindow(window)
957+
.setTimestamp(timestamp)
958+
.setPaneInfo(PaneInfo.NO_FIRING);
934959
}
935960

936961
@Override
@@ -1117,8 +1142,13 @@ public <T> void outputWindowedValue(
11171142
Collection<? extends BoundedWindow> windows,
11181143
PaneInfo paneInfo) {
11191144
checkTimestamp(this.timestamp, timestamp);
1120-
SimpleDoFnRunner.this.outputWindowedValue(
1121-
tag, WindowedValues.of(output, timestamp, windows, paneInfo));
1145+
WindowedValues.builder(outputTemplate)
1146+
.withValue(output)
1147+
.setTimestamp(timestamp)
1148+
.setWindows(windows)
1149+
.setPaneInfo(paneInfo)
1150+
.setReceiver(wv -> SimpleDoFnRunner.this.outputWindowedValue(tag, wv))
1151+
.output();
11221152
}
11231153

11241154
@Override

‎runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -246,8 +246,8 @@ private BundleWindowedValueReceiver(UncommittedBundle<KV<K, Iterable<V>>> bundle
246246
}
247247

248248
@Override
249-
public void output(WindowedValue<KV<K, Iterable<V>>> valueWithMetadata) {
250-
bundle.add(valueWithMetadata);
249+
public void output(WindowedValue<KV<K, Iterable<V>>> windowedValue) {
250+
bundle.add(windowedValue);
251251
}
252252
}
253253
}

‎runners/direct-java/src/main/java/org/apache/beam/runners/direct/WindowEvaluatorFactory.java‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -90,9 +90,7 @@ public WindowIntoEvaluator(
9090
public void processElement(WindowedValue<InputT> compressedElement) throws Exception {
9191
for (WindowedValue<InputT> element : compressedElement.explodeWindows()) {
9292
Collection<? extends BoundedWindow> windows = assignWindows(windowFn, element);
93-
outputBundle.add(
94-
WindowedValues.of(
95-
element.getValue(), element.getTimestamp(), windows, element.getPaneInfo()));
93+
WindowedValues.builder(element).setWindows(windows).setReceiver(outputBundle::add).output();
9694
}
9795
}
9896

‎runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java‎

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1403,20 +1403,19 @@ private SourceContextWrapper(SourceContext<WindowedValue<OutputT>> ctx) {
14031403
@Override
14041404
public void collect(WindowedValue<ValueWithRecordId<OutputT>> element) {
14051405
OutputT originalValue = element.getValue().getValue();
1406-
WindowedValue<OutputT> output =
1407-
WindowedValues.of(
1408-
originalValue, element.getTimestamp(), element.getWindows(), element.getPaneInfo());
1409-
ctx.collect(output);
1406+
WindowedValues.builder(element)
1407+
.withValue(originalValue)
1408+
.setReceiver(ctx::collect)
1409+
.output();
14101410
}
14111411

14121412
@Override
14131413
public void collectWithTimestamp(
14141414
WindowedValue<ValueWithRecordId<OutputT>> element, long timestamp) {
14151415
OutputT originalValue = element.getValue().getValue();
1416-
WindowedValue<OutputT> output =
1417-
WindowedValues.of(
1418-
originalValue, element.getTimestamp(), element.getWindows(), element.getPaneInfo());
1419-
ctx.collectWithTimestamp(output, timestamp);
1416+
WindowedValues.builder(element)
1417+
.withValue(originalValue)
1418+
.setReceiver(wv -> ctx.collectWithTimestamp(wv, timestamp));
14201419
}
14211420

14221421
@Override

‎runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkDoFnFunction.java‎

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -223,12 +223,10 @@ public void setCollector(Collector<WindowedValue<RawUnionValue>> collector) {
223223
@Override
224224
public <T> void output(TupleTag<T> tag, WindowedValue<T> output) {
225225
checkStateNotNull(collector);
226-
collector.collect(
227-
WindowedValues.of(
228-
new RawUnionValue(0 /* single output */, output.getValue()),
229-
output.getTimestamp(),
230-
output.getWindows(),
231-
output.getPaneInfo()));
226+
WindowedValues.builder(output)
227+
.withValue(new RawUnionValue(0 /* single output */, output.getValue()))
228+
.setReceiver(collector::collect)
229+
.output();
232230
}
233231
}
234232

@@ -257,13 +255,10 @@ public void setCollector(Collector<WindowedValue<RawUnionValue>> collector) {
257255
@Override
258256
public <T> void output(TupleTag<T> tag, WindowedValue<T> output) {
259257
checkStateNotNull(collector);
260-
261-
collector.collect(
262-
WindowedValues.of(
263-
new RawUnionValue(outputMap.get(tag), output.getValue()),
264-
output.getTimestamp(),
265-
output.getWindows(),
266-
output.getPaneInfo()));
258+
WindowedValues.builder(output)
259+
.withValue(new RawUnionValue(outputMap.get(tag), output.getValue()))
260+
.setReceiver(collector::collect)
261+
.output();
267262
}
268263
}
269264
}

‎runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkNonMergingReduceFunction.java‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -101,11 +101,11 @@ public void reduce(
101101
(WindowedValue<KV<K, InputT>> wv) ->
102102
Objects.requireNonNull(wv).getValue().getValue()));
103103
}
104-
coll.collect(
105-
WindowedValues.of(
106-
KV.of(first.getValue().getKey(), values),
107-
combinedTimestamp,
108-
first.getWindows(),
109-
PaneInfo.ON_TIME_AND_ONLY_FIRING));
104+
WindowedValues.builder(first)
105+
.withValue(KV.of(first.getValue().getKey(), values))
106+
.setReceiver(coll::collect)
107+
.setPaneInfo(PaneInfo.ON_TIME_AND_ONLY_FIRING)
108+
.setTimestamp(combinedTimestamp)
109+
.output();
110110
}
111111
}

‎runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/RedistributeByKeyOverrideFactory.java‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
*/
1818
package org.apache.beam.runners.dataflow;
1919

20-
import java.util.Collections;
2120
import org.apache.beam.runners.dataflow.internal.DataflowGroupByKey;
2221
import org.apache.beam.sdk.runners.AppliedPTransform;
2322
import org.apache.beam.sdk.runners.PTransformOverrideFactory;
@@ -134,12 +133,14 @@ public Duration getAllowedTimestampSkew() {
134133

135134
@ProcessElement
136135
public void processElement(
137-
@Element KV<K, ValueInSingleWindow<V>> kv, OutputReceiver<KV<K, V>> r) {
138-
r.outputWindowedValue(
139-
KV.of(kv.getKey(), kv.getValue().getValue()),
140-
kv.getValue().getTimestamp(),
141-
Collections.singleton(kv.getValue().getWindow()),
142-
kv.getValue().getPaneInfo());
136+
@Element KV<K, ValueInSingleWindow<V>> kv,
137+
OutputReceiver<KV<K, V>> outputReceiver) {
138+
outputReceiver
139+
.builder(KV.of(kv.getKey(), kv.getValue().getValue()))
140+
.setTimestamp(kv.getValue().getTimestamp())
141+
.setWindow(kv.getValue().getWindow())
142+
.setPaneInfo(kv.getValue().getPaneInfo())
143+
.output();
143144
}
144145
}));
145146
}

0 commit comments

Comments
 (0)