You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by al...@apache.org on 2018/03/27 04:37:10 UTC
[beam] branch master updated: [BEAM-622] Add checkpointing tests
for DoFnOperator and WindowDoFnOperator
This is an automated email from the ASF dual-hosted git repository.
aljoscha pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new ebc0fb0 [BEAM-622] Add checkpointing tests for DoFnOperator and WindowDoFnOperator
new 302b247 Merge pull request #4937: [BEAM-622] Add checkpointing tests for DoFnOperator and WindowDoFnOperator
ebc0fb0 is described below
commit ebc0fb0e3d0f7340400fe07d7d508907da351424
Author: Grzegorz KoĊakowski <gr...@getindata.com>
AuthorDate: Thu Mar 22 09:51:44 2018 +0100
[BEAM-622] Add checkpointing tests for DoFnOperator and WindowDoFnOperator
Timers restore test for DoFnOperator already exists.
---
.../flink/streaming/DedupingOperatorTest.java | 30 +--
.../runners/flink/streaming/DoFnOperatorTest.java | 177 ++++++++++------
.../flink/streaming/StreamRecordStripper.java | 49 +++++
.../flink/streaming/WindowDoFnOperatorTest.java | 227 +++++++++++++++++++++
4 files changed, 397 insertions(+), 86 deletions(-)
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/DedupingOperatorTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/DedupingOperatorTest.java
index 0b376e9..73be0ef 100644
--- a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/DedupingOperatorTest.java
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/DedupingOperatorTest.java
@@ -17,13 +17,11 @@
*/
package org.apache.beam.runners.flink.streaming;
+import static org.apache.beam.runners.flink.streaming.StreamRecordStripper.stripStreamRecordFromWindowedValue;
import static org.hamcrest.collection.IsIterableContainingInOrder.contains;
import static org.junit.Assert.assertThat;
-import com.google.common.base.Function;
-import com.google.common.collect.FluentIterable;
import java.nio.ByteBuffer;
-import javax.annotation.Nullable;
import org.apache.beam.runners.flink.translation.wrappers.streaming.io.DedupingOperator;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.ValueWithRecordId;
@@ -64,7 +62,7 @@ public class DedupingOperatorTest {
WindowedValue.valueInGlobalWindow(new ValueWithRecordId<>(key1, key1.getBytes()))));
assertThat(
- this.<String>stripStreamRecordFromWindowedValue(harness.getOutput()),
+ stripStreamRecordFromWindowedValue(harness.getOutput()),
contains(WindowedValue.valueInGlobalWindow(key1),
WindowedValue.valueInGlobalWindow(key2)));
@@ -86,7 +84,7 @@ public class DedupingOperatorTest {
WindowedValue.valueInGlobalWindow(new ValueWithRecordId<>(key3, key3.getBytes()))));
assertThat(
- this.<String>stripStreamRecordFromWindowedValue(harness.getOutput()),
+ stripStreamRecordFromWindowedValue(harness.getOutput()),
contains(WindowedValue.valueInGlobalWindow(key3)));
harness.close();
@@ -102,26 +100,4 @@ public class DedupingOperatorTest {
value -> ByteBuffer.wrap(value.getValue().getId()),
TypeInformation.of(ByteBuffer.class));
}
-
- private <T> Iterable<WindowedValue<T>> stripStreamRecordFromWindowedValue(
- Iterable<Object> input) {
-
- return FluentIterable.from(input)
- .filter(
- o ->
- o instanceof StreamRecord && ((StreamRecord) o).getValue() instanceof WindowedValue)
- .transform(
- new Function<Object, WindowedValue<T>>() {
- @Nullable
- @Override
- @SuppressWarnings({"unchecked", "rawtypes"})
- public WindowedValue<T> apply(@Nullable Object o) {
- if (o instanceof StreamRecord
- && ((StreamRecord) o).getValue() instanceof WindowedValue) {
- return (WindowedValue) ((StreamRecord) o).getValue();
- }
- throw new RuntimeException("unreachable");
- }
- });
- }
}
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/DoFnOperatorTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/DoFnOperatorTest.java
index 73a0a08..4d6fca6 100644
--- a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/DoFnOperatorTest.java
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/DoFnOperatorTest.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.runners.flink.streaming;
+import static org.apache.beam.runners.flink.streaming.StreamRecordStripper.stripStreamRecordFromWindowedValue;
import static org.hamcrest.Matchers.emptyIterable;
import static org.hamcrest.collection.IsIterableContainingInOrder.contains;
import static org.junit.Assert.assertEquals;
@@ -28,6 +29,7 @@ import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import java.util.Collections;
import java.util.HashMap;
+import java.util.Optional;
import javax.annotation.Nullable;
import org.apache.beam.runners.core.StatefulDoFnRunner;
import org.apache.beam.runners.flink.FlinkPipelineOptions;
@@ -39,6 +41,7 @@ import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.coders.VarLongCoder;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.state.StateSpecs;
@@ -65,6 +68,7 @@ import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.OperatorStateHandles;
@@ -133,7 +137,7 @@ public class DoFnOperatorTest {
testHarness.processElement(new StreamRecord<>(WindowedValue.valueInGlobalWindow("Hello")));
assertThat(
- this.<String>stripStreamRecordFromWindowedValue(testHarness.getOutput()),
+ stripStreamRecordFromWindowedValue(testHarness.getOutput()),
contains(WindowedValue.valueInGlobalWindow("Hello")));
testHarness.close();
@@ -259,7 +263,6 @@ public class DoFnOperatorTest {
StringUtf8Coder.of(),
windowingStrategy.getWindowFn().windowCoder());
-
TupleTag<String> outputTag = new TupleTag<>("main-output");
DoFnOperator<Integer, String> doFnOperator =
@@ -293,14 +296,14 @@ public class DoFnOperatorTest {
new StreamRecord<>(WindowedValue.of(13, new Instant(0), window1, PaneInfo.NO_FIRING)));
assertThat(
- this.<String>stripStreamRecordFromWindowedValue(testHarness.getOutput()),
+ stripStreamRecordFromWindowedValue(testHarness.getOutput()),
emptyIterable());
// this does not yet fire the timer (in vanilla Flink it would)
testHarness.processWatermark(timerTimestamp.getMillis());
assertThat(
- this.<String>stripStreamRecordFromWindowedValue(testHarness.getOutput()),
+ stripStreamRecordFromWindowedValue(testHarness.getOutput()),
emptyIterable());
testHarness.getOutput().clear();
@@ -309,7 +312,7 @@ public class DoFnOperatorTest {
testHarness.processWatermark(timerTimestamp.getMillis() + 1);
assertThat(
- this.<String>stripStreamRecordFromWindowedValue(testHarness.getOutput()),
+ stripStreamRecordFromWindowedValue(testHarness.getOutput()),
contains(
WindowedValue.of(
outputMessage, new Instant(timerTimestamp), window1, PaneInfo.NO_FIRING)));
@@ -372,7 +375,7 @@ public class DoFnOperatorTest {
new StreamRecord<>(WindowedValue.of(13, new Instant(0), window1, PaneInfo.NO_FIRING)));
assertThat(
- this.<String>stripStreamRecordFromWindowedValue(testHarness.getOutput()),
+ stripStreamRecordFromWindowedValue(testHarness.getOutput()),
contains(WindowedValue.of("13", new Instant(0), window1, PaneInfo.NO_FIRING)));
testHarness.getOutput().clear();
@@ -384,7 +387,7 @@ public class DoFnOperatorTest {
new StreamRecord<>(WindowedValue.of(17, new Instant(0), window1, PaneInfo.NO_FIRING)));
assertThat(
- this.<String>stripStreamRecordFromWindowedValue(testHarness.getOutput()),
+ stripStreamRecordFromWindowedValue(testHarness.getOutput()),
contains(WindowedValue.of("17", new Instant(0), window1, PaneInfo.NO_FIRING)));
testHarness.getOutput().clear();
@@ -396,7 +399,7 @@ public class DoFnOperatorTest {
new StreamRecord<>(WindowedValue.of(17, new Instant(0), window1, PaneInfo.NO_FIRING)));
assertThat(
- this.<String>stripStreamRecordFromWindowedValue(testHarness.getOutput()),
+ stripStreamRecordFromWindowedValue(testHarness.getOutput()),
emptyIterable());
testHarness.close();
@@ -488,7 +491,7 @@ public class DoFnOperatorTest {
WindowedValue.of(KV.of("key2", 7), new Instant(3), window1, PaneInfo.NO_FIRING)));
assertThat(
- this.<KV<String, Integer>>stripStreamRecordFromWindowedValue(testHarness.getOutput()),
+ stripStreamRecordFromWindowedValue(testHarness.getOutput()),
contains(
WindowedValue.of(
KV.of("key1", 5 + offset), new Instant(1), window1, PaneInfo.NO_FIRING),
@@ -510,7 +513,7 @@ public class DoFnOperatorTest {
.getMillis() + 1);
assertThat(
- this.<KV<String, Integer>>stripStreamRecordFromWindowedValue(testHarness.getOutput()),
+ stripStreamRecordFromWindowedValue(testHarness.getOutput()),
contains(
WindowedValue.of(
KV.of("key1", timerOutput), new Instant(9), window1, PaneInfo.NO_FIRING),
@@ -603,7 +606,7 @@ public class DoFnOperatorTest {
valuesInWindow(ImmutableList.of("foo", "bar"), new Instant(1000), secondWindow))));
assertThat(
- this.<String>stripStreamRecordFromWindowedValue(testHarness.getOutput()),
+ stripStreamRecordFromWindowedValue(testHarness.getOutput()),
contains(helloElement, worldElement));
testHarness.close();
@@ -611,6 +614,80 @@ public class DoFnOperatorTest {
}
@Test
+ public void testStateRestore() throws Exception {
+ DoFn<KV<String, Long>, KV<String, Long>> filterElementsEqualToCountFn =
+ new DoFn<KV<String, Long>, KV<String, Long>>() {
+
+ @StateId("counter")
+ private final StateSpec<ValueState<Long>> counterSpec = StateSpecs
+ .value(VarLongCoder.of());
+
+ @ProcessElement
+ public void processElement(ProcessContext context,
+ @StateId("counter") ValueState<Long> count) {
+ long currentCount = Optional.ofNullable(count.read()).orElse(0L);
+ currentCount = currentCount + 1;
+ count.write(currentCount);
+
+ KV<String, Long> currentElement = context.element();
+ if (currentCount == currentElement.getValue()) {
+ context.output(currentElement);
+ }
+ }
+ };
+
+ WindowingStrategy<Object, GlobalWindow> windowingStrategy = WindowingStrategy.globalDefault();
+
+ TupleTag<KV<String, Long>> outputTag = new TupleTag<>("main-output");
+
+ FullWindowedValueCoder<KV<String, Long>> kvCoder = WindowedValue.getFullCoder(
+ KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of()),
+ windowingStrategy.getWindowFn().windowCoder()
+ );
+
+ CoderTypeInformation<String> keyCoderInfo = new CoderTypeInformation<>(StringUtf8Coder.of());
+ KeySelector<WindowedValue<KV<String, Long>>, String> keySelector = e -> e.getValue().getKey();
+
+ OneInputStreamOperatorTestHarness<WindowedValue<KV<String, Long>>,
+ WindowedValue<KV<String, Long>>> testHarness = createTestHarness(windowingStrategy,
+ filterElementsEqualToCountFn, kvCoder, kvCoder, outputTag, keyCoderInfo, keySelector);
+ testHarness.open();
+
+ testHarness
+ .processElement(new StreamRecord<>(WindowedValue.valueInGlobalWindow(KV.of("a", 100L))));
+ testHarness
+ .processElement(new StreamRecord<>(WindowedValue.valueInGlobalWindow(KV.of("a", 100L))));
+
+ final OperatorStateHandles snapshot = testHarness.snapshot(0, 0);
+ testHarness.close();
+
+ testHarness = createTestHarness(windowingStrategy, filterElementsEqualToCountFn, kvCoder,
+ kvCoder, outputTag, keyCoderInfo, keySelector);
+ testHarness.initializeState(snapshot);
+ testHarness.open();
+
+ // after restore: counter = 2
+ testHarness
+ .processElement(new StreamRecord<>(WindowedValue.valueInGlobalWindow(KV.of("a", 100L))));
+ testHarness
+ .processElement(new StreamRecord<>(WindowedValue.valueInGlobalWindow(KV.of("a", 4L))));
+ testHarness
+ .processElement(new StreamRecord<>(WindowedValue.valueInGlobalWindow(KV.of("a", 5L))));
+ testHarness
+ .processElement(new StreamRecord<>(WindowedValue.valueInGlobalWindow(KV.of("a", 100L))));
+
+ assertThat(
+ stripStreamRecordFromWindowedValue(testHarness.getOutput()),
+ contains(
+ WindowedValue.valueInGlobalWindow(KV.of("a", 4L)),
+ WindowedValue.valueInGlobalWindow(KV.of("a", 5L))
+ )
+ );
+
+ testHarness.close();
+ }
+
+ @Test
public void testTimersRestore() throws Exception {
final Instant timerTimestamp = new Instant(1000);
final String outputMessage = "Timer fired";
@@ -649,11 +726,15 @@ public class DoFnOperatorTest {
TupleTag<String> outputTag = new TupleTag<>("main-output");
+
final CoderTypeSerializer<WindowedValue<String>> outputSerializer = new CoderTypeSerializer<>(
outputCoder);
+ CoderTypeInformation<Integer> keyCoderInfo = new CoderTypeInformation<>(VarIntCoder.of());
+ KeySelector<WindowedValue<Integer>, Integer> keySelector = WindowedValue::getValue;
OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<String>> testHarness =
- createTestHarness(windowingStrategy, fn, inputCoder, outputCoder, outputTag);
+ createTestHarness(windowingStrategy, fn, inputCoder, outputCoder, outputTag, keyCoderInfo,
+ keySelector);
testHarness.setup(outputSerializer);
@@ -668,14 +749,15 @@ public class DoFnOperatorTest {
new StreamRecord<>(WindowedValue.of(13, new Instant(0), window1, PaneInfo.NO_FIRING)));
assertThat(
- this.<String>stripStreamRecordFromWindowedValue(testHarness.getOutput()),
+ stripStreamRecordFromWindowedValue(testHarness.getOutput()),
emptyIterable());
// snapshot and restore
final OperatorStateHandles snapshot = testHarness.snapshot(0, 0);
testHarness.close();
- testHarness = createTestHarness(windowingStrategy, fn, inputCoder, outputCoder, outputTag);
+ testHarness = createTestHarness(windowingStrategy, fn, inputCoder, outputCoder, outputTag,
+ keyCoderInfo, keySelector);
testHarness.setup(outputSerializer);
testHarness.initializeState(snapshot);
testHarness.open();
@@ -684,7 +766,7 @@ public class DoFnOperatorTest {
testHarness.processWatermark(timerTimestamp.getMillis() + 1);
assertThat(
- this.<String>stripStreamRecordFromWindowedValue(testHarness.getOutput()),
+ stripStreamRecordFromWindowedValue(testHarness.getOutput()),
contains(
WindowedValue.of(
outputMessage, new Instant(timerTimestamp), window1, PaneInfo.NO_FIRING)));
@@ -692,26 +774,25 @@ public class DoFnOperatorTest {
testHarness.close();
}
- private OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<String>>
- createTestHarness(WindowingStrategy<Object, IntervalWindow> windowingStrategy,
- DoFn<Integer, String> fn, FullWindowedValueCoder<Integer> inputCoder,
- FullWindowedValueCoder<String> outputCoder, TupleTag<String> outputTag) throws Exception {
- DoFnOperator<Integer, String> doFnOperator =
- new DoFnOperator<>(
- fn,
- "stepName",
- inputCoder,
- outputTag,
- Collections.emptyList(),
- new DoFnOperator.MultiOutputOutputManagerFactory<>(outputTag, outputCoder),
- windowingStrategy,
- new HashMap<>(), /* side-input mapping */
- Collections.emptyList(), /* side inputs */
- PipelineOptionsFactory.as(FlinkPipelineOptions.class),
- VarIntCoder.of() /* key coder */);
-
- return new KeyedOneInputStreamOperatorTestHarness<>(
- doFnOperator, WindowedValue::getValue, new CoderTypeInformation<>(VarIntCoder.of()));
+ private <K, InT, OutT> OneInputStreamOperatorTestHarness<WindowedValue<InT>, WindowedValue<OutT>>
+ createTestHarness(WindowingStrategy<Object, ?> windowingStrategy, DoFn<InT, OutT> fn,
+ FullWindowedValueCoder<InT> inputCoder, FullWindowedValueCoder<OutT> outputCoder,
+ TupleTag<OutT> outputTag, TypeInformation<K> keyCoderInfo,
+ KeySelector<WindowedValue<InT>, K> keySelector) throws Exception {
+ DoFnOperator<InT, OutT> doFnOperator = new DoFnOperator<>(
+ fn,
+ "stepName",
+ inputCoder,
+ outputTag,
+ Collections.emptyList(),
+ new DoFnOperator.MultiOutputOutputManagerFactory<>(outputTag, outputCoder),
+ windowingStrategy,
+ new HashMap<>(), /* side-input mapping */
+ Collections.emptyList(), /* side inputs */
+ PipelineOptionsFactory.as(FlinkPipelineOptions.class),
+ VarIntCoder.of() /* key coder */);
+
+ return new KeyedOneInputStreamOperatorTestHarness<>(doFnOperator, keySelector, keyCoderInfo);
}
/**
@@ -783,7 +864,7 @@ public class DoFnOperatorTest {
// There is a finishBundle in snapshot()
// Elements will be buffered as part of finishing a bundle in snapshot()
assertThat(
- this.<String>stripStreamRecordFromWindowedValue(testHarness.getOutput()),
+ stripStreamRecordFromWindowedValue(testHarness.getOutput()),
contains(
WindowedValue.valueInGlobalWindow("a"),
WindowedValue.valueInGlobalWindow("b"),
@@ -821,7 +902,7 @@ public class DoFnOperatorTest {
newHarness.setProcessingTime(10);
assertThat(
- this.<String>stripStreamRecordFromWindowedValue(newHarness.getOutput()),
+ stripStreamRecordFromWindowedValue(newHarness.getOutput()),
contains(
WindowedValue.valueInGlobalWindow("finishBundle"),
WindowedValue.valueInGlobalWindow("d"),
@@ -830,28 +911,6 @@ public class DoFnOperatorTest {
newHarness.close();
}
- private <T> Iterable<WindowedValue<T>> stripStreamRecordFromWindowedValue(
- Iterable<Object> input) {
-
- return FluentIterable.from(input)
- .filter(
- o ->
- o instanceof StreamRecord && ((StreamRecord) o).getValue() instanceof WindowedValue)
- .transform(
- new Function<Object, WindowedValue<T>>() {
- @Nullable
- @Override
- @SuppressWarnings({"unchecked", "rawtypes"})
- public WindowedValue<T> apply(@Nullable Object o) {
- if (o instanceof StreamRecord
- && ((StreamRecord) o).getValue() instanceof WindowedValue) {
- return (WindowedValue) ((StreamRecord) o).getValue();
- }
- throw new RuntimeException("unreachable");
- }
- });
- }
-
private Iterable<WindowedValue<String>> stripStreamRecord(Iterable<?> input) {
return FluentIterable.from(input)
.filter(o -> o instanceof StreamRecord)
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/StreamRecordStripper.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/StreamRecordStripper.java
new file mode 100644
index 0000000..aebb414
--- /dev/null
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/StreamRecordStripper.java
@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.flink.streaming;
+
+import com.google.common.base.Function;
+import com.google.common.collect.FluentIterable;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+
+class StreamRecordStripper {
+
+ @SuppressWarnings("Guava")
+ static <T> Iterable<WindowedValue<T>> stripStreamRecordFromWindowedValue(Iterable<Object> input) {
+ return FluentIterable.from(input)
+ .filter(
+ o ->
+ o instanceof StreamRecord && ((StreamRecord) o).getValue() instanceof WindowedValue)
+ .transform(
+ new Function<Object, WindowedValue<T>>() {
+ @Nullable
+ @Override
+ @SuppressWarnings({"unchecked", "rawtypes"})
+ public WindowedValue<T> apply(@Nullable Object o) {
+ if (o instanceof StreamRecord
+ && ((StreamRecord) o).getValue() instanceof WindowedValue) {
+ return (WindowedValue) ((StreamRecord) o).getValue();
+ }
+ throw new RuntimeException("unreachable");
+ }
+ });
+ }
+
+}
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/WindowDoFnOperatorTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/WindowDoFnOperatorTest.java
new file mode 100644
index 0000000..5a3717d
--- /dev/null
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/WindowDoFnOperatorTest.java
@@ -0,0 +1,227 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.flink.streaming;
+
+import static java.util.Collections.emptyList;
+import static java.util.Collections.emptyMap;
+import static org.apache.beam.runners.flink.streaming.StreamRecordStripper.stripStreamRecordFromWindowedValue;
+import static org.apache.beam.sdk.transforms.windowing.PaneInfo.NO_FIRING;
+import static org.apache.beam.sdk.transforms.windowing.PaneInfo.Timing.ON_TIME;
+import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.joda.time.Duration.standardMinutes;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThat;
+
+import com.google.common.collect.Iterables;
+import java.io.ByteArrayOutputStream;
+import java.nio.ByteBuffer;
+import org.apache.beam.runners.core.KeyedWorkItem;
+import org.apache.beam.runners.core.SystemReduceFn;
+import org.apache.beam.runners.flink.FlinkPipelineOptions;
+import org.apache.beam.runners.flink.translation.wrappers.streaming.DoFnOperator.MultiOutputOutputManagerFactory;
+import org.apache.beam.runners.flink.translation.wrappers.streaming.SingletonKeyedWorkItem;
+import org.apache.beam.runners.flink.translation.wrappers.streaming.SingletonKeyedWorkItemCoder;
+import org.apache.beam.runners.flink.translation.wrappers.streaming.WindowDoFnOperator;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.coders.VarLongCoder;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.transforms.Sum;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.FixedWindows;
+import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.AppliedCombineFn;
+import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.beam.sdk.util.WindowedValue.FullWindowedValueCoder;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.WindowingStrategy;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.GenericTypeInfo;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.OperatorStateHandles;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Tests for {@link WindowDoFnOperator}.
+ */
+@RunWith(JUnit4.class)
+public class WindowDoFnOperatorTest {
+
+ @Test
+ public void testRestore() throws Exception {
+ // test harness
+ KeyedOneInputStreamOperatorTestHarness<ByteBuffer, WindowedValue<KeyedWorkItem<Long, Long>>,
+ WindowedValue<KV<Long, Long>>> testHarness = createTestHarness(getWindowDoFnOperator());
+ testHarness.open();
+
+ // process elements
+ IntervalWindow window = new IntervalWindow(new Instant(0), Duration.millis(10_000));
+ testHarness.processWatermark(0L);
+ testHarness.processElement(
+ Item.builder().key(1L).timestamp(1L).value(100L).window(window).build().toStreamRecord());
+ testHarness.processElement(
+ Item.builder().key(1L).timestamp(2L).value(20L).window(window).build().toStreamRecord());
+ testHarness.processElement(
+ Item.builder().key(2L).timestamp(3L).value(77L).window(window).build().toStreamRecord());
+
+ // create snapshot
+ OperatorStateHandles snapshot = testHarness.snapshot(0, 0);
+ testHarness.close();
+
+ // restore from the snapshot
+ testHarness = createTestHarness(getWindowDoFnOperator());
+ testHarness.initializeState(snapshot);
+ testHarness.open();
+
+ // close window
+ testHarness.processWatermark(10_000L);
+
+ Iterable<WindowedValue<KV<Long, Long>>> output = stripStreamRecordFromWindowedValue(
+ testHarness.getOutput());
+
+ assertEquals(2, Iterables.size(output));
+ assertThat(output, containsInAnyOrder(
+ WindowedValue.of(KV.of(1L, 120L), new Instant(9_999), window,
+ PaneInfo.createPane(true, true, ON_TIME)),
+ WindowedValue.of(KV.of(2L, 77L), new Instant(9_999), window,
+ PaneInfo.createPane(true, true, ON_TIME))
+ )
+ );
+ // cleanup
+ testHarness.close();
+ }
+
+ private WindowDoFnOperator<Long, Long, Long> getWindowDoFnOperator() {
+ WindowingStrategy<Object, IntervalWindow> windowingStrategy =
+ WindowingStrategy.of(FixedWindows.of(standardMinutes(1)));
+
+ TupleTag<KV<Long, Long>> outputTag = new TupleTag<>("main-output");
+
+ SystemReduceFn<Long, Long, long[], Long, BoundedWindow> reduceFn = SystemReduceFn.combining(
+ VarLongCoder.of(),
+ AppliedCombineFn.withInputCoder(
+ Sum.ofLongs(),
+ CoderRegistry.createDefault(),
+ KvCoder.of(VarLongCoder.of(), VarLongCoder.of())
+ )
+ );
+
+ Coder<IntervalWindow> windowCoder = windowingStrategy.getWindowFn().windowCoder();
+ SingletonKeyedWorkItemCoder<Long, Long> workItemCoder = SingletonKeyedWorkItemCoder
+ .of(VarLongCoder.of(), VarLongCoder.of(), windowCoder);
+ FullWindowedValueCoder<SingletonKeyedWorkItem<Long, Long>> inputCoder = WindowedValue
+ .getFullCoder(workItemCoder, windowCoder);
+ FullWindowedValueCoder<KV<Long, Long>> outputCoder = WindowedValue
+ .getFullCoder(KvCoder.of(VarLongCoder.of(), VarLongCoder.of()), windowCoder);
+
+ return new WindowDoFnOperator<Long, Long, Long>(
+ reduceFn,
+ "stepName",
+ (Coder) inputCoder,
+ outputTag,
+ emptyList(),
+ new MultiOutputOutputManagerFactory<>(outputTag, outputCoder),
+ windowingStrategy,
+ emptyMap(),
+ emptyList(),
+ PipelineOptionsFactory.as(FlinkPipelineOptions.class),
+ VarLongCoder.of()
+ );
+ }
+
+ private KeyedOneInputStreamOperatorTestHarness<ByteBuffer,
+ WindowedValue<KeyedWorkItem<Long, Long>>, WindowedValue<KV<Long, Long>>> createTestHarness(
+ WindowDoFnOperator<Long, Long, Long> windowDoFnOperator) throws Exception {
+ return new KeyedOneInputStreamOperatorTestHarness<>(
+ windowDoFnOperator,
+ (KeySelector<WindowedValue<KeyedWorkItem<Long, Long>>, ByteBuffer>) o -> {
+ try (ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
+ VarLongCoder.of().encode(o.getValue().key(), baos);
+ return ByteBuffer.wrap(baos.toByteArray());
+ }
+ },
+ new GenericTypeInfo<>(ByteBuffer.class)
+ );
+ }
+
+ private static class Item {
+
+ static ItemBuilder builder() {
+ return new ItemBuilder();
+ }
+
+ private long key;
+ private long value;
+ private long timestamp;
+ private IntervalWindow window;
+
+ StreamRecord<WindowedValue<KeyedWorkItem<Long, Long>>> toStreamRecord() {
+ WindowedValue<Long> item = WindowedValue.of(value, new Instant(timestamp), window, NO_FIRING);
+ WindowedValue<KeyedWorkItem<Long, Long>> keyedItem = WindowedValue
+ .of(new SingletonKeyedWorkItem<>(key, item), new Instant(timestamp), window, NO_FIRING);
+ return new StreamRecord<>(keyedItem);
+ }
+
+ private static final class ItemBuilder {
+
+ private long key;
+ private long value;
+ private long timestamp;
+ private IntervalWindow window;
+
+ ItemBuilder key(long key) {
+ this.key = key;
+ return this;
+ }
+
+ ItemBuilder value(long value) {
+ this.value = value;
+ return this;
+ }
+
+ ItemBuilder timestamp(long timestamp) {
+ this.timestamp = timestamp;
+ return this;
+ }
+
+ ItemBuilder window(IntervalWindow window) {
+ this.window = window;
+ return this;
+ }
+
+ Item build() {
+ Item item = new Item();
+ item.key = this.key;
+ item.value = this.value;
+ item.window = this.window;
+ item.timestamp = this.timestamp;
+ return item;
+ }
+ }
+
+ }
+
+}
--
To stop receiving notification emails like this one, please contact
aljoscha@apache.org.