Skip to content

Commit 679adb8

Browse files
ahmedabu98ahmedabu98
authored
Pipe Valuekind through DoFn and output builders (#38490)
* add valuekind * add tests and DirectRunner support * add some doc to model * add some doc to model * add missing noteOutput * add to ParamWindowedValueCoder * add to windmill keyed work item * add missing pieces to Dataflow runner, fnapi, spark * remove premature dataflow enabliung * address comments * remove duplicate method * remove extra methods from DoFn * cleanup * cleanup * create a valuekind helper in df worker * cleanup
1 parent 590c9c2 commit 679adb8

25 files changed

Lines changed: 1059 additions & 22 deletions

File tree

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
import org.apache.beam.sdk.values.PCollectionView;
5353
import org.apache.beam.sdk.values.Row;
5454
import org.apache.beam.sdk.values.TupleTag;
55+
import org.apache.beam.sdk.values.ValueKind;
5556
import org.apache.beam.sdk.values.WindowedValue;
5657
import org.apache.beam.sdk.values.WindowedValues;
5758
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
@@ -200,6 +201,11 @@ public CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn) {
200201
return processContext.causedByDrain();
201202
}
202203

204+
@Override
205+
public ValueKind valueKind(DoFn<InputT, OutputT> doFn) {
206+
return processContext.valueKind();
207+
}
208+
203209
@Override
204210
public RestrictionTracker<?, ?> restrictionTracker() {
205211
return processContext.tracker;
@@ -436,6 +442,11 @@ public CausedByDrain causedByDrain() {
436442
return element.causedByDrain();
437443
}
438444

445+
@Override
446+
public ValueKind valueKind() {
447+
return element.getValueKind();
448+
}
449+
439450
@Override
440451
public PipelineOptions getPipelineOptions() {
441452
return pipelineOptions;

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@
6161
import org.apache.beam.sdk.values.PCollectionView;
6262
import org.apache.beam.sdk.values.Row;
6363
import org.apache.beam.sdk.values.TupleTag;
64+
import org.apache.beam.sdk.values.ValueKind;
6465
import org.apache.beam.sdk.values.WindowedValue;
6566
import org.apache.beam.sdk.values.WindowedValues;
6667
import org.apache.beam.sdk.values.WindowingStrategy;
@@ -506,6 +507,11 @@ public Instant timestamp() {
506507
return elem.getRecordOffset();
507508
}
508509

510+
@Override
511+
public ValueKind valueKind() {
512+
return elem.getValueKind();
513+
}
514+
509515
public Collection<? extends BoundedWindow> windows() {
510516
return elem.getWindows();
511517
}
@@ -597,6 +603,11 @@ public CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn) {
597603
return elem.causedByDrain();
598604
}
599605

606+
@Override
607+
public ValueKind valueKind(DoFn<InputT, OutputT> doFn) {
608+
return elem.getValueKind();
609+
}
610+
600611
@Override
601612
public String timerId(DoFn<InputT, OutputT> doFn) {
602613
throw new UnsupportedOperationException(
@@ -895,6 +906,11 @@ public CausedByDrain causedByDrain(DoFn<InputT, OutputT> doFn) {
895906
return causedByDrain;
896907
}
897908

909+
@Override
910+
public ValueKind valueKind(DoFn<InputT, OutputT> doFn) {
911+
throw new UnsupportedOperationException("ValueKind parameters are not supported.");
912+
}
913+
898914
@Override
899915
public String timerId(DoFn<InputT, OutputT> doFn) {
900916
return timerId;
@@ -1228,6 +1244,11 @@ public TimeDomain timeDomain(DoFn<InputT, OutputT> doFn) {
12281244
"Cannot access time domain outside of @ProcessTimer method.");
12291245
}
12301246

1247+
@Override
1248+
public ValueKind valueKind(DoFn<InputT, OutputT> doFn) {
1249+
throw new UnsupportedOperationException("ValueKind parameters are not supported.");
1250+
}
1251+
12311252
@Override
12321253
public KeyT key() {
12331254
return key;

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1381,6 +1381,11 @@ public T getValue() {
13811381
return value;
13821382
}
13831383

1384+
@Override
1385+
public ValueKind getValueKind() {
1386+
return ValueKind.INSERT;
1387+
}
1388+
13841389
@Override
13851390
public CausedByDrain causedByDrain() {
13861391
return CausedByDrain.NORMAL;
@@ -1416,11 +1421,6 @@ public PaneInfo getPaneInfo() {
14161421
return null;
14171422
}
14181423

1419-
@Override
1420-
public ValueKind getValueKind() {
1421-
return ValueKind.INSERT;
1422-
}
1423-
14241424
@Override
14251425
public Iterable<WindowedValue<T>> explodeWindows() {
14261426
return Collections.emptyList();

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,7 @@ public void processElement(
141141
.setTimestamp(kv.getValue().getTimestamp())
142142
.setWindow(kv.getValue().getWindow())
143143
.setPaneInfo(kv.getValue().getPaneInfo())
144+
.setValueKind(kv.getValue().getValueKind())
144145
.output();
145146
}
146147
}));

‎runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,7 @@
154154
import org.apache.beam.sdk.transforms.PTransform;
155155
import org.apache.beam.sdk.transforms.ParDo;
156156
import org.apache.beam.sdk.transforms.Redistribute;
157+
import org.apache.beam.sdk.transforms.Reshuffle;
157158
import org.apache.beam.sdk.transforms.SerializableFunctions;
158159
import org.apache.beam.sdk.transforms.SimpleFunction;
159160
import org.apache.beam.sdk.transforms.resourcehints.ResourceHints;
@@ -171,11 +172,17 @@
171172
import org.apache.beam.sdk.util.construction.PipelineTranslation;
172173
import org.apache.beam.sdk.util.construction.SdkComponents;
173174
import org.apache.beam.sdk.util.construction.TransformPayloadTranslatorRegistrar;
175+
import org.apache.beam.sdk.values.CausedByDrain;
174176
import org.apache.beam.sdk.values.KV;
175177
import org.apache.beam.sdk.values.PCollection;
178+
import org.apache.beam.sdk.values.PCollectionTuple;
176179
import org.apache.beam.sdk.values.PValues;
177180
import org.apache.beam.sdk.values.TimestampedValue;
181+
import org.apache.beam.sdk.values.TupleTag;
182+
import org.apache.beam.sdk.values.TupleTagList;
178183
import org.apache.beam.sdk.values.TypeDescriptors;
184+
import org.apache.beam.sdk.values.ValueKind;
185+
import org.apache.beam.sdk.values.WindowedValues;
179186
import org.apache.beam.sdk.values.WindowingStrategy;
180187
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.InvalidProtocolBufferException;
181188
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
@@ -191,6 +198,7 @@
191198
import org.joda.time.Duration;
192199
import org.joda.time.Instant;
193200
import org.junit.Before;
201+
import org.junit.Ignore;
194202
import org.junit.Rule;
195203
import org.junit.Test;
196204
import org.junit.experimental.categories.Category;
@@ -2777,4 +2785,139 @@ public void process() {}
27772785
}));
27782786
p.run();
27792787
}
2788+
2789+
@Test
2790+
@Category({ValidatesRunner.class})
2791+
public void testValueKindParameterAndOutputWithKind() {
2792+
boolean isRunnerV2 = false;
2793+
@Nullable
2794+
List<String> experiments =
2795+
pipeline.getOptions().as(DataflowPipelineOptions.class).getExperiments();
2796+
if (experiments != null
2797+
&& (experiments.contains("use_unified_worker") || experiments.contains("use_runner_v2"))) {
2798+
isRunnerV2 = true;
2799+
}
2800+
// Skipp runner v2 because its Create uses a splittable DoFn, which contains a shuffle.
2801+
// ValueKind is not supported in Dataflow shuffle yet
2802+
assumeFalse(isRunnerV2);
2803+
2804+
PCollection<String> input = pipeline.apply(Create.of("a", "b", "c", "d"));
2805+
TupleTag<String> mainTag = new TupleTag<String>() {};
2806+
TupleTag<String> sideTag = new TupleTag<String>() {};
2807+
2808+
PCollectionTuple tuple =
2809+
input.apply(
2810+
"SetKind",
2811+
ParDo.of(
2812+
new DoFn<String, String>() {
2813+
@ProcessElement
2814+
public void processElement(
2815+
@Element String element,
2816+
@Timestamp org.joda.time.Instant timestamp,
2817+
BoundedWindow window,
2818+
PaneInfo paneInfo,
2819+
ProcessContext c,
2820+
MultiOutputReceiver outputReceiver) {
2821+
switch (element) {
2822+
case "a":
2823+
c.output(element); // default: INSERT
2824+
return;
2825+
case "b":
2826+
c.outputWindowedValue(
2827+
WindowedValues.of(
2828+
element,
2829+
timestamp,
2830+
Collections.singleton(window),
2831+
paneInfo,
2832+
null,
2833+
null,
2834+
CausedByDrain.NORMAL,
2835+
null,
2836+
ValueKind.UPDATE_BEFORE));
2837+
return;
2838+
case "c":
2839+
outputReceiver
2840+
.get(mainTag)
2841+
.builder(element)
2842+
.setValueKind(ValueKind.UPDATE_AFTER)
2843+
.output();
2844+
return;
2845+
case "d":
2846+
outputReceiver
2847+
.get(sideTag)
2848+
.builder(element)
2849+
.setValueKind(ValueKind.DELETE)
2850+
.output();
2851+
}
2852+
}
2853+
})
2854+
.withOutputTags(mainTag, TupleTagList.of(sideTag)));
2855+
2856+
PCollection<String> main =
2857+
tuple
2858+
.get(mainTag)
2859+
.apply(
2860+
"ReadKind",
2861+
ParDo.of(
2862+
new DoFn<String, String>() {
2863+
@ProcessElement
2864+
public void processElement(
2865+
@Element String element, ProcessContext c, ValueKind kind) {
2866+
c.output(element + ":" + kind);
2867+
}
2868+
}));
2869+
2870+
PCollection<String> side =
2871+
tuple
2872+
.get(sideTag)
2873+
.apply(
2874+
"ReadKind-SideTag",
2875+
ParDo.of(
2876+
new DoFn<String, String>() {
2877+
@ProcessElement
2878+
public void processElement(
2879+
@Element String element, ProcessContext c, ValueKind kind) {
2880+
c.output(element + ":" + kind);
2881+
}
2882+
}));
2883+
2884+
PAssert.that(main).containsInAnyOrder("a:INSERT", "b:UPDATE_BEFORE", "c:UPDATE_AFTER");
2885+
PAssert.that(side).containsInAnyOrder("d:DELETE");
2886+
pipeline.run();
2887+
}
2888+
2889+
@Test
2890+
@Ignore("enable once when element metadata is supported in shuffle")
2891+
@Category({ValidatesRunner.class})
2892+
public void testValueKindPreservedAcrossShuffle() {
2893+
PCollection<KV<String, String>> input = pipeline.apply(Create.of(KV.of("key", "value")));
2894+
2895+
PCollection<String> output =
2896+
input
2897+
.apply(
2898+
"SetKind",
2899+
ParDo.of(
2900+
new DoFn<KV<String, String>, KV<String, String>>() {
2901+
@ProcessElement
2902+
public void processElement(
2903+
@Element KV<String, String> element,
2904+
OutputReceiver<KV<String, String>> out) {
2905+
out.builder(element).setValueKind(ValueKind.UPDATE_BEFORE).output();
2906+
}
2907+
}))
2908+
.apply(Reshuffle.of())
2909+
.apply(
2910+
"ReadKind",
2911+
ParDo.of(
2912+
new DoFn<KV<String, String>, String>() {
2913+
@ProcessElement
2914+
public void processElement(
2915+
@Element KV<String, String> element, ProcessContext c, ValueKind kind) {
2916+
c.output(element.getValue() + ":" + kind);
2917+
}
2918+
}));
2919+
2920+
PAssert.that(output).containsInAnyOrder("value:UPDATE_BEFORE");
2921+
pipeline.run();
2922+
}
27802923
}

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,6 @@
3939
import org.apache.beam.sdk.values.CausedByDrain;
4040
import org.apache.beam.sdk.values.KV;
4141
import org.apache.beam.sdk.values.ValueKind;
42-
import org.apache.beam.sdk.values.ValueKindUtil;
4342
import org.apache.beam.sdk.values.WindowedValue;
4443
import org.apache.beam.sdk.values.WindowedValues;
4544
import org.apache.beam.sdk.values.WindowedValues.FullWindowedValueCoder;
@@ -149,7 +148,7 @@ protected WindowedValue<T> decodeMessage(Windmill.Message message) throws IOExce
149148
elementMetadata.getDrain() == BeamFnApi.Elements.DrainMode.Enum.DRAINING
150149
? CausedByDrain.CAUSED_BY_DRAIN
151150
: CausedByDrain.NORMAL;
152-
valueKind = ValueKindUtil.fromProto(elementMetadata.getValueKind());
151+
valueKind = WindmillValueKindHelper.fromProto(elementMetadata.getValueKind());
153152
}
154153
if (valueCoder instanceof KvCoder) {
155154
KvCoder<?, ?> kvCoder = (KvCoder<?, ?>) valueCoder;

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,6 @@
4343
import org.apache.beam.sdk.util.common.ElementByteSizeObserver;
4444
import org.apache.beam.sdk.values.CausedByDrain;
4545
import org.apache.beam.sdk.values.ValueKind;
46-
import org.apache.beam.sdk.values.ValueKindUtil;
4746
import org.apache.beam.sdk.values.WindowedValue;
4847
import org.apache.beam.sdk.values.WindowedValues;
4948
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Predicate;
@@ -158,7 +157,7 @@ public Iterable<TimerData> timersIterable() {
158157
elementMetadata.getDrain() == BeamFnApi.Elements.DrainMode.Enum.DRAINING
159158
? CausedByDrain.CAUSED_BY_DRAIN
160159
: CausedByDrain.NORMAL;
161-
valueKind = ValueKindUtil.fromProto(elementMetadata.getValueKind());
160+
valueKind = WindmillValueKindHelper.fromProto(elementMetadata.getValueKind());
162161
}
163162
InputStream inputStream = message.getData().newInput();
164163
ElemT value = valueCoder.decode(inputStream, Coder.Context.OUTER);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,7 @@ public long add(WindowedValue<T> data) throws IOException {
226226
data.causedByDrain() == CausedByDrain.CAUSED_BY_DRAIN
227227
? BeamFnApi.Elements.DrainMode.Enum.DRAINING
228228
: BeamFnApi.Elements.DrainMode.Enum.NOT_DRAINING)
229+
.setValueKind(WindmillValueKindHelper.toProto(data.getValueKind()))
229230
.build();
230231
ByteString metadata =
231232
encodeMetadata(

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

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.beam.runners.dataflow.worker;
19+
20+
import org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements;
21+
import org.apache.beam.sdk.values.ValueKind;
22+
23+
/**
24+
* Utility class for converting between {@link ValueKind} and {@link Elements.ValueKind.Enum} inside
25+
* the Dataflow Worker. This is a duplicate of the utility inside core sdk because core utility is
26+
* compiled against unrepackaged BeamFnApi, whereas Dataflow Worker relocates BeamFnApi to dataflow
27+
* repackaged version causing NoSuchMethodErrors at runtime.
28+
*/
29+
final class WindmillValueKindHelper {
30+
public static Elements.ValueKind.Enum toProto(ValueKind valueKind) {
31+
switch (valueKind) {
32+
case INSERT:
33+
return Elements.ValueKind.Enum.INSERT;
34+
case UPDATE_BEFORE:
35+
return Elements.ValueKind.Enum.UPDATE_BEFORE;
36+
case UPDATE_AFTER:
37+
return Elements.ValueKind.Enum.UPDATE_AFTER;
38+
case DELETE:
39+
return Elements.ValueKind.Enum.DELETE;
40+
default:
41+
throw new IllegalArgumentException("Unknown ValueKind: " + valueKind);
42+
}
43+
}
44+
45+
public static ValueKind fromProto(Elements.ValueKind.Enum proto) {
46+
switch (proto) {
47+
case VALUE_KIND_UNSPECIFIED:
48+
case INSERT:
49+
return ValueKind.INSERT;
50+
case UPDATE_BEFORE:
51+
return ValueKind.UPDATE_BEFORE;
52+
case UPDATE_AFTER:
53+
return ValueKind.UPDATE_AFTER;
54+
case DELETE:
55+
return ValueKind.DELETE;
56+
default:
57+
throw new IllegalArgumentException("Unknown ValueKind: " + proto);
58+
}
59+
}
60+
}

0 commit comments

Comments
 (0)