Skip to content

Commit a888fb9

Browse files
kennknowleskennknowles
committed
WIP: add OutputBuilder
1 parent 60307b4 commit a888fb9

91 files changed

Lines changed: 1157 additions & 837 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

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

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

11991199
List<String> skipDefRegexes = []
12001200
skipDefRegexes << "AutoValue_.*"
1201+
skipDefRegexes << "AutoBuilder_.*"
12011202
skipDefRegexes << "AutoOneOf_.*"
12021203
skipDefRegexes << ".*\\.jmh_generated\\..*"
12031204
skipDefRegexes += configuration.generatedClassPatterns
@@ -1291,7 +1292,8 @@ class BeamModulePlugin implements Plugin<Project> {
12911292
'**/org/apache/beam/gradle/**',
12921293
'**/org/apache/beam/model/**',
12931294
'**/org/apache/beam/runners/dataflow/worker/windmill/**',
1294-
'**/AutoValue_*'
1295+
'**/AutoValue_*',
1296+
'**/AutoBuilder_*',
12951297
]
12961298

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

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

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import org.apache.beam.sdk.transforms.DoFnSchemaInformation;
3030
import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
3131
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
32+
import org.apache.beam.sdk.util.ValueWithMetadataMultiReceiver;
3233
import org.apache.beam.sdk.util.WindowedValue;
3334
import org.apache.beam.sdk.values.KV;
3435
import org.apache.beam.sdk.values.PCollectionView;
@@ -41,12 +42,6 @@
4142
"nullness" // TODO(https://github.com/apache/beam/issues/20497)
4243
})
4344
public class DoFnRunners {
44-
/** Information about how to create output receivers and output to them. */
45-
public interface OutputManager {
46-
/** Outputs a single element to the receiver indicated by the given {@link TupleTag}. */
47-
<T> void output(TupleTag<T> tag, WindowedValue<T> output);
48-
}
49-
5045
/**
5146
* Returns an implementation of {@link DoFnRunner} that for a {@link DoFn}.
5247
*
@@ -58,7 +53,7 @@ public static <InputT, OutputT> DoFnRunner<InputT, OutputT> simpleRunner(
5853
PipelineOptions options,
5954
DoFn<InputT, OutputT> fn,
6055
SideInputReader sideInputReader,
61-
OutputManager outputManager,
56+
ValueWithMetadataMultiReceiver outputManager,
6257
TupleTag<OutputT> mainOutputTag,
6358
List<TupleTag<?>> additionalOutputTags,
6459
StepContext stepContext,
@@ -168,7 +163,7 @@ ProcessFnRunner<InputT, OutputT, RestrictionT> newProcessFnRunner(
168163
PipelineOptions options,
169164
Collection<PCollectionView<?>> views,
170165
ReadyCheckingSideInputReader sideInputReader,
171-
OutputManager outputManager,
166+
ValueWithMetadataMultiReceiver outputManager,
172167
TupleTag<OutputT> mainOutputTag,
173168
List<TupleTag<?>> additionalOutputTags,
174169
StepContext stepContext,

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

Lines changed: 5 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -17,20 +17,17 @@
1717
*/
1818
package org.apache.beam.runners.core;
1919

20-
import java.util.Collection;
2120
import org.apache.beam.model.pipeline.v1.RunnerApi;
2221
import org.apache.beam.runners.core.triggers.ExecutableTriggerStateMachine;
2322
import org.apache.beam.runners.core.triggers.TriggerStateMachines;
2423
import org.apache.beam.sdk.transforms.DoFn;
2524
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
26-
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
2725
import org.apache.beam.sdk.util.SystemDoFnInternal;
28-
import org.apache.beam.sdk.util.WindowedValue;
26+
import org.apache.beam.sdk.util.ValueWithMetadataMultiReceiver;
2927
import org.apache.beam.sdk.util.construction.TriggerTranslation;
3028
import org.apache.beam.sdk.values.KV;
3129
import org.apache.beam.sdk.values.TupleTag;
3230
import org.apache.beam.sdk.values.WindowingStrategy;
33-
import org.joda.time.Instant;
3431

3532
/**
3633
* A general {@link GroupAlsoByWindowsAggregators}. This delegates all of the logic to the {@link
@@ -51,7 +48,7 @@ DoFn<KeyedWorkItem<K, InputT>, KV<K, OutputT>> create(
5148
TimerInternalsFactory<K> timerInternalsFactory,
5249
SideInputReader sideInputReader,
5350
SystemReduceFn<K, InputT, ?, OutputT, W> reduceFn,
54-
DoFnRunners.OutputManager outputManager,
51+
ValueWithMetadataMultiReceiver outputManager,
5552
TupleTag<KV<K, OutputT>> mainTag) {
5653
return new GroupAlsoByWindowViaWindowSetNewDoFn<>(
5754
strategy,
@@ -68,7 +65,7 @@ DoFn<KeyedWorkItem<K, InputT>, KV<K, OutputT>> create(
6865
private transient StateInternalsFactory<K> stateInternalsFactory;
6966
private transient TimerInternalsFactory<K> timerInternalsFactory;
7067
private transient SideInputReader sideInputReader;
71-
private transient DoFnRunners.OutputManager outputManager;
68+
private transient ValueWithMetadataMultiReceiver outputManager;
7269
private TupleTag<KV<K, OutputT>> mainTag;
7370

7471
public GroupAlsoByWindowViaWindowSetNewDoFn(
@@ -77,7 +74,7 @@ public GroupAlsoByWindowViaWindowSetNewDoFn(
7774
TimerInternalsFactory<K> timerInternalsFactory,
7875
SideInputReader sideInputReader,
7976
SystemReduceFn<K, InputT, ?, OutputT, W> reduceFn,
80-
DoFnRunners.OutputManager outputManager,
77+
ValueWithMetadataMultiReceiver outputManager,
8178
TupleTag<KV<K, OutputT>> mainTag) {
8279
this.timerInternalsFactory = timerInternalsFactory;
8380
this.sideInputReader = sideInputReader;
@@ -91,29 +88,6 @@ public GroupAlsoByWindowViaWindowSetNewDoFn(
9188
this.triggerProto = TriggerTranslation.toProto(windowingStrategy.getTrigger());
9289
}
9390

94-
private OutputWindowedValue<KV<K, OutputT>> outputWindowedValue() {
95-
return new OutputWindowedValue<KV<K, OutputT>>() {
96-
@Override
97-
public void outputWindowedValue(
98-
KV<K, OutputT> output,
99-
Instant timestamp,
100-
Collection<? extends BoundedWindow> windows,
101-
PaneInfo pane) {
102-
outputManager.output(mainTag, WindowedValue.of(output, timestamp, windows, pane));
103-
}
104-
105-
@Override
106-
public <AdditionalOutputT> void outputWindowedValue(
107-
TupleTag<AdditionalOutputT> tag,
108-
AdditionalOutputT output,
109-
Instant timestamp,
110-
Collection<? extends BoundedWindow> windows,
111-
PaneInfo pane) {
112-
outputManager.output(tag, WindowedValue.of(output, timestamp, windows, pane));
113-
}
114-
};
115-
}
116-
11791
@ProcessElement
11892
public void processElement(ProcessContext c) throws Exception {
11993
KeyedWorkItem<K, InputT> keyedWorkItem = c.element();
@@ -130,7 +104,7 @@ public void processElement(ProcessContext c) throws Exception {
130104
TriggerStateMachines.stateMachineForTrigger(triggerProto)),
131105
stateInternals,
132106
timerInternals,
133-
outputWindowedValue(),
107+
outputManager.forTag(mainTag),
134108
sideInputReader,
135109
reduceFn,
136110
c.getPipelineOptions());

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,7 @@ public <K, InputT> Iterable<WindowedValue<InputT>> filter(
141141
} else {
142142
nonLateElements.add(
143143
WindowedValue.of(
144-
element.getValue(), element.getTimestamp(), window, element.getPane()));
144+
element.getValue(), element.getTimestamp(), window, element.getPaneInfo()));
145145
}
146146
}
147147
}

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

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
4646
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
4747
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
48+
import org.apache.beam.sdk.util.ValueWithMetadataMultiReceiver;
4849
import org.apache.beam.sdk.util.WindowedValue;
4950
import org.apache.beam.sdk.values.KV;
5051
import org.apache.beam.sdk.values.PCollectionView;
@@ -72,19 +73,20 @@ public class OutputAndTimeBoundedSplittableProcessElementInvoker<
7273
InputT, OutputT, RestrictionT, PositionT, WatermarkEstimatorStateT> {
7374
private final DoFn<InputT, OutputT> fn;
7475
private final PipelineOptions pipelineOptions;
75-
private final OutputWindowedValue<OutputT> output;
76+
private final ValueWithMetadataMultiReceiver outputReceiver;
7677
private final SideInputReader sideInputReader;
7778
private final ScheduledExecutorService executor;
7879
private final int maxNumOutputs;
7980
private final Duration maxDuration;
8081
private final Supplier<BundleFinalizer> bundleFinalizer;
82+
private final TupleTag<OutputT> mainOutputTag;
8183

8284
/**
8385
* Creates a new invoker from components.
8486
*
8587
* @param fn The original {@link DoFn}.
8688
* @param pipelineOptions {@link PipelineOptions} to include in the {@link DoFn.ProcessContext}.
87-
* @param output Hook for outputting from the {@link DoFn.ProcessElement} method.
89+
* @param outputReceiver Hook for outputting from the {@link DoFn.ProcessElement} method.
8890
* @param sideInputReader Hook for accessing side inputs.
8991
* @param executor Executor on which a checkpoint will be scheduled after the given duration.
9092
* @param maxNumOutputs Maximum number of outputs, in total over all output tags, after which a
@@ -98,15 +100,17 @@ public class OutputAndTimeBoundedSplittableProcessElementInvoker<
98100
public OutputAndTimeBoundedSplittableProcessElementInvoker(
99101
DoFn<InputT, OutputT> fn,
100102
PipelineOptions pipelineOptions,
101-
OutputWindowedValue<OutputT> output,
103+
ValueWithMetadataMultiReceiver outputReceiver,
104+
TupleTag<OutputT> mainOutputTag,
102105
SideInputReader sideInputReader,
103106
ScheduledExecutorService executor,
104107
int maxNumOutputs,
105108
Duration maxDuration,
106109
Supplier<BundleFinalizer> bundleFinalizer) {
107110
this.fn = fn;
108111
this.pipelineOptions = pipelineOptions;
109-
this.output = output;
112+
this.outputReceiver = outputReceiver;
113+
this.mainOutputTag = mainOutputTag;
110114
this.sideInputReader = sideInputReader;
111115
this.executor = executor;
112116
this.maxNumOutputs = maxNumOutputs;
@@ -375,7 +379,7 @@ public Instant timestamp() {
375379

376380
@Override
377381
public PaneInfo pane() {
378-
return element.getPane();
382+
return element.getPaneInfo();
379383
}
380384

381385
@Override
@@ -390,7 +394,7 @@ public void output(OutputT output) {
390394

391395
@Override
392396
public void outputWithTimestamp(OutputT value, Instant timestamp) {
393-
outputWindowedValue(value, timestamp, element.getWindows(), element.getPane());
397+
outputWindowedValue(value, timestamp, element.getWindows(), element.getPaneInfo());
394398
}
395399

396400
@Override
@@ -403,7 +407,9 @@ public void outputWindowedValue(
403407
if (watermarkEstimator instanceof TimestampObservingWatermarkEstimator) {
404408
((TimestampObservingWatermarkEstimator) watermarkEstimator).observeTimestamp(timestamp);
405409
}
406-
output.outputWindowedValue(value, timestamp, windows, paneInfo);
410+
outputReceiver
411+
.forTag(mainOutputTag)
412+
.output(WindowedValue.of(value, timestamp, windows, paneInfo));
407413
}
408414

409415
@Override
@@ -413,7 +419,9 @@ public <T> void output(TupleTag<T> tag, T value) {
413419

414420
@Override
415421
public <T> void outputWithTimestamp(TupleTag<T> tag, T value, Instant timestamp) {
416-
outputWindowedValue(tag, value, timestamp, element.getWindows(), element.getPane());
422+
outputReceiver
423+
.forTag(tag)
424+
.output(WindowedValue.of(value, timestamp, element.getWindows(), element.getPaneInfo()));
417425
}
418426

419427
@Override
@@ -427,7 +435,7 @@ public <T> void outputWindowedValue(
427435
if (watermarkEstimator instanceof TimestampObservingWatermarkEstimator) {
428436
((TimestampObservingWatermarkEstimator) watermarkEstimator).observeTimestamp(timestamp);
429437
}
430-
output.outputWindowedValue(tag, value, timestamp, windows, paneInfo);
438+
outputReceiver.forTag(tag).output(WindowedValue.of(value, timestamp, windows, paneInfo));
431439
}
432440

433441
private void noteOutput() {

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

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import org.apache.beam.sdk.transforms.windowing.Window;
4747
import org.apache.beam.sdk.transforms.windowing.Window.ClosingBehavior;
4848
import org.apache.beam.sdk.transforms.windowing.WindowFn;
49+
import org.apache.beam.sdk.util.ValueWithMetadataReceiver;
4950
import org.apache.beam.sdk.util.WindowTracing;
5051
import org.apache.beam.sdk.util.WindowedValue;
5152
import org.apache.beam.sdk.values.KV;
@@ -106,7 +107,7 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
106107
*/
107108
private final WindowingStrategy<Object, W> windowingStrategy;
108109

109-
private final OutputWindowedValue<KV<K, OutputT>> outputter;
110+
private final ValueWithMetadataReceiver<KV<K, OutputT>> outputter;
110111

111112
private final StateInternals stateInternals;
112113

@@ -214,7 +215,7 @@ public ReduceFnRunner(
214215
ExecutableTriggerStateMachine triggerStateMachine,
215216
StateInternals stateInternals,
216217
TimerInternals timerInternals,
217-
OutputWindowedValue<KV<K, OutputT>> outputter,
218+
ValueWithMetadataReceiver<KV<K, OutputT>> outputter,
218219
@Nullable SideInputReader sideInputReader,
219220
ReduceFn<K, InputT, OutputT, W> reduceFn,
220221
@Nullable PipelineOptions options) {
@@ -1055,7 +1056,8 @@ private void prefetchOnTrigger(
10551056
}
10561057

10571058
// Output the actual value.
1058-
outputter.outputWindowedValue(KV.of(key, toOutput), outputTimestamp, windows, pane);
1059+
outputter.output(
1060+
WindowedValue.of(KV.of(key, toOutput), outputTimestamp, windows, pane));
10591061
});
10601062

10611063
reduceFn.onTrigger(renamedTriggerContext);

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

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@
2626
import java.util.List;
2727
import java.util.Map;
2828
import java.util.Set;
29-
import org.apache.beam.runners.core.DoFnRunners.OutputManager;
3029
import org.apache.beam.sdk.coders.Coder;
3130
import org.apache.beam.sdk.options.PipelineOptions;
3231
import org.apache.beam.sdk.schemas.SchemaCoder;
@@ -54,6 +53,7 @@
5453
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
5554
import org.apache.beam.sdk.util.SystemDoFnInternal;
5655
import org.apache.beam.sdk.util.UserCodeException;
56+
import org.apache.beam.sdk.util.ValueWithMetadataMultiReceiver;
5757
import org.apache.beam.sdk.util.WindowedValue;
5858
import org.apache.beam.sdk.values.PCollectionView;
5959
import org.apache.beam.sdk.values.Row;
@@ -93,7 +93,7 @@ public class SimpleDoFnRunner<InputT, OutputT> implements DoFnRunner<InputT, Out
9393
private final DoFnInvoker<InputT, OutputT> invoker;
9494

9595
private final SideInputReader sideInputReader;
96-
private final OutputManager outputManager;
96+
private final ValueWithMetadataMultiReceiver outputManager;
9797

9898
private final TupleTag<OutputT> mainOutputTag;
9999
/** The set of known output tags. */
@@ -123,7 +123,7 @@ public SimpleDoFnRunner(
123123
PipelineOptions options,
124124
DoFn<InputT, OutputT> fn,
125125
SideInputReader sideInputReader,
126-
OutputManager outputManager,
126+
ValueWithMetadataMultiReceiver outputManager,
127127
TupleTag<OutputT> mainOutputTag,
128128
List<TupleTag<?>> additionalOutputTags,
129129
StepContext stepContext,
@@ -273,7 +273,7 @@ private void checkTimestamp(Instant elemTimestamp, Instant timestamp) {
273273

274274
private <T> void outputWindowedValue(TupleTag<T> tag, WindowedValue<T> windowedElem) {
275275
checkArgument(outputTags.contains(tag), "Unknown output tag %s", tag);
276-
outputManager.output(tag, windowedElem);
276+
outputManager.forTag(tag).output(windowedElem);
277277
}
278278

279279
/** An {@link DoFnInvoker.ArgumentProvider} for {@link DoFn.StartBundle @StartBundle}. */
@@ -404,7 +404,7 @@ public <T> T sideInput(PCollectionView<T> view) {
404404

405405
@Override
406406
public PaneInfo pane() {
407-
return elem.getPane();
407+
return elem.getPaneInfo();
408408
}
409409

410410
@Override
@@ -436,7 +436,7 @@ public <T> void output(TupleTag<T> tag, T output) {
436436
public <T> void outputWithTimestamp(TupleTag<T> tag, T output, Instant timestamp) {
437437
checkNotNull(tag, "Tag passed to outputWithTimestamp cannot be null");
438438
checkTimestamp(elem.getTimestamp(), timestamp);
439-
outputWindowedValue(tag, output, timestamp, elem.getWindows(), elem.getPane());
439+
outputWindowedValue(tag, output, timestamp, elem.getWindows(), elem.getPaneInfo());
440440
}
441441

442442
@Override
@@ -446,8 +446,14 @@ public <T> void outputWindowedValue(
446446
Instant timestamp,
447447
Collection<? extends BoundedWindow> windows,
448448
PaneInfo paneInfo) {
449-
SimpleDoFnRunner.this.outputWindowedValue(
450-
tag, WindowedValue.of(output, timestamp, windows, paneInfo));
449+
WindowedValue.Builder<T> builder =
450+
((WindowedValue.Builder<T>) elem.toBuilder())
451+
.setValue(output)
452+
.setTimestamp(timestamp)
453+
.setWindows(windows)
454+
.setPaneInfo(paneInfo);
455+
456+
SimpleDoFnRunner.this.outputWindowedValue(tag, builder.build());
451457
}
452458

453459
@Override

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -430,7 +430,7 @@ public PipelineOptions pipelineOptions() {
430430

431431
@Override
432432
public PaneInfo paneInfo(DoFn<InputT, OutputT> doFn) {
433-
return elementAndRestriction.getKey().getPane();
433+
return elementAndRestriction.getKey().getPaneInfo();
434434
}
435435

436436
@Override
@@ -490,7 +490,7 @@ public PipelineOptions pipelineOptions() {
490490

491491
@Override
492492
public PaneInfo paneInfo(DoFn<InputT, OutputT> doFn) {
493-
return elementAndRestriction.getKey().getPane();
493+
return elementAndRestriction.getKey().getPaneInfo();
494494
}
495495

496496
@Override
@@ -544,7 +544,7 @@ public PipelineOptions pipelineOptions() {
544544

545545
@Override
546546
public PaneInfo paneInfo(DoFn<InputT, OutputT> doFn) {
547-
return elementAndRestriction.getKey().getPane();
547+
return elementAndRestriction.getKey().getPaneInfo();
548548
}
549549

550550
@Override

0 commit comments

Comments
 (0)