You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by jk...@apache.org on 2017/05/18 23:30:16 UTC
[1/4] beam git commit: Splits SplittableParDo into a
core-construction part and a KWI-related part
Repository: beam
Updated Branches:
refs/heads/master 2e0cf0083 -> c93b3aa3d
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SplittableDoFnOperator.java
----------------------------------------------------------------------
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SplittableDoFnOperator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SplittableDoFnOperator.java
index 968fc0a..4ac2ff5 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SplittableDoFnOperator.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/SplittableDoFnOperator.java
@@ -26,16 +26,16 @@ import java.util.Map;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
-import org.apache.beam.runners.core.ElementAndRestriction;
import org.apache.beam.runners.core.KeyedWorkItem;
import org.apache.beam.runners.core.KeyedWorkItems;
import org.apache.beam.runners.core.OutputAndTimeBoundedSplittableProcessElementInvoker;
import org.apache.beam.runners.core.OutputWindowedValue;
-import org.apache.beam.runners.core.SplittableParDo;
+import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems.ProcessFn;
import org.apache.beam.runners.core.StateInternals;
import org.apache.beam.runners.core.StateInternalsFactory;
import org.apache.beam.runners.core.TimerInternals;
import org.apache.beam.runners.core.TimerInternalsFactory;
+import org.apache.beam.runners.core.construction.ElementAndRestriction;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.transforms.DoFn;
@@ -94,7 +94,7 @@ public class SplittableDoFnOperator<
public void open() throws Exception {
super.open();
- checkState(doFn instanceof SplittableParDo.ProcessFn);
+ checkState(doFn instanceof ProcessFn);
StateInternalsFactory<String> stateInternalsFactory = new StateInternalsFactory<String>() {
@Override
@@ -114,9 +114,9 @@ public class SplittableDoFnOperator<
executorService = Executors.newSingleThreadScheduledExecutor(Executors.defaultThreadFactory());
- ((SplittableParDo.ProcessFn) doFn).setStateInternalsFactory(stateInternalsFactory);
- ((SplittableParDo.ProcessFn) doFn).setTimerInternalsFactory(timerInternalsFactory);
- ((SplittableParDo.ProcessFn) doFn).setProcessElementInvoker(
+ ((ProcessFn) doFn).setStateInternalsFactory(stateInternalsFactory);
+ ((ProcessFn) doFn).setTimerInternalsFactory(timerInternalsFactory);
+ ((ProcessFn) doFn).setProcessElementInvoker(
new OutputAndTimeBoundedSplittableProcessElementInvoker<>(
doFn,
serializedOptions.getPipelineOptions(),
[2/4] beam git commit: Splits SplittableParDo into a
core-construction part and a KWI-related part
Posted by jk...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java
new file mode 100644
index 0000000..b38e364
--- /dev/null
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java
@@ -0,0 +1,418 @@
+/*
+ * 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.core;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.Iterables;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.runners.core.construction.ElementAndRestriction;
+import org.apache.beam.runners.core.construction.PTransformReplacements;
+import org.apache.beam.runners.core.construction.ReplacementOutputs;
+import org.apache.beam.runners.core.construction.SplittableParDo.ProcessKeyedElements;
+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.options.PipelineOptions;
+import org.apache.beam.sdk.runners.AppliedPTransform;
+import org.apache.beam.sdk.runners.PTransformOverrideFactory;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.ValueState;
+import org.apache.beam.sdk.state.WatermarkHoldState;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupByKey;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
+import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PValue;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.WindowingStrategy;
+import org.joda.time.Instant;
+
+/**
+ * Utilities for implementing {@link ProcessKeyedElements} using {@link KeyedWorkItem} and
+ * runner-specific {@link StateInternals} and {@link TimerInternals}.
+ */
+public class SplittableParDoViaKeyedWorkItems {
+ /**
+ * Runner-specific primitive {@link GroupByKey GroupByKey-like} {@link PTransform} that produces
+ * {@link KeyedWorkItem KeyedWorkItems} so that downstream transforms can access state and timers.
+ *
+ * <p>Unlike a real {@link GroupByKey}, ignores the input's windowing and triggering strategy and
+ * emits output immediately.
+ */
+ public static class GBKIntoKeyedWorkItems<KeyT, InputT>
+ extends PTransform<PCollection<KV<KeyT, InputT>>, PCollection<KeyedWorkItem<KeyT, InputT>>> {
+ @Override
+ public PCollection<KeyedWorkItem<KeyT, InputT>> expand(PCollection<KV<KeyT, InputT>> input) {
+ return PCollection.createPrimitiveOutputInternal(
+ input.getPipeline(), WindowingStrategy.globalDefault(), input.isBounded());
+ }
+ }
+
+ /** Overrides a {@link ProcessKeyedElements} into {@link SplittableProcessViaKeyedWorkItems}. */
+ public static class OverrideFactory<InputT, OutputT, RestrictionT>
+ implements PTransformOverrideFactory<
+ PCollection<KV<String, ElementAndRestriction<InputT, RestrictionT>>>, PCollectionTuple,
+ ProcessKeyedElements<InputT, OutputT, RestrictionT>> {
+ @Override
+ public PTransformReplacement<
+ PCollection<KV<String, ElementAndRestriction<InputT, RestrictionT>>>, PCollectionTuple>
+ getReplacementTransform(
+ AppliedPTransform<
+ PCollection<KV<String, ElementAndRestriction<InputT, RestrictionT>>>,
+ PCollectionTuple, ProcessKeyedElements<InputT, OutputT, RestrictionT>>
+ transform) {
+ return PTransformReplacement.of(
+ PTransformReplacements.getSingletonMainInput(transform),
+ new SplittableProcessViaKeyedWorkItems<>(transform.getTransform()));
+ }
+
+ @Override
+ public Map<PValue, ReplacementOutput> mapOutputs(
+ Map<TupleTag<?>, PValue> outputs, PCollectionTuple newOutput) {
+ return ReplacementOutputs.tagged(outputs, newOutput);
+ }
+ }
+
+ /**
+ * Runner-specific primitive {@link PTransform} that invokes the {@link DoFn.ProcessElement}
+ * method for a splittable {@link DoFn}.
+ */
+ public static class SplittableProcessViaKeyedWorkItems<InputT, OutputT, RestrictionT>
+ extends PTransform<
+ PCollection<KV<String, ElementAndRestriction<InputT, RestrictionT>>>, PCollectionTuple> {
+ private final ProcessKeyedElements<InputT, OutputT, RestrictionT> original;
+
+ public SplittableProcessViaKeyedWorkItems(
+ ProcessKeyedElements<InputT, OutputT, RestrictionT> original) {
+ this.original = original;
+ }
+
+ @Override
+ public PCollectionTuple expand(
+ PCollection<KV<String, ElementAndRestriction<InputT, RestrictionT>>> input) {
+ return input
+ .apply(new GBKIntoKeyedWorkItems<String, ElementAndRestriction<InputT, RestrictionT>>())
+ .setCoder(
+ KeyedWorkItemCoder.of(
+ StringUtf8Coder.of(),
+ ((KvCoder<String, ElementAndRestriction<InputT, RestrictionT>>) input.getCoder())
+ .getValueCoder(),
+ input.getWindowingStrategy().getWindowFn().windowCoder()))
+ .apply(new ProcessElements<>(original));
+ }
+ }
+
+ /** A primitive transform wrapping around {@link ProcessFn}. */
+ public static class ProcessElements<
+ InputT, OutputT, RestrictionT, TrackerT extends RestrictionTracker<RestrictionT>>
+ extends PTransform<
+ PCollection<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>>,
+ PCollectionTuple> {
+ private final ProcessKeyedElements<InputT, OutputT, RestrictionT> original;
+
+ public ProcessElements(ProcessKeyedElements<InputT, OutputT, RestrictionT> original) {
+ this.original = original;
+ }
+
+ public ProcessFn<InputT, OutputT, RestrictionT, TrackerT> newProcessFn(
+ DoFn<InputT, OutputT> fn) {
+ return new ProcessFn<>(
+ fn,
+ original.getElementCoder(),
+ original.getRestrictionCoder(),
+ original.getInputWindowingStrategy());
+ }
+
+ public DoFn<InputT, OutputT> getFn() {
+ return original.getFn();
+ }
+
+ public List<PCollectionView<?>> getSideInputs() {
+ return original.getSideInputs();
+ }
+
+ public TupleTag<OutputT> getMainOutputTag() {
+ return original.getMainOutputTag();
+ }
+
+ public TupleTagList getAdditionalOutputTags() {
+ return original.getAdditionalOutputTags();
+ }
+
+ @Override
+ public PCollectionTuple expand(
+ PCollection<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>> input) {
+ return ProcessKeyedElements.createPrimitiveOutputFor(
+ input,
+ original.getFn(),
+ original.getMainOutputTag(),
+ original.getAdditionalOutputTags(),
+ original.getInputWindowingStrategy());
+ }
+ }
+
+ /**
+ * The heart of splittable {@link DoFn} execution: processes a single (element, restriction) pair
+ * by creating a tracker for the restriction and checkpointing/resuming processing later if
+ * necessary.
+ *
+ * <p>Takes {@link KeyedWorkItem} and assumes that the KeyedWorkItem contains a single element (or
+ * a single timer set by {@link ProcessFn itself}, in a single window. This is necessary because
+ * {@link ProcessFn} sets timers, and timers are namespaced to a single window and it should be
+ * the window of the input element.
+ *
+ * <p>See also: https://issues.apache.org/jira/browse/BEAM-1983
+ */
+ @VisibleForTesting
+ public static class ProcessFn<
+ InputT, OutputT, RestrictionT, TrackerT extends RestrictionTracker<RestrictionT>>
+ extends DoFn<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT> {
+ /**
+ * The state cell containing a watermark hold for the output of this {@link DoFn}. The hold is
+ * acquired during the first {@link DoFn.ProcessElement} call for each element and restriction,
+ * and is released when the {@link DoFn.ProcessElement} call returns and there is no residual
+ * restriction captured by the {@link SplittableProcessElementInvoker}.
+ *
+ * <p>A hold is needed to avoid letting the output watermark immediately progress together with
+ * the input watermark when the first {@link DoFn.ProcessElement} call for this element
+ * completes.
+ */
+ private static final StateTag<WatermarkHoldState> watermarkHoldTag =
+ StateTags.makeSystemTagInternal(
+ StateTags.<GlobalWindow>watermarkStateInternal("hold", TimestampCombiner.LATEST));
+
+ /**
+ * The state cell containing a copy of the element. Written during the first {@link
+ * DoFn.ProcessElement} call and read during subsequent calls in response to timer firings, when
+ * the original element is no longer available.
+ */
+ private final StateTag<ValueState<WindowedValue<InputT>>> elementTag;
+
+ /**
+ * The state cell containing a restriction representing the unprocessed part of work for this
+ * element.
+ */
+ private StateTag<ValueState<RestrictionT>> restrictionTag;
+
+ private final DoFn<InputT, OutputT> fn;
+ private final Coder<InputT> elementCoder;
+ private final Coder<RestrictionT> restrictionCoder;
+ private final WindowingStrategy<InputT, ?> inputWindowingStrategy;
+
+ private transient StateInternalsFactory<String> stateInternalsFactory;
+ private transient TimerInternalsFactory<String> timerInternalsFactory;
+ private transient SplittableProcessElementInvoker<InputT, OutputT, RestrictionT, TrackerT>
+ processElementInvoker;
+
+ private transient DoFnInvoker<InputT, OutputT> invoker;
+
+ public ProcessFn(
+ DoFn<InputT, OutputT> fn,
+ Coder<InputT> elementCoder,
+ Coder<RestrictionT> restrictionCoder,
+ WindowingStrategy<InputT, ?> inputWindowingStrategy) {
+ this.fn = fn;
+ this.elementCoder = elementCoder;
+ this.restrictionCoder = restrictionCoder;
+ this.inputWindowingStrategy = inputWindowingStrategy;
+ this.elementTag =
+ StateTags.value(
+ "element",
+ WindowedValue.getFullCoder(
+ elementCoder, inputWindowingStrategy.getWindowFn().windowCoder()));
+ this.restrictionTag = StateTags.value("restriction", restrictionCoder);
+ }
+
+ public void setStateInternalsFactory(StateInternalsFactory<String> stateInternalsFactory) {
+ this.stateInternalsFactory = stateInternalsFactory;
+ }
+
+ public void setTimerInternalsFactory(TimerInternalsFactory<String> timerInternalsFactory) {
+ this.timerInternalsFactory = timerInternalsFactory;
+ }
+
+ public void setProcessElementInvoker(
+ SplittableProcessElementInvoker<InputT, OutputT, RestrictionT, TrackerT> invoker) {
+ this.processElementInvoker = invoker;
+ }
+
+ public DoFn<InputT, OutputT> getFn() {
+ return fn;
+ }
+
+ public Coder<InputT> getElementCoder() {
+ return elementCoder;
+ }
+
+ public Coder<RestrictionT> getRestrictionCoder() {
+ return restrictionCoder;
+ }
+
+ public WindowingStrategy<InputT, ?> getInputWindowingStrategy() {
+ return inputWindowingStrategy;
+ }
+
+ @Setup
+ public void setup() throws Exception {
+ invoker = DoFnInvokers.invokerFor(fn);
+ invoker.invokeSetup();
+ }
+
+ @Teardown
+ public void tearDown() throws Exception {
+ invoker.invokeTeardown();
+ }
+
+ @StartBundle
+ public void startBundle(StartBundleContext c) throws Exception {
+ invoker.invokeStartBundle(wrapContextAsStartBundle(c));
+ }
+
+ @FinishBundle
+ public void finishBundle(FinishBundleContext c) throws Exception {
+ invoker.invokeFinishBundle(wrapContextAsFinishBundle(c));
+ }
+
+ @ProcessElement
+ public void processElement(final ProcessContext c) {
+ String key = c.element().key();
+ StateInternals stateInternals = stateInternalsFactory.stateInternalsForKey(key);
+ TimerInternals timerInternals = timerInternalsFactory.timerInternalsForKey(key);
+
+ // Initialize state (element and restriction) depending on whether this is the seed call.
+ // The seed call is the first call for this element, which actually has the element.
+ // Subsequent calls are timer firings and the element has to be retrieved from the state.
+ TimerInternals.TimerData timer = Iterables.getOnlyElement(c.element().timersIterable(), null);
+ boolean isSeedCall = (timer == null);
+ StateNamespace stateNamespace;
+ if (isSeedCall) {
+ WindowedValue<ElementAndRestriction<InputT, RestrictionT>> windowedValue =
+ Iterables.getOnlyElement(c.element().elementsIterable());
+ BoundedWindow window = Iterables.getOnlyElement(windowedValue.getWindows());
+ stateNamespace =
+ StateNamespaces.window(
+ (Coder<BoundedWindow>) inputWindowingStrategy.getWindowFn().windowCoder(), window);
+ } else {
+ stateNamespace = timer.getNamespace();
+ }
+
+ ValueState<WindowedValue<InputT>> elementState =
+ stateInternals.state(stateNamespace, elementTag);
+ ValueState<RestrictionT> restrictionState =
+ stateInternals.state(stateNamespace, restrictionTag);
+ WatermarkHoldState holdState = stateInternals.state(stateNamespace, watermarkHoldTag);
+
+ ElementAndRestriction<WindowedValue<InputT>, RestrictionT> elementAndRestriction;
+ if (isSeedCall) {
+ WindowedValue<ElementAndRestriction<InputT, RestrictionT>> windowedValue =
+ Iterables.getOnlyElement(c.element().elementsIterable());
+ WindowedValue<InputT> element = windowedValue.withValue(windowedValue.getValue().element());
+ elementState.write(element);
+ elementAndRestriction =
+ ElementAndRestriction.of(element, windowedValue.getValue().restriction());
+ } else {
+ // This is not the first ProcessElement call for this element/restriction - rather,
+ // this is a timer firing, so we need to fetch the element and restriction from state.
+ elementState.readLater();
+ restrictionState.readLater();
+ elementAndRestriction =
+ ElementAndRestriction.of(elementState.read(), restrictionState.read());
+ }
+
+ final TrackerT tracker = invoker.invokeNewTracker(elementAndRestriction.restriction());
+ SplittableProcessElementInvoker<InputT, OutputT, RestrictionT, TrackerT>.Result result =
+ processElementInvoker.invokeProcessElement(
+ invoker, elementAndRestriction.element(), tracker);
+
+ // Save state for resuming.
+ if (result.getResidualRestriction() == null) {
+ // All work for this element/restriction is completed. Clear state and release hold.
+ elementState.clear();
+ restrictionState.clear();
+ holdState.clear();
+ return;
+ }
+ restrictionState.write(result.getResidualRestriction());
+ Instant futureOutputWatermark = result.getFutureOutputWatermark();
+ if (futureOutputWatermark == null) {
+ futureOutputWatermark = elementAndRestriction.element().getTimestamp();
+ }
+ holdState.add(futureOutputWatermark);
+ // Set a timer to continue processing this element.
+ timerInternals.setTimer(
+ TimerInternals.TimerData.of(
+ stateNamespace, timerInternals.currentProcessingTime(), TimeDomain.PROCESSING_TIME));
+ }
+
+ private DoFn<InputT, OutputT>.StartBundleContext wrapContextAsStartBundle(
+ final StartBundleContext baseContext) {
+ return fn.new StartBundleContext() {
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return baseContext.getPipelineOptions();
+ }
+
+ private void throwUnsupportedOutput() {
+ throw new UnsupportedOperationException(
+ String.format(
+ "Splittable DoFn can only output from @%s",
+ ProcessElement.class.getSimpleName()));
+ }
+ };
+ }
+
+ private DoFn<InputT, OutputT>.FinishBundleContext wrapContextAsFinishBundle(
+ final FinishBundleContext baseContext) {
+ return fn.new FinishBundleContext() {
+ @Override
+ public void output(OutputT output, Instant timestamp, BoundedWindow window) {
+ throwUnsupportedOutput();
+ }
+
+ @Override
+ public <T> void output(TupleTag<T> tag, T output, Instant timestamp, BoundedWindow window) {
+ throwUnsupportedOutput();
+ }
+
+ @Override
+ public PipelineOptions getPipelineOptions() {
+ return baseContext.getPipelineOptions();
+ }
+
+ private void throwUnsupportedOutput() {
+ throw new UnsupportedOperationException(
+ String.format(
+ "Splittable DoFn can only output from @%s",
+ ProcessElement.class.getSimpleName()));
+ }
+ };
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/core-java/src/test/java/org/apache/beam/runners/core/ElementAndRestrictionCoderTest.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/ElementAndRestrictionCoderTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/ElementAndRestrictionCoderTest.java
deleted file mode 100644
index f516046..0000000
--- a/runners/core-java/src/test/java/org/apache/beam/runners/core/ElementAndRestrictionCoderTest.java
+++ /dev/null
@@ -1,127 +0,0 @@
-/*
- * 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.core;
-
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.List;
-import org.apache.beam.sdk.coders.BigEndianLongCoder;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.CoderException;
-import org.apache.beam.sdk.coders.ListCoder;
-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.testing.CoderProperties;
-import org.apache.beam.sdk.util.CoderUtils;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.rules.ExpectedException;
-import org.junit.runner.RunWith;
-import org.junit.runners.Parameterized;
-import org.junit.runners.Parameterized.Parameter;
-
-/**
- * Tests for {@link ElementAndRestrictionCoder}. Parroted from {@link
- * org.apache.beam.sdk.coders.KvCoderTest}.
- */
-@RunWith(Parameterized.class)
-public class ElementAndRestrictionCoderTest<K, V> {
- private static class CoderAndData<T> {
- Coder<T> coder;
- List<T> data;
- }
-
- private static class AnyCoderAndData {
- private CoderAndData<?> coderAndData;
- }
-
- private static <T> AnyCoderAndData coderAndData(Coder<T> coder, List<T> data) {
- CoderAndData<T> coderAndData = new CoderAndData<>();
- coderAndData.coder = coder;
- coderAndData.data = data;
- AnyCoderAndData res = new AnyCoderAndData();
- res.coderAndData = coderAndData;
- return res;
- }
-
- private static final List<AnyCoderAndData> TEST_DATA =
- Arrays.asList(
- coderAndData(
- VarIntCoder.of(), Arrays.asList(-1, 0, 1, 13, Integer.MAX_VALUE, Integer.MIN_VALUE)),
- coderAndData(
- BigEndianLongCoder.of(),
- Arrays.asList(-1L, 0L, 1L, 13L, Long.MAX_VALUE, Long.MIN_VALUE)),
- coderAndData(StringUtf8Coder.of(), Arrays.asList("", "hello", "goodbye", "1")),
- coderAndData(
- ElementAndRestrictionCoder.of(StringUtf8Coder.of(), VarIntCoder.of()),
- Arrays.asList(
- ElementAndRestriction.of("", -1),
- ElementAndRestriction.of("hello", 0),
- ElementAndRestriction.of("goodbye", Integer.MAX_VALUE))),
- coderAndData(
- ListCoder.of(VarLongCoder.of()),
- Arrays.asList(Arrays.asList(1L, 2L, 3L), Collections.<Long>emptyList())));
-
- @Parameterized.Parameters(name = "{index}: keyCoder={0} key={1} valueCoder={2} value={3}")
- public static Collection<Object[]> data() {
- List<Object[]> parameters = new ArrayList<>();
- for (AnyCoderAndData keyCoderAndData : TEST_DATA) {
- Coder keyCoder = keyCoderAndData.coderAndData.coder;
- for (Object key : keyCoderAndData.coderAndData.data) {
- for (AnyCoderAndData valueCoderAndData : TEST_DATA) {
- Coder valueCoder = valueCoderAndData.coderAndData.coder;
- for (Object value : valueCoderAndData.coderAndData.data) {
- parameters.add(new Object[] {keyCoder, key, valueCoder, value});
- }
- }
- }
- }
- return parameters;
- }
-
- @Parameter(0)
- public Coder<K> keyCoder;
- @Parameter(1)
- public K key;
- @Parameter(2)
- public Coder<V> valueCoder;
- @Parameter(3)
- public V value;
-
- @Test
- @SuppressWarnings("rawtypes")
- public void testDecodeEncodeEqual() throws Exception {
- CoderProperties.coderDecodeEncodeEqual(
- ElementAndRestrictionCoder.of(keyCoder, valueCoder),
- ElementAndRestriction.of(key, value));
- }
-
- @Rule public ExpectedException thrown = ExpectedException.none();
-
- @Test
- public void encodeNullThrowsCoderException() throws Exception {
- thrown.expect(CoderException.class);
- thrown.expectMessage("cannot encode a null ElementAndRestriction");
-
- CoderUtils.encodeToBase64(
- ElementAndRestrictionCoder.of(StringUtf8Coder.of(), VarIntCoder.of()), null);
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/core-java/src/test/java/org/apache/beam/runners/core/SplittableParDoProcessFnTest.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/SplittableParDoProcessFnTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/SplittableParDoProcessFnTest.java
new file mode 100644
index 0000000..d242431
--- /dev/null
+++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/SplittableParDoProcessFnTest.java
@@ -0,0 +1,515 @@
+/*
+ * 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.core;
+
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+import static org.hamcrest.Matchers.hasItem;
+import static org.hamcrest.Matchers.hasItems;
+import static org.hamcrest.Matchers.not;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
+
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.NoSuchElementException;
+import java.util.concurrent.Executors;
+import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems.ProcessFn;
+import org.apache.beam.runners.core.construction.ElementAndRestriction;
+import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.InstantCoder;
+import org.apache.beam.sdk.coders.SerializableCoder;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFnTester;
+import org.apache.beam.sdk.transforms.splittabledofn.HasDefaultTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRange;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.FixedWindows;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
+import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.TimestampedValue;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.ValueInSingleWindow;
+import org.apache.beam.sdk.values.WindowingStrategy;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link SplittableParDoViaKeyedWorkItems.ProcessFn}. */
+@RunWith(JUnit4.class)
+public class SplittableParDoProcessFnTest {
+ private static final int MAX_OUTPUTS_PER_BUNDLE = 10000;
+ private static final Duration MAX_BUNDLE_DURATION = Duration.standardSeconds(5);
+
+ // ----------------- Tests for whether the transform sets boundedness correctly --------------
+ private static class SomeRestriction
+ implements Serializable, HasDefaultTracker<SomeRestriction, SomeRestrictionTracker> {
+ @Override
+ public SomeRestrictionTracker newTracker() {
+ return new SomeRestrictionTracker(this);
+ }
+ }
+
+ private static class SomeRestrictionTracker implements RestrictionTracker<SomeRestriction> {
+ private final SomeRestriction someRestriction;
+
+ public SomeRestrictionTracker(SomeRestriction someRestriction) {
+ this.someRestriction = someRestriction;
+ }
+
+ @Override
+ public SomeRestriction currentRestriction() {
+ return someRestriction;
+ }
+
+ @Override
+ public SomeRestriction checkpoint() {
+ return someRestriction;
+ }
+
+ @Override
+ public void checkDone() {}
+ }
+
+ @Rule
+ public TestPipeline pipeline = TestPipeline.create();
+
+ /**
+ * A helper for testing {@link ProcessFn} on 1 element (but
+ * possibly over multiple {@link DoFn.ProcessElement} calls).
+ */
+ private static class ProcessFnTester<
+ InputT, OutputT, RestrictionT, TrackerT extends RestrictionTracker<RestrictionT>>
+ implements AutoCloseable {
+ private final DoFnTester<
+ KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT>
+ tester;
+ private Instant currentProcessingTime;
+
+ private InMemoryTimerInternals timerInternals;
+ private TestInMemoryStateInternals<String> stateInternals;
+
+ ProcessFnTester(
+ Instant currentProcessingTime,
+ final DoFn<InputT, OutputT> fn,
+ Coder<InputT> inputCoder,
+ Coder<RestrictionT> restrictionCoder,
+ int maxOutputsPerBundle,
+ Duration maxBundleDuration)
+ throws Exception {
+ // The exact windowing strategy doesn't matter in this test, but it should be able to
+ // encode IntervalWindow's because that's what all tests here use.
+ WindowingStrategy<InputT, BoundedWindow> windowingStrategy =
+ (WindowingStrategy) WindowingStrategy.of(FixedWindows.of(Duration.standardSeconds(1)));
+ final ProcessFn<InputT, OutputT, RestrictionT, TrackerT> processFn =
+ new ProcessFn<>(
+ fn, inputCoder, restrictionCoder, windowingStrategy);
+ this.tester = DoFnTester.of(processFn);
+ this.timerInternals = new InMemoryTimerInternals();
+ this.stateInternals = new TestInMemoryStateInternals<>("dummy");
+ processFn.setStateInternalsFactory(
+ new StateInternalsFactory<String>() {
+ @Override
+ public StateInternals stateInternalsForKey(String key) {
+ return stateInternals;
+ }
+ });
+ processFn.setTimerInternalsFactory(
+ new TimerInternalsFactory<String>() {
+ @Override
+ public TimerInternals timerInternalsForKey(String key) {
+ return timerInternals;
+ }
+ });
+ processFn.setProcessElementInvoker(
+ new OutputAndTimeBoundedSplittableProcessElementInvoker<
+ InputT, OutputT, RestrictionT, TrackerT>(
+ fn,
+ tester.getPipelineOptions(),
+ new OutputWindowedValueToDoFnTester<>(tester),
+ new SideInputReader() {
+ @Override
+ public <T> T get(PCollectionView<T> view, BoundedWindow window) {
+ throw new NoSuchElementException();
+ }
+
+ @Override
+ public <T> boolean contains(PCollectionView<T> view) {
+ return false;
+ }
+
+ @Override
+ public boolean isEmpty() {
+ return true;
+ }
+ },
+ Executors.newSingleThreadScheduledExecutor(Executors.defaultThreadFactory()),
+ maxOutputsPerBundle,
+ maxBundleDuration));
+ // Do not clone since ProcessFn references non-serializable DoFnTester itself
+ // through the state/timer/output callbacks.
+ this.tester.setCloningBehavior(DoFnTester.CloningBehavior.DO_NOT_CLONE);
+ this.tester.startBundle();
+ timerInternals.advanceProcessingTime(currentProcessingTime);
+
+ this.currentProcessingTime = currentProcessingTime;
+ }
+
+ @Override
+ public void close() throws Exception {
+ tester.close();
+ }
+
+ /** Performs a seed {@link DoFn.ProcessElement} call feeding the element and restriction. */
+ void startElement(InputT element, RestrictionT restriction) throws Exception {
+ startElement(
+ WindowedValue.of(
+ ElementAndRestriction.of(element, restriction),
+ currentProcessingTime,
+ GlobalWindow.INSTANCE,
+ PaneInfo.ON_TIME_AND_ONLY_FIRING));
+ }
+
+ void startElement(WindowedValue<ElementAndRestriction<InputT, RestrictionT>> windowedValue)
+ throws Exception {
+ tester.processElement(
+ KeyedWorkItems.elementsWorkItem("key", Collections.singletonList(windowedValue)));
+ }
+
+ /**
+ * Advances processing time by a given duration and, if any timers fired, performs a non-seed
+ * {@link DoFn.ProcessElement} call, feeding it the timers.
+ */
+ boolean advanceProcessingTimeBy(Duration duration) throws Exception {
+ currentProcessingTime = currentProcessingTime.plus(duration);
+ timerInternals.advanceProcessingTime(currentProcessingTime);
+
+ List<TimerInternals.TimerData> timers = new ArrayList<>();
+ TimerInternals.TimerData nextTimer;
+ while ((nextTimer = timerInternals.removeNextProcessingTimer()) != null) {
+ timers.add(nextTimer);
+ }
+ if (timers.isEmpty()) {
+ return false;
+ }
+ tester.processElement(
+ KeyedWorkItems.<String, ElementAndRestriction<InputT, RestrictionT>>timersWorkItem(
+ "key", timers));
+ return true;
+ }
+
+ List<TimestampedValue<OutputT>> peekOutputElementsInWindow(BoundedWindow window) {
+ return tester.peekOutputElementsInWindow(window);
+ }
+
+ List<OutputT> takeOutputElements() {
+ return tester.takeOutputElements();
+ }
+
+ public Instant getWatermarkHold() {
+ return stateInternals.earliestWatermarkHold();
+ }
+ }
+
+ private static class OutputWindowedValueToDoFnTester<OutputT>
+ implements OutputWindowedValue<OutputT> {
+ private final DoFnTester<?, OutputT> tester;
+
+ private OutputWindowedValueToDoFnTester(DoFnTester<?, OutputT> tester) {
+ this.tester = tester;
+ }
+
+ @Override
+ public void outputWindowedValue(
+ OutputT output,
+ Instant timestamp,
+ Collection<? extends BoundedWindow> windows,
+ PaneInfo pane) {
+ outputWindowedValue(tester.getMainOutputTag(), output, timestamp, windows, pane);
+ }
+
+ @Override
+ public <AdditionalOutputT> void outputWindowedValue(
+ TupleTag<AdditionalOutputT> tag,
+ AdditionalOutputT output,
+ Instant timestamp,
+ Collection<? extends BoundedWindow> windows,
+ PaneInfo pane) {
+ for (BoundedWindow window : windows) {
+ tester.getMutableOutput(tag).add(ValueInSingleWindow.of(output, timestamp, window, pane));
+ }
+ }
+ }
+
+ /** A simple splittable {@link DoFn} that's actually monolithic. */
+ private static class ToStringFn extends DoFn<Integer, String> {
+ @ProcessElement
+ public void process(ProcessContext c, SomeRestrictionTracker tracker) {
+ c.output(c.element().toString() + "a");
+ c.output(c.element().toString() + "b");
+ c.output(c.element().toString() + "c");
+ }
+
+ @GetInitialRestriction
+ public SomeRestriction getInitialRestriction(Integer elem) {
+ return new SomeRestriction();
+ }
+ }
+
+ @Test
+ public void testTrivialProcessFnPropagatesOutputWindowAndTimestamp() throws Exception {
+ // Tests that ProcessFn correctly propagates the window and timestamp of the element
+ // inside the KeyedWorkItem.
+ // The underlying DoFn is actually monolithic, so this doesn't test splitting.
+ DoFn<Integer, String> fn = new ToStringFn();
+
+ Instant base = Instant.now();
+
+ IntervalWindow w =
+ new IntervalWindow(
+ base.minus(Duration.standardMinutes(1)), base.plus(Duration.standardMinutes(1)));
+
+ ProcessFnTester<Integer, String, SomeRestriction, SomeRestrictionTracker> tester =
+ new ProcessFnTester<>(
+ base,
+ fn,
+ BigEndianIntegerCoder.of(),
+ SerializableCoder.of(SomeRestriction.class),
+ MAX_OUTPUTS_PER_BUNDLE,
+ MAX_BUNDLE_DURATION);
+ tester.startElement(
+ WindowedValue.of(
+ ElementAndRestriction.of(42, new SomeRestriction()),
+ base,
+ Collections.singletonList(w),
+ PaneInfo.ON_TIME_AND_ONLY_FIRING));
+
+ assertEquals(
+ Arrays.asList(
+ TimestampedValue.of("42a", base),
+ TimestampedValue.of("42b", base),
+ TimestampedValue.of("42c", base)),
+ tester.peekOutputElementsInWindow(w));
+ }
+
+ private static class WatermarkUpdateFn extends DoFn<Instant, String> {
+ @ProcessElement
+ public void process(ProcessContext c, OffsetRangeTracker tracker) {
+ for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) {
+ c.updateWatermark(c.element().plus(Duration.standardSeconds(i)));
+ c.output(String.valueOf(i));
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange getInitialRestriction(Instant elem) {
+ throw new IllegalStateException("Expected to be supplied explicitly in this test");
+ }
+
+ @NewTracker
+ public OffsetRangeTracker newTracker(OffsetRange range) {
+ return new OffsetRangeTracker(range);
+ }
+ }
+
+ @Test
+ public void testUpdatesWatermark() throws Exception {
+ DoFn<Instant, String> fn = new WatermarkUpdateFn();
+ Instant base = Instant.now();
+
+ ProcessFnTester<Instant, String, OffsetRange, OffsetRangeTracker> tester =
+ new ProcessFnTester<>(
+ base,
+ fn,
+ InstantCoder.of(),
+ SerializableCoder.of(OffsetRange.class),
+ 3,
+ MAX_BUNDLE_DURATION);
+
+ tester.startElement(base, new OffsetRange(0, 8));
+ assertThat(tester.takeOutputElements(), hasItems("0", "1", "2"));
+ assertEquals(base.plus(Duration.standardSeconds(2)), tester.getWatermarkHold());
+
+ assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
+ assertThat(tester.takeOutputElements(), hasItems("3", "4", "5"));
+ assertEquals(base.plus(Duration.standardSeconds(5)), tester.getWatermarkHold());
+
+ assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
+ assertThat(tester.takeOutputElements(), hasItems("6", "7"));
+ assertEquals(null, tester.getWatermarkHold());
+ }
+
+ /**
+ * A splittable {@link DoFn} that generates the sequence [init, init + total).
+ */
+ private static class CounterFn extends DoFn<Integer, String> {
+ @ProcessElement
+ public void process(ProcessContext c, OffsetRangeTracker tracker) {
+ for (long i = tracker.currentRestriction().getFrom();
+ tracker.tryClaim(i); ++i) {
+ c.output(String.valueOf(c.element() + i));
+ }
+ }
+
+ @GetInitialRestriction
+ public OffsetRange getInitialRestriction(Integer elem) {
+ throw new UnsupportedOperationException("Expected to be supplied explicitly in this test");
+ }
+ }
+
+ @Test
+ public void testCheckpointsAfterNumOutputs() throws Exception {
+ int max = 100;
+ DoFn<Integer, String> fn = new CounterFn();
+ Instant base = Instant.now();
+ int baseIndex = 42;
+
+ ProcessFnTester<Integer, String, OffsetRange, OffsetRangeTracker> tester =
+ new ProcessFnTester<>(
+ base, fn, BigEndianIntegerCoder.of(), SerializableCoder.of(OffsetRange.class),
+ max, MAX_BUNDLE_DURATION);
+
+ List<String> elements;
+
+ // Create an fn that attempts to 2x output more than checkpointing allows.
+ tester.startElement(baseIndex, new OffsetRange(0, 2 * max + max / 2));
+ elements = tester.takeOutputElements();
+ assertEquals(max, elements.size());
+ // Should output the range [0, max)
+ assertThat(elements, hasItem(String.valueOf(baseIndex)));
+ assertThat(elements, hasItem(String.valueOf(baseIndex + max - 1)));
+
+ assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
+ elements = tester.takeOutputElements();
+ assertEquals(max, elements.size());
+ // Should output the range [max, 2*max)
+ assertThat(elements, hasItem(String.valueOf(baseIndex + max)));
+ assertThat(elements, hasItem(String.valueOf(baseIndex + 2 * max - 1)));
+
+ assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
+ elements = tester.takeOutputElements();
+ assertEquals(max / 2, elements.size());
+ // Should output the range [2*max, 2*max + max/2)
+ assertThat(elements, hasItem(String.valueOf(baseIndex + 2 * max)));
+ assertThat(elements, hasItem(String.valueOf(baseIndex + 2 * max + max / 2 - 1)));
+ assertThat(elements, not(hasItem((String.valueOf(baseIndex + 2 * max + max / 2)))));
+ }
+
+ @Test
+ public void testCheckpointsAfterDuration() throws Exception {
+ // Don't bound number of outputs.
+ int max = Integer.MAX_VALUE;
+ // But bound bundle duration - the bundle should terminate.
+ Duration maxBundleDuration = Duration.standardSeconds(1);
+ // Create an fn that attempts to 2x output more than checkpointing allows.
+ DoFn<Integer, String> fn = new CounterFn();
+ Instant base = Instant.now();
+ int baseIndex = 42;
+
+ ProcessFnTester<Integer, String, OffsetRange, OffsetRangeTracker> tester =
+ new ProcessFnTester<>(
+ base, fn, BigEndianIntegerCoder.of(), SerializableCoder.of(OffsetRange.class),
+ max, maxBundleDuration);
+
+ List<String> elements;
+
+ tester.startElement(baseIndex, new OffsetRange(0, Long.MAX_VALUE));
+ // Bundle should terminate, and should do at least some processing.
+ elements = tester.takeOutputElements();
+ assertFalse(elements.isEmpty());
+ // Bundle should have run for at least the requested duration.
+ assertThat(
+ Instant.now().getMillis() - base.getMillis(),
+ greaterThanOrEqualTo(maxBundleDuration.getMillis()));
+ }
+
+ private static class LifecycleVerifyingFn extends DoFn<Integer, String> {
+ private enum State {
+ BEFORE_SETUP,
+ OUTSIDE_BUNDLE,
+ INSIDE_BUNDLE,
+ TORN_DOWN
+ }
+
+ private State state = State.BEFORE_SETUP;
+
+ @ProcessElement
+ public void process(ProcessContext c, SomeRestrictionTracker tracker) {
+ assertEquals(State.INSIDE_BUNDLE, state);
+ }
+
+ @GetInitialRestriction
+ public SomeRestriction getInitialRestriction(Integer element) {
+ return new SomeRestriction();
+ }
+
+ @Setup
+ public void setup() {
+ assertEquals(State.BEFORE_SETUP, state);
+ state = State.OUTSIDE_BUNDLE;
+ }
+
+ @Teardown
+ public void tearDown() {
+ assertEquals(State.OUTSIDE_BUNDLE, state);
+ state = State.TORN_DOWN;
+ }
+
+ @StartBundle
+ public void startBundle() {
+ assertEquals(State.OUTSIDE_BUNDLE, state);
+ state = State.INSIDE_BUNDLE;
+ }
+
+ @FinishBundle
+ public void finishBundle() {
+ assertEquals(State.INSIDE_BUNDLE, state);
+ state = State.OUTSIDE_BUNDLE;
+ }
+ }
+
+ @Test
+ public void testInvokesLifecycleMethods() throws Exception {
+ DoFn<Integer, String> fn = new LifecycleVerifyingFn();
+ try (ProcessFnTester<Integer, String, SomeRestriction, SomeRestrictionTracker> tester =
+ new ProcessFnTester<>(
+ Instant.now(),
+ fn,
+ BigEndianIntegerCoder.of(),
+ SerializableCoder.of(SomeRestriction.class),
+ MAX_OUTPUTS_PER_BUNDLE,
+ MAX_BUNDLE_DURATION)) {
+ tester.startElement(42, new SomeRestriction());
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/core-java/src/test/java/org/apache/beam/runners/core/SplittableParDoTest.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/SplittableParDoTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/SplittableParDoTest.java
deleted file mode 100644
index be4cf08..0000000
--- a/runners/core-java/src/test/java/org/apache/beam/runners/core/SplittableParDoTest.java
+++ /dev/null
@@ -1,606 +0,0 @@
-/*
- * 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.core;
-
-import static org.hamcrest.Matchers.greaterThanOrEqualTo;
-import static org.hamcrest.Matchers.hasItem;
-import static org.hamcrest.Matchers.hasItems;
-import static org.hamcrest.Matchers.not;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertThat;
-import static org.junit.Assert.assertTrue;
-
-import java.io.Serializable;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.List;
-import java.util.NoSuchElementException;
-import java.util.concurrent.Executors;
-import javax.annotation.Nullable;
-import org.apache.beam.sdk.Pipeline;
-import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.InstantCoder;
-import org.apache.beam.sdk.coders.SerializableCoder;
-import org.apache.beam.sdk.testing.TestPipeline;
-import org.apache.beam.sdk.transforms.Create;
-import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.DoFn.BoundedPerElement;
-import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
-import org.apache.beam.sdk.transforms.DoFnTester;
-import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.transforms.splittabledofn.HasDefaultTracker;
-import org.apache.beam.sdk.transforms.splittabledofn.OffsetRange;
-import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
-import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.transforms.windowing.FixedWindows;
-import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
-import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
-import org.apache.beam.sdk.transforms.windowing.PaneInfo;
-import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.TimestampedValue;
-import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.sdk.values.TupleTagList;
-import org.apache.beam.sdk.values.ValueInSingleWindow;
-import org.apache.beam.sdk.values.WindowingStrategy;
-import org.joda.time.Duration;
-import org.joda.time.Instant;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-/** Tests for {@link SplittableParDo}. */
-@RunWith(JUnit4.class)
-public class SplittableParDoTest {
- private static final int MAX_OUTPUTS_PER_BUNDLE = 10000;
- private static final Duration MAX_BUNDLE_DURATION = Duration.standardSeconds(5);
-
- // ----------------- Tests for whether the transform sets boundedness correctly --------------
- private static class SomeRestriction
- implements Serializable, HasDefaultTracker<SomeRestriction, SomeRestrictionTracker> {
- @Override
- public SomeRestrictionTracker newTracker() {
- return new SomeRestrictionTracker(this);
- }
- }
-
- private static class SomeRestrictionTracker implements RestrictionTracker<SomeRestriction> {
- private final SomeRestriction someRestriction;
-
- public SomeRestrictionTracker(SomeRestriction someRestriction) {
- this.someRestriction = someRestriction;
- }
-
- @Override
- public SomeRestriction currentRestriction() {
- return someRestriction;
- }
-
- @Override
- public SomeRestriction checkpoint() {
- return someRestriction;
- }
-
- @Override
- public void checkDone() {}
- }
-
- @BoundedPerElement
- private static class BoundedFakeFn extends DoFn<Integer, String> {
- @ProcessElement
- public void processElement(ProcessContext context, SomeRestrictionTracker tracker) {}
-
- @GetInitialRestriction
- public SomeRestriction getInitialRestriction(Integer element) {
- return null;
- }
- }
-
- @UnboundedPerElement
- private static class UnboundedFakeFn extends DoFn<Integer, String> {
- @ProcessElement
- public void processElement(ProcessContext context, SomeRestrictionTracker tracker) {}
-
- @GetInitialRestriction
- public SomeRestriction getInitialRestriction(Integer element) {
- return null;
- }
- }
-
- private static PCollection<Integer> makeUnboundedCollection(Pipeline pipeline) {
- return pipeline
- .apply("unbounded", Create.of(1, 2, 3))
- .setIsBoundedInternal(PCollection.IsBounded.UNBOUNDED);
- }
-
- private static PCollection<Integer> makeBoundedCollection(Pipeline pipeline) {
- return pipeline
- .apply("bounded", Create.of(1, 2, 3))
- .setIsBoundedInternal(PCollection.IsBounded.BOUNDED);
- }
-
- private static final TupleTag<String> MAIN_OUTPUT_TAG = new TupleTag<String>() {};
-
- private ParDo.MultiOutput<Integer, String> makeParDo(DoFn<Integer, String> fn) {
- return ParDo.of(fn).withOutputTags(MAIN_OUTPUT_TAG, TupleTagList.empty());
- }
-
- @Rule
- public TestPipeline pipeline = TestPipeline.create();
-
- @Test
- public void testBoundednessForBoundedFn() {
- pipeline.enableAbandonedNodeEnforcement(false);
-
- DoFn<Integer, String> boundedFn = new BoundedFakeFn();
- assertEquals(
- "Applying a bounded SDF to a bounded collection produces a bounded collection",
- PCollection.IsBounded.BOUNDED,
- makeBoundedCollection(pipeline)
- .apply("bounded to bounded", new SplittableParDo<>(makeParDo(boundedFn)))
- .get(MAIN_OUTPUT_TAG)
- .isBounded());
- assertEquals(
- "Applying a bounded SDF to an unbounded collection produces an unbounded collection",
- PCollection.IsBounded.UNBOUNDED,
- makeUnboundedCollection(pipeline)
- .apply("bounded to unbounded", new SplittableParDo<>(makeParDo(boundedFn)))
- .get(MAIN_OUTPUT_TAG)
- .isBounded());
- }
-
- @Test
- public void testBoundednessForUnboundedFn() {
- pipeline.enableAbandonedNodeEnforcement(false);
-
- DoFn<Integer, String> unboundedFn = new UnboundedFakeFn();
- assertEquals(
- "Applying an unbounded SDF to a bounded collection produces a bounded collection",
- PCollection.IsBounded.UNBOUNDED,
- makeBoundedCollection(pipeline)
- .apply("unbounded to bounded", new SplittableParDo<>(makeParDo(unboundedFn)))
- .get(MAIN_OUTPUT_TAG)
- .isBounded());
- assertEquals(
- "Applying an unbounded SDF to an unbounded collection produces an unbounded collection",
- PCollection.IsBounded.UNBOUNDED,
- makeUnboundedCollection(pipeline)
- .apply("unbounded to unbounded", new SplittableParDo<>(makeParDo(unboundedFn)))
- .get(MAIN_OUTPUT_TAG)
- .isBounded());
- }
-
- // ------------------------------- Tests for ProcessFn ---------------------------------
-
- /**
- * A helper for testing {@link SplittableParDo.ProcessFn} on 1 element (but possibly over multiple
- * {@link DoFn.ProcessElement} calls).
- */
- private static class ProcessFnTester<
- InputT, OutputT, RestrictionT, TrackerT extends RestrictionTracker<RestrictionT>>
- implements AutoCloseable {
- private final DoFnTester<
- KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT>
- tester;
- private Instant currentProcessingTime;
-
- private InMemoryTimerInternals timerInternals;
- private TestInMemoryStateInternals<String> stateInternals;
-
- ProcessFnTester(
- Instant currentProcessingTime,
- final DoFn<InputT, OutputT> fn,
- Coder<InputT> inputCoder,
- Coder<RestrictionT> restrictionCoder,
- int maxOutputsPerBundle,
- Duration maxBundleDuration)
- throws Exception {
- // The exact windowing strategy doesn't matter in this test, but it should be able to
- // encode IntervalWindow's because that's what all tests here use.
- WindowingStrategy<InputT, BoundedWindow> windowingStrategy =
- (WindowingStrategy) WindowingStrategy.of(FixedWindows.of(Duration.standardSeconds(1)));
- final SplittableParDo.ProcessFn<InputT, OutputT, RestrictionT, TrackerT> processFn =
- new SplittableParDo.ProcessFn<>(
- fn, inputCoder, restrictionCoder, windowingStrategy);
- this.tester = DoFnTester.of(processFn);
- this.timerInternals = new InMemoryTimerInternals();
- this.stateInternals = new TestInMemoryStateInternals<>("dummy");
- processFn.setStateInternalsFactory(
- new StateInternalsFactory<String>() {
- @Override
- public StateInternals stateInternalsForKey(String key) {
- return stateInternals;
- }
- });
- processFn.setTimerInternalsFactory(
- new TimerInternalsFactory<String>() {
- @Override
- public TimerInternals timerInternalsForKey(String key) {
- return timerInternals;
- }
- });
- processFn.setProcessElementInvoker(
- new OutputAndTimeBoundedSplittableProcessElementInvoker<
- InputT, OutputT, RestrictionT, TrackerT>(
- fn,
- tester.getPipelineOptions(),
- new OutputWindowedValueToDoFnTester<>(tester),
- new SideInputReader() {
- @Nullable
- @Override
- public <T> T get(PCollectionView<T> view, BoundedWindow window) {
- throw new NoSuchElementException();
- }
-
- @Override
- public <T> boolean contains(PCollectionView<T> view) {
- return false;
- }
-
- @Override
- public boolean isEmpty() {
- return true;
- }
- },
- Executors.newSingleThreadScheduledExecutor(Executors.defaultThreadFactory()),
- maxOutputsPerBundle,
- maxBundleDuration));
- // Do not clone since ProcessFn references non-serializable DoFnTester itself
- // through the state/timer/output callbacks.
- this.tester.setCloningBehavior(DoFnTester.CloningBehavior.DO_NOT_CLONE);
- this.tester.startBundle();
- timerInternals.advanceProcessingTime(currentProcessingTime);
-
- this.currentProcessingTime = currentProcessingTime;
- }
-
- @Override
- public void close() throws Exception {
- tester.close();
- }
-
- /** Performs a seed {@link DoFn.ProcessElement} call feeding the element and restriction. */
- void startElement(InputT element, RestrictionT restriction) throws Exception {
- startElement(
- WindowedValue.of(
- ElementAndRestriction.of(element, restriction),
- currentProcessingTime,
- GlobalWindow.INSTANCE,
- PaneInfo.ON_TIME_AND_ONLY_FIRING));
- }
-
- void startElement(WindowedValue<ElementAndRestriction<InputT, RestrictionT>> windowedValue)
- throws Exception {
- tester.processElement(
- KeyedWorkItems.elementsWorkItem("key", Collections.singletonList(windowedValue)));
- }
-
- /**
- * Advances processing time by a given duration and, if any timers fired, performs a non-seed
- * {@link DoFn.ProcessElement} call, feeding it the timers.
- */
- boolean advanceProcessingTimeBy(Duration duration) throws Exception {
- currentProcessingTime = currentProcessingTime.plus(duration);
- timerInternals.advanceProcessingTime(currentProcessingTime);
-
- List<TimerInternals.TimerData> timers = new ArrayList<>();
- TimerInternals.TimerData nextTimer;
- while ((nextTimer = timerInternals.removeNextProcessingTimer()) != null) {
- timers.add(nextTimer);
- }
- if (timers.isEmpty()) {
- return false;
- }
- tester.processElement(
- KeyedWorkItems.<String, ElementAndRestriction<InputT, RestrictionT>>timersWorkItem(
- "key", timers));
- return true;
- }
-
- List<TimestampedValue<OutputT>> peekOutputElementsInWindow(BoundedWindow window) {
- return tester.peekOutputElementsInWindow(window);
- }
-
- List<OutputT> takeOutputElements() {
- return tester.takeOutputElements();
- }
-
- public Instant getWatermarkHold() {
- return stateInternals.earliestWatermarkHold();
- }
- }
-
- private static class OutputWindowedValueToDoFnTester<OutputT>
- implements OutputWindowedValue<OutputT> {
- private final DoFnTester<?, OutputT> tester;
-
- private OutputWindowedValueToDoFnTester(DoFnTester<?, OutputT> tester) {
- this.tester = tester;
- }
-
- @Override
- public void outputWindowedValue(
- OutputT output,
- Instant timestamp,
- Collection<? extends BoundedWindow> windows,
- PaneInfo pane) {
- outputWindowedValue(tester.getMainOutputTag(), output, timestamp, windows, pane);
- }
-
- @Override
- public <AdditionalOutputT> void outputWindowedValue(
- TupleTag<AdditionalOutputT> tag,
- AdditionalOutputT output,
- Instant timestamp,
- Collection<? extends BoundedWindow> windows,
- PaneInfo pane) {
- for (BoundedWindow window : windows) {
- tester.getMutableOutput(tag).add(ValueInSingleWindow.of(output, timestamp, window, pane));
- }
- }
- }
-
- /** A simple splittable {@link DoFn} that's actually monolithic. */
- private static class ToStringFn extends DoFn<Integer, String> {
- @ProcessElement
- public void process(ProcessContext c, SomeRestrictionTracker tracker) {
- c.output(c.element().toString() + "a");
- c.output(c.element().toString() + "b");
- c.output(c.element().toString() + "c");
- }
-
- @GetInitialRestriction
- public SomeRestriction getInitialRestriction(Integer elem) {
- return new SomeRestriction();
- }
- }
-
- @Test
- public void testTrivialProcessFnPropagatesOutputWindowAndTimestamp() throws Exception {
- // Tests that ProcessFn correctly propagates the window and timestamp of the element
- // inside the KeyedWorkItem.
- // The underlying DoFn is actually monolithic, so this doesn't test splitting.
- DoFn<Integer, String> fn = new ToStringFn();
-
- Instant base = Instant.now();
-
- IntervalWindow w =
- new IntervalWindow(
- base.minus(Duration.standardMinutes(1)), base.plus(Duration.standardMinutes(1)));
-
- ProcessFnTester<Integer, String, SomeRestriction, SomeRestrictionTracker> tester =
- new ProcessFnTester<>(
- base,
- fn,
- BigEndianIntegerCoder.of(),
- SerializableCoder.of(SomeRestriction.class),
- MAX_OUTPUTS_PER_BUNDLE,
- MAX_BUNDLE_DURATION);
- tester.startElement(
- WindowedValue.of(
- ElementAndRestriction.of(42, new SomeRestriction()),
- base,
- Collections.singletonList(w),
- PaneInfo.ON_TIME_AND_ONLY_FIRING));
-
- assertEquals(
- Arrays.asList(
- TimestampedValue.of("42a", base),
- TimestampedValue.of("42b", base),
- TimestampedValue.of("42c", base)),
- tester.peekOutputElementsInWindow(w));
- }
-
- private static class WatermarkUpdateFn extends DoFn<Instant, String> {
- @ProcessElement
- public void process(ProcessContext c, OffsetRangeTracker tracker) {
- for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) {
- c.updateWatermark(c.element().plus(Duration.standardSeconds(i)));
- c.output(String.valueOf(i));
- }
- }
-
- @GetInitialRestriction
- public OffsetRange getInitialRestriction(Instant elem) {
- throw new IllegalStateException("Expected to be supplied explicitly in this test");
- }
-
- @NewTracker
- public OffsetRangeTracker newTracker(OffsetRange range) {
- return new OffsetRangeTracker(range);
- }
- }
-
- @Test
- public void testUpdatesWatermark() throws Exception {
- DoFn<Instant, String> fn = new WatermarkUpdateFn();
- Instant base = Instant.now();
-
- ProcessFnTester<Instant, String, OffsetRange, OffsetRangeTracker> tester =
- new ProcessFnTester<>(
- base,
- fn,
- InstantCoder.of(),
- SerializableCoder.of(OffsetRange.class),
- 3,
- MAX_BUNDLE_DURATION);
-
- tester.startElement(base, new OffsetRange(0, 8));
- assertThat(tester.takeOutputElements(), hasItems("0", "1", "2"));
- assertEquals(base.plus(Duration.standardSeconds(2)), tester.getWatermarkHold());
-
- assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
- assertThat(tester.takeOutputElements(), hasItems("3", "4", "5"));
- assertEquals(base.plus(Duration.standardSeconds(5)), tester.getWatermarkHold());
-
- assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
- assertThat(tester.takeOutputElements(), hasItems("6", "7"));
- assertEquals(null, tester.getWatermarkHold());
- }
-
- /**
- * A splittable {@link DoFn} that generates the sequence [init, init + total).
- */
- private static class CounterFn extends DoFn<Integer, String> {
- @ProcessElement
- public void process(ProcessContext c, OffsetRangeTracker tracker) {
- for (long i = tracker.currentRestriction().getFrom();
- tracker.tryClaim(i); ++i) {
- c.output(String.valueOf(c.element() + i));
- }
- }
-
- @GetInitialRestriction
- public OffsetRange getInitialRestriction(Integer elem) {
- throw new UnsupportedOperationException("Expected to be supplied explicitly in this test");
- }
- }
-
- @Test
- public void testCheckpointsAfterNumOutputs() throws Exception {
- int max = 100;
- DoFn<Integer, String> fn = new CounterFn();
- Instant base = Instant.now();
- int baseIndex = 42;
-
- ProcessFnTester<Integer, String, OffsetRange, OffsetRangeTracker> tester =
- new ProcessFnTester<>(
- base, fn, BigEndianIntegerCoder.of(), SerializableCoder.of(OffsetRange.class),
- max, MAX_BUNDLE_DURATION);
-
- List<String> elements;
-
- // Create an fn that attempts to 2x output more than checkpointing allows.
- tester.startElement(baseIndex, new OffsetRange(0, 2 * max + max / 2));
- elements = tester.takeOutputElements();
- assertEquals(max, elements.size());
- // Should output the range [0, max)
- assertThat(elements, hasItem(String.valueOf(baseIndex)));
- assertThat(elements, hasItem(String.valueOf(baseIndex + max - 1)));
-
- assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
- elements = tester.takeOutputElements();
- assertEquals(max, elements.size());
- // Should output the range [max, 2*max)
- assertThat(elements, hasItem(String.valueOf(baseIndex + max)));
- assertThat(elements, hasItem(String.valueOf(baseIndex + 2 * max - 1)));
-
- assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
- elements = tester.takeOutputElements();
- assertEquals(max / 2, elements.size());
- // Should output the range [2*max, 2*max + max/2)
- assertThat(elements, hasItem(String.valueOf(baseIndex + 2 * max)));
- assertThat(elements, hasItem(String.valueOf(baseIndex + 2 * max + max / 2 - 1)));
- assertThat(elements, not(hasItem((String.valueOf(baseIndex + 2 * max + max / 2)))));
- }
-
- @Test
- public void testCheckpointsAfterDuration() throws Exception {
- // Don't bound number of outputs.
- int max = Integer.MAX_VALUE;
- // But bound bundle duration - the bundle should terminate.
- Duration maxBundleDuration = Duration.standardSeconds(1);
- // Create an fn that attempts to 2x output more than checkpointing allows.
- DoFn<Integer, String> fn = new CounterFn();
- Instant base = Instant.now();
- int baseIndex = 42;
-
- ProcessFnTester<Integer, String, OffsetRange, OffsetRangeTracker> tester =
- new ProcessFnTester<>(
- base, fn, BigEndianIntegerCoder.of(), SerializableCoder.of(OffsetRange.class),
- max, maxBundleDuration);
-
- List<String> elements;
-
- tester.startElement(baseIndex, new OffsetRange(0, Long.MAX_VALUE));
- // Bundle should terminate, and should do at least some processing.
- elements = tester.takeOutputElements();
- assertFalse(elements.isEmpty());
- // Bundle should have run for at least the requested duration.
- assertThat(
- Instant.now().getMillis() - base.getMillis(),
- greaterThanOrEqualTo(maxBundleDuration.getMillis()));
- }
-
- private static class LifecycleVerifyingFn extends DoFn<Integer, String> {
- private enum State {
- BEFORE_SETUP,
- OUTSIDE_BUNDLE,
- INSIDE_BUNDLE,
- TORN_DOWN
- }
-
- private State state = State.BEFORE_SETUP;
-
- @ProcessElement
- public void process(ProcessContext c, SomeRestrictionTracker tracker) {
- assertEquals(State.INSIDE_BUNDLE, state);
- }
-
- @GetInitialRestriction
- public SomeRestriction getInitialRestriction(Integer element) {
- return new SomeRestriction();
- }
-
- @Setup
- public void setup() {
- assertEquals(State.BEFORE_SETUP, state);
- state = State.OUTSIDE_BUNDLE;
- }
-
- @Teardown
- public void tearDown() {
- assertEquals(State.OUTSIDE_BUNDLE, state);
- state = State.TORN_DOWN;
- }
-
- @StartBundle
- public void startBundle() {
- assertEquals(State.OUTSIDE_BUNDLE, state);
- state = State.INSIDE_BUNDLE;
- }
-
- @FinishBundle
- public void finishBundle() {
- assertEquals(State.INSIDE_BUNDLE, state);
- state = State.OUTSIDE_BUNDLE;
- }
- }
-
- @Test
- public void testInvokesLifecycleMethods() throws Exception {
- DoFn<Integer, String> fn = new LifecycleVerifyingFn();
- try (ProcessFnTester<Integer, String, SomeRestriction, SomeRestrictionTracker> tester =
- new ProcessFnTester<>(
- Instant.now(),
- fn,
- BigEndianIntegerCoder.of(),
- SerializableCoder.of(SomeRestriction.class),
- MAX_OUTPUTS_PER_BUNDLE,
- MAX_BUNDLE_DURATION)) {
- tester.startElement(42, new SomeRestriction());
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGBKIntoKeyedWorkItemsOverrideFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGBKIntoKeyedWorkItemsOverrideFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGBKIntoKeyedWorkItemsOverrideFactory.java
index 64eecc8..3fefe20 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGBKIntoKeyedWorkItemsOverrideFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGBKIntoKeyedWorkItemsOverrideFactory.java
@@ -18,17 +18,14 @@
package org.apache.beam.runners.direct;
import org.apache.beam.runners.core.KeyedWorkItem;
-import org.apache.beam.runners.core.SplittableParDo.GBKIntoKeyedWorkItems;
+import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems;
import org.apache.beam.runners.core.construction.PTransformReplacements;
import org.apache.beam.runners.core.construction.SingleInputOutputOverrideFactory;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
-/**
- * Provides an implementation of {@link SplittableParDo.GBKIntoKeyedWorkItems} for the Direct
- * Runner.
- */
+/** Provides an implementation of {@link GBKIntoKeyedWorkItems} for the Direct Runner. */
class DirectGBKIntoKeyedWorkItemsOverrideFactory<KeyT, InputT>
extends SingleInputOutputOverrideFactory<
PCollection<KV<KeyT, InputT>>, PCollection<KeyedWorkItem<KeyT, InputT>>,
@@ -39,7 +36,7 @@ class DirectGBKIntoKeyedWorkItemsOverrideFactory<KeyT, InputT>
getReplacementTransform(
AppliedPTransform<
PCollection<KV<KeyT, InputT>>, PCollection<KeyedWorkItem<KeyT, InputT>>,
- GBKIntoKeyedWorkItems<KeyT, InputT>>
+ GBKIntoKeyedWorkItems<KeyT, InputT>>
transform) {
return PTransformReplacement.of(
PTransformReplacements.getSingletonMainInput(transform),
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRunner.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRunner.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRunner.java
index 181896f..69dea8f 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRunner.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRunner.java
@@ -27,8 +27,9 @@ import java.util.EnumSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
-import org.apache.beam.runners.core.SplittableParDo.GBKIntoKeyedWorkItems;
+import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems;
import org.apache.beam.runners.core.construction.PTransformMatchers;
+import org.apache.beam.runners.core.construction.SplittableParDo;
import org.apache.beam.runners.direct.DirectRunner.DirectPipelineResult;
import org.apache.beam.runners.direct.TestStreamEvaluatorFactory.DirectTestStreamFactory;
import org.apache.beam.sdk.Pipeline;
@@ -244,7 +245,12 @@ public class DirectRunner extends PipelineRunner<DirectPipelineResult> {
PTransformMatchers.stateOrTimerParDoMulti(), new ParDoMultiOverrideFactory()))
.add(
PTransformOverride.of(
- PTransformMatchers.classEqualTo(GBKIntoKeyedWorkItems.class),
+ PTransformMatchers.classEqualTo(SplittableParDo.ProcessKeyedElements.class),
+ new SplittableParDoViaKeyedWorkItems.OverrideFactory()))
+ .add(
+ PTransformOverride.of(
+ PTransformMatchers.classEqualTo(
+ SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems.class),
new DirectGBKIntoKeyedWorkItemsOverrideFactory())) /* Returns a GBKO */
.add(
PTransformOverride.of(
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java
index f9b6eba..347f313 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java
@@ -23,7 +23,7 @@ import com.google.common.collect.ImmutableSet;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
-import org.apache.beam.runners.core.SplittableParDo;
+import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems;
import org.apache.beam.runners.direct.DirectGroupByKey.DirectGroupAlsoByWindow;
import org.apache.beam.runners.direct.DirectGroupByKey.DirectGroupByKeyOnly;
import org.apache.beam.sdk.Pipeline.PipelineVisitor;
@@ -48,7 +48,7 @@ class KeyedPValueTrackingVisitor implements PipelineVisitor {
private static final Set<Class<? extends PTransform>> PRODUCES_KEYED_OUTPUTS =
ImmutableSet.of(
- SplittableParDo.GBKIntoKeyedWorkItems.class,
+ SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems.class,
DirectGroupByKeyOnly.class,
DirectGroupAlsoByWindow.class);
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoMultiOverrideFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoMultiOverrideFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoMultiOverrideFactory.java
index 89903da..be433db 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoMultiOverrideFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoMultiOverrideFactory.java
@@ -23,9 +23,9 @@ import java.util.Map;
import org.apache.beam.runners.core.KeyedWorkItem;
import org.apache.beam.runners.core.KeyedWorkItemCoder;
import org.apache.beam.runners.core.KeyedWorkItems;
-import org.apache.beam.runners.core.SplittableParDo;
import org.apache.beam.runners.core.construction.PTransformReplacements;
import org.apache.beam.runners.core.construction.ReplacementOutputs;
+import org.apache.beam.runners.core.construction.SplittableParDo;
import org.apache.beam.sdk.coders.CannotProvideCoderException;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/direct-java/src/main/java/org/apache/beam/runners/direct/SplittableProcessElementsEvaluatorFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/SplittableProcessElementsEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/SplittableProcessElementsEvaluatorFactory.java
index f490b0b..dc85d87 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/SplittableProcessElementsEvaluatorFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/SplittableProcessElementsEvaluatorFactory.java
@@ -24,18 +24,18 @@ import java.util.List;
import java.util.concurrent.Executors;
import org.apache.beam.runners.core.DoFnRunners;
import org.apache.beam.runners.core.DoFnRunners.OutputManager;
-import org.apache.beam.runners.core.ElementAndRestriction;
import org.apache.beam.runners.core.KeyedWorkItem;
import org.apache.beam.runners.core.OutputAndTimeBoundedSplittableProcessElementInvoker;
import org.apache.beam.runners.core.OutputWindowedValue;
import org.apache.beam.runners.core.PushbackSideInputDoFnRunner;
import org.apache.beam.runners.core.ReadyCheckingSideInputReader;
-import org.apache.beam.runners.core.SplittableParDo;
-import org.apache.beam.runners.core.SplittableParDo.ProcessFn;
+import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems.ProcessElements;
+import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems.ProcessFn;
import org.apache.beam.runners.core.StateInternals;
import org.apache.beam.runners.core.StateInternalsFactory;
import org.apache.beam.runners.core.TimerInternals;
import org.apache.beam.runners.core.TimerInternalsFactory;
+import org.apache.beam.runners.core.construction.ElementAndRestriction;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.transforms.DoFn;
@@ -88,12 +88,11 @@ class SplittableProcessElementsEvaluatorFactory<
createEvaluator(
AppliedPTransform<
PCollection<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>>,
- PCollectionTuple,
- SplittableParDo.ProcessElements<InputT, OutputT, RestrictionT, TrackerT>>
+ PCollectionTuple, ProcessElements<InputT, OutputT, RestrictionT, TrackerT>>
application,
CommittedBundle<InputT> inputBundle)
throws Exception {
- final SplittableParDo.ProcessElements<InputT, OutputT, RestrictionT, TrackerT> transform =
+ final ProcessElements<InputT, OutputT, RestrictionT, TrackerT> transform =
application.getTransform();
ProcessFn<InputT, OutputT, RestrictionT, TrackerT> processFn =
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformEvaluatorRegistry.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformEvaluatorRegistry.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformEvaluatorRegistry.java
index d0e622d..718cca2 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformEvaluatorRegistry.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformEvaluatorRegistry.java
@@ -25,7 +25,7 @@ import java.util.ArrayList;
import java.util.Collection;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
-import org.apache.beam.runners.core.SplittableParDo;
+import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems;
import org.apache.beam.runners.direct.DirectGroupByKey.DirectGroupAlsoByWindow;
import org.apache.beam.runners.direct.DirectGroupByKey.DirectGroupByKeyOnly;
import org.apache.beam.runners.direct.ParDoMultiOverrideFactory.StatefulParDo;
@@ -66,7 +66,7 @@ class TransformEvaluatorRegistry implements TransformEvaluatorFactory {
new TestStreamEvaluatorFactory(ctxt))
// Runner-specific primitive used in expansion of SplittableParDo
.put(
- SplittableParDo.ProcessElements.class,
+ SplittableParDoViaKeyedWorkItems.ProcessElements.class,
new SplittableProcessElementsEvaluatorFactory<>(ctxt))
.build();
return new TransformEvaluatorRegistry(primitives);
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslator.java
----------------------------------------------------------------------
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslator.java
index 35d1bcd..53a1fa1 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslator.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslator.java
@@ -20,11 +20,12 @@ package org.apache.beam.runners.flink;
import com.google.common.collect.ImmutableList;
import java.util.List;
import java.util.Map;
-import org.apache.beam.runners.core.SplittableParDo;
+import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems;
import org.apache.beam.runners.core.construction.PTransformMatchers;
import org.apache.beam.runners.core.construction.PTransformReplacements;
import org.apache.beam.runners.core.construction.ReplacementOutputs;
import org.apache.beam.runners.core.construction.SingleInputOutputOverrideFactory;
+import org.apache.beam.runners.core.construction.SplittableParDo;
import org.apache.beam.runners.core.construction.UnconsumedReads;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.options.PipelineOptions;
@@ -80,6 +81,10 @@ class FlinkStreamingPipelineTranslator extends FlinkPipelineTranslator {
new SplittableParDoOverrideFactory()))
.add(
PTransformOverride.of(
+ PTransformMatchers.classEqualTo(SplittableParDo.ProcessKeyedElements.class),
+ new SplittableParDoViaKeyedWorkItems.OverrideFactory()))
+ .add(
+ PTransformOverride.of(
PTransformMatchers.classEqualTo(View.AsIterable.class),
new ReflectiveOneToOneOverrideFactory(
FlinkStreamingViewOverrides.StreamingViewAsIterable.class, flinkRunner)))
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java
----------------------------------------------------------------------
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java
index 9a93205..00e9934 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java
@@ -30,10 +30,10 @@ import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
-import org.apache.beam.runners.core.ElementAndRestriction;
import org.apache.beam.runners.core.KeyedWorkItem;
-import org.apache.beam.runners.core.SplittableParDo;
+import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems;
import org.apache.beam.runners.core.SystemReduceFn;
+import org.apache.beam.runners.core.construction.ElementAndRestriction;
import org.apache.beam.runners.flink.translation.functions.FlinkAssignWindows;
import org.apache.beam.runners.flink.translation.types.CoderTypeInformation;
import org.apache.beam.runners.flink.translation.wrappers.streaming.DoFnOperator;
@@ -119,9 +119,11 @@ class FlinkStreamingTransformTranslators {
TRANSLATORS.put(ParDo.MultiOutput.class, new ParDoStreamingTranslator());
TRANSLATORS.put(
- SplittableParDo.ProcessElements.class, new SplittableProcessElementsStreamingTranslator());
+ SplittableParDoViaKeyedWorkItems.ProcessElements.class,
+ new SplittableProcessElementsStreamingTranslator());
TRANSLATORS.put(
- SplittableParDo.GBKIntoKeyedWorkItems.class, new GBKIntoKeyedWorkItemsTranslator());
+ SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems.class,
+ new GBKIntoKeyedWorkItemsTranslator());
TRANSLATORS.put(Window.Assign.class, new WindowAssignTranslator());
@@ -329,7 +331,8 @@ class FlinkStreamingTransformTranslators {
}
/**
- * Helper for translating {@link ParDo.MultiOutput} and {@link SplittableParDo.ProcessElements}.
+ * Helper for translating {@link ParDo.MultiOutput} and {@link
+ * SplittableParDoViaKeyedWorkItems.ProcessElements}.
*/
static class ParDoTranslationHelper {
@@ -382,7 +385,7 @@ class FlinkStreamingTransformTranslators {
keyCoder = ((KvCoder) input.getCoder()).getKeyCoder();
inputDataStream = inputDataStream.keyBy(new KvToByteBufferKeySelector(keyCoder));
stateful = true;
- } else if (doFn instanceof SplittableParDo.ProcessFn) {
+ } else if (doFn instanceof SplittableParDoViaKeyedWorkItems.ProcessFn) {
// we know that it is keyed on String
keyCoder = StringUtf8Coder.of();
stateful = true;
@@ -578,19 +581,19 @@ class FlinkStreamingTransformTranslators {
private static class SplittableProcessElementsStreamingTranslator<
InputT, OutputT, RestrictionT, TrackerT extends RestrictionTracker<RestrictionT>>
extends FlinkStreamingPipelineTranslator.StreamTransformTranslator<
- SplittableParDo.ProcessElements<InputT, OutputT, RestrictionT, TrackerT>> {
+ SplittableParDoViaKeyedWorkItems.ProcessElements<InputT, OutputT, RestrictionT, TrackerT>> {
@Override
public void translateNode(
- SplittableParDo.ProcessElements<InputT, OutputT, RestrictionT, TrackerT> transform,
+ SplittableParDoViaKeyedWorkItems.ProcessElements<InputT, OutputT, RestrictionT, TrackerT>
+ transform,
FlinkStreamingTranslationContext context) {
ParDoTranslationHelper.translateParDo(
transform.getName(),
transform.newProcessFn(transform.getFn()),
context.getCurrentTransform().getFullName(),
- (PCollection<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>>)
- context.getInput(transform),
+ context.getInput(transform),
transform.getSideInputs(),
context.getOutputs(transform),
transform.getMainOutputTag(),
@@ -930,18 +933,18 @@ class FlinkStreamingTransformTranslators {
private static class GBKIntoKeyedWorkItemsTranslator<K, InputT>
extends FlinkStreamingPipelineTranslator.StreamTransformTranslator<
- SplittableParDo.GBKIntoKeyedWorkItems<K, InputT>> {
+ SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems<K, InputT>> {
@Override
boolean canTranslate(
- SplittableParDo.GBKIntoKeyedWorkItems<K, InputT> transform,
+ SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems<K, InputT> transform,
FlinkStreamingTranslationContext context) {
return true;
}
@Override
public void translateNode(
- SplittableParDo.GBKIntoKeyedWorkItems<K, InputT> transform,
+ SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems<K, InputT> transform,
FlinkStreamingTranslationContext context) {
PCollection<KV<K, InputT>> input = context.getInput(transform);
[3/4] beam git commit: Splits SplittableParDo into a
core-construction part and a KWI-related part
Posted by jk...@apache.org.
Splits SplittableParDo into a core-construction part and a KWI-related part
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/7a2e829c
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/7a2e829c
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/7a2e829c
Branch: refs/heads/master
Commit: 7a2e829c849d3828353994b8eddca4784a900aca
Parents: 2e0cf00
Author: Eugene Kirpichov <ki...@google.com>
Authored: Mon May 15 15:54:03 2017 -0700
Committer: Eugene Kirpichov <ki...@google.com>
Committed: Thu May 18 16:11:09 2017 -0700
----------------------------------------------------------------------
.../apache/beam/runners/apex/ApexRunner.java | 7 +-
.../translation/ApexPipelineTranslator.java | 11 +-
.../apex/translation/ParDoTranslator.java | 11 +-
.../operators/ApexParDoOperator.java | 10 +-
.../construction/ElementAndRestriction.java | 42 ++
.../ElementAndRestrictionCoder.java | 88 +++
.../core/construction/SplittableParDo.java | 298 +++++++++
.../ElementAndRestrictionCoderTest.java | 126 ++++
.../core/construction/SplittableParDoTest.java | 157 +++++
.../apache/beam/runners/core/DoFnRunners.java | 2 +-
.../runners/core/ElementAndRestriction.java | 42 --
.../core/ElementAndRestrictionCoder.java | 88 ---
.../beam/runners/core/ProcessFnRunner.java | 7 +-
.../beam/runners/core/SplittableParDo.java | 602 ------------------
.../core/SplittableParDoViaKeyedWorkItems.java | 418 +++++++++++++
.../core/ElementAndRestrictionCoderTest.java | 127 ----
.../core/SplittableParDoProcessFnTest.java | 515 ++++++++++++++++
.../beam/runners/core/SplittableParDoTest.java | 606 -------------------
...ectGBKIntoKeyedWorkItemsOverrideFactory.java | 9 +-
.../beam/runners/direct/DirectRunner.java | 10 +-
.../direct/KeyedPValueTrackingVisitor.java | 4 +-
.../direct/ParDoMultiOverrideFactory.java | 2 +-
...littableProcessElementsEvaluatorFactory.java | 11 +-
.../direct/TransformEvaluatorRegistry.java | 4 +-
.../flink/FlinkStreamingPipelineTranslator.java | 7 +-
.../FlinkStreamingTransformTranslators.java | 29 +-
.../streaming/SplittableDoFnOperator.java | 12 +-
27 files changed, 1721 insertions(+), 1524 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/apex/src/main/java/org/apache/beam/runners/apex/ApexRunner.java
----------------------------------------------------------------------
diff --git a/runners/apex/src/main/java/org/apache/beam/runners/apex/ApexRunner.java b/runners/apex/src/main/java/org/apache/beam/runners/apex/ApexRunner.java
index 2fd0b22..f91d8e5 100644
--- a/runners/apex/src/main/java/org/apache/beam/runners/apex/ApexRunner.java
+++ b/runners/apex/src/main/java/org/apache/beam/runners/apex/ApexRunner.java
@@ -39,12 +39,13 @@ import org.apache.apex.api.Launcher;
import org.apache.apex.api.Launcher.AppHandle;
import org.apache.apex.api.Launcher.LaunchMode;
import org.apache.beam.runners.apex.translation.ApexPipelineTranslator;
-import org.apache.beam.runners.core.SplittableParDo;
+import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems;
import org.apache.beam.runners.core.construction.PTransformMatchers;
import org.apache.beam.runners.core.construction.PTransformReplacements;
import org.apache.beam.runners.core.construction.PrimitiveCreate;
import org.apache.beam.runners.core.construction.ReplacementOutputs;
import org.apache.beam.runners.core.construction.SingleInputOutputOverrideFactory;
+import org.apache.beam.runners.core.construction.SplittableParDo;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineRunner;
import org.apache.beam.sdk.coders.Coder;
@@ -124,6 +125,10 @@ public class ApexRunner extends PipelineRunner<ApexRunnerResult> {
PTransformOverride.of(
PTransformMatchers.splittableParDoMulti(),
new SplittableParDoOverrideFactory<>()))
+ .add(
+ PTransformOverride.of(
+ PTransformMatchers.classEqualTo(SplittableParDo.ProcessKeyedElements.class),
+ new SplittableParDoViaKeyedWorkItems.OverrideFactory<>()))
.build();
}
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/ApexPipelineTranslator.java
----------------------------------------------------------------------
diff --git a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/ApexPipelineTranslator.java b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/ApexPipelineTranslator.java
index b3a6d1c..32e470f 100644
--- a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/ApexPipelineTranslator.java
+++ b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/ApexPipelineTranslator.java
@@ -25,7 +25,8 @@ import org.apache.beam.runners.apex.ApexPipelineOptions;
import org.apache.beam.runners.apex.ApexRunner.CreateApexPCollectionView;
import org.apache.beam.runners.apex.translation.operators.ApexProcessFnOperator;
import org.apache.beam.runners.apex.translation.operators.ApexReadUnboundedInputOperator;
-import org.apache.beam.runners.core.SplittableParDo;
+import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems;
+import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems;
import org.apache.beam.runners.core.construction.PrimitiveCreate;
import org.apache.beam.runners.core.construction.UnboundedReadFromBoundedSource.BoundedToUnboundedSourceAdapter;
import org.apache.beam.sdk.Pipeline;
@@ -63,9 +64,9 @@ public class ApexPipelineTranslator implements Pipeline.PipelineVisitor {
static {
// register TransformTranslators
registerTransformTranslator(ParDo.MultiOutput.class, new ParDoTranslator<>());
- registerTransformTranslator(SplittableParDo.ProcessElements.class,
+ registerTransformTranslator(SplittableParDoViaKeyedWorkItems.ProcessElements.class,
new ParDoTranslator.SplittableProcessElementsTranslator());
- registerTransformTranslator(SplittableParDo.GBKIntoKeyedWorkItems.class,
+ registerTransformTranslator(GBKIntoKeyedWorkItems.class,
new GBKIntoKeyedWorkItemsTranslator());
registerTransformTranslator(Read.Unbounded.class, new ReadUnboundedTranslator());
registerTransformTranslator(Read.Bounded.class, new ReadBoundedTranslator());
@@ -183,11 +184,11 @@ public class ApexPipelineTranslator implements Pipeline.PipelineVisitor {
}
private static class GBKIntoKeyedWorkItemsTranslator<K, InputT>
- implements TransformTranslator<SplittableParDo.GBKIntoKeyedWorkItems<K, InputT>> {
+ implements TransformTranslator<GBKIntoKeyedWorkItems<K, InputT>> {
@Override
public void translate(
- SplittableParDo.GBKIntoKeyedWorkItems<K, InputT> transform, TranslationContext context) {
+ GBKIntoKeyedWorkItems<K, InputT> transform, TranslationContext context) {
// https://issues.apache.org/jira/browse/BEAM-1850
ApexProcessFnOperator<KV<K, InputT>> operator = ApexProcessFnOperator.toKeyedWorkItems(
context.getPipelineOptions());
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/ParDoTranslator.java
----------------------------------------------------------------------
diff --git a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/ParDoTranslator.java b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/ParDoTranslator.java
index 9133cb6..e46687a 100644
--- a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/ParDoTranslator.java
+++ b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/ParDoTranslator.java
@@ -30,7 +30,7 @@ import java.util.Map;
import java.util.Map.Entry;
import org.apache.beam.runners.apex.ApexRunner;
import org.apache.beam.runners.apex.translation.operators.ApexParDoOperator;
-import org.apache.beam.runners.core.SplittableParDo;
+import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems.ProcessElements;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
@@ -132,14 +132,13 @@ class ParDoTranslator<InputT, OutputT>
}
}
- static class SplittableProcessElementsTranslator<InputT, OutputT,
- RestrictionT, TrackerT extends RestrictionTracker<RestrictionT>>
- implements TransformTranslator<SplittableParDo.ProcessElements<InputT, OutputT,
- RestrictionT, TrackerT>> {
+ static class SplittableProcessElementsTranslator<
+ InputT, OutputT, RestrictionT, TrackerT extends RestrictionTracker<RestrictionT>>
+ implements TransformTranslator<ProcessElements<InputT, OutputT, RestrictionT, TrackerT>> {
@Override
public void translate(
- SplittableParDo.ProcessElements<InputT, OutputT, RestrictionT, TrackerT> transform,
+ ProcessElements<InputT, OutputT, RestrictionT, TrackerT> transform,
TranslationContext context) {
Map<TupleTag<?>, PValue> outputs = context.getOutputs();
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/operators/ApexParDoOperator.java
----------------------------------------------------------------------
diff --git a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/operators/ApexParDoOperator.java b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/operators/ApexParDoOperator.java
index 7fee0d5..809ca2a 100644
--- a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/operators/ApexParDoOperator.java
+++ b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/operators/ApexParDoOperator.java
@@ -55,7 +55,7 @@ import org.apache.beam.runners.core.PushbackSideInputDoFnRunner;
import org.apache.beam.runners.core.SideInputHandler;
import org.apache.beam.runners.core.SideInputReader;
import org.apache.beam.runners.core.SimplePushbackSideInputDoFnRunner;
-import org.apache.beam.runners.core.SplittableParDo;
+import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems.ProcessFn;
import org.apache.beam.runners.core.StateInternals;
import org.apache.beam.runners.core.StateInternalsFactory;
import org.apache.beam.runners.core.StateNamespace;
@@ -160,7 +160,7 @@ public class ApexParDoOperator<InputT, OutputT> extends BaseOperator implements
TimerInternals.TimerDataCoder.of(windowingStrategy.getWindowFn().windowCoder());
this.currentKeyTimerInternals = new ApexTimerInternals<>(timerCoder);
- if (doFn instanceof SplittableParDo.ProcessFn) {
+ if (doFn instanceof ProcessFn) {
// we know that it is keyed on String
Coder<?> keyCoder = StringUtf8Coder.of();
this.currentKeyStateInternals = new StateInternalsProxy<>(
@@ -445,15 +445,15 @@ public class ApexParDoOperator<InputT, OutputT> extends BaseOperator implements
pushbackDoFnRunner =
SimplePushbackSideInputDoFnRunner.create(doFnRunner, sideInputs, sideInputHandler);
- if (doFn instanceof SplittableParDo.ProcessFn) {
+ if (doFn instanceof ProcessFn) {
@SuppressWarnings("unchecked")
StateInternalsFactory<String> stateInternalsFactory =
(StateInternalsFactory<String>) this.currentKeyStateInternals.getFactory();
@SuppressWarnings({ "rawtypes", "unchecked" })
- SplittableParDo.ProcessFn<InputT, OutputT, Object, RestrictionTracker<Object>>
- splittableDoFn = (SplittableParDo.ProcessFn) doFn;
+ ProcessFn<InputT, OutputT, Object, RestrictionTracker<Object>>
+ splittableDoFn = (ProcessFn) doFn;
splittableDoFn.setStateInternalsFactory(stateInternalsFactory);
TimerInternalsFactory<String> timerInternalsFactory = new TimerInternalsFactory<String>() {
@Override
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ElementAndRestriction.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ElementAndRestriction.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ElementAndRestriction.java
new file mode 100644
index 0000000..53a86b1
--- /dev/null
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ElementAndRestriction.java
@@ -0,0 +1,42 @@
+/*
+ * 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.core.construction;
+
+import com.google.auto.value.AutoValue;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.transforms.DoFn;
+
+/**
+ * A tuple of an element and a restriction applied to processing it with a
+ * <a href="https://s.apache.org/splittable-do-fn">splittable</a> {@link DoFn}.
+ */
+@Experimental(Experimental.Kind.SPLITTABLE_DO_FN)
+@AutoValue
+public abstract class ElementAndRestriction<ElementT, RestrictionT> {
+ /** The element to process. */
+ public abstract ElementT element();
+
+ /** The restriction applied to processing the element. */
+ public abstract RestrictionT restriction();
+
+ /** Constructs the {@link ElementAndRestriction}. */
+ public static <InputT, RestrictionT> ElementAndRestriction<InputT, RestrictionT> of(
+ InputT element, RestrictionT restriction) {
+ return new AutoValue_ElementAndRestriction<>(element, restriction);
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ElementAndRestrictionCoder.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ElementAndRestrictionCoder.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ElementAndRestrictionCoder.java
new file mode 100644
index 0000000..5ff0aae
--- /dev/null
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ElementAndRestrictionCoder.java
@@ -0,0 +1,88 @@
+/*
+ * 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.core.construction;
+
+import com.google.common.collect.ImmutableList;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.List;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderException;
+import org.apache.beam.sdk.coders.StructuredCoder;
+
+/** A {@link Coder} for {@link ElementAndRestriction}. */
+@Experimental(Experimental.Kind.SPLITTABLE_DO_FN)
+public class ElementAndRestrictionCoder<ElementT, RestrictionT>
+ extends StructuredCoder<ElementAndRestriction<ElementT, RestrictionT>> {
+ private final Coder<ElementT> elementCoder;
+ private final Coder<RestrictionT> restrictionCoder;
+
+ /**
+ * Creates an {@link ElementAndRestrictionCoder} from an element coder and a restriction coder.
+ */
+ public static <ElementT, RestrictionT> ElementAndRestrictionCoder<ElementT, RestrictionT> of(
+ Coder<ElementT> elementCoder, Coder<RestrictionT> restrictionCoder) {
+ return new ElementAndRestrictionCoder<>(elementCoder, restrictionCoder);
+ }
+
+ private ElementAndRestrictionCoder(
+ Coder<ElementT> elementCoder, Coder<RestrictionT> restrictionCoder) {
+ this.elementCoder = elementCoder;
+ this.restrictionCoder = restrictionCoder;
+ }
+
+ @Override
+ public void encode(
+ ElementAndRestriction<ElementT, RestrictionT> value, OutputStream outStream)
+ throws IOException {
+ if (value == null) {
+ throw new CoderException("cannot encode a null ElementAndRestriction");
+ }
+ elementCoder.encode(value.element(), outStream);
+ restrictionCoder.encode(value.restriction(), outStream);
+ }
+
+ @Override
+ public ElementAndRestriction<ElementT, RestrictionT> decode(InputStream inStream)
+ throws IOException {
+ ElementT key = elementCoder.decode(inStream);
+ RestrictionT value = restrictionCoder.decode(inStream);
+ return ElementAndRestriction.of(key, value);
+ }
+
+ @Override
+ public List<? extends Coder<?>> getCoderArguments() {
+ return ImmutableList.of(elementCoder, restrictionCoder);
+ }
+
+ @Override
+ public void verifyDeterministic() throws NonDeterministicException {
+ elementCoder.verifyDeterministic();
+ restrictionCoder.verifyDeterministic();
+ }
+
+ public Coder<ElementT> getElementCoder() {
+ return elementCoder;
+ }
+
+ public Coder<RestrictionT> getRestrictionCoder() {
+ return restrictionCoder;
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SplittableParDo.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SplittableParDo.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SplittableParDo.java
new file mode 100644
index 0000000..23ba66a
--- /dev/null
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SplittableParDo.java
@@ -0,0 +1,298 @@
+/*
+ * 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.core.construction;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import java.util.List;
+import java.util.UUID;
+import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.WithKeys;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
+import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.reflect.DoFnSignature;
+import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.apache.beam.sdk.values.WindowingStrategy;
+
+/**
+ * A utility transform that executes a <a
+ * href="https://s.apache.org/splittable-do-fn">splittable</a> {@link DoFn} by expanding it into a
+ * network of simpler transforms:
+ *
+ * <ol>
+ * <li>Pair each element with an initial restriction
+ * <li>Split each restriction into sub-restrictions
+ * <li>Explode windows, since splitting within each window has to happen independently
+ * <li>Assign a unique key to each element/restriction pair
+ * <li>Process the keyed element/restriction pairs in a runner-specific way with the splittable
+ * {@link DoFn}'s {@link DoFn.ProcessElement} method.
+ * </ol>
+ *
+ * <p>This transform is intended as a helper for internal use by runners when implementing {@code
+ * ParDo.of(splittable DoFn)}, but not for direct use by pipeline writers.
+ */
+@Experimental(Experimental.Kind.SPLITTABLE_DO_FN)
+public class SplittableParDo<InputT, OutputT, RestrictionT>
+ extends PTransform<PCollection<InputT>, PCollectionTuple> {
+ private final ParDo.MultiOutput<InputT, OutputT> parDo;
+
+ /**
+ * Creates the transform for the given original multi-output {@link ParDo}.
+ *
+ * @param parDo The splittable {@link ParDo} transform.
+ */
+ public SplittableParDo(ParDo.MultiOutput<InputT, OutputT> parDo) {
+ checkNotNull(parDo, "parDo must not be null");
+ this.parDo = parDo;
+ checkArgument(
+ DoFnSignatures.getSignature(parDo.getFn().getClass()).processElement().isSplittable(),
+ "fn must be a splittable DoFn");
+ }
+
+ @Override
+ public PCollectionTuple expand(PCollection<InputT> input) {
+ DoFn<InputT, OutputT> fn = parDo.getFn();
+ Coder<RestrictionT> restrictionCoder =
+ DoFnInvokers.invokerFor(fn)
+ .invokeGetRestrictionCoder(input.getPipeline().getCoderRegistry());
+ Coder<ElementAndRestriction<InputT, RestrictionT>> splitCoder =
+ ElementAndRestrictionCoder.of(input.getCoder(), restrictionCoder);
+
+ PCollection<KV<String, ElementAndRestriction<InputT, RestrictionT>>> keyedRestrictions =
+ input
+ .apply(
+ "Pair with initial restriction",
+ ParDo.of(new PairWithRestrictionFn<InputT, OutputT, RestrictionT>(fn)))
+ .setCoder(splitCoder)
+ .apply("Split restriction", ParDo.of(new SplitRestrictionFn<InputT, RestrictionT>(fn)))
+ .setCoder(splitCoder)
+ // ProcessFn requires all input elements to be in a single window and have a single
+ // element per work item. This must precede the unique keying so each key has a single
+ // associated element.
+ .apply(
+ "Explode windows",
+ ParDo.of(new ExplodeWindowsFn<ElementAndRestriction<InputT, RestrictionT>>()))
+ .apply(
+ "Assign unique key",
+ WithKeys.of(new RandomUniqueKeyFn<ElementAndRestriction<InputT, RestrictionT>>()));
+
+ return keyedRestrictions.apply(
+ "ProcessKeyedElements",
+ new ProcessKeyedElements<>(
+ fn,
+ input.getCoder(),
+ restrictionCoder,
+ (WindowingStrategy<InputT, ?>) input.getWindowingStrategy(),
+ parDo.getSideInputs(),
+ parDo.getMainOutputTag(),
+ parDo.getAdditionalOutputTags()));
+ }
+
+ /**
+ * A {@link DoFn} that forces each of its outputs to be in a single window, by indicating to the
+ * runner that it observes the window of its input element, so the runner is forced to apply it to
+ * each input in a single window and thus its output is also in a single window.
+ */
+ private static class ExplodeWindowsFn<InputT> extends DoFn<InputT, InputT> {
+ @ProcessElement
+ public void process(ProcessContext c, BoundedWindow window) {
+ c.output(c.element());
+ }
+ }
+
+ /**
+ * Runner-specific primitive {@link PTransform} that invokes the {@link DoFn.ProcessElement}
+ * method for a splittable {@link DoFn} on each {@link ElementAndRestriction} of the input
+ * {@link PCollection} of {@link KV KVs} keyed with arbitrary but globally unique keys.
+ */
+ public static class ProcessKeyedElements<InputT, OutputT, RestrictionT>
+ extends PTransform<
+ PCollection<KV<String, ElementAndRestriction<InputT, RestrictionT>>>, PCollectionTuple> {
+ private final DoFn<InputT, OutputT> fn;
+ private final Coder<InputT> elementCoder;
+ private final Coder<RestrictionT> restrictionCoder;
+ private final WindowingStrategy<InputT, ?> windowingStrategy;
+ private final List<PCollectionView<?>> sideInputs;
+ private final TupleTag<OutputT> mainOutputTag;
+ private final TupleTagList additionalOutputTags;
+
+ /**
+ * @param fn the splittable {@link DoFn}.
+ * @param windowingStrategy the {@link WindowingStrategy} of the input collection.
+ * @param sideInputs list of side inputs that should be available to the {@link DoFn}.
+ * @param mainOutputTag {@link TupleTag Tag} of the {@link DoFn DoFn's} main output.
+ * @param additionalOutputTags {@link TupleTagList Tags} of the {@link DoFn DoFn's} additional
+ * outputs.
+ */
+ public ProcessKeyedElements(
+ DoFn<InputT, OutputT> fn,
+ Coder<InputT> elementCoder,
+ Coder<RestrictionT> restrictionCoder,
+ WindowingStrategy<InputT, ?> windowingStrategy,
+ List<PCollectionView<?>> sideInputs,
+ TupleTag<OutputT> mainOutputTag,
+ TupleTagList additionalOutputTags) {
+ this.fn = fn;
+ this.elementCoder = elementCoder;
+ this.restrictionCoder = restrictionCoder;
+ this.windowingStrategy = windowingStrategy;
+ this.sideInputs = sideInputs;
+ this.mainOutputTag = mainOutputTag;
+ this.additionalOutputTags = additionalOutputTags;
+ }
+
+ public DoFn<InputT, OutputT> getFn() {
+ return fn;
+ }
+
+ public Coder<InputT> getElementCoder() {
+ return elementCoder;
+ }
+
+ public Coder<RestrictionT> getRestrictionCoder() {
+ return restrictionCoder;
+ }
+
+ public WindowingStrategy<InputT, ?> getInputWindowingStrategy() {
+ return windowingStrategy;
+ }
+
+ public List<PCollectionView<?>> getSideInputs() {
+ return sideInputs;
+ }
+
+ public TupleTag<OutputT> getMainOutputTag() {
+ return mainOutputTag;
+ }
+
+ public TupleTagList getAdditionalOutputTags() {
+ return additionalOutputTags;
+ }
+
+ @Override
+ public PCollectionTuple expand(
+ PCollection<KV<String, ElementAndRestriction<InputT, RestrictionT>>>
+ input) {
+ return createPrimitiveOutputFor(
+ input, fn, mainOutputTag, additionalOutputTags, windowingStrategy);
+ }
+
+ public static <OutputT> PCollectionTuple createPrimitiveOutputFor(
+ PCollection<?> input,
+ DoFn<?, OutputT> fn,
+ TupleTag<OutputT> mainOutputTag,
+ TupleTagList additionalOutputTags,
+ WindowingStrategy<?, ?> windowingStrategy) {
+ DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());
+ PCollectionTuple outputs =
+ PCollectionTuple.ofPrimitiveOutputsInternal(
+ input.getPipeline(),
+ TupleTagList.of(mainOutputTag).and(additionalOutputTags.getAll()),
+ windowingStrategy,
+ input.isBounded().and(signature.isBoundedPerElement()));
+
+ // Set output type descriptor similarly to how ParDo.MultiOutput does it.
+ outputs.get(mainOutputTag).setTypeDescriptor(fn.getOutputTypeDescriptor());
+
+ return outputs;
+ }
+ }
+
+ /**
+ * Assigns a random unique key to each element of the input collection, so that the output
+ * collection is effectively the same elements as input, but the per-key state and timers are now
+ * effectively per-element.
+ */
+ private static class RandomUniqueKeyFn<T> implements SerializableFunction<T, String> {
+ @Override
+ public String apply(T input) {
+ return UUID.randomUUID().toString();
+ }
+ }
+
+ /**
+ * Pairs each input element with its initial restriction using the given splittable {@link DoFn}.
+ */
+ private static class PairWithRestrictionFn<InputT, OutputT, RestrictionT>
+ extends DoFn<InputT, ElementAndRestriction<InputT, RestrictionT>> {
+ private DoFn<InputT, OutputT> fn;
+ private transient DoFnInvoker<InputT, OutputT> invoker;
+
+ PairWithRestrictionFn(DoFn<InputT, OutputT> fn) {
+ this.fn = fn;
+ }
+
+ @Setup
+ public void setup() {
+ invoker = DoFnInvokers.invokerFor(fn);
+ }
+
+ @ProcessElement
+ public void processElement(ProcessContext context) {
+ context.output(
+ ElementAndRestriction.of(
+ context.element(),
+ invoker.<RestrictionT>invokeGetInitialRestriction(context.element())));
+ }
+ }
+
+ /** Splits the restriction using the given {@link SplitRestriction} method. */
+ private static class SplitRestrictionFn<InputT, RestrictionT>
+ extends DoFn<
+ ElementAndRestriction<InputT, RestrictionT>,
+ ElementAndRestriction<InputT, RestrictionT>> {
+ private final DoFn<InputT, ?> splittableFn;
+ private transient DoFnInvoker<InputT, ?> invoker;
+
+ SplitRestrictionFn(DoFn<InputT, ?> splittableFn) {
+ this.splittableFn = splittableFn;
+ }
+
+ @Setup
+ public void setup() {
+ invoker = DoFnInvokers.invokerFor(splittableFn);
+ }
+
+ @ProcessElement
+ public void processElement(final ProcessContext c) {
+ final InputT element = c.element().element();
+ invoker.invokeSplitRestriction(
+ element,
+ c.element().restriction(),
+ new OutputReceiver<RestrictionT>() {
+ @Override
+ public void output(RestrictionT part) {
+ c.output(ElementAndRestriction.of(element, part));
+ }
+ });
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ElementAndRestrictionCoderTest.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ElementAndRestrictionCoderTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ElementAndRestrictionCoderTest.java
new file mode 100644
index 0000000..051cbaa
--- /dev/null
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ElementAndRestrictionCoderTest.java
@@ -0,0 +1,126 @@
+/*
+ * 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.core.construction;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import org.apache.beam.sdk.coders.BigEndianLongCoder;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderException;
+import org.apache.beam.sdk.coders.ListCoder;
+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.testing.CoderProperties;
+import org.apache.beam.sdk.util.CoderUtils;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+
+/**
+ * Tests for {@link ElementAndRestrictionCoder}.
+ */
+@RunWith(Parameterized.class)
+public class ElementAndRestrictionCoderTest<K, V> {
+ private static class CoderAndData<T> {
+ Coder<T> coder;
+ List<T> data;
+ }
+
+ private static class AnyCoderAndData {
+ private CoderAndData<?> coderAndData;
+ }
+
+ private static <T> AnyCoderAndData coderAndData(Coder<T> coder, List<T> data) {
+ CoderAndData<T> coderAndData = new CoderAndData<>();
+ coderAndData.coder = coder;
+ coderAndData.data = data;
+ AnyCoderAndData res = new AnyCoderAndData();
+ res.coderAndData = coderAndData;
+ return res;
+ }
+
+ private static final List<AnyCoderAndData> TEST_DATA =
+ Arrays.asList(
+ coderAndData(
+ VarIntCoder.of(), Arrays.asList(-1, 0, 1, 13, Integer.MAX_VALUE, Integer.MIN_VALUE)),
+ coderAndData(
+ BigEndianLongCoder.of(),
+ Arrays.asList(-1L, 0L, 1L, 13L, Long.MAX_VALUE, Long.MIN_VALUE)),
+ coderAndData(StringUtf8Coder.of(), Arrays.asList("", "hello", "goodbye", "1")),
+ coderAndData(
+ ElementAndRestrictionCoder.of(StringUtf8Coder.of(), VarIntCoder.of()),
+ Arrays.asList(
+ ElementAndRestriction.of("", -1),
+ ElementAndRestriction.of("hello", 0),
+ ElementAndRestriction.of("goodbye", Integer.MAX_VALUE))),
+ coderAndData(
+ ListCoder.of(VarLongCoder.of()),
+ Arrays.asList(Arrays.asList(1L, 2L, 3L), Collections.<Long>emptyList())));
+
+ @Parameterized.Parameters(name = "{index}: keyCoder={0} key={1} valueCoder={2} value={3}")
+ public static Collection<Object[]> data() {
+ List<Object[]> parameters = new ArrayList<>();
+ for (AnyCoderAndData keyCoderAndData : TEST_DATA) {
+ Coder keyCoder = keyCoderAndData.coderAndData.coder;
+ for (Object key : keyCoderAndData.coderAndData.data) {
+ for (AnyCoderAndData valueCoderAndData : TEST_DATA) {
+ Coder valueCoder = valueCoderAndData.coderAndData.coder;
+ for (Object value : valueCoderAndData.coderAndData.data) {
+ parameters.add(new Object[] {keyCoder, key, valueCoder, value});
+ }
+ }
+ }
+ }
+ return parameters;
+ }
+
+ @Parameter(0)
+ public Coder<K> keyCoder;
+ @Parameter(1)
+ public K key;
+ @Parameter(2)
+ public Coder<V> valueCoder;
+ @Parameter(3)
+ public V value;
+
+ @Test
+ @SuppressWarnings("rawtypes")
+ public void testDecodeEncodeEqual() throws Exception {
+ CoderProperties.coderDecodeEncodeEqual(
+ ElementAndRestrictionCoder.of(keyCoder, valueCoder),
+ ElementAndRestriction.of(key, value));
+ }
+
+ @Rule public ExpectedException thrown = ExpectedException.none();
+
+ @Test
+ public void encodeNullThrowsCoderException() throws Exception {
+ thrown.expect(CoderException.class);
+ thrown.expectMessage("cannot encode a null ElementAndRestriction");
+
+ CoderUtils.encodeToBase64(
+ ElementAndRestrictionCoder.of(StringUtf8Coder.of(), VarIntCoder.of()), null);
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SplittableParDoTest.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SplittableParDoTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SplittableParDoTest.java
new file mode 100644
index 0000000..6e4d6c4
--- /dev/null
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SplittableParDoTest.java
@@ -0,0 +1,157 @@
+/*
+ * 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.core.construction;
+
+import static org.junit.Assert.assertEquals;
+
+import java.io.Serializable;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.DoFn.BoundedPerElement;
+import org.apache.beam.sdk.transforms.DoFn.UnboundedPerElement;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.splittabledofn.HasDefaultTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TupleTagList;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link SplittableParDo}. */
+@RunWith(JUnit4.class)
+public class SplittableParDoTest {
+ // ----------------- Tests for whether the transform sets boundedness correctly --------------
+ private static class SomeRestriction
+ implements Serializable, HasDefaultTracker<SomeRestriction, SomeRestrictionTracker> {
+ @Override
+ public SomeRestrictionTracker newTracker() {
+ return new SomeRestrictionTracker(this);
+ }
+ }
+
+ private static class SomeRestrictionTracker implements RestrictionTracker<SomeRestriction> {
+ private final SomeRestriction someRestriction;
+
+ public SomeRestrictionTracker(SomeRestriction someRestriction) {
+ this.someRestriction = someRestriction;
+ }
+
+ @Override
+ public SomeRestriction currentRestriction() {
+ return someRestriction;
+ }
+
+ @Override
+ public SomeRestriction checkpoint() {
+ return someRestriction;
+ }
+
+ @Override
+ public void checkDone() {}
+ }
+
+ @BoundedPerElement
+ private static class BoundedFakeFn extends DoFn<Integer, String> {
+ @ProcessElement
+ public void processElement(ProcessContext context, SomeRestrictionTracker tracker) {}
+
+ @GetInitialRestriction
+ public SomeRestriction getInitialRestriction(Integer element) {
+ return null;
+ }
+ }
+
+ @UnboundedPerElement
+ private static class UnboundedFakeFn extends DoFn<Integer, String> {
+ @ProcessElement
+ public void processElement(ProcessContext context, SomeRestrictionTracker tracker) {}
+
+ @GetInitialRestriction
+ public SomeRestriction getInitialRestriction(Integer element) {
+ return null;
+ }
+ }
+
+ private static PCollection<Integer> makeUnboundedCollection(Pipeline pipeline) {
+ return pipeline
+ .apply("unbounded", Create.of(1, 2, 3))
+ .setIsBoundedInternal(PCollection.IsBounded.UNBOUNDED);
+ }
+
+ private static PCollection<Integer> makeBoundedCollection(Pipeline pipeline) {
+ return pipeline
+ .apply("bounded", Create.of(1, 2, 3))
+ .setIsBoundedInternal(PCollection.IsBounded.BOUNDED);
+ }
+
+ private static final TupleTag<String> MAIN_OUTPUT_TAG = new TupleTag<String>() {};
+
+ private ParDo.MultiOutput<Integer, String> makeParDo(DoFn<Integer, String> fn) {
+ return ParDo.of(fn).withOutputTags(MAIN_OUTPUT_TAG, TupleTagList.empty());
+ }
+
+ @Rule
+ public TestPipeline pipeline = TestPipeline.create();
+
+ @Test
+ public void testBoundednessForBoundedFn() {
+ pipeline.enableAbandonedNodeEnforcement(false);
+
+ DoFn<Integer, String> boundedFn = new BoundedFakeFn();
+ assertEquals(
+ "Applying a bounded SDF to a bounded collection produces a bounded collection",
+ PCollection.IsBounded.BOUNDED,
+ makeBoundedCollection(pipeline)
+ .apply("bounded to bounded", new SplittableParDo<>(makeParDo(boundedFn)))
+ .get(MAIN_OUTPUT_TAG)
+ .isBounded());
+ assertEquals(
+ "Applying a bounded SDF to an unbounded collection produces an unbounded collection",
+ PCollection.IsBounded.UNBOUNDED,
+ makeUnboundedCollection(pipeline)
+ .apply("bounded to unbounded", new SplittableParDo<>(makeParDo(boundedFn)))
+ .get(MAIN_OUTPUT_TAG)
+ .isBounded());
+ }
+
+ @Test
+ public void testBoundednessForUnboundedFn() {
+ pipeline.enableAbandonedNodeEnforcement(false);
+
+ DoFn<Integer, String> unboundedFn = new UnboundedFakeFn();
+ assertEquals(
+ "Applying an unbounded SDF to a bounded collection produces a bounded collection",
+ PCollection.IsBounded.UNBOUNDED,
+ makeBoundedCollection(pipeline)
+ .apply("unbounded to bounded", new SplittableParDo<>(makeParDo(unboundedFn)))
+ .get(MAIN_OUTPUT_TAG)
+ .isBounded());
+ assertEquals(
+ "Applying an unbounded SDF to an unbounded collection produces an unbounded collection",
+ PCollection.IsBounded.UNBOUNDED,
+ makeUnboundedCollection(pipeline)
+ .apply("unbounded to unbounded", new SplittableParDo<>(makeParDo(unboundedFn)))
+ .get(MAIN_OUTPUT_TAG)
+ .isBounded());
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunners.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunners.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunners.java
index f3cca6f..71dfd11 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunners.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunners.java
@@ -20,7 +20,7 @@ package org.apache.beam.runners.core;
import java.util.Collection;
import java.util.List;
import org.apache.beam.runners.core.ExecutionContext.StepContext;
-import org.apache.beam.runners.core.SplittableParDo.ProcessFn;
+import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems.ProcessFn;
import org.apache.beam.runners.core.StatefulDoFnRunner.CleanupTimer;
import org.apache.beam.runners.core.StatefulDoFnRunner.StateCleaner;
import org.apache.beam.sdk.options.PipelineOptions;
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/core-java/src/main/java/org/apache/beam/runners/core/ElementAndRestriction.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/ElementAndRestriction.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/ElementAndRestriction.java
deleted file mode 100644
index 4a5d0c4..0000000
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/ElementAndRestriction.java
+++ /dev/null
@@ -1,42 +0,0 @@
-/*
- * 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.core;
-
-import com.google.auto.value.AutoValue;
-import org.apache.beam.sdk.annotations.Experimental;
-import org.apache.beam.sdk.transforms.DoFn;
-
-/**
- * A tuple of an element and a restriction applied to processing it with a
- * <a href="https://s.apache.org/splittable-do-fn">splittable</a> {@link DoFn}.
- */
-@Experimental(Experimental.Kind.SPLITTABLE_DO_FN)
-@AutoValue
-public abstract class ElementAndRestriction<ElementT, RestrictionT> {
- /** The element to process. */
- public abstract ElementT element();
-
- /** The restriction applied to processing the element. */
- public abstract RestrictionT restriction();
-
- /** Constructs the {@link ElementAndRestriction}. */
- public static <InputT, RestrictionT> ElementAndRestriction<InputT, RestrictionT> of(
- InputT element, RestrictionT restriction) {
- return new AutoValue_ElementAndRestriction<>(element, restriction);
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/core-java/src/main/java/org/apache/beam/runners/core/ElementAndRestrictionCoder.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/ElementAndRestrictionCoder.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/ElementAndRestrictionCoder.java
deleted file mode 100644
index 4440b85..0000000
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/ElementAndRestrictionCoder.java
+++ /dev/null
@@ -1,88 +0,0 @@
-/*
- * 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.core;
-
-import com.google.common.collect.ImmutableList;
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
-import java.util.List;
-import org.apache.beam.sdk.annotations.Experimental;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.CoderException;
-import org.apache.beam.sdk.coders.StructuredCoder;
-
-/** A {@link Coder} for {@link ElementAndRestriction}. */
-@Experimental(Experimental.Kind.SPLITTABLE_DO_FN)
-public class ElementAndRestrictionCoder<ElementT, RestrictionT>
- extends StructuredCoder<ElementAndRestriction<ElementT, RestrictionT>> {
- private final Coder<ElementT> elementCoder;
- private final Coder<RestrictionT> restrictionCoder;
-
- /**
- * Creates an {@link ElementAndRestrictionCoder} from an element coder and a restriction coder.
- */
- public static <ElementT, RestrictionT> ElementAndRestrictionCoder<ElementT, RestrictionT> of(
- Coder<ElementT> elementCoder, Coder<RestrictionT> restrictionCoder) {
- return new ElementAndRestrictionCoder<>(elementCoder, restrictionCoder);
- }
-
- private ElementAndRestrictionCoder(
- Coder<ElementT> elementCoder, Coder<RestrictionT> restrictionCoder) {
- this.elementCoder = elementCoder;
- this.restrictionCoder = restrictionCoder;
- }
-
- @Override
- public void encode(
- ElementAndRestriction<ElementT, RestrictionT> value, OutputStream outStream)
- throws IOException {
- if (value == null) {
- throw new CoderException("cannot encode a null ElementAndRestriction");
- }
- elementCoder.encode(value.element(), outStream);
- restrictionCoder.encode(value.restriction(), outStream);
- }
-
- @Override
- public ElementAndRestriction<ElementT, RestrictionT> decode(InputStream inStream)
- throws IOException {
- ElementT key = elementCoder.decode(inStream);
- RestrictionT value = restrictionCoder.decode(inStream);
- return ElementAndRestriction.of(key, value);
- }
-
- @Override
- public List<? extends Coder<?>> getCoderArguments() {
- return ImmutableList.of(elementCoder, restrictionCoder);
- }
-
- @Override
- public void verifyDeterministic() throws NonDeterministicException {
- elementCoder.verifyDeterministic();
- restrictionCoder.verifyDeterministic();
- }
-
- public Coder<ElementT> getElementCoder() {
- return elementCoder;
- }
-
- public Coder<RestrictionT> getRestrictionCoder() {
- return restrictionCoder;
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/core-java/src/main/java/org/apache/beam/runners/core/ProcessFnRunner.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/ProcessFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/ProcessFnRunner.java
index 61f413f..31e86bd 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/ProcessFnRunner.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/ProcessFnRunner.java
@@ -18,13 +18,13 @@
package org.apache.beam.runners.core;
import static com.google.common.base.Preconditions.checkArgument;
-import static org.apache.beam.runners.core.SplittableParDo.ProcessFn;
import com.google.common.collect.Iterables;
import java.util.Collection;
import java.util.Collections;
import org.apache.beam.runners.core.StateNamespaces.WindowNamespace;
import org.apache.beam.runners.core.TimerInternals.TimerData;
+import org.apache.beam.runners.core.construction.ElementAndRestriction;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
@@ -32,7 +32,10 @@ import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.PCollectionView;
import org.joda.time.Instant;
-/** Runs a {@link ProcessFn} by constructing the appropriate contexts and passing them in. */
+/**
+ * Runs a {@link SplittableParDoViaKeyedWorkItems.ProcessFn} by constructing the appropriate
+ * contexts and passing them in.
+ */
public class ProcessFnRunner<InputT, OutputT, RestrictionT>
implements PushbackSideInputDoFnRunner<
KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT> {
http://git-wip-us.apache.org/repos/asf/beam/blob/7a2e829c/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDo.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDo.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDo.java
deleted file mode 100644
index 6503fa2..0000000
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDo.java
+++ /dev/null
@@ -1,602 +0,0 @@
-/*
- * 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.core;
-
-import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Preconditions.checkNotNull;
-
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.Iterables;
-import java.util.List;
-import java.util.UUID;
-import org.apache.beam.sdk.annotations.Experimental;
-import org.apache.beam.sdk.coders.CannotProvideCoderException;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.state.TimeDomain;
-import org.apache.beam.sdk.state.ValueState;
-import org.apache.beam.sdk.state.WatermarkHoldState;
-import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.GroupByKey;
-import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.transforms.SerializableFunction;
-import org.apache.beam.sdk.transforms.WithKeys;
-import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
-import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
-import org.apache.beam.sdk.transforms.reflect.DoFnSignature;
-import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
-import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
-import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
-import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
-import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.sdk.values.KV;
-import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PCollectionTuple;
-import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.sdk.values.TupleTagList;
-import org.apache.beam.sdk.values.WindowingStrategy;
-import org.joda.time.Instant;
-
-/**
- * A utility transform that executes a <a
- * href="https://s.apache.org/splittable-do-fn">splittable</a> {@link DoFn} by expanding it into a
- * network of simpler transforms:
- *
- * <ol>
- * <li>Pair each element with an initial restriction
- * <li>Split each restriction into sub-restrictions
- * <li>Assign a unique key to each element/restriction pair
- * <li>Group by key (so that work is partitioned by key and we can access state/timers)
- * <li>Process each keyed element/restriction pair with the splittable {@link DoFn}'s {@link
- * DoFn.ProcessElement} method, using state and timers API.
- * </ol>
- *
- * <p>This transform is intended as a helper for internal use by runners when implementing {@code
- * ParDo.of(splittable DoFn)}, but not for direct use by pipeline writers.
- */
-@Experimental(Experimental.Kind.SPLITTABLE_DO_FN)
-public class SplittableParDo<InputT, OutputT, RestrictionT>
- extends PTransform<PCollection<InputT>, PCollectionTuple> {
- private final ParDo.MultiOutput<InputT, OutputT> parDo;
-
- /**
- * Creates the transform for the given original multi-output {@link ParDo}.
- *
- * @param parDo The splittable {@link ParDo} transform.
- */
- public SplittableParDo(ParDo.MultiOutput<InputT, OutputT> parDo) {
- checkNotNull(parDo, "parDo must not be null");
- this.parDo = parDo;
- checkArgument(
- DoFnSignatures.getSignature(parDo.getFn().getClass()).processElement().isSplittable(),
- "fn must be a splittable DoFn");
- }
-
- @Override
- public PCollectionTuple expand(PCollection<InputT> input) {
- return applyTyped(input);
- }
-
- private PCollectionTuple applyTyped(PCollection<InputT> input) {
- DoFn<InputT, OutputT> fn = parDo.getFn();
- Coder<RestrictionT> restrictionCoder =
- DoFnInvokers.invokerFor(fn)
- .invokeGetRestrictionCoder(input.getPipeline().getCoderRegistry());
- PCollection<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>> keyedWorkItems =
- applySplitIntoKeyedWorkItems(input, fn, restrictionCoder);
- return keyedWorkItems.apply(
- "Process",
- new ProcessElements<>(
- fn,
- input.getCoder(),
- restrictionCoder,
- (WindowingStrategy<InputT, ?>) input.getWindowingStrategy(),
- parDo.getSideInputs(),
- parDo.getMainOutputTag(),
- parDo.getAdditionalOutputTags()));
- }
-
- private static <InputT, OutputT, RestrictionT>
- PCollection<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>>
- applySplitIntoKeyedWorkItems(
- PCollection<InputT> input,
- DoFn<InputT, OutputT> fn,
- Coder<RestrictionT> restrictionCoder) {
- Coder<ElementAndRestriction<InputT, RestrictionT>> splitCoder =
- ElementAndRestrictionCoder.of(input.getCoder(), restrictionCoder);
-
- PCollection<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>> keyedWorkItems =
- input
- .apply(
- "Pair with initial restriction",
- ParDo.of(new PairWithRestrictionFn<InputT, OutputT, RestrictionT>(fn)))
- .setCoder(splitCoder)
- .apply("Split restriction", ParDo.of(new SplitRestrictionFn<InputT, RestrictionT>(fn)))
- .setCoder(splitCoder)
- // ProcessFn requires all input elements to be in a single window and have a single
- // element per work item. This must precede the unique keying so each key has a single
- // associated element.
- .apply(
- "Explode windows",
- ParDo.of(new ExplodeWindowsFn<ElementAndRestriction<InputT, RestrictionT>>()))
- .apply(
- "Assign unique key",
- WithKeys.of(new RandomUniqueKeyFn<ElementAndRestriction<InputT, RestrictionT>>()))
- .apply(
- "Group by key",
- new GBKIntoKeyedWorkItems<String, ElementAndRestriction<InputT, RestrictionT>>())
- .setCoder(
- KeyedWorkItemCoder.of(
- StringUtf8Coder.of(),
- splitCoder,
- input.getWindowingStrategy().getWindowFn().windowCoder()));
- checkArgument(
- keyedWorkItems.getWindowingStrategy().getWindowFn() instanceof GlobalWindows,
- "GBKIntoKeyedWorkItems must produce a globally windowed collection, "
- + "but windowing strategy was: %s",
- keyedWorkItems.getWindowingStrategy());
- return keyedWorkItems;
- }
-
- /**
- * A {@link DoFn} that forces each of its outputs to be in a single window, by indicating to the
- * runner that it observes the window of its input element, so the runner is forced to apply it to
- * each input in a single window and thus its output is also in a single window.
- */
- private static class ExplodeWindowsFn<InputT> extends DoFn<InputT, InputT> {
- @ProcessElement
- public void process(ProcessContext c, BoundedWindow window) {
- c.output(c.element());
- }
- }
-
- /**
- * Runner-specific primitive {@link GroupByKey GroupByKey-like} {@link PTransform} that produces
- * {@link KeyedWorkItem KeyedWorkItems} so that downstream transforms can access state and timers.
- *
- * <p>Unlike a real {@link GroupByKey}, ignores the input's windowing and triggering strategy and
- * emits output immediately.
- */
- public static class GBKIntoKeyedWorkItems<KeyT, InputT>
- extends PTransform<PCollection<KV<KeyT, InputT>>, PCollection<KeyedWorkItem<KeyT, InputT>>> {
- @Override
- public PCollection<KeyedWorkItem<KeyT, InputT>> expand(PCollection<KV<KeyT, InputT>> input) {
- return PCollection.createPrimitiveOutputInternal(
- input.getPipeline(), WindowingStrategy.globalDefault(), input.isBounded());
- }
- }
-
- /**
- * Runner-specific primitive {@link PTransform} that invokes the {@link DoFn.ProcessElement}
- * method for a splittable {@link DoFn}.
- */
- public static class ProcessElements<
- InputT, OutputT, RestrictionT, TrackerT extends RestrictionTracker<RestrictionT>>
- extends PTransform<
- PCollection<? extends KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>>,
- PCollectionTuple> {
- private final DoFn<InputT, OutputT> fn;
- private final Coder<InputT> elementCoder;
- private final Coder<RestrictionT> restrictionCoder;
- private final WindowingStrategy<InputT, ?> windowingStrategy;
- private final List<PCollectionView<?>> sideInputs;
- private final TupleTag<OutputT> mainOutputTag;
- private final TupleTagList additionalOutputTags;
-
- /**
- * @param fn the splittable {@link DoFn}.
- * @param windowingStrategy the {@link WindowingStrategy} of the input collection.
- * @param sideInputs list of side inputs that should be available to the {@link DoFn}.
- * @param mainOutputTag {@link TupleTag Tag} of the {@link DoFn DoFn's} main output.
- * @param additionalOutputTags {@link TupleTagList Tags} of the {@link DoFn DoFn's} additional
- * outputs.
- */
- public ProcessElements(
- DoFn<InputT, OutputT> fn,
- Coder<InputT> elementCoder,
- Coder<RestrictionT> restrictionCoder,
- WindowingStrategy<InputT, ?> windowingStrategy,
- List<PCollectionView<?>> sideInputs,
- TupleTag<OutputT> mainOutputTag,
- TupleTagList additionalOutputTags) {
- this.fn = fn;
- this.elementCoder = elementCoder;
- this.restrictionCoder = restrictionCoder;
- this.windowingStrategy = windowingStrategy;
- this.sideInputs = sideInputs;
- this.mainOutputTag = mainOutputTag;
- this.additionalOutputTags = additionalOutputTags;
- }
-
- public DoFn<InputT, OutputT> getFn() {
- return fn;
- }
-
- public List<PCollectionView<?>> getSideInputs() {
- return sideInputs;
- }
-
- public TupleTag<OutputT> getMainOutputTag() {
- return mainOutputTag;
- }
-
- public TupleTagList getAdditionalOutputTags() {
- return additionalOutputTags;
- }
-
- public ProcessFn<InputT, OutputT, RestrictionT, TrackerT> newProcessFn(
- DoFn<InputT, OutputT> fn) {
- return new SplittableParDo.ProcessFn<>(
- fn, elementCoder, restrictionCoder, windowingStrategy);
- }
-
- @Override
- public PCollectionTuple expand(
- PCollection<? extends KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>>
- input) {
- DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());
- PCollectionTuple outputs =
- PCollectionTuple.ofPrimitiveOutputsInternal(
- input.getPipeline(),
- TupleTagList.of(mainOutputTag).and(additionalOutputTags.getAll()),
- windowingStrategy,
- input.isBounded().and(signature.isBoundedPerElement()));
-
- // Set output type descriptor similarly to how ParDo.MultiOutput does it.
- outputs.get(mainOutputTag).setTypeDescriptor(fn.getOutputTypeDescriptor());
-
- return outputs;
- }
-
- @Override
- public <T> Coder<T> getDefaultOutputCoder(
- PCollection<? extends KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>>
- input,
- PCollection<T> output)
- throws CannotProvideCoderException {
- // Similar logic to ParDo.MultiOutput.getOutputCoder.
- @SuppressWarnings("unchecked")
- KeyedWorkItemCoder<String, ElementAndRestriction<InputT, RestrictionT>> kwiCoder =
- (KeyedWorkItemCoder) input.getCoder();
- Coder<InputT> inputCoder =
- ((ElementAndRestrictionCoder<InputT, RestrictionT>) kwiCoder.getElementCoder())
- .getElementCoder();
- return input
- .getPipeline()
- .getCoderRegistry()
- .getCoder(output.getTypeDescriptor(), fn.getInputTypeDescriptor(), inputCoder);
- }
- }
-
- /**
- * Assigns a random unique key to each element of the input collection, so that the output
- * collection is effectively the same elements as input, but the per-key state and timers are now
- * effectively per-element.
- */
- private static class RandomUniqueKeyFn<T> implements SerializableFunction<T, String> {
- @Override
- public String apply(T input) {
- return UUID.randomUUID().toString();
- }
- }
-
- /**
- * Pairs each input element with its initial restriction using the given splittable {@link DoFn}.
- */
- private static class PairWithRestrictionFn<InputT, OutputT, RestrictionT>
- extends DoFn<InputT, ElementAndRestriction<InputT, RestrictionT>> {
- private DoFn<InputT, OutputT> fn;
- private transient DoFnInvoker<InputT, OutputT> invoker;
-
- PairWithRestrictionFn(DoFn<InputT, OutputT> fn) {
- this.fn = fn;
- }
-
- @Setup
- public void setup() {
- invoker = DoFnInvokers.invokerFor(fn);
- }
-
- @ProcessElement
- public void processElement(ProcessContext context) {
- context.output(
- ElementAndRestriction.of(
- context.element(),
- invoker.<RestrictionT>invokeGetInitialRestriction(context.element())));
- }
- }
-
- /**
- * The heart of splittable {@link DoFn} execution: processes a single (element, restriction) pair
- * by creating a tracker for the restriction and checkpointing/resuming processing later if
- * necessary.
- *
- * <p>Takes {@link KeyedWorkItem} and assumes that the KeyedWorkItem contains a single element
- * (or a single timer set by {@link ProcessFn itself}, in a single window. This is necessary
- * because {@link ProcessFn} sets timers, and timers are namespaced to a single window and it
- * should be the window of the input element.
- *
- * <p>See also: https://issues.apache.org/jira/browse/BEAM-1983
- */
- @VisibleForTesting
- public static class ProcessFn<
- InputT, OutputT, RestrictionT, TrackerT extends RestrictionTracker<RestrictionT>>
- extends DoFn<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT> {
- /**
- * The state cell containing a watermark hold for the output of this {@link DoFn}. The hold is
- * acquired during the first {@link DoFn.ProcessElement} call for each element and restriction,
- * and is released when the {@link DoFn.ProcessElement} call returns and there is no residual
- * restriction captured by the {@link SplittableProcessElementInvoker}.
- *
- * <p>A hold is needed to avoid letting the output watermark immediately progress together with
- * the input watermark when the first {@link DoFn.ProcessElement} call for this element
- * completes.
- */
- private static final StateTag<WatermarkHoldState> watermarkHoldTag =
- StateTags.makeSystemTagInternal(
- StateTags.<GlobalWindow>watermarkStateInternal(
- "hold", TimestampCombiner.LATEST));
-
- /**
- * The state cell containing a copy of the element. Written during the first {@link
- * DoFn.ProcessElement} call and read during subsequent calls in response to timer firings, when
- * the original element is no longer available.
- */
- private final StateTag<ValueState<WindowedValue<InputT>>> elementTag;
-
- /**
- * The state cell containing a restriction representing the unprocessed part of work for this
- * element.
- */
- private StateTag<ValueState<RestrictionT>> restrictionTag;
-
- private final DoFn<InputT, OutputT> fn;
- private final Coder<InputT> elementCoder;
- private final Coder<RestrictionT> restrictionCoder;
- private final WindowingStrategy<InputT, ?> inputWindowingStrategy;
-
- private transient StateInternalsFactory<String> stateInternalsFactory;
- private transient TimerInternalsFactory<String> timerInternalsFactory;
- private transient SplittableProcessElementInvoker<InputT, OutputT, RestrictionT, TrackerT>
- processElementInvoker;
-
- private transient DoFnInvoker<InputT, OutputT> invoker;
-
- public ProcessFn(
- DoFn<InputT, OutputT> fn,
- Coder<InputT> elementCoder,
- Coder<RestrictionT> restrictionCoder,
- WindowingStrategy<InputT, ?> inputWindowingStrategy) {
- this.fn = fn;
- this.elementCoder = elementCoder;
- this.restrictionCoder = restrictionCoder;
- this.inputWindowingStrategy = inputWindowingStrategy;
- this.elementTag =
- StateTags.value(
- "element",
- WindowedValue.getFullCoder(
- elementCoder, inputWindowingStrategy.getWindowFn().windowCoder()));
- this.restrictionTag = StateTags.value("restriction", restrictionCoder);
- }
-
- public void setStateInternalsFactory(StateInternalsFactory<String> stateInternalsFactory) {
- this.stateInternalsFactory = stateInternalsFactory;
- }
-
- public void setTimerInternalsFactory(TimerInternalsFactory<String> timerInternalsFactory) {
- this.timerInternalsFactory = timerInternalsFactory;
- }
-
- public void setProcessElementInvoker(
- SplittableProcessElementInvoker<InputT, OutputT, RestrictionT, TrackerT> invoker) {
- this.processElementInvoker = invoker;
- }
-
- public DoFn<InputT, OutputT> getFn() {
- return fn;
- }
-
- public Coder<InputT> getElementCoder() {
- return elementCoder;
- }
-
- public Coder<RestrictionT> getRestrictionCoder() {
- return restrictionCoder;
- }
-
- public WindowingStrategy<InputT, ?> getInputWindowingStrategy() {
- return inputWindowingStrategy;
- }
-
- @Setup
- public void setup() throws Exception {
- invoker = DoFnInvokers.invokerFor(fn);
- invoker.invokeSetup();
- }
-
- @Teardown
- public void tearDown() throws Exception {
- invoker.invokeTeardown();
- }
-
- @StartBundle
- public void startBundle(StartBundleContext c) throws Exception {
- invoker.invokeStartBundle(wrapContextAsStartBundle(c));
- }
-
- @FinishBundle
- public void finishBundle(FinishBundleContext c) throws Exception {
- invoker.invokeFinishBundle(wrapContextAsFinishBundle(c));
- }
-
- @ProcessElement
- public void processElement(final ProcessContext c) {
- String key = c.element().key();
- StateInternals stateInternals =
- stateInternalsFactory.stateInternalsForKey(key);
- TimerInternals timerInternals = timerInternalsFactory.timerInternalsForKey(key);
-
- // Initialize state (element and restriction) depending on whether this is the seed call.
- // The seed call is the first call for this element, which actually has the element.
- // Subsequent calls are timer firings and the element has to be retrieved from the state.
- TimerInternals.TimerData timer = Iterables.getOnlyElement(c.element().timersIterable(), null);
- boolean isSeedCall = (timer == null);
- StateNamespace stateNamespace;
- if (isSeedCall) {
- WindowedValue<ElementAndRestriction<InputT, RestrictionT>> windowedValue =
- Iterables.getOnlyElement(c.element().elementsIterable());
- BoundedWindow window = Iterables.getOnlyElement(windowedValue.getWindows());
- stateNamespace =
- StateNamespaces.window(
- (Coder<BoundedWindow>) inputWindowingStrategy.getWindowFn().windowCoder(), window);
- } else {
- stateNamespace = timer.getNamespace();
- }
-
- ValueState<WindowedValue<InputT>> elementState =
- stateInternals.state(stateNamespace, elementTag);
- ValueState<RestrictionT> restrictionState =
- stateInternals.state(stateNamespace, restrictionTag);
- WatermarkHoldState holdState =
- stateInternals.state(stateNamespace, watermarkHoldTag);
-
- ElementAndRestriction<WindowedValue<InputT>, RestrictionT> elementAndRestriction;
- if (isSeedCall) {
- WindowedValue<ElementAndRestriction<InputT, RestrictionT>> windowedValue =
- Iterables.getOnlyElement(c.element().elementsIterable());
- WindowedValue<InputT> element = windowedValue.withValue(windowedValue.getValue().element());
- elementState.write(element);
- elementAndRestriction =
- ElementAndRestriction.of(element, windowedValue.getValue().restriction());
- } else {
- // This is not the first ProcessElement call for this element/restriction - rather,
- // this is a timer firing, so we need to fetch the element and restriction from state.
- elementState.readLater();
- restrictionState.readLater();
- elementAndRestriction =
- ElementAndRestriction.of(elementState.read(), restrictionState.read());
- }
-
- final TrackerT tracker = invoker.invokeNewTracker(elementAndRestriction.restriction());
- SplittableProcessElementInvoker<InputT, OutputT, RestrictionT, TrackerT>.Result result =
- processElementInvoker.invokeProcessElement(
- invoker, elementAndRestriction.element(), tracker);
-
- // Save state for resuming.
- if (result.getResidualRestriction() == null) {
- // All work for this element/restriction is completed. Clear state and release hold.
- elementState.clear();
- restrictionState.clear();
- holdState.clear();
- return;
- }
- restrictionState.write(result.getResidualRestriction());
- Instant futureOutputWatermark = result.getFutureOutputWatermark();
- if (futureOutputWatermark == null) {
- futureOutputWatermark = elementAndRestriction.element().getTimestamp();
- }
- holdState.add(futureOutputWatermark);
- // Set a timer to continue processing this element.
- timerInternals.setTimer(
- TimerInternals.TimerData.of(
- stateNamespace, timerInternals.currentProcessingTime(), TimeDomain.PROCESSING_TIME));
- }
-
- private DoFn<InputT, OutputT>.StartBundleContext wrapContextAsStartBundle(
- final StartBundleContext baseContext) {
- return fn.new StartBundleContext() {
- @Override
- public PipelineOptions getPipelineOptions() {
- return baseContext.getPipelineOptions();
- }
-
- private void throwUnsupportedOutput() {
- throw new UnsupportedOperationException(
- String.format(
- "Splittable DoFn can only output from @%s",
- ProcessElement.class.getSimpleName()));
- }
- };
- }
-
- private DoFn<InputT, OutputT>.FinishBundleContext wrapContextAsFinishBundle(
- final FinishBundleContext baseContext) {
- return fn.new FinishBundleContext() {
- @Override
- public void output(OutputT output, Instant timestamp, BoundedWindow window) {
- throwUnsupportedOutput();
- }
-
- @Override
- public <T> void output(TupleTag<T> tag, T output, Instant timestamp, BoundedWindow window) {
- throwUnsupportedOutput();
- }
-
- @Override
- public PipelineOptions getPipelineOptions() {
- return baseContext.getPipelineOptions();
- }
-
- private void throwUnsupportedOutput() {
- throw new UnsupportedOperationException(
- String.format(
- "Splittable DoFn can only output from @%s",
- ProcessElement.class.getSimpleName()));
- }
- };
- }
-
- }
-
- /** Splits the restriction using the given {@link DoFn.SplitRestriction} method. */
- private static class SplitRestrictionFn<InputT, RestrictionT>
- extends DoFn<
- ElementAndRestriction<InputT, RestrictionT>,
- ElementAndRestriction<InputT, RestrictionT>> {
- private final DoFn<InputT, ?> splittableFn;
- private transient DoFnInvoker<InputT, ?> invoker;
-
- SplitRestrictionFn(DoFn<InputT, ?> splittableFn) {
- this.splittableFn = splittableFn;
- }
-
- @Setup
- public void setup() {
- invoker = DoFnInvokers.invokerFor(splittableFn);
- }
-
- @ProcessElement
- public void processElement(final ProcessContext c) {
- final InputT element = c.element().element();
- invoker.invokeSplitRestriction(
- element,
- c.element().restriction(),
- new OutputReceiver<RestrictionT>() {
- @Override
- public void output(RestrictionT part) {
- c.output(ElementAndRestriction.of(element, part));
- }
- });
- }
- }
-}
[4/4] beam git commit: This closes #3156
Posted by jk...@apache.org.
This closes #3156
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/c93b3aa3
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/c93b3aa3
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/c93b3aa3
Branch: refs/heads/master
Commit: c93b3aa3d2081558a776fa12c564efb5f15090ab
Parents: 2e0cf00 7a2e829
Author: Eugene Kirpichov <ki...@google.com>
Authored: Thu May 18 16:15:00 2017 -0700
Committer: Eugene Kirpichov <ki...@google.com>
Committed: Thu May 18 16:15:00 2017 -0700
----------------------------------------------------------------------
.../apache/beam/runners/apex/ApexRunner.java | 7 +-
.../translation/ApexPipelineTranslator.java | 11 +-
.../apex/translation/ParDoTranslator.java | 11 +-
.../operators/ApexParDoOperator.java | 10 +-
.../construction/ElementAndRestriction.java | 42 ++
.../ElementAndRestrictionCoder.java | 88 +++
.../core/construction/SplittableParDo.java | 298 +++++++++
.../ElementAndRestrictionCoderTest.java | 126 ++++
.../core/construction/SplittableParDoTest.java | 157 +++++
.../apache/beam/runners/core/DoFnRunners.java | 2 +-
.../runners/core/ElementAndRestriction.java | 42 --
.../core/ElementAndRestrictionCoder.java | 88 ---
.../beam/runners/core/ProcessFnRunner.java | 7 +-
.../beam/runners/core/SplittableParDo.java | 602 ------------------
.../core/SplittableParDoViaKeyedWorkItems.java | 418 +++++++++++++
.../core/ElementAndRestrictionCoderTest.java | 127 ----
.../core/SplittableParDoProcessFnTest.java | 515 ++++++++++++++++
.../beam/runners/core/SplittableParDoTest.java | 606 -------------------
...ectGBKIntoKeyedWorkItemsOverrideFactory.java | 9 +-
.../beam/runners/direct/DirectRunner.java | 10 +-
.../direct/KeyedPValueTrackingVisitor.java | 4 +-
.../direct/ParDoMultiOverrideFactory.java | 2 +-
...littableProcessElementsEvaluatorFactory.java | 11 +-
.../direct/TransformEvaluatorRegistry.java | 4 +-
.../flink/FlinkStreamingPipelineTranslator.java | 7 +-
.../FlinkStreamingTransformTranslators.java | 29 +-
.../streaming/SplittableDoFnOperator.java | 12 +-
27 files changed, 1721 insertions(+), 1524 deletions(-)
----------------------------------------------------------------------