You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by jb...@apache.org on 2017/07/20 17:09:28 UTC
[01/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Repository: beam
Updated Branches:
refs/heads/DSL_SQL 97a156c30 -> 152115e89
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/typehints/trivial_inference.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/typehints/trivial_inference.py b/sdks/python/apache_beam/typehints/trivial_inference.py
index c740596..977ea06 100644
--- a/sdks/python/apache_beam/typehints/trivial_inference.py
+++ b/sdks/python/apache_beam/typehints/trivial_inference.py
@@ -40,7 +40,8 @@ def instance_to_type(o):
"""
t = type(o)
if o is None:
- return type(None)
+ # TODO(robertwb): Eliminate inconsistent use of None vs. NoneType.
+ return None
elif t not in typehints.DISALLOWED_PRIMITIVE_TYPES:
if t == types.InstanceType:
return o.__class__
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/typehints/trivial_inference_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/typehints/trivial_inference_test.py b/sdks/python/apache_beam/typehints/trivial_inference_test.py
index e7f451d..ac00baa 100644
--- a/sdks/python/apache_beam/typehints/trivial_inference_test.py
+++ b/sdks/python/apache_beam/typehints/trivial_inference_test.py
@@ -60,13 +60,6 @@ class TrivialInferenceTest(unittest.TestCase):
self.assertReturnType(any_tuple,
reverse, [trivial_inference.Const((1, 2, 3))])
- def testNoneReturn(self):
- def func(a):
- if a == 5:
- return a
- return None
- self.assertReturnType(typehints.Union[int, type(None)], func, [int])
-
def testListComprehension(self):
self.assertReturnType(
typehints.List[int],
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/utils/plugin.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/utils/plugin.py b/sdks/python/apache_beam/utils/plugin.py
deleted file mode 100644
index 563b93c..0000000
--- a/sdks/python/apache_beam/utils/plugin.py
+++ /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.
-#
-
-"""A BeamPlugin base class.
-
-For experimental usage only; no backwards-compatibility guarantees.
-"""
-
-
-class BeamPlugin(object):
- """Plugin base class to be extended by dependent users such as FileSystem.
- Any instantiated subclass will be imported at worker startup time."""
-
- @classmethod
- def get_all_subclasses(cls):
- """Get all the subclasses of the BeamPlugin class."""
- all_subclasses = []
- for subclass in cls.__subclasses__():
- all_subclasses.append(subclass)
- all_subclasses.extend(subclass.get_all_subclasses())
- return all_subclasses
-
- @classmethod
- def get_all_plugin_paths(cls):
- """Get full import paths of the BeamPlugin subclass."""
- def fullname(o):
- return o.__module__ + "." + o.__name__
- return [fullname(o) for o in cls.get_all_subclasses()]
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/utils/timestamp.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/utils/timestamp.py b/sdks/python/apache_beam/utils/timestamp.py
index b3e840e..5d1b48c 100644
--- a/sdks/python/apache_beam/utils/timestamp.py
+++ b/sdks/python/apache_beam/utils/timestamp.py
@@ -208,8 +208,3 @@ class Duration(object):
def __mod__(self, other):
other = Duration.of(other)
return Duration(micros=self.micros % other.micros)
-
-
-# The minimum granularity / interval expressible in a Timestamp / Duration
-# object.
-TIME_GRANULARITY = Duration(micros=1)
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/utils/urns.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/utils/urns.py b/sdks/python/apache_beam/utils/urns.py
index e553eea..849b8e3 100644
--- a/sdks/python/apache_beam/utils/urns.py
+++ b/sdks/python/apache_beam/utils/urns.py
@@ -102,7 +102,7 @@ class RunnerApiFn(object):
Prefer overriding self.to_runner_api_parameter.
"""
- from apache_beam.portability.api import beam_runner_api_pb2
+ from apache_beam.runners.api import beam_runner_api_pb2
urn, typed_param = self.to_runner_api_parameter(context)
return beam_runner_api_pb2.SdkFunctionSpec(
spec=beam_runner_api_pb2.FunctionSpec(
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/version.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/version.py b/sdks/python/apache_beam/version.py
index 8b0a430..ae92a23 100644
--- a/sdks/python/apache_beam/version.py
+++ b/sdks/python/apache_beam/version.py
@@ -18,4 +18,4 @@
"""Apache Beam SDK version information and utilities."""
-__version__ = '2.2.0.dev'
+__version__ = '2.1.0.dev'
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/gen_protos.py
----------------------------------------------------------------------
diff --git a/sdks/python/gen_protos.py b/sdks/python/gen_protos.py
index a3d963d..92b8414 100644
--- a/sdks/python/gen_protos.py
+++ b/sdks/python/gen_protos.py
@@ -35,7 +35,7 @@ BEAM_PROTO_PATHS = [
os.path.join('..', 'common', 'fn-api', 'src', 'main', 'proto')
]
-PYTHON_OUTPUT_PATH = os.path.join('apache_beam', 'portability', 'api')
+PYTHON_OUTPUT_PATH = os.path.join('apache_beam', 'runners', 'api')
def generate_proto_files():
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/python/pom.xml b/sdks/python/pom.xml
index 1077689..1295654 100644
--- a/sdks/python/pom.xml
+++ b/sdks/python/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/run_pylint.sh
----------------------------------------------------------------------
diff --git a/sdks/python/run_pylint.sh b/sdks/python/run_pylint.sh
index 2691be4..7808136 100755
--- a/sdks/python/run_pylint.sh
+++ b/sdks/python/run_pylint.sh
@@ -46,7 +46,7 @@ EXCLUDED_GENERATED_FILES=(
"apache_beam/io/gcp/internal/clients/storage/storage_v1_client.py"
"apache_beam/io/gcp/internal/clients/storage/storage_v1_messages.py"
"apache_beam/coders/proto2_coder_test_messages_pb2.py"
-apache_beam/portability/api/*pb2*.py
+apache_beam/runners/api/*pb2*.py
)
FILES_TO_IGNORE=""
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/setup.py
----------------------------------------------------------------------
diff --git a/sdks/python/setup.py b/sdks/python/setup.py
index da82466..051043b 100644
--- a/sdks/python/setup.py
+++ b/sdks/python/setup.py
@@ -112,17 +112,14 @@ REQUIRED_SETUP_PACKAGES = [
REQUIRED_TEST_PACKAGES = [
'pyhamcrest>=1.9,<2.0',
- # Six required by nose plugins management.
- 'six>=1.9',
]
GCP_REQUIREMENTS = [
'google-apitools>=0.5.10,<=0.5.11',
'proto-google-cloud-datastore-v1>=0.90.0,<=0.90.4',
'googledatastore==7.0.1',
- 'google-cloud-pubsub==0.26.0',
# GCP packages required by tests
- 'google-cloud-bigquery==0.25.0',
+ 'google-cloud-bigquery>=0.23.0,<0.25.0',
]
[25/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ParDoTranslationTest.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ParDoTranslationTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ParDoTranslationTest.java
index a87a16d..a8490bf 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ParDoTranslationTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ParDoTranslationTest.java
@@ -23,9 +23,9 @@ import static org.hamcrest.Matchers.instanceOf;
import static org.junit.Assert.assertThat;
import com.google.common.collect.ImmutableList;
+import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
-import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.VarIntCoder;
@@ -143,30 +143,22 @@ public class ParDoTranslationTest {
inputs.putAll(parDo.getAdditionalInputs());
PCollectionTuple output = mainInput.apply(parDo);
- SdkComponents sdkComponents = SdkComponents.create();
-
- // Encode
- RunnerApi.PTransform protoTransform =
- PTransformTranslation.toProto(
+ SdkComponents components = SdkComponents.create();
+ String transformId =
+ components.registerPTransform(
AppliedPTransform.<PCollection<KV<Long, String>>, PCollection<Void>, MultiOutput>of(
"foo", inputs, output.expand(), parDo, p),
- sdkComponents);
- Components protoComponents = sdkComponents.toComponents();
-
- // Decode
- Pipeline rehydratedPipeline = Pipeline.create();
+ Collections.<AppliedPTransform<?, ?, ?>>emptyList());
+ Components protoComponents = components.toComponents();
+ RunnerApi.PTransform protoTransform = protoComponents.getTransformsOrThrow(transformId);
ParDoPayload parDoPayload =
protoTransform.getSpec().getParameter().unpack(ParDoPayload.class);
for (PCollectionView<?> view : parDo.getSideInputs()) {
SideInput sideInput = parDoPayload.getSideInputsOrThrow(view.getTagInternal().getId());
PCollectionView<?> restoredView =
- ParDoTranslation.viewFromProto(
- sideInput,
- view.getTagInternal().getId(),
- view.getPCollection(),
- protoTransform,
- protoComponents);
+ ParDoTranslation.fromProto(
+ sideInput, view.getTagInternal().getId(), protoTransform, protoComponents);
assertThat(restoredView.getTagInternal(), equalTo(view.getTagInternal()));
assertThat(restoredView.getViewFn(), instanceOf(view.getViewFn().getClass()));
assertThat(
@@ -177,7 +169,7 @@ public class ParDoTranslationTest {
view.getWindowingStrategyInternal().fixDefaults()));
assertThat(restoredView.getCoderInternal(), equalTo(view.getCoderInternal()));
}
- String mainInputId = sdkComponents.registerPCollection(mainInput);
+ String mainInputId = components.registerPCollection(mainInput);
assertThat(
ParDoTranslation.getMainInput(protoTransform, protoComponents),
equalTo(protoComponents.getPcollectionsOrThrow(mainInputId)));
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/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
index 267232c..6e4d6c4 100644
--- 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
@@ -17,7 +17,6 @@
*/
package org.apache.beam.runners.core.construction;
-import static org.apache.beam.sdk.transforms.DoFn.ProcessContinuation.stop;
import static org.junit.Assert.assertEquals;
import java.io.Serializable;
@@ -25,6 +24,8 @@ 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;
@@ -69,6 +70,7 @@ public class SplittableParDoTest {
public void checkDone() {}
}
+ @BoundedPerElement
private static class BoundedFakeFn extends DoFn<Integer, String> {
@ProcessElement
public void processElement(ProcessContext context, SomeRestrictionTracker tracker) {}
@@ -79,12 +81,10 @@ public class SplittableParDoTest {
}
}
+ @UnboundedPerElement
private static class UnboundedFakeFn extends DoFn<Integer, String> {
@ProcessElement
- public ProcessContinuation processElement(
- ProcessContext context, SomeRestrictionTracker tracker) {
- return stop();
- }
+ public void processElement(ProcessContext context, SomeRestrictionTracker tracker) {}
@GetInitialRestriction
public SomeRestriction getInitialRestriction(Integer element) {
@@ -122,14 +122,14 @@ public class SplittableParDoTest {
"Applying a bounded SDF to a bounded collection produces a bounded collection",
PCollection.IsBounded.BOUNDED,
makeBoundedCollection(pipeline)
- .apply("bounded to bounded", SplittableParDo.forJavaParDo(makeParDo(boundedFn)))
+ .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", SplittableParDo.forJavaParDo(makeParDo(boundedFn)))
+ .apply("bounded to unbounded", new SplittableParDo<>(makeParDo(boundedFn)))
.get(MAIN_OUTPUT_TAG)
.isBounded());
}
@@ -143,14 +143,14 @@ public class SplittableParDoTest {
"Applying an unbounded SDF to a bounded collection produces a bounded collection",
PCollection.IsBounded.UNBOUNDED,
makeBoundedCollection(pipeline)
- .apply("unbounded to bounded", SplittableParDo.forJavaParDo(makeParDo(unboundedFn)))
+ .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", SplittableParDo.forJavaParDo(makeParDo(unboundedFn)))
+ .apply("unbounded to unbounded", new SplittableParDo<>(makeParDo(unboundedFn)))
.get(MAIN_OUTPUT_TAG)
.isBounded());
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/TransformInputsTest.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/TransformInputsTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/TransformInputsTest.java
deleted file mode 100644
index f5b2c11..0000000
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/TransformInputsTest.java
+++ /dev/null
@@ -1,166 +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.construction;
-
-import static org.junit.Assert.assertThat;
-
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.Map;
-import org.apache.beam.sdk.coders.VoidCoder;
-import org.apache.beam.sdk.io.GenerateSequence;
-import org.apache.beam.sdk.runners.AppliedPTransform;
-import org.apache.beam.sdk.testing.TestPipeline;
-import org.apache.beam.sdk.transforms.Create;
-import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PDone;
-import org.apache.beam.sdk.values.PInput;
-import org.apache.beam.sdk.values.POutput;
-import org.apache.beam.sdk.values.PValue;
-import org.apache.beam.sdk.values.TupleTag;
-import org.hamcrest.Matchers;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.rules.ExpectedException;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-/** Tests for {@link TransformInputs}. */
-@RunWith(JUnit4.class)
-public class TransformInputsTest {
- @Rule public TestPipeline pipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false);
- @Rule public ExpectedException thrown = ExpectedException.none();
-
- @Test
- public void nonAdditionalInputsWithNoInputSucceeds() {
- AppliedPTransform<PInput, POutput, TestTransform> transform =
- AppliedPTransform.of(
- "input-free",
- Collections.<TupleTag<?>, PValue>emptyMap(),
- Collections.<TupleTag<?>, PValue>emptyMap(),
- new TestTransform(),
- pipeline);
-
- assertThat(TransformInputs.nonAdditionalInputs(transform), Matchers.<PValue>empty());
- }
-
- @Test
- public void nonAdditionalInputsWithOneMainInputSucceeds() {
- PCollection<Long> input = pipeline.apply(GenerateSequence.from(1L));
- AppliedPTransform<PInput, POutput, TestTransform> transform =
- AppliedPTransform.of(
- "input-single",
- Collections.<TupleTag<?>, PValue>singletonMap(new TupleTag<Long>() {}, input),
- Collections.<TupleTag<?>, PValue>emptyMap(),
- new TestTransform(),
- pipeline);
-
- assertThat(
- TransformInputs.nonAdditionalInputs(transform), Matchers.<PValue>containsInAnyOrder(input));
- }
-
- @Test
- public void nonAdditionalInputsWithMultipleNonAdditionalInputsSucceeds() {
- Map<TupleTag<?>, PValue> allInputs = new HashMap<>();
- PCollection<Integer> mainInts = pipeline.apply("MainInput", Create.of(12, 3));
- allInputs.put(new TupleTag<Integer>() {}, mainInts);
- PCollection<Void> voids = pipeline.apply("VoidInput", Create.empty(VoidCoder.of()));
- allInputs.put(new TupleTag<Void>() {}, voids);
- AppliedPTransform<PInput, POutput, TestTransform> transform =
- AppliedPTransform.of(
- "additional-free",
- allInputs,
- Collections.<TupleTag<?>, PValue>emptyMap(),
- new TestTransform(),
- pipeline);
-
- assertThat(
- TransformInputs.nonAdditionalInputs(transform),
- Matchers.<PValue>containsInAnyOrder(voids, mainInts));
- }
-
- @Test
- public void nonAdditionalInputsWithAdditionalInputsSucceeds() {
- Map<TupleTag<?>, PValue> additionalInputs = new HashMap<>();
- additionalInputs.put(new TupleTag<String>() {}, pipeline.apply(Create.of("1, 2", "3")));
- additionalInputs.put(new TupleTag<Long>() {}, pipeline.apply(GenerateSequence.from(3L)));
-
- Map<TupleTag<?>, PValue> allInputs = new HashMap<>();
- PCollection<Integer> mainInts = pipeline.apply("MainInput", Create.of(12, 3));
- allInputs.put(new TupleTag<Integer>() {}, mainInts);
- PCollection<Void> voids = pipeline.apply("VoidInput", Create.empty(VoidCoder.of()));
- allInputs.put(
- new TupleTag<Void>() {}, voids);
- allInputs.putAll(additionalInputs);
-
- AppliedPTransform<PInput, POutput, TestTransform> transform =
- AppliedPTransform.of(
- "additional",
- allInputs,
- Collections.<TupleTag<?>, PValue>emptyMap(),
- new TestTransform(additionalInputs),
- pipeline);
-
- assertThat(
- TransformInputs.nonAdditionalInputs(transform),
- Matchers.<PValue>containsInAnyOrder(mainInts, voids));
- }
-
- @Test
- public void nonAdditionalInputsWithOnlyAdditionalInputsThrows() {
- Map<TupleTag<?>, PValue> additionalInputs = new HashMap<>();
- additionalInputs.put(new TupleTag<String>() {}, pipeline.apply(Create.of("1, 2", "3")));
- additionalInputs.put(new TupleTag<Long>() {}, pipeline.apply(GenerateSequence.from(3L)));
-
- AppliedPTransform<PInput, POutput, TestTransform> transform =
- AppliedPTransform.of(
- "additional-only",
- additionalInputs,
- Collections.<TupleTag<?>, PValue>emptyMap(),
- new TestTransform(additionalInputs),
- pipeline);
-
- thrown.expect(IllegalArgumentException.class);
- thrown.expectMessage("at least one");
- TransformInputs.nonAdditionalInputs(transform);
- }
-
- private static class TestTransform extends PTransform<PInput, POutput> {
- private final Map<TupleTag<?>, PValue> additionalInputs;
-
- private TestTransform() {
- this(Collections.<TupleTag<?>, PValue>emptyMap());
- }
-
- private TestTransform(Map<TupleTag<?>, PValue> additionalInputs) {
- this.additionalInputs = additionalInputs;
- }
-
- @Override
- public POutput expand(PInput input) {
- return PDone.in(input.getPipeline());
- }
-
- @Override
- public Map<TupleTag<?>, PValue> getAdditionalInputs() {
- return additionalInputs;
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WindowingStrategyTranslationTest.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WindowingStrategyTranslationTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WindowingStrategyTranslationTest.java
index 7a57fd7..e406545 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WindowingStrategyTranslationTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WindowingStrategyTranslationTest.java
@@ -116,8 +116,5 @@ public class WindowingStrategyTranslationTest {
protoComponents.getCodersOrThrow(
components.registerCoder(windowingStrategy.getWindowFn().windowCoder()));
- assertThat(
- proto.getAssignsToOneWindow(),
- equalTo(windowingStrategy.getWindowFn().assignsToOneWindow()));
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WriteFilesTranslationTest.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WriteFilesTranslationTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WriteFilesTranslationTest.java
index 283df16..739034c 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WriteFilesTranslationTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WriteFilesTranslationTest.java
@@ -26,10 +26,8 @@ import java.util.Objects;
import javax.annotation.Nullable;
import org.apache.beam.sdk.common.runner.v1.RunnerApi;
import org.apache.beam.sdk.common.runner.v1.RunnerApi.ParDoPayload;
-import org.apache.beam.sdk.io.DynamicFileDestinations;
import org.apache.beam.sdk.io.FileBasedSink;
import org.apache.beam.sdk.io.FileBasedSink.FilenamePolicy;
-import org.apache.beam.sdk.io.FileBasedSink.OutputFileHints;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.WriteFiles;
import org.apache.beam.sdk.io.fs.ResourceId;
@@ -38,8 +36,6 @@ import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.transforms.SerializableFunction;
-import org.apache.beam.sdk.transforms.SerializableFunctions;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PDone;
import org.junit.Test;
@@ -60,17 +56,16 @@ public class WriteFilesTranslationTest {
@RunWith(Parameterized.class)
public static class TestWriteFilesPayloadTranslation {
@Parameters(name = "{index}: {0}")
- public static Iterable<WriteFiles<Object, Void, Object>> data() {
- SerializableFunction<Object, Object> format = SerializableFunctions.constant(null);
- return ImmutableList.of(
- WriteFiles.to(new DummySink(), format),
- WriteFiles.to(new DummySink(), format).withWindowedWrites(),
- WriteFiles.to(new DummySink(), format).withNumShards(17),
- WriteFiles.to(new DummySink(), format).withWindowedWrites().withNumShards(42));
+ public static Iterable<WriteFiles<?>> data() {
+ return ImmutableList.<WriteFiles<?>>of(
+ WriteFiles.to(new DummySink()),
+ WriteFiles.to(new DummySink()).withWindowedWrites(),
+ WriteFiles.to(new DummySink()).withNumShards(17),
+ WriteFiles.to(new DummySink()).withWindowedWrites().withNumShards(42));
}
@Parameter(0)
- public WriteFiles<String, Void, String> writeFiles;
+ public WriteFiles<String> writeFiles;
public static TestPipeline p = TestPipeline.create().enableAbandonedNodeEnforcement(false);
@@ -85,7 +80,7 @@ public class WriteFilesTranslationTest {
assertThat(payload.getWindowedWrites(), equalTo(writeFiles.isWindowedWrites()));
assertThat(
- (FileBasedSink<String, Void>) WriteFilesTranslation.sinkFromProto(payload.getSink()),
+ (FileBasedSink<String>) WriteFilesTranslation.sinkFromProto(payload.getSink()),
equalTo(writeFiles.getSink()));
}
@@ -94,9 +89,9 @@ public class WriteFilesTranslationTest {
PCollection<String> input = p.apply(Create.of("hello"));
PDone output = input.apply(writeFiles);
- AppliedPTransform<PCollection<String>, PDone, WriteFiles<String, Void, String>>
- appliedPTransform =
- AppliedPTransform.of("foo", input.expand(), output.expand(), writeFiles, p);
+ AppliedPTransform<PCollection<String>, PDone, WriteFiles<String>> appliedPTransform =
+ AppliedPTransform.<PCollection<String>, PDone, WriteFiles<String>>of(
+ "foo", input.expand(), output.expand(), writeFiles, p);
assertThat(
WriteFilesTranslation.isRunnerDeterminedSharding(appliedPTransform),
@@ -106,9 +101,7 @@ public class WriteFilesTranslationTest {
WriteFilesTranslation.isWindowedWrites(appliedPTransform),
equalTo(writeFiles.isWindowedWrites()));
- assertThat(
- WriteFilesTranslation.<String, Void, String>getSink(appliedPTransform),
- equalTo(writeFiles.getSink()));
+ assertThat(WriteFilesTranslation.getSink(appliedPTransform), equalTo(writeFiles.getSink()));
}
}
@@ -116,16 +109,16 @@ public class WriteFilesTranslationTest {
* A simple {@link FileBasedSink} for testing serialization/deserialization. Not mocked to avoid
* any issues serializing mocks.
*/
- private static class DummySink extends FileBasedSink<Object, Void> {
+ private static class DummySink extends FileBasedSink<String> {
DummySink() {
super(
StaticValueProvider.of(FileSystems.matchNewResource("nowhere", false)),
- DynamicFileDestinations.constant(new DummyFilenamePolicy()));
+ new DummyFilenamePolicy());
}
@Override
- public WriteOperation<Object, Void> createWriteOperation() {
+ public WriteOperation<String> createWriteOperation() {
return new DummyWriteOperation(this);
}
@@ -137,39 +130,46 @@ public class WriteFilesTranslationTest {
DummySink that = (DummySink) other;
- return getTempDirectoryProvider().isAccessible()
- && that.getTempDirectoryProvider().isAccessible()
- && getTempDirectoryProvider().get().equals(that.getTempDirectoryProvider().get());
+ return getFilenamePolicy().equals(((DummySink) other).getFilenamePolicy())
+ && getBaseOutputDirectoryProvider().isAccessible()
+ && that.getBaseOutputDirectoryProvider().isAccessible()
+ && getBaseOutputDirectoryProvider()
+ .get()
+ .equals(that.getBaseOutputDirectoryProvider().get());
}
@Override
public int hashCode() {
return Objects.hash(
DummySink.class,
- getTempDirectoryProvider().isAccessible() ? getTempDirectoryProvider().get() : null);
+ getFilenamePolicy(),
+ getBaseOutputDirectoryProvider().isAccessible()
+ ? getBaseOutputDirectoryProvider().get()
+ : null);
}
}
- private static class DummyWriteOperation extends FileBasedSink.WriteOperation<Object, Void> {
- public DummyWriteOperation(FileBasedSink<Object, Void> sink) {
+ private static class DummyWriteOperation extends FileBasedSink.WriteOperation<String> {
+ public DummyWriteOperation(FileBasedSink<String> sink) {
super(sink);
}
@Override
- public FileBasedSink.Writer<Object, Void> createWriter() throws Exception {
+ public FileBasedSink.Writer<String> createWriter() throws Exception {
throw new UnsupportedOperationException("Should never be called.");
}
}
private static class DummyFilenamePolicy extends FilenamePolicy {
@Override
- public ResourceId windowedFilename(WindowedContext c, OutputFileHints outputFileHints) {
+ public ResourceId windowedFilename(
+ ResourceId outputDirectory, WindowedContext c, String extension) {
throw new UnsupportedOperationException("Should never be called.");
}
@Nullable
@Override
- public ResourceId unwindowedFilename(Context c, OutputFileHints outputFileHints) {
+ public ResourceId unwindowedFilename(ResourceId outputDirectory, Context c, String extension) {
throw new UnsupportedOperationException("Should never be called.");
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/pom.xml
----------------------------------------------------------------------
diff --git a/runners/core-java/pom.xml b/runners/core-java/pom.xml
index 8c8e599..c3a8d25 100644
--- a/runners/core-java/pom.xml
+++ b/runners/core-java/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-runners-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java
index 28938c1..1cf1509 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/LateDataDroppingDoFnRunner.java
@@ -134,27 +134,26 @@ public class LateDataDroppingDoFnRunner<K, InputT, OutputT, W extends BoundedWin
// The element is too late for this window.
droppedDueToLateness.inc();
WindowTracing.debug(
- "{}: Dropping element at {} for key:{}; window:{} "
- + "since too far behind inputWatermark:{}; outputWatermark:{}",
- LateDataFilter.class.getSimpleName(),
- input.getTimestamp(),
- key,
- window,
- timerInternals.currentInputWatermarkTime(),
+ "ReduceFnRunner.processElement: Dropping element at {} for key:{}; window:{} "
+ + "since too far behind inputWatermark:{}; outputWatermark:{}",
+ input.getTimestamp(), key, window, timerInternals.currentInputWatermarkTime(),
timerInternals.currentOutputWatermarkTime());
}
}
- Iterable<WindowedValue<InputT>> nonLateElements =
- Iterables.filter(
- concatElements,
- new Predicate<WindowedValue<InputT>>() {
- @Override
- public boolean apply(WindowedValue<InputT> input) {
- BoundedWindow window = Iterables.getOnlyElement(input.getWindows());
- return !canDropDueToExpiredWindow(window);
- }
- });
+ Iterable<WindowedValue<InputT>> nonLateElements = Iterables.filter(
+ concatElements,
+ new Predicate<WindowedValue<InputT>>() {
+ @Override
+ public boolean apply(WindowedValue<InputT> input) {
+ BoundedWindow window = Iterables.getOnlyElement(input.getWindows());
+ if (canDropDueToExpiredWindow(window)) {
+ return false;
+ } else {
+ return true;
+ }
+ }
+ });
return nonLateElements;
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java
index 0c956d5..2db6531 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java
@@ -96,7 +96,7 @@ public class OutputAndTimeBoundedSplittableProcessElementInvoker<
final WindowedValue<InputT> element,
final TrackerT tracker) {
final ProcessContext processContext = new ProcessContext(element, tracker);
- DoFn.ProcessContinuation cont = invoker.invokeProcessElement(
+ invoker.invokeProcessElement(
new DoFnInvoker.ArgumentProvider<InputT, OutputT>() {
@Override
public DoFn<InputT, OutputT>.ProcessContext processContext(
@@ -118,11 +118,6 @@ public class OutputAndTimeBoundedSplittableProcessElementInvoker<
}
@Override
- public PipelineOptions pipelineOptions() {
- return pipelineOptions;
- }
-
- @Override
public StartBundleContext startBundleContext(DoFn<InputT, OutputT> doFn) {
throw new IllegalStateException(
"Should not access startBundleContext() from @"
@@ -155,37 +150,10 @@ public class OutputAndTimeBoundedSplittableProcessElementInvoker<
"Access to timers not supported in Splittable DoFn");
}
});
- // TODO: verify that if there was a failed tryClaim() call, then cont.shouldResume() is false.
- // Currently we can't verify this because there are no hooks into tryClaim().
- // See https://issues.apache.org/jira/browse/BEAM-2607
- RestrictionT residual = processContext.extractCheckpoint();
- if (cont.shouldResume()) {
- if (residual == null) {
- // No checkpoint had been taken by the runner while the ProcessElement call ran, however
- // the call says that not the whole restriction has been processed. So we need to take
- // a checkpoint now: checkpoint() guarantees that the primary restriction describes exactly
- // the work that was done in the current ProcessElement call, and returns a residual
- // restriction that describes exactly the work that wasn't done in the current call.
- residual = tracker.checkpoint();
- } else {
- // A checkpoint was taken by the runner, and then the ProcessElement call returned resume()
- // without making more tryClaim() calls (since no tryClaim() calls can succeed after
- // checkpoint(), and since if it had made a failed tryClaim() call, it should have returned
- // stop()).
- // This means that the resulting primary restriction and the taken checkpoint already
- // accurately describe respectively the work that was and wasn't done in the current
- // ProcessElement call.
- // In other words, if we took a checkpoint *after* ProcessElement completed (like in the
- // branch above), it would have been equivalent to this one.
- }
- } else {
- // The ProcessElement call returned stop() - that means the tracker's current restriction
- // has been fully processed by the call. A checkpoint may or may not have been taken in
- // "residual"; if it was, then we'll need to process it; if no, then we don't - nothing
- // special needs to be done.
- }
+
tracker.checkDone();
- return new Result(residual, cont, processContext.getLastReportedWatermark());
+ return new Result(
+ processContext.extractCheckpoint(), processContext.getLastReportedWatermark());
}
private class ProcessContext extends DoFn<InputT, OutputT>.ProcessContext {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/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 88275d6..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
@@ -24,11 +24,11 @@ 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;
import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
import org.joda.time.Instant;
@@ -38,13 +38,16 @@ import org.joda.time.Instant;
*/
public class ProcessFnRunner<InputT, OutputT, RestrictionT>
implements PushbackSideInputDoFnRunner<
- KeyedWorkItem<String, KV<InputT, RestrictionT>>, OutputT> {
- private final DoFnRunner<KeyedWorkItem<String, KV<InputT, RestrictionT>>, OutputT> underlying;
+ KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT> {
+ private final DoFnRunner<
+ KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT>
+ underlying;
private final Collection<PCollectionView<?>> views;
private final ReadyCheckingSideInputReader sideInputReader;
ProcessFnRunner(
- DoFnRunner<KeyedWorkItem<String, KV<InputT, RestrictionT>>, OutputT> underlying,
+ DoFnRunner<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT>
+ underlying,
Collection<PCollectionView<?>> views,
ReadyCheckingSideInputReader sideInputReader) {
this.underlying = underlying;
@@ -58,9 +61,10 @@ public class ProcessFnRunner<InputT, OutputT, RestrictionT>
}
@Override
- public Iterable<WindowedValue<KeyedWorkItem<String, KV<InputT, RestrictionT>>>>
+ public Iterable<WindowedValue<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>>>
processElementInReadyWindows(
- WindowedValue<KeyedWorkItem<String, KV<InputT, RestrictionT>>> windowedKWI) {
+ WindowedValue<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>>
+ windowedKWI) {
checkTrivialOuterWindows(windowedKWI);
BoundedWindow window = getUnderlyingWindow(windowedKWI.getValue());
if (!isReady(window)) {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java
index 634a2d1..62d519f 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java
@@ -29,6 +29,7 @@ import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
+import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;
@@ -50,7 +51,6 @@ import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.transforms.windowing.PaneInfo.Timing;
import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
-import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.transforms.windowing.Window.ClosingBehavior;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.util.WindowTracing;
@@ -637,9 +637,11 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
}
/**
- * A descriptor of the activation for a window based on a timer.
+ * Enriches TimerData with state necessary for processing a timer as well as
+ * common queries about a timer.
*/
- private class WindowActivation {
+ private class EnrichedTimerData {
+ public final Instant timestamp;
public final ReduceFn<K, InputT, OutputT, W>.Context directContext;
public final ReduceFn<K, InputT, OutputT, W>.Context renamedContext;
// If this is an end-of-window timer then we may need to set a garbage collection timer
@@ -650,34 +652,18 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
// end-of-window time to be a signal to garbage collect.
public final boolean isGarbageCollection;
- WindowActivation(
+ EnrichedTimerData(
+ TimerData timer,
ReduceFn<K, InputT, OutputT, W>.Context directContext,
ReduceFn<K, InputT, OutputT, W>.Context renamedContext) {
+ this.timestamp = timer.getTimestamp();
this.directContext = directContext;
this.renamedContext = renamedContext;
W window = directContext.window();
-
- // The output watermark is before the end of the window if it is either unknown
- // or it is known to be before it. If it is unknown, that means that there hasn't been
- // enough data to advance it.
- boolean outputWatermarkBeforeEOW =
- timerInternals.currentOutputWatermarkTime() == null
- || !timerInternals.currentOutputWatermarkTime().isAfter(window.maxTimestamp());
-
- // The "end of the window" is reached when the local input watermark (for this key) surpasses
- // it but the local output watermark (also for this key) has not. After data is emitted and
- // the output watermark hold is released, the output watermark on this key will immediately
- // exceed the end of the window (otherwise we could see multiple ON_TIME outputs)
- this.isEndOfWindow =
- timerInternals.currentInputWatermarkTime().isAfter(window.maxTimestamp())
- && outputWatermarkBeforeEOW;
-
- // The "GC time" is reached when the input watermark surpasses the end of the window
- // plus allowed lateness. After this, the window is expired and expunged.
- this.isGarbageCollection =
- timerInternals
- .currentInputWatermarkTime()
- .isAfter(LateDataUtils.garbageCollectionTime(window, windowingStrategy));
+ this.isEndOfWindow = TimeDomain.EVENT_TIME == timer.getDomain()
+ && timer.getTimestamp().equals(window.maxTimestamp());
+ Instant cleanupTime = LateDataUtils.garbageCollectionTime(window, windowingStrategy);
+ this.isGarbageCollection = !timer.getTimestamp().isBefore(cleanupTime);
}
// Has this window had its trigger finish?
@@ -696,47 +682,24 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
return;
}
- // Create a reusable context for each window and begin prefetching necessary
+ // Create a reusable context for each timer and begin prefetching necessary
// state.
- Map<BoundedWindow, WindowActivation> windowActivations = new HashMap();
-
+ List<EnrichedTimerData> enrichedTimers = new LinkedList();
for (TimerData timer : timers) {
checkArgument(timer.getNamespace() instanceof WindowNamespace,
"Expected timer to be in WindowNamespace, but was in %s", timer.getNamespace());
@SuppressWarnings("unchecked")
WindowNamespace<W> windowNamespace = (WindowNamespace<W>) timer.getNamespace();
W window = windowNamespace.getWindow();
-
- WindowTracing.debug("{}: Received timer key:{}; window:{}; data:{} with "
- + "inputWatermark:{}; outputWatermark:{}",
- ReduceFnRunner.class.getSimpleName(),
- key, window, timer,
- timerInternals.currentInputWatermarkTime(),
- timerInternals.currentOutputWatermarkTime());
-
- // Processing time timers for an expired window are ignored, just like elements
- // that show up too late. Window GC is management by an event time timer
- if (TimeDomain.EVENT_TIME != timer.getDomain() && windowIsExpired(window)) {
- continue;
- }
-
- // How a window is processed is a function only of the current state, not the details
- // of the timer. This makes us robust to large leaps in processing time and watermark
- // time, where both EOW and GC timers come in together and we need to GC and emit
- // the final pane.
- if (windowActivations.containsKey(window)) {
- continue;
- }
-
ReduceFn<K, InputT, OutputT, W>.Context directContext =
contextFactory.base(window, StateStyle.DIRECT);
ReduceFn<K, InputT, OutputT, W>.Context renamedContext =
contextFactory.base(window, StateStyle.RENAMED);
- WindowActivation windowActivation = new WindowActivation(directContext, renamedContext);
- windowActivations.put(window, windowActivation);
+ EnrichedTimerData enrichedTimer = new EnrichedTimerData(timer, directContext, renamedContext);
+ enrichedTimers.add(enrichedTimer);
// Perform prefetching of state to determine if the trigger should fire.
- if (windowActivation.isGarbageCollection) {
+ if (enrichedTimer.isGarbageCollection) {
triggerRunner.prefetchIsClosed(directContext.state());
} else {
triggerRunner.prefetchShouldFire(directContext.window(), directContext.state());
@@ -744,7 +707,7 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
}
// For those windows that are active and open, prefetch the triggering or emitting state.
- for (WindowActivation timer : windowActivations.values()) {
+ for (EnrichedTimerData timer : enrichedTimers) {
if (timer.windowIsActiveAndOpen()) {
ReduceFn<K, InputT, OutputT, W>.Context directContext = timer.directContext;
if (timer.isGarbageCollection) {
@@ -757,27 +720,25 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
}
// Perform processing now that everything is prefetched.
- for (WindowActivation windowActivation : windowActivations.values()) {
- ReduceFn<K, InputT, OutputT, W>.Context directContext = windowActivation.directContext;
- ReduceFn<K, InputT, OutputT, W>.Context renamedContext = windowActivation.renamedContext;
+ for (EnrichedTimerData timer : enrichedTimers) {
+ ReduceFn<K, InputT, OutputT, W>.Context directContext = timer.directContext;
+ ReduceFn<K, InputT, OutputT, W>.Context renamedContext = timer.renamedContext;
- if (windowActivation.isGarbageCollection) {
- WindowTracing.debug(
- "{}: Cleaning up for key:{}; window:{} with inputWatermark:{}; outputWatermark:{}",
- ReduceFnRunner.class.getSimpleName(),
- key,
- directContext.window(),
+ if (timer.isGarbageCollection) {
+ WindowTracing.debug("ReduceFnRunner.onTimer: Cleaning up for key:{}; window:{} at {} with "
+ + "inputWatermark:{}; outputWatermark:{}",
+ key, directContext.window(), timer.timestamp,
timerInternals.currentInputWatermarkTime(),
timerInternals.currentOutputWatermarkTime());
- boolean windowIsActiveAndOpen = windowActivation.windowIsActiveAndOpen();
+ boolean windowIsActiveAndOpen = timer.windowIsActiveAndOpen();
if (windowIsActiveAndOpen) {
// We need to call onTrigger to emit the final pane if required.
// The final pane *may* be ON_TIME if no prior ON_TIME pane has been emitted,
// and the watermark has passed the end of the window.
@Nullable
Instant newHold = onTrigger(
- directContext, renamedContext, true /* isFinished */, windowActivation.isEndOfWindow);
+ directContext, renamedContext, true /* isFinished */, timer.isEndOfWindow);
checkState(newHold == null, "Hold placed at %s despite isFinished being true.", newHold);
}
@@ -785,20 +746,18 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
// see elements for it again.
clearAllState(directContext, renamedContext, windowIsActiveAndOpen);
} else {
- WindowTracing.debug(
- "{}.onTimers: Triggering for key:{}; window:{} at {} with "
+ WindowTracing.debug("ReduceFnRunner.onTimer: Triggering for key:{}; window:{} at {} with "
+ "inputWatermark:{}; outputWatermark:{}",
- key,
- directContext.window(),
+ key, directContext.window(), timer.timestamp,
timerInternals.currentInputWatermarkTime(),
timerInternals.currentOutputWatermarkTime());
- if (windowActivation.windowIsActiveAndOpen()
+ if (timer.windowIsActiveAndOpen()
&& triggerRunner.shouldFire(
directContext.window(), directContext.timers(), directContext.state())) {
emit(directContext, renamedContext);
}
- if (windowActivation.isEndOfWindow) {
+ if (timer.isEndOfWindow) {
// If the window strategy trigger includes a watermark trigger then at this point
// there should be no data holds, either because we'd already cleared them on an
// earlier onTrigger, or because we just cleared them on the above emit.
@@ -960,9 +919,8 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
// The pane has elements.
return true;
}
- if (timing == Timing.ON_TIME
- && windowingStrategy.getOnTimeBehavior() == Window.OnTimeBehavior.FIRE_ALWAYS) {
- // This is an empty ON_TIME pane.
+ if (timing == Timing.ON_TIME) {
+ // This is the unique ON_TIME pane.
return true;
}
if (isFinished && windowingStrategy.getClosingBehavior() == ClosingBehavior.FIRE_ALWAYS) {
@@ -990,8 +948,13 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
private Instant onTrigger(
final ReduceFn<K, InputT, OutputT, W>.Context directContext,
ReduceFn<K, InputT, OutputT, W>.Context renamedContext,
- final boolean isFinished, boolean isEndOfWindow)
+ boolean isFinished, boolean isEndOfWindow)
throws Exception {
+ Instant inputWM = timerInternals.currentInputWatermarkTime();
+
+ // Calculate the pane info.
+ final PaneInfo pane = paneInfoTracker.getNextPaneInfo(directContext, isFinished).read();
+
// Extract the window hold, and as a side effect clear it.
final WatermarkHold.OldAndNewHolds pair =
watermarkHold.extractAndRelease(renamedContext, isFinished).read();
@@ -1000,13 +963,7 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
@Nullable Instant newHold = pair.newHold;
final boolean isEmpty = nonEmptyPanes.isEmpty(renamedContext.state()).read();
- if (isEmpty
- && windowingStrategy.getClosingBehavior() == ClosingBehavior.FIRE_IF_NON_EMPTY
- && windowingStrategy.getOnTimeBehavior() == Window.OnTimeBehavior.FIRE_IF_NON_EMPTY) {
- return newHold;
- }
- Instant inputWM = timerInternals.currentInputWatermarkTime();
if (newHold != null) {
// We can't be finished yet.
checkState(
@@ -1038,9 +995,6 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
}
}
- // Calculate the pane info.
- final PaneInfo pane = paneInfoTracker.getNextPaneInfo(directContext, isFinished).read();
-
// Only emit a pane if it has data or empty panes are observable.
if (needToEmit(isEmpty, isFinished, pane.getTiming())) {
// Run reduceFn.onTrigger method.
@@ -1051,11 +1005,9 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
@Override
public void output(OutputT toOutput) {
// We're going to output panes, so commit the (now used) PaneInfo.
- // This is unnecessary if the trigger isFinished since the saved
+ // TODO: This is unnecessary if the trigger isFinished since the saved
// state will be immediately deleted.
- if (!isFinished) {
- paneInfoTracker.storeCurrentPaneInfo(directContext, pane);
- }
+ paneInfoTracker.storeCurrentPaneInfo(directContext, pane);
// Output the actual value.
outputter.outputWindowedValue(
@@ -1129,9 +1081,4 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
}
}
- private boolean windowIsExpired(BoundedWindow w) {
- return timerInternals
- .currentInputWatermarkTime()
- .isAfter(w.maxTimestamp().plus(windowingStrategy.getAllowedLateness()));
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java
index c3bfef6..7d7babd 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java
@@ -233,11 +233,6 @@ public class SimpleDoFnRunner<InputT, OutputT> implements DoFnRunner<InputT, Out
}
@Override
- public PipelineOptions pipelineOptions() {
- return getPipelineOptions();
- }
-
- @Override
public DoFn<InputT, OutputT>.StartBundleContext startBundleContext(DoFn<InputT, OutputT> doFn) {
return this;
}
@@ -303,11 +298,6 @@ public class SimpleDoFnRunner<InputT, OutputT> implements DoFnRunner<InputT, Out
}
@Override
- public PipelineOptions pipelineOptions() {
- return getPipelineOptions();
- }
-
- @Override
public DoFn<InputT, OutputT>.StartBundleContext startBundleContext(DoFn<InputT, OutputT> doFn) {
throw new UnsupportedOperationException(
"Cannot access StartBundleContext outside of @StartBundle method.");
@@ -477,11 +467,6 @@ public class SimpleDoFnRunner<InputT, OutputT> implements DoFnRunner<InputT, Out
}
@Override
- public PipelineOptions pipelineOptions() {
- return getPipelineOptions();
- }
-
- @Override
public DoFn<InputT, OutputT>.StartBundleContext startBundleContext(DoFn<InputT, OutputT> doFn) {
throw new UnsupportedOperationException("StartBundleContext parameters are not supported.");
}
@@ -583,11 +568,6 @@ public class SimpleDoFnRunner<InputT, OutputT> implements DoFnRunner<InputT, Out
}
@Override
- public PipelineOptions pipelineOptions() {
- return getPipelineOptions();
- }
-
- @Override
public DoFn<InputT, OutputT>.StartBundleContext startBundleContext(DoFn<InputT, OutputT> doFn) {
throw new UnsupportedOperationException("StartBundleContext parameters are not supported.");
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/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
index 6e97645..c4b086a 100644
--- 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
@@ -21,6 +21,7 @@ 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.PTransformTranslation.RawPTransform;
import org.apache.beam.runners.core.construction.ReplacementOutputs;
@@ -85,15 +86,15 @@ public class SplittableParDoViaKeyedWorkItems {
/** Overrides a {@link ProcessKeyedElements} into {@link SplittableProcessViaKeyedWorkItems}. */
public static class OverrideFactory<InputT, OutputT, RestrictionT>
implements PTransformOverrideFactory<
- PCollection<KV<String, KV<InputT, RestrictionT>>>, PCollectionTuple,
- ProcessKeyedElements<InputT, OutputT, RestrictionT>> {
+ PCollection<KV<String, ElementAndRestriction<InputT, RestrictionT>>>, PCollectionTuple,
+ ProcessKeyedElements<InputT, OutputT, RestrictionT>> {
@Override
public PTransformReplacement<
- PCollection<KV<String, KV<InputT, RestrictionT>>>, PCollectionTuple>
+ PCollection<KV<String, ElementAndRestriction<InputT, RestrictionT>>>, PCollectionTuple>
getReplacementTransform(
AppliedPTransform<
- PCollection<KV<String, KV<InputT, RestrictionT>>>, PCollectionTuple,
- ProcessKeyedElements<InputT, OutputT, RestrictionT>>
+ PCollection<KV<String, ElementAndRestriction<InputT, RestrictionT>>>,
+ PCollectionTuple, ProcessKeyedElements<InputT, OutputT, RestrictionT>>
transform) {
return PTransformReplacement.of(
PTransformReplacements.getSingletonMainInput(transform),
@@ -112,7 +113,8 @@ public class SplittableParDoViaKeyedWorkItems {
* method for a splittable {@link DoFn}.
*/
public static class SplittableProcessViaKeyedWorkItems<InputT, OutputT, RestrictionT>
- extends PTransform<PCollection<KV<String, KV<InputT, RestrictionT>>>, PCollectionTuple> {
+ extends PTransform<
+ PCollection<KV<String, ElementAndRestriction<InputT, RestrictionT>>>, PCollectionTuple> {
private final ProcessKeyedElements<InputT, OutputT, RestrictionT> original;
public SplittableProcessViaKeyedWorkItems(
@@ -121,13 +123,15 @@ public class SplittableParDoViaKeyedWorkItems {
}
@Override
- public PCollectionTuple expand(PCollection<KV<String, KV<InputT, RestrictionT>>> input) {
+ public PCollectionTuple expand(
+ PCollection<KV<String, ElementAndRestriction<InputT, RestrictionT>>> input) {
return input
- .apply(new GBKIntoKeyedWorkItems<String, KV<InputT, RestrictionT>>())
+ .apply(new GBKIntoKeyedWorkItems<String, ElementAndRestriction<InputT, RestrictionT>>())
.setCoder(
KeyedWorkItemCoder.of(
StringUtf8Coder.of(),
- ((KvCoder<String, KV<InputT, RestrictionT>>) input.getCoder()).getValueCoder(),
+ ((KvCoder<String, ElementAndRestriction<InputT, RestrictionT>>) input.getCoder())
+ .getValueCoder(),
input.getWindowingStrategy().getWindowFn().windowCoder()))
.apply(new ProcessElements<>(original));
}
@@ -137,7 +141,8 @@ public class SplittableParDoViaKeyedWorkItems {
public static class ProcessElements<
InputT, OutputT, RestrictionT, TrackerT extends RestrictionTracker<RestrictionT>>
extends PTransform<
- PCollection<KeyedWorkItem<String, KV<InputT, RestrictionT>>>, PCollectionTuple> {
+ PCollection<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>>,
+ PCollectionTuple> {
private final ProcessKeyedElements<InputT, OutputT, RestrictionT> original;
public ProcessElements(ProcessKeyedElements<InputT, OutputT, RestrictionT> original) {
@@ -171,7 +176,7 @@ public class SplittableParDoViaKeyedWorkItems {
@Override
public PCollectionTuple expand(
- PCollection<KeyedWorkItem<String, KV<InputT, RestrictionT>>> input) {
+ PCollection<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>> input) {
return ProcessKeyedElements.createPrimitiveOutputFor(
input,
original.getFn(),
@@ -196,12 +201,12 @@ public class SplittableParDoViaKeyedWorkItems {
@VisibleForTesting
public static class ProcessFn<
InputT, OutputT, RestrictionT, TrackerT extends RestrictionTracker<RestrictionT>>
- extends DoFn<KeyedWorkItem<String, KV<InputT, RestrictionT>>, OutputT> {
+ 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 {@link
- * ProcessContinuation#stop()}.
+ * 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
@@ -316,7 +321,7 @@ public class SplittableParDoViaKeyedWorkItems {
boolean isSeedCall = (timer == null);
StateNamespace stateNamespace;
if (isSeedCall) {
- WindowedValue<KV<InputT, RestrictionT>> windowedValue =
+ WindowedValue<ElementAndRestriction<InputT, RestrictionT>> windowedValue =
Iterables.getOnlyElement(c.element().elementsIterable());
BoundedWindow window = Iterables.getOnlyElement(windowedValue.getWindows());
stateNamespace =
@@ -332,25 +337,27 @@ public class SplittableParDoViaKeyedWorkItems {
stateInternals.state(stateNamespace, restrictionTag);
WatermarkHoldState holdState = stateInternals.state(stateNamespace, watermarkHoldTag);
- KV<WindowedValue<InputT>, RestrictionT> elementAndRestriction;
+ ElementAndRestriction<WindowedValue<InputT>, RestrictionT> elementAndRestriction;
if (isSeedCall) {
- WindowedValue<KV<InputT, RestrictionT>> windowedValue =
+ WindowedValue<ElementAndRestriction<InputT, RestrictionT>> windowedValue =
Iterables.getOnlyElement(c.element().elementsIterable());
- WindowedValue<InputT> element = windowedValue.withValue(windowedValue.getValue().getKey());
+ WindowedValue<InputT> element = windowedValue.withValue(windowedValue.getValue().element());
elementState.write(element);
- elementAndRestriction = KV.of(element, windowedValue.getValue().getValue());
+ 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 = KV.of(elementState.read(), restrictionState.read());
+ elementAndRestriction =
+ ElementAndRestriction.of(elementState.read(), restrictionState.read());
}
- final TrackerT tracker = invoker.invokeNewTracker(elementAndRestriction.getValue());
+ final TrackerT tracker = invoker.invokeNewTracker(elementAndRestriction.restriction());
SplittableProcessElementInvoker<InputT, OutputT, RestrictionT, TrackerT>.Result result =
processElementInvoker.invokeProcessElement(
- invoker, elementAndRestriction.getKey(), tracker);
+ invoker, elementAndRestriction.element(), tracker);
// Save state for resuming.
if (result.getResidualRestriction() == null) {
@@ -363,14 +370,13 @@ public class SplittableParDoViaKeyedWorkItems {
restrictionState.write(result.getResidualRestriction());
Instant futureOutputWatermark = result.getFutureOutputWatermark();
if (futureOutputWatermark == null) {
- futureOutputWatermark = elementAndRestriction.getKey().getTimestamp();
+ futureOutputWatermark = elementAndRestriction.element().getTimestamp();
}
- Instant wakeupTime =
- timerInternals.currentProcessingTime().plus(result.getContinuation().resumeDelay());
holdState.add(futureOutputWatermark);
// Set a timer to continue processing this element.
timerInternals.setTimer(
- TimerInternals.TimerData.of(stateNamespace, wakeupTime, TimeDomain.PROCESSING_TIME));
+ TimerInternals.TimerData.of(
+ stateNamespace, timerInternals.currentProcessingTime(), TimeDomain.PROCESSING_TIME));
}
private DoFn<InputT, OutputT>.StartBundleContext wrapContextAsStartBundle(
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableProcessElementInvoker.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableProcessElementInvoker.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableProcessElementInvoker.java
index 7732df3..ced6c01 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableProcessElementInvoker.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableProcessElementInvoker.java
@@ -17,8 +17,6 @@
*/
package org.apache.beam.runners.core;
-import static com.google.common.base.Preconditions.checkNotNull;
-
import javax.annotation.Nullable;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.reflect.DoFnInvoker;
@@ -36,35 +34,20 @@ public abstract class SplittableProcessElementInvoker<
public class Result {
@Nullable
private final RestrictionT residualRestriction;
- private final DoFn.ProcessContinuation continuation;
private final Instant futureOutputWatermark;
public Result(
- @Nullable RestrictionT residualRestriction,
- DoFn.ProcessContinuation continuation,
- Instant futureOutputWatermark) {
- this.continuation = checkNotNull(continuation);
- if (continuation.shouldResume()) {
- checkNotNull(residualRestriction);
- }
+ @Nullable RestrictionT residualRestriction, Instant futureOutputWatermark) {
this.residualRestriction = residualRestriction;
this.futureOutputWatermark = futureOutputWatermark;
}
- /**
- * Can be {@code null} only if {@link #getContinuation} specifies the call should not resume.
- * However, the converse is not true: this can be non-null even if {@link #getContinuation}
- * is {@link DoFn.ProcessContinuation#stop()}.
- */
+ /** If {@code null}, means the call should not resume. */
@Nullable
public RestrictionT getResidualRestriction() {
return residualRestriction;
}
- public DoFn.ProcessContinuation getContinuation() {
- return continuation;
- }
-
public Instant getFutureOutputWatermark() {
return futureOutputWatermark;
}
@@ -74,8 +57,8 @@ public abstract class SplittableProcessElementInvoker<
* Invokes the {@link DoFn.ProcessElement} method using the given {@link DoFnInvoker} for the
* original {@link DoFn}, on the given element and with the given {@link RestrictionTracker}.
*
- * @return Information on how to resume the call: residual restriction, a {@link
- * DoFn.ProcessContinuation}, and a future output watermark.
+ * @return Information on how to resume the call: residual restriction and a
+ * future output watermark.
*/
public abstract Result invokeProcessElement(
DoFnInvoker<InputT, OutputT> invoker, WindowedValue<InputT> element, TrackerT tracker);
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/main/java/org/apache/beam/runners/core/SystemReduceFn.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SystemReduceFn.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SystemReduceFn.java
index 3144bd6..c189b0d 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/SystemReduceFn.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/SystemReduceFn.java
@@ -18,7 +18,6 @@
package org.apache.beam.runners.core;
-import com.google.common.annotations.VisibleForTesting;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.state.BagState;
import org.apache.beam.sdk.state.CombiningState;
@@ -104,11 +103,6 @@ public abstract class SystemReduceFn<K, InputT, AccumT, OutputT, W extends Bound
this.bufferTag = bufferTag;
}
- @VisibleForTesting
- StateTag<? extends GroupingState<InputT, OutputT>> getBufferTag() {
- return bufferTag;
- }
-
@Override
public void processValue(ProcessValueContext c) throws Exception {
c.state().access(bufferTag).add(c.value());
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterAllStateMachine.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterAllStateMachine.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterAllStateMachine.java
index 3530ed1..0f0c17c 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterAllStateMachine.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterAllStateMachine.java
@@ -23,6 +23,7 @@ import com.google.common.base.Joiner;
import com.google.common.collect.ImmutableList;
import java.util.Arrays;
import java.util.List;
+import org.apache.beam.runners.core.triggers.TriggerStateMachine.OnceTriggerStateMachine;
import org.apache.beam.sdk.annotations.Experimental;
/**
@@ -30,7 +31,7 @@ import org.apache.beam.sdk.annotations.Experimental;
* have fired.
*/
@Experimental(Experimental.Kind.TRIGGER)
-public class AfterAllStateMachine extends TriggerStateMachine {
+public class AfterAllStateMachine extends OnceTriggerStateMachine {
private AfterAllStateMachine(List<TriggerStateMachine> subTriggers) {
super(subTriggers);
@@ -41,11 +42,11 @@ public class AfterAllStateMachine extends TriggerStateMachine {
* Returns an {@code AfterAll} {@code Trigger} with the given subtriggers.
*/
@SafeVarargs
- public static TriggerStateMachine of(TriggerStateMachine... triggers) {
+ public static OnceTriggerStateMachine of(TriggerStateMachine... triggers) {
return new AfterAllStateMachine(Arrays.<TriggerStateMachine>asList(triggers));
}
- public static TriggerStateMachine of(Iterable<? extends TriggerStateMachine> triggers) {
+ public static OnceTriggerStateMachine of(Iterable<? extends TriggerStateMachine> triggers) {
return new AfterAllStateMachine(ImmutableList.copyOf(triggers));
}
@@ -77,21 +78,24 @@ public class AfterAllStateMachine extends TriggerStateMachine {
*/
@Override
public boolean shouldFire(TriggerContext context) throws Exception {
- for (ExecutableTriggerStateMachine subTrigger : context.trigger().subTriggers()) {
- if (!context.forTrigger(subTrigger).trigger().isFinished()
- && !subTrigger.invokeShouldFire(context)) {
+ for (ExecutableTriggerStateMachine subtrigger : context.trigger().subTriggers()) {
+ if (!context.forTrigger(subtrigger).trigger().isFinished()
+ && !subtrigger.invokeShouldFire(context)) {
return false;
}
}
return true;
}
+ /**
+ * Invokes {@link #onFire} for all subtriggers, eliding redundant calls to {@link #shouldFire}
+ * because they all must be ready to fire.
+ */
@Override
- public void onFire(TriggerContext context) throws Exception {
- for (ExecutableTriggerStateMachine subTrigger : context.trigger().subTriggers()) {
- subTrigger.invokeOnFire(context);
+ public void onOnlyFiring(TriggerContext context) throws Exception {
+ for (ExecutableTriggerStateMachine subtrigger : context.trigger().subTriggers()) {
+ subtrigger.invokeOnFire(context);
}
- context.trigger().setFinished(true);
}
@Override
@@ -99,6 +103,7 @@ public class AfterAllStateMachine extends TriggerStateMachine {
StringBuilder builder = new StringBuilder("AfterAll.of(");
Joiner.on(", ").appendTo(builder, subTriggers);
builder.append(")");
+
return builder.toString();
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterDelayFromFirstElementStateMachine.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterDelayFromFirstElementStateMachine.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterDelayFromFirstElementStateMachine.java
index 06c2066..8d8d0de 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterDelayFromFirstElementStateMachine.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterDelayFromFirstElementStateMachine.java
@@ -27,6 +27,7 @@ import org.apache.beam.runners.core.StateAccessor;
import org.apache.beam.runners.core.StateMerging;
import org.apache.beam.runners.core.StateTag;
import org.apache.beam.runners.core.StateTags;
+import org.apache.beam.runners.core.triggers.TriggerStateMachine.OnceTriggerStateMachine;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.coders.InstantCoder;
import org.apache.beam.sdk.state.CombiningState;
@@ -49,7 +50,7 @@ import org.joda.time.format.PeriodFormatter;
// This class should be inlined to subclasses and deleted, simplifying them too
// https://issues.apache.org/jira/browse/BEAM-1486
@Experimental(Experimental.Kind.TRIGGER)
-public abstract class AfterDelayFromFirstElementStateMachine extends TriggerStateMachine {
+public abstract class AfterDelayFromFirstElementStateMachine extends OnceTriggerStateMachine {
protected static final List<SerializableFunction<Instant, Instant>> IDENTITY =
ImmutableList.<SerializableFunction<Instant, Instant>>of();
@@ -236,9 +237,8 @@ public abstract class AfterDelayFromFirstElementStateMachine extends TriggerStat
}
@Override
- public final void onFire(TriggerContext context) throws Exception {
+ protected void onOnlyFiring(TriggerStateMachine.TriggerContext context) throws Exception {
clear(context);
- context.trigger().setFinished(true);
}
protected Instant computeTargetTimestamp(Instant time) {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterFirstStateMachine.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterFirstStateMachine.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterFirstStateMachine.java
index 58c24c5..840a65c 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterFirstStateMachine.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterFirstStateMachine.java
@@ -23,6 +23,7 @@ import com.google.common.base.Joiner;
import com.google.common.collect.ImmutableList;
import java.util.Arrays;
import java.util.List;
+import org.apache.beam.runners.core.triggers.TriggerStateMachine.OnceTriggerStateMachine;
import org.apache.beam.sdk.annotations.Experimental;
/**
@@ -30,7 +31,7 @@ import org.apache.beam.sdk.annotations.Experimental;
* sub-triggers have fired.
*/
@Experimental(Experimental.Kind.TRIGGER)
-public class AfterFirstStateMachine extends TriggerStateMachine {
+public class AfterFirstStateMachine extends OnceTriggerStateMachine {
AfterFirstStateMachine(List<TriggerStateMachine> subTriggers) {
super(subTriggers);
@@ -41,12 +42,12 @@ public class AfterFirstStateMachine extends TriggerStateMachine {
* Returns an {@code AfterFirst} {@code Trigger} with the given subtriggers.
*/
@SafeVarargs
- public static TriggerStateMachine of(
+ public static OnceTriggerStateMachine of(
TriggerStateMachine... triggers) {
return new AfterFirstStateMachine(Arrays.<TriggerStateMachine>asList(triggers));
}
- public static TriggerStateMachine of(
+ public static OnceTriggerStateMachine of(
Iterable<? extends TriggerStateMachine> triggers) {
return new AfterFirstStateMachine(ImmutableList.copyOf(triggers));
}
@@ -78,19 +79,18 @@ public class AfterFirstStateMachine extends TriggerStateMachine {
}
@Override
- public void onFire(TriggerContext context) throws Exception {
- for (ExecutableTriggerStateMachine subTrigger : context.trigger().subTriggers()) {
- TriggerContext subContext = context.forTrigger(subTrigger);
- if (subTrigger.invokeShouldFire(subContext)) {
+ protected void onOnlyFiring(TriggerContext context) throws Exception {
+ for (ExecutableTriggerStateMachine subtrigger : context.trigger().subTriggers()) {
+ TriggerContext subContext = context.forTrigger(subtrigger);
+ if (subtrigger.invokeShouldFire(subContext)) {
// If the trigger is ready to fire, then do whatever it needs to do.
- subTrigger.invokeOnFire(subContext);
+ subtrigger.invokeOnFire(subContext);
} else {
// If the trigger is not ready to fire, it is nonetheless true that whatever
// pending pane it was tracking is now gone.
- subTrigger.invokeClear(subContext);
+ subtrigger.invokeClear(subContext);
}
}
- context.trigger().setFinished(true);
}
@Override
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterPaneStateMachine.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterPaneStateMachine.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterPaneStateMachine.java
index 1ce035a..b9fbac3 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterPaneStateMachine.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterPaneStateMachine.java
@@ -23,6 +23,7 @@ import org.apache.beam.runners.core.StateAccessor;
import org.apache.beam.runners.core.StateMerging;
import org.apache.beam.runners.core.StateTag;
import org.apache.beam.runners.core.StateTags;
+import org.apache.beam.runners.core.triggers.TriggerStateMachine.OnceTriggerStateMachine;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.coders.VarLongCoder;
import org.apache.beam.sdk.state.CombiningState;
@@ -32,7 +33,7 @@ import org.apache.beam.sdk.transforms.Sum;
* {@link TriggerStateMachine}s that fire based on properties of the elements in the current pane.
*/
@Experimental(Experimental.Kind.TRIGGER)
-public class AfterPaneStateMachine extends TriggerStateMachine {
+public class AfterPaneStateMachine extends OnceTriggerStateMachine {
private static final StateTag<CombiningState<Long, long[], Long>>
ELEMENTS_IN_PANE_TAG =
@@ -129,8 +130,7 @@ private static final StateTag<CombiningState<Long, long[], Long>>
}
@Override
- public void onFire(TriggerStateMachine.TriggerContext context) throws Exception {
+ protected void onOnlyFiring(TriggerStateMachine.TriggerContext context) throws Exception {
clear(context);
- context.trigger().setFinished(true);
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterWatermarkStateMachine.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterWatermarkStateMachine.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterWatermarkStateMachine.java
index 509c96b..c9eee15 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterWatermarkStateMachine.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/AfterWatermarkStateMachine.java
@@ -22,6 +22,7 @@ import static com.google.common.base.Preconditions.checkNotNull;
import com.google.common.collect.ImmutableList;
import java.util.Objects;
import javax.annotation.Nullable;
+import org.apache.beam.runners.core.triggers.TriggerStateMachine.OnceTriggerStateMachine;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.state.TimeDomain;
@@ -241,7 +242,7 @@ public class AfterWatermarkStateMachine {
/**
* A watermark trigger targeted relative to the end of the window.
*/
- public static class FromEndOfWindow extends TriggerStateMachine {
+ public static class FromEndOfWindow extends OnceTriggerStateMachine {
private FromEndOfWindow() {
super(null);
@@ -318,8 +319,6 @@ public class AfterWatermarkStateMachine {
}
@Override
- public void onFire(TriggerStateMachine.TriggerContext context) throws Exception {
- context.trigger().setFinished(true);
- }
+ protected void onOnlyFiring(TriggerStateMachine.TriggerContext context) throws Exception { }
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/ExecutableTriggerStateMachine.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/ExecutableTriggerStateMachine.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/ExecutableTriggerStateMachine.java
index cdcff64..c4d89c2 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/ExecutableTriggerStateMachine.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/ExecutableTriggerStateMachine.java
@@ -23,6 +23,7 @@ import static com.google.common.base.Preconditions.checkState;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.List;
+import org.apache.beam.runners.core.triggers.TriggerStateMachine.OnceTriggerStateMachine;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
/**
@@ -45,14 +46,17 @@ public class ExecutableTriggerStateMachine implements Serializable {
private static <W extends BoundedWindow> ExecutableTriggerStateMachine create(
TriggerStateMachine trigger, int nextUnusedIndex) {
-
+ if (trigger instanceof OnceTriggerStateMachine) {
+ return new ExecutableOnceTriggerStateMachine(
+ (OnceTriggerStateMachine) trigger, nextUnusedIndex);
+ } else {
return new ExecutableTriggerStateMachine(trigger, nextUnusedIndex);
-
+ }
}
public static <W extends BoundedWindow> ExecutableTriggerStateMachine createForOnceTrigger(
- TriggerStateMachine trigger, int nextUnusedIndex) {
- return new ExecutableTriggerStateMachine(trigger, nextUnusedIndex);
+ OnceTriggerStateMachine trigger, int nextUnusedIndex) {
+ return new ExecutableOnceTriggerStateMachine(trigger, nextUnusedIndex);
}
private ExecutableTriggerStateMachine(TriggerStateMachine trigger, int nextUnusedIndex) {
@@ -142,4 +146,15 @@ public class ExecutableTriggerStateMachine implements Serializable {
public void invokeClear(TriggerStateMachine.TriggerContext c) throws Exception {
trigger.clear(c.forTrigger(this));
}
+
+ /**
+ * {@link ExecutableTriggerStateMachine} that enforces the fact that the trigger should always
+ * FIRE_AND_FINISH and never just FIRE.
+ */
+ private static class ExecutableOnceTriggerStateMachine extends ExecutableTriggerStateMachine {
+
+ public ExecutableOnceTriggerStateMachine(OnceTriggerStateMachine trigger, int nextUnusedIndex) {
+ super(trigger, nextUnusedIndex);
+ }
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/NeverStateMachine.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/NeverStateMachine.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/NeverStateMachine.java
index f8c5e8b..f32c7a8 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/NeverStateMachine.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/NeverStateMachine.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.runners.core.triggers;
+import org.apache.beam.runners.core.triggers.TriggerStateMachine.OnceTriggerStateMachine;
import org.apache.beam.sdk.transforms.GroupByKey;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
@@ -26,7 +27,7 @@ import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
* <p>Using this trigger will only produce output when the watermark passes the end of the
* {@link BoundedWindow window} plus the allowed lateness.
*/
-public final class NeverStateMachine extends TriggerStateMachine {
+public final class NeverStateMachine extends OnceTriggerStateMachine {
/**
* Returns a trigger which never fires. Output will be produced from the using {@link GroupByKey}
* when the {@link BoundedWindow} closes.
@@ -52,7 +53,7 @@ public final class NeverStateMachine extends TriggerStateMachine {
}
@Override
- public void onFire(TriggerStateMachine.TriggerContext context) {
+ protected void onOnlyFiring(TriggerStateMachine.TriggerContext context) {
throw new UnsupportedOperationException(
String.format("%s should never fire", getClass().getSimpleName()));
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/TriggerStateMachine.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/TriggerStateMachine.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/TriggerStateMachine.java
index 880aa48..6a2cf0c 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/TriggerStateMachine.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/triggers/TriggerStateMachine.java
@@ -453,8 +453,35 @@ public abstract class TriggerStateMachine implements Serializable {
* }
* </pre>
*
+ * <p>Note that if {@code t1} is {@link OnceTriggerStateMachine}, then {@code t1.orFinally(t2)} is
+ * the same as {@code AfterFirst.of(t1, t2)}.
*/
public TriggerStateMachine orFinally(TriggerStateMachine until) {
return new OrFinallyStateMachine(this, until);
}
+
+ /**
+ * {@link TriggerStateMachine}s that are guaranteed to fire at most once should extend from this,
+ * rather than the general {@link TriggerStateMachine} class to indicate that behavior.
+ */
+ public abstract static class OnceTriggerStateMachine extends TriggerStateMachine {
+ protected OnceTriggerStateMachine(List<TriggerStateMachine> subTriggers) {
+ super(subTriggers);
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public final void onFire(TriggerContext context) throws Exception {
+ onOnlyFiring(context);
+ context.trigger().setFinished(true);
+ }
+
+ /**
+ * Called exactly once by {@link #onFire} when the trigger is fired. By default,
+ * invokes {@link #onFire} on all subtriggers for which {@link #shouldFire} is {@code true}.
+ */
+ protected abstract void onOnlyFiring(TriggerContext context) throws Exception;
+ }
}
[15/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/SerializableFunctions.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/SerializableFunctions.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/SerializableFunctions.java
deleted file mode 100644
index d057d81..0000000
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/SerializableFunctions.java
+++ /dev/null
@@ -1,50 +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.sdk.transforms;
-
-/** Useful {@link SerializableFunction} overrides. */
-public class SerializableFunctions {
- private static class Identity<T> implements SerializableFunction<T, T> {
- @Override
- public T apply(T input) {
- return input;
- }
- }
-
- private static class Constant<InT, OutT> implements SerializableFunction<InT, OutT> {
- OutT value;
-
- Constant(OutT value) {
- this.value = value;
- }
-
- @Override
- public OutT apply(InT input) {
- return value;
- }
- }
-
- public static <T> SerializableFunction<T, T> identity() {
- return new Identity<>();
- }
-
- public static <InT, OutT> SerializableFunction<InT, OutT> constant(OutT value) {
- return new Constant<>(value);
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/View.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/View.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/View.java
index 331b143..073c750 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/View.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/View.java
@@ -257,10 +257,8 @@ public class View {
throw new IllegalStateException("Unable to create a side-input view from input", e);
}
- PCollectionView<List<T>> view =
- PCollectionViews.listView(input, input.getWindowingStrategy(), input.getCoder());
- input.apply(CreatePCollectionView.<T, List<T>>of(view));
- return view;
+ return input.apply(CreatePCollectionView.<T, List<T>>of(PCollectionViews.listView(
+ input, input.getWindowingStrategy(), input.getCoder())));
}
}
@@ -284,10 +282,8 @@ public class View {
throw new IllegalStateException("Unable to create a side-input view from input", e);
}
- PCollectionView<Iterable<T>> view =
- PCollectionViews.iterableView(input, input.getWindowingStrategy(), input.getCoder());
- input.apply(CreatePCollectionView.<T, Iterable<T>>of(view));
- return view;
+ return input.apply(CreatePCollectionView.<T, Iterable<T>>of(PCollectionViews.iterableView(
+ input, input.getWindowingStrategy(), input.getCoder())));
}
}
@@ -427,10 +423,11 @@ public class View {
throw new IllegalStateException("Unable to create a side-input view from input", e);
}
- PCollectionView<Map<K, Iterable<V>>> view =
- PCollectionViews.multimapView(input, input.getWindowingStrategy(), input.getCoder());
- input.apply(CreatePCollectionView.<KV<K, V>, Map<K, Iterable<V>>>of(view));
- return view;
+ return input.apply(CreatePCollectionView.<KV<K, V>, Map<K, Iterable<V>>>of(
+ PCollectionViews.multimapView(
+ input,
+ input.getWindowingStrategy(),
+ input.getCoder())));
}
}
@@ -462,10 +459,11 @@ public class View {
throw new IllegalStateException("Unable to create a side-input view from input", e);
}
- PCollectionView<Map<K, V>> view =
- PCollectionViews.mapView(input, input.getWindowingStrategy(), input.getCoder());
- input.apply(CreatePCollectionView.<KV<K, V>, Map<K, V>>of(view));
- return view;
+ return input.apply(CreatePCollectionView.<KV<K, V>, Map<K, V>>of(
+ PCollectionViews.mapView(
+ input,
+ input.getWindowingStrategy(),
+ input.getCoder())));
}
}
@@ -482,7 +480,7 @@ public class View {
*/
@Internal
public static class CreatePCollectionView<ElemT, ViewT>
- extends PTransform<PCollection<ElemT>, PCollection<ElemT>> {
+ extends PTransform<PCollection<ElemT>, PCollectionView<ViewT>> {
private PCollectionView<ViewT> view;
private CreatePCollectionView(PCollectionView<ViewT> view) {
@@ -508,10 +506,8 @@ public class View {
}
@Override
- public PCollection<ElemT> expand(PCollection<ElemT> input) {
- return PCollection.<ElemT>createPrimitiveOutputInternal(
- input.getPipeline(), input.getWindowingStrategy(), input.isBounded())
- .setCoder(input.getCoder());
+ public PCollectionView<ViewT> expand(PCollection<ElemT> input) {
+ return view;
}
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java
index cf96c9b..5d5887a 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyDoFnInvokerFactory.java
@@ -89,7 +89,6 @@ public class ByteBuddyDoFnInvokerFactory implements DoFnInvokerFactory {
public static final String PROCESS_CONTEXT_PARAMETER_METHOD = "processContext";
public static final String ON_TIMER_CONTEXT_PARAMETER_METHOD = "onTimerContext";
public static final String WINDOW_PARAMETER_METHOD = "window";
- public static final String PIPELINE_OPTIONS_PARAMETER_METHOD = "pipelineOptions";
public static final String RESTRICTION_TRACKER_PARAMETER_METHOD = "restrictionTracker";
public static final String STATE_PARAMETER_METHOD = "state";
public static final String TIMER_PARAMETER_METHOD = "timer";
@@ -627,11 +626,6 @@ public class ByteBuddyDoFnInvokerFactory implements DoFnInvokerFactory {
getExtraContextFactoryMethodDescription(TIMER_PARAMETER_METHOD, String.class)),
TypeCasting.to(new TypeDescription.ForLoadedType(Timer.class)));
}
-
- @Override
- public StackManipulation dispatch(DoFnSignature.Parameter.PipelineOptionsParameter p) {
- return simpleExtraContextParameter(PIPELINE_OPTIONS_PARAMETER_METHOD);
- }
});
}
@@ -640,17 +634,6 @@ public class ByteBuddyDoFnInvokerFactory implements DoFnInvokerFactory {
* {@link ProcessElement} method.
*/
private static final class ProcessElementDelegation extends DoFnMethodDelegation {
- private static final MethodDescription PROCESS_CONTINUATION_STOP_METHOD;
-
- static {
- try {
- PROCESS_CONTINUATION_STOP_METHOD =
- new MethodDescription.ForLoadedMethod(DoFn.ProcessContinuation.class.getMethod("stop"));
- } catch (NoSuchMethodException e) {
- throw new RuntimeException("Failed to locate ProcessContinuation.stop()");
- }
- }
-
private final DoFnSignature.ProcessElementMethod signature;
/** Implementation of {@link MethodDelegation} for the {@link ProcessElement} method. */
@@ -684,16 +667,6 @@ public class ByteBuddyDoFnInvokerFactory implements DoFnInvokerFactory {
}
return new StackManipulation.Compound(pushParameters);
}
-
- @Override
- protected StackManipulation afterDelegation(MethodDescription instrumentedMethod) {
- if (TypeDescription.VOID.equals(targetMethod.getReturnType().asErasure())) {
- return new StackManipulation.Compound(
- MethodInvocation.invoke(PROCESS_CONTINUATION_STOP_METHOD), MethodReturn.REFERENCE);
- } else {
- return MethodReturn.of(targetMethod.getReturnType().asErasure());
- }
- }
}
private static class UserCodeMethodInvocation implements StackManipulation {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyOnTimerInvokerFactory.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyOnTimerInvokerFactory.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyOnTimerInvokerFactory.java
index 5e31f2e..e031337 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyOnTimerInvokerFactory.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/ByteBuddyOnTimerInvokerFactory.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.sdk.transforms.reflect;
-
import com.google.common.base.CharMatcher;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
@@ -62,14 +61,13 @@ class ByteBuddyOnTimerInvokerFactory implements OnTimerInvokerFactory {
@SuppressWarnings("unchecked")
Class<? extends DoFn<?, ?>> fnClass = (Class<? extends DoFn<?, ?>>) fn.getClass();
- try {
- OnTimerMethodSpecifier onTimerMethodSpecifier =
- OnTimerMethodSpecifier.forClassAndTimerId(fnClass, timerId);
- Constructor<?> constructor = constructorCache.get(onTimerMethodSpecifier);
- OnTimerInvoker<InputT, OutputT> invoker =
+ try {
+ Constructor<?> constructor = constructorCache.get(fnClass).get(timerId);
+ @SuppressWarnings("unchecked")
+ OnTimerInvoker<InputT, OutputT> invoker =
(OnTimerInvoker<InputT, OutputT>) constructor.newInstance(fn);
- return invoker;
+ return invoker;
} catch (InstantiationException
| IllegalAccessException
| IllegalArgumentException
@@ -99,31 +97,50 @@ class ByteBuddyOnTimerInvokerFactory implements OnTimerInvokerFactory {
private static final String FN_DELEGATE_FIELD_NAME = "delegate";
/**
- * A cache of constructors of generated {@link OnTimerInvoker} classes,
- * keyed by {@link OnTimerMethodSpecifier}.
+ * A cache of constructors of generated {@link OnTimerInvoker} classes, keyed by {@link DoFn}
+ * class and then by {@link TimerId}.
*
* <p>Needed because generating an invoker class is expensive, and to avoid generating an
* excessive number of classes consuming PermGen memory in Java's that still have PermGen.
*/
- private final LoadingCache<OnTimerMethodSpecifier, Constructor<?>> constructorCache =
- CacheBuilder.newBuilder().build(
- new CacheLoader<OnTimerMethodSpecifier, Constructor<?>>() {
- @Override
- public Constructor<?> load(final OnTimerMethodSpecifier onTimerMethodSpecifier)
- throws Exception {
- DoFnSignature signature =
- DoFnSignatures.getSignature(onTimerMethodSpecifier.fnClass());
- Class<? extends OnTimerInvoker<?, ?>> invokerClass =
- generateOnTimerInvokerClass(signature, onTimerMethodSpecifier.timerId());
- try {
- return invokerClass.getConstructor(signature.fnClass());
- } catch (IllegalArgumentException | NoSuchMethodException | SecurityException e) {
- throw new RuntimeException(e);
- }
-
- }
- });
- /**
+ private final LoadingCache<Class<? extends DoFn<?, ?>>, LoadingCache<String, Constructor<?>>>
+ constructorCache =
+ CacheBuilder.newBuilder()
+ .build(
+ new CacheLoader<
+ Class<? extends DoFn<?, ?>>, LoadingCache<String, Constructor<?>>>() {
+ @Override
+ public LoadingCache<String, Constructor<?>> load(
+ final Class<? extends DoFn<?, ?>> fnClass) throws Exception {
+ return CacheBuilder.newBuilder().build(new OnTimerConstructorLoader(fnClass));
+ }
+ });
+
+ /**
+ * A cache loader fixed to a particular {@link DoFn} class that loads constructors for the
+ * invokers for its {@link OnTimer @OnTimer} methods.
+ */
+ private static class OnTimerConstructorLoader extends CacheLoader<String, Constructor<?>> {
+
+ private final DoFnSignature signature;
+
+ public OnTimerConstructorLoader(Class<? extends DoFn<?, ?>> clazz) {
+ this.signature = DoFnSignatures.getSignature(clazz);
+ }
+
+ @Override
+ public Constructor<?> load(String timerId) throws Exception {
+ Class<? extends OnTimerInvoker<?, ?>> invokerClass =
+ generateOnTimerInvokerClass(signature, timerId);
+ try {
+ return invokerClass.getConstructor(signature.fnClass());
+ } catch (IllegalArgumentException | NoSuchMethodException | SecurityException e) {
+ throw new RuntimeException(e);
+ }
+ }
+ }
+
+ /**
* Generates a {@link OnTimerInvoker} class for the given {@link DoFnSignature} and {@link
* TimerId}.
*/
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvoker.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvoker.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvoker.java
index 8b41fee..6fd4052 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvoker.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnInvoker.java
@@ -19,7 +19,6 @@ package org.apache.beam.sdk.transforms.reflect;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderRegistry;
-import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.state.State;
import org.apache.beam.sdk.state.Timer;
import org.apache.beam.sdk.transforms.DoFn;
@@ -54,10 +53,8 @@ public interface DoFnInvoker<InputT, OutputT> {
* Invoke the {@link DoFn.ProcessElement} method on the bound {@link DoFn}.
*
* @param extra Factory for producing extra parameter objects (such as window), if necessary.
- * @return The {@link DoFn.ProcessContinuation} returned by the underlying method, or {@link
- * DoFn.ProcessContinuation#stop()} if it returns {@code void}.
*/
- DoFn.ProcessContinuation invokeProcessElement(ArgumentProvider<InputT, OutputT> extra);
+ void invokeProcessElement(ArgumentProvider<InputT, OutputT> extra);
/** Invoke the appropriate {@link DoFn.OnTimer} method on the bound {@link DoFn}. */
void invokeOnTimer(String timerId, ArgumentProvider<InputT, OutputT> arguments);
@@ -103,12 +100,7 @@ public interface DoFnInvoker<InputT, OutputT> {
*/
BoundedWindow window();
- /** Provide {@link PipelineOptions}. */
- PipelineOptions pipelineOptions();
-
- /**
- * Provide a {@link DoFn.StartBundleContext} to use with the given {@link DoFn}.
- */
+ /** Provide a {@link DoFn.StartBundleContext} to use with the given {@link DoFn}. */
DoFn<InputT, OutputT>.StartBundleContext startBundleContext(DoFn<InputT, OutputT> doFn);
/** Provide a {@link DoFn.FinishBundleContext} to use with the given {@link DoFn}. */
@@ -146,11 +138,6 @@ public interface DoFnInvoker<InputT, OutputT> {
}
@Override
- public PipelineOptions pipelineOptions() {
- return null;
- }
-
- @Override
public DoFn<InputT, OutputT>.StartBundleContext startBundleContext(DoFn<InputT, OutputT> doFn) {
return null;
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignature.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignature.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignature.java
index bfad69e..0b4bf90 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignature.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignature.java
@@ -27,13 +27,11 @@ import java.util.List;
import java.util.Map;
import javax.annotation.Nullable;
import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.state.State;
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.state.Timer;
import org.apache.beam.sdk.state.TimerSpec;
import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.DoFn.ProcessContinuation;
import org.apache.beam.sdk.transforms.DoFn.StateId;
import org.apache.beam.sdk.transforms.DoFn.TimerId;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.RestrictionTrackerParameter;
@@ -195,8 +193,6 @@ public abstract class DoFnSignature {
return cases.dispatch((StateParameter) this);
} else if (this instanceof TimerParameter) {
return cases.dispatch((TimerParameter) this);
- } else if (this instanceof PipelineOptionsParameter) {
- return cases.dispatch((PipelineOptionsParameter) this);
} else {
throw new IllegalStateException(
String.format("Attempt to case match on unknown %s subclass %s",
@@ -216,7 +212,6 @@ public abstract class DoFnSignature {
ResultT dispatch(RestrictionTrackerParameter p);
ResultT dispatch(StateParameter p);
ResultT dispatch(TimerParameter p);
- ResultT dispatch(PipelineOptionsParameter p);
/**
* A base class for a visitor with a default method for cases it is not interested in.
@@ -264,11 +259,6 @@ public abstract class DoFnSignature {
public ResultT dispatch(TimerParameter p) {
return dispatchDefault(p);
}
-
- @Override
- public ResultT dispatch(PipelineOptionsParameter p) {
- return dispatchDefault(p);
- }
}
}
@@ -297,11 +287,6 @@ public abstract class DoFnSignature {
return new AutoValue_DoFnSignature_Parameter_WindowParameter(windowT);
}
- /** Returns a {@link PipelineOptionsParameter}. */
- public static PipelineOptionsParameter pipelineOptions() {
- return new AutoValue_DoFnSignature_Parameter_PipelineOptionsParameter();
- }
-
/**
* Returns a {@link RestrictionTrackerParameter}.
*/
@@ -321,14 +306,6 @@ public abstract class DoFnSignature {
}
/**
- * Descriptor for a {@link Parameter} of a subtype of {@link PipelineOptions}.
- */
- @AutoValue
- public abstract static class PipelineOptionsParameter extends Parameter {
- PipelineOptionsParameter() {}
- }
-
- /**
* Descriptor for a {@link Parameter} of type {@link DoFn.StartBundleContext}.
*
* <p>All such descriptors are equal.
@@ -337,7 +314,6 @@ public abstract class DoFnSignature {
public abstract static class StartBundleContextParameter extends Parameter {
StartBundleContextParameter() {}
}
-
/**
* Descriptor for a {@link Parameter} of type {@link DoFn.FinishBundleContext}.
*
@@ -434,21 +410,16 @@ public abstract class DoFnSignature {
@Nullable
public abstract TypeDescriptor<? extends BoundedWindow> windowT();
- /** Whether this {@link DoFn} returns a {@link ProcessContinuation} or void. */
- public abstract boolean hasReturnValue();
-
static ProcessElementMethod create(
Method targetMethod,
List<Parameter> extraParameters,
TypeDescriptor<?> trackerT,
- @Nullable TypeDescriptor<? extends BoundedWindow> windowT,
- boolean hasReturnValue) {
+ @Nullable TypeDescriptor<? extends BoundedWindow> windowT) {
return new AutoValue_DoFnSignature_ProcessElementMethod(
targetMethod,
Collections.unmodifiableList(extraParameters),
trackerT,
- windowT,
- hasReturnValue);
+ windowT);
}
/**
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java
index de57c3b..bb191b1 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java
@@ -17,8 +17,6 @@
*/
package org.apache.beam.sdk.transforms.reflect;
-import static com.google.common.base.Preconditions.checkState;
-
import com.google.auto.value.AutoValue;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Predicates;
@@ -44,7 +42,6 @@ import java.util.List;
import java.util.Map;
import javax.annotation.Nullable;
import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.state.State;
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.state.Timer;
@@ -81,23 +78,19 @@ public class DoFnSignatures {
ImmutableList.of(
Parameter.ProcessContextParameter.class,
Parameter.WindowParameter.class,
- Parameter.PipelineOptionsParameter.class,
Parameter.TimerParameter.class,
Parameter.StateParameter.class);
private static final Collection<Class<? extends Parameter>>
ALLOWED_SPLITTABLE_PROCESS_ELEMENT_PARAMETERS =
ImmutableList.of(
- Parameter.PipelineOptionsParameter.class,
- Parameter.ProcessContextParameter.class,
- Parameter.RestrictionTrackerParameter.class);
+ Parameter.ProcessContextParameter.class, Parameter.RestrictionTrackerParameter.class);
private static final Collection<Class<? extends Parameter>>
ALLOWED_ON_TIMER_PARAMETERS =
ImmutableList.of(
Parameter.OnTimerContextParameter.class,
Parameter.WindowParameter.class,
- Parameter.PipelineOptionsParameter.class,
Parameter.TimerParameter.class,
Parameter.StateParameter.class);
@@ -194,15 +187,6 @@ public class DoFnSignatures {
extraParameters, Predicates.instanceOf(WindowParameter.class));
}
- /**
- * Indicates whether a {@link Parameter.PipelineOptionsParameter} is
- * known in this context.
- */
- public boolean hasPipelineOptionsParamter() {
- return Iterables.any(
- extraParameters, Predicates.instanceOf(Parameter.PipelineOptionsParameter.class));
- }
-
/** The window type, if any, used by this method. */
@Nullable
public TypeDescriptor<? extends BoundedWindow> getWindowType() {
@@ -442,8 +426,6 @@ public class DoFnSignatures {
* <li>If the {@link DoFn} (or any of its supertypes) is annotated as {@link
* DoFn.BoundedPerElement} or {@link DoFn.UnboundedPerElement}, use that. Only one of
* these must be specified.
- * <li>If {@link DoFn.ProcessElement} returns {@link DoFn.ProcessContinuation}, assume it is
- * unbounded. Otherwise (if it returns {@code void}), assume it is bounded.
* <li>If {@link DoFn.ProcessElement} returns {@code void}, but the {@link DoFn} is annotated
* {@link DoFn.UnboundedPerElement}, this is an error.
* </ol>
@@ -469,10 +451,7 @@ public class DoFnSignatures {
}
if (processElement.isSplittable()) {
if (isBounded == null) {
- isBounded =
- processElement.hasReturnValue()
- ? PCollection.IsBounded.UNBOUNDED
- : PCollection.IsBounded.BOUNDED;
+ isBounded = PCollection.IsBounded.BOUNDED;
}
} else {
errors.checkArgument(
@@ -481,7 +460,6 @@ public class DoFnSignatures {
+ ((isBounded == PCollection.IsBounded.BOUNDED)
? DoFn.BoundedPerElement.class.getSimpleName()
: DoFn.UnboundedPerElement.class.getSimpleName()));
- checkState(!processElement.hasReturnValue(), "Should have been inferred splittable");
isBounded = PCollection.IsBounded.BOUNDED;
}
return isBounded;
@@ -718,10 +696,8 @@ public class DoFnSignatures {
TypeDescriptor<?> outputT,
FnAnalysisContext fnContext) {
errors.checkArgument(
- void.class.equals(m.getReturnType())
- || DoFn.ProcessContinuation.class.equals(m.getReturnType()),
- "Must return void or %s",
- DoFn.ProcessContinuation.class.getSimpleName());
+ void.class.equals(m.getReturnType()),
+ "Must return void");
MethodAnalysisContext methodContext = MethodAnalysisContext.create();
@@ -761,11 +737,7 @@ public class DoFnSignatures {
}
return DoFnSignature.ProcessElementMethod.create(
- m,
- methodContext.getExtraParameters(),
- trackerT,
- windowT,
- DoFn.ProcessContinuation.class.equals(m.getReturnType()));
+ m, methodContext.getExtraParameters(), trackerT, windowT);
}
private static void checkParameterOneOf(
@@ -817,12 +789,6 @@ public class DoFnSignatures {
"Multiple %s parameters",
BoundedWindow.class.getSimpleName());
return Parameter.boundedWindow((TypeDescriptor<? extends BoundedWindow>) paramT);
- } else if (PipelineOptions.class.equals(rawType)) {
- methodErrors.checkArgument(
- !methodContext.hasPipelineOptionsParamter(),
- "Multiple %s parameters",
- PipelineOptions.class.getSimpleName());
- return Parameter.pipelineOptions();
} else if (RestrictionTracker.class.isAssignableFrom(rawType)) {
methodErrors.checkArgument(
!methodContext.hasRestrictionTrackerParameter(),
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/OnTimerMethodSpecifier.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/OnTimerMethodSpecifier.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/OnTimerMethodSpecifier.java
deleted file mode 100644
index edf7e3c..0000000
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/OnTimerMethodSpecifier.java
+++ /dev/null
@@ -1,37 +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.sdk.transforms.reflect;
-
-import com.google.auto.value.AutoValue;
-import org.apache.beam.sdk.transforms.DoFn;
-
-/**
- * Used by {@link ByteBuddyOnTimerInvokerFactory} to Dynamically generate
- * {@link OnTimerInvoker} instances for invoking a particular
- * {@link DoFn.TimerId} on a particular {@link DoFn}.
- */
-
-@AutoValue
-abstract class OnTimerMethodSpecifier {
- public abstract Class<? extends DoFn<?, ?>> fnClass();
- public abstract String timerId();
- public static OnTimerMethodSpecifier
- forClassAndTimerId(Class<? extends DoFn<?, ?>> fnClass, String timerId){
- return new AutoValue_OnTimerMethodSpecifier(fnClass, timerId);
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/OffsetRange.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/OffsetRange.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/OffsetRange.java
new file mode 100644
index 0000000..104f5f2
--- /dev/null
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/OffsetRange.java
@@ -0,0 +1,77 @@
+/*
+ * 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.sdk.transforms.splittabledofn;
+
+import static com.google.common.base.Preconditions.checkArgument;
+
+import java.io.Serializable;
+
+/** A restriction represented by a range of integers [from, to). */
+public class OffsetRange
+ implements Serializable, HasDefaultTracker<OffsetRange, OffsetRangeTracker> {
+ private final long from;
+ private final long to;
+
+ public OffsetRange(long from, long to) {
+ checkArgument(from <= to, "Malformed range [%s, %s)", from, to);
+ this.from = from;
+ this.to = to;
+ }
+
+ public long getFrom() {
+ return from;
+ }
+
+ public long getTo() {
+ return to;
+ }
+
+ @Override
+ public OffsetRangeTracker newTracker() {
+ return new OffsetRangeTracker(this);
+ }
+
+ @Override
+ public String toString() {
+ return "[" + from + ", " + to + ')';
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+
+ OffsetRange that = (OffsetRange) o;
+
+ if (from != that.from) {
+ return false;
+ }
+ return to == that.to;
+ }
+
+ @Override
+ public int hashCode() {
+ int result = (int) (from ^ (from >>> 32));
+ result = 31 * result + (int) (to ^ (to >>> 32));
+ return result;
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/OffsetRangeTracker.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/OffsetRangeTracker.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/OffsetRangeTracker.java
index 4987409..0271a0d 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/OffsetRangeTracker.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/OffsetRangeTracker.java
@@ -21,8 +21,6 @@ import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;
-import com.google.common.base.MoreObjects;
-import org.apache.beam.sdk.io.range.OffsetRange;
import org.apache.beam.sdk.transforms.DoFn;
/**
@@ -101,13 +99,4 @@ public class OffsetRangeTracker implements RestrictionTracker<OffsetRange> {
lastAttemptedOffset + 1,
range.getTo());
}
-
- @Override
- public String toString() {
- return MoreObjects.toStringHelper(this)
- .add("range", range)
- .add("lastClaimedOffset", lastClaimedOffset)
- .add("lastAttemptedOffset", lastAttemptedOffset)
- .toString();
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/RestrictionTracker.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/RestrictionTracker.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/RestrictionTracker.java
index 8cb0a6b..27ef68f 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/RestrictionTracker.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/RestrictionTracker.java
@@ -31,13 +31,10 @@ public interface RestrictionTracker<RestrictionT> {
RestrictionT currentRestriction();
/**
- * Signals that the current {@link DoFn.ProcessElement} call should terminate as soon as possible:
- * after this method returns, the tracker MUST refuse all future claim calls, and {@link
- * #checkDone} MUST succeed.
- *
- * <p>Modifies {@link #currentRestriction}. Returns a restriction representing the rest of the
- * work: the old value of {@link #currentRestriction} is equivalent to the new value and the
- * return value of this method combined. Must be called at most once on a given object.
+ * Signals that the current {@link DoFn.ProcessElement} call should terminate as soon as possible.
+ * Modifies {@link #currentRestriction}. Returns a restriction representing the rest of the work:
+ * the old value of {@link #currentRestriction} is equivalent to the new value and the return
+ * value of this method combined. Must be called at most once on a given object.
*/
RestrictionT checkpoint();
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/GlobalWindows.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/GlobalWindows.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/GlobalWindows.java
index c68c497..d48d26b 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/GlobalWindows.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/GlobalWindows.java
@@ -79,11 +79,6 @@ public class GlobalWindows extends NonMergingWindowFn<Object, GlobalWindow> {
}
@Override
- public boolean assignsToOneWindow() {
- return true;
- }
-
- @Override
public boolean equals(Object other) {
return other instanceof GlobalWindows;
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PartitioningWindowFn.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PartitioningWindowFn.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PartitioningWindowFn.java
index 341ba27..40ee68a 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PartitioningWindowFn.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PartitioningWindowFn.java
@@ -58,9 +58,4 @@ public abstract class PartitioningWindowFn<T, W extends BoundedWindow>
public Instant getOutputTime(Instant inputTimestamp, W window) {
return inputTimestamp;
}
-
- @Override
- public final boolean assignsToOneWindow() {
- return true;
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/SlidingWindows.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/SlidingWindows.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/SlidingWindows.java
index 150b956..f657884 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/SlidingWindows.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/SlidingWindows.java
@@ -148,11 +148,6 @@ public class SlidingWindows extends NonMergingWindowFn<Object, IntervalWindow> {
}
@Override
- public boolean assignsToOneWindow() {
- return !this.period.isShorterThan(this.size);
- }
-
- @Override
public void verifyCompatibility(WindowFn<?, ?> other) throws IncompatibleWindowException {
if (!this.isCompatible(other)) {
throw new IncompatibleWindowException(
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/Window.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/Window.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/Window.java
index a12be6d..105ebfb 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/Window.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/Window.java
@@ -163,24 +163,6 @@ public abstract class Window<T> extends PTransform<PCollection<T>, PCollection<T
}
/**
- * Specifies the conditions under which an on-time pane will be created when a window is closed.
- */
- public enum OnTimeBehavior {
- /**
- * Always fire the on-time pane. Even if there is no new data since the previous firing,
- * an element will be produced.
- *
- * <p>This is the default behavior.
- */
- FIRE_ALWAYS,
- /**
- * Only fire the on-time pane if there is new data since the previous firing.
- */
- FIRE_IF_NON_EMPTY
-
- }
-
- /**
* Creates a {@code Window} {@code PTransform} that uses the given
* {@link WindowFn} to window the data.
*
@@ -213,7 +195,6 @@ public abstract class Window<T> extends PTransform<PCollection<T>, PCollection<T
@Nullable abstract AccumulationMode getAccumulationMode();
@Nullable abstract Duration getAllowedLateness();
@Nullable abstract ClosingBehavior getClosingBehavior();
- @Nullable abstract OnTimeBehavior getOnTimeBehavior();
@Nullable abstract TimestampCombiner getTimestampCombiner();
abstract Builder<T> toBuilder();
@@ -225,7 +206,6 @@ public abstract class Window<T> extends PTransform<PCollection<T>, PCollection<T
abstract Builder<T> setAccumulationMode(AccumulationMode mode);
abstract Builder<T> setAllowedLateness(Duration allowedLateness);
abstract Builder<T> setClosingBehavior(ClosingBehavior closingBehavior);
- abstract Builder<T> setOnTimeBehavior(OnTimeBehavior onTimeBehavior);
abstract Builder<T> setTimestampCombiner(TimestampCombiner timestampCombiner);
abstract Window<T> build();
@@ -319,15 +299,6 @@ public abstract class Window<T> extends PTransform<PCollection<T>, PCollection<T
}
/**
- * <b><i>(Experimental)</i></b> Override the default {@link OnTimeBehavior}, to control
- * whether to output an empty on-time pane.
- */
- @Experimental(Kind.TRIGGER)
- public Window<T> withOnTimeBehavior(OnTimeBehavior behavior) {
- return toBuilder().setOnTimeBehavior(behavior).build();
- }
-
- /**
* Get the output strategy of this {@link Window Window PTransform}. For internal use
* only.
*/
@@ -350,9 +321,6 @@ public abstract class Window<T> extends PTransform<PCollection<T>, PCollection<T
if (getClosingBehavior() != null) {
result = result.withClosingBehavior(getClosingBehavior());
}
- if (getOnTimeBehavior() != null) {
- result = result.withOnTimeBehavior(getOnTimeBehavior());
- }
if (getTimestampCombiner() != null) {
result = result.withTimestampCombiner(getTimestampCombiner());
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/WindowFn.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/WindowFn.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/WindowFn.java
index ffe85f3..001d630 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/WindowFn.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/WindowFn.java
@@ -180,17 +180,6 @@ public abstract class WindowFn<T, W extends BoundedWindow>
}
/**
- * Returns true if this {@link WindowFn} always assigns an element to exactly one window.
- *
- * <p>If this varies per-element, or cannot be determined, conservatively return false.
- *
- * <p>By default, returns false.
- */
- public boolean assignsToOneWindow() {
- return false;
- }
-
- /**
* Returns a {@link TypeDescriptor} capturing what is known statically about the window type of
* this {@link WindowFn} instance's most-derived class.
*
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/IdentityWindowFn.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/IdentityWindowFn.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/IdentityWindowFn.java
index ef6d833..a4bfdda 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/IdentityWindowFn.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/IdentityWindowFn.java
@@ -116,9 +116,4 @@ public class IdentityWindowFn<T> extends NonMergingWindowFn<T, BoundedWindow> {
public Instant getOutputTime(Instant inputTimestamp, BoundedWindow window) {
return inputTimestamp;
}
-
- @Override
- public boolean assignsToOneWindow() {
- return true;
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollection.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollection.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollection.java
index 4063d11..f210fd8 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollection.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollection.java
@@ -20,8 +20,6 @@ package org.apache.beam.sdk.values;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
-import java.util.Collections;
-import java.util.Map;
import javax.annotation.Nullable;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.annotations.Internal;
@@ -228,11 +226,6 @@ public class PCollection<T> extends PValueBase implements PValue {
return super.getName();
}
- @Override
- public final Map<TupleTag<?>, PValue> expand() {
- return Collections.<TupleTag<?>, PValue>singletonMap(tag, this);
- }
-
/**
* Sets the name of this {@link PCollection}. Returns {@code this}.
*
@@ -321,11 +314,6 @@ public class PCollection<T> extends PValueBase implements PValue {
private IsBounded isBounded;
- /**
- * A local {@link TupleTag} used in the expansion of this {@link PValueBase}.
- */
- private final TupleTag<?> tag = new TupleTag<>();
-
private PCollection(Pipeline p) {
super(p);
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionViews.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionViews.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionViews.java
index e17e146..74887c7 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionViews.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionViews.java
@@ -21,7 +21,6 @@ import com.google.common.base.Function;
import com.google.common.base.MoreObjects;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import com.google.common.collect.Multimap;
import java.io.IOException;
@@ -39,7 +38,6 @@ import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.IterableCoder;
import org.apache.beam.sdk.transforms.Materialization;
import org.apache.beam.sdk.transforms.Materializations;
-import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ViewFn;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.InvalidWindows;
@@ -141,18 +139,6 @@ public class PCollectionViews {
}
/**
- * Expands a list of {@link PCollectionView} into the form needed for
- * {@link PTransform#getAdditionalInputs()}.
- */
- public static Map<TupleTag<?>, PValue> toAdditionalInputs(Iterable<PCollectionView<?>> views) {
- ImmutableMap.Builder<TupleTag<?>, PValue> additionalInputs = ImmutableMap.builder();
- for (PCollectionView<?> view : views) {
- additionalInputs.put(view.getTagInternal(), view.getPCollection());
- }
- return additionalInputs.build();
- }
-
- /**
* Implementation of conversion of singleton {@code Iterable<WindowedValue<T>>} to {@code T}.
*
* <p>For internal use only.
@@ -184,15 +170,6 @@ public class PCollectionViews {
}
/**
- * Returns if a default value was specified.
- */
- @Deprecated
- @Internal
- public boolean hasDefault() {
- return hasDefault;
- }
-
- /**
* Returns the default value that was specified.
*
* <p>For internal use only.
@@ -296,16 +273,6 @@ public class PCollectionViews {
}
}));
}
-
- @Override
- public boolean equals(Object other) {
- return other instanceof ListViewFn;
- }
-
- @Override
- public int hashCode() {
- return ListViewFn.class.hashCode();
- }
}
/**
@@ -524,10 +491,5 @@ public class PCollectionViews {
public String toString() {
return MoreObjects.toStringHelper(this).add("tag", tag).toString();
}
-
- @Override
- public Map<TupleTag<?>, PValue> expand() {
- return Collections.<TupleTag<?>, PValue>singletonMap(tag, pCollection);
- }
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PValueBase.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PValueBase.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PValueBase.java
index f312eac..6f638d7 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PValueBase.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PValueBase.java
@@ -19,6 +19,8 @@ package org.apache.beam.sdk.values;
import static com.google.common.base.Preconditions.checkState;
+import java.util.Collections;
+import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.annotations.Internal;
import org.apache.beam.sdk.transforms.PTransform;
@@ -85,6 +87,11 @@ public abstract class PValueBase implements PValue {
private String name;
/**
+ * A local {@link TupleTag} used in the expansion of this {@link PValueBase}.
+ */
+ private TupleTag<?> tag = new TupleTag<>();
+
+ /**
* Whether this {@link PValueBase} has been finalized, and its core
* properties, e.g., name, can no longer be changed.
*/
@@ -101,6 +108,11 @@ public abstract class PValueBase implements PValue {
}
@Override
+ public final Map<TupleTag<?>, PValue> expand() {
+ return Collections.<TupleTag<?>, PValue>singletonMap(tag, this);
+ }
+
+ @Override
public void finishSpecifying(PInput input, PTransform<?, ?> transform) {
finishedSpecifying = true;
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ShardedKey.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ShardedKey.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ShardedKey.java
deleted file mode 100644
index e56af13..0000000
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ShardedKey.java
+++ /dev/null
@@ -1,65 +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.sdk.values;
-
-import java.io.Serializable;
-import java.util.Objects;
-
-/** A key and a shard number. */
-public class ShardedKey<K> implements Serializable {
- private static final long serialVersionUID = 1L;
- private final K key;
- private final int shardNumber;
-
- public static <K> ShardedKey<K> of(K key, int shardNumber) {
- return new ShardedKey<>(key, shardNumber);
- }
-
- private ShardedKey(K key, int shardNumber) {
- this.key = key;
- this.shardNumber = shardNumber;
- }
-
- public K getKey() {
- return key;
- }
-
- public int getShardNumber() {
- return shardNumber;
- }
-
- @Override
- public String toString() {
- return "key: " + key + " shard: " + shardNumber;
- }
-
- @Override
- public boolean equals(Object o) {
- if (!(o instanceof ShardedKey)) {
- return false;
- }
- ShardedKey<K> other = (ShardedKey<K>) o;
- return Objects.equals(key, other.key) && Objects.equals(shardNumber, other.shardNumber);
- }
-
- @Override
- public int hashCode() {
- return Objects.hash(key, shardNumber);
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowingStrategy.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowingStrategy.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowingStrategy.java
index 3b74e69..8a773e2 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowingStrategy.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowingStrategy.java
@@ -29,7 +29,6 @@ import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
import org.apache.beam.sdk.transforms.windowing.Trigger;
import org.apache.beam.sdk.transforms.windowing.Window.ClosingBehavior;
-import org.apache.beam.sdk.transforms.windowing.Window.OnTimeBehavior;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.joda.time.Duration;
@@ -60,7 +59,6 @@ public class WindowingStrategy<T, W extends BoundedWindow> implements Serializab
private final AccumulationMode mode;
private final Duration allowedLateness;
private final ClosingBehavior closingBehavior;
- private final OnTimeBehavior onTimeBehavior;
private final TimestampCombiner timestampCombiner;
private final boolean triggerSpecified;
private final boolean modeSpecified;
@@ -73,8 +71,7 @@ public class WindowingStrategy<T, W extends BoundedWindow> implements Serializab
AccumulationMode mode, boolean modeSpecified,
Duration allowedLateness, boolean allowedLatenessSpecified,
TimestampCombiner timestampCombiner, boolean timestampCombinerSpecified,
- ClosingBehavior closingBehavior,
- OnTimeBehavior onTimeBehavior) {
+ ClosingBehavior closingBehavior) {
this.windowFn = windowFn;
this.trigger = trigger;
this.triggerSpecified = triggerSpecified;
@@ -83,7 +80,6 @@ public class WindowingStrategy<T, W extends BoundedWindow> implements Serializab
this.allowedLateness = allowedLateness;
this.allowedLatenessSpecified = allowedLatenessSpecified;
this.closingBehavior = closingBehavior;
- this.onTimeBehavior = onTimeBehavior;
this.timestampCombiner = timestampCombiner;
this.timestampCombinerSpecified = timestampCombinerSpecified;
}
@@ -102,8 +98,7 @@ public class WindowingStrategy<T, W extends BoundedWindow> implements Serializab
AccumulationMode.DISCARDING_FIRED_PANES, false,
DEFAULT_ALLOWED_LATENESS, false,
TimestampCombiner.END_OF_WINDOW, false,
- ClosingBehavior.FIRE_IF_NON_EMPTY,
- OnTimeBehavior.FIRE_ALWAYS);
+ ClosingBehavior.FIRE_IF_NON_EMPTY);
}
public WindowFn<T, W> getWindowFn() {
@@ -138,10 +133,6 @@ public class WindowingStrategy<T, W extends BoundedWindow> implements Serializab
return closingBehavior;
}
- public OnTimeBehavior getOnTimeBehavior() {
- return onTimeBehavior;
- }
-
public TimestampCombiner getTimestampCombiner() {
return timestampCombiner;
}
@@ -161,8 +152,7 @@ public class WindowingStrategy<T, W extends BoundedWindow> implements Serializab
mode, modeSpecified,
allowedLateness, allowedLatenessSpecified,
timestampCombiner, timestampCombinerSpecified,
- closingBehavior,
- onTimeBehavior);
+ closingBehavior);
}
/**
@@ -176,8 +166,7 @@ public class WindowingStrategy<T, W extends BoundedWindow> implements Serializab
mode, true,
allowedLateness, allowedLatenessSpecified,
timestampCombiner, timestampCombinerSpecified,
- closingBehavior,
- onTimeBehavior);
+ closingBehavior);
}
/**
@@ -194,8 +183,7 @@ public class WindowingStrategy<T, W extends BoundedWindow> implements Serializab
mode, modeSpecified,
allowedLateness, allowedLatenessSpecified,
timestampCombiner, timestampCombinerSpecified,
- closingBehavior,
- onTimeBehavior);
+ closingBehavior);
}
/**
@@ -209,8 +197,7 @@ public class WindowingStrategy<T, W extends BoundedWindow> implements Serializab
mode, modeSpecified,
allowedLateness, true,
timestampCombiner, timestampCombinerSpecified,
- closingBehavior,
- onTimeBehavior);
+ closingBehavior);
}
public WindowingStrategy<T, W> withClosingBehavior(ClosingBehavior closingBehavior) {
@@ -220,19 +207,7 @@ public class WindowingStrategy<T, W extends BoundedWindow> implements Serializab
mode, modeSpecified,
allowedLateness, allowedLatenessSpecified,
timestampCombiner, timestampCombinerSpecified,
- closingBehavior,
- onTimeBehavior);
- }
-
- public WindowingStrategy<T, W> withOnTimeBehavior(OnTimeBehavior onTimeBehavior) {
- return new WindowingStrategy<T, W>(
- windowFn,
- trigger, triggerSpecified,
- mode, modeSpecified,
- allowedLateness, allowedLatenessSpecified,
- timestampCombiner, timestampCombinerSpecified,
- closingBehavior,
- onTimeBehavior);
+ closingBehavior);
}
@Experimental(Experimental.Kind.OUTPUT_TIME)
@@ -244,8 +219,7 @@ public class WindowingStrategy<T, W extends BoundedWindow> implements Serializab
mode, modeSpecified,
allowedLateness, allowedLatenessSpecified,
timestampCombiner, true,
- closingBehavior,
- onTimeBehavior);
+ closingBehavior);
}
@Override
@@ -272,7 +246,6 @@ public class WindowingStrategy<T, W extends BoundedWindow> implements Serializab
&& getMode().equals(other.getMode())
&& getAllowedLateness().equals(other.getAllowedLateness())
&& getClosingBehavior().equals(other.getClosingBehavior())
- && getOnTimeBehavior().equals(other.getOnTimeBehavior())
&& getTrigger().equals(other.getTrigger())
&& getTimestampCombiner().equals(other.getTimestampCombiner())
&& getWindowFn().equals(other.getWindowFn());
@@ -305,7 +278,6 @@ public class WindowingStrategy<T, W extends BoundedWindow> implements Serializab
mode, true,
allowedLateness, true,
timestampCombiner, true,
- closingBehavior,
- onTimeBehavior);
+ closingBehavior);
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/AvroIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/AvroIOTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/AvroIOTest.java
index 260e47a..6d01d32 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/AvroIOTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/AvroIOTest.java
@@ -54,11 +54,10 @@ import org.apache.avro.reflect.ReflectData;
import org.apache.avro.reflect.ReflectDatumReader;
import org.apache.beam.sdk.coders.AvroCoder;
import org.apache.beam.sdk.coders.DefaultCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.io.FileBasedSink.FilenamePolicy;
-import org.apache.beam.sdk.io.FileBasedSink.OutputFileHints;
import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions;
import org.apache.beam.sdk.io.fs.ResourceId;
-import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
import org.apache.beam.sdk.testing.NeedsRunner;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
@@ -277,42 +276,37 @@ public class AvroIOTest {
}
private static class WindowedFilenamePolicy extends FilenamePolicy {
- final ResourceId outputFilePrefix;
+ final String outputFilePrefix;
- WindowedFilenamePolicy(ResourceId outputFilePrefix) {
+ WindowedFilenamePolicy(String outputFilePrefix) {
this.outputFilePrefix = outputFilePrefix;
}
@Override
- public ResourceId windowedFilename(WindowedContext input, OutputFileHints outputFileHints) {
- String filenamePrefix =
- outputFilePrefix.isDirectory() ? "" : firstNonNull(outputFilePrefix.getFilename(), "");
-
- String filename =
- String.format(
- "%s-%s-%s-of-%s-pane-%s%s%s",
- filenamePrefix,
- input.getWindow(),
- input.getShardNumber(),
- input.getNumShards() - 1,
- input.getPaneInfo().getIndex(),
- input.getPaneInfo().isLast() ? "-final" : "",
- outputFileHints.getSuggestedFilenameSuffix());
- return outputFilePrefix
- .getCurrentDirectory()
- .resolve(filename, StandardResolveOptions.RESOLVE_FILE);
+ public ResourceId windowedFilename(
+ ResourceId outputDirectory, WindowedContext input, String extension) {
+ String filename = String.format(
+ "%s-%s-%s-of-%s-pane-%s%s%s",
+ outputFilePrefix,
+ input.getWindow(),
+ input.getShardNumber(),
+ input.getNumShards() - 1,
+ input.getPaneInfo().getIndex(),
+ input.getPaneInfo().isLast() ? "-final" : "",
+ extension);
+ return outputDirectory.resolve(filename, StandardResolveOptions.RESOLVE_FILE);
}
@Override
- public ResourceId unwindowedFilename(Context input, OutputFileHints outputFileHints) {
+ public ResourceId unwindowedFilename(
+ ResourceId outputDirectory, Context input, String extension) {
throw new UnsupportedOperationException("Expecting windowed outputs only");
}
@Override
public void populateDisplayData(DisplayData.Builder builder) {
- builder.add(
- DisplayData.item("fileNamePrefix", outputFilePrefix.toString())
- .withLabel("File Name Prefix"));
+ builder.add(DisplayData.item("fileNamePrefix", outputFilePrefix)
+ .withLabel("File Name Prefix"));
}
}
@@ -365,18 +359,15 @@ public class AvroIOTest {
Arrays.copyOfRange(secondWindowArray, 1, secondWindowArray.length))
.advanceWatermarkToInfinity();
- FilenamePolicy policy =
- new WindowedFilenamePolicy(FileBasedSink.convertToFileResourceIfPossible(baseFilename));
+ FilenamePolicy policy = new WindowedFilenamePolicy(baseFilename);
windowedAvroWritePipeline
.apply(values)
.apply(Window.<GenericClass>into(FixedWindows.of(Duration.standardMinutes(1))))
- .apply(
- AvroIO.write(GenericClass.class)
- .to(policy)
- .withTempDirectory(
- StaticValueProvider.of(FileSystems.matchNewResource(baseDir.toString(), true)))
- .withWindowedWrites()
- .withNumShards(2));
+ .apply(AvroIO.write(GenericClass.class)
+ .to(baseFilename)
+ .withFilenamePolicy(policy)
+ .withWindowedWrites()
+ .withNumShards(2));
windowedAvroWritePipeline.run();
// Validate that the data written matches the expected elements in the expected order
@@ -503,14 +494,13 @@ public class AvroIOTest {
expectedFiles.add(
new File(
DefaultFilenamePolicy.constructName(
- FileBasedSink.convertToFileResourceIfPossible(outputFilePrefix),
- shardNameTemplate,
- "" /* no suffix */,
- i,
- numShards,
- null,
- null)
- .toString()));
+ outputFilePrefix,
+ shardNameTemplate,
+ "" /* no suffix */,
+ i,
+ numShards,
+ null,
+ null)));
}
List<String> actualElements = new ArrayList<>();
@@ -582,4 +572,15 @@ public class AvroIOTest {
assertThat(displayData, hasDisplayItem("numShards", 100));
assertThat(displayData, hasDisplayItem("codec", CodecFactory.snappyCodec().toString()));
}
+
+ @Test
+ public void testWindowedWriteRequiresFilenamePolicy() {
+ PCollection<String> emptyInput = p.apply(Create.empty(StringUtf8Coder.of()));
+ AvroIO.Write write = AvroIO.write(String.class).to("/tmp/some/file").withWindowedWrites();
+
+ expectedException.expect(IllegalStateException.class);
+ expectedException.expectMessage(
+ "When using windowed writes, a filename policy must be set via withFilenamePolicy()");
+ emptyInput.apply(write);
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/DefaultFilenamePolicyTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/DefaultFilenamePolicyTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/DefaultFilenamePolicyTest.java
index 9dc6d33..217420c 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/DefaultFilenamePolicyTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/DefaultFilenamePolicyTest.java
@@ -17,9 +17,9 @@
*/
package org.apache.beam.sdk.io;
+import static org.apache.beam.sdk.io.DefaultFilenamePolicy.constructName;
import static org.junit.Assert.assertEquals;
-import org.apache.beam.sdk.io.fs.ResourceId;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
@@ -30,108 +30,69 @@ import org.junit.runners.JUnit4;
@RunWith(JUnit4.class)
public class DefaultFilenamePolicyTest {
- private static String constructName(
- String baseFilename,
- String shardTemplate,
- String suffix,
- int shardNum,
- int numShards,
- String paneStr,
- String windowStr) {
- ResourceId constructed =
- DefaultFilenamePolicy.constructName(
- FileSystems.matchNewResource(baseFilename, false),
- shardTemplate,
- suffix,
- shardNum,
- numShards,
- paneStr,
- windowStr);
- return constructed.toString();
- }
-
@Test
public void testConstructName() {
- assertEquals(
- "/path/to/output-001-of-123.txt",
- constructName("/path/to/output", "-SSS-of-NNN", ".txt", 1, 123, null, null));
+ assertEquals("output-001-of-123.txt",
+ constructName("output", "-SSS-of-NNN", ".txt", 1, 123, null, null));
- assertEquals(
- "/path/to/out.txt/part-00042",
- constructName("/path/to/out.txt", "/part-SSSSS", "", 42, 100, null, null));
+ assertEquals("out.txt/part-00042",
+ constructName("out.txt", "/part-SSSSS", "", 42, 100, null, null));
- assertEquals("/path/to/out.txt", constructName("/path/to/ou", "t.t", "xt", 1, 1, null, null));
+ assertEquals("out.txt",
+ constructName("ou", "t.t", "xt", 1, 1, null, null));
- assertEquals(
- "/path/to/out0102shard.txt",
- constructName("/path/to/out", "SSNNshard", ".txt", 1, 2, null, null));
+ assertEquals("out0102shard.txt",
+ constructName("out", "SSNNshard", ".txt", 1, 2, null, null));
- assertEquals(
- "/path/to/out-2/1.part-1-of-2.txt",
- constructName("/path/to/out", "-N/S.part-S-of-N", ".txt", 1, 2, null, null));
+ assertEquals("out-2/1.part-1-of-2.txt",
+ constructName("out", "-N/S.part-S-of-N", ".txt", 1, 2, null, null));
}
@Test
public void testConstructNameWithLargeShardCount() {
- assertEquals(
- "/out-100-of-5000.txt", constructName("/out", "-SS-of-NN", ".txt", 100, 5000, null, null));
+ assertEquals("out-100-of-5000.txt",
+ constructName("out", "-SS-of-NN", ".txt", 100, 5000, null, null));
}
@Test
public void testConstructWindowedName() {
- assertEquals(
- "/path/to/output-001-of-123.txt",
- constructName("/path/to/output", "-SSS-of-NNN", ".txt", 1, 123, null, null));
-
- assertEquals(
- "/path/to/output-001-of-123-PPP-W.txt",
- constructName("/path/to/output", "-SSS-of-NNN-PPP-W", ".txt", 1, 123, null, null));
-
- assertEquals(
- "/path/to/out" + ".txt/part-00042-myPaneStr-myWindowStr",
- constructName(
- "/path/to/out.txt", "/part-SSSSS-P-W", "", 42, 100, "myPaneStr", "myWindowStr"));
-
- assertEquals(
- "/path/to/out.txt",
- constructName("/path/to/ou", "t.t", "xt", 1, 1, "myPaneStr2", "anotherWindowStr"));
-
- assertEquals(
- "/path/to/out0102shard-oneMoreWindowStr-anotherPaneStr.txt",
- constructName(
- "/path/to/out", "SSNNshard-W-P", ".txt", 1, 2, "anotherPaneStr", "oneMoreWindowStr"));
-
- assertEquals(
- "/out-2/1.part-1-of-2-slidingWindow1-myPaneStr3-windowslidingWindow1-"
- + "panemyPaneStr3.txt",
- constructName(
- "/out",
- "-N/S.part-S-of-N-W-P-windowW-paneP",
- ".txt",
- 1,
- 2,
- "myPaneStr3",
- "slidingWindow1"));
+ assertEquals("output-001-of-123.txt",
+ constructName("output", "-SSS-of-NNN", ".txt", 1, 123, null, null));
+
+ assertEquals("output-001-of-123-PPP-W.txt",
+ constructName("output", "-SSS-of-NNN-PPP-W", ".txt", 1, 123, null, null));
+
+ assertEquals("out.txt/part-00042-myPaneStr-myWindowStr",
+ constructName("out.txt", "/part-SSSSS-P-W", "", 42, 100, "myPaneStr",
+ "myWindowStr"));
+
+ assertEquals("out.txt", constructName("ou", "t.t", "xt", 1, 1, "myPaneStr2",
+ "anotherWindowStr"));
+
+ assertEquals("out0102shard-oneMoreWindowStr-anotherPaneStr.txt",
+ constructName("out", "SSNNshard-W-P", ".txt", 1, 2, "anotherPaneStr",
+ "oneMoreWindowStr"));
+
+ assertEquals("out-2/1.part-1-of-2-slidingWindow1-myPaneStr3-windowslidingWindow1-"
+ + "panemyPaneStr3.txt",
+ constructName("out", "-N/S.part-S-of-N-W-P-windowW-paneP", ".txt", 1, 2, "myPaneStr3",
+ "slidingWindow1"));
// test first/last pane
- assertEquals(
- "/out.txt/part-00042-myWindowStr-pane-11-true-false",
- constructName(
- "/out.txt", "/part-SSSSS-W-P", "", 42, 100, "pane-11-true-false", "myWindowStr"));
-
- assertEquals(
- "/path/to/out.txt",
- constructName("/path/to/ou", "t.t", "xt", 1, 1, "pane", "anotherWindowStr"));
-
- assertEquals(
- "/out0102shard-oneMoreWindowStr-pane--1-false-false-pane--1-false-false.txt",
- constructName(
- "/out", "SSNNshard-W-P-P", ".txt", 1, 2, "pane--1-false-false", "oneMoreWindowStr"));
-
- assertEquals(
- "/path/to/out-2/1.part-1-of-2-sWindow1-winsWindow1-ppaneL.txt",
- constructName(
- "/path/to/out", "-N/S.part-S-of-N-W-winW-pP", ".txt", 1, 2, "paneL", "sWindow1"));
+ assertEquals("out.txt/part-00042-myWindowStr-pane-11-true-false",
+ constructName("out.txt", "/part-SSSSS-W-P", "", 42, 100, "pane-11-true-false",
+ "myWindowStr"));
+
+ assertEquals("out.txt", constructName("ou", "t.t", "xt", 1, 1, "pane",
+ "anotherWindowStr"));
+
+ assertEquals("out0102shard-oneMoreWindowStr-pane--1-false-false-pane--1-false-false.txt",
+ constructName("out", "SSNNshard-W-P-P", ".txt", 1, 2, "pane--1-false-false",
+ "oneMoreWindowStr"));
+
+ assertEquals("out-2/1.part-1-of-2-sWindow1-winsWindow1-ppaneL.txt",
+ constructName("out",
+ "-N/S.part-S-of-N-W-winW-pP", ".txt", 1, 2, "paneL", "sWindow1"));
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/DrunkWritableByteChannelFactory.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/DrunkWritableByteChannelFactory.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/DrunkWritableByteChannelFactory.java
index a7644b6..6615a2e 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/DrunkWritableByteChannelFactory.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/DrunkWritableByteChannelFactory.java
@@ -39,7 +39,7 @@ public class DrunkWritableByteChannelFactory implements WritableByteChannelFacto
}
@Override
- public String getSuggestedFilenameSuffix() {
+ public String getFilenameSuffix() {
return ".drunk";
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileBasedSinkTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileBasedSinkTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileBasedSinkTest.java
index 755bb59..caad759 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileBasedSinkTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileBasedSinkTest.java
@@ -103,7 +103,7 @@ public class FileBasedSinkTest {
SimpleSink.SimpleWriter writer =
buildWriteOperationWithTempDir(getBaseTempDirectory()).createWriter();
- writer.openUnwindowed(testUid, -1, null);
+ writer.openUnwindowed(testUid, -1);
for (String value : values) {
writer.write(value);
}
@@ -198,27 +198,23 @@ public class FileBasedSinkTest {
throws Exception {
int numFiles = temporaryFiles.size();
- List<FileResult<Void>> fileResults = new ArrayList<>();
+ List<FileResult> fileResults = new ArrayList<>();
// Create temporary output bundles and output File objects.
for (int i = 0; i < numFiles; i++) {
fileResults.add(
- new FileResult<Void>(
+ new FileResult(
LocalResources.fromFile(temporaryFiles.get(i), false),
WriteFiles.UNKNOWN_SHARDNUM,
null,
- null,
null));
}
writeOp.finalize(fileResults);
+ ResourceId outputDirectory = writeOp.getSink().getBaseOutputDirectoryProvider().get();
for (int i = 0; i < numFiles; i++) {
- ResourceId outputFilename =
- writeOp
- .getSink()
- .getDynamicDestinations()
- .getFilenamePolicy(null)
- .unwindowedFilename(new Context(i, numFiles), CompressionType.UNCOMPRESSED);
+ ResourceId outputFilename = writeOp.getSink().getFilenamePolicy()
+ .unwindowedFilename(outputDirectory, new Context(i, numFiles), "");
assertTrue(new File(outputFilename.toString()).exists());
assertFalse(temporaryFiles.get(i).exists());
}
@@ -235,12 +231,11 @@ public class FileBasedSinkTest {
private void testRemoveTemporaryFiles(int numFiles, ResourceId tempDirectory)
throws Exception {
String prefix = "file";
- SimpleSink<Void> sink =
- SimpleSink.makeSimpleSink(
- getBaseOutputDirectory(), prefix, "", "", CompressionType.UNCOMPRESSED);
+ SimpleSink sink =
+ new SimpleSink(getBaseOutputDirectory(), prefix, "", "");
- WriteOperation<String, Void> writeOp =
- new SimpleSink.SimpleWriteOperation<>(sink, tempDirectory);
+ WriteOperation<String> writeOp =
+ new SimpleSink.SimpleWriteOperation(sink, tempDirectory);
List<File> temporaryFiles = new ArrayList<>();
List<File> outputFiles = new ArrayList<>();
@@ -277,6 +272,8 @@ public class FileBasedSinkTest {
@Test
public void testCopyToOutputFiles() throws Exception {
SimpleSink.SimpleWriteOperation writeOp = buildWriteOperation();
+ ResourceId outputDirectory = writeOp.getSink().getBaseOutputDirectoryProvider().get();
+
List<String> inputFilenames = Arrays.asList("input-1", "input-2", "input-3");
List<String> inputContents = Arrays.asList("1", "2", "3");
List<String> expectedOutputFilenames = Arrays.asList(
@@ -295,14 +292,9 @@ public class FileBasedSinkTest {
File inputTmpFile = tmpFolder.newFile(inputFilenames.get(i));
List<String> lines = Collections.singletonList(inputContents.get(i));
writeFile(lines, inputTmpFile);
- inputFilePaths.put(
- LocalResources.fromFile(inputTmpFile, false),
- writeOp
- .getSink()
- .getDynamicDestinations()
- .getFilenamePolicy(null)
- .unwindowedFilename(
- new Context(i, inputFilenames.size()), CompressionType.UNCOMPRESSED));
+ inputFilePaths.put(LocalResources.fromFile(inputTmpFile, false),
+ writeOp.getSink().getFilenamePolicy()
+ .unwindowedFilename(outputDirectory, new Context(i, inputFilenames.size()), ""));
}
// Copy input files to output files.
@@ -319,8 +311,7 @@ public class FileBasedSinkTest {
ResourceId outputDirectory, FilenamePolicy policy, int numFiles) {
List<ResourceId> filenames = new ArrayList<>();
for (int i = 0; i < numFiles; i++) {
- filenames.add(
- policy.unwindowedFilename(new Context(i, numFiles), CompressionType.UNCOMPRESSED));
+ filenames.add(policy.unwindowedFilename(outputDirectory, new Context(i, numFiles), ""));
}
return filenames;
}
@@ -335,10 +326,8 @@ public class FileBasedSinkTest {
List<ResourceId> actual;
ResourceId root = getBaseOutputDirectory();
- SimpleSink<Void> sink =
- SimpleSink.makeSimpleSink(
- root, "file", ".SSSSS.of.NNNNN", ".test", CompressionType.UNCOMPRESSED);
- FilenamePolicy policy = sink.getDynamicDestinations().getFilenamePolicy(null);
+ SimpleSink sink = new SimpleSink(root, "file", ".SSSSS.of.NNNNN", ".test");
+ FilenamePolicy policy = sink.getFilenamePolicy();
expected = Arrays.asList(
root.resolve("file.00000.of.00003.test", StandardResolveOptions.RESOLVE_FILE),
@@ -363,9 +352,8 @@ public class FileBasedSinkTest {
@Test
public void testCollidingOutputFilenames() throws IOException {
ResourceId root = getBaseOutputDirectory();
- SimpleSink<Void> sink =
- SimpleSink.makeSimpleSink(root, "file", "-NN", "test", CompressionType.UNCOMPRESSED);
- SimpleSink.SimpleWriteOperation<Void> writeOp = new SimpleSink.SimpleWriteOperation<>(sink);
+ SimpleSink sink = new SimpleSink(root, "file", "-NN", "test");
+ SimpleSink.SimpleWriteOperation writeOp = new SimpleSink.SimpleWriteOperation(sink);
ResourceId temp1 = root.resolve("temp1", StandardResolveOptions.RESOLVE_FILE);
ResourceId temp2 = root.resolve("temp2", StandardResolveOptions.RESOLVE_FILE);
@@ -373,11 +361,11 @@ public class FileBasedSinkTest {
ResourceId output = root.resolve("file-03.test", StandardResolveOptions.RESOLVE_FILE);
// More than one shard does.
try {
- Iterable<FileResult<Void>> results =
+ Iterable<FileResult> results =
Lists.newArrayList(
- new FileResult<Void>(temp1, 1, null, null, null),
- new FileResult<Void>(temp2, 1, null, null, null),
- new FileResult<Void>(temp3, 1, null, null, null));
+ new FileResult(temp1, 1, null, null),
+ new FileResult(temp2, 1, null, null),
+ new FileResult(temp3, 1, null, null));
writeOp.buildOutputFilenames(results);
fail("Should have failed.");
} catch (IllegalStateException exn) {
@@ -391,10 +379,8 @@ public class FileBasedSinkTest {
List<ResourceId> expected;
List<ResourceId> actual;
ResourceId root = getBaseOutputDirectory();
- SimpleSink<Void> sink =
- SimpleSink.makeSimpleSink(
- root, "file", "-SSSSS-of-NNNNN", "", CompressionType.UNCOMPRESSED);
- FilenamePolicy policy = sink.getDynamicDestinations().getFilenamePolicy(null);
+ SimpleSink sink = new SimpleSink(root, "file", "-SSSSS-of-NNNNN", "");
+ FilenamePolicy policy = sink.getFilenamePolicy();
expected = Arrays.asList(
root.resolve("file-00000-of-00003", StandardResolveOptions.RESOLVE_FILE),
@@ -500,11 +486,10 @@ public class FileBasedSinkTest {
public void testFileBasedWriterWithWritableByteChannelFactory() throws Exception {
final String testUid = "testId";
ResourceId root = getBaseOutputDirectory();
- WriteOperation<String, Void> writeOp =
- SimpleSink.makeSimpleSink(
- root, "file", "-SS-of-NN", "txt", new DrunkWritableByteChannelFactory())
+ WriteOperation<String> writeOp =
+ new SimpleSink(root, "file", "-SS-of-NN", "txt", new DrunkWritableByteChannelFactory())
.createWriteOperation();
- final Writer<String, Void> writer = writeOp.createWriter();
+ final Writer<String> writer = writeOp.createWriter();
final ResourceId expectedFile =
writeOp.tempDirectory.get().resolve(testUid, StandardResolveOptions.RESOLVE_FILE);
@@ -518,7 +503,7 @@ public class FileBasedSinkTest {
expected.add("footer");
expected.add("footer");
- writer.openUnwindowed(testUid, -1, null);
+ writer.openUnwindowed(testUid, -1);
writer.write("a");
writer.write("b");
final FileResult result = writer.close();
@@ -528,20 +513,20 @@ public class FileBasedSinkTest {
}
/** Build a SimpleSink with default options. */
- private SimpleSink<Void> buildSink() {
- return SimpleSink.makeSimpleSink(
- getBaseOutputDirectory(), "file", "-SS-of-NN", ".test", CompressionType.UNCOMPRESSED);
+ private SimpleSink buildSink() {
+ return new SimpleSink(getBaseOutputDirectory(), "file", "-SS-of-NN", ".test");
}
- /** Build a SimpleWriteOperation with default options and the given temporary directory. */
- private SimpleSink.SimpleWriteOperation<Void> buildWriteOperationWithTempDir(
- ResourceId tempDirectory) {
- SimpleSink<Void> sink = buildSink();
- return new SimpleSink.SimpleWriteOperation<>(sink, tempDirectory);
+ /**
+ * Build a SimpleWriteOperation with default options and the given temporary directory.
+ */
+ private SimpleSink.SimpleWriteOperation buildWriteOperationWithTempDir(ResourceId tempDirectory) {
+ SimpleSink sink = buildSink();
+ return new SimpleSink.SimpleWriteOperation(sink, tempDirectory);
}
/** Build a write operation with the default options for it and its parent sink. */
- private SimpleSink.SimpleWriteOperation<Void> buildWriteOperation() {
+ private SimpleSink.SimpleWriteOperation buildWriteOperation() {
return buildSink().createWriteOperation();
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/SimpleSink.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/SimpleSink.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/SimpleSink.java
index 9196178..c97313d 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/SimpleSink.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/SimpleSink.java
@@ -19,55 +19,33 @@ package org.apache.beam.sdk.io;
import java.nio.ByteBuffer;
import java.nio.channels.WritableByteChannel;
-import org.apache.beam.sdk.io.DefaultFilenamePolicy.Params;
-import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions;
import org.apache.beam.sdk.io.fs.ResourceId;
import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
import org.apache.beam.sdk.util.MimeTypes;
/**
- * A simple {@link FileBasedSink} that writes {@link String} values as lines with header and footer.
+ * A simple {@link FileBasedSink} that writes {@link String} values as lines with
+ * header and footer.
*/
-class SimpleSink<DestinationT> extends FileBasedSink<String, DestinationT> {
- public SimpleSink(
- ResourceId tempDirectory,
- DynamicDestinations<String, DestinationT> dynamicDestinations,
- WritableByteChannelFactory writableByteChannelFactory) {
- super(StaticValueProvider.of(tempDirectory), dynamicDestinations, writableByteChannelFactory);
+class SimpleSink extends FileBasedSink<String> {
+ public SimpleSink(ResourceId baseOutputDirectory, String prefix, String template, String suffix) {
+ this(baseOutputDirectory, prefix, template, suffix, CompressionType.UNCOMPRESSED);
}
- public static SimpleSink<Void> makeSimpleSink(
- ResourceId tempDirectory, FilenamePolicy filenamePolicy) {
- return new SimpleSink<>(
- tempDirectory,
- DynamicFileDestinations.<String>constant(filenamePolicy),
- CompressionType.UNCOMPRESSED);
- }
-
- public static SimpleSink<Void> makeSimpleSink(
- ResourceId baseDirectory,
- String prefix,
- String shardTemplate,
- String suffix,
- WritableByteChannelFactory writableByteChannelFactory) {
- DynamicDestinations<String, Void> dynamicDestinations =
- DynamicFileDestinations.constant(
- DefaultFilenamePolicy.fromParams(
- new Params()
- .withBaseFilename(
- baseDirectory.resolve(prefix, StandardResolveOptions.RESOLVE_FILE))
- .withShardTemplate(shardTemplate)
- .withSuffix(suffix)));
- return new SimpleSink<>(baseDirectory, dynamicDestinations, writableByteChannelFactory);
+ public SimpleSink(ResourceId baseOutputDirectory, String prefix, String template, String suffix,
+ WritableByteChannelFactory writableByteChannelFactory) {
+ super(
+ StaticValueProvider.of(baseOutputDirectory),
+ new DefaultFilenamePolicy(StaticValueProvider.of(prefix), template, suffix),
+ writableByteChannelFactory);
}
@Override
- public SimpleWriteOperation<DestinationT> createWriteOperation() {
- return new SimpleWriteOperation<>(this);
+ public SimpleWriteOperation createWriteOperation() {
+ return new SimpleWriteOperation(this);
}
- static final class SimpleWriteOperation<DestinationT>
- extends WriteOperation<String, DestinationT> {
+ static final class SimpleWriteOperation extends WriteOperation<String> {
public SimpleWriteOperation(SimpleSink sink, ResourceId tempOutputDirectory) {
super(sink, tempOutputDirectory);
}
@@ -77,12 +55,12 @@ class SimpleSink<DestinationT> extends FileBasedSink<String, DestinationT> {
}
@Override
- public SimpleWriter<DestinationT> createWriter() throws Exception {
- return new SimpleWriter<>(this);
+ public SimpleWriter createWriter() throws Exception {
+ return new SimpleWriter(this);
}
}
- static final class SimpleWriter<DestinationT> extends Writer<String, DestinationT> {
+ static final class SimpleWriter extends Writer<String> {
static final String HEADER = "header";
static final String FOOTER = "footer";
[09/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/GcpApiSurfaceTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/GcpApiSurfaceTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/GcpApiSurfaceTest.java
index 8aac417..91caded 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/GcpApiSurfaceTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/GcpApiSurfaceTest.java
@@ -52,7 +52,6 @@ public class GcpApiSurfaceTest {
@SuppressWarnings("unchecked")
final Set<Matcher<Class<?>>> allowedClasses =
ImmutableSet.of(
- classesInPackage("com.google.api.core"),
classesInPackage("com.google.api.client.googleapis"),
classesInPackage("com.google.api.client.http"),
classesInPackage("com.google.api.client.json"),
@@ -61,18 +60,9 @@ public class GcpApiSurfaceTest {
classesInPackage("com.google.auth"),
classesInPackage("com.google.bigtable.v2"),
classesInPackage("com.google.cloud.bigtable.config"),
- classesInPackage("com.google.spanner.v1"),
- Matchers.<Class<?>>equalTo(com.google.api.gax.grpc.ApiException.class),
Matchers.<Class<?>>equalTo(com.google.cloud.bigtable.grpc.BigtableClusterName.class),
Matchers.<Class<?>>equalTo(com.google.cloud.bigtable.grpc.BigtableInstanceName.class),
Matchers.<Class<?>>equalTo(com.google.cloud.bigtable.grpc.BigtableTableName.class),
- Matchers.<Class<?>>equalTo(com.google.cloud.BaseServiceException.class),
- Matchers.<Class<?>>equalTo(com.google.cloud.BaseServiceException.Error.class),
- Matchers.<Class<?>>equalTo(com.google.cloud.BaseServiceException.ExceptionData.class),
- Matchers.<Class<?>>equalTo(com.google.cloud.BaseServiceException.ExceptionData.Builder
- .class),
- Matchers.<Class<?>>equalTo(com.google.cloud.RetryHelper.RetryHelperException.class),
- Matchers.<Class<?>>equalTo(com.google.cloud.grpc.BaseGrpcServiceException.class),
Matchers.<Class<?>>equalTo(com.google.cloud.ByteArray.class),
Matchers.<Class<?>>equalTo(com.google.cloud.Date.class),
Matchers.<Class<?>>equalTo(com.google.cloud.Timestamp.class),
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTest.java
index d31f3a0..bfd260a 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTest.java
@@ -82,7 +82,6 @@ import org.apache.beam.sdk.coders.Coder.Context;
import org.apache.beam.sdk.coders.CoderException;
import org.apache.beam.sdk.coders.IterableCoder;
import org.apache.beam.sdk.coders.KvCoder;
-import org.apache.beam.sdk.coders.ShardedKeyCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.io.BoundedSource;
@@ -132,7 +131,6 @@ import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollection.IsBounded;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PCollectionViews;
-import org.apache.beam.sdk.values.ShardedKey;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.sdk.values.ValueInSingleWindow;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableReadIT.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableReadIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableReadIT.java
index 91f0bae..a064bd6 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableReadIT.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableReadIT.java
@@ -20,7 +20,6 @@ package org.apache.beam.sdk.io.gcp.bigtable;
import com.google.bigtable.v2.Row;
import com.google.cloud.bigtable.config.BigtableOptions;
import org.apache.beam.sdk.Pipeline;
-import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
@@ -42,10 +41,8 @@ public class BigtableReadIT {
BigtableTestOptions options = TestPipeline.testingPipelineOptions()
.as(BigtableTestOptions.class);
- String project = options.as(GcpOptions.class).getProject();
-
BigtableOptions.Builder bigtableOptionsBuilder = new BigtableOptions.Builder()
- .setProjectId(project)
+ .setProjectId(options.getProjectId())
.setInstanceId(options.getInstanceId());
final String tableId = "BigtableReadTest";
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableTestOptions.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableTestOptions.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableTestOptions.java
index 03cb697..0ab7576 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableTestOptions.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableTestOptions.java
@@ -25,6 +25,11 @@ import org.apache.beam.sdk.testing.TestPipelineOptions;
* Properties needed when using Bigtable with the Beam SDK.
*/
public interface BigtableTestOptions extends TestPipelineOptions {
+ @Description("Project ID for Bigtable")
+ @Default.String("apache-beam-testing")
+ String getProjectId();
+ void setProjectId(String value);
+
@Description("Instance ID for Bigtable")
@Default.String("beam-test")
String getInstanceId();
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteIT.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteIT.java
index 010bcc4..1d168f1 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteIT.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableWriteIT.java
@@ -73,17 +73,15 @@ public class BigtableWriteIT implements Serializable {
private static BigtableTableAdminClient tableAdminClient;
private final String tableId =
String.format("BigtableWriteIT-%tF-%<tH-%<tM-%<tS-%<tL", new Date());
- private String project;
@Before
public void setup() throws Exception {
PipelineOptionsFactory.register(BigtableTestOptions.class);
options = TestPipeline.testingPipelineOptions().as(BigtableTestOptions.class);
- project = options.as(GcpOptions.class).getProject();
bigtableOptions =
new Builder()
- .setProjectId(project)
+ .setProjectId(options.getProjectId())
.setInstanceId(options.getInstanceId())
.setUserAgent("apache-beam-test")
.build();
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottlerTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottlerTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottlerTest.java
deleted file mode 100644
index c12cf55..0000000
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottlerTest.java
+++ /dev/null
@@ -1,111 +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.sdk.io.gcp.datastore;
-
-import static org.hamcrest.Matchers.closeTo;
-import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.greaterThan;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertThat;
-
-import java.util.Random;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-import org.mockito.Mockito;
-
-/**
- * Tests for {@link AdaptiveThrottler}.
- */
-@RunWith(JUnit4.class)
-public class AdaptiveThrottlerTest {
-
- static final long START_TIME_MS = 0;
- static final long SAMPLE_PERIOD_MS = 60000;
- static final long SAMPLE_BUCKET_MS = 1000;
- static final double OVERLOAD_RATIO = 2;
-
- /** Returns a throttler configured with the standard parameters above. */
- AdaptiveThrottler getThrottler() {
- return new AdaptiveThrottler(SAMPLE_PERIOD_MS, SAMPLE_BUCKET_MS, OVERLOAD_RATIO);
- }
-
- @Test
- public void testNoInitialThrottling() throws Exception {
- AdaptiveThrottler throttler = getThrottler();
- assertThat(throttler.throttlingProbability(START_TIME_MS), equalTo(0.0));
- assertThat("first request is not throttled",
- throttler.throttleRequest(START_TIME_MS), equalTo(false));
- }
-
- @Test
- public void testNoThrottlingIfNoErrors() throws Exception {
- AdaptiveThrottler throttler = getThrottler();
- long t = START_TIME_MS;
- for (; t < START_TIME_MS + 20; t++) {
- assertFalse(throttler.throttleRequest(t));
- throttler.successfulRequest(t);
- }
- assertThat(throttler.throttlingProbability(t), equalTo(0.0));
- }
-
- @Test
- public void testNoThrottlingAfterErrorsExpire() throws Exception {
- AdaptiveThrottler throttler = getThrottler();
- long t = START_TIME_MS;
- for (; t < START_TIME_MS + SAMPLE_PERIOD_MS; t++) {
- throttler.throttleRequest(t);
- // and no successfulRequest.
- }
- assertThat("check that we set up a non-zero probability of throttling",
- throttler.throttlingProbability(t), greaterThan(0.0));
- for (; t < START_TIME_MS + 2 * SAMPLE_PERIOD_MS; t++) {
- throttler.throttleRequest(t);
- throttler.successfulRequest(t);
- }
- assertThat(throttler.throttlingProbability(t), equalTo(0.0));
- }
-
- @Test
- public void testThrottlingAfterErrors() throws Exception {
- Random mockRandom = Mockito.mock(Random.class);
- Mockito.when(mockRandom.nextDouble()).thenReturn(
- 0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9,
- 0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9);
- AdaptiveThrottler throttler = new AdaptiveThrottler(
- SAMPLE_PERIOD_MS, SAMPLE_BUCKET_MS, OVERLOAD_RATIO, mockRandom);
- for (int i = 0; i < 20; i++) {
- boolean throttled = throttler.throttleRequest(START_TIME_MS + i);
- // 1/3rd of requests succeeding.
- if (i % 3 == 1) {
- throttler.successfulRequest(START_TIME_MS + i);
- }
-
- // Once we have some history in place, check what throttling happens.
- if (i >= 10) {
- // Expect 1/3rd of requests to be throttled. (So 1/3rd throttled, 1/3rd succeeding, 1/3rd
- // tried and failing).
- assertThat(String.format("for i=%d", i),
- throttler.throttlingProbability(START_TIME_MS + i), closeTo(0.33, /*error=*/ 0.1));
- // Requests 10..13 should be throttled, 14..19 not throttled given the mocked random numbers
- // that we fed to throttler.
- assertThat(String.format("for i=%d", i), throttled, equalTo(i < 14));
- }
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java
index a3f5d38..460049e 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java
@@ -27,7 +27,7 @@ import static com.google.datastore.v1.client.DatastoreHelper.makeOrder;
import static com.google.datastore.v1.client.DatastoreHelper.makeUpsert;
import static com.google.datastore.v1.client.DatastoreHelper.makeValue;
import static org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DATASTORE_BATCH_UPDATE_BYTES_LIMIT;
-import static org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DATASTORE_BATCH_UPDATE_ENTITIES_START;
+import static org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DATASTORE_BATCH_UPDATE_LIMIT;
import static org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read.DEFAULT_BUNDLE_SIZE_BYTES;
import static org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read.QUERY_BATCH_LIMIT;
import static org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read.getEstimatedSizeBytes;
@@ -606,7 +606,7 @@ public class DatastoreV1Test {
/** Tests {@link DatastoreWriterFn} with entities of more than one batches, but not a multiple. */
@Test
public void testDatatoreWriterFnWithMultipleBatches() throws Exception {
- datastoreWriterFnTest(DATASTORE_BATCH_UPDATE_ENTITIES_START * 3 + 100);
+ datastoreWriterFnTest(DATASTORE_BATCH_UPDATE_LIMIT * 3 + 100);
}
/**
@@ -615,7 +615,7 @@ public class DatastoreV1Test {
*/
@Test
public void testDatatoreWriterFnWithBatchesExactMultiple() throws Exception {
- datastoreWriterFnTest(DATASTORE_BATCH_UPDATE_ENTITIES_START * 2);
+ datastoreWriterFnTest(DATASTORE_BATCH_UPDATE_LIMIT * 2);
}
// A helper method to test DatastoreWriterFn for various batch sizes.
@@ -628,14 +628,14 @@ public class DatastoreV1Test {
}
DatastoreWriterFn datastoreWriter = new DatastoreWriterFn(StaticValueProvider.of(PROJECT_ID),
- null, mockDatastoreFactory, new FakeWriteBatcher());
+ null, mockDatastoreFactory);
DoFnTester<Mutation, Void> doFnTester = DoFnTester.of(datastoreWriter);
doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
doFnTester.processBundle(mutations);
int start = 0;
while (start < numMutations) {
- int end = Math.min(numMutations, start + DATASTORE_BATCH_UPDATE_ENTITIES_START);
+ int end = Math.min(numMutations, start + DATASTORE_BATCH_UPDATE_LIMIT);
CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
commitRequest.addAllMutations(mutations.subList(start, end));
@@ -651,28 +651,26 @@ public class DatastoreV1Test {
@Test
public void testDatatoreWriterFnWithLargeEntities() throws Exception {
List<Mutation> mutations = new ArrayList<>();
- int entitySize = 0;
+ int propertySize = 900_000;
for (int i = 0; i < 12; ++i) {
- Entity entity = Entity.newBuilder().setKey(makeKey("key" + i, i + 1))
- .putProperties("long", makeValue(new String(new char[900_000])
- ).setExcludeFromIndexes(true).build())
- .build();
- entitySize = entity.getSerializedSize(); // Take the size of any one entity.
- mutations.add(makeUpsert(entity).build());
+ Entity.Builder entity = Entity.newBuilder().setKey(makeKey("key" + i, i + 1));
+ entity.putProperties("long", makeValue(new String(new char[propertySize])
+ ).setExcludeFromIndexes(true).build());
+ mutations.add(makeUpsert(entity.build()).build());
}
DatastoreWriterFn datastoreWriter = new DatastoreWriterFn(StaticValueProvider.of(PROJECT_ID),
- null, mockDatastoreFactory, new FakeWriteBatcher());
+ null, mockDatastoreFactory);
DoFnTester<Mutation, Void> doFnTester = DoFnTester.of(datastoreWriter);
doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
doFnTester.processBundle(mutations);
// This test is over-specific currently; it requires that we split the 12 entity writes into 3
// requests, but we only need each CommitRequest to be less than 10MB in size.
- int entitiesPerRpc = DATASTORE_BATCH_UPDATE_BYTES_LIMIT / entitySize;
+ int propertiesPerRpc = DATASTORE_BATCH_UPDATE_BYTES_LIMIT / propertySize;
int start = 0;
while (start < mutations.size()) {
- int end = Math.min(mutations.size(), start + entitiesPerRpc);
+ int end = Math.min(mutations.size(), start + propertiesPerRpc);
CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
commitRequest.addAllMutations(mutations.subList(start, end));
@@ -783,7 +781,7 @@ public class DatastoreV1Test {
*/
@Test
public void testSplitQueryFnWithQueryLimit() throws Exception {
- Query queryWithLimit = QUERY.toBuilder()
+ Query queryWithLimit = QUERY.toBuilder().clone()
.setLimit(Int32Value.newBuilder().setValue(1))
.build();
@@ -896,50 +894,6 @@ public class DatastoreV1Test {
.apply(DatastoreIO.v1().write().withProjectId(options.getDatastoreProject()));
}
- @Test
- public void testWriteBatcherWithoutData() {
- DatastoreV1.WriteBatcher writeBatcher = new DatastoreV1.WriteBatcherImpl();
- writeBatcher.start();
- assertEquals(DatastoreV1.DATASTORE_BATCH_UPDATE_ENTITIES_START, writeBatcher.nextBatchSize(0));
- }
-
- @Test
- public void testWriteBatcherFastQueries() {
- DatastoreV1.WriteBatcher writeBatcher = new DatastoreV1.WriteBatcherImpl();
- writeBatcher.start();
- writeBatcher.addRequestLatency(0, 1000, 200);
- writeBatcher.addRequestLatency(0, 1000, 200);
- assertEquals(DatastoreV1.DATASTORE_BATCH_UPDATE_ENTITIES_LIMIT, writeBatcher.nextBatchSize(0));
- }
-
- @Test
- public void testWriteBatcherSlowQueries() {
- DatastoreV1.WriteBatcher writeBatcher = new DatastoreV1.WriteBatcherImpl();
- writeBatcher.start();
- writeBatcher.addRequestLatency(0, 10000, 200);
- writeBatcher.addRequestLatency(0, 10000, 200);
- assertEquals(100, writeBatcher.nextBatchSize(0));
- }
-
- @Test
- public void testWriteBatcherSizeNotBelowMinimum() {
- DatastoreV1.WriteBatcher writeBatcher = new DatastoreV1.WriteBatcherImpl();
- writeBatcher.start();
- writeBatcher.addRequestLatency(0, 30000, 50);
- writeBatcher.addRequestLatency(0, 30000, 50);
- assertEquals(DatastoreV1.DATASTORE_BATCH_UPDATE_ENTITIES_MIN, writeBatcher.nextBatchSize(0));
- }
-
- @Test
- public void testWriteBatcherSlidingWindow() {
- DatastoreV1.WriteBatcher writeBatcher = new DatastoreV1.WriteBatcherImpl();
- writeBatcher.start();
- writeBatcher.addRequestLatency(0, 30000, 50);
- writeBatcher.addRequestLatency(50000, 5000, 200);
- writeBatcher.addRequestLatency(100000, 5000, 200);
- assertEquals(200, writeBatcher.nextBatchSize(150000));
- }
-
/** Helper Methods */
/** A helper function that verifies if all the queries have unique keys. */
@@ -1079,24 +1033,8 @@ public class DatastoreV1Test {
private List<Query> splitQuery(Query query, int numSplits) {
List<Query> queries = new LinkedList<>();
for (int i = 0; i < numSplits; i++) {
- queries.add(query.toBuilder().build());
+ queries.add(query.toBuilder().clone().build());
}
return queries;
}
-
- /**
- * A WriteBatcher for unit tests, which does no timing-based adjustments (so unit tests have
- * consistent results).
- */
- static class FakeWriteBatcher implements DatastoreV1.WriteBatcher {
- @Override
- public void start() {}
- @Override
- public void addRequestLatency(long timeSinceEpochMillis, long latencyMillis, int numMutations) {
- }
- @Override
- public int nextBatchSize(long timeSinceEpochMillis) {
- return DatastoreV1.DATASTORE_BATCH_UPDATE_ENTITIES_START;
- }
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/V1TestUtil.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/V1TestUtil.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/V1TestUtil.java
index cd61229..5e618df 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/V1TestUtil.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/V1TestUtil.java
@@ -374,7 +374,7 @@ class V1TestUtil {
// Read the next batch of query results.
private Iterator<EntityResult> getIteratorAndMoveCursor() throws DatastoreException {
- Query.Builder query = this.query.toBuilder();
+ Query.Builder query = this.query.toBuilder().clone();
query.setLimit(Int32Value.newBuilder().setValue(QUERY_BATCH_LIMIT));
if (currentBatch != null && !currentBatch.getEndCursor().isEmpty()) {
query.setStartCursor(currentBatch.getEndCursor());
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/FakeServiceFactory.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/FakeServiceFactory.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/FakeServiceFactory.java
deleted file mode 100644
index 753d807..0000000
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/FakeServiceFactory.java
+++ /dev/null
@@ -1,82 +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.sdk.io.gcp.spanner;
-
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
-import static org.mockito.Mockito.withSettings;
-
-import com.google.cloud.ServiceFactory;
-import com.google.cloud.spanner.DatabaseClient;
-import com.google.cloud.spanner.DatabaseId;
-import com.google.cloud.spanner.Spanner;
-import com.google.cloud.spanner.SpannerOptions;
-import java.io.Serializable;
-import java.util.ArrayList;
-import java.util.List;
-import javax.annotation.concurrent.GuardedBy;
-import org.mockito.Matchers;
-
-/**
- * A serialization friendly type service factory that maintains a mock {@link Spanner} and
- * {@link DatabaseClient}.
- * */
-class FakeServiceFactory
- implements ServiceFactory<Spanner, SpannerOptions>, Serializable {
-
- // Marked as static so they could be returned by serviceFactory, which is serializable.
- private static final Object lock = new Object();
-
- @GuardedBy("lock")
- private static final List<Spanner> mockSpanners = new ArrayList<>();
-
- @GuardedBy("lock")
- private static final List<DatabaseClient> mockDatabaseClients = new ArrayList<>();
-
- @GuardedBy("lock")
- private static int count = 0;
-
- private final int index;
-
- public FakeServiceFactory() {
- synchronized (lock) {
- index = count++;
- mockSpanners.add(mock(Spanner.class, withSettings().serializable()));
- mockDatabaseClients.add(mock(DatabaseClient.class, withSettings().serializable()));
- }
- when(mockSpanner().getDatabaseClient(Matchers.any(DatabaseId.class)))
- .thenReturn(mockDatabaseClient());
- }
-
- DatabaseClient mockDatabaseClient() {
- synchronized (lock) {
- return mockDatabaseClients.get(index);
- }
- }
-
- Spanner mockSpanner() {
- synchronized (lock) {
- return mockSpanners.get(index);
- }
- }
-
- @Override
- public Spanner create(SpannerOptions serviceOptions) {
- return mockSpanner();
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimatorTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimatorTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimatorTest.java
index 013b83d..03eb28e 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimatorTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimatorTest.java
@@ -135,16 +135,4 @@ public class MutationSizeEstimatorTest {
assertThat(MutationSizeEstimator.sizeOf(timestampArray), is(24L));
assertThat(MutationSizeEstimator.sizeOf(dateArray), is(48L));
}
-
- @Test
- public void group() throws Exception {
- Mutation int64 = Mutation.newInsertOrUpdateBuilder("test").set("one").to(1).build();
- Mutation float64 = Mutation.newInsertOrUpdateBuilder("test").set("one").to(2.9).build();
- Mutation bool = Mutation.newInsertOrUpdateBuilder("test").set("one").to(false).build();
-
- MutationGroup group = MutationGroup.create(int64, float64, bool);
-
- assertThat(MutationSizeEstimator.sizeOf(group), is(17L));
- }
-
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/RandomUtils.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/RandomUtils.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/RandomUtils.java
deleted file mode 100644
index f479b4a..0000000
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/RandomUtils.java
+++ /dev/null
@@ -1,41 +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.sdk.io.gcp.spanner;
-
-import java.util.Random;
-
-/**
- * Useful randomness related utilities.
- */
-public class RandomUtils {
-
- private static final char[] ALPHANUMERIC = "1234567890abcdefghijklmnopqrstuvwxyz".toCharArray();
-
- private RandomUtils() {
- }
-
- public static String randomAlphaNumeric(int length) {
- Random random = new Random();
- char[] result = new char[length];
- for (int i = 0; i < length; i++) {
- result[i] = ALPHANUMERIC[random.nextInt(ALPHANUMERIC.length)];
- }
- return new String(result);
- }
-
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOReadTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOReadTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOReadTest.java
deleted file mode 100644
index 5ba2da0..0000000
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOReadTest.java
+++ /dev/null
@@ -1,281 +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.sdk.io.gcp.spanner;
-
-import static org.junit.Assert.assertThat;
-import static org.mockito.Matchers.any;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
-import com.google.cloud.Timestamp;
-import com.google.cloud.spanner.KeySet;
-import com.google.cloud.spanner.ReadOnlyTransaction;
-import com.google.cloud.spanner.ResultSets;
-import com.google.cloud.spanner.Statement;
-import com.google.cloud.spanner.Struct;
-import com.google.cloud.spanner.TimestampBound;
-import com.google.cloud.spanner.Type;
-import com.google.cloud.spanner.Value;
-import java.io.Serializable;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.List;
-import org.apache.beam.sdk.testing.NeedsRunner;
-import org.apache.beam.sdk.testing.PAssert;
-import org.apache.beam.sdk.testing.TestPipeline;
-import org.apache.beam.sdk.transforms.DoFnTester;
-import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PCollectionView;
-import org.hamcrest.Matchers;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.experimental.categories.Category;
-import org.junit.rules.ExpectedException;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-import org.mockito.Mockito;
-
-/** Unit tests for {@link SpannerIO}. */
-@RunWith(JUnit4.class)
-public class SpannerIOReadTest implements Serializable {
- @Rule
- public final transient TestPipeline pipeline = TestPipeline.create();
- @Rule
- public final transient ExpectedException thrown = ExpectedException.none();
-
- private FakeServiceFactory serviceFactory;
- private ReadOnlyTransaction mockTx;
-
- private Type fakeType = Type.struct(Type.StructField.of("id", Type.int64()),
- Type.StructField.of("name", Type.string()));
-
- private List<Struct> fakeRows = Arrays.asList(
- Struct.newBuilder().add("id", Value.int64(1)).add("name", Value.string("Alice")).build(),
- Struct.newBuilder().add("id", Value.int64(2)).add("name", Value.string("Bob")).build());
-
- @Before
- @SuppressWarnings("unchecked")
- public void setUp() throws Exception {
- serviceFactory = new FakeServiceFactory();
- mockTx = Mockito.mock(ReadOnlyTransaction.class);
- }
-
- @Test
- public void emptyTransform() throws Exception {
- SpannerIO.Read read = SpannerIO.read();
- thrown.expect(NullPointerException.class);
- thrown.expectMessage("requires instance id to be set with");
- read.validate(null);
- }
-
- @Test
- public void emptyInstanceId() throws Exception {
- SpannerIO.Read read = SpannerIO.read().withDatabaseId("123");
- thrown.expect(NullPointerException.class);
- thrown.expectMessage("requires instance id to be set with");
- read.validate(null);
- }
-
- @Test
- public void emptyDatabaseId() throws Exception {
- SpannerIO.Read read = SpannerIO.read().withInstanceId("123");
- thrown.expect(NullPointerException.class);
- thrown.expectMessage("requires database id to be set with");
- read.validate(null);
- }
-
- @Test
- public void emptyQuery() throws Exception {
- SpannerIO.Read read =
- SpannerIO.read().withInstanceId("123").withDatabaseId("aaa").withTimestamp(Timestamp.now());
- thrown.expect(IllegalArgumentException.class);
- thrown.expectMessage("requires configuring query or read operation");
- read.validate(null);
- }
-
- @Test
- public void emptyColumns() throws Exception {
- SpannerIO.Read read =
- SpannerIO.read()
- .withInstanceId("123")
- .withDatabaseId("aaa")
- .withTimestamp(Timestamp.now())
- .withTable("users");
- thrown.expect(NullPointerException.class);
- thrown.expectMessage("requires a list of columns");
- read.validate(null);
- }
-
- @Test
- public void validRead() throws Exception {
- SpannerIO.Read read =
- SpannerIO.read()
- .withInstanceId("123")
- .withDatabaseId("aaa")
- .withTimestamp(Timestamp.now())
- .withTable("users")
- .withColumns("id", "name", "email");
- read.validate(null);
- }
-
- @Test
- public void validQuery() throws Exception {
- SpannerIO.Read read =
- SpannerIO.read()
- .withInstanceId("123")
- .withDatabaseId("aaa")
- .withTimestamp(Timestamp.now())
- .withQuery("SELECT * FROM users");
- read.validate(null);
- }
-
- @Test
- public void runQuery() throws Exception {
- SpannerIO.Read read =
- SpannerIO.read()
- .withProjectId("test")
- .withInstanceId("123")
- .withDatabaseId("aaa")
- .withTimestamp(Timestamp.now())
- .withQuery("SELECT * FROM users")
- .withServiceFactory(serviceFactory);
-
- NaiveSpannerReadFn readFn = new NaiveSpannerReadFn(read);
- DoFnTester<Object, Struct> fnTester = DoFnTester.of(readFn);
-
- when(serviceFactory.mockDatabaseClient().readOnlyTransaction(any(TimestampBound.class)))
- .thenReturn(mockTx);
- when(mockTx.executeQuery(any(Statement.class)))
- .thenReturn(ResultSets.forRows(fakeType, fakeRows));
-
- List<Struct> result = fnTester.processBundle(1);
- assertThat(result, Matchers.<Struct>iterableWithSize(2));
-
- verify(serviceFactory.mockDatabaseClient()).readOnlyTransaction(TimestampBound
- .strong());
- verify(mockTx).executeQuery(Statement.of("SELECT * FROM users"));
- }
-
- @Test
- public void runRead() throws Exception {
- SpannerIO.Read read =
- SpannerIO.read()
- .withProjectId("test")
- .withInstanceId("123")
- .withDatabaseId("aaa")
- .withTimestamp(Timestamp.now())
- .withTable("users")
- .withColumns("id", "name")
- .withServiceFactory(serviceFactory);
-
- NaiveSpannerReadFn readFn = new NaiveSpannerReadFn(read);
- DoFnTester<Object, Struct> fnTester = DoFnTester.of(readFn);
-
- when(serviceFactory.mockDatabaseClient().readOnlyTransaction(any(TimestampBound.class)))
- .thenReturn(mockTx);
- when(mockTx.read("users", KeySet.all(), Arrays.asList("id", "name")))
- .thenReturn(ResultSets.forRows(fakeType, fakeRows));
-
- List<Struct> result = fnTester.processBundle(1);
- assertThat(result, Matchers.<Struct>iterableWithSize(2));
-
- verify(serviceFactory.mockDatabaseClient()).readOnlyTransaction(TimestampBound.strong());
- verify(mockTx).read("users", KeySet.all(), Arrays.asList("id", "name"));
- }
-
- @Test
- public void runReadUsingIndex() throws Exception {
- SpannerIO.Read read =
- SpannerIO.read()
- .withProjectId("test")
- .withInstanceId("123")
- .withDatabaseId("aaa")
- .withTimestamp(Timestamp.now())
- .withTable("users")
- .withColumns("id", "name")
- .withIndex("theindex")
- .withServiceFactory(serviceFactory);
-
- NaiveSpannerReadFn readFn = new NaiveSpannerReadFn(read);
- DoFnTester<Object, Struct> fnTester = DoFnTester.of(readFn);
-
- when(serviceFactory.mockDatabaseClient().readOnlyTransaction(any(TimestampBound.class)))
- .thenReturn(mockTx);
- when(mockTx.readUsingIndex("users", "theindex", KeySet.all(), Arrays.asList("id", "name")))
- .thenReturn(ResultSets.forRows(fakeType, fakeRows));
-
- List<Struct> result = fnTester.processBundle(1);
- assertThat(result, Matchers.<Struct>iterableWithSize(2));
-
- verify(serviceFactory.mockDatabaseClient()).readOnlyTransaction(TimestampBound.strong());
- verify(mockTx).readUsingIndex("users", "theindex", KeySet.all(), Arrays.asList("id", "name"));
- }
-
- @Test
- @Category(NeedsRunner.class)
- public void readPipeline() throws Exception {
- Timestamp timestamp = Timestamp.ofTimeMicroseconds(12345);
-
- PCollectionView<Transaction> tx = pipeline
- .apply("tx", SpannerIO.createTransaction()
- .withProjectId("test")
- .withInstanceId("123")
- .withDatabaseId("aaa")
- .withServiceFactory(serviceFactory));
-
- PCollection<Struct> one = pipeline.apply("read q", SpannerIO.read()
- .withProjectId("test")
- .withInstanceId("123")
- .withDatabaseId("aaa")
- .withTimestamp(Timestamp.now())
- .withQuery("SELECT * FROM users")
- .withServiceFactory(serviceFactory)
- .withTransaction(tx));
- PCollection<Struct> two = pipeline.apply("read r", SpannerIO.read()
- .withProjectId("test")
- .withInstanceId("123")
- .withDatabaseId("aaa")
- .withTimestamp(Timestamp.now())
- .withTable("users")
- .withColumns("id", "name")
- .withServiceFactory(serviceFactory)
- .withTransaction(tx));
-
- when(serviceFactory.mockDatabaseClient().readOnlyTransaction(any(TimestampBound.class)))
- .thenReturn(mockTx);
-
- when(mockTx.executeQuery(Statement.of("SELECT 1"))).thenReturn(ResultSets.forRows(Type.struct(),
- Collections.<Struct>emptyList()));
-
- when(mockTx.executeQuery(Statement.of("SELECT * FROM users")))
- .thenReturn(ResultSets.forRows(fakeType, fakeRows));
- when(mockTx.read("users", KeySet.all(), Arrays.asList("id", "name")))
- .thenReturn(ResultSets.forRows(fakeType, fakeRows));
- when(mockTx.getReadTimestamp()).thenReturn(timestamp);
-
- PAssert.that(one).containsInAnyOrder(fakeRows);
- PAssert.that(two).containsInAnyOrder(fakeRows);
-
- pipeline.run();
-
- verify(serviceFactory.mockDatabaseClient(), times(2))
- .readOnlyTransaction(TimestampBound.ofReadTimestamp(timestamp));
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOTest.java
new file mode 100644
index 0000000..5bdfea5
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOTest.java
@@ -0,0 +1,244 @@
+/*
+ * 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.sdk.io.gcp.spanner;
+
+import static org.mockito.Mockito.argThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.withSettings;
+
+import com.google.api.core.ApiFuture;
+import com.google.cloud.ServiceFactory;
+import com.google.cloud.spanner.DatabaseClient;
+import com.google.cloud.spanner.DatabaseId;
+import com.google.cloud.spanner.Mutation;
+import com.google.cloud.spanner.Spanner;
+import com.google.cloud.spanner.SpannerOptions;
+import com.google.common.collect.Iterables;
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import javax.annotation.concurrent.GuardedBy;
+
+import org.apache.beam.sdk.testing.NeedsRunner;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.DoFnTester;
+import org.apache.beam.sdk.values.PCollection;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.rules.ExpectedException;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.ArgumentMatcher;
+import org.mockito.Matchers;
+
+
+/**
+ * Unit tests for {@link SpannerIO}.
+ */
+@RunWith(JUnit4.class)
+public class SpannerIOTest implements Serializable {
+ @Rule public final transient TestPipeline pipeline = TestPipeline.create();
+ @Rule public transient ExpectedException thrown = ExpectedException.none();
+
+ private FakeServiceFactory serviceFactory;
+
+ @Before
+ @SuppressWarnings("unchecked")
+ public void setUp() throws Exception {
+ serviceFactory = new FakeServiceFactory();
+ }
+
+ @Test
+ public void emptyTransform() throws Exception {
+ SpannerIO.Write write = SpannerIO.write();
+ thrown.expect(NullPointerException.class);
+ thrown.expectMessage("requires instance id to be set with");
+ write.validate(null);
+ }
+
+ @Test
+ public void emptyInstanceId() throws Exception {
+ SpannerIO.Write write = SpannerIO.write().withDatabaseId("123");
+ thrown.expect(NullPointerException.class);
+ thrown.expectMessage("requires instance id to be set with");
+ write.validate(null);
+ }
+
+ @Test
+ public void emptyDatabaseId() throws Exception {
+ SpannerIO.Write write = SpannerIO.write().withInstanceId("123");
+ thrown.expect(NullPointerException.class);
+ thrown.expectMessage("requires database id to be set with");
+ write.validate(null);
+ }
+
+ @Test
+ @Category(NeedsRunner.class)
+ public void singleMutationPipeline() throws Exception {
+ Mutation mutation = Mutation.newInsertOrUpdateBuilder("test").set("one").to(2).build();
+ PCollection<Mutation> mutations = pipeline.apply(Create.of(mutation));
+
+ mutations.apply(
+ SpannerIO.write()
+ .withProjectId("test-project")
+ .withInstanceId("test-instance")
+ .withDatabaseId("test-database")
+ .withServiceFactory(serviceFactory));
+ pipeline.run();
+ verify(serviceFactory.mockSpanner())
+ .getDatabaseClient(DatabaseId.of("test-project", "test-instance", "test-database"));
+ verify(serviceFactory.mockDatabaseClient(), times(1))
+ .writeAtLeastOnce(argThat(new IterableOfSize(1)));
+ }
+
+ @Test
+ public void batching() throws Exception {
+ Mutation one = Mutation.newInsertOrUpdateBuilder("test").set("one").to(1).build();
+ Mutation two = Mutation.newInsertOrUpdateBuilder("test").set("two").to(2).build();
+ SpannerIO.Write write =
+ SpannerIO.write()
+ .withProjectId("test-project")
+ .withInstanceId("test-instance")
+ .withDatabaseId("test-database")
+ .withBatchSizeBytes(1000000000)
+ .withServiceFactory(serviceFactory);
+ SpannerIO.SpannerWriteFn writerFn = new SpannerIO.SpannerWriteFn(write);
+ DoFnTester<Mutation, Void> fnTester = DoFnTester.of(writerFn);
+ fnTester.processBundle(Arrays.asList(one, two));
+
+ verify(serviceFactory.mockSpanner())
+ .getDatabaseClient(DatabaseId.of("test-project", "test-instance", "test-database"));
+ verify(serviceFactory.mockDatabaseClient(), times(1))
+ .writeAtLeastOnce(argThat(new IterableOfSize(2)));
+ }
+
+ @Test
+ public void batchingGroups() throws Exception {
+ Mutation one = Mutation.newInsertOrUpdateBuilder("test").set("one").to(1).build();
+ Mutation two = Mutation.newInsertOrUpdateBuilder("test").set("two").to(2).build();
+ Mutation three = Mutation.newInsertOrUpdateBuilder("test").set("three").to(3).build();
+
+ // Have a room to accumulate one more item.
+ long batchSize = MutationSizeEstimator.sizeOf(one) + 1;
+
+ SpannerIO.Write write =
+ SpannerIO.write()
+ .withProjectId("test-project")
+ .withInstanceId("test-instance")
+ .withDatabaseId("test-database")
+ .withBatchSizeBytes(batchSize)
+ .withServiceFactory(serviceFactory);
+ SpannerIO.SpannerWriteFn writerFn = new SpannerIO.SpannerWriteFn(write);
+ DoFnTester<Mutation, Void> fnTester = DoFnTester.of(writerFn);
+ fnTester.processBundle(Arrays.asList(one, two, three));
+
+ verify(serviceFactory.mockSpanner())
+ .getDatabaseClient(DatabaseId.of("test-project", "test-instance", "test-database"));
+ verify(serviceFactory.mockDatabaseClient(), times(1))
+ .writeAtLeastOnce(argThat(new IterableOfSize(2)));
+ verify(serviceFactory.mockDatabaseClient(), times(1))
+ .writeAtLeastOnce(argThat(new IterableOfSize(1)));
+ }
+
+ @Test
+ public void noBatching() throws Exception {
+ Mutation one = Mutation.newInsertOrUpdateBuilder("test").set("one").to(1).build();
+ Mutation two = Mutation.newInsertOrUpdateBuilder("test").set("two").to(2).build();
+ SpannerIO.Write write =
+ SpannerIO.write()
+ .withProjectId("test-project")
+ .withInstanceId("test-instance")
+ .withDatabaseId("test-database")
+ .withBatchSizeBytes(0) // turn off batching.
+ .withServiceFactory(serviceFactory);
+ SpannerIO.SpannerWriteFn writerFn = new SpannerIO.SpannerWriteFn(write);
+ DoFnTester<Mutation, Void> fnTester = DoFnTester.of(writerFn);
+ fnTester.processBundle(Arrays.asList(one, two));
+
+ verify(serviceFactory.mockSpanner())
+ .getDatabaseClient(DatabaseId.of("test-project", "test-instance", "test-database"));
+ verify(serviceFactory.mockDatabaseClient(), times(2))
+ .writeAtLeastOnce(argThat(new IterableOfSize(1)));
+ }
+
+ private static class FakeServiceFactory
+ implements ServiceFactory<Spanner, SpannerOptions>, Serializable {
+ // Marked as static so they could be returned by serviceFactory, which is serializable.
+ private static final Object lock = new Object();
+
+ @GuardedBy("lock")
+ private static final List<Spanner> mockSpanners = new ArrayList<>();
+
+ @GuardedBy("lock")
+ private static final List<DatabaseClient> mockDatabaseClients = new ArrayList<>();
+
+ @GuardedBy("lock")
+ private static int count = 0;
+
+ private final int index;
+
+ public FakeServiceFactory() {
+ synchronized (lock) {
+ index = count++;
+ mockSpanners.add(mock(Spanner.class, withSettings().serializable()));
+ mockDatabaseClients.add(mock(DatabaseClient.class, withSettings().serializable()));
+ }
+ ApiFuture voidFuture = mock(ApiFuture.class, withSettings().serializable());
+ when(mockSpanner().getDatabaseClient(Matchers.any(DatabaseId.class)))
+ .thenReturn(mockDatabaseClient());
+ when(mockSpanner().closeAsync()).thenReturn(voidFuture);
+ }
+
+ DatabaseClient mockDatabaseClient() {
+ synchronized (lock) {
+ return mockDatabaseClients.get(index);
+ }
+ }
+
+ Spanner mockSpanner() {
+ synchronized (lock) {
+ return mockSpanners.get(index);
+ }
+ }
+
+ @Override
+ public Spanner create(SpannerOptions serviceOptions) {
+ return mockSpanner();
+ }
+ }
+
+ private static class IterableOfSize extends ArgumentMatcher<Iterable<Mutation>> {
+ private final int size;
+
+ private IterableOfSize(int size) {
+ this.size = size;
+ }
+
+ @Override
+ public boolean matches(Object argument) {
+ return argument instanceof Iterable && Iterables.size((Iterable<?>) argument) == size;
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOWriteTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOWriteTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOWriteTest.java
deleted file mode 100644
index 09cdb8e..0000000
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIOWriteTest.java
+++ /dev/null
@@ -1,258 +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.sdk.io.gcp.spanner;
-
-import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem;
-import static org.hamcrest.Matchers.hasSize;
-import static org.junit.Assert.assertThat;
-import static org.mockito.Mockito.argThat;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-
-import com.google.cloud.spanner.DatabaseId;
-import com.google.cloud.spanner.Mutation;
-import com.google.common.collect.Iterables;
-import java.io.Serializable;
-import java.util.Arrays;
-
-import org.apache.beam.sdk.testing.NeedsRunner;
-import org.apache.beam.sdk.testing.TestPipeline;
-import org.apache.beam.sdk.transforms.Create;
-import org.apache.beam.sdk.transforms.DoFnTester;
-import org.apache.beam.sdk.transforms.display.DisplayData;
-import org.apache.beam.sdk.values.PCollection;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.experimental.categories.Category;
-import org.junit.rules.ExpectedException;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-import org.mockito.ArgumentMatcher;
-
-/**
- * Unit tests for {@link SpannerIO}.
- */
-@RunWith(JUnit4.class)
-public class SpannerIOWriteTest implements Serializable {
- @Rule public final transient TestPipeline pipeline = TestPipeline.create();
- @Rule public transient ExpectedException thrown = ExpectedException.none();
-
- private FakeServiceFactory serviceFactory;
-
- @Before
- @SuppressWarnings("unchecked")
- public void setUp() throws Exception {
- serviceFactory = new FakeServiceFactory();
- }
-
- @Test
- public void emptyTransform() throws Exception {
- SpannerIO.Write write = SpannerIO.write();
- thrown.expect(NullPointerException.class);
- thrown.expectMessage("requires instance id to be set with");
- write.validate(null);
- }
-
- @Test
- public void emptyInstanceId() throws Exception {
- SpannerIO.Write write = SpannerIO.write().withDatabaseId("123");
- thrown.expect(NullPointerException.class);
- thrown.expectMessage("requires instance id to be set with");
- write.validate(null);
- }
-
- @Test
- public void emptyDatabaseId() throws Exception {
- SpannerIO.Write write = SpannerIO.write().withInstanceId("123");
- thrown.expect(NullPointerException.class);
- thrown.expectMessage("requires database id to be set with");
- write.validate(null);
- }
-
- @Test
- @Category(NeedsRunner.class)
- public void singleMutationPipeline() throws Exception {
- Mutation mutation = Mutation.newInsertOrUpdateBuilder("test").set("one").to(2).build();
- PCollection<Mutation> mutations = pipeline.apply(Create.of(mutation));
-
- mutations.apply(
- SpannerIO.write()
- .withProjectId("test-project")
- .withInstanceId("test-instance")
- .withDatabaseId("test-database")
- .withServiceFactory(serviceFactory));
- pipeline.run();
- verify(serviceFactory.mockSpanner())
- .getDatabaseClient(DatabaseId.of("test-project", "test-instance", "test-database"));
- verify(serviceFactory.mockDatabaseClient(), times(1))
- .writeAtLeastOnce(argThat(new IterableOfSize(1)));
- }
-
- @Test
- @Category(NeedsRunner.class)
- public void singleMutationGroupPipeline() throws Exception {
- Mutation one = Mutation.newInsertOrUpdateBuilder("test").set("one").to(1).build();
- Mutation two = Mutation.newInsertOrUpdateBuilder("test").set("two").to(2).build();
- Mutation three = Mutation.newInsertOrUpdateBuilder("test").set("three").to(3).build();
- PCollection<MutationGroup> mutations = pipeline
- .apply(Create.<MutationGroup>of(g(one, two, three)));
- mutations.apply(
- SpannerIO.write()
- .withProjectId("test-project")
- .withInstanceId("test-instance")
- .withDatabaseId("test-database")
- .withServiceFactory(serviceFactory)
- .grouped());
- pipeline.run();
- verify(serviceFactory.mockSpanner())
- .getDatabaseClient(DatabaseId.of("test-project", "test-instance", "test-database"));
- verify(serviceFactory.mockDatabaseClient(), times(1))
- .writeAtLeastOnce(argThat(new IterableOfSize(3)));
- }
-
- @Test
- public void batching() throws Exception {
- MutationGroup one = g(Mutation.newInsertOrUpdateBuilder("test").set("one").to(1).build());
- MutationGroup two = g(Mutation.newInsertOrUpdateBuilder("test").set("two").to(2).build());
- SpannerIO.Write write =
- SpannerIO.write()
- .withProjectId("test-project")
- .withInstanceId("test-instance")
- .withDatabaseId("test-database")
- .withBatchSizeBytes(1000000000)
- .withServiceFactory(serviceFactory);
- SpannerWriteGroupFn writerFn = new SpannerWriteGroupFn(write);
- DoFnTester<MutationGroup, Void> fnTester = DoFnTester.of(writerFn);
- fnTester.processBundle(Arrays.asList(one, two));
-
- verify(serviceFactory.mockSpanner())
- .getDatabaseClient(DatabaseId.of("test-project", "test-instance", "test-database"));
- verify(serviceFactory.mockDatabaseClient(), times(1))
- .writeAtLeastOnce(argThat(new IterableOfSize(2)));
- }
-
- @Test
- public void batchingGroups() throws Exception {
- MutationGroup one = g(Mutation.newInsertOrUpdateBuilder("test").set("one").to(1).build());
- MutationGroup two = g(Mutation.newInsertOrUpdateBuilder("test").set("two").to(2).build());
- MutationGroup three = g(Mutation.newInsertOrUpdateBuilder("test").set("three").to(3).build());
-
- // Have a room to accumulate one more item.
- long batchSize = MutationSizeEstimator.sizeOf(one) + 1;
-
- SpannerIO.Write write =
- SpannerIO.write()
- .withProjectId("test-project")
- .withInstanceId("test-instance")
- .withDatabaseId("test-database")
- .withBatchSizeBytes(batchSize)
- .withServiceFactory(serviceFactory);
- SpannerWriteGroupFn writerFn = new SpannerWriteGroupFn(write);
- DoFnTester<MutationGroup, Void> fnTester = DoFnTester.of(writerFn);
- fnTester.processBundle(Arrays.asList(one, two, three));
-
- verify(serviceFactory.mockSpanner())
- .getDatabaseClient(DatabaseId.of("test-project", "test-instance", "test-database"));
- verify(serviceFactory.mockDatabaseClient(), times(1))
- .writeAtLeastOnce(argThat(new IterableOfSize(2)));
- verify(serviceFactory.mockDatabaseClient(), times(1))
- .writeAtLeastOnce(argThat(new IterableOfSize(1)));
- }
-
- @Test
- public void noBatching() throws Exception {
- MutationGroup one = g(Mutation.newInsertOrUpdateBuilder("test").set("one").to(1).build());
- MutationGroup two = g(Mutation.newInsertOrUpdateBuilder("test").set("two").to(2).build());
- SpannerIO.Write write =
- SpannerIO.write()
- .withProjectId("test-project")
- .withInstanceId("test-instance")
- .withDatabaseId("test-database")
- .withBatchSizeBytes(0) // turn off batching.
- .withServiceFactory(serviceFactory);
- SpannerWriteGroupFn writerFn = new SpannerWriteGroupFn(write);
- DoFnTester<MutationGroup, Void> fnTester = DoFnTester.of(writerFn);
- fnTester.processBundle(Arrays.asList(one, two));
-
- verify(serviceFactory.mockSpanner())
- .getDatabaseClient(DatabaseId.of("test-project", "test-instance", "test-database"));
- verify(serviceFactory.mockDatabaseClient(), times(2))
- .writeAtLeastOnce(argThat(new IterableOfSize(1)));
- }
-
- @Test
- public void groups() throws Exception {
- Mutation one = Mutation.newInsertOrUpdateBuilder("test").set("one").to(1).build();
- Mutation two = Mutation.newInsertOrUpdateBuilder("test").set("two").to(2).build();
- Mutation three = Mutation.newInsertOrUpdateBuilder("test").set("three").to(3).build();
-
- // Smallest batch size
- long batchSize = 1;
-
- SpannerIO.Write write =
- SpannerIO.write()
- .withProjectId("test-project")
- .withInstanceId("test-instance")
- .withDatabaseId("test-database")
- .withBatchSizeBytes(batchSize)
- .withServiceFactory(serviceFactory);
- SpannerWriteGroupFn writerFn = new SpannerWriteGroupFn(write);
- DoFnTester<MutationGroup, Void> fnTester = DoFnTester.of(writerFn);
- fnTester.processBundle(Arrays.asList(g(one, two, three)));
-
- verify(serviceFactory.mockSpanner())
- .getDatabaseClient(DatabaseId.of("test-project", "test-instance", "test-database"));
- verify(serviceFactory.mockDatabaseClient(), times(1))
- .writeAtLeastOnce(argThat(new IterableOfSize(3)));
- }
-
- @Test
- public void displayData() throws Exception {
- SpannerIO.Write write =
- SpannerIO.write()
- .withProjectId("test-project")
- .withInstanceId("test-instance")
- .withDatabaseId("test-database")
- .withBatchSizeBytes(123);
-
- DisplayData data = DisplayData.from(write);
- assertThat(data.items(), hasSize(4));
- assertThat(data, hasDisplayItem("projectId", "test-project"));
- assertThat(data, hasDisplayItem("instanceId", "test-instance"));
- assertThat(data, hasDisplayItem("databaseId", "test-database"));
- assertThat(data, hasDisplayItem("batchSizeBytes", 123));
- }
-
- private static class IterableOfSize extends ArgumentMatcher<Iterable<Mutation>> {
- private final int size;
-
- private IterableOfSize(int size) {
- this.size = size;
- }
-
- @Override
- public boolean matches(Object argument) {
- return argument instanceof Iterable && Iterables.size((Iterable<?>) argument) == size;
- }
- }
-
- private static MutationGroup g(Mutation m, Mutation... other) {
- return MutationGroup.create(m, other);
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerReadIT.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerReadIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerReadIT.java
deleted file mode 100644
index d866975..0000000
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerReadIT.java
+++ /dev/null
@@ -1,166 +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.sdk.io.gcp.spanner;
-
-import com.google.cloud.spanner.Database;
-import com.google.cloud.spanner.DatabaseAdminClient;
-import com.google.cloud.spanner.DatabaseClient;
-import com.google.cloud.spanner.DatabaseId;
-import com.google.cloud.spanner.Mutation;
-import com.google.cloud.spanner.Operation;
-import com.google.cloud.spanner.Spanner;
-import com.google.cloud.spanner.SpannerOptions;
-import com.google.cloud.spanner.Struct;
-import com.google.cloud.spanner.TimestampBound;
-import com.google.spanner.admin.database.v1.CreateDatabaseMetadata;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
-
-import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
-import org.apache.beam.sdk.options.Default;
-import org.apache.beam.sdk.options.Description;
-import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.sdk.testing.PAssert;
-import org.apache.beam.sdk.testing.TestPipeline;
-import org.apache.beam.sdk.testing.TestPipelineOptions;
-import org.apache.beam.sdk.transforms.Count;
-import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PCollectionView;
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-/** End-to-end test of Cloud Spanner Source. */
-@RunWith(JUnit4.class)
-public class SpannerReadIT {
-
- private static final int MAX_DB_NAME_LENGTH = 30;
-
- @Rule public final transient TestPipeline p = TestPipeline.create();
-
- /** Pipeline options for this test. */
- public interface SpannerTestPipelineOptions extends TestPipelineOptions {
- @Description("Instance ID to write to in Spanner")
- @Default.String("beam-test")
- String getInstanceId();
- void setInstanceId(String value);
-
- @Description("Database ID prefix to write to in Spanner")
- @Default.String("beam-testdb")
- String getDatabaseIdPrefix();
- void setDatabaseIdPrefix(String value);
-
- @Description("Table name")
- @Default.String("users")
- String getTable();
- void setTable(String value);
- }
-
- private Spanner spanner;
- private DatabaseAdminClient databaseAdminClient;
- private SpannerTestPipelineOptions options;
- private String databaseName;
- private String project;
-
- @Before
- public void setUp() throws Exception {
- PipelineOptionsFactory.register(SpannerTestPipelineOptions.class);
- options = TestPipeline.testingPipelineOptions().as(SpannerTestPipelineOptions.class);
-
- project = options.as(GcpOptions.class).getProject();
-
- spanner = SpannerOptions.newBuilder().setProjectId(project).build().getService();
-
- databaseName = generateDatabaseName();
-
- databaseAdminClient = spanner.getDatabaseAdminClient();
-
- // Delete database if exists.
- databaseAdminClient.dropDatabase(options.getInstanceId(), databaseName);
-
- Operation<Database, CreateDatabaseMetadata> op =
- databaseAdminClient.createDatabase(
- options.getInstanceId(),
- databaseName,
- Collections.singleton(
- "CREATE TABLE "
- + options.getTable()
- + " ("
- + " Key INT64,"
- + " Value STRING(MAX),"
- + ") PRIMARY KEY (Key)"));
- op.waitFor();
- }
-
- @Test
- public void testRead() throws Exception {
- DatabaseClient databaseClient =
- spanner.getDatabaseClient(
- DatabaseId.of(
- project, options.getInstanceId(), databaseName));
-
- List<Mutation> mutations = new ArrayList<>();
- for (int i = 0; i < 5L; i++) {
- mutations.add(
- Mutation.newInsertOrUpdateBuilder(options.getTable())
- .set("key")
- .to((long) i)
- .set("value")
- .to(RandomUtils.randomAlphaNumeric(100))
- .build());
- }
-
- databaseClient.writeAtLeastOnce(mutations);
-
- SpannerConfig spannerConfig = SpannerConfig.create()
- .withProjectId(project)
- .withInstanceId(options.getInstanceId())
- .withDatabaseId(databaseName);
-
- PCollectionView<Transaction> tx =
- p.apply(
- SpannerIO.createTransaction()
- .withSpannerConfig(spannerConfig)
- .withTimestampBound(TimestampBound.strong()));
-
- PCollection<Struct> output =
- p.apply(
- SpannerIO.read()
- .withSpannerConfig(spannerConfig)
- .withQuery("SELECT * FROM " + options.getTable())
- .withTransaction(tx));
- PAssert.thatSingleton(output.apply("Count rows", Count.<Struct>globally())).isEqualTo(5L);
- p.run();
- }
-
- @After
- public void tearDown() throws Exception {
- databaseAdminClient.dropDatabase(options.getInstanceId(), databaseName);
- spanner.close();
- }
-
- private String generateDatabaseName() {
- String random = RandomUtils
- .randomAlphaNumeric(MAX_DB_NAME_LENGTH - 1 - options.getDatabaseIdPrefix().length());
- return options.getDatabaseIdPrefix() + "-" + random;
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteIT.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteIT.java
index d208f5c..8df224b 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteIT.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteIT.java
@@ -33,7 +33,7 @@ import com.google.cloud.spanner.SpannerOptions;
import com.google.cloud.spanner.Statement;
import com.google.spanner.admin.database.v1.CreateDatabaseMetadata;
import java.util.Collections;
-import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
+
import org.apache.beam.sdk.io.GenerateSequence;
import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.Description;
@@ -42,6 +42,7 @@ import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.testing.TestPipelineOptions;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.commons.lang3.RandomStringUtils;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
@@ -59,6 +60,11 @@ public class SpannerWriteIT {
/** Pipeline options for this test. */
public interface SpannerTestPipelineOptions extends TestPipelineOptions {
+ @Description("Project ID for Spanner")
+ @Default.String("apache-beam-testing")
+ String getProjectId();
+ void setProjectId(String value);
+
@Description("Instance ID to write to in Spanner")
@Default.String("beam-test")
String getInstanceId();
@@ -79,16 +85,13 @@ public class SpannerWriteIT {
private DatabaseAdminClient databaseAdminClient;
private SpannerTestPipelineOptions options;
private String databaseName;
- private String project;
@Before
public void setUp() throws Exception {
PipelineOptionsFactory.register(SpannerTestPipelineOptions.class);
options = TestPipeline.testingPipelineOptions().as(SpannerTestPipelineOptions.class);
- project = options.as(GcpOptions.class).getProject();
-
- spanner = SpannerOptions.newBuilder().setProjectId(project).build().getService();
+ spanner = SpannerOptions.newBuilder().setProjectId(options.getProjectId()).build().getService();
databaseName = generateDatabaseName();
@@ -112,8 +115,9 @@ public class SpannerWriteIT {
}
private String generateDatabaseName() {
- String random = RandomUtils
- .randomAlphaNumeric(MAX_DB_NAME_LENGTH - 1 - options.getDatabaseIdPrefix().length());
+ String random = RandomStringUtils
+ .randomAlphanumeric(MAX_DB_NAME_LENGTH - 1 - options.getDatabaseIdPrefix().length())
+ .toLowerCase();
return options.getDatabaseIdPrefix() + "-" + random;
}
@@ -123,7 +127,7 @@ public class SpannerWriteIT {
.apply(ParDo.of(new GenerateMutations(options.getTable())))
.apply(
SpannerIO.write()
- .withProjectId(project)
+ .withProjectId(options.getProjectId())
.withInstanceId(options.getInstanceId())
.withDatabaseId(databaseName));
@@ -131,7 +135,7 @@ public class SpannerWriteIT {
DatabaseClient databaseClient =
spanner.getDatabaseClient(
DatabaseId.of(
- project, options.getInstanceId(), databaseName));
+ options.getProjectId(), options.getInstanceId(), databaseName));
ResultSet resultSet =
databaseClient
@@ -145,7 +149,7 @@ public class SpannerWriteIT {
@After
public void tearDown() throws Exception {
databaseAdminClient.dropDatabase(options.getInstanceId(), databaseName);
- spanner.close();
+ spanner.closeAsync().get();
}
private static class GenerateMutations extends DoFn<Long, Mutation> {
@@ -161,7 +165,7 @@ public class SpannerWriteIT {
Mutation.WriteBuilder builder = Mutation.newInsertOrUpdateBuilder(table);
Long key = c.element();
builder.set("Key").to(key);
- builder.set("Value").to(RandomUtils.randomAlphaNumeric(valueSize));
+ builder.set("Value").to(RandomStringUtils.randomAlphabetic(valueSize));
Mutation mutation = builder.build();
c.output(mutation);
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hadoop-common/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/hadoop-common/pom.xml b/sdks/java/io/hadoop-common/pom.xml
index 4bcbcd7..8749243 100644
--- a/sdks/java/io/hadoop-common/pom.xml
+++ b/sdks/java/io/hadoop-common/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hadoop-file-system/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/hadoop-file-system/pom.xml b/sdks/java/io/hadoop-file-system/pom.xml
index a9c2e57..db5a1db 100644
--- a/sdks/java/io/hadoop-file-system/pom.xml
+++ b/sdks/java/io/hadoop-file-system/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
@@ -44,6 +44,37 @@
</plugins>
</build>
+ <properties>
+ <!--
+ This is the version of Hadoop used to compile the hadoop-common module.
+ This dependency is defined with a provided scope.
+ Users must supply their own Hadoop version at runtime.
+ -->
+ <hadoop.version>2.7.3</hadoop.version>
+ </properties>
+
+ <dependencyManagement>
+ <!--
+ We define dependencies here instead of sdks/java/io because
+ of a version mimatch between this Hadoop version and other
+ Hadoop versions declared in other io submodules.
+ -->
+ <dependencies>
+ <dependency>
+ <groupId>org.apache.hadoop</groupId>
+ <artifactId>hadoop-hdfs</artifactId>
+ <classifier>tests</classifier>
+ <version>${hadoop.version}</version>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.hadoop</groupId>
+ <artifactId>hadoop-minicluster</artifactId>
+ <version>${hadoop.version}</version>
+ </dependency>
+ </dependencies>
+ </dependencyManagement>
+
<dependencies>
<dependency>
<groupId>org.apache.beam</groupId>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hadoop/input-format/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/hadoop/input-format/pom.xml b/sdks/java/io/hadoop/input-format/pom.xml
index 0953119..06f9f11 100644
--- a/sdks/java/io/hadoop/input-format/pom.xml
+++ b/sdks/java/io/hadoop/input-format/pom.xml
@@ -20,7 +20,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-hadoop-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
<artifactId>beam-sdks-java-io-hadoop-input-format</artifactId>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hadoop/input-format/src/main/java/org/apache/beam/sdk/io/hadoop/inputformat/HadoopInputFormatIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hadoop/input-format/src/main/java/org/apache/beam/sdk/io/hadoop/inputformat/HadoopInputFormatIO.java b/sdks/java/io/hadoop/input-format/src/main/java/org/apache/beam/sdk/io/hadoop/inputformat/HadoopInputFormatIO.java
index 0b4c23f..efd47fd 100644
--- a/sdks/java/io/hadoop/input-format/src/main/java/org/apache/beam/sdk/io/hadoop/inputformat/HadoopInputFormatIO.java
+++ b/sdks/java/io/hadoop/input-format/src/main/java/org/apache/beam/sdk/io/hadoop/inputformat/HadoopInputFormatIO.java
@@ -166,7 +166,7 @@ import org.slf4j.LoggerFactory;
* }
* </pre>
*/
-@Experimental(Experimental.Kind.SOURCE_SINK)
+@Experimental
public class HadoopInputFormatIO {
private static final Logger LOG = LoggerFactory.getLogger(HadoopInputFormatIO.class);
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hadoop/jdk1.8-tests/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/hadoop/jdk1.8-tests/pom.xml b/sdks/java/io/hadoop/jdk1.8-tests/pom.xml
index 12944f4..9f84e88 100644
--- a/sdks/java/io/hadoop/jdk1.8-tests/pom.xml
+++ b/sdks/java/io/hadoop/jdk1.8-tests/pom.xml
@@ -26,7 +26,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-hadoop-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
<artifactId>beam-sdks-java-io-hadoop-jdk1.8-tests</artifactId>
@@ -108,11 +108,13 @@
<dependency>
<groupId>org.apache.spark</groupId>
<artifactId>spark-streaming_2.10</artifactId>
+ <version>${spark.version}</version>
<scope>runtime</scope>
</dependency>
<dependency>
<groupId>org.apache.spark</groupId>
<artifactId>spark-core_2.10</artifactId>
+ <version>${spark.version}</version>
<scope>runtime</scope>
<exclusions>
<exclusion>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithElasticTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithElasticTest.java b/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithElasticTest.java
index 3f866a4..8745521 100644
--- a/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithElasticTest.java
+++ b/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithElasticTest.java
@@ -20,7 +20,6 @@ package org.apache.beam.sdk.io.hadoop.inputformat;
import java.io.File;
import java.io.IOException;
import java.io.Serializable;
-import java.net.ServerSocket;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
@@ -77,7 +76,7 @@ public class HIFIOWithElasticTest implements Serializable {
private static final long serialVersionUID = 1L;
private static final Logger LOG = LoggerFactory.getLogger(HIFIOWithElasticTest.class);
private static final String ELASTIC_IN_MEM_HOSTNAME = "127.0.0.1";
- private static String elasticInMemPort = "9200";
+ private static final String ELASTIC_IN_MEM_PORT = "9200";
private static final String ELASTIC_INTERNAL_VERSION = "5.x";
private static final String TRUE = "true";
private static final String ELASTIC_INDEX_NAME = "beamdb";
@@ -95,10 +94,6 @@ public class HIFIOWithElasticTest implements Serializable {
@BeforeClass
public static void startServer()
throws NodeValidationException, InterruptedException, IOException {
- ServerSocket serverSocket = new ServerSocket(0);
- int port = serverSocket.getLocalPort();
- serverSocket.close();
- elasticInMemPort = String.valueOf(port);
ElasticEmbeddedServer.startElasticEmbeddedServer();
}
@@ -178,7 +173,7 @@ public class HIFIOWithElasticTest implements Serializable {
public Configuration getConfiguration() {
Configuration conf = new Configuration();
conf.set(ConfigurationOptions.ES_NODES, ELASTIC_IN_MEM_HOSTNAME);
- conf.set(ConfigurationOptions.ES_PORT, String.format("%s", elasticInMemPort));
+ conf.set(ConfigurationOptions.ES_PORT, String.format("%s", ELASTIC_IN_MEM_PORT));
conf.set(ConfigurationOptions.ES_RESOURCE, ELASTIC_RESOURCE);
conf.set("es.internal.es.version", ELASTIC_INTERNAL_VERSION);
conf.set(ConfigurationOptions.ES_NODES_DISCOVERY, TRUE);
@@ -214,7 +209,7 @@ public class HIFIOWithElasticTest implements Serializable {
Settings settings = Settings.builder()
.put("node.data", TRUE)
.put("network.host", ELASTIC_IN_MEM_HOSTNAME)
- .put("http.port", elasticInMemPort)
+ .put("http.port", ELASTIC_IN_MEM_PORT)
.put("path.data", elasticTempFolder.getRoot().getPath())
.put("path.home", elasticTempFolder.getRoot().getPath())
.put("transport.type", "local")
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hadoop/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/hadoop/pom.xml b/sdks/java/io/hadoop/pom.xml
index bc3569d..a1c7a2e 100644
--- a/sdks/java/io/hadoop/pom.xml
+++ b/sdks/java/io/hadoop/pom.xml
@@ -20,7 +20,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
<packaging>pom</packaging>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hbase/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/hbase/pom.xml b/sdks/java/io/hbase/pom.xml
index 40f516a..746b993 100644
--- a/sdks/java/io/hbase/pom.xml
+++ b/sdks/java/io/hbase/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
@@ -31,7 +31,8 @@
<description>Library to read and write from/to HBase</description>
<properties>
- <hbase.version>1.2.6</hbase.version>
+ <hbase.version>1.2.5</hbase.version>
+ <hbase.hadoop.version>2.5.1</hbase.hadoop.version>
</properties>
<build>
@@ -63,12 +64,6 @@
</dependency>
<dependency>
- <groupId>com.google.auto.service</groupId>
- <artifactId>auto-service</artifactId>
- <optional>true</optional>
- </dependency>
-
- <dependency>
<groupId>org.apache.hbase</groupId>
<artifactId>hbase-shaded-client</artifactId>
<version>${hbase.version}</version>
@@ -108,26 +103,15 @@
<dependency>
<groupId>org.apache.hadoop</groupId>
<artifactId>hadoop-minicluster</artifactId>
- <scope>test</scope>
- </dependency>
-
- <dependency>
- <groupId>org.apache.hadoop</groupId>
- <artifactId>hadoop-hdfs</artifactId>
+ <version>${hbase.hadoop.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.apache.hadoop</groupId>
<artifactId>hadoop-common</artifactId>
+ <version>${hbase.hadoop.version}</version>
<scope>test</scope>
- <exclusions>
- <!-- Fix build on JDK-9 -->
- <exclusion>
- <groupId>jdk.tools</groupId>
- <artifactId>jdk.tools</artifactId>
- </exclusion>
- </exclusions>
</dependency>
<dependency>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseCoderProviderRegistrar.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseCoderProviderRegistrar.java b/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseCoderProviderRegistrar.java
deleted file mode 100644
index 2973d1b..0000000
--- a/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseCoderProviderRegistrar.java
+++ /dev/null
@@ -1,40 +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.sdk.io.hbase;
-
-import com.google.auto.service.AutoService;
-import com.google.common.collect.ImmutableList;
-import java.util.List;
-import org.apache.beam.sdk.coders.CoderProvider;
-import org.apache.beam.sdk.coders.CoderProviderRegistrar;
-import org.apache.beam.sdk.coders.CoderProviders;
-import org.apache.beam.sdk.values.TypeDescriptor;
-import org.apache.hadoop.hbase.client.Result;
-
-/**
- * A {@link CoderProviderRegistrar} for standard types used with {@link HBaseIO}.
- */
-@AutoService(CoderProviderRegistrar.class)
-public class HBaseCoderProviderRegistrar implements CoderProviderRegistrar {
- @Override
- public List<CoderProvider> getCoderProviders() {
- return ImmutableList.of(
- HBaseMutationCoder.getCoderProvider(),
- CoderProviders.forCoder(TypeDescriptor.of(Result.class), HBaseResultCoder.of()));
- }
-}
[26/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/apex/src/test/java/org/apache/beam/runners/apex/translation/utils/ApexStateInternalsTest.java
----------------------------------------------------------------------
diff --git a/runners/apex/src/test/java/org/apache/beam/runners/apex/translation/utils/ApexStateInternalsTest.java b/runners/apex/src/test/java/org/apache/beam/runners/apex/translation/utils/ApexStateInternalsTest.java
index 87aa8c2..a7e64af 100644
--- a/runners/apex/src/test/java/org/apache/beam/runners/apex/translation/utils/ApexStateInternalsTest.java
+++ b/runners/apex/src/test/java/org/apache/beam/runners/apex/translation/utils/ApexStateInternalsTest.java
@@ -18,109 +18,350 @@
package org.apache.beam.runners.apex.translation.utils;
import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertThat;
import com.datatorrent.lib.util.KryoCloneUtils;
-import org.apache.beam.runners.core.StateInternals;
-import org.apache.beam.runners.core.StateInternalsTest;
+import java.util.Arrays;
+import org.apache.beam.runners.apex.translation.utils.ApexStateInternals.ApexStateBackend;
+import org.apache.beam.runners.apex.translation.utils.ApexStateInternals.ApexStateInternalsFactory;
+import org.apache.beam.runners.core.StateMerging;
import org.apache.beam.runners.core.StateNamespace;
import org.apache.beam.runners.core.StateNamespaceForTest;
import org.apache.beam.runners.core.StateTag;
import org.apache.beam.runners.core.StateTags;
import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.CombiningState;
+import org.apache.beam.sdk.state.GroupingState;
+import org.apache.beam.sdk.state.ReadableState;
import org.apache.beam.sdk.state.ValueState;
+import org.apache.beam.sdk.state.WatermarkHoldState;
+import org.apache.beam.sdk.transforms.Sum;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
+import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
import org.hamcrest.Matchers;
-import org.junit.Ignore;
+import org.joda.time.Instant;
+import org.junit.Before;
import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-import org.junit.runners.Suite;
/**
* Tests for {@link ApexStateInternals}. This is based on the tests for
- * {@code StateInternalsTest}.
+ * {@code InMemoryStateInternals}.
*/
-@RunWith(Suite.class)
-@Suite.SuiteClasses({
- ApexStateInternalsTest.StandardStateInternalsTests.class,
- ApexStateInternalsTest.OtherTests.class
-})
public class ApexStateInternalsTest {
+ private static final BoundedWindow WINDOW_1 = new IntervalWindow(new Instant(0), new Instant(10));
+ private static final StateNamespace NAMESPACE_1 = new StateNamespaceForTest("ns1");
+ private static final StateNamespace NAMESPACE_2 = new StateNamespaceForTest("ns2");
+ private static final StateNamespace NAMESPACE_3 = new StateNamespaceForTest("ns3");
- private static StateInternals newStateInternals() {
- return new ApexStateInternals.ApexStateBackend()
+ private static final StateTag<ValueState<String>> STRING_VALUE_ADDR =
+ StateTags.value("stringValue", StringUtf8Coder.of());
+ private static final StateTag<CombiningState<Integer, int[], Integer>>
+ SUM_INTEGER_ADDR = StateTags.combiningValueFromInputInternal(
+ "sumInteger", VarIntCoder.of(), Sum.ofIntegers());
+ private static final StateTag<BagState<String>> STRING_BAG_ADDR =
+ StateTags.bag("stringBag", StringUtf8Coder.of());
+ private static final StateTag<WatermarkHoldState>
+ WATERMARK_EARLIEST_ADDR =
+ StateTags.watermarkStateInternal("watermark", TimestampCombiner.EARLIEST);
+ private static final StateTag<WatermarkHoldState> WATERMARK_LATEST_ADDR =
+ StateTags.watermarkStateInternal("watermark", TimestampCombiner.LATEST);
+ private static final StateTag<WatermarkHoldState> WATERMARK_EOW_ADDR =
+ StateTags.watermarkStateInternal("watermark", TimestampCombiner.END_OF_WINDOW);
+
+ private ApexStateInternals<String> underTest;
+
+ @Before
+ public void initStateInternals() {
+ underTest = new ApexStateInternals.ApexStateBackend()
.newStateInternalsFactory(StringUtf8Coder.of())
- .stateInternalsForKey("dummyKey");
- }
-
- /**
- * A standard StateInternals test. Ignore set and map tests.
- */
- @RunWith(JUnit4.class)
- public static class StandardStateInternalsTests extends StateInternalsTest {
- @Override
- protected StateInternals createStateInternals() {
- return newStateInternals();
- }
-
- @Override
- @Ignore
- public void testSet() {}
-
- @Override
- @Ignore
- public void testSetIsEmpty() {}
-
- @Override
- @Ignore
- public void testMergeSetIntoSource() {}
-
- @Override
- @Ignore
- public void testMergeSetIntoNewNamespace() {}
-
- @Override
- @Ignore
- public void testMap() {}
-
- @Override
- @Ignore
- public void testSetReadable() {}
-
- @Override
- @Ignore
- public void testMapReadable() {}
- }
-
- /**
- * A specific test of ApexStateInternalsTest.
- */
- @RunWith(JUnit4.class)
- public static class OtherTests {
-
- private static final StateNamespace NAMESPACE = new StateNamespaceForTest("ns");
- private static final StateTag<ValueState<String>> STRING_VALUE_ADDR =
- StateTags.value("stringValue", StringUtf8Coder.of());
-
- @Test
- public void testSerialization() throws Exception {
- ApexStateInternals.ApexStateInternalsFactory<String> sif =
- new ApexStateInternals.ApexStateBackend().
- newStateInternalsFactory(StringUtf8Coder.of());
- ApexStateInternals<String> keyAndState = sif.stateInternalsForKey("dummy");
-
- ValueState<String> value = keyAndState.state(NAMESPACE, STRING_VALUE_ADDR);
- assertEquals(keyAndState.state(NAMESPACE, STRING_VALUE_ADDR), value);
- value.write("hello");
-
- ApexStateInternals.ApexStateInternalsFactory<String> cloned;
- assertNotNull("Serialization", cloned = KryoCloneUtils.cloneObject(sif));
- ApexStateInternals<String> clonedKeyAndState = cloned.stateInternalsForKey("dummy");
-
- ValueState<String> clonedValue = clonedKeyAndState.state(NAMESPACE, STRING_VALUE_ADDR);
- assertThat(clonedValue.read(), Matchers.equalTo("hello"));
- assertEquals(clonedKeyAndState.state(NAMESPACE, STRING_VALUE_ADDR), value);
- }
+ .stateInternalsForKey((String) null);
+ }
+
+ @Test
+ public void testBag() throws Exception {
+ BagState<String> value = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
+
+ assertEquals(value, underTest.state(NAMESPACE_1, STRING_BAG_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, STRING_BAG_ADDR)));
+
+ assertThat(value.read(), Matchers.emptyIterable());
+ value.add("hello");
+ assertThat(value.read(), Matchers.containsInAnyOrder("hello"));
+
+ value.add("world");
+ assertThat(value.read(), Matchers.containsInAnyOrder("hello", "world"));
+
+ value.clear();
+ assertThat(value.read(), Matchers.emptyIterable());
+ assertEquals(underTest.state(NAMESPACE_1, STRING_BAG_ADDR), value);
+
+ }
+
+ @Test
+ public void testBagIsEmpty() throws Exception {
+ BagState<String> value = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
+
+ assertThat(value.isEmpty().read(), Matchers.is(true));
+ ReadableState<Boolean> readFuture = value.isEmpty();
+ value.add("hello");
+ assertThat(readFuture.read(), Matchers.is(false));
+
+ value.clear();
+ assertThat(readFuture.read(), Matchers.is(true));
+ }
+
+ @Test
+ public void testMergeBagIntoSource() throws Exception {
+ BagState<String> bag1 = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
+ BagState<String> bag2 = underTest.state(NAMESPACE_2, STRING_BAG_ADDR);
+
+ bag1.add("Hello");
+ bag2.add("World");
+ bag1.add("!");
+
+ StateMerging.mergeBags(Arrays.asList(bag1, bag2), bag1);
+
+ // Reading the merged bag gets both the contents
+ assertThat(bag1.read(), Matchers.containsInAnyOrder("Hello", "World", "!"));
+ assertThat(bag2.read(), Matchers.emptyIterable());
+ }
+
+ @Test
+ public void testMergeBagIntoNewNamespace() throws Exception {
+ BagState<String> bag1 = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
+ BagState<String> bag2 = underTest.state(NAMESPACE_2, STRING_BAG_ADDR);
+ BagState<String> bag3 = underTest.state(NAMESPACE_3, STRING_BAG_ADDR);
+
+ bag1.add("Hello");
+ bag2.add("World");
+ bag1.add("!");
+
+ StateMerging.mergeBags(Arrays.asList(bag1, bag2, bag3), bag3);
+
+ // Reading the merged bag gets both the contents
+ assertThat(bag3.read(), Matchers.containsInAnyOrder("Hello", "World", "!"));
+ assertThat(bag1.read(), Matchers.emptyIterable());
+ assertThat(bag2.read(), Matchers.emptyIterable());
}
+
+ @Test
+ public void testCombiningValue() throws Exception {
+ GroupingState<Integer, Integer> value = underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
+
+ // State instances are cached, but depend on the namespace.
+ assertEquals(value, underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, SUM_INTEGER_ADDR)));
+
+ assertThat(value.read(), Matchers.equalTo(0));
+ value.add(2);
+ assertThat(value.read(), Matchers.equalTo(2));
+
+ value.add(3);
+ assertThat(value.read(), Matchers.equalTo(5));
+
+ value.clear();
+ assertThat(value.read(), Matchers.equalTo(0));
+ assertEquals(underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR), value);
+ }
+
+ @Test
+ public void testCombiningIsEmpty() throws Exception {
+ GroupingState<Integer, Integer> value = underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
+
+ assertThat(value.isEmpty().read(), Matchers.is(true));
+ ReadableState<Boolean> readFuture = value.isEmpty();
+ value.add(5);
+ assertThat(readFuture.read(), Matchers.is(false));
+
+ value.clear();
+ assertThat(readFuture.read(), Matchers.is(true));
+ }
+
+ @Test
+ public void testMergeCombiningValueIntoSource() throws Exception {
+ CombiningState<Integer, int[], Integer> value1 =
+ underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
+ CombiningState<Integer, int[], Integer> value2 =
+ underTest.state(NAMESPACE_2, SUM_INTEGER_ADDR);
+
+ value1.add(5);
+ value2.add(10);
+ value1.add(6);
+
+ assertThat(value1.read(), Matchers.equalTo(11));
+ assertThat(value2.read(), Matchers.equalTo(10));
+
+ // Merging clears the old values and updates the result value.
+ StateMerging.mergeCombiningValues(Arrays.asList(value1, value2), value1);
+
+ assertThat(value1.read(), Matchers.equalTo(21));
+ assertThat(value2.read(), Matchers.equalTo(0));
+ }
+
+ @Test
+ public void testMergeCombiningValueIntoNewNamespace() throws Exception {
+ CombiningState<Integer, int[], Integer> value1 =
+ underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
+ CombiningState<Integer, int[], Integer> value2 =
+ underTest.state(NAMESPACE_2, SUM_INTEGER_ADDR);
+ CombiningState<Integer, int[], Integer> value3 =
+ underTest.state(NAMESPACE_3, SUM_INTEGER_ADDR);
+
+ value1.add(5);
+ value2.add(10);
+ value1.add(6);
+
+ StateMerging.mergeCombiningValues(Arrays.asList(value1, value2), value3);
+
+ // Merging clears the old values and updates the result value.
+ assertThat(value1.read(), Matchers.equalTo(0));
+ assertThat(value2.read(), Matchers.equalTo(0));
+ assertThat(value3.read(), Matchers.equalTo(21));
+ }
+
+ @Test
+ public void testWatermarkEarliestState() throws Exception {
+ WatermarkHoldState value =
+ underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR);
+
+ // State instances are cached, but depend on the namespace.
+ assertEquals(value, underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, WATERMARK_EARLIEST_ADDR)));
+
+ assertThat(value.read(), Matchers.nullValue());
+ value.add(new Instant(2000));
+ assertThat(value.read(), Matchers.equalTo(new Instant(2000)));
+
+ value.add(new Instant(3000));
+ assertThat(value.read(), Matchers.equalTo(new Instant(2000)));
+
+ value.add(new Instant(1000));
+ assertThat(value.read(), Matchers.equalTo(new Instant(1000)));
+
+ value.clear();
+ assertThat(value.read(), Matchers.equalTo(null));
+ assertEquals(underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR), value);
+ }
+
+ @Test
+ public void testWatermarkLatestState() throws Exception {
+ WatermarkHoldState value =
+ underTest.state(NAMESPACE_1, WATERMARK_LATEST_ADDR);
+
+ // State instances are cached, but depend on the namespace.
+ assertEquals(value, underTest.state(NAMESPACE_1, WATERMARK_LATEST_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, WATERMARK_LATEST_ADDR)));
+
+ assertThat(value.read(), Matchers.nullValue());
+ value.add(new Instant(2000));
+ assertThat(value.read(), Matchers.equalTo(new Instant(2000)));
+
+ value.add(new Instant(3000));
+ assertThat(value.read(), Matchers.equalTo(new Instant(3000)));
+
+ value.add(new Instant(1000));
+ assertThat(value.read(), Matchers.equalTo(new Instant(3000)));
+
+ value.clear();
+ assertThat(value.read(), Matchers.equalTo(null));
+ assertEquals(underTest.state(NAMESPACE_1, WATERMARK_LATEST_ADDR), value);
+ }
+
+ @Test
+ public void testWatermarkEndOfWindowState() throws Exception {
+ WatermarkHoldState value = underTest.state(NAMESPACE_1, WATERMARK_EOW_ADDR);
+
+ // State instances are cached, but depend on the namespace.
+ assertEquals(value, underTest.state(NAMESPACE_1, WATERMARK_EOW_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, WATERMARK_EOW_ADDR)));
+
+ assertThat(value.read(), Matchers.nullValue());
+ value.add(new Instant(2000));
+ assertThat(value.read(), Matchers.equalTo(new Instant(2000)));
+
+ value.clear();
+ assertThat(value.read(), Matchers.equalTo(null));
+ assertEquals(underTest.state(NAMESPACE_1, WATERMARK_EOW_ADDR), value);
+ }
+
+ @Test
+ public void testWatermarkStateIsEmpty() throws Exception {
+ WatermarkHoldState value =
+ underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR);
+
+ assertThat(value.isEmpty().read(), Matchers.is(true));
+ ReadableState<Boolean> readFuture = value.isEmpty();
+ value.add(new Instant(1000));
+ assertThat(readFuture.read(), Matchers.is(false));
+
+ value.clear();
+ assertThat(readFuture.read(), Matchers.is(true));
+ }
+
+ @Test
+ public void testMergeEarliestWatermarkIntoSource() throws Exception {
+ WatermarkHoldState value1 =
+ underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR);
+ WatermarkHoldState value2 =
+ underTest.state(NAMESPACE_2, WATERMARK_EARLIEST_ADDR);
+
+ value1.add(new Instant(3000));
+ value2.add(new Instant(5000));
+ value1.add(new Instant(4000));
+ value2.add(new Instant(2000));
+
+ // Merging clears the old values and updates the merged value.
+ StateMerging.mergeWatermarks(Arrays.asList(value1, value2), value1, WINDOW_1);
+
+ assertThat(value1.read(), Matchers.equalTo(new Instant(2000)));
+ assertThat(value2.read(), Matchers.equalTo(null));
+ }
+
+ @Test
+ public void testMergeLatestWatermarkIntoSource() throws Exception {
+ WatermarkHoldState value1 =
+ underTest.state(NAMESPACE_1, WATERMARK_LATEST_ADDR);
+ WatermarkHoldState value2 =
+ underTest.state(NAMESPACE_2, WATERMARK_LATEST_ADDR);
+ WatermarkHoldState value3 =
+ underTest.state(NAMESPACE_3, WATERMARK_LATEST_ADDR);
+
+ value1.add(new Instant(3000));
+ value2.add(new Instant(5000));
+ value1.add(new Instant(4000));
+ value2.add(new Instant(2000));
+
+ // Merging clears the old values and updates the result value.
+ StateMerging.mergeWatermarks(Arrays.asList(value1, value2), value3, WINDOW_1);
+
+ // Merging clears the old values and updates the result value.
+ assertThat(value3.read(), Matchers.equalTo(new Instant(5000)));
+ assertThat(value1.read(), Matchers.equalTo(null));
+ assertThat(value2.read(), Matchers.equalTo(null));
+ }
+
+ @Test
+ public void testSerialization() throws Exception {
+ ApexStateInternalsFactory<String> sif = new ApexStateBackend().
+ newStateInternalsFactory(StringUtf8Coder.of());
+ ApexStateInternals<String> keyAndState = sif.stateInternalsForKey("dummy");
+
+ ValueState<String> value = keyAndState.state(NAMESPACE_1, STRING_VALUE_ADDR);
+ assertEquals(keyAndState.state(NAMESPACE_1, STRING_VALUE_ADDR), value);
+ value.write("hello");
+
+ ApexStateInternalsFactory<String> cloned;
+ assertNotNull("Serialization", cloned = KryoCloneUtils.cloneObject(sif));
+ ApexStateInternals<String> clonedKeyAndState = cloned.stateInternalsForKey("dummy");
+
+ ValueState<String> clonedValue = clonedKeyAndState.state(NAMESPACE_1, STRING_VALUE_ADDR);
+ assertThat(clonedValue.read(), Matchers.equalTo("hello"));
+ assertEquals(clonedKeyAndState.state(NAMESPACE_1, STRING_VALUE_ADDR), value);
+ }
+
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-construction-java/pom.xml
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/pom.xml b/runners/core-construction-java/pom.xml
index b85b5f5..67951e9 100644
--- a/runners/core-construction-java/pom.xml
+++ b/runners/core-construction-java/pom.xml
@@ -24,7 +24,7 @@
<parent>
<artifactId>beam-runners-parent</artifactId>
<groupId>org.apache.beam</groupId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CreatePCollectionViewTranslation.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CreatePCollectionViewTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CreatePCollectionViewTranslation.java
index 8fc99b9..aa24909 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CreatePCollectionViewTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/CreatePCollectionViewTranslation.java
@@ -56,8 +56,8 @@ public class CreatePCollectionViewTranslation {
@Deprecated
public static <ElemT, ViewT> PCollectionView<ViewT> getView(
AppliedPTransform<
- PCollection<ElemT>, PCollection<ElemT>,
- PTransform<PCollection<ElemT>, PCollection<ElemT>>>
+ PCollection<ElemT>, PCollectionView<ViewT>,
+ PTransform<PCollection<ElemT>, PCollectionView<ViewT>>>
application)
throws IOException {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/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/c1b2b96a/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/c1b2b96a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PCollectionTranslation.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PCollectionTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PCollectionTranslation.java
index 52526bb..968966f 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PCollectionTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PCollectionTranslation.java
@@ -20,7 +20,6 @@ package org.apache.beam.runners.core.construction;
import com.google.protobuf.InvalidProtocolBufferException;
import java.io.IOException;
-import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.common.runner.v1.RunnerApi;
import org.apache.beam.sdk.values.PCollection;
@@ -48,21 +47,6 @@ public class PCollectionTranslation {
.build();
}
- public static PCollection<?> fromProto(
- Pipeline pipeline, RunnerApi.PCollection pCollection, RunnerApi.Components components)
- throws IOException {
- return PCollection.createPrimitiveOutputInternal(
- pipeline,
- WindowingStrategyTranslation.fromProto(
- components.getWindowingStrategiesOrThrow(pCollection.getWindowingStrategyId()),
- components),
- fromProto(pCollection.getIsBounded()))
- .setCoder(
- (Coder)
- CoderTranslation.fromProto(
- components.getCodersOrThrow(pCollection.getCoderId()), components));
- }
-
public static IsBounded isBounded(RunnerApi.PCollection pCollection) {
return fromProto(pCollection.getIsBounded());
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformMatchers.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformMatchers.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformMatchers.java
index 0d27241..bfe24a0 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformMatchers.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformMatchers.java
@@ -17,14 +17,12 @@
*/
package org.apache.beam.runners.core.construction;
-import static org.apache.beam.runners.core.construction.PTransformTranslation.WRITE_FILES_TRANSFORM_URN;
-
import com.google.common.base.MoreObjects;
-import java.io.IOException;
import java.util.HashSet;
import java.util.Set;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.io.WriteFiles;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.runners.PTransformMatcher;
import org.apache.beam.sdk.transforms.DoFn;
@@ -52,34 +50,6 @@ public class PTransformMatchers {
private PTransformMatchers() {}
/**
- * Returns a {@link PTransformMatcher} that matches a {@link PTransform} if the URN of the
- * {@link PTransform} is equal to the URN provided ot this matcher.
- */
- public static PTransformMatcher urnEqualTo(String urn) {
- return new EqualUrnPTransformMatcher(urn);
- }
-
- private static class EqualUrnPTransformMatcher implements PTransformMatcher {
- private final String urn;
-
- private EqualUrnPTransformMatcher(String urn) {
- this.urn = urn;
- }
-
- @Override
- public boolean matches(AppliedPTransform<?, ?, ?> application) {
- return urn.equals(PTransformTranslation.urnForTransformOrNull(application.getTransform()));
- }
-
- @Override
- public String toString() {
- return MoreObjects.toStringHelper(this)
- .add("urn", urn)
- .toString();
- }
- }
-
- /**
* Returns a {@link PTransformMatcher} that matches a {@link PTransform} if the class of the
* {@link PTransform} is equal to the {@link Class} provided ot this matcher.
*/
@@ -181,68 +151,6 @@ public class PTransformMatchers {
}
/**
- * A {@link PTransformMatcher} that matches a {@link ParDo} by URN if it has a splittable {@link
- * DoFn}.
- */
- public static PTransformMatcher splittableParDo() {
- return new PTransformMatcher() {
- @Override
- public boolean matches(AppliedPTransform<?, ?, ?> application) {
- if (PTransformTranslation.PAR_DO_TRANSFORM_URN.equals(
- PTransformTranslation.urnForTransformOrNull(application.getTransform()))) {
-
- try {
- return ParDoTranslation.isSplittable(application);
- } catch (IOException e) {
- throw new RuntimeException(
- String.format(
- "Transform with URN %s could not be translated",
- PTransformTranslation.PAR_DO_TRANSFORM_URN),
- e);
- }
- }
- return false;
- }
-
- @Override
- public String toString() {
- return MoreObjects.toStringHelper("SplittableParDoMultiMatcher").toString();
- }
- };
- }
-
- /**
- * A {@link PTransformMatcher} that matches a {@link ParDo} transform by URN
- * and whether it contains state or timers as specified by {@link ParDoTranslation}.
- */
- public static PTransformMatcher stateOrTimerParDo() {
- return new PTransformMatcher() {
- @Override
- public boolean matches(AppliedPTransform<?, ?, ?> application) {
- if (PTransformTranslation.PAR_DO_TRANSFORM_URN.equals(
- PTransformTranslation.urnForTransformOrNull(application.getTransform()))) {
-
- try {
- return ParDoTranslation.usesStateOrTimers(application);
- } catch (IOException e) {
- throw new RuntimeException(
- String.format(
- "Transform with URN %s could not be translated",
- PTransformTranslation.PAR_DO_TRANSFORM_URN),
- e);
- }
- }
- return false;
- }
-
- @Override
- public String toString() {
- return MoreObjects.toStringHelper("StateOrTimerParDoMatcher").toString();
- }
- };
- }
-
- /**
* A {@link PTransformMatcher} that matches a {@link ParDo.MultiOutput} containing a {@link DoFn}
* that uses state or timers, as specified by {@link DoFnSignature#usesState()} and
* {@link DoFnSignature#usesTimers()}.
@@ -360,18 +268,9 @@ public class PTransformMatchers {
return new PTransformMatcher() {
@Override
public boolean matches(AppliedPTransform<?, ?, ?> application) {
- if (WRITE_FILES_TRANSFORM_URN.equals(
- PTransformTranslation.urnForTransformOrNull(application.getTransform()))) {
- try {
- return WriteFilesTranslation.isRunnerDeterminedSharding(
- (AppliedPTransform) application);
- } catch (IOException exc) {
- throw new RuntimeException(
- String.format(
- "Transform with URN %s failed to parse: %s",
- WRITE_FILES_TRANSFORM_URN, application.getTransform()),
- exc);
- }
+ if (application.getTransform() instanceof WriteFiles) {
+ WriteFiles write = (WriteFiles) application.getTransform();
+ return write.getSharding() == null && write.getNumShards() == null;
}
return false;
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java
index bae7b05..32ecf43 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java
@@ -179,12 +179,13 @@ public class PTransformTranslation {
* Returns the URN for the transform if it is known, otherwise throws.
*/
public static String urnForTransform(PTransform<?, ?> transform) {
- String urn = urnForTransformOrNull(transform);
- if (urn == null) {
+ TransformPayloadTranslator translator = KNOWN_PAYLOAD_TRANSLATORS.get(transform.getClass());
+ if (translator == null) {
throw new IllegalStateException(
String.format("No translator known for %s", transform.getClass().getName()));
}
- return urn;
+
+ return translator.getUrn(transform);
}
/**
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java
index 03f29ff..34e0d86 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java
@@ -19,7 +19,6 @@
package org.apache.beam.runners.core.construction;
import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;
import static org.apache.beam.runners.core.construction.PTransformTranslation.PAR_DO_TRANSFORM_URN;
@@ -35,11 +34,9 @@ import com.google.protobuf.BytesValue;
import com.google.protobuf.InvalidProtocolBufferException;
import java.io.IOException;
import java.io.Serializable;
-import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
-import java.util.Set;
import org.apache.beam.runners.core.construction.PTransformTranslation.TransformPayloadTranslator;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.IterableCoder;
@@ -75,10 +72,8 @@ import org.apache.beam.sdk.transforms.windowing.WindowMappingFn;
import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.util.WindowedValue.FullWindowedValueCoder;
-import org.apache.beam.sdk.values.PCollection;
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;
/**
@@ -220,67 +215,11 @@ public class ParDoTranslation {
return doFnAndMainOutputTagFromProto(payload.getDoFn()).getDoFn();
}
- public static DoFn<?, ?> getDoFn(AppliedPTransform<?, ?, ?> application) throws IOException {
- return getDoFn(getParDoPayload(application));
- }
-
public static TupleTag<?> getMainOutputTag(ParDoPayload payload)
throws InvalidProtocolBufferException {
return doFnAndMainOutputTagFromProto(payload.getDoFn()).getMainOutputTag();
}
- public static TupleTag<?> getMainOutputTag(AppliedPTransform<?, ?, ?> application)
- throws IOException {
- return getMainOutputTag(getParDoPayload(application));
- }
-
- public static TupleTagList getAdditionalOutputTags(AppliedPTransform<?, ?, ?> application)
- throws IOException {
-
- RunnerApi.PTransform protoTransform =
- PTransformTranslation.toProto(application, SdkComponents.create());
-
- ParDoPayload payload = protoTransform.getSpec().getParameter().unpack(ParDoPayload.class);
- TupleTag<?> mainOutputTag = getMainOutputTag(payload);
- Set<String> outputTags =
- Sets.difference(
- protoTransform.getOutputsMap().keySet(), Collections.singleton(mainOutputTag.getId()));
-
- ArrayList<TupleTag<?>> additionalOutputTags = new ArrayList<>();
- for (String outputTag : outputTags) {
- additionalOutputTags.add(new TupleTag<>(outputTag));
- }
- return TupleTagList.of(additionalOutputTags);
- }
-
- public static List<PCollectionView<?>> getSideInputs(AppliedPTransform<?, ?, ?> application)
- throws IOException {
-
- SdkComponents sdkComponents = SdkComponents.create();
- RunnerApi.PTransform parDoProto =
- PTransformTranslation.toProto(application, sdkComponents);
- ParDoPayload payload = parDoProto.getSpec().getParameter().unpack(ParDoPayload.class);
-
- List<PCollectionView<?>> views = new ArrayList<>();
- for (Map.Entry<String, SideInput> sideInputEntry : payload.getSideInputsMap().entrySet()) {
- String sideInputTag = sideInputEntry.getKey();
- RunnerApi.SideInput sideInput = sideInputEntry.getValue();
- PCollection<?> originalPCollection =
- checkNotNull(
- (PCollection<?>) application.getInputs().get(new TupleTag<>(sideInputTag)),
- "no input with tag %s",
- sideInputTag);
- views.add(
- viewFromProto(
- sideInput,
- sideInputTag,
- originalPCollection,
- parDoProto,
- sdkComponents.toComponents()));
- }
- return views;
- }
-
public static RunnerApi.PCollection getMainInput(
RunnerApi.PTransform ptransform, Components components) throws IOException {
checkArgument(
@@ -508,27 +447,15 @@ public class ParDoTranslation {
return builder.build();
}
- /**
- * Create a {@link PCollectionView} from a side input spec and an already-deserialized {@link
- * PCollection} that should be wired up.
- */
- public static PCollectionView<?> viewFromProto(
- SideInput sideInput,
- String localName,
- PCollection<?> pCollection,
- RunnerApi.PTransform parDoTransform,
- Components components)
+ public static PCollectionView<?> fromProto(
+ SideInput sideInput, String id, RunnerApi.PTransform parDoTransform, Components components)
throws IOException {
- checkArgument(
- localName != null,
- "%s.viewFromProto: localName must not be null",
- ParDoTranslation.class.getSimpleName());
- TupleTag<?> tag = new TupleTag<>(localName);
+ TupleTag<?> tag = new TupleTag<>(id);
WindowMappingFn<?> windowMappingFn = windowMappingFnFromProto(sideInput.getWindowMappingFn());
ViewFn<?, ?> viewFn = viewFnFromProto(sideInput.getViewFn());
RunnerApi.PCollection inputCollection =
- components.getPcollectionsOrThrow(parDoTransform.getInputsOrThrow(localName));
+ components.getPcollectionsOrThrow(parDoTransform.getInputsOrThrow(id));
WindowingStrategy<?, ?> windowingStrategy =
WindowingStrategyTranslation.fromProto(
components.getWindowingStrategiesOrThrow(inputCollection.getWindowingStrategyId()),
@@ -548,7 +475,6 @@ public class ParDoTranslation {
PCollectionView<?> view =
new RunnerPCollectionView<>(
- pCollection,
(TupleTag<Iterable<WindowedValue<?>>>) tag,
(ViewFn<Iterable<WindowedValue<?>>, ?>) viewFn,
windowMappingFn,
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/RunnerPCollectionView.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/RunnerPCollectionView.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/RunnerPCollectionView.java
index 85139e8..89e8784 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/RunnerPCollectionView.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/RunnerPCollectionView.java
@@ -18,8 +18,6 @@
package org.apache.beam.runners.core.construction;
-import java.util.Map;
-import java.util.Objects;
import javax.annotation.Nullable;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.common.runner.v1.RunnerApi.SideInput;
@@ -28,7 +26,6 @@ import org.apache.beam.sdk.transforms.windowing.WindowMappingFn;
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.PValue;
import org.apache.beam.sdk.values.PValueBase;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
@@ -40,19 +37,16 @@ class RunnerPCollectionView<T> extends PValueBase implements PCollectionView<T>
private final WindowMappingFn<?> windowMappingFn;
private final WindowingStrategy<?, ?> windowingStrategy;
private final Coder<Iterable<WindowedValue<?>>> coder;
- private final transient PCollection<?> pCollection;
/**
* Create a new {@link RunnerPCollectionView} from the provided components.
*/
RunnerPCollectionView(
- PCollection<?> pCollection,
TupleTag<Iterable<WindowedValue<?>>> tag,
ViewFn<Iterable<WindowedValue<?>>, T> viewFn,
WindowMappingFn<?> windowMappingFn,
@Nullable WindowingStrategy<?, ?> windowingStrategy,
@Nullable Coder<Iterable<WindowedValue<?>>> coder) {
- this.pCollection = pCollection;
this.tag = tag;
this.viewFn = viewFn;
this.windowMappingFn = windowMappingFn;
@@ -60,9 +54,11 @@ class RunnerPCollectionView<T> extends PValueBase implements PCollectionView<T>
this.coder = coder;
}
+ @Nullable
@Override
public PCollection<?> getPCollection() {
- return pCollection;
+ throw new IllegalStateException(
+ String.format("Cannot call getPCollection on a %s", getClass().getSimpleName()));
}
@Override
@@ -89,25 +85,4 @@ class RunnerPCollectionView<T> extends PValueBase implements PCollectionView<T>
public Coder<Iterable<WindowedValue<?>>> getCoderInternal() {
return coder;
}
-
- @Override
- public Map<TupleTag<?>, PValue> expand() {
- throw new UnsupportedOperationException(String.format(
- "A %s cannot be expanded", RunnerPCollectionView.class.getSimpleName()));
- }
-
- @Override
- public boolean equals(Object other) {
- if (!(other instanceof PCollectionView)) {
- return false;
- }
- @SuppressWarnings("unchecked")
- PCollectionView<?> otherView = (PCollectionView<?>) other;
- return tag.equals(otherView.getTagInternal());
- }
-
- @Override
- public int hashCode() {
- return Objects.hash(tag);
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/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
index e71187b..665e39d 100644
--- 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
@@ -18,16 +18,13 @@
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.io.IOException;
import java.util.List;
-import java.util.Map;
import java.util.UUID;
import org.apache.beam.runners.core.construction.PTransformTranslation.RawPTransform;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.KvCoder;
-import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
@@ -42,8 +39,6 @@ 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.PCollectionViews;
-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;
@@ -59,7 +54,7 @@ import org.apache.beam.sdk.values.WindowingStrategy;
* <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.
+ * {@link DoFn}'s {@link DoFn.ProcessElement} method.
* </ol>
*
* <p>This transform is intended as a helper for internal use by runners when implementing {@code
@@ -68,11 +63,7 @@ import org.apache.beam.sdk.values.WindowingStrategy;
@Experimental(Experimental.Kind.SPLITTABLE_DO_FN)
public class SplittableParDo<InputT, OutputT, RestrictionT>
extends PTransform<PCollection<InputT>, PCollectionTuple> {
-
- private final DoFn<InputT, OutputT> doFn;
- private final List<PCollectionView<?>> sideInputs;
- private final TupleTag<OutputT> mainOutputTag;
- private final TupleTagList additionalOutputTags;
+ private final ParDo.MultiOutput<InputT, OutputT> parDo;
public static final String SPLITTABLE_PROCESS_URN =
"urn:beam:runners_core:transforms:splittable_process:v1";
@@ -83,97 +74,56 @@ public class SplittableParDo<InputT, OutputT, RestrictionT>
public static final String SPLITTABLE_GBKIKWI_URN =
"urn:beam:runners_core:transforms:splittable_gbkikwi:v1";
- private SplittableParDo(
- DoFn<InputT, OutputT> doFn,
- TupleTag<OutputT> mainOutputTag,
- List<PCollectionView<?>> sideInputs,
- TupleTagList additionalOutputTags) {
- checkArgument(
- DoFnSignatures.getSignature(doFn.getClass()).processElement().isSplittable(),
- "fn must be a splittable DoFn");
- this.doFn = doFn;
- this.mainOutputTag = mainOutputTag;
- this.sideInputs = sideInputs;
- this.additionalOutputTags = additionalOutputTags;
- }
-
/**
- * Creates a {@link SplittableParDo} from an original Java {@link ParDo}.
+ * Creates the transform for the given original multi-output {@link ParDo}.
*
* @param parDo The splittable {@link ParDo} transform.
*/
- public static <InputT, OutputT> SplittableParDo<InputT, OutputT, ?> forJavaParDo(
- ParDo.MultiOutput<InputT, OutputT> parDo) {
- checkArgument(parDo != null, "parDo must not be null");
+ 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");
- return new SplittableParDo(
- parDo.getFn(),
- parDo.getMainOutputTag(),
- parDo.getSideInputs(),
- parDo.getAdditionalOutputTags());
- }
-
- /**
- * Creates the transform for a {@link ParDo}-compatible {@link AppliedPTransform}.
- *
- * <p>The input may generally be a deserialized transform so it may not actually be a {@link
- * ParDo}. Instead {@link ParDoTranslation} will be used to extract fields.
- */
- public static SplittableParDo<?, ?, ?> forAppliedParDo(AppliedPTransform<?, ?, ?> parDo) {
- checkArgument(parDo != null, "parDo must not be null");
-
- try {
- return new SplittableParDo<>(
- ParDoTranslation.getDoFn(parDo),
- (TupleTag) ParDoTranslation.getMainOutputTag(parDo),
- ParDoTranslation.getSideInputs(parDo),
- ParDoTranslation.getAdditionalOutputTags(parDo));
- } catch (IOException exc) {
- throw new RuntimeException(exc);
- }
}
@Override
public PCollectionTuple expand(PCollection<InputT> input) {
+ DoFn<InputT, OutputT> fn = parDo.getFn();
Coder<RestrictionT> restrictionCoder =
- DoFnInvokers.invokerFor(doFn)
+ DoFnInvokers.invokerFor(fn)
.invokeGetRestrictionCoder(input.getPipeline().getCoderRegistry());
- Coder<KV<InputT, RestrictionT>> splitCoder = KvCoder.of(input.getCoder(), restrictionCoder);
+ Coder<ElementAndRestriction<InputT, RestrictionT>> splitCoder =
+ ElementAndRestrictionCoder.of(input.getCoder(), restrictionCoder);
- PCollection<KV<String, KV<InputT, RestrictionT>>> keyedRestrictions =
+ PCollection<KV<String, ElementAndRestriction<InputT, RestrictionT>>> keyedRestrictions =
input
.apply(
"Pair with initial restriction",
- ParDo.of(new PairWithRestrictionFn<InputT, OutputT, RestrictionT>(doFn)))
+ ParDo.of(new PairWithRestrictionFn<InputT, OutputT, RestrictionT>(fn)))
.setCoder(splitCoder)
- .apply(
- "Split restriction", ParDo.of(new SplitRestrictionFn<InputT, RestrictionT>(doFn)))
+ .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<KV<InputT, RestrictionT>>()))
+ .apply(
+ "Explode windows",
+ ParDo.of(new ExplodeWindowsFn<ElementAndRestriction<InputT, RestrictionT>>()))
.apply(
"Assign unique key",
- WithKeys.of(new RandomUniqueKeyFn<KV<InputT, RestrictionT>>()));
+ WithKeys.of(new RandomUniqueKeyFn<ElementAndRestriction<InputT, RestrictionT>>()));
return keyedRestrictions.apply(
"ProcessKeyedElements",
new ProcessKeyedElements<>(
- doFn,
+ fn,
input.getCoder(),
restrictionCoder,
(WindowingStrategy<InputT, ?>) input.getWindowingStrategy(),
- sideInputs,
- mainOutputTag,
- additionalOutputTags));
- }
-
- @Override
- public Map<TupleTag<?>, PValue> getAdditionalInputs() {
- return PCollectionViews.toAdditionalInputs(sideInputs);
+ parDo.getSideInputs(),
+ parDo.getMainOutputTag(),
+ parDo.getAdditionalOutputTags()));
}
/**
@@ -190,11 +140,12 @@ public class SplittableParDo<InputT, OutputT, RestrictionT>
/**
* Runner-specific primitive {@link PTransform} that invokes the {@link DoFn.ProcessElement}
- * method for a splittable {@link DoFn} on each {@link KV} of the input {@link PCollection} of
- * {@link KV KVs} keyed with arbitrary but globally unique keys.
+ * 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 RawPTransform<PCollection<KV<String, KV<InputT, RestrictionT>>>, PCollectionTuple> {
+ extends RawPTransform<
+ PCollection<KV<String, ElementAndRestriction<InputT, RestrictionT>>>, PCollectionTuple> {
private final DoFn<InputT, OutputT> fn;
private final Coder<InputT> elementCoder;
private final Coder<RestrictionT> restrictionCoder;
@@ -257,7 +208,9 @@ public class SplittableParDo<InputT, OutputT, RestrictionT>
}
@Override
- public PCollectionTuple expand(PCollection<KV<String, KV<InputT, RestrictionT>>> input) {
+ public PCollectionTuple expand(
+ PCollection<KV<String, ElementAndRestriction<InputT, RestrictionT>>>
+ input) {
return createPrimitiveOutputFor(
input, fn, mainOutputTag, additionalOutputTags, windowingStrategy);
}
@@ -283,11 +236,6 @@ public class SplittableParDo<InputT, OutputT, RestrictionT>
}
@Override
- public Map<TupleTag<?>, PValue> getAdditionalInputs() {
- return PCollectionViews.toAdditionalInputs(sideInputs);
- }
-
- @Override
public String getUrn() {
return SPLITTABLE_PROCESS_KEYED_ELEMENTS_URN;
}
@@ -309,7 +257,7 @@ public class SplittableParDo<InputT, OutputT, RestrictionT>
* Pairs each input element with its initial restriction using the given splittable {@link DoFn}.
*/
private static class PairWithRestrictionFn<InputT, OutputT, RestrictionT>
- extends DoFn<InputT, KV<InputT, RestrictionT>> {
+ extends DoFn<InputT, ElementAndRestriction<InputT, RestrictionT>> {
private DoFn<InputT, OutputT> fn;
private transient DoFnInvoker<InputT, OutputT> invoker;
@@ -325,7 +273,7 @@ public class SplittableParDo<InputT, OutputT, RestrictionT>
@ProcessElement
public void processElement(ProcessContext context) {
context.output(
- KV.of(
+ ElementAndRestriction.of(
context.element(),
invoker.<RestrictionT>invokeGetInitialRestriction(context.element())));
}
@@ -333,7 +281,9 @@ public class SplittableParDo<InputT, OutputT, RestrictionT>
/** Splits the restriction using the given {@link SplitRestriction} method. */
private static class SplitRestrictionFn<InputT, RestrictionT>
- extends DoFn<KV<InputT, RestrictionT>, KV<InputT, RestrictionT>> {
+ extends DoFn<
+ ElementAndRestriction<InputT, RestrictionT>,
+ ElementAndRestriction<InputT, RestrictionT>> {
private final DoFn<InputT, ?> splittableFn;
private transient DoFnInvoker<InputT, ?> invoker;
@@ -348,14 +298,14 @@ public class SplittableParDo<InputT, OutputT, RestrictionT>
@ProcessElement
public void processElement(final ProcessContext c) {
- final InputT element = c.element().getKey();
+ final InputT element = c.element().element();
invoker.invokeSplitRestriction(
element,
- c.element().getValue(),
+ c.element().restriction(),
new OutputReceiver<RestrictionT>() {
@Override
public void output(RestrictionT part) {
- c.output(KV.of(element, part));
+ c.output(ElementAndRestriction.of(element, part));
}
});
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/TestStreamTranslation.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/TestStreamTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/TestStreamTranslation.java
index 515de57..90e6304 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/TestStreamTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/TestStreamTranslation.java
@@ -18,9 +18,6 @@
package org.apache.beam.runners.core.construction;
-import static com.google.common.base.Preconditions.checkArgument;
-import static org.apache.beam.runners.core.construction.PTransformTranslation.TEST_STREAM_TRANSFORM_URN;
-
import com.google.auto.service.AutoService;
import com.google.protobuf.Any;
import com.google.protobuf.ByteString;
@@ -36,8 +33,6 @@ import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.testing.TestStream;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.util.CoderUtils;
-import org.apache.beam.sdk.values.PBegin;
-import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.TimestampedValue;
import org.joda.time.Duration;
import org.joda.time.Instant;
@@ -62,48 +57,6 @@ public class TestStreamTranslation {
return builder.build();
}
- private static TestStream<?> fromProto(
- RunnerApi.TestStreamPayload testStreamPayload, RunnerApi.Components components)
- throws IOException {
-
- Coder<Object> coder =
- (Coder<Object>)
- CoderTranslation.fromProto(
- components.getCodersOrThrow(testStreamPayload.getCoderId()), components);
-
- List<TestStream.Event<Object>> events = new ArrayList<>();
-
- for (RunnerApi.TestStreamPayload.Event event : testStreamPayload.getEventsList()) {
- events.add(fromProto(event, coder));
- }
- return TestStream.fromRawEvents(coder, events);
- }
-
- /**
- * Converts an {@link AppliedPTransform}, which may be a rehydrated transform or an original
- * {@link TestStream}, to a {@link TestStream}.
- */
- public static <T> TestStream<T> getTestStream(
- AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>> application)
- throws IOException {
- // For robustness, we don't take this shortcut:
- // if (application.getTransform() instanceof TestStream) {
- // return application.getTransform()
- // }
-
- SdkComponents sdkComponents = SdkComponents.create();
- RunnerApi.PTransform transformProto = PTransformTranslation.toProto(application, sdkComponents);
- checkArgument(
- TEST_STREAM_TRANSFORM_URN.equals(transformProto.getSpec().getUrn()),
- "Attempt to get %s from a transform with wrong URN %s",
- TestStream.class.getSimpleName(),
- transformProto.getSpec().getUrn());
- RunnerApi.TestStreamPayload testStreamPayload =
- transformProto.getSpec().getParameter().unpack(RunnerApi.TestStreamPayload.class);
-
- return (TestStream<T>) fromProto(testStreamPayload, sdkComponents.toComponents());
- }
-
static <T> RunnerApi.TestStreamPayload.Event toProto(TestStream.Event<T> event, Coder<T> coder)
throws IOException {
switch (event.getType()) {
@@ -177,7 +130,7 @@ public class TestStreamTranslation {
static class TestStreamTranslator implements TransformPayloadTranslator<TestStream<?>> {
@Override
public String getUrn(TestStream<?> transform) {
- return TEST_STREAM_TRANSFORM_URN;
+ return PTransformTranslation.TEST_STREAM_TRANSFORM_URN;
}
@Override
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/TransformInputs.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/TransformInputs.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/TransformInputs.java
deleted file mode 100644
index 2baf93a..0000000
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/TransformInputs.java
+++ /dev/null
@@ -1,50 +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.construction;
-
-import static com.google.common.base.Preconditions.checkArgument;
-
-import com.google.common.collect.ImmutableList;
-import java.util.Collection;
-import java.util.Map;
-import org.apache.beam.sdk.runners.AppliedPTransform;
-import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.values.PValue;
-import org.apache.beam.sdk.values.TupleTag;
-
-/** Utilities for extracting subsets of inputs from an {@link AppliedPTransform}. */
-public class TransformInputs {
- /**
- * Gets all inputs of the {@link AppliedPTransform} that are not returned by {@link
- * PTransform#getAdditionalInputs()}.
- */
- public static Collection<PValue> nonAdditionalInputs(AppliedPTransform<?, ?, ?> application) {
- ImmutableList.Builder<PValue> mainInputs = ImmutableList.builder();
- PTransform<?, ?> transform = application.getTransform();
- for (Map.Entry<TupleTag<?>, PValue> input : application.getInputs().entrySet()) {
- if (!transform.getAdditionalInputs().containsKey(input.getKey())) {
- mainInputs.add(input.getValue());
- }
- }
- checkArgument(
- !mainInputs.build().isEmpty() || application.getInputs().isEmpty(),
- "Expected at least one main input if any inputs exist");
- return mainInputs.build();
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowingStrategyTranslation.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowingStrategyTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowingStrategyTranslation.java
index 1456a3f..718efe7 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowingStrategyTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowingStrategyTranslation.java
@@ -38,7 +38,6 @@ import org.apache.beam.sdk.transforms.windowing.SlidingWindows;
import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
import org.apache.beam.sdk.transforms.windowing.Trigger;
import org.apache.beam.sdk.transforms.windowing.Window.ClosingBehavior;
-import org.apache.beam.sdk.transforms.windowing.Window.OnTimeBehavior;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.values.WindowingStrategy;
@@ -120,27 +119,6 @@ public class WindowingStrategyTranslation implements Serializable {
}
}
-
- public static OnTimeBehavior fromProto(RunnerApi.OnTimeBehavior proto) {
- switch (proto) {
- case FIRE_ALWAYS:
- return OnTimeBehavior.FIRE_ALWAYS;
- case FIRE_IF_NONEMPTY:
- return OnTimeBehavior.FIRE_IF_NON_EMPTY;
- case UNRECOGNIZED:
- default:
- // Whether or not it is proto that cannot recognize it (due to the version of the
- // generated code we link to) or the switch hasn't been updated to handle it,
- // the situation is the same: we don't know what this OutputTime means
- throw new IllegalArgumentException(
- String.format(
- "Cannot convert unknown %s to %s: %s",
- RunnerApi.OnTimeBehavior.class.getCanonicalName(),
- OnTimeBehavior.class.getCanonicalName(),
- proto));
- }
- }
-
public static RunnerApi.OutputTime toProto(TimestampCombiner timestampCombiner) {
switch(timestampCombiner) {
case EARLIEST:
@@ -307,7 +285,6 @@ public class WindowingStrategyTranslation implements Serializable {
.setAllowedLateness(windowingStrategy.getAllowedLateness().getMillis())
.setTrigger(TriggerTranslation.toProto(windowingStrategy.getTrigger()))
.setWindowFn(windowFnSpec)
- .setAssignsToOneWindow(windowingStrategy.getWindowFn().assignsToOneWindow())
.setWindowCoderId(
components.registerCoder(windowingStrategy.getWindowFn().windowCoder()));
@@ -346,15 +323,13 @@ public class WindowingStrategyTranslation implements Serializable {
Trigger trigger = TriggerTranslation.fromProto(proto.getTrigger());
ClosingBehavior closingBehavior = fromProto(proto.getClosingBehavior());
Duration allowedLateness = Duration.millis(proto.getAllowedLateness());
- OnTimeBehavior onTimeBehavior = fromProto(proto.getOnTimeBehavior());
return WindowingStrategy.of(windowFn)
.withAllowedLateness(allowedLateness)
.withMode(accumulationMode)
.withTrigger(trigger)
.withTimestampCombiner(timestampCombiner)
- .withClosingBehavior(closingBehavior)
- .withOnTimeBehavior(onTimeBehavior);
+ .withClosingBehavior(closingBehavior);
}
public static WindowFn<?, ?> windowFnFromProto(SdkFunctionSpec windowFnSpec)
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WriteFilesTranslation.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WriteFilesTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WriteFilesTranslation.java
index b1d2da4..99b77ef 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WriteFilesTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WriteFilesTranslation.java
@@ -26,7 +26,6 @@ import com.google.protobuf.Any;
import com.google.protobuf.ByteString;
import com.google.protobuf.BytesValue;
import java.io.IOException;
-import java.io.Serializable;
import java.util.Collections;
import java.util.Map;
import org.apache.beam.runners.core.construction.PTransformTranslation.TransformPayloadTranslator;
@@ -38,7 +37,6 @@ import org.apache.beam.sdk.io.FileBasedSink;
import org.apache.beam.sdk.io.WriteFiles;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PDone;
@@ -53,45 +51,32 @@ public class WriteFilesTranslation {
public static final String CUSTOM_JAVA_FILE_BASED_SINK_URN =
"urn:beam:file_based_sink:javasdk:0.1";
- public static final String CUSTOM_JAVA_FILE_BASED_SINK_FORMAT_FUNCTION_URN =
- "urn:beam:file_based_sink_format_function:javasdk:0.1";
-
@VisibleForTesting
- static WriteFilesPayload toProto(WriteFiles<?, ?, ?> transform) {
+ static WriteFilesPayload toProto(WriteFiles<?> transform) {
return WriteFilesPayload.newBuilder()
.setSink(toProto(transform.getSink()))
- .setFormatFunction(toProto(transform.getFormatFunction()))
.setWindowedWrites(transform.isWindowedWrites())
.setRunnerDeterminedSharding(
transform.getNumShards() == null && transform.getSharding() == null)
.build();
}
- private static SdkFunctionSpec toProto(FileBasedSink<?, ?> sink) {
- return toProto(CUSTOM_JAVA_FILE_BASED_SINK_URN, sink);
- }
-
- private static SdkFunctionSpec toProto(SerializableFunction<?, ?> serializableFunction) {
- return toProto(CUSTOM_JAVA_FILE_BASED_SINK_FORMAT_FUNCTION_URN, serializableFunction);
- }
-
- private static SdkFunctionSpec toProto(String urn, Serializable serializable) {
+ private static SdkFunctionSpec toProto(FileBasedSink<?> sink) {
return SdkFunctionSpec.newBuilder()
.setSpec(
FunctionSpec.newBuilder()
- .setUrn(urn)
+ .setUrn(CUSTOM_JAVA_FILE_BASED_SINK_URN)
.setParameter(
Any.pack(
BytesValue.newBuilder()
.setValue(
- ByteString.copyFrom(
- SerializableUtils.serializeToByteArray(serializable)))
+ ByteString.copyFrom(SerializableUtils.serializeToByteArray(sink)))
.build())))
.build();
}
@VisibleForTesting
- static FileBasedSink<?, ?> sinkFromProto(SdkFunctionSpec sinkProto) throws IOException {
+ static FileBasedSink<?> sinkFromProto(SdkFunctionSpec sinkProto) throws IOException {
checkArgument(
sinkProto.getSpec().getUrn().equals(CUSTOM_JAVA_FILE_BASED_SINK_URN),
"Cannot extract %s instance from %s with URN %s",
@@ -102,44 +87,16 @@ public class WriteFilesTranslation {
byte[] serializedSink =
sinkProto.getSpec().getParameter().unpack(BytesValue.class).getValue().toByteArray();
- return (FileBasedSink<?, ?>)
+ return (FileBasedSink<?>)
SerializableUtils.deserializeFromByteArray(
serializedSink, FileBasedSink.class.getSimpleName());
}
- @VisibleForTesting
- static <InputT, OutputT> SerializableFunction<InputT, OutputT> formatFunctionFromProto(
- SdkFunctionSpec sinkProto) throws IOException {
- checkArgument(
- sinkProto.getSpec().getUrn().equals(CUSTOM_JAVA_FILE_BASED_SINK_FORMAT_FUNCTION_URN),
- "Cannot extract %s instance from %s with URN %s",
- SerializableFunction.class.getSimpleName(),
- FunctionSpec.class.getSimpleName(),
- sinkProto.getSpec().getUrn());
-
- byte[] serializedFunction =
- sinkProto.getSpec().getParameter().unpack(BytesValue.class).getValue().toByteArray();
-
- return (SerializableFunction<InputT, OutputT>)
- SerializableUtils.deserializeFromByteArray(
- serializedFunction, FileBasedSink.class.getSimpleName());
- }
-
- public static <UserT, DestinationT, OutputT> FileBasedSink<OutputT, DestinationT> getSink(
- AppliedPTransform<PCollection<UserT>, PDone, ? extends PTransform<PCollection<UserT>, PDone>>
- transform)
- throws IOException {
- return (FileBasedSink<OutputT, DestinationT>)
- sinkFromProto(getWriteFilesPayload(transform).getSink());
- }
-
- public static <InputT, OutputT> SerializableFunction<InputT, OutputT> getFormatFunction(
- AppliedPTransform<
- PCollection<InputT>, PDone, ? extends PTransform<PCollection<InputT>, PDone>>
+ public static <T> FileBasedSink<T> getSink(
+ AppliedPTransform<PCollection<T>, PDone, ? extends PTransform<PCollection<T>, PDone>>
transform)
throws IOException {
- return formatFunctionFromProto(
- getWriteFilesPayload(transform).<InputT, OutputT>getFormatFunction());
+ return (FileBasedSink<T>) sinkFromProto(getWriteFilesPayload(transform).getSink());
}
public static <T> boolean isWindowedWrites(
@@ -167,15 +124,15 @@ public class WriteFilesTranslation {
.unpack(WriteFilesPayload.class);
}
- static class WriteFilesTranslator implements TransformPayloadTranslator<WriteFiles<?, ?, ?>> {
+ static class WriteFilesTranslator implements TransformPayloadTranslator<WriteFiles<?>> {
@Override
- public String getUrn(WriteFiles<?, ?, ?> transform) {
+ public String getUrn(WriteFiles<?> transform) {
return PTransformTranslation.WRITE_FILES_TRANSFORM_URN;
}
@Override
public FunctionSpec translate(
- AppliedPTransform<?, ?, WriteFiles<?, ?, ?>> transform, SdkComponents components) {
+ AppliedPTransform<?, ?, WriteFiles<?>> transform, SdkComponents components) {
return FunctionSpec.newBuilder()
.setUrn(getUrn(transform.getTransform()))
.setParameter(Any.pack(toProto(transform.getTransform())))
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/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/c1b2b96a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PCollectionTranslationTest.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PCollectionTranslationTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PCollectionTranslationTest.java
index 5c45487..3b94220 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PCollectionTranslationTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PCollectionTranslationTest.java
@@ -113,28 +113,6 @@ public class PCollectionTranslationTest {
@Test
public void testEncodeDecodeCycle() throws Exception {
- // Encode
- SdkComponents sdkComponents = SdkComponents.create();
- RunnerApi.PCollection protoCollection =
- PCollectionTranslation.toProto(testCollection, sdkComponents);
- RunnerApi.Components protoComponents = sdkComponents.toComponents();
-
- // Decode
- Pipeline pipeline = Pipeline.create();
- PCollection<?> decodedCollection =
- PCollectionTranslation.fromProto(pipeline, protoCollection, protoComponents);
-
- // Verify
- assertThat(decodedCollection.getCoder(), Matchers.<Coder<?>>equalTo(testCollection.getCoder()));
- assertThat(
- decodedCollection.getWindowingStrategy(),
- Matchers.<WindowingStrategy<?, ?>>equalTo(
- testCollection.getWindowingStrategy().fixDefaults()));
- assertThat(decodedCollection.isBounded(), equalTo(testCollection.isBounded()));
- }
-
- @Test
- public void testEncodeDecodeFields() throws Exception {
SdkComponents sdkComponents = SdkComponents.create();
RunnerApi.PCollection protoCollection = PCollectionTranslation
.toProto(testCollection, sdkComponents);
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PTransformMatchersTest.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PTransformMatchersTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PTransformMatchersTest.java
index 99d3dd1..2497598 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PTransformMatchersTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PTransformMatchersTest.java
@@ -27,12 +27,9 @@ import com.google.common.base.MoreObjects;
import com.google.common.collect.ImmutableMap;
import java.io.Serializable;
import java.util.Collections;
-import org.apache.beam.sdk.coders.KvCoder;
-import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.io.DefaultFilenamePolicy;
-import org.apache.beam.sdk.io.DynamicFileDestinations;
import org.apache.beam.sdk.io.FileBasedSink;
import org.apache.beam.sdk.io.FileBasedSink.FilenamePolicy;
import org.apache.beam.sdk.io.LocalResources;
@@ -56,7 +53,6 @@ import org.apache.beam.sdk.transforms.Materialization;
import org.apache.beam.sdk.transforms.Materializations;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.transforms.SerializableFunctions;
import org.apache.beam.sdk.transforms.Sum;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.View.CreatePCollectionView;
@@ -99,14 +95,9 @@ public class PTransformMatchersTest implements Serializable {
PCollection<KV<String, Integer>> input =
PCollection.createPrimitiveOutputInternal(
p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED);
- input.setName("dummy input");
- input.setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()));
-
PCollection<Integer> output =
PCollection.createPrimitiveOutputInternal(
p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED);
- output.setName("dummy output");
- output.setCoder(VarIntCoder.of());
return AppliedPTransform.of("pardo", input.expand(), output.expand(), pardo, p);
}
@@ -281,18 +272,6 @@ public class PTransformMatchersTest implements Serializable {
}
@Test
- public void parDoSplittable() {
- AppliedPTransform<?, ?, ?> parDoApplication =
- getAppliedTransform(
- ParDo.of(splittableDoFn).withOutputTags(new TupleTag<Integer>(), TupleTagList.empty()));
- assertThat(PTransformMatchers.splittableParDo().matches(parDoApplication), is(true));
-
- assertThat(PTransformMatchers.stateOrTimerParDoMulti().matches(parDoApplication), is(false));
- assertThat(PTransformMatchers.splittableParDoSingle().matches(parDoApplication), is(false));
- assertThat(PTransformMatchers.stateOrTimerParDoSingle().matches(parDoApplication), is(false));
- }
-
- @Test
public void parDoMultiWithState() {
AppliedPTransform<?, ?, ?> parDoApplication =
getAppliedTransform(
@@ -305,19 +284,6 @@ public class PTransformMatchersTest implements Serializable {
}
@Test
- public void parDoWithState() {
- AppliedPTransform<?, ?, ?> statefulApplication =
- getAppliedTransform(
- ParDo.of(doFnWithState).withOutputTags(new TupleTag<Integer>(), TupleTagList.empty()));
- assertThat(PTransformMatchers.stateOrTimerParDo().matches(statefulApplication), is(true));
-
- AppliedPTransform<?, ?, ?> splittableApplication =
- getAppliedTransform(
- ParDo.of(splittableDoFn).withOutputTags(new TupleTag<Integer>(), TupleTagList.empty()));
- assertThat(PTransformMatchers.stateOrTimerParDo().matches(splittableApplication), is(false));
- }
-
- @Test
public void parDoMultiWithTimers() {
AppliedPTransform<?, ?, ?> parDoApplication =
getAppliedTransform(
@@ -539,32 +505,30 @@ public class PTransformMatchersTest implements Serializable {
public void writeWithRunnerDeterminedSharding() {
ResourceId outputDirectory = LocalResources.fromString("/foo/bar", true /* isDirectory */);
FilenamePolicy policy =
- DefaultFilenamePolicy.fromStandardParameters(
+ DefaultFilenamePolicy.constructUsingStandardParameters(
StaticValueProvider.of(outputDirectory),
DefaultFilenamePolicy.DEFAULT_UNWINDOWED_SHARD_TEMPLATE,
"",
false);
- WriteFiles<Integer, Void, Integer> write =
+ WriteFiles<Integer> write =
WriteFiles.to(
- new FileBasedSink<Integer, Void>(
- StaticValueProvider.of(outputDirectory), DynamicFileDestinations.constant(null)) {
+ new FileBasedSink<Integer>(StaticValueProvider.of(outputDirectory), policy) {
@Override
- public WriteOperation<Integer, Void> createWriteOperation() {
+ public WriteOperation<Integer> createWriteOperation() {
return null;
}
- },
- SerializableFunctions.<Integer>identity());
+ });
assertThat(
PTransformMatchers.writeWithRunnerDeterminedSharding().matches(appliedWrite(write)),
is(true));
- WriteFiles<Integer, Void, Integer> withStaticSharding = write.withNumShards(3);
+ WriteFiles<Integer> withStaticSharding = write.withNumShards(3);
assertThat(
PTransformMatchers.writeWithRunnerDeterminedSharding()
.matches(appliedWrite(withStaticSharding)),
is(false));
- WriteFiles<Integer, Void, Integer> withCustomSharding =
+ WriteFiles<Integer> withCustomSharding =
write.withSharding(Sum.integersGlobally().asSingletonView());
assertThat(
PTransformMatchers.writeWithRunnerDeterminedSharding()
@@ -572,8 +536,8 @@ public class PTransformMatchersTest implements Serializable {
is(false));
}
- private AppliedPTransform<?, ?, ?> appliedWrite(WriteFiles<Integer, Void, Integer> write) {
- return AppliedPTransform.<PCollection<Integer>, PDone, WriteFiles<Integer, Void, Integer>>of(
+ private AppliedPTransform<?, ?, ?> appliedWrite(WriteFiles<Integer> write) {
+ return AppliedPTransform.<PCollection<Integer>, PDone, WriteFiles<Integer>>of(
"WriteFiles",
Collections.<TupleTag<?>, PValue>emptyMap(),
Collections.<TupleTag<?>, PValue>emptyMap(),
[27/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
Revert "[BEAM-2610] This closes #3553"
This reverts commit ec494f675aa73fbdc7929f9592f33951941962b0, reversing
changes made to d89d1ee1a3085269cdf44ec50e29a95c8f43757b.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/c1b2b96a
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/c1b2b96a
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/c1b2b96a
Branch: refs/heads/DSL_SQL
Commit: c1b2b96a438b86a8b0023f6943dcf0a4f238ba39
Parents: 97a156c
Author: mingmxu <mi...@ebay.com>
Authored: Wed Jul 19 14:26:11 2017 -0700
Committer: mingmxu <mi...@ebay.com>
Committed: Wed Jul 19 14:26:11 2017 -0700
----------------------------------------------------------------------
.gitignore | 2 +-
.../jenkins/common_job_properties.groovy | 9 +-
.../job_beam_PerformanceTests_Python.groovy | 58 --
..._beam_PostCommit_Java_JDKVersionsTest.groovy | 2 -
..._PostCommit_Java_MavenInstall_Windows.groovy | 3 +-
.../job_beam_PreCommit_Website_Merge.groovy | 59 --
examples/java/pom.xml | 20 +-
.../examples/common/WriteOneFilePerWindow.java | 52 +-
.../beam/examples/WindowedWordCountIT.java | 4 +-
examples/java8/pom.xml | 20 +-
.../complete/game/utils/WriteToText.java | 43 +-
.../examples/complete/game/LeaderBoardTest.java | 2 -
examples/pom.xml | 2 +-
pom.xml | 123 +---
runners/apex/pom.xml | 20 +-
.../apache/beam/runners/apex/ApexRunner.java | 61 +-
.../translation/ApexPipelineTranslator.java | 16 +-
.../apex/translation/TranslationContext.java | 4 +-
.../operators/ApexParDoOperator.java | 21 +-
.../runners/apex/examples/WordCountTest.java | 8 +-
.../utils/ApexStateInternalsTest.java | 411 ++++++++---
runners/core-construction-java/pom.xml | 2 +-
.../CreatePCollectionViewTranslation.java | 4 +-
.../construction/ElementAndRestriction.java | 42 ++
.../ElementAndRestrictionCoder.java | 88 +++
.../construction/PCollectionTranslation.java | 16 -
.../core/construction/PTransformMatchers.java | 109 +--
.../construction/PTransformTranslation.java | 7 +-
.../core/construction/ParDoTranslation.java | 82 +--
.../construction/RunnerPCollectionView.java | 31 +-
.../core/construction/SplittableParDo.java | 124 +---
.../construction/TestStreamTranslation.java | 49 +-
.../core/construction/TransformInputs.java | 50 --
.../WindowingStrategyTranslation.java | 27 +-
.../construction/WriteFilesTranslation.java | 67 +-
.../ElementAndRestrictionCoderTest.java | 126 ++++
.../PCollectionTranslationTest.java | 22 -
.../construction/PTransformMatchersTest.java | 54 +-
.../core/construction/ParDoTranslationTest.java | 28 +-
.../core/construction/SplittableParDoTest.java | 18 +-
.../core/construction/TransformInputsTest.java | 166 -----
.../WindowingStrategyTranslationTest.java | 3 -
.../construction/WriteFilesTranslationTest.java | 62 +-
runners/core-java/pom.xml | 2 +-
.../core/LateDataDroppingDoFnRunner.java | 33 +-
...eBoundedSplittableProcessElementInvoker.java | 40 +-
.../beam/runners/core/ProcessFnRunner.java | 16 +-
.../beam/runners/core/ReduceFnRunner.java | 135 ++--
.../beam/runners/core/SimpleDoFnRunner.java | 20 -
.../core/SplittableParDoViaKeyedWorkItems.java | 58 +-
.../core/SplittableProcessElementInvoker.java | 25 +-
.../beam/runners/core/SystemReduceFn.java | 6 -
.../core/triggers/AfterAllStateMachine.java | 25 +-
.../AfterDelayFromFirstElementStateMachine.java | 6 +-
.../core/triggers/AfterFirstStateMachine.java | 20 +-
.../core/triggers/AfterPaneStateMachine.java | 6 +-
.../triggers/AfterWatermarkStateMachine.java | 7 +-
.../triggers/ExecutableTriggerStateMachine.java | 23 +-
.../core/triggers/NeverStateMachine.java | 5 +-
.../core/triggers/TriggerStateMachine.java | 27 +
.../core/InMemoryStateInternalsTest.java | 569 +++++++++++++--
...ndedSplittableProcessElementInvokerTest.java | 47 +-
.../beam/runners/core/ReduceFnRunnerTest.java | 374 +---------
.../beam/runners/core/ReduceFnTester.java | 48 +-
.../core/SplittableParDoProcessFnTest.java | 117 +--
.../beam/runners/core/StateInternalsTest.java | 613 ----------------
.../beam/runners/core/WindowMatchers.java | 15 -
.../triggers/AfterFirstStateMachineTest.java | 5 +-
.../AfterWatermarkStateMachineTest.java | 7 +-
.../core/triggers/StubTriggerStateMachine.java | 7 +-
runners/direct-java/pom.xml | 7 +-
.../beam/runners/direct/CommittedResult.java | 12 +-
.../apache/beam/runners/direct/DirectGraph.java | 38 +-
.../beam/runners/direct/DirectGraphVisitor.java | 48 +-
.../beam/runners/direct/DirectGroupByKey.java | 13 +-
.../direct/DirectGroupByKeyOverrideFactory.java | 14 +-
.../beam/runners/direct/DirectRegistrar.java | 2 +-
.../beam/runners/direct/DirectRunner.java | 64 +-
.../beam/runners/direct/DirectTestOptions.java | 42 --
.../beam/runners/direct/EvaluationContext.java | 26 +-
.../direct/ExecutorServiceParallelExecutor.java | 27 +-
.../runners/direct/ParDoEvaluatorFactory.java | 9 +-
.../direct/ParDoMultiOverrideFactory.java | 121 +---
...littableProcessElementsEvaluatorFactory.java | 37 +-
.../direct/StatefulParDoEvaluatorFactory.java | 12 +-
.../direct/TestStreamEvaluatorFactory.java | 20 +-
.../runners/direct/ViewEvaluatorFactory.java | 8 +-
.../runners/direct/ViewOverrideFactory.java | 69 +-
.../beam/runners/direct/WatermarkManager.java | 18 +-
.../direct/WriteWithShardingFactory.java | 34 +-
.../runners/direct/CommittedResultTest.java | 17 +-
.../runners/direct/DirectGraphVisitorTest.java | 10 +-
.../beam/runners/direct/DirectGraphs.java | 7 -
.../runners/direct/DirectRegistrarTest.java | 2 +-
.../runners/direct/EvaluationContextTest.java | 7 +-
.../ImmutabilityEnforcementFactoryTest.java | 4 +-
.../beam/runners/direct/ParDoEvaluatorTest.java | 7 +-
.../StatefulParDoEvaluatorFactoryTest.java | 65 +-
.../runners/direct/TransformExecutorTest.java | 12 +-
.../direct/ViewEvaluatorFactoryTest.java | 8 +-
.../runners/direct/ViewOverrideFactoryTest.java | 37 +-
.../direct/WatermarkCallbackExecutorTest.java | 1 -
.../runners/direct/WatermarkManagerTest.java | 16 +-
.../direct/WriteWithShardingFactoryTest.java | 44 +-
runners/flink/pom.xml | 11 +-
.../runners/flink/CreateStreamingFlinkView.java | 154 ----
.../flink/FlinkBatchTranslationContext.java | 3 +-
.../FlinkPipelineExecutionEnvironment.java | 2 -
.../flink/FlinkStreamingPipelineTranslator.java | 86 ++-
.../FlinkStreamingTransformTranslators.java | 36 +-
.../flink/FlinkStreamingTranslationContext.java | 3 +-
.../flink/FlinkStreamingViewOverrides.java | 372 ++++++++++
.../runners/flink/FlinkTransformOverrides.java | 53 --
.../streaming/SplittableDoFnOperator.java | 16 +-
.../streaming/state/FlinkStateInternals.java | 425 ++++++-----
.../FlinkBroadcastStateInternalsTest.java | 242 +++++--
.../FlinkKeyGroupStateInternalsTest.java | 359 +++++-----
.../streaming/FlinkSplitStateInternalsTest.java | 132 ++--
.../streaming/FlinkStateInternalsTest.java | 343 ++++++++-
runners/google-cloud-dataflow-java/pom.xml | 10 +-
.../dataflow/BatchStatefulParDoOverrides.java | 4 -
.../runners/dataflow/BatchViewOverrides.java | 182 +++--
.../runners/dataflow/CreateDataflowView.java | 8 +-
.../dataflow/DataflowPipelineTranslator.java | 62 +-
.../beam/runners/dataflow/DataflowRunner.java | 133 +---
.../dataflow/SplittableParDoOverrides.java | 76 --
.../dataflow/StreamingViewOverrides.java | 10 +-
.../runners/dataflow/TransformTranslator.java | 4 +-
.../runners/dataflow/util/PropertyNames.java | 1 -
.../beam/runners/dataflow/util/TimeUtil.java | 24 +-
.../DataflowPipelineTranslatorTest.java | 95 +--
.../runners/dataflow/DataflowRunnerTest.java | 198 +-----
.../runners/dataflow/util/TimeUtilTest.java | 6 -
runners/pom.xml | 2 +-
runners/spark/pom.xml | 70 +-
.../spark/SparkNativePipelineVisitor.java | 3 +-
.../apache/beam/runners/spark/SparkRunner.java | 9 +-
.../beam/runners/spark/TestSparkRunner.java | 2 +-
.../SparkGroupAlsoByWindowViaWindowSet.java | 6 +-
.../spark/stateful/SparkTimerInternals.java | 18 +-
.../spark/translation/EvaluationContext.java | 4 +-
.../spark/translation/TransformTranslator.java | 50 +-
.../spark/util/GlobalWatermarkHolder.java | 127 +---
.../spark/GlobalWatermarkHolderTest.java | 18 +-
.../runners/spark/SparkRunnerDebuggerTest.java | 26 +-
.../spark/stateful/SparkStateInternalsTest.java | 66 --
.../spark/translation/StorageLevelTest.java | 4 +-
sdks/common/fn-api/pom.xml | 2 +-
.../fn-api/src/main/proto/beam_fn_api.proto | 237 ++++--
sdks/common/pom.xml | 2 +-
sdks/common/runner-api/pom.xml | 2 +-
.../src/main/proto/beam_runner_api.proto | 26 +-
sdks/java/build-tools/pom.xml | 2 +-
.../src/main/resources/beam/findbugs-filter.xml | 9 -
sdks/java/core/pom.xml | 2 +-
.../apache/beam/sdk/coders/ShardedKeyCoder.java | 66 --
.../java/org/apache/beam/sdk/io/AvroIO.java | 220 +++---
.../java/org/apache/beam/sdk/io/AvroSink.java | 32 +-
.../apache/beam/sdk/io/CompressedSource.java | 40 +-
.../beam/sdk/io/DefaultFilenamePolicy.java | 274 ++-----
.../beam/sdk/io/DynamicFileDestinations.java | 115 ---
.../org/apache/beam/sdk/io/FileBasedSink.java | 513 ++++++-------
.../apache/beam/sdk/io/OffsetBasedSource.java | 22 +-
.../java/org/apache/beam/sdk/io/TFRecordIO.java | 44 +-
.../java/org/apache/beam/sdk/io/TextIO.java | 712 ++++---------------
.../java/org/apache/beam/sdk/io/TextSink.java | 22 +-
.../java/org/apache/beam/sdk/io/WriteFiles.java | 647 ++++++-----------
.../beam/sdk/io/range/ByteKeyRangeTracker.java | 22 +-
.../apache/beam/sdk/io/range/OffsetRange.java | 101 ---
.../beam/sdk/io/range/OffsetRangeTracker.java | 3 -
.../sdk/options/PipelineOptionsFactory.java | 18 +-
.../sdk/options/PipelineOptionsValidator.java | 34 +-
.../sdk/options/ProxyInvocationHandler.java | 19 +-
.../beam/sdk/runners/TransformHierarchy.java | 165 +----
.../apache/beam/sdk/testing/StaticWindows.java | 5 -
.../org/apache/beam/sdk/testing/TestStream.java | 12 -
.../org/apache/beam/sdk/transforms/Combine.java | 30 +-
.../org/apache/beam/sdk/transforms/DoFn.java | 52 +-
.../apache/beam/sdk/transforms/DoFnTester.java | 21 +-
.../org/apache/beam/sdk/transforms/ParDo.java | 41 +-
.../sdk/transforms/SerializableFunctions.java | 50 --
.../org/apache/beam/sdk/transforms/View.java | 38 +-
.../reflect/ByteBuddyDoFnInvokerFactory.java | 27 -
.../reflect/ByteBuddyOnTimerInvokerFactory.java | 73 +-
.../sdk/transforms/reflect/DoFnInvoker.java | 17 +-
.../sdk/transforms/reflect/DoFnSignature.java | 33 +-
.../sdk/transforms/reflect/DoFnSignatures.java | 44 +-
.../reflect/OnTimerMethodSpecifier.java | 37 -
.../transforms/splittabledofn/OffsetRange.java | 77 ++
.../splittabledofn/OffsetRangeTracker.java | 11 -
.../splittabledofn/RestrictionTracker.java | 11 +-
.../sdk/transforms/windowing/GlobalWindows.java | 5 -
.../windowing/PartitioningWindowFn.java | 5 -
.../transforms/windowing/SlidingWindows.java | 5 -
.../beam/sdk/transforms/windowing/Window.java | 32 -
.../beam/sdk/transforms/windowing/WindowFn.java | 11 -
.../apache/beam/sdk/util/IdentityWindowFn.java | 5 -
.../org/apache/beam/sdk/values/PCollection.java | 12 -
.../beam/sdk/values/PCollectionViews.java | 38 -
.../org/apache/beam/sdk/values/PValueBase.java | 12 +
.../org/apache/beam/sdk/values/ShardedKey.java | 65 --
.../beam/sdk/values/WindowingStrategy.java | 46 +-
.../java/org/apache/beam/sdk/io/AvroIOTest.java | 85 +--
.../beam/sdk/io/DefaultFilenamePolicyTest.java | 135 ++--
.../sdk/io/DrunkWritableByteChannelFactory.java | 2 +-
.../apache/beam/sdk/io/FileBasedSinkTest.java | 93 +--
.../java/org/apache/beam/sdk/io/SimpleSink.java | 56 +-
.../java/org/apache/beam/sdk/io/TextIOTest.java | 326 +--------
.../org/apache/beam/sdk/io/WriteFilesTest.java | 366 ++--------
.../options/PipelineOptionsValidatorTest.java | 44 --
.../sdk/options/ProxyInvocationHandlerTest.java | 19 -
.../sdk/runners/TransformHierarchyTest.java | 197 -----
.../sdk/testing/PCollectionViewTesting.java | 8 -
.../apache/beam/sdk/transforms/CombineTest.java | 365 ++++------
.../beam/sdk/transforms/DoFnTesterTest.java | 32 -
.../beam/sdk/transforms/GroupByKeyTest.java | 39 -
.../apache/beam/sdk/transforms/ParDoTest.java | 165 -----
.../beam/sdk/transforms/SplittableDoFnTest.java | 155 +---
.../transforms/reflect/DoFnInvokersTest.java | 93 +--
.../DoFnSignaturesProcessElementTest.java | 2 +-
.../DoFnSignaturesSplittableDoFnTest.java | 83 +--
.../transforms/reflect/DoFnSignaturesTest.java | 14 -
.../splittabledofn/OffsetRangeTrackerTest.java | 1 -
.../windowing/SlidingWindowsTest.java | 30 +-
.../google-cloud-platform-core/pom.xml | 2 +-
.../java/org/apache/beam/sdk/util/GcsUtil.java | 2 +-
.../sdk/util/RetryHttpRequestInitializer.java | 147 ++--
.../extensions/gcp/GcpCoreApiSurfaceTest.java | 48 +-
.../util/RetryHttpRequestInitializerTest.java | 31 +-
sdks/java/extensions/jackson/pom.xml | 2 +-
sdks/java/extensions/join-library/pom.xml | 2 +-
sdks/java/extensions/pom.xml | 2 +-
sdks/java/extensions/protobuf/pom.xml | 2 +-
sdks/java/extensions/sorter/pom.xml | 8 +-
sdks/java/harness/pom.xml | 18 +-
.../harness/control/ProcessBundleHandler.java | 295 ++++++--
.../fn/harness/control/RegisterHandler.java | 2 +-
.../beam/runners/core/BeamFnDataReadRunner.java | 70 +-
.../runners/core/BeamFnDataWriteRunner.java | 67 +-
.../beam/runners/core/BoundedSourceRunner.java | 74 +-
.../beam/runners/core/FnApiDoFnRunner.java | 547 --------------
.../runners/core/PTransformRunnerFactory.java | 81 ---
.../control/ProcessBundleHandlerTest.java | 521 ++++++++++++--
.../fn/harness/control/RegisterHandlerTest.java | 8 +-
.../runners/core/BeamFnDataReadRunnerTest.java | 112 +--
.../runners/core/BeamFnDataWriteRunnerTest.java | 120 +---
.../runners/core/BoundedSourceRunnerTest.java | 124 +---
.../beam/runners/core/FnApiDoFnRunnerTest.java | 210 ------
sdks/java/io/amqp/pom.xml | 100 ---
.../org/apache/beam/sdk/io/amqp/AmqpIO.java | 399 -----------
.../beam/sdk/io/amqp/AmqpMessageCoder.java | 79 --
.../amqp/AmqpMessageCoderProviderRegistrar.java | 44 --
.../apache/beam/sdk/io/amqp/package-info.java | 22 -
.../org/apache/beam/sdk/io/amqp/AmqpIOTest.java | 148 ----
.../beam/sdk/io/amqp/AmqpMessageCoderTest.java | 89 ---
sdks/java/io/cassandra/pom.xml | 2 +-
.../beam/sdk/io/cassandra/CassandraIO.java | 2 +-
sdks/java/io/common/pom.xml | 2 +-
.../sdk/io/common/IOTestPipelineOptions.java | 6 +-
sdks/java/io/elasticsearch/pom.xml | 10 +-
.../sdk/io/elasticsearch/ElasticsearchIO.java | 17 +-
.../elasticsearch/ElasticSearchIOTestUtils.java | 81 +--
.../sdk/io/elasticsearch/ElasticsearchIOIT.java | 14 +-
.../io/elasticsearch/ElasticsearchIOTest.java | 36 +-
.../elasticsearch/ElasticsearchTestDataSet.java | 37 +-
sdks/java/io/google-cloud-platform/pom.xml | 14 +-
.../beam/sdk/io/gcp/bigquery/BatchLoads.java | 2 -
.../io/gcp/bigquery/DynamicDestinations.java | 29 +-
.../io/gcp/bigquery/GenerateShardedTable.java | 1 -
.../beam/sdk/io/gcp/bigquery/ShardedKey.java | 67 ++
.../sdk/io/gcp/bigquery/ShardedKeyCoder.java | 74 ++
.../sdk/io/gcp/bigquery/StreamingWriteFn.java | 1 -
.../io/gcp/bigquery/StreamingWriteTables.java | 2 -
.../sdk/io/gcp/bigquery/TagWithUniqueIds.java | 1 -
.../io/gcp/bigquery/WriteBundlesToFiles.java | 2 -
.../bigquery/WriteGroupedRecordsToFiles.java | 1 -
.../sdk/io/gcp/bigquery/WritePartition.java | 1 -
.../beam/sdk/io/gcp/bigquery/WriteTables.java | 1 -
.../beam/sdk/io/gcp/bigtable/BigtableIO.java | 8 +-
.../io/gcp/bigtable/BigtableServiceImpl.java | 9 +-
.../sdk/io/gcp/datastore/AdaptiveThrottler.java | 103 ---
.../beam/sdk/io/gcp/datastore/DatastoreV1.java | 149 +---
.../sdk/io/gcp/datastore/MovingAverage.java | 50 --
.../sdk/io/gcp/spanner/AbstractSpannerFn.java | 58 --
.../sdk/io/gcp/spanner/CreateTransactionFn.java | 51 --
.../beam/sdk/io/gcp/spanner/MutationGroup.java | 67 --
.../io/gcp/spanner/MutationSizeEstimator.java | 9 -
.../sdk/io/gcp/spanner/NaiveSpannerReadFn.java | 65 --
.../beam/sdk/io/gcp/spanner/SpannerConfig.java | 137 ----
.../beam/sdk/io/gcp/spanner/SpannerIO.java | 616 +++++-----------
.../sdk/io/gcp/spanner/SpannerWriteGroupFn.java | 125 ----
.../beam/sdk/io/gcp/spanner/Transaction.java | 33 -
.../beam/sdk/io/gcp/GcpApiSurfaceTest.java | 10 -
.../sdk/io/gcp/bigquery/BigQueryIOTest.java | 2 -
.../sdk/io/gcp/bigtable/BigtableReadIT.java | 5 +-
.../io/gcp/bigtable/BigtableTestOptions.java | 5 +
.../sdk/io/gcp/bigtable/BigtableWriteIT.java | 4 +-
.../io/gcp/datastore/AdaptiveThrottlerTest.java | 111 ---
.../sdk/io/gcp/datastore/DatastoreV1Test.java | 92 +--
.../beam/sdk/io/gcp/datastore/V1TestUtil.java | 2 +-
.../sdk/io/gcp/spanner/FakeServiceFactory.java | 82 ---
.../gcp/spanner/MutationSizeEstimatorTest.java | 12 -
.../beam/sdk/io/gcp/spanner/RandomUtils.java | 41 --
.../sdk/io/gcp/spanner/SpannerIOReadTest.java | 281 --------
.../beam/sdk/io/gcp/spanner/SpannerIOTest.java | 244 +++++++
.../sdk/io/gcp/spanner/SpannerIOWriteTest.java | 258 -------
.../beam/sdk/io/gcp/spanner/SpannerReadIT.java | 166 -----
.../beam/sdk/io/gcp/spanner/SpannerWriteIT.java | 26 +-
sdks/java/io/hadoop-common/pom.xml | 2 +-
sdks/java/io/hadoop-file-system/pom.xml | 33 +-
sdks/java/io/hadoop/input-format/pom.xml | 2 +-
.../hadoop/inputformat/HadoopInputFormatIO.java | 2 +-
sdks/java/io/hadoop/jdk1.8-tests/pom.xml | 4 +-
.../inputformat/HIFIOWithElasticTest.java | 11 +-
sdks/java/io/hadoop/pom.xml | 2 +-
sdks/java/io/hbase/pom.xml | 26 +-
.../io/hbase/HBaseCoderProviderRegistrar.java | 40 --
.../org/apache/beam/sdk/io/hbase/HBaseIO.java | 48 +-
.../beam/sdk/io/hbase/HBaseMutationCoder.java | 42 --
.../hbase/HBaseCoderProviderRegistrarTest.java | 45 --
.../apache/beam/sdk/io/hbase/HBaseIOTest.java | 49 +-
sdks/java/io/hcatalog/pom.xml | 175 -----
.../apache/beam/sdk/io/hcatalog/HCatalogIO.java | 492 -------------
.../beam/sdk/io/hcatalog/package-info.java | 22 -
.../io/hcatalog/EmbeddedMetastoreService.java | 87 ---
.../beam/sdk/io/hcatalog/HCatalogIOTest.java | 277 --------
.../sdk/io/hcatalog/HCatalogIOTestUtils.java | 108 ---
.../hcatalog/src/test/resources/hive-site.xml | 301 --------
sdks/java/io/jdbc/pom.xml | 4 +-
.../org/apache/beam/sdk/io/jdbc/JdbcIO.java | 2 +-
sdks/java/io/jms/pom.xml | 2 +-
.../java/org/apache/beam/sdk/io/jms/JmsIO.java | 2 +-
sdks/java/io/kafka/pom.xml | 2 +-
.../org/apache/beam/sdk/io/kafka/KafkaIO.java | 132 ++--
.../apache/beam/sdk/io/kafka/KafkaIOTest.java | 30 -
sdks/java/io/kinesis/pom.xml | 2 +-
.../sdk/io/kinesis/CheckpointGenerator.java | 6 +-
.../beam/sdk/io/kinesis/CustomOptional.java | 111 ++-
.../io/kinesis/DynamicCheckpointGenerator.java | 52 +-
.../sdk/io/kinesis/GetKinesisRecordsResult.java | 49 +-
.../sdk/io/kinesis/KinesisClientProvider.java | 4 +-
.../apache/beam/sdk/io/kinesis/KinesisIO.java | 281 ++++----
.../beam/sdk/io/kinesis/KinesisReader.java | 206 +++---
.../sdk/io/kinesis/KinesisReaderCheckpoint.java | 97 ++-
.../beam/sdk/io/kinesis/KinesisRecord.java | 177 +++--
.../beam/sdk/io/kinesis/KinesisRecordCoder.java | 68 +-
.../beam/sdk/io/kinesis/KinesisSource.java | 147 ++--
.../beam/sdk/io/kinesis/RecordFilter.java | 18 +-
.../apache/beam/sdk/io/kinesis/RoundRobin.java | 37 +-
.../beam/sdk/io/kinesis/ShardCheckpoint.java | 241 ++++---
.../sdk/io/kinesis/ShardRecordsIterator.java | 106 ++-
.../sdk/io/kinesis/SimplifiedKinesisClient.java | 215 +++---
.../beam/sdk/io/kinesis/StartingPoint.java | 84 ++-
.../io/kinesis/StaticCheckpointGenerator.java | 27 +-
.../io/kinesis/TransientKinesisException.java | 7 +-
.../beam/sdk/io/kinesis/AmazonKinesisMock.java | 539 +++++++-------
.../beam/sdk/io/kinesis/CustomOptionalTest.java | 27 +-
.../kinesis/DynamicCheckpointGeneratorTest.java | 33 +-
.../sdk/io/kinesis/KinesisMockReadTest.java | 97 ++-
.../io/kinesis/KinesisReaderCheckpointTest.java | 52 +-
.../beam/sdk/io/kinesis/KinesisReaderIT.java | 127 ++--
.../beam/sdk/io/kinesis/KinesisReaderTest.java | 166 +++--
.../sdk/io/kinesis/KinesisRecordCoderTest.java | 34 +-
.../beam/sdk/io/kinesis/KinesisTestOptions.java | 43 +-
.../beam/sdk/io/kinesis/KinesisUploader.java | 70 +-
.../beam/sdk/io/kinesis/RecordFilterTest.java | 52 +-
.../beam/sdk/io/kinesis/RoundRobinTest.java | 42 +-
.../sdk/io/kinesis/ShardCheckpointTest.java | 203 +++---
.../io/kinesis/ShardRecordsIteratorTest.java | 216 +++---
.../io/kinesis/SimplifiedKinesisClientTest.java | 351 +++++----
sdks/java/io/mongodb/pom.xml | 2 +-
.../beam/sdk/io/mongodb/MongoDbGridFSIO.java | 2 +-
.../apache/beam/sdk/io/mongodb/MongoDbIO.java | 317 ++-------
.../beam/sdk/io/mongodb/MongoDbIOTest.java | 37 -
sdks/java/io/mqtt/pom.xml | 2 +-
.../org/apache/beam/sdk/io/mqtt/MqttIO.java | 2 +-
sdks/java/io/pom.xml | 35 +-
sdks/java/io/xml/pom.xml | 2 +-
.../java/org/apache/beam/sdk/io/xml/XmlIO.java | 4 +-
.../org/apache/beam/sdk/io/xml/XmlSink.java | 21 +-
.../org/apache/beam/sdk/io/xml/XmlSinkTest.java | 4 +-
sdks/java/java8tests/pom.xml | 2 +-
sdks/java/javadoc/pom.xml | 19 +-
.../maven-archetypes/examples-java8/pom.xml | 2 +-
.../main/resources/archetype-resources/pom.xml | 1 +
sdks/java/maven-archetypes/examples/pom.xml | 2 +-
.../main/resources/archetype-resources/pom.xml | 1 +
sdks/java/maven-archetypes/pom.xml | 2 +-
sdks/java/maven-archetypes/starter/pom.xml | 2 +-
.../resources/projects/basic/reference/pom.xml | 2 +-
sdks/java/pom.xml | 2 +-
sdks/pom.xml | 2 +-
sdks/python/apache_beam/coders/coder_impl.py | 4 -
sdks/python/apache_beam/coders/coders.py | 7 +-
.../apache_beam/coders/coders_test_common.py | 8 -
.../examples/snippets/snippets_test.py | 16 -
.../apache_beam/examples/streaming_wordcount.py | 25 +-
.../apache_beam/examples/windowed_wordcount.py | 93 ---
sdks/python/apache_beam/io/filesystem.py | 22 +-
sdks/python/apache_beam/io/gcp/gcsio.py | 10 +-
sdks/python/apache_beam/io/gcp/pubsub.py | 180 ++---
sdks/python/apache_beam/io/gcp/pubsub_test.py | 101 +--
.../io/gcp/tests/bigquery_matcher.py | 6 +-
.../io/gcp/tests/bigquery_matcher_test.py | 2 +-
sdks/python/apache_beam/io/range_trackers.py | 130 ++++
.../apache_beam/io/range_trackers_test.py | 186 +++++
.../apache_beam/options/pipeline_options.py | 35 +-
.../options/pipeline_options_test.py | 39 +-
.../apache_beam/options/value_provider_test.py | 93 ++-
sdks/python/apache_beam/pipeline.py | 230 +-----
sdks/python/apache_beam/pipeline_test.py | 53 --
sdks/python/apache_beam/portability/__init__.py | 18 -
.../apache_beam/portability/api/__init__.py | 21 -
sdks/python/apache_beam/pvalue.py | 2 +-
sdks/python/apache_beam/runners/api/__init__.py | 21 +
.../runners/dataflow/dataflow_runner.py | 91 +--
.../runners/dataflow/dataflow_runner_test.py | 24 +-
.../runners/dataflow/internal/apiclient.py | 35 +-
.../runners/dataflow/internal/apiclient_test.py | 29 +-
.../runners/dataflow/internal/dependency.py | 69 +-
.../runners/dataflow/native_io/iobase_test.py | 39 +-
.../dataflow/native_io/streaming_create.py | 72 --
.../runners/dataflow/ptransform_overrides.py | 52 --
.../runners/direct/bundle_factory.py | 2 +-
.../apache_beam/runners/direct/direct_runner.py | 108 ---
.../runners/direct/evaluation_context.py | 73 +-
.../apache_beam/runners/direct/executor.py | 135 ++--
.../runners/direct/transform_evaluator.py | 447 +-----------
.../runners/direct/transform_result.py | 41 ++
sdks/python/apache_beam/runners/direct/util.py | 67 --
.../runners/direct/watermark_manager.py | 100 +--
.../apache_beam/runners/pipeline_context.py | 19 +-
.../runners/portability/fn_api_runner.py | 306 ++++----
.../runners/portability/fn_api_runner_test.py | 31 +-
.../runners/worker/bundle_processor.py | 426 -----------
.../apache_beam/runners/worker/data_plane.py | 28 +-
.../runners/worker/data_plane_test.py | 2 +-
.../apache_beam/runners/worker/log_handler.py | 2 +-
.../runners/worker/log_handler_test.py | 2 +-
.../runners/worker/operation_specs.py | 9 +-
.../apache_beam/runners/worker/operations.py | 1 -
.../apache_beam/runners/worker/sdk_worker.py | 370 +++++++++-
.../runners/worker/sdk_worker_main.py | 2 +-
.../runners/worker/sdk_worker_test.py | 95 ++-
sdks/python/apache_beam/testing/test_stream.py | 5 -
.../apache_beam/testing/test_stream_test.py | 68 --
sdks/python/apache_beam/transforms/combiners.py | 8 -
.../apache_beam/transforms/combiners_test.py | 7 +-
sdks/python/apache_beam/transforms/core.py | 102 ++-
.../python/apache_beam/transforms/ptransform.py | 43 +-
sdks/python/apache_beam/transforms/trigger.py | 28 +-
sdks/python/apache_beam/transforms/window.py | 4 +-
.../apache_beam/typehints/trivial_inference.py | 3 +-
.../typehints/trivial_inference_test.py | 7 -
sdks/python/apache_beam/utils/plugin.py | 42 --
sdks/python/apache_beam/utils/timestamp.py | 5 -
sdks/python/apache_beam/utils/urns.py | 2 +-
sdks/python/apache_beam/version.py | 2 +-
sdks/python/gen_protos.py | 2 +-
sdks/python/pom.xml | 2 +-
sdks/python/run_pylint.sh | 2 +-
sdks/python/setup.py | 5 +-
462 files changed, 10754 insertions(+), 21718 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/.gitignore
----------------------------------------------------------------------
diff --git a/.gitignore b/.gitignore
index 36c5cc8..bd419a7 100644
--- a/.gitignore
+++ b/.gitignore
@@ -25,7 +25,7 @@ sdks/python/**/*.egg
sdks/python/LICENSE
sdks/python/NOTICE
sdks/python/README.md
-sdks/python/apache_beam/portability/api/*pb2*.*
+sdks/python/apache_beam/runners/api/*pb2*.*
# Ignore IntelliJ files.
.idea/
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/.test-infra/jenkins/common_job_properties.groovy
----------------------------------------------------------------------
diff --git a/.test-infra/jenkins/common_job_properties.groovy b/.test-infra/jenkins/common_job_properties.groovy
index 70534c6..6d4d68b 100644
--- a/.test-infra/jenkins/common_job_properties.groovy
+++ b/.test-infra/jenkins/common_job_properties.groovy
@@ -23,12 +23,11 @@
class common_job_properties {
// Sets common top-level job properties for website repository jobs.
- static void setTopLevelWebsiteJobProperties(context,
- String branch = 'asf-site') {
+ static void setTopLevelWebsiteJobProperties(context) {
setTopLevelJobProperties(
context,
'beam-site',
- branch,
+ 'asf-site',
'beam',
30)
}
@@ -265,10 +264,8 @@ class common_job_properties {
shell('rm -rf PerfKitBenchmarker')
// Clone appropriate perfkit branch
shell('git clone https://github.com/GoogleCloudPlatform/PerfKitBenchmarker.git')
- // Install Perfkit benchmark requirements.
+ // Install job requirements.
shell('pip install --user -r PerfKitBenchmarker/requirements.txt')
- // Install job requirements for Python SDK.
- shell('pip install --user -e sdks/python/[gcp,test]')
// Launch performance test.
shell("python PerfKitBenchmarker/pkb.py $pkbArgs")
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/.test-infra/jenkins/job_beam_PerformanceTests_Python.groovy
----------------------------------------------------------------------
diff --git a/.test-infra/jenkins/job_beam_PerformanceTests_Python.groovy b/.test-infra/jenkins/job_beam_PerformanceTests_Python.groovy
deleted file mode 100644
index 6a71bda..0000000
--- a/.test-infra/jenkins/job_beam_PerformanceTests_Python.groovy
+++ /dev/null
@@ -1,58 +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.
- */
-
-import common_job_properties
-
-// This job runs the Beam Python performance tests on PerfKit Benchmarker.
-job('beam_PerformanceTests_Python'){
- // Set default Beam job properties.
- common_job_properties.setTopLevelMainJobProperties(delegate)
-
- // Run job in postcommit every 6 hours, don't trigger every push.
- common_job_properties.setPostCommit(
- delegate,
- '0 */6 * * *',
- false,
- 'commits@beam.apache.org')
-
- // Allows triggering this build against pull requests.
- common_job_properties.enablePhraseTriggeringFromPullRequest(
- delegate,
- 'Python SDK Performance Test',
- 'Run Python Performance Test')
-
- def pipelineArgs = [
- project: 'apache-beam-testing',
- staging_location: 'gs://temp-storage-for-end-to-end-tests/staging-it',
- temp_location: 'gs://temp-storage-for-end-to-end-tests/temp-it',
- output: 'gs://temp-storage-for-end-to-end-tests/py-it-cloud/output'
- ]
- def pipelineArgList = []
- pipelineArgs.each({
- key, value -> pipelineArgList.add("--$key=$value")
- })
- def pipelineArgsJoined = pipelineArgList.join(',')
-
- def argMap = [
- beam_sdk : 'python',
- benchmarks: 'beam_integration_benchmark',
- beam_it_args: pipelineArgsJoined
- ]
-
- common_job_properties.buildPerformanceTest(delegate, argMap)
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/.test-infra/jenkins/job_beam_PostCommit_Java_JDKVersionsTest.groovy
----------------------------------------------------------------------
diff --git a/.test-infra/jenkins/job_beam_PostCommit_Java_JDKVersionsTest.groovy b/.test-infra/jenkins/job_beam_PostCommit_Java_JDKVersionsTest.groovy
index df0a2c7..f23e741 100644
--- a/.test-infra/jenkins/job_beam_PostCommit_Java_JDKVersionsTest.groovy
+++ b/.test-infra/jenkins/job_beam_PostCommit_Java_JDKVersionsTest.groovy
@@ -37,8 +37,6 @@ matrixJob('beam_PostCommit_Java_JDK_Versions_Test') {
common_job_properties.setPostCommit(
delegate,
'0 */6 * * *',
- false,
- '', // TODO: Remove last two args once test is stable again.
false)
// Allows triggering this build against pull requests.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/.test-infra/jenkins/job_beam_PostCommit_Java_MavenInstall_Windows.groovy
----------------------------------------------------------------------
diff --git a/.test-infra/jenkins/job_beam_PostCommit_Java_MavenInstall_Windows.groovy b/.test-infra/jenkins/job_beam_PostCommit_Java_MavenInstall_Windows.groovy
index 6ef272c..f781b4e 100644
--- a/.test-infra/jenkins/job_beam_PostCommit_Java_MavenInstall_Windows.groovy
+++ b/.test-infra/jenkins/job_beam_PostCommit_Java_MavenInstall_Windows.groovy
@@ -32,8 +32,7 @@ mavenJob('beam_PostCommit_Java_MavenInstall_Windows') {
common_job_properties.setMavenConfig(delegate, 'Maven 3.3.3 (Windows)')
// Sets that this is a PostCommit job.
- // TODO(BEAM-1042, BEAM-1045, BEAM-2269, BEAM-2299) Turn notifications back on once fixed.
- common_job_properties.setPostCommit(delegate, '0 */6 * * *', false, '', false)
+ common_job_properties.setPostCommit(delegate, '0 */6 * * *', false)
// Allows triggering this build against pull requests.
common_job_properties.enablePhraseTriggeringFromPullRequest(
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/.test-infra/jenkins/job_beam_PreCommit_Website_Merge.groovy
----------------------------------------------------------------------
diff --git a/.test-infra/jenkins/job_beam_PreCommit_Website_Merge.groovy b/.test-infra/jenkins/job_beam_PreCommit_Website_Merge.groovy
deleted file mode 100644
index 0e2ae3f..0000000
--- a/.test-infra/jenkins/job_beam_PreCommit_Website_Merge.groovy
+++ /dev/null
@@ -1,59 +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.
- */
-
-import common_job_properties
-
-// Defines a job.
-job('beam_PreCommit_Website_Merge') {
- description('Runs website tests for mergebot.')
-
- // Set common parameters.
- common_job_properties.setTopLevelWebsiteJobProperties(delegate, 'mergebot')
-
- triggers {
- githubPush()
- }
-
- steps {
- // Run the following shell script as a build step.
- shell '''
- # Install RVM per instructions at https://rvm.io/rvm/install.
- RVM_GPG_KEY=409B6B1796C275462A1703113804BB82D39DC0E3
- gpg --keyserver hkp://keys.gnupg.net --recv-keys $RVM_GPG_KEY
-
- \\curl -sSL https://get.rvm.io | bash
- source /home/jenkins/.rvm/scripts/rvm
-
- # Install Ruby.
- RUBY_VERSION_NUM=2.3.0
- rvm install ruby $RUBY_VERSION_NUM --autolibs=read-only
-
- # Install Bundler gem
- PATH=~/.gem/ruby/$RUBY_VERSION_NUM/bin:$PATH
- GEM_PATH=~/.gem/ruby/$RUBY_VERSION_NUM/:$GEM_PATH
- gem install bundler --user-install
-
- # Install all needed gems.
- bundle install --path ~/.gem/
-
- # Build the new site and test it.
- rm -fr ./content/
- bundle exec rake test
- '''.stripIndent().trim()
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/examples/java/pom.xml
----------------------------------------------------------------------
diff --git a/examples/java/pom.xml b/examples/java/pom.xml
index ae64a79..701e4fe 100644
--- a/examples/java/pom.xml
+++ b/examples/java/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-examples-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
@@ -34,6 +34,10 @@
<packaging>jar</packaging>
+ <properties>
+ <spark.version>1.6.2</spark.version>
+ </properties>
+
<profiles>
<!--
@@ -62,12 +66,6 @@
<groupId>org.apache.beam</groupId>
<artifactId>beam-runners-apex</artifactId>
<scope>runtime</scope>
- <exclusions>
- <exclusion>
- <groupId>javax.servlet</groupId>
- <artifactId>servlet-api</artifactId>
- </exclusion>
- </exclusions>
</dependency>
<!--
Apex depends on httpclient version 4.3.5, project has a transitive dependency to httpclient 4.0.1 from
@@ -97,12 +95,6 @@
<groupId>org.apache.beam</groupId>
<artifactId>beam-runners-flink_2.10</artifactId>
<scope>runtime</scope>
- <exclusions>
- <exclusion>
- <groupId>javax.servlet</groupId>
- <artifactId>servlet-api</artifactId>
- </exclusion>
- </exclusions>
</dependency>
</dependencies>
</profile>
@@ -124,11 +116,13 @@
<dependency>
<groupId>org.apache.spark</groupId>
<artifactId>spark-streaming_2.10</artifactId>
+ <version>${spark.version}</version>
<scope>runtime</scope>
</dependency>
<dependency>
<groupId>org.apache.spark</groupId>
<artifactId>spark-core_2.10</artifactId>
+ <version>${spark.version}</version>
<scope>runtime</scope>
<exclusions>
<exclusion>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/examples/java/src/main/java/org/apache/beam/examples/common/WriteOneFilePerWindow.java
----------------------------------------------------------------------
diff --git a/examples/java/src/main/java/org/apache/beam/examples/common/WriteOneFilePerWindow.java b/examples/java/src/main/java/org/apache/beam/examples/common/WriteOneFilePerWindow.java
index 49865ba..5e6df9c 100644
--- a/examples/java/src/main/java/org/apache/beam/examples/common/WriteOneFilePerWindow.java
+++ b/examples/java/src/main/java/org/apache/beam/examples/common/WriteOneFilePerWindow.java
@@ -17,12 +17,11 @@
*/
package org.apache.beam.examples.common;
-import static com.google.common.base.MoreObjects.firstNonNull;
+import static com.google.common.base.Verify.verifyNotNull;
import javax.annotation.Nullable;
import org.apache.beam.sdk.io.FileBasedSink;
import org.apache.beam.sdk.io.FileBasedSink.FilenamePolicy;
-import org.apache.beam.sdk.io.FileBasedSink.OutputFileHints;
import org.apache.beam.sdk.io.TextIO;
import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions;
import org.apache.beam.sdk.io.fs.ResourceId;
@@ -54,12 +53,22 @@ public class WriteOneFilePerWindow extends PTransform<PCollection<String>, PDone
@Override
public PDone expand(PCollection<String> input) {
+ // filenamePrefix may contain a directory and a filename component. Pull out only the filename
+ // component from that path for the PerWindowFiles.
+ String prefix = "";
ResourceId resource = FileBasedSink.convertToFileResourceIfPossible(filenamePrefix);
- TextIO.Write write =
- TextIO.write()
- .to(new PerWindowFiles(resource))
- .withTempDirectory(resource.getCurrentDirectory())
- .withWindowedWrites();
+ if (!resource.isDirectory()) {
+ prefix = verifyNotNull(
+ resource.getFilename(),
+ "A non-directory resource should have a non-null filename: %s",
+ resource);
+ }
+
+
+ TextIO.Write write = TextIO.write()
+ .to(resource.getCurrentDirectory())
+ .withFilenamePolicy(new PerWindowFiles(prefix))
+ .withWindowedWrites();
if (numShards != null) {
write = write.withNumShards(numShards);
}
@@ -74,36 +83,31 @@ public class WriteOneFilePerWindow extends PTransform<PCollection<String>, PDone
*/
public static class PerWindowFiles extends FilenamePolicy {
- private final ResourceId baseFilename;
+ private final String prefix;
- public PerWindowFiles(ResourceId baseFilename) {
- this.baseFilename = baseFilename;
+ public PerWindowFiles(String prefix) {
+ this.prefix = prefix;
}
public String filenamePrefixForWindow(IntervalWindow window) {
- String prefix =
- baseFilename.isDirectory() ? "" : firstNonNull(baseFilename.getFilename(), "");
return String.format("%s-%s-%s",
prefix, FORMATTER.print(window.start()), FORMATTER.print(window.end()));
}
@Override
- public ResourceId windowedFilename(WindowedContext context, OutputFileHints outputFileHints) {
+ public ResourceId windowedFilename(
+ ResourceId outputDirectory, WindowedContext context, String extension) {
IntervalWindow window = (IntervalWindow) context.getWindow();
- String filename =
- String.format(
- "%s-%s-of-%s%s",
- filenamePrefixForWindow(window),
- context.getShardNumber(),
- context.getNumShards(),
- outputFileHints.getSuggestedFilenameSuffix());
- return baseFilename
- .getCurrentDirectory()
- .resolve(filename, StandardResolveOptions.RESOLVE_FILE);
+ String filename = String.format(
+ "%s-%s-of-%s%s",
+ filenamePrefixForWindow(window), context.getShardNumber(), context.getNumShards(),
+ extension);
+ return outputDirectory.resolve(filename, StandardResolveOptions.RESOLVE_FILE);
}
@Override
- public ResourceId unwindowedFilename(Context context, OutputFileHints outputFileHints) {
+ public ResourceId unwindowedFilename(
+ ResourceId outputDirectory, Context context, String extension) {
throw new UnsupportedOperationException("Unsupported.");
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/examples/java/src/test/java/org/apache/beam/examples/WindowedWordCountIT.java
----------------------------------------------------------------------
diff --git a/examples/java/src/test/java/org/apache/beam/examples/WindowedWordCountIT.java b/examples/java/src/test/java/org/apache/beam/examples/WindowedWordCountIT.java
index bec7952..eb7e4c4 100644
--- a/examples/java/src/test/java/org/apache/beam/examples/WindowedWordCountIT.java
+++ b/examples/java/src/test/java/org/apache/beam/examples/WindowedWordCountIT.java
@@ -32,7 +32,6 @@ import java.util.concurrent.ThreadLocalRandom;
import org.apache.beam.examples.common.ExampleUtils;
import org.apache.beam.examples.common.WriteOneFilePerWindow.PerWindowFiles;
import org.apache.beam.sdk.PipelineResult;
-import org.apache.beam.sdk.io.FileBasedSink;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
@@ -150,8 +149,7 @@ public class WindowedWordCountIT {
String outputPrefix = options.getOutput();
- PerWindowFiles filenamePolicy =
- new PerWindowFiles(FileBasedSink.convertToFileResourceIfPossible(outputPrefix));
+ PerWindowFiles filenamePolicy = new PerWindowFiles(outputPrefix);
List<ShardedFile> expectedOutputFiles = Lists.newArrayListWithCapacity(6);
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/examples/java8/pom.xml
----------------------------------------------------------------------
diff --git a/examples/java8/pom.xml b/examples/java8/pom.xml
index 6fd29a4..56295a4 100644
--- a/examples/java8/pom.xml
+++ b/examples/java8/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-examples-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
@@ -35,6 +35,10 @@
<packaging>jar</packaging>
+ <properties>
+ <spark.version>1.6.2</spark.version>
+ </properties>
+
<profiles>
<!--
The direct runner is available by default.
@@ -62,12 +66,6 @@
<groupId>org.apache.beam</groupId>
<artifactId>beam-runners-apex</artifactId>
<scope>runtime</scope>
- <exclusions>
- <exclusion>
- <groupId>javax.servlet</groupId>
- <artifactId>servlet-api</artifactId>
- </exclusion>
- </exclusions>
</dependency>
<!--
Apex depends on httpclient version 4.3.5, project has a transitive dependency to httpclient 4.0.1 from
@@ -97,12 +95,6 @@
<groupId>org.apache.beam</groupId>
<artifactId>beam-runners-flink_2.10</artifactId>
<scope>runtime</scope>
- <exclusions>
- <exclusion>
- <groupId>javax.servlet</groupId>
- <artifactId>servlet-api</artifactId>
- </exclusion>
- </exclusions>
</dependency>
</dependencies>
</profile>
@@ -124,11 +116,13 @@
<dependency>
<groupId>org.apache.spark</groupId>
<artifactId>spark-streaming_2.10</artifactId>
+ <version>${spark.version}</version>
<scope>runtime</scope>
</dependency>
<dependency>
<groupId>org.apache.spark</groupId>
<artifactId>spark-core_2.10</artifactId>
+ <version>${spark.version}</version>
<scope>runtime</scope>
<exclusions>
<exclusion>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/examples/java8/src/main/java/org/apache/beam/examples/complete/game/utils/WriteToText.java
----------------------------------------------------------------------
diff --git a/examples/java8/src/main/java/org/apache/beam/examples/complete/game/utils/WriteToText.java b/examples/java8/src/main/java/org/apache/beam/examples/complete/game/utils/WriteToText.java
index 1d60198..e6c8ddb 100644
--- a/examples/java8/src/main/java/org/apache/beam/examples/complete/game/utils/WriteToText.java
+++ b/examples/java8/src/main/java/org/apache/beam/examples/complete/game/utils/WriteToText.java
@@ -18,6 +18,7 @@
package org.apache.beam.examples.complete.game.utils;
import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Verify.verifyNotNull;
import java.io.Serializable;
import java.util.ArrayList;
@@ -27,7 +28,6 @@ import java.util.TimeZone;
import java.util.stream.Collectors;
import org.apache.beam.sdk.io.FileBasedSink;
import org.apache.beam.sdk.io.FileBasedSink.FilenamePolicy;
-import org.apache.beam.sdk.io.FileBasedSink.OutputFileHints;
import org.apache.beam.sdk.io.TextIO;
import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions;
import org.apache.beam.sdk.io.fs.ResourceId;
@@ -111,12 +111,21 @@ public class WriteToText<InputT>
checkArgument(
input.getWindowingStrategy().getWindowFn().windowCoder() == IntervalWindow.getCoder());
+ // filenamePrefix may contain a directory and a filename component. Pull out only the filename
+ // component from that path for the PerWindowFiles.
+ String prefix = "";
ResourceId resource = FileBasedSink.convertToFileResourceIfPossible(filenamePrefix);
+ if (!resource.isDirectory()) {
+ prefix = verifyNotNull(
+ resource.getFilename(),
+ "A non-directory resource should have a non-null filename: %s",
+ resource);
+ }
return input.apply(
TextIO.write()
- .to(new PerWindowFiles(resource))
- .withTempDirectory(resource.getCurrentDirectory())
+ .to(resource.getCurrentDirectory())
+ .withFilenamePolicy(new PerWindowFiles(prefix))
.withWindowedWrites()
.withNumShards(3));
}
@@ -130,33 +139,31 @@ public class WriteToText<InputT>
*/
protected static class PerWindowFiles extends FilenamePolicy {
- private final ResourceId prefix;
+ private final String prefix;
- public PerWindowFiles(ResourceId prefix) {
+ public PerWindowFiles(String prefix) {
this.prefix = prefix;
}
public String filenamePrefixForWindow(IntervalWindow window) {
- String filePrefix = prefix.isDirectory() ? "" : prefix.getFilename();
- return String.format(
- "%s-%s-%s", filePrefix, formatter.print(window.start()), formatter.print(window.end()));
+ return String.format("%s-%s-%s",
+ prefix, formatter.print(window.start()), formatter.print(window.end()));
}
@Override
- public ResourceId windowedFilename(WindowedContext context, OutputFileHints outputFileHints) {
+ public ResourceId windowedFilename(
+ ResourceId outputDirectory, WindowedContext context, String extension) {
IntervalWindow window = (IntervalWindow) context.getWindow();
- String filename =
- String.format(
- "%s-%s-of-%s%s",
- filenamePrefixForWindow(window),
- context.getShardNumber(),
- context.getNumShards(),
- outputFileHints.getSuggestedFilenameSuffix());
- return prefix.getCurrentDirectory().resolve(filename, StandardResolveOptions.RESOLVE_FILE);
+ String filename = String.format(
+ "%s-%s-of-%s%s",
+ filenamePrefixForWindow(window), context.getShardNumber(), context.getNumShards(),
+ extension);
+ return outputDirectory.resolve(filename, StandardResolveOptions.RESOLVE_FILE);
}
@Override
- public ResourceId unwindowedFilename(Context context, OutputFileHints outputFileHints) {
+ public ResourceId unwindowedFilename(
+ ResourceId outputDirectory, Context context, String extension) {
throw new UnsupportedOperationException("Unsupported.");
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/examples/java8/src/test/java/org/apache/beam/examples/complete/game/LeaderBoardTest.java
----------------------------------------------------------------------
diff --git a/examples/java8/src/test/java/org/apache/beam/examples/complete/game/LeaderBoardTest.java b/examples/java8/src/test/java/org/apache/beam/examples/complete/game/LeaderBoardTest.java
index 611e2b3..745c210 100644
--- a/examples/java8/src/test/java/org/apache/beam/examples/complete/game/LeaderBoardTest.java
+++ b/examples/java8/src/test/java/org/apache/beam/examples/complete/game/LeaderBoardTest.java
@@ -276,8 +276,6 @@ public class LeaderBoardTest implements Serializable {
.addElements(event(TestUser.RED_ONE, 4, Duration.standardMinutes(2)),
event(TestUser.BLUE_TWO, 3, Duration.ZERO),
event(TestUser.BLUE_ONE, 3, Duration.standardMinutes(3)))
- // Move the watermark to the end of the window to output on time
- .advanceWatermarkTo(baseTime.plus(TEAM_WINDOW_DURATION))
// Move the watermark past the end of the allowed lateness plus the end of the window
.advanceWatermarkTo(baseTime.plus(ALLOWED_LATENESS)
.plus(TEAM_WINDOW_DURATION).plus(Duration.standardMinutes(1)))
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/examples/pom.xml
----------------------------------------------------------------------
diff --git a/examples/pom.xml b/examples/pom.xml
index 51f4c35..a7e61dd 100644
--- a/examples/pom.xml
+++ b/examples/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/pom.xml
----------------------------------------------------------------------
diff --git a/pom.xml b/pom.xml
index 987760f..1d8d4b0 100644
--- a/pom.xml
+++ b/pom.xml
@@ -34,7 +34,7 @@
<url>http://beam.apache.org/</url>
<inceptionYear>2016</inceptionYear>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<licenses>
<license>
@@ -101,14 +101,13 @@
<beamSurefireArgline />
<!-- If updating dependencies, please update any relevant javadoc offlineLinks -->
- <apache.commons.compress.version>1.14</apache.commons.compress.version>
- <apache.commons.lang.version>3.6</apache.commons.lang.version>
- <apache.commons.text.version>1.1</apache.commons.text.version>
+ <apache.commons.lang.version>3.5</apache.commons.lang.version>
+ <apache.commons.compress.version>1.9</apache.commons.compress.version>
<apex.kryo.version>2.24.0</apex.kryo.version>
<api-common.version>1.0.0-rc2</api-common.version>
<avro.version>1.8.2</avro.version>
<bigquery.version>v2-rev295-1.22.0</bigquery.version>
- <bigtable.version>0.9.7.1</bigtable.version>
+ <bigtable.version>0.9.6.2</bigtable.version>
<cloudresourcemanager.version>v1-rev6-1.22.0</cloudresourcemanager.version>
<pubsubgrpc.version>0.1.0</pubsubgrpc.version>
<clouddebugger.version>v2-rev8-1.22.0</clouddebugger.version>
@@ -127,14 +126,8 @@
<guava.version>20.0</guava.version>
<grpc.version>1.2.0</grpc.version>
<grpc-google-common-protos.version>0.1.9</grpc-google-common-protos.version>
- <!--
- This is the version of Hadoop used to compile the module that depend on Hadoop.
- This dependency is defined with a provided scope.
- Users must supply their own Hadoop version at runtime.
- -->
- <hadoop.version>2.7.3</hadoop.version>
<hamcrest.version>1.3</hamcrest.version>
- <jackson.version>2.8.9</jackson.version>
+ <jackson.version>2.8.8</jackson.version>
<findbugs.version>3.0.1</findbugs.version>
<joda.version>2.4</joda.version>
<junit.version>4.12</junit.version>
@@ -144,8 +137,8 @@
<protobuf.version>3.2.0</protobuf.version>
<pubsub.version>v1-rev10-1.22.0</pubsub.version>
<slf4j.version>1.7.14</slf4j.version>
- <spanner.version>0.20.0-beta</spanner.version>
- <spark.version>1.6.3</spark.version>
+ <spanner.version>0.16.0-beta</spanner.version>
+ <spark.version>1.6.2</spark.version>
<spring.version>4.3.5.RELEASE</spring.version>
<stax2.version>3.1.4</stax2.version>
<storage.version>v1-rev71-1.22.0</storage.version>
@@ -159,7 +152,7 @@
<failsafe-plugin.version>2.20</failsafe-plugin.version>
<maven-compiler-plugin.version>3.6.1</maven-compiler-plugin.version>
<maven-dependency-plugin.version>3.0.1</maven-dependency-plugin.version>
- <maven-exec-plugin.version>1.6.0</maven-exec-plugin.version>
+ <maven-exec-plugin.version>1.4.0</maven-exec-plugin.version>
<maven-jar-plugin.version>3.0.2</maven-jar-plugin.version>
<maven-resources-plugin.version>3.0.2</maven-resources-plugin.version>
<maven-shade-plugin.version>3.0.0</maven-shade-plugin.version>
@@ -167,7 +160,6 @@
<compiler.error.flag>-Werror</compiler.error.flag>
<compiler.default.pkginfo.flag>-Xpkginfo:always</compiler.default.pkginfo.flag>
<compiler.default.exclude>nothing</compiler.default.exclude>
- <gax-grpc.version>0.20.0</gax-grpc.version>
</properties>
<packaging>pom</packaging>
@@ -429,18 +421,6 @@
<dependency>
<groupId>org.apache.beam</groupId>
- <artifactId>beam-sdks-java-io-amqp</artifactId>
- <version>${project.version}</version>
- </dependency>
-
- <dependency>
- <groupId>org.apache.beam</groupId>
- <artifactId>beam-sdks-java-io-cassandra</artifactId>
- <version>${project.version}</version>
- </dependency>
-
- <dependency>
- <groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-elasticsearch</artifactId>
<version>${project.version}</version>
</dependency>
@@ -478,12 +458,6 @@
<dependency>
<groupId>org.apache.beam</groupId>
- <artifactId>beam-sdks-java-io-hcatalog</artifactId>
- <version>${project.version}</version>
- </dependency>
-
- <dependency>
- <groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-jdbc</artifactId>
<version>${project.version}</version>
</dependency>
@@ -538,13 +512,6 @@
<dependency>
<groupId>org.apache.beam</groupId>
- <artifactId>beam-runners-core-java</artifactId>
- <version>${project.version}</version>
- <type>test-jar</type>
- </dependency>
-
- <dependency>
- <groupId>org.apache.beam</groupId>
<artifactId>beam-runners-direct-java</artifactId>
<version>${project.version}</version>
</dependency>
@@ -598,12 +565,6 @@
</dependency>
<dependency>
- <groupId>org.apache.commons</groupId>
- <artifactId>commons-text</artifactId>
- <version>${apache.commons.text.version}</version>
- </dependency>
-
- <dependency>
<groupId>io.grpc</groupId>
<artifactId>grpc-all</artifactId>
<version>${grpc.version}</version>
@@ -664,12 +625,6 @@
</dependency>
<dependency>
- <groupId>com.google.api</groupId>
- <artifactId>gax-grpc</artifactId>
- <version>${gax-grpc.version}</version>
- </dependency>
-
- <dependency>
<groupId>com.google.api-client</groupId>
<artifactId>google-api-client</artifactId>
<version>${google-clients.version}</version>
@@ -884,11 +839,6 @@
</dependency>
<dependency>
- <groupId>com.google.cloud</groupId>
- <artifactId>google-cloud-core-grpc</artifactId>
- <version>${grpc.version}</version>
- </dependency>
- <dependency>
<groupId>com.google.cloud.bigtable</groupId>
<artifactId>bigtable-protos</artifactId>
<version>${bigtable.version}</version>
@@ -1100,42 +1050,6 @@
<version>${snappy-java.version}</version>
</dependency>
- <dependency>
- <groupId>org.apache.hadoop</groupId>
- <artifactId>hadoop-client</artifactId>
- <version>${hadoop.version}</version>
- </dependency>
-
- <dependency>
- <groupId>org.apache.hadoop</groupId>
- <artifactId>hadoop-common</artifactId>
- <version>${hadoop.version}</version>
- </dependency>
-
- <dependency>
- <groupId>org.apache.hadoop</groupId>
- <artifactId>hadoop-mapreduce-client-core</artifactId>
- <version>${hadoop.version}</version>
- </dependency>
-
- <dependency>
- <groupId>org.apache.spark</groupId>
- <artifactId>spark-core_2.10</artifactId>
- <version>${spark.version}</version>
- </dependency>
-
- <dependency>
- <groupId>org.apache.spark</groupId>
- <artifactId>spark-streaming_2.10</artifactId>
- <version>${spark.version}</version>
- </dependency>
-
- <dependency>
- <groupId>org.apache.spark</groupId>
- <artifactId>spark-network-common_2.10</artifactId>
- <version>${spark.version}</version>
- </dependency>
-
<!-- Testing -->
<dependency>
@@ -1205,27 +1119,6 @@
<scope>test</scope>
</dependency>
- <dependency>
- <groupId>org.apache.hadoop</groupId>
- <artifactId>hadoop-minicluster</artifactId>
- <version>${hadoop.version}</version>
- <scope>test</scope>
- </dependency>
-
- <dependency>
- <groupId>org.apache.hadoop</groupId>
- <artifactId>hadoop-hdfs</artifactId>
- <version>${hadoop.version}</version>
- <scope>test</scope>
- </dependency>
-
- <dependency>
- <groupId>org.apache.hadoop</groupId>
- <artifactId>hadoop-hdfs</artifactId>
- <version>${hadoop.version}</version>
- <classifier>tests</classifier>
- <scope>test</scope>
- </dependency>
</dependencies>
</dependencyManagement>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/apex/pom.xml
----------------------------------------------------------------------
diff --git a/runners/apex/pom.xml b/runners/apex/pom.xml
index fd5aafb..4a36bec 100644
--- a/runners/apex/pom.xml
+++ b/runners/apex/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-runners-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
@@ -75,13 +75,6 @@
<artifactId>apex-engine</artifactId>
<version>${apex.core.version}</version>
<scope>runtime</scope>
- <exclusions>
- <!-- Fix build on JDK-9 -->
- <exclusion>
- <groupId>jdk.tools</groupId>
- <artifactId>jdk.tools</artifactId>
- </exclusion>
- </exclusions>
</dependency>
<!-- Beam -->
@@ -191,13 +184,6 @@
<type>test-jar</type>
<scope>test</scope>
</dependency>
-
- <dependency>
- <groupId>org.apache.beam</groupId>
- <artifactId>beam-runners-core-java</artifactId>
- <type>test-jar</type>
- <scope>test</scope>
- </dependency>
</dependencies>
<build>
@@ -263,12 +249,12 @@
<configuration>
<ignoredUsedUndeclaredDependencies>
<ignoredUsedUndeclaredDependency>org.apache.apex:apex-api:jar:${apex.core.version}</ignoredUsedUndeclaredDependency>
- <ignoredUsedUndeclaredDependency>org.apache.commons:commons-lang3::${apache.commons.lang.version}</ignoredUsedUndeclaredDependency>
+ <ignoredUsedUndeclaredDependency>org.apache.commons:commons-lang3::3.1</ignoredUsedUndeclaredDependency>
<ignoredUsedUndeclaredDependency>commons-io:commons-io:jar:2.4</ignoredUsedUndeclaredDependency>
<ignoredUsedUndeclaredDependency>com.esotericsoftware.kryo:kryo::${apex.kryo.version}</ignoredUsedUndeclaredDependency>
<ignoredUsedUndeclaredDependency>com.datatorrent:netlet::1.3.0</ignoredUsedUndeclaredDependency>
<ignoredUsedUndeclaredDependency>org.slf4j:slf4j-api:jar:1.7.14</ignoredUsedUndeclaredDependency>
- <ignoredUsedUndeclaredDependency>org.apache.hadoop:hadoop-common:jar:${hadoop.version}</ignoredUsedUndeclaredDependency>
+ <ignoredUsedUndeclaredDependency>org.apache.hadoop:hadoop-common:jar:2.6.0</ignoredUsedUndeclaredDependency>
<ignoredUsedUndeclaredDependency>joda-time:joda-time:jar:2.4</ignoredUsedUndeclaredDependency>
<ignoredUsedUndeclaredDependency>com.google.guava:guava:jar:20.0</ignoredUsedUndeclaredDependency>
</ignoredUsedUndeclaredDependencies>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/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 fd0a1c9..c595b3f 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
@@ -62,6 +62,8 @@ 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.ParDo.MultiOutput;
+import org.apache.beam.sdk.transforms.View;
+import org.apache.beam.sdk.transforms.View.AsIterable;
import org.apache.beam.sdk.transforms.View.CreatePCollectionView;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionTuple;
@@ -212,7 +214,7 @@ public class ApexRunner extends PipelineRunner<ApexRunnerResult> {
* @param <ViewT> The type associated with the {@link PCollectionView} used as a side input
*/
public static class CreateApexPCollectionView<ElemT, ViewT>
- extends PTransform<PCollection<ElemT>, PCollection<ElemT>> {
+ extends PTransform<PCollection<List<ElemT>>, PCollectionView<ViewT>> {
private static final long serialVersionUID = 1L;
private PCollectionView<ViewT> view;
@@ -226,13 +228,7 @@ public class ApexRunner extends PipelineRunner<ApexRunnerResult> {
}
@Override
- public PCollection<ElemT> expand(PCollection<ElemT> input) {
- return PCollection.<ElemT>createPrimitiveOutputInternal(
- input.getPipeline(), input.getWindowingStrategy(), input.isBounded())
- .setCoder(input.getCoder());
- }
-
- public PCollectionView<ViewT> getView() {
+ public PCollectionView<ViewT> expand(PCollection<List<ElemT>> input) {
return view;
}
}
@@ -245,7 +241,7 @@ public class ApexRunner extends PipelineRunner<ApexRunnerResult> {
}
private static class StreamingWrapSingletonInList<T>
- extends PTransform<PCollection<T>, PCollection<T>> {
+ extends PTransform<PCollection<T>, PCollectionView<T>> {
private static final long serialVersionUID = 1L;
CreatePCollectionView<T, T> transform;
@@ -258,11 +254,10 @@ public class ApexRunner extends PipelineRunner<ApexRunnerResult> {
}
@Override
- public PCollection<T> expand(PCollection<T> input) {
- input
+ public PCollectionView<T> expand(PCollection<T> input) {
+ return input
.apply(ParDo.of(new WrapAsList<T>()))
- .apply(CreateApexPCollectionView.<List<T>, T>of(transform.getView()));
- return input;
+ .apply(CreateApexPCollectionView.<T, T>of(transform.getView()));
}
@Override
@@ -272,12 +267,15 @@ public class ApexRunner extends PipelineRunner<ApexRunnerResult> {
static class Factory<T>
extends SingleInputOutputOverrideFactory<
- PCollection<T>, PCollection<T>,
+ PCollection<T>, PCollectionView<T>,
CreatePCollectionView<T, T>> {
@Override
- public PTransformReplacement<PCollection<T>, PCollection<T>> getReplacementTransform(
- AppliedPTransform<PCollection<T>, PCollection<T>, CreatePCollectionView<T, T>>
- transform) {
+ public PTransformReplacement<PCollection<T>, PCollectionView<T>>
+ getReplacementTransform(
+ AppliedPTransform<
+ PCollection<T>, PCollectionView<T>,
+ CreatePCollectionView<T, T>>
+ transform) {
return PTransformReplacement.of(
PTransformReplacements.getSingletonMainInput(transform),
new StreamingWrapSingletonInList<>(transform.getTransform()));
@@ -286,19 +284,18 @@ public class ApexRunner extends PipelineRunner<ApexRunnerResult> {
}
private static class StreamingViewAsIterable<T>
- extends PTransform<PCollection<T>, PCollection<T>> {
+ extends PTransform<PCollection<T>, PCollectionView<Iterable<T>>> {
private static final long serialVersionUID = 1L;
- private final PCollectionView<Iterable<T>> view;
- private StreamingViewAsIterable(PCollectionView<Iterable<T>> view) {
- this.view = view;
- }
+ private StreamingViewAsIterable() {}
@Override
- public PCollection<T> expand(PCollection<T> input) {
- return ((PCollection<T>)
- input.apply(Combine.globally(new Concatenate<T>()).withoutDefaults()))
- .apply(CreateApexPCollectionView.<T, Iterable<T>>of(view));
+ public PCollectionView<Iterable<T>> expand(PCollection<T> input) {
+ PCollectionView<Iterable<T>> view =
+ PCollectionViews.iterableView(input, input.getWindowingStrategy(), input.getCoder());
+
+ return input.apply(Combine.globally(new Concatenate<T>()).withoutDefaults())
+ .apply(CreateApexPCollectionView.<T, Iterable<T>> of(view));
}
@Override
@@ -308,17 +305,15 @@ public class ApexRunner extends PipelineRunner<ApexRunnerResult> {
static class Factory<T>
extends SingleInputOutputOverrideFactory<
- PCollection<T>, PCollection<T>, CreatePCollectionView<T, Iterable<T>>> {
+ PCollection<T>, PCollectionView<Iterable<T>>, View.AsIterable<T>> {
@Override
- public PTransformReplacement<PCollection<T>, PCollection<T>>
+ public PTransformReplacement<PCollection<T>, PCollectionView<Iterable<T>>>
getReplacementTransform(
- AppliedPTransform<
- PCollection<T>, PCollection<T>,
- CreatePCollectionView<T, Iterable<T>>>
+ AppliedPTransform<PCollection<T>, PCollectionView<Iterable<T>>, AsIterable<T>>
transform) {
return PTransformReplacement.of(
PTransformReplacements.getSingletonMainInput(transform),
- new StreamingViewAsIterable<T>(transform.getTransform().getView()));
+ new StreamingViewAsIterable<T>());
}
}
}
@@ -381,7 +376,7 @@ public class ApexRunner extends PipelineRunner<ApexRunnerResult> {
AppliedPTransform<PCollection<InputT>, PCollectionTuple, MultiOutput<InputT, OutputT>>
transform) {
return PTransformReplacement.of(PTransformReplacements.getSingletonMainInput(transform),
- SplittableParDo.forJavaParDo(transform.getTransform()));
+ new SplittableParDo<>(transform.getTransform()));
}
@Override
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/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 02f53ec..bda074b 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
@@ -39,6 +39,7 @@ import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.View.CreatePCollectionView;
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PValue;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -153,6 +154,7 @@ public class ApexPipelineTranslator extends Pipeline.PipelineVisitor.Defaults {
unboundedSource, true, context.getPipelineOptions());
context.addOperator(operator, operator.output);
}
+
}
private static class CreateApexPCollectionViewTranslator<ElemT, ViewT>
@@ -160,10 +162,11 @@ public class ApexPipelineTranslator extends Pipeline.PipelineVisitor.Defaults {
private static final long serialVersionUID = 1L;
@Override
- public void translate(
- CreateApexPCollectionView<ElemT, ViewT> transform, TranslationContext context) {
- context.addView(transform.getView());
- LOG.debug("view {}", transform.getView().getName());
+ public void translate(CreateApexPCollectionView<ElemT, ViewT> transform,
+ TranslationContext context) {
+ PCollectionView<ViewT> view = (PCollectionView<ViewT>) context.getOutput();
+ context.addView(view);
+ LOG.debug("view {}", view.getName());
}
}
@@ -174,8 +177,9 @@ public class ApexPipelineTranslator extends Pipeline.PipelineVisitor.Defaults {
@Override
public void translate(
CreatePCollectionView<ElemT, ViewT> transform, TranslationContext context) {
- context.addView(transform.getView());
- LOG.debug("view {}", transform.getView().getName());
+ PCollectionView<ViewT> view = (PCollectionView<ViewT>) context.getOutput();
+ context.addView(view);
+ LOG.debug("view {}", view.getName());
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/TranslationContext.java
----------------------------------------------------------------------
diff --git a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/TranslationContext.java b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/TranslationContext.java
index 94d13e1..aff3863 100644
--- a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/TranslationContext.java
+++ b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/TranslationContext.java
@@ -34,7 +34,6 @@ import org.apache.beam.runners.apex.translation.utils.ApexStateInternals;
import org.apache.beam.runners.apex.translation.utils.ApexStateInternals.ApexStateBackend;
import org.apache.beam.runners.apex.translation.utils.ApexStreamTuple;
import org.apache.beam.runners.apex.translation.utils.CoderAdapterStreamCodec;
-import org.apache.beam.runners.core.construction.TransformInputs;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.transforms.PTransform;
@@ -94,8 +93,7 @@ class TranslationContext {
}
public <InputT extends PValue> InputT getInput() {
- return (InputT)
- Iterables.getOnlyElement(TransformInputs.nonAdditionalInputs(getCurrentTransform()));
+ return (InputT) Iterables.getOnlyElement(getCurrentTransform().getInputs().values());
}
public Map<TupleTag<?>, PValue> getOutputs() {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/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 c3cbab2..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
@@ -359,7 +359,10 @@ public class ApexParDoOperator<InputT, OutputT> extends BaseOperator implements
}
}
if (sideInputs.isEmpty()) {
- outputWatermark(mark);
+ if (traceTuples) {
+ LOG.debug("\nemitting watermark {}\n", mark);
+ }
+ output.emit(mark);
return;
}
@@ -367,20 +370,10 @@ public class ApexParDoOperator<InputT, OutputT> extends BaseOperator implements
Math.min(pushedBackWatermark.get(), currentInputWatermark);
if (potentialOutputWatermark > currentOutputWatermark) {
currentOutputWatermark = potentialOutputWatermark;
- outputWatermark(ApexStreamTuple.WatermarkTuple.of(currentOutputWatermark));
- }
- }
-
- private void outputWatermark(ApexStreamTuple.WatermarkTuple<?> mark) {
- if (traceTuples) {
- LOG.debug("\nemitting {}\n", mark);
- }
- output.emit(mark);
- if (!additionalOutputPortMapping.isEmpty()) {
- for (DefaultOutputPort<ApexStreamTuple<?>> additionalOutput :
- additionalOutputPortMapping.values()) {
- additionalOutput.emit(mark);
+ if (traceTuples) {
+ LOG.debug("\nemitting watermark {}\n", currentOutputWatermark);
}
+ output.emit(ApexStreamTuple.WatermarkTuple.of(currentOutputWatermark));
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/apex/src/test/java/org/apache/beam/runners/apex/examples/WordCountTest.java
----------------------------------------------------------------------
diff --git a/runners/apex/src/test/java/org/apache/beam/runners/apex/examples/WordCountTest.java b/runners/apex/src/test/java/org/apache/beam/runners/apex/examples/WordCountTest.java
index ba75746..e76096e 100644
--- a/runners/apex/src/test/java/org/apache/beam/runners/apex/examples/WordCountTest.java
+++ b/runners/apex/src/test/java/org/apache/beam/runners/apex/examples/WordCountTest.java
@@ -123,15 +123,11 @@ public class WordCountTest {
options.setInputFile(new File(inputFile).getAbsolutePath());
String outputFilePrefix = "target/wordcountresult.txt";
options.setOutput(outputFilePrefix);
+ WordCountTest.main(TestPipeline.convertToArgs(options));
File outFile1 = new File(outputFilePrefix + "-00000-of-00002");
File outFile2 = new File(outputFilePrefix + "-00001-of-00002");
- Assert.assertTrue(!outFile1.exists() || outFile1.delete());
- Assert.assertTrue(!outFile2.exists() || outFile2.delete());
-
- WordCountTest.main(TestPipeline.convertToArgs(options));
-
- Assert.assertTrue("result files exist", outFile1.exists() && outFile2.exists());
+ Assert.assertTrue(outFile1.exists() && outFile2.exists());
HashSet<String> results = new HashSet<>();
results.addAll(FileUtils.readLines(outFile1));
results.addAll(FileUtils.readLines(outFile2));
[17/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java
index 78340f3..f1eb7c0 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java
@@ -22,38 +22,22 @@ import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;
import com.google.auto.value.AutoValue;
-import com.google.common.annotations.VisibleForTesting;
-import java.io.IOException;
-import java.util.concurrent.ThreadLocalRandom;
import javax.annotation.Nullable;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.annotations.Experimental.Kind;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.VoidCoder;
-import org.apache.beam.sdk.io.CompressedSource.CompressionMode;
-import org.apache.beam.sdk.io.DefaultFilenamePolicy.Params;
-import org.apache.beam.sdk.io.FileBasedSink.DynamicDestinations;
import org.apache.beam.sdk.io.FileBasedSink.FilenamePolicy;
import org.apache.beam.sdk.io.FileBasedSink.WritableByteChannelFactory;
import org.apache.beam.sdk.io.Read.Bounded;
-import org.apache.beam.sdk.io.fs.MatchResult;
-import org.apache.beam.sdk.io.fs.MatchResult.Metadata;
-import org.apache.beam.sdk.io.fs.MatchResult.Status;
import org.apache.beam.sdk.io.fs.ResourceId;
-import org.apache.beam.sdk.io.range.OffsetRange;
import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider;
import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
-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.Reshuffle;
import org.apache.beam.sdk.transforms.SerializableFunction;
-import org.apache.beam.sdk.transforms.SerializableFunctions;
-import org.apache.beam.sdk.transforms.Values;
import org.apache.beam.sdk.transforms.display.DisplayData;
-import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PDone;
@@ -63,14 +47,13 @@ import org.apache.beam.sdk.values.PDone;
*
* <p>To read a {@link PCollection} from one or more text files, use {@code TextIO.read()} to
* instantiate a transform and use {@link TextIO.Read#from(String)} to specify the path of the
- * file(s) to be read. Alternatively, if the filenames to be read are themselves in a
- * {@link PCollection}, apply {@link TextIO#readAll()}.
+ * file(s) to be read.
*
* <p>{@link TextIO.Read} returns a {@link PCollection} of {@link String Strings}, each
* corresponding to one line of an input UTF-8 text file (split into lines delimited by '\n', '\r',
* or '\r\n').
*
- * <p>Example 1: reading a file or filepattern.
+ * <p>Example:
*
* <pre>{@code
* Pipeline p = ...;
@@ -79,24 +62,22 @@ import org.apache.beam.sdk.values.PDone;
* PCollection<String> lines = p.apply(TextIO.read().from("/local/path/to/file.txt"));
* }</pre>
*
- * <p>Example 2: reading a PCollection of filenames.
- *
- * <pre>{@code
- * Pipeline p = ...;
- *
- * // E.g. the filenames might be computed from other data in the pipeline, or
- * // read from a data source.
- * PCollection<String> filenames = ...;
- *
- * // Read all files in the collection.
- * PCollection<String> lines = filenames.apply(TextIO.readAll());
- * }</pre>
- *
* <p>To write a {@link PCollection} to one or more text files, use {@code TextIO.write()}, using
* {@link TextIO.Write#to(String)} to specify the output prefix of the files to write.
*
- * <p>For example:
+ * <p>By default, all input is put into the global window before writing. If per-window writes are
+ * desired - for example, when using a streaming runner -
+ * {@link TextIO.Write#withWindowedWrites()} will cause windowing and triggering to be
+ * preserved. When producing windowed writes, the number of output shards must be set explicitly
+ * using {@link TextIO.Write#withNumShards(int)}; some runners may set this for you to a
+ * runner-chosen value, so you may need not set it yourself. A {@link FilenamePolicy} can also be
+ * set in case you need better control over naming files created by unique windows.
+ * {@link DefaultFilenamePolicy} policy for producing unique filenames might not be appropriate
+ * for your use case.
+ *
+ * <p>Any existing files with the same names as generated output files will be overwritten.
*
+ * <p>For example:
* <pre>{@code
* // A simple Write to a local file (only runs locally):
* PCollection<String> lines = ...;
@@ -104,49 +85,10 @@ import org.apache.beam.sdk.values.PDone;
*
* // Same as above, only with Gzip compression:
* PCollection<String> lines = ...;
- * lines.apply(TextIO.write().to("/path/to/file.txt"))
+ * lines.apply(TextIO.write().to("/path/to/file.txt"));
* .withSuffix(".txt")
* .withWritableByteChannelFactory(FileBasedSink.CompressionType.GZIP));
* }</pre>
- *
- * <p>By default, all input is put into the global window before writing. If per-window writes are
- * desired - for example, when using a streaming runner - {@link TextIO.Write#withWindowedWrites()}
- * will cause windowing and triggering to be preserved. When producing windowed writes with a
- * streaming runner that supports triggers, the number of output shards must be set explicitly using
- * {@link TextIO.Write#withNumShards(int)}; some runners may set this for you to a runner-chosen
- * value, so you may need not set it yourself. If setting an explicit template using {@link
- * TextIO.Write#withShardNameTemplate(String)}, make sure that the template contains placeholders
- * for the window and the pane; W is expanded into the window text, and P into the pane; the default
- * template will include both the window and the pane in the filename.
- *
- * <p>If you want better control over how filenames are generated than the default policy allows, a
- * custom {@link FilenamePolicy} can also be set using {@link TextIO.Write#to(FilenamePolicy)}.
- *
- * <p>TextIO also supports dynamic, value-dependent file destinations. The most general form of this
- * is done via {@link TextIO.Write#to(DynamicDestinations)}. A {@link DynamicDestinations} class
- * allows you to convert any input value into a custom destination object, and map that destination
- * object to a {@link FilenamePolicy}. This allows using different filename policies (or more
- * commonly, differently-configured instances of the same policy) based on the input record. Often
- * this is used in conjunction with {@link TextIO#writeCustomType(SerializableFunction)}, which
- * allows your {@link DynamicDestinations} object to examine the input type and takes a format
- * function to convert that type to a string for writing.
- *
- * <p>A convenience shortcut is provided for the case where the default naming policy is used, but
- * different configurations of this policy are wanted based on the input record. Default naming
- * policies can be configured using the {@link DefaultFilenamePolicy.Params} object.
- *
- * <pre>{@code
- * PCollection<UserEvent>> lines = ...;
- * lines.apply(TextIO.<UserEvent>writeCustomType(new FormatEvent())
- * .to(new SerializableFunction<UserEvent, Params>() {
- * public String apply(UserEvent value) {
- * return new Params().withBaseFilename(baseDirectory + "/" + value.country());
- * }
- * }),
- * new Params().withBaseFilename(baseDirectory + "/empty");
- * }</pre>
- *
- * <p>Any existing files with the same names as generated output files will be overwritten.
*/
public class TextIO {
/**
@@ -158,54 +100,16 @@ public class TextIO {
}
/**
- * A {@link PTransform} that works like {@link #read}, but reads each file in a {@link
- * PCollection} of filepatterns.
- *
- * <p>Can be applied to both bounded and unbounded {@link PCollection PCollections}, so this is
- * suitable for reading a {@link PCollection} of filepatterns arriving as a stream. However, every
- * filepattern is expanded once at the moment it is processed, rather than watched for new files
- * matching the filepattern to appear. Likewise, every file is read once, rather than watched for
- * new entries.
- */
- public static ReadAll readAll() {
- return new AutoValue_TextIO_ReadAll.Builder()
- .setCompressionType(CompressionType.AUTO)
- // 64MB is a reasonable value that allows to amortize the cost of opening files,
- // but is not so large as to exhaust a typical runner's maximum amount of output per
- // ProcessElement call.
- .setDesiredBundleSizeBytes(64 * 1024 * 1024L)
- .build();
- }
-
- /**
* A {@link PTransform} that writes a {@link PCollection} to a text file (or multiple text files
* matching a sharding pattern), with each element of the input collection encoded into its own
* line.
*/
public static Write write() {
- return new TextIO.Write();
- }
-
- /**
- * A {@link PTransform} that writes a {@link PCollection} to a text file (or multiple text files
- * matching a sharding pattern), with each element of the input collection encoded into its own
- * line.
- *
- * <p>This version allows you to apply {@link TextIO} writes to a PCollection of a custom type
- * {@link T}, along with a format function that converts the input type {@link T} to the String
- * that will be written to the file. The advantage of this is it allows a user-provided {@link
- * DynamicDestinations} object, set via {@link Write#to(DynamicDestinations)} to examine the
- * user's custom type when choosing a destination.
- */
- public static <T> TypedWrite<T> writeCustomType(SerializableFunction<T, String> formatFunction) {
- return new AutoValue_TextIO_TypedWrite.Builder<T>()
+ return new AutoValue_TextIO_Write.Builder()
.setFilenamePrefix(null)
- .setTempDirectory(null)
.setShardTemplate(null)
.setFilenameSuffix(null)
.setFilenamePolicy(null)
- .setDynamicDestinations(null)
- .setFormatFunction(formatFunction)
.setWritableByteChannelFactory(FileBasedSink.CompressionType.UNCOMPRESSED)
.setWindowedWrites(false)
.setNumShards(0)
@@ -274,34 +178,29 @@ public class TextIO {
// Helper to create a source specific to the requested compression type.
protected FileBasedSource<String> getSource() {
- return wrapWithCompression(new TextSource(getFilepattern()), getCompressionType());
- }
-
- private static FileBasedSource<String> wrapWithCompression(
- FileBasedSource<String> source, CompressionType compressionType) {
- switch (compressionType) {
+ switch (getCompressionType()) {
case UNCOMPRESSED:
- return source;
+ return new TextSource(getFilepattern());
case AUTO:
- return CompressedSource.from(source);
+ return CompressedSource.from(new TextSource(getFilepattern()));
case BZIP2:
return
- CompressedSource.from(source)
- .withDecompression(CompressionMode.BZIP2);
+ CompressedSource.from(new TextSource(getFilepattern()))
+ .withDecompression(CompressedSource.CompressionMode.BZIP2);
case GZIP:
return
- CompressedSource.from(source)
- .withDecompression(CompressionMode.GZIP);
+ CompressedSource.from(new TextSource(getFilepattern()))
+ .withDecompression(CompressedSource.CompressionMode.GZIP);
case ZIP:
return
- CompressedSource.from(source)
- .withDecompression(CompressionMode.ZIP);
+ CompressedSource.from(new TextSource(getFilepattern()))
+ .withDecompression(CompressedSource.CompressionMode.ZIP);
case DEFLATE:
return
- CompressedSource.from(source)
- .withDecompression(CompressionMode.DEFLATE);
+ CompressedSource.from(new TextSource(getFilepattern()))
+ .withDecompression(CompressedSource.CompressionMode.DEFLATE);
default:
- throw new IllegalArgumentException("Unknown compression type: " + compressionType);
+ throw new IllegalArgumentException("Unknown compression type: " + getFilepattern());
}
}
@@ -324,170 +223,18 @@ public class TextIO {
}
}
- /////////////////////////////////////////////////////////////////////////////
-
- /** Implementation of {@link #readAll}. */
- @AutoValue
- public abstract static class ReadAll
- extends PTransform<PCollection<String>, PCollection<String>> {
- abstract CompressionType getCompressionType();
- abstract long getDesiredBundleSizeBytes();
-
- abstract Builder toBuilder();
-
- @AutoValue.Builder
- abstract static class Builder {
- abstract Builder setCompressionType(CompressionType compressionType);
- abstract Builder setDesiredBundleSizeBytes(long desiredBundleSizeBytes);
-
- abstract ReadAll build();
- }
-
- /** Same as {@link Read#withCompressionType(CompressionType)}. */
- public ReadAll withCompressionType(CompressionType compressionType) {
- return toBuilder().setCompressionType(compressionType).build();
- }
-
- @VisibleForTesting
- ReadAll withDesiredBundleSizeBytes(long desiredBundleSizeBytes) {
- return toBuilder().setDesiredBundleSizeBytes(desiredBundleSizeBytes).build();
- }
-
- @Override
- public PCollection<String> expand(PCollection<String> input) {
- return input
- .apply("Expand glob", ParDo.of(new ExpandGlobFn()))
- .apply(
- "Split into ranges",
- ParDo.of(new SplitIntoRangesFn(getCompressionType(), getDesiredBundleSizeBytes())))
- .apply("Reshuffle", new ReshuffleWithUniqueKey<KV<Metadata, OffsetRange>>())
- .apply("Read", ParDo.of(new ReadTextFn(this)));
- }
-
- private static class ReshuffleWithUniqueKey<T>
- extends PTransform<PCollection<T>, PCollection<T>> {
- @Override
- public PCollection<T> expand(PCollection<T> input) {
- return input
- .apply("Unique key", ParDo.of(new AssignUniqueKeyFn<T>()))
- .apply("Reshuffle", Reshuffle.<Integer, T>of())
- .apply("Values", Values.<T>create());
- }
- }
-
- private static class AssignUniqueKeyFn<T> extends DoFn<T, KV<Integer, T>> {
- private int index;
-
- @Setup
- public void setup() {
- this.index = ThreadLocalRandom.current().nextInt();
- }
-
- @ProcessElement
- public void process(ProcessContext c) {
- c.output(KV.of(++index, c.element()));
- }
- }
-
- private static class ExpandGlobFn extends DoFn<String, Metadata> {
- @ProcessElement
- public void process(ProcessContext c) throws Exception {
- MatchResult match = FileSystems.match(c.element());
- checkArgument(
- match.status().equals(Status.OK),
- "Failed to match filepattern %s: %s",
- c.element(),
- match.status());
- for (Metadata metadata : match.metadata()) {
- c.output(metadata);
- }
- }
- }
-
- private static class SplitIntoRangesFn extends DoFn<Metadata, KV<Metadata, OffsetRange>> {
- private final CompressionType compressionType;
- private final long desiredBundleSize;
-
- private SplitIntoRangesFn(CompressionType compressionType, long desiredBundleSize) {
- this.compressionType = compressionType;
- this.desiredBundleSize = desiredBundleSize;
- }
-
- @ProcessElement
- public void process(ProcessContext c) {
- Metadata metadata = c.element();
- final boolean isSplittable = isSplittable(metadata, compressionType);
- if (!isSplittable) {
- c.output(KV.of(metadata, new OffsetRange(0, metadata.sizeBytes())));
- return;
- }
- for (OffsetRange range :
- new OffsetRange(0, metadata.sizeBytes()).split(desiredBundleSize, 0)) {
- c.output(KV.of(metadata, range));
- }
- }
-
- static boolean isSplittable(Metadata metadata, CompressionType compressionType) {
- if (!metadata.isReadSeekEfficient()) {
- return false;
- }
- switch (compressionType) {
- case AUTO:
- return !CompressionMode.isCompressed(metadata.resourceId().toString());
- case UNCOMPRESSED:
- return true;
- case GZIP:
- case BZIP2:
- case ZIP:
- case DEFLATE:
- return false;
- default:
- throw new UnsupportedOperationException("Unknown compression type: " + compressionType);
- }
- }
- }
-
- private static class ReadTextFn extends DoFn<KV<Metadata, OffsetRange>, String> {
- private final TextIO.ReadAll spec;
-
- private ReadTextFn(ReadAll spec) {
- this.spec = spec;
- }
-
- @ProcessElement
- public void process(ProcessContext c) throws IOException {
- Metadata metadata = c.element().getKey();
- OffsetRange range = c.element().getValue();
- FileBasedSource<String> source =
- TextIO.Read.wrapWithCompression(
- new TextSource(StaticValueProvider.of(metadata.toString())),
- spec.getCompressionType());
- BoundedSource.BoundedReader<String> reader =
- source
- .createForSubrangeOfFile(metadata, range.getFrom(), range.getTo())
- .createReader(c.getPipelineOptions());
- for (boolean more = reader.start(); more; more = reader.advance()) {
- c.output(reader.getCurrent());
- }
- }
- }
- }
/////////////////////////////////////////////////////////////////////////////
/** Implementation of {@link #write}. */
@AutoValue
- public abstract static class TypedWrite<T> extends PTransform<PCollection<T>, PDone> {
+ public abstract static class Write extends PTransform<PCollection<String>, PDone> {
/** The prefix of each file written, combined with suffix and shardTemplate. */
@Nullable abstract ValueProvider<ResourceId> getFilenamePrefix();
/** The suffix of each file written, combined with prefix and shardTemplate. */
@Nullable abstract String getFilenameSuffix();
- /** The base directory used for generating temporary files. */
- @Nullable
- abstract ValueProvider<ResourceId> getTempDirectory();
-
/** An optional header to add to each file. */
@Nullable abstract String getHeader();
@@ -503,13 +250,6 @@ public class TextIO {
/** A policy for naming output files. */
@Nullable abstract FilenamePolicy getFilenamePolicy();
- /** Allows for value-dependent {@link DynamicDestinations} to be vended. */
- @Nullable
- abstract DynamicDestinations<T, ?> getDynamicDestinations();
-
- /** A function that converts T to a String, for writing to the file. */
- abstract SerializableFunction<T, String> getFormatFunction();
-
/** Whether to write windowed output files. */
abstract boolean getWindowedWrites();
@@ -519,68 +259,66 @@ public class TextIO {
*/
abstract WritableByteChannelFactory getWritableByteChannelFactory();
- abstract Builder<T> toBuilder();
+ abstract Builder toBuilder();
@AutoValue.Builder
- abstract static class Builder<T> {
- abstract Builder<T> setFilenamePrefix(ValueProvider<ResourceId> filenamePrefix);
-
- abstract Builder<T> setTempDirectory(ValueProvider<ResourceId> tempDirectory);
-
- abstract Builder<T> setShardTemplate(@Nullable String shardTemplate);
-
- abstract Builder<T> setFilenameSuffix(@Nullable String filenameSuffix);
-
- abstract Builder<T> setHeader(@Nullable String header);
-
- abstract Builder<T> setFooter(@Nullable String footer);
-
- abstract Builder<T> setFilenamePolicy(@Nullable FilenamePolicy filenamePolicy);
-
- abstract Builder<T> setDynamicDestinations(
- @Nullable DynamicDestinations<T, ?> dynamicDestinations);
-
- abstract Builder<T> setFormatFunction(SerializableFunction<T, String> formatFunction);
-
- abstract Builder<T> setNumShards(int numShards);
-
- abstract Builder<T> setWindowedWrites(boolean windowedWrites);
-
- abstract Builder<T> setWritableByteChannelFactory(
+ abstract static class Builder {
+ abstract Builder setFilenamePrefix(ValueProvider<ResourceId> filenamePrefix);
+ abstract Builder setShardTemplate(@Nullable String shardTemplate);
+ abstract Builder setFilenameSuffix(@Nullable String filenameSuffix);
+ abstract Builder setHeader(@Nullable String header);
+ abstract Builder setFooter(@Nullable String footer);
+ abstract Builder setFilenamePolicy(@Nullable FilenamePolicy filenamePolicy);
+ abstract Builder setNumShards(int numShards);
+ abstract Builder setWindowedWrites(boolean windowedWrites);
+ abstract Builder setWritableByteChannelFactory(
WritableByteChannelFactory writableByteChannelFactory);
- abstract TypedWrite<T> build();
+ abstract Write build();
}
/**
- * Writes to text files with the given prefix. The given {@code prefix} can reference any {@link
- * FileSystem} on the classpath. This prefix is used by the {@link DefaultFilenamePolicy} to
- * generate filenames.
+ * Writes to text files with the given prefix. The given {@code prefix} can reference any
+ * {@link FileSystem} on the classpath.
*
- * <p>By default, a {@link DefaultFilenamePolicy} will be used built using the specified prefix
- * to define the base output directory and file prefix, a shard identifier (see {@link
- * #withNumShards(int)}), and a common suffix (if supplied using {@link #withSuffix(String)}).
+ * <p>The name of the output files will be determined by the {@link FilenamePolicy} used.
*
- * <p>This default policy can be overridden using {@link #to(FilenamePolicy)}, in which case
- * {@link #withShardNameTemplate(String)} and {@link #withSuffix(String)} should not be set.
- * Custom filename policies do not automatically see this prefix - you should explicitly pass
- * the prefix into your {@link FilenamePolicy} object if you need this.
+ * <p>By default, a {@link DefaultFilenamePolicy} will be used built using the specified prefix
+ * to define the base output directory and file prefix, a shard identifier (see
+ * {@link #withNumShards(int)}), and a common suffix (if supplied using
+ * {@link #withSuffix(String)}).
*
- * <p>If {@link #withTempDirectory} has not been called, this filename prefix will be used to
- * infer a directory for temporary files.
+ * <p>This default policy can be overridden using {@link #withFilenamePolicy(FilenamePolicy)},
+ * in which case {@link #withShardNameTemplate(String)} and {@link #withSuffix(String)} should
+ * not be set.
*/
- public TypedWrite<T> to(String filenamePrefix) {
+ public Write to(String filenamePrefix) {
return to(FileBasedSink.convertToFileResourceIfPossible(filenamePrefix));
}
- /** Like {@link #to(String)}. */
+ /**
+ * Writes to text files with prefix from the given resource.
+ *
+ * <p>The name of the output files will be determined by the {@link FilenamePolicy} used.
+ *
+ * <p>By default, a {@link DefaultFilenamePolicy} will be used built using the specified prefix
+ * to define the base output directory and file prefix, a shard identifier (see
+ * {@link #withNumShards(int)}), and a common suffix (if supplied using
+ * {@link #withSuffix(String)}).
+ *
+ * <p>This default policy can be overridden using {@link #withFilenamePolicy(FilenamePolicy)},
+ * in which case {@link #withShardNameTemplate(String)} and {@link #withSuffix(String)} should
+ * not be set.
+ */
@Experimental(Kind.FILESYSTEM)
- public TypedWrite<T> to(ResourceId filenamePrefix) {
+ public Write to(ResourceId filenamePrefix) {
return toResource(StaticValueProvider.of(filenamePrefix));
}
- /** Like {@link #to(String)}. */
- public TypedWrite<T> to(ValueProvider<String> outputPrefix) {
+ /**
+ * Like {@link #to(String)}.
+ */
+ public Write to(ValueProvider<String> outputPrefix) {
return toResource(NestedValueProvider.of(outputPrefix,
new SerializableFunction<String, ResourceId>() {
@Override
@@ -591,77 +329,43 @@ public class TextIO {
}
/**
- * Writes to files named according to the given {@link FileBasedSink.FilenamePolicy}. A
- * directory for temporary files must be specified using {@link #withTempDirectory}.
+ * Like {@link #to(ResourceId)}.
*/
- public TypedWrite<T> to(FilenamePolicy filenamePolicy) {
- return toBuilder().setFilenamePolicy(filenamePolicy).build();
- }
-
- /**
- * Use a {@link DynamicDestinations} object to vend {@link FilenamePolicy} objects. These
- * objects can examine the input record when creating a {@link FilenamePolicy}. A directory for
- * temporary files must be specified using {@link #withTempDirectory}.
- */
- public TypedWrite<T> to(DynamicDestinations<T, ?> dynamicDestinations) {
- return toBuilder().setDynamicDestinations(dynamicDestinations).build();
- }
-
- /**
- * Write to dynamic destinations using the default filename policy. The destinationFunction maps
- * the input record to a {@link DefaultFilenamePolicy.Params} object that specifies where the
- * records should be written (base filename, file suffix, and shard template). The
- * emptyDestination parameter specified where empty files should be written for when the written
- * {@link PCollection} is empty.
- */
- public TypedWrite<T> to(
- SerializableFunction<T, Params> destinationFunction, Params emptyDestination) {
- return to(DynamicFileDestinations.toDefaultPolicies(destinationFunction, emptyDestination));
- }
-
- /** Like {@link #to(ResourceId)}. */
@Experimental(Kind.FILESYSTEM)
- public TypedWrite<T> toResource(ValueProvider<ResourceId> filenamePrefix) {
+ public Write toResource(ValueProvider<ResourceId> filenamePrefix) {
return toBuilder().setFilenamePrefix(filenamePrefix).build();
}
- /** Set the base directory used to generate temporary files. */
- @Experimental(Kind.FILESYSTEM)
- public TypedWrite<T> withTempDirectory(ValueProvider<ResourceId> tempDirectory) {
- return toBuilder().setTempDirectory(tempDirectory).build();
- }
-
- /** Set the base directory used to generate temporary files. */
- @Experimental(Kind.FILESYSTEM)
- public TypedWrite<T> withTempDirectory(ResourceId tempDirectory) {
- return withTempDirectory(StaticValueProvider.of(tempDirectory));
- }
-
/**
* Uses the given {@link ShardNameTemplate} for naming output files. This option may only be
- * used when using one of the default filename-prefix to() overrides - i.e. not when using
- * either {@link #to(FilenamePolicy)} or {@link #to(DynamicDestinations)}.
+ * used when {@link #withFilenamePolicy(FilenamePolicy)} has not been configured.
*
* <p>See {@link DefaultFilenamePolicy} for how the prefix, shard name template, and suffix are
* used.
*/
- public TypedWrite<T> withShardNameTemplate(String shardTemplate) {
+ public Write withShardNameTemplate(String shardTemplate) {
return toBuilder().setShardTemplate(shardTemplate).build();
}
/**
- * Configures the filename suffix for written files. This option may only be used when using one
- * of the default filename-prefix to() overrides - i.e. not when using either {@link
- * #to(FilenamePolicy)} or {@link #to(DynamicDestinations)}.
+ * Configures the filename suffix for written files. This option may only be used when
+ * {@link #withFilenamePolicy(FilenamePolicy)} has not been configured.
*
* <p>See {@link DefaultFilenamePolicy} for how the prefix, shard name template, and suffix are
* used.
*/
- public TypedWrite<T> withSuffix(String filenameSuffix) {
+ public Write withSuffix(String filenameSuffix) {
return toBuilder().setFilenameSuffix(filenameSuffix).build();
}
/**
+ * Configures the {@link FileBasedSink.FilenamePolicy} that will be used to name written files.
+ */
+ public Write withFilenamePolicy(FilenamePolicy filenamePolicy) {
+ return toBuilder().setFilenamePolicy(filenamePolicy).build();
+ }
+
+ /**
* Configures the number of output shards produced overall (when using unwindowed writes) or
* per-window (when using windowed writes).
*
@@ -671,13 +375,14 @@ public class TextIO {
*
* @param numShards the number of shards to use, or 0 to let the system decide.
*/
- public TypedWrite<T> withNumShards(int numShards) {
+ public Write withNumShards(int numShards) {
checkArgument(numShards >= 0);
return toBuilder().setNumShards(numShards).build();
}
/**
- * Forces a single file as output and empty shard name template.
+ * Forces a single file as output and empty shard name template. This option is only compatible
+ * with unwindowed writes.
*
* <p>For unwindowed writes, constraining the number of shards is likely to reduce the
* performance of a pipeline. Setting this value is not recommended unless you require a
@@ -685,7 +390,7 @@ public class TextIO {
*
* <p>This is equivalent to {@code .withNumShards(1).withShardNameTemplate("")}
*/
- public TypedWrite<T> withoutSharding() {
+ public Write withoutSharding() {
return withNumShards(1).withShardNameTemplate("");
}
@@ -694,7 +399,7 @@ public class TextIO {
*
* <p>A {@code null} value will clear any previously configured header.
*/
- public TypedWrite<T> withHeader(@Nullable String header) {
+ public Write withHeader(@Nullable String header) {
return toBuilder().setHeader(header).build();
}
@@ -703,82 +408,48 @@ public class TextIO {
*
* <p>A {@code null} value will clear any previously configured footer.
*/
- public TypedWrite<T> withFooter(@Nullable String footer) {
+ public Write withFooter(@Nullable String footer) {
return toBuilder().setFooter(footer).build();
}
/**
- * Returns a transform for writing to text files like this one but that has the given {@link
- * WritableByteChannelFactory} to be used by the {@link FileBasedSink} during output. The
- * default is value is {@link FileBasedSink.CompressionType#UNCOMPRESSED}.
+ * Returns a transform for writing to text files like this one but that has the given
+ * {@link WritableByteChannelFactory} to be used by the {@link FileBasedSink} during output.
+ * The default is value is {@link FileBasedSink.CompressionType#UNCOMPRESSED}.
*
* <p>A {@code null} value will reset the value to the default value mentioned above.
*/
- public TypedWrite<T> withWritableByteChannelFactory(
+ public Write withWritableByteChannelFactory(
WritableByteChannelFactory writableByteChannelFactory) {
return toBuilder().setWritableByteChannelFactory(writableByteChannelFactory).build();
}
- /**
- * Preserves windowing of input elements and writes them to files based on the element's window.
- *
- * <p>If using {@link #to(FileBasedSink.FilenamePolicy)}. Filenames will be generated using
- * {@link FilenamePolicy#windowedFilename}. See also {@link WriteFiles#withWindowedWrites()}.
- */
- public TypedWrite<T> withWindowedWrites() {
+ public Write withWindowedWrites() {
return toBuilder().setWindowedWrites(true).build();
}
- private DynamicDestinations<T, ?> resolveDynamicDestinations() {
- DynamicDestinations<T, ?> dynamicDestinations = getDynamicDestinations();
- if (dynamicDestinations == null) {
- FilenamePolicy usedFilenamePolicy = getFilenamePolicy();
- if (usedFilenamePolicy == null) {
- usedFilenamePolicy =
- DefaultFilenamePolicy.fromStandardParameters(
- getFilenamePrefix(),
- getShardTemplate(),
- getFilenameSuffix(),
- getWindowedWrites());
- }
- dynamicDestinations = DynamicFileDestinations.constant(usedFilenamePolicy);
- }
- return dynamicDestinations;
- }
-
@Override
- public PDone expand(PCollection<T> input) {
- checkState(
- getFilenamePrefix() != null || getTempDirectory() != null,
- "Need to set either the filename prefix or the tempDirectory of a TextIO.Write "
- + "transform.");
+ public PDone expand(PCollection<String> input) {
+ checkState(getFilenamePrefix() != null,
+ "Need to set the filename prefix of a TextIO.Write transform.");
checkState(
- getFilenamePolicy() == null || getDynamicDestinations() == null,
- "Cannot specify both a filename policy and dynamic destinations");
- if (getFilenamePolicy() != null || getDynamicDestinations() != null) {
- checkState(
- getShardTemplate() == null && getFilenameSuffix() == null,
- "shardTemplate and filenameSuffix should only be used with the default "
- + "filename policy");
+ (getFilenamePolicy() == null)
+ || (getShardTemplate() == null && getFilenameSuffix() == null),
+ "Cannot set a filename policy and also a filename template or suffix.");
+
+ FilenamePolicy usedFilenamePolicy = getFilenamePolicy();
+ if (usedFilenamePolicy == null) {
+ usedFilenamePolicy = DefaultFilenamePolicy.constructUsingStandardParameters(
+ getFilenamePrefix(), getShardTemplate(), getFilenameSuffix(), getWindowedWrites());
}
- return expandTyped(input, resolveDynamicDestinations());
- }
-
- public <DestinationT> PDone expandTyped(
- PCollection<T> input, DynamicDestinations<T, DestinationT> dynamicDestinations) {
- ValueProvider<ResourceId> tempDirectory = getTempDirectory();
- if (tempDirectory == null) {
- tempDirectory = getFilenamePrefix();
- }
- WriteFiles<T, DestinationT, String> write =
+ WriteFiles<String> write =
WriteFiles.to(
- new TextSink<>(
- tempDirectory,
- dynamicDestinations,
+ new TextSink(
+ getFilenamePrefix(),
+ usedFilenamePolicy,
getHeader(),
getFooter(),
- getWritableByteChannelFactory()),
- getFormatFunction());
+ getWritableByteChannelFactory()));
if (getNumShards() > 0) {
write = write.withNumShards(getNumShards());
}
@@ -792,26 +463,27 @@ public class TextIO {
public void populateDisplayData(DisplayData.Builder builder) {
super.populateDisplayData(builder);
- resolveDynamicDestinations().populateDisplayData(builder);
- String tempDirectory = null;
- if (getTempDirectory() != null) {
- tempDirectory =
- getTempDirectory().isAccessible()
- ? getTempDirectory().get().toString()
- : getTempDirectory().toString();
+ String prefixString = "";
+ if (getFilenamePrefix() != null) {
+ prefixString = getFilenamePrefix().isAccessible()
+ ? getFilenamePrefix().get().toString() : getFilenamePrefix().toString();
}
builder
- .addIfNotDefault(
- DisplayData.item("numShards", getNumShards()).withLabel("Maximum Output Shards"), 0)
- .addIfNotNull(
- DisplayData.item("tempDirectory", tempDirectory)
- .withLabel("Directory for temporary files"))
- .addIfNotNull(DisplayData.item("fileHeader", getHeader()).withLabel("File Header"))
- .addIfNotNull(DisplayData.item("fileFooter", getFooter()).withLabel("File Footer"))
- .add(
- DisplayData.item(
- "writableByteChannelFactory", getWritableByteChannelFactory().toString())
- .withLabel("Compression/Transformation Type"));
+ .addIfNotNull(DisplayData.item("filePrefix", prefixString)
+ .withLabel("Output File Prefix"))
+ .addIfNotNull(DisplayData.item("fileSuffix", getFilenameSuffix())
+ .withLabel("Output File Suffix"))
+ .addIfNotNull(DisplayData.item("shardNameTemplate", getShardTemplate())
+ .withLabel("Output Shard Name Template"))
+ .addIfNotDefault(DisplayData.item("numShards", getNumShards())
+ .withLabel("Maximum Output Shards"), 0)
+ .addIfNotNull(DisplayData.item("fileHeader", getHeader())
+ .withLabel("File Header"))
+ .addIfNotNull(DisplayData.item("fileFooter", getFooter())
+ .withLabel("File Footer"))
+ .add(DisplayData
+ .item("writableByteChannelFactory", getWritableByteChannelFactory().toString())
+ .withLabel("Compression/Transformation Type"));
}
@Override
@@ -821,128 +493,6 @@ public class TextIO {
}
/**
- * This class is used as the default return value of {@link TextIO#write()}.
- *
- * <p>All methods in this class delegate to the appropriate method of {@link TextIO.TypedWrite}.
- * This class exists for backwards compatibility, and will be removed in Beam 3.0.
- */
- public static class Write extends PTransform<PCollection<String>, PDone> {
- @VisibleForTesting TypedWrite<String> inner;
-
- Write() {
- this(TextIO.writeCustomType(SerializableFunctions.<String>identity()));
- }
-
- Write(TypedWrite<String> inner) {
- this.inner = inner;
- }
-
- /** See {@link TypedWrite#to(String)}. */
- public Write to(String filenamePrefix) {
- return new Write(inner.to(filenamePrefix));
- }
-
- /** See {@link TypedWrite#to(ResourceId)}. */
- @Experimental(Kind.FILESYSTEM)
- public Write to(ResourceId filenamePrefix) {
- return new Write(inner.to(filenamePrefix));
- }
-
- /** See {@link TypedWrite#to(ValueProvider)}. */
- public Write to(ValueProvider<String> outputPrefix) {
- return new Write(inner.to(outputPrefix));
- }
-
- /** See {@link TypedWrite#toResource(ValueProvider)}. */
- @Experimental(Kind.FILESYSTEM)
- public Write toResource(ValueProvider<ResourceId> filenamePrefix) {
- return new Write(inner.toResource(filenamePrefix));
- }
-
- /** See {@link TypedWrite#to(FilenamePolicy)}. */
- @Experimental(Kind.FILESYSTEM)
- public Write to(FilenamePolicy filenamePolicy) {
- return new Write(inner.to(filenamePolicy));
- }
-
- /** See {@link TypedWrite#to(DynamicDestinations)}. */
- @Experimental(Kind.FILESYSTEM)
- public Write to(DynamicDestinations<String, ?> dynamicDestinations) {
- return new Write(inner.to(dynamicDestinations));
- }
-
- /** See {@link TypedWrite#to(SerializableFunction, Params)}. */
- @Experimental(Kind.FILESYSTEM)
- public Write to(
- SerializableFunction<String, Params> destinationFunction, Params emptyDestination) {
- return new Write(inner.to(destinationFunction, emptyDestination));
- }
-
- /** See {@link TypedWrite#withTempDirectory(ValueProvider)}. */
- @Experimental(Kind.FILESYSTEM)
- public Write withTempDirectory(ValueProvider<ResourceId> tempDirectory) {
- return new Write(inner.withTempDirectory(tempDirectory));
- }
-
- /** See {@link TypedWrite#withTempDirectory(ResourceId)}. */
- @Experimental(Kind.FILESYSTEM)
- public Write withTempDirectory(ResourceId tempDirectory) {
- return new Write(inner.withTempDirectory(tempDirectory));
- }
-
- /** See {@link TypedWrite#withShardNameTemplate(String)}. */
- public Write withShardNameTemplate(String shardTemplate) {
- return new Write(inner.withShardNameTemplate(shardTemplate));
- }
-
- /** See {@link TypedWrite#withSuffix(String)}. */
- public Write withSuffix(String filenameSuffix) {
- return new Write(inner.withSuffix(filenameSuffix));
- }
-
- /** See {@link TypedWrite#withNumShards(int)}. */
- public Write withNumShards(int numShards) {
- return new Write(inner.withNumShards(numShards));
- }
-
- /** See {@link TypedWrite#withoutSharding()}. */
- public Write withoutSharding() {
- return new Write(inner.withoutSharding());
- }
-
- /** See {@link TypedWrite#withHeader(String)}. */
- public Write withHeader(@Nullable String header) {
- return new Write(inner.withHeader(header));
- }
-
- /** See {@link TypedWrite#withFooter(String)}. */
- public Write withFooter(@Nullable String footer) {
- return new Write(inner.withFooter(footer));
- }
-
- /** See {@link TypedWrite#withWritableByteChannelFactory(WritableByteChannelFactory)}. */
- public Write withWritableByteChannelFactory(
- WritableByteChannelFactory writableByteChannelFactory) {
- return new Write(inner.withWritableByteChannelFactory(writableByteChannelFactory));
- }
-
- /** See {@link TypedWrite#withWindowedWrites}. */
- public Write withWindowedWrites() {
- return new Write(inner.withWindowedWrites());
- }
-
- @Override
- public void populateDisplayData(DisplayData.Builder builder) {
- inner.populateDisplayData(builder);
- }
-
- @Override
- public PDone expand(PCollection<String> input) {
- return inner.expand(input);
- }
- }
-
- /**
* Possible text file compression types.
*/
public enum CompressionType {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextSink.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextSink.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextSink.java
index b57b28c..511d697 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextSink.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextSink.java
@@ -34,29 +34,27 @@ import org.apache.beam.sdk.util.MimeTypes;
* '\n'} represented in {@code UTF-8} format as the record separator. Each record (including the
* last) is terminated.
*/
-class TextSink<UserT, DestinationT> extends FileBasedSink<String, DestinationT> {
+class TextSink extends FileBasedSink<String> {
@Nullable private final String header;
@Nullable private final String footer;
TextSink(
ValueProvider<ResourceId> baseOutputFilename,
- DynamicDestinations<UserT, DestinationT> dynamicDestinations,
+ FilenamePolicy filenamePolicy,
@Nullable String header,
@Nullable String footer,
WritableByteChannelFactory writableByteChannelFactory) {
- super(baseOutputFilename, dynamicDestinations, writableByteChannelFactory);
+ super(baseOutputFilename, filenamePolicy, writableByteChannelFactory);
this.header = header;
this.footer = footer;
}
-
@Override
- public WriteOperation<String, DestinationT> createWriteOperation() {
- return new TextWriteOperation<>(this, header, footer);
+ public WriteOperation<String> createWriteOperation() {
+ return new TextWriteOperation(this, header, footer);
}
/** A {@link WriteOperation WriteOperation} for text files. */
- private static class TextWriteOperation<DestinationT>
- extends WriteOperation<String, DestinationT> {
+ private static class TextWriteOperation extends WriteOperation<String> {
@Nullable private final String header;
@Nullable private final String footer;
@@ -67,20 +65,20 @@ class TextSink<UserT, DestinationT> extends FileBasedSink<String, DestinationT>
}
@Override
- public Writer<String, DestinationT> createWriter() throws Exception {
- return new TextWriter<>(this, header, footer);
+ public Writer<String> createWriter() throws Exception {
+ return new TextWriter(this, header, footer);
}
}
/** A {@link Writer Writer} for text files. */
- private static class TextWriter<DestinationT> extends Writer<String, DestinationT> {
+ private static class TextWriter extends Writer<String> {
private static final String NEWLINE = "\n";
@Nullable private final String header;
@Nullable private final String footer;
private OutputStreamWriter out;
public TextWriter(
- WriteOperation<String, DestinationT> writeOperation,
+ WriteOperation<String> writeOperation,
@Nullable String header,
@Nullable String footer) {
super(writeOperation, MimeTypes.TEXT);
[28/28] beam git commit: This closes #3602
Posted by jb...@apache.org.
This closes #3602
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/152115e8
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/152115e8
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/152115e8
Branch: refs/heads/DSL_SQL
Commit: 152115e893b25d0a45dd208ce613a5bdad5da8cd
Parents: 97a156c c1b2b96
Author: Jean-Baptiste Onofré <jb...@apache.org>
Authored: Thu Jul 20 19:09:10 2017 +0200
Committer: Jean-Baptiste Onofré <jb...@apache.org>
Committed: Thu Jul 20 19:09:10 2017 +0200
----------------------------------------------------------------------
.gitignore | 2 +-
.../jenkins/common_job_properties.groovy | 9 +-
.../job_beam_PerformanceTests_Python.groovy | 58 --
..._beam_PostCommit_Java_JDKVersionsTest.groovy | 2 -
..._PostCommit_Java_MavenInstall_Windows.groovy | 3 +-
.../job_beam_PreCommit_Website_Merge.groovy | 59 --
examples/java/pom.xml | 20 +-
.../examples/common/WriteOneFilePerWindow.java | 52 +-
.../beam/examples/WindowedWordCountIT.java | 4 +-
examples/java8/pom.xml | 20 +-
.../complete/game/utils/WriteToText.java | 43 +-
.../examples/complete/game/LeaderBoardTest.java | 2 -
examples/pom.xml | 2 +-
pom.xml | 123 +---
runners/apex/pom.xml | 20 +-
.../apache/beam/runners/apex/ApexRunner.java | 61 +-
.../translation/ApexPipelineTranslator.java | 16 +-
.../apex/translation/TranslationContext.java | 4 +-
.../operators/ApexParDoOperator.java | 21 +-
.../runners/apex/examples/WordCountTest.java | 8 +-
.../utils/ApexStateInternalsTest.java | 411 ++++++++---
runners/core-construction-java/pom.xml | 2 +-
.../CreatePCollectionViewTranslation.java | 4 +-
.../construction/ElementAndRestriction.java | 42 ++
.../ElementAndRestrictionCoder.java | 88 +++
.../construction/PCollectionTranslation.java | 16 -
.../core/construction/PTransformMatchers.java | 109 +--
.../construction/PTransformTranslation.java | 7 +-
.../core/construction/ParDoTranslation.java | 82 +--
.../construction/RunnerPCollectionView.java | 31 +-
.../core/construction/SplittableParDo.java | 124 +---
.../construction/TestStreamTranslation.java | 49 +-
.../core/construction/TransformInputs.java | 50 --
.../WindowingStrategyTranslation.java | 27 +-
.../construction/WriteFilesTranslation.java | 67 +-
.../ElementAndRestrictionCoderTest.java | 126 ++++
.../PCollectionTranslationTest.java | 22 -
.../construction/PTransformMatchersTest.java | 54 +-
.../core/construction/ParDoTranslationTest.java | 28 +-
.../core/construction/SplittableParDoTest.java | 18 +-
.../core/construction/TransformInputsTest.java | 166 -----
.../WindowingStrategyTranslationTest.java | 3 -
.../construction/WriteFilesTranslationTest.java | 62 +-
runners/core-java/pom.xml | 2 +-
.../core/LateDataDroppingDoFnRunner.java | 33 +-
...eBoundedSplittableProcessElementInvoker.java | 40 +-
.../beam/runners/core/ProcessFnRunner.java | 16 +-
.../beam/runners/core/ReduceFnRunner.java | 135 ++--
.../beam/runners/core/SimpleDoFnRunner.java | 20 -
.../core/SplittableParDoViaKeyedWorkItems.java | 58 +-
.../core/SplittableProcessElementInvoker.java | 25 +-
.../beam/runners/core/SystemReduceFn.java | 6 -
.../core/triggers/AfterAllStateMachine.java | 25 +-
.../AfterDelayFromFirstElementStateMachine.java | 6 +-
.../core/triggers/AfterFirstStateMachine.java | 20 +-
.../core/triggers/AfterPaneStateMachine.java | 6 +-
.../triggers/AfterWatermarkStateMachine.java | 7 +-
.../triggers/ExecutableTriggerStateMachine.java | 23 +-
.../core/triggers/NeverStateMachine.java | 5 +-
.../core/triggers/TriggerStateMachine.java | 27 +
.../core/InMemoryStateInternalsTest.java | 569 +++++++++++++--
...ndedSplittableProcessElementInvokerTest.java | 47 +-
.../beam/runners/core/ReduceFnRunnerTest.java | 374 +---------
.../beam/runners/core/ReduceFnTester.java | 48 +-
.../core/SplittableParDoProcessFnTest.java | 117 +--
.../beam/runners/core/StateInternalsTest.java | 613 ----------------
.../beam/runners/core/WindowMatchers.java | 15 -
.../triggers/AfterFirstStateMachineTest.java | 5 +-
.../AfterWatermarkStateMachineTest.java | 7 +-
.../core/triggers/StubTriggerStateMachine.java | 7 +-
runners/direct-java/pom.xml | 7 +-
.../beam/runners/direct/CommittedResult.java | 12 +-
.../apache/beam/runners/direct/DirectGraph.java | 38 +-
.../beam/runners/direct/DirectGraphVisitor.java | 48 +-
.../beam/runners/direct/DirectGroupByKey.java | 13 +-
.../direct/DirectGroupByKeyOverrideFactory.java | 14 +-
.../beam/runners/direct/DirectRegistrar.java | 2 +-
.../beam/runners/direct/DirectRunner.java | 64 +-
.../beam/runners/direct/DirectTestOptions.java | 42 --
.../beam/runners/direct/EvaluationContext.java | 26 +-
.../direct/ExecutorServiceParallelExecutor.java | 27 +-
.../runners/direct/ParDoEvaluatorFactory.java | 9 +-
.../direct/ParDoMultiOverrideFactory.java | 121 +---
...littableProcessElementsEvaluatorFactory.java | 37 +-
.../direct/StatefulParDoEvaluatorFactory.java | 12 +-
.../direct/TestStreamEvaluatorFactory.java | 20 +-
.../runners/direct/ViewEvaluatorFactory.java | 8 +-
.../runners/direct/ViewOverrideFactory.java | 69 +-
.../beam/runners/direct/WatermarkManager.java | 18 +-
.../direct/WriteWithShardingFactory.java | 34 +-
.../runners/direct/CommittedResultTest.java | 17 +-
.../runners/direct/DirectGraphVisitorTest.java | 10 +-
.../beam/runners/direct/DirectGraphs.java | 7 -
.../runners/direct/DirectRegistrarTest.java | 2 +-
.../runners/direct/EvaluationContextTest.java | 7 +-
.../ImmutabilityEnforcementFactoryTest.java | 4 +-
.../beam/runners/direct/ParDoEvaluatorTest.java | 7 +-
.../StatefulParDoEvaluatorFactoryTest.java | 65 +-
.../runners/direct/TransformExecutorTest.java | 12 +-
.../direct/ViewEvaluatorFactoryTest.java | 8 +-
.../runners/direct/ViewOverrideFactoryTest.java | 37 +-
.../direct/WatermarkCallbackExecutorTest.java | 1 -
.../runners/direct/WatermarkManagerTest.java | 16 +-
.../direct/WriteWithShardingFactoryTest.java | 44 +-
runners/flink/pom.xml | 11 +-
.../runners/flink/CreateStreamingFlinkView.java | 154 ----
.../flink/FlinkBatchTranslationContext.java | 3 +-
.../FlinkPipelineExecutionEnvironment.java | 2 -
.../flink/FlinkStreamingPipelineTranslator.java | 86 ++-
.../FlinkStreamingTransformTranslators.java | 36 +-
.../flink/FlinkStreamingTranslationContext.java | 3 +-
.../flink/FlinkStreamingViewOverrides.java | 372 ++++++++++
.../runners/flink/FlinkTransformOverrides.java | 53 --
.../streaming/SplittableDoFnOperator.java | 16 +-
.../streaming/state/FlinkStateInternals.java | 425 ++++++-----
.../FlinkBroadcastStateInternalsTest.java | 242 +++++--
.../FlinkKeyGroupStateInternalsTest.java | 359 +++++-----
.../streaming/FlinkSplitStateInternalsTest.java | 132 ++--
.../streaming/FlinkStateInternalsTest.java | 343 ++++++++-
runners/google-cloud-dataflow-java/pom.xml | 10 +-
.../dataflow/BatchStatefulParDoOverrides.java | 4 -
.../runners/dataflow/BatchViewOverrides.java | 182 +++--
.../runners/dataflow/CreateDataflowView.java | 8 +-
.../dataflow/DataflowPipelineTranslator.java | 62 +-
.../beam/runners/dataflow/DataflowRunner.java | 133 +---
.../dataflow/SplittableParDoOverrides.java | 76 --
.../dataflow/StreamingViewOverrides.java | 10 +-
.../runners/dataflow/TransformTranslator.java | 4 +-
.../runners/dataflow/util/PropertyNames.java | 1 -
.../beam/runners/dataflow/util/TimeUtil.java | 24 +-
.../DataflowPipelineTranslatorTest.java | 95 +--
.../runners/dataflow/DataflowRunnerTest.java | 198 +-----
.../runners/dataflow/util/TimeUtilTest.java | 6 -
runners/pom.xml | 2 +-
runners/spark/pom.xml | 70 +-
.../spark/SparkNativePipelineVisitor.java | 3 +-
.../apache/beam/runners/spark/SparkRunner.java | 9 +-
.../beam/runners/spark/TestSparkRunner.java | 2 +-
.../SparkGroupAlsoByWindowViaWindowSet.java | 6 +-
.../spark/stateful/SparkTimerInternals.java | 18 +-
.../spark/translation/EvaluationContext.java | 4 +-
.../spark/translation/TransformTranslator.java | 50 +-
.../spark/util/GlobalWatermarkHolder.java | 127 +---
.../spark/GlobalWatermarkHolderTest.java | 18 +-
.../runners/spark/SparkRunnerDebuggerTest.java | 26 +-
.../spark/stateful/SparkStateInternalsTest.java | 66 --
.../spark/translation/StorageLevelTest.java | 4 +-
sdks/common/fn-api/pom.xml | 2 +-
.../fn-api/src/main/proto/beam_fn_api.proto | 237 ++++--
sdks/common/pom.xml | 2 +-
sdks/common/runner-api/pom.xml | 2 +-
.../src/main/proto/beam_runner_api.proto | 26 +-
sdks/java/build-tools/pom.xml | 2 +-
.../src/main/resources/beam/findbugs-filter.xml | 9 -
sdks/java/core/pom.xml | 2 +-
.../apache/beam/sdk/coders/ShardedKeyCoder.java | 66 --
.../java/org/apache/beam/sdk/io/AvroIO.java | 220 +++---
.../java/org/apache/beam/sdk/io/AvroSink.java | 32 +-
.../apache/beam/sdk/io/CompressedSource.java | 40 +-
.../beam/sdk/io/DefaultFilenamePolicy.java | 274 ++-----
.../beam/sdk/io/DynamicFileDestinations.java | 115 ---
.../org/apache/beam/sdk/io/FileBasedSink.java | 513 ++++++-------
.../apache/beam/sdk/io/OffsetBasedSource.java | 22 +-
.../java/org/apache/beam/sdk/io/TFRecordIO.java | 44 +-
.../java/org/apache/beam/sdk/io/TextIO.java | 712 ++++---------------
.../java/org/apache/beam/sdk/io/TextSink.java | 22 +-
.../java/org/apache/beam/sdk/io/WriteFiles.java | 647 ++++++-----------
.../beam/sdk/io/range/ByteKeyRangeTracker.java | 22 +-
.../apache/beam/sdk/io/range/OffsetRange.java | 101 ---
.../beam/sdk/io/range/OffsetRangeTracker.java | 3 -
.../sdk/options/PipelineOptionsFactory.java | 18 +-
.../sdk/options/PipelineOptionsValidator.java | 34 +-
.../sdk/options/ProxyInvocationHandler.java | 19 +-
.../beam/sdk/runners/TransformHierarchy.java | 165 +----
.../apache/beam/sdk/testing/StaticWindows.java | 5 -
.../org/apache/beam/sdk/testing/TestStream.java | 12 -
.../org/apache/beam/sdk/transforms/Combine.java | 30 +-
.../org/apache/beam/sdk/transforms/DoFn.java | 52 +-
.../apache/beam/sdk/transforms/DoFnTester.java | 21 +-
.../org/apache/beam/sdk/transforms/ParDo.java | 41 +-
.../sdk/transforms/SerializableFunctions.java | 50 --
.../org/apache/beam/sdk/transforms/View.java | 38 +-
.../reflect/ByteBuddyDoFnInvokerFactory.java | 27 -
.../reflect/ByteBuddyOnTimerInvokerFactory.java | 73 +-
.../sdk/transforms/reflect/DoFnInvoker.java | 17 +-
.../sdk/transforms/reflect/DoFnSignature.java | 33 +-
.../sdk/transforms/reflect/DoFnSignatures.java | 44 +-
.../reflect/OnTimerMethodSpecifier.java | 37 -
.../transforms/splittabledofn/OffsetRange.java | 77 ++
.../splittabledofn/OffsetRangeTracker.java | 11 -
.../splittabledofn/RestrictionTracker.java | 11 +-
.../sdk/transforms/windowing/GlobalWindows.java | 5 -
.../windowing/PartitioningWindowFn.java | 5 -
.../transforms/windowing/SlidingWindows.java | 5 -
.../beam/sdk/transforms/windowing/Window.java | 32 -
.../beam/sdk/transforms/windowing/WindowFn.java | 11 -
.../apache/beam/sdk/util/IdentityWindowFn.java | 5 -
.../org/apache/beam/sdk/values/PCollection.java | 12 -
.../beam/sdk/values/PCollectionViews.java | 38 -
.../org/apache/beam/sdk/values/PValueBase.java | 12 +
.../org/apache/beam/sdk/values/ShardedKey.java | 65 --
.../beam/sdk/values/WindowingStrategy.java | 46 +-
.../java/org/apache/beam/sdk/io/AvroIOTest.java | 85 +--
.../beam/sdk/io/DefaultFilenamePolicyTest.java | 135 ++--
.../sdk/io/DrunkWritableByteChannelFactory.java | 2 +-
.../apache/beam/sdk/io/FileBasedSinkTest.java | 93 +--
.../java/org/apache/beam/sdk/io/SimpleSink.java | 56 +-
.../java/org/apache/beam/sdk/io/TextIOTest.java | 326 +--------
.../org/apache/beam/sdk/io/WriteFilesTest.java | 366 ++--------
.../options/PipelineOptionsValidatorTest.java | 44 --
.../sdk/options/ProxyInvocationHandlerTest.java | 19 -
.../sdk/runners/TransformHierarchyTest.java | 197 -----
.../sdk/testing/PCollectionViewTesting.java | 8 -
.../apache/beam/sdk/transforms/CombineTest.java | 365 ++++------
.../beam/sdk/transforms/DoFnTesterTest.java | 32 -
.../beam/sdk/transforms/GroupByKeyTest.java | 39 -
.../apache/beam/sdk/transforms/ParDoTest.java | 165 -----
.../beam/sdk/transforms/SplittableDoFnTest.java | 155 +---
.../transforms/reflect/DoFnInvokersTest.java | 93 +--
.../DoFnSignaturesProcessElementTest.java | 2 +-
.../DoFnSignaturesSplittableDoFnTest.java | 83 +--
.../transforms/reflect/DoFnSignaturesTest.java | 14 -
.../splittabledofn/OffsetRangeTrackerTest.java | 1 -
.../windowing/SlidingWindowsTest.java | 30 +-
.../google-cloud-platform-core/pom.xml | 2 +-
.../java/org/apache/beam/sdk/util/GcsUtil.java | 2 +-
.../sdk/util/RetryHttpRequestInitializer.java | 147 ++--
.../extensions/gcp/GcpCoreApiSurfaceTest.java | 48 +-
.../util/RetryHttpRequestInitializerTest.java | 31 +-
sdks/java/extensions/jackson/pom.xml | 2 +-
sdks/java/extensions/join-library/pom.xml | 2 +-
sdks/java/extensions/pom.xml | 2 +-
sdks/java/extensions/protobuf/pom.xml | 2 +-
sdks/java/extensions/sorter/pom.xml | 8 +-
sdks/java/harness/pom.xml | 18 +-
.../harness/control/ProcessBundleHandler.java | 295 ++++++--
.../fn/harness/control/RegisterHandler.java | 2 +-
.../beam/runners/core/BeamFnDataReadRunner.java | 70 +-
.../runners/core/BeamFnDataWriteRunner.java | 67 +-
.../beam/runners/core/BoundedSourceRunner.java | 74 +-
.../beam/runners/core/FnApiDoFnRunner.java | 547 --------------
.../runners/core/PTransformRunnerFactory.java | 81 ---
.../control/ProcessBundleHandlerTest.java | 521 ++++++++++++--
.../fn/harness/control/RegisterHandlerTest.java | 8 +-
.../runners/core/BeamFnDataReadRunnerTest.java | 112 +--
.../runners/core/BeamFnDataWriteRunnerTest.java | 120 +---
.../runners/core/BoundedSourceRunnerTest.java | 124 +---
.../beam/runners/core/FnApiDoFnRunnerTest.java | 210 ------
sdks/java/io/amqp/pom.xml | 100 ---
.../org/apache/beam/sdk/io/amqp/AmqpIO.java | 399 -----------
.../beam/sdk/io/amqp/AmqpMessageCoder.java | 79 --
.../amqp/AmqpMessageCoderProviderRegistrar.java | 44 --
.../apache/beam/sdk/io/amqp/package-info.java | 22 -
.../org/apache/beam/sdk/io/amqp/AmqpIOTest.java | 148 ----
.../beam/sdk/io/amqp/AmqpMessageCoderTest.java | 89 ---
sdks/java/io/cassandra/pom.xml | 2 +-
.../beam/sdk/io/cassandra/CassandraIO.java | 2 +-
sdks/java/io/common/pom.xml | 2 +-
.../sdk/io/common/IOTestPipelineOptions.java | 6 +-
sdks/java/io/elasticsearch/pom.xml | 10 +-
.../sdk/io/elasticsearch/ElasticsearchIO.java | 17 +-
.../elasticsearch/ElasticSearchIOTestUtils.java | 81 +--
.../sdk/io/elasticsearch/ElasticsearchIOIT.java | 14 +-
.../io/elasticsearch/ElasticsearchIOTest.java | 36 +-
.../elasticsearch/ElasticsearchTestDataSet.java | 37 +-
sdks/java/io/google-cloud-platform/pom.xml | 14 +-
.../beam/sdk/io/gcp/bigquery/BatchLoads.java | 2 -
.../io/gcp/bigquery/DynamicDestinations.java | 29 +-
.../io/gcp/bigquery/GenerateShardedTable.java | 1 -
.../beam/sdk/io/gcp/bigquery/ShardedKey.java | 67 ++
.../sdk/io/gcp/bigquery/ShardedKeyCoder.java | 74 ++
.../sdk/io/gcp/bigquery/StreamingWriteFn.java | 1 -
.../io/gcp/bigquery/StreamingWriteTables.java | 2 -
.../sdk/io/gcp/bigquery/TagWithUniqueIds.java | 1 -
.../io/gcp/bigquery/WriteBundlesToFiles.java | 2 -
.../bigquery/WriteGroupedRecordsToFiles.java | 1 -
.../sdk/io/gcp/bigquery/WritePartition.java | 1 -
.../beam/sdk/io/gcp/bigquery/WriteTables.java | 1 -
.../beam/sdk/io/gcp/bigtable/BigtableIO.java | 8 +-
.../io/gcp/bigtable/BigtableServiceImpl.java | 9 +-
.../sdk/io/gcp/datastore/AdaptiveThrottler.java | 103 ---
.../beam/sdk/io/gcp/datastore/DatastoreV1.java | 149 +---
.../sdk/io/gcp/datastore/MovingAverage.java | 50 --
.../sdk/io/gcp/spanner/AbstractSpannerFn.java | 58 --
.../sdk/io/gcp/spanner/CreateTransactionFn.java | 51 --
.../beam/sdk/io/gcp/spanner/MutationGroup.java | 67 --
.../io/gcp/spanner/MutationSizeEstimator.java | 9 -
.../sdk/io/gcp/spanner/NaiveSpannerReadFn.java | 65 --
.../beam/sdk/io/gcp/spanner/SpannerConfig.java | 137 ----
.../beam/sdk/io/gcp/spanner/SpannerIO.java | 616 +++++-----------
.../sdk/io/gcp/spanner/SpannerWriteGroupFn.java | 125 ----
.../beam/sdk/io/gcp/spanner/Transaction.java | 33 -
.../beam/sdk/io/gcp/GcpApiSurfaceTest.java | 10 -
.../sdk/io/gcp/bigquery/BigQueryIOTest.java | 2 -
.../sdk/io/gcp/bigtable/BigtableReadIT.java | 5 +-
.../io/gcp/bigtable/BigtableTestOptions.java | 5 +
.../sdk/io/gcp/bigtable/BigtableWriteIT.java | 4 +-
.../io/gcp/datastore/AdaptiveThrottlerTest.java | 111 ---
.../sdk/io/gcp/datastore/DatastoreV1Test.java | 92 +--
.../beam/sdk/io/gcp/datastore/V1TestUtil.java | 2 +-
.../sdk/io/gcp/spanner/FakeServiceFactory.java | 82 ---
.../gcp/spanner/MutationSizeEstimatorTest.java | 12 -
.../beam/sdk/io/gcp/spanner/RandomUtils.java | 41 --
.../sdk/io/gcp/spanner/SpannerIOReadTest.java | 281 --------
.../beam/sdk/io/gcp/spanner/SpannerIOTest.java | 244 +++++++
.../sdk/io/gcp/spanner/SpannerIOWriteTest.java | 258 -------
.../beam/sdk/io/gcp/spanner/SpannerReadIT.java | 166 -----
.../beam/sdk/io/gcp/spanner/SpannerWriteIT.java | 26 +-
sdks/java/io/hadoop-common/pom.xml | 2 +-
sdks/java/io/hadoop-file-system/pom.xml | 33 +-
sdks/java/io/hadoop/input-format/pom.xml | 2 +-
.../hadoop/inputformat/HadoopInputFormatIO.java | 2 +-
sdks/java/io/hadoop/jdk1.8-tests/pom.xml | 4 +-
.../inputformat/HIFIOWithElasticTest.java | 11 +-
sdks/java/io/hadoop/pom.xml | 2 +-
sdks/java/io/hbase/pom.xml | 26 +-
.../io/hbase/HBaseCoderProviderRegistrar.java | 40 --
.../org/apache/beam/sdk/io/hbase/HBaseIO.java | 48 +-
.../beam/sdk/io/hbase/HBaseMutationCoder.java | 42 --
.../hbase/HBaseCoderProviderRegistrarTest.java | 45 --
.../apache/beam/sdk/io/hbase/HBaseIOTest.java | 49 +-
sdks/java/io/hcatalog/pom.xml | 175 -----
.../apache/beam/sdk/io/hcatalog/HCatalogIO.java | 492 -------------
.../beam/sdk/io/hcatalog/package-info.java | 22 -
.../io/hcatalog/EmbeddedMetastoreService.java | 87 ---
.../beam/sdk/io/hcatalog/HCatalogIOTest.java | 277 --------
.../sdk/io/hcatalog/HCatalogIOTestUtils.java | 108 ---
.../hcatalog/src/test/resources/hive-site.xml | 301 --------
sdks/java/io/jdbc/pom.xml | 4 +-
.../org/apache/beam/sdk/io/jdbc/JdbcIO.java | 2 +-
sdks/java/io/jms/pom.xml | 2 +-
.../java/org/apache/beam/sdk/io/jms/JmsIO.java | 2 +-
sdks/java/io/kafka/pom.xml | 2 +-
.../org/apache/beam/sdk/io/kafka/KafkaIO.java | 132 ++--
.../apache/beam/sdk/io/kafka/KafkaIOTest.java | 30 -
sdks/java/io/kinesis/pom.xml | 2 +-
.../sdk/io/kinesis/CheckpointGenerator.java | 6 +-
.../beam/sdk/io/kinesis/CustomOptional.java | 111 ++-
.../io/kinesis/DynamicCheckpointGenerator.java | 52 +-
.../sdk/io/kinesis/GetKinesisRecordsResult.java | 49 +-
.../sdk/io/kinesis/KinesisClientProvider.java | 4 +-
.../apache/beam/sdk/io/kinesis/KinesisIO.java | 281 ++++----
.../beam/sdk/io/kinesis/KinesisReader.java | 206 +++---
.../sdk/io/kinesis/KinesisReaderCheckpoint.java | 97 ++-
.../beam/sdk/io/kinesis/KinesisRecord.java | 177 +++--
.../beam/sdk/io/kinesis/KinesisRecordCoder.java | 68 +-
.../beam/sdk/io/kinesis/KinesisSource.java | 147 ++--
.../beam/sdk/io/kinesis/RecordFilter.java | 18 +-
.../apache/beam/sdk/io/kinesis/RoundRobin.java | 37 +-
.../beam/sdk/io/kinesis/ShardCheckpoint.java | 241 ++++---
.../sdk/io/kinesis/ShardRecordsIterator.java | 106 ++-
.../sdk/io/kinesis/SimplifiedKinesisClient.java | 215 +++---
.../beam/sdk/io/kinesis/StartingPoint.java | 84 ++-
.../io/kinesis/StaticCheckpointGenerator.java | 27 +-
.../io/kinesis/TransientKinesisException.java | 7 +-
.../beam/sdk/io/kinesis/AmazonKinesisMock.java | 539 +++++++-------
.../beam/sdk/io/kinesis/CustomOptionalTest.java | 27 +-
.../kinesis/DynamicCheckpointGeneratorTest.java | 33 +-
.../sdk/io/kinesis/KinesisMockReadTest.java | 97 ++-
.../io/kinesis/KinesisReaderCheckpointTest.java | 52 +-
.../beam/sdk/io/kinesis/KinesisReaderIT.java | 127 ++--
.../beam/sdk/io/kinesis/KinesisReaderTest.java | 166 +++--
.../sdk/io/kinesis/KinesisRecordCoderTest.java | 34 +-
.../beam/sdk/io/kinesis/KinesisTestOptions.java | 43 +-
.../beam/sdk/io/kinesis/KinesisUploader.java | 70 +-
.../beam/sdk/io/kinesis/RecordFilterTest.java | 52 +-
.../beam/sdk/io/kinesis/RoundRobinTest.java | 42 +-
.../sdk/io/kinesis/ShardCheckpointTest.java | 203 +++---
.../io/kinesis/ShardRecordsIteratorTest.java | 216 +++---
.../io/kinesis/SimplifiedKinesisClientTest.java | 351 +++++----
sdks/java/io/mongodb/pom.xml | 2 +-
.../beam/sdk/io/mongodb/MongoDbGridFSIO.java | 2 +-
.../apache/beam/sdk/io/mongodb/MongoDbIO.java | 317 ++-------
.../beam/sdk/io/mongodb/MongoDbIOTest.java | 37 -
sdks/java/io/mqtt/pom.xml | 2 +-
.../org/apache/beam/sdk/io/mqtt/MqttIO.java | 2 +-
sdks/java/io/pom.xml | 35 +-
sdks/java/io/xml/pom.xml | 2 +-
.../java/org/apache/beam/sdk/io/xml/XmlIO.java | 4 +-
.../org/apache/beam/sdk/io/xml/XmlSink.java | 21 +-
.../org/apache/beam/sdk/io/xml/XmlSinkTest.java | 4 +-
sdks/java/java8tests/pom.xml | 2 +-
sdks/java/javadoc/pom.xml | 19 +-
.../maven-archetypes/examples-java8/pom.xml | 2 +-
.../main/resources/archetype-resources/pom.xml | 1 +
sdks/java/maven-archetypes/examples/pom.xml | 2 +-
.../main/resources/archetype-resources/pom.xml | 1 +
sdks/java/maven-archetypes/pom.xml | 2 +-
sdks/java/maven-archetypes/starter/pom.xml | 2 +-
.../resources/projects/basic/reference/pom.xml | 2 +-
sdks/java/pom.xml | 2 +-
sdks/pom.xml | 2 +-
sdks/python/apache_beam/coders/coder_impl.py | 4 -
sdks/python/apache_beam/coders/coders.py | 7 +-
.../apache_beam/coders/coders_test_common.py | 8 -
.../examples/snippets/snippets_test.py | 16 -
.../apache_beam/examples/streaming_wordcount.py | 25 +-
.../apache_beam/examples/windowed_wordcount.py | 93 ---
sdks/python/apache_beam/io/filesystem.py | 22 +-
sdks/python/apache_beam/io/gcp/gcsio.py | 10 +-
sdks/python/apache_beam/io/gcp/pubsub.py | 180 ++---
sdks/python/apache_beam/io/gcp/pubsub_test.py | 101 +--
.../io/gcp/tests/bigquery_matcher.py | 6 +-
.../io/gcp/tests/bigquery_matcher_test.py | 2 +-
sdks/python/apache_beam/io/range_trackers.py | 130 ++++
.../apache_beam/io/range_trackers_test.py | 186 +++++
.../apache_beam/options/pipeline_options.py | 35 +-
.../options/pipeline_options_test.py | 39 +-
.../apache_beam/options/value_provider_test.py | 93 ++-
sdks/python/apache_beam/pipeline.py | 230 +-----
sdks/python/apache_beam/pipeline_test.py | 53 --
sdks/python/apache_beam/portability/__init__.py | 18 -
.../apache_beam/portability/api/__init__.py | 21 -
sdks/python/apache_beam/pvalue.py | 2 +-
sdks/python/apache_beam/runners/api/__init__.py | 21 +
.../runners/dataflow/dataflow_runner.py | 91 +--
.../runners/dataflow/dataflow_runner_test.py | 24 +-
.../runners/dataflow/internal/apiclient.py | 35 +-
.../runners/dataflow/internal/apiclient_test.py | 29 +-
.../runners/dataflow/internal/dependency.py | 69 +-
.../runners/dataflow/native_io/iobase_test.py | 39 +-
.../dataflow/native_io/streaming_create.py | 72 --
.../runners/dataflow/ptransform_overrides.py | 52 --
.../runners/direct/bundle_factory.py | 2 +-
.../apache_beam/runners/direct/direct_runner.py | 108 ---
.../runners/direct/evaluation_context.py | 73 +-
.../apache_beam/runners/direct/executor.py | 135 ++--
.../runners/direct/transform_evaluator.py | 447 +-----------
.../runners/direct/transform_result.py | 41 ++
sdks/python/apache_beam/runners/direct/util.py | 67 --
.../runners/direct/watermark_manager.py | 100 +--
.../apache_beam/runners/pipeline_context.py | 19 +-
.../runners/portability/fn_api_runner.py | 306 ++++----
.../runners/portability/fn_api_runner_test.py | 31 +-
.../runners/worker/bundle_processor.py | 426 -----------
.../apache_beam/runners/worker/data_plane.py | 28 +-
.../runners/worker/data_plane_test.py | 2 +-
.../apache_beam/runners/worker/log_handler.py | 2 +-
.../runners/worker/log_handler_test.py | 2 +-
.../runners/worker/operation_specs.py | 9 +-
.../apache_beam/runners/worker/operations.py | 1 -
.../apache_beam/runners/worker/sdk_worker.py | 370 +++++++++-
.../runners/worker/sdk_worker_main.py | 2 +-
.../runners/worker/sdk_worker_test.py | 95 ++-
sdks/python/apache_beam/testing/test_stream.py | 5 -
.../apache_beam/testing/test_stream_test.py | 68 --
sdks/python/apache_beam/transforms/combiners.py | 8 -
.../apache_beam/transforms/combiners_test.py | 7 +-
sdks/python/apache_beam/transforms/core.py | 102 ++-
.../python/apache_beam/transforms/ptransform.py | 43 +-
sdks/python/apache_beam/transforms/trigger.py | 28 +-
sdks/python/apache_beam/transforms/window.py | 4 +-
.../apache_beam/typehints/trivial_inference.py | 3 +-
.../typehints/trivial_inference_test.py | 7 -
sdks/python/apache_beam/utils/plugin.py | 42 --
sdks/python/apache_beam/utils/timestamp.py | 5 -
sdks/python/apache_beam/utils/urns.py | 2 +-
sdks/python/apache_beam/version.py | 2 +-
sdks/python/gen_protos.py | 2 +-
sdks/python/pom.xml | 2 +-
sdks/python/run_pylint.sh | 2 +-
sdks/python/setup.py | 5 +-
462 files changed, 10754 insertions(+), 21718 deletions(-)
----------------------------------------------------------------------
[12/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/harness/src/main/java/org/apache/beam/runners/core/BeamFnDataReadRunner.java
----------------------------------------------------------------------
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/runners/core/BeamFnDataReadRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/runners/core/BeamFnDataReadRunner.java
index 9339347..f0fe274 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/runners/core/BeamFnDataReadRunner.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/runners/core/BeamFnDataReadRunner.java
@@ -18,28 +18,22 @@
package org.apache.beam.runners.core;
-import static com.google.common.collect.Iterables.getOnlyElement;
-
import com.fasterxml.jackson.databind.ObjectMapper;
-import com.google.auto.service.AutoService;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Multimap;
+import com.google.common.collect.FluentIterable;
+import com.google.common.collect.ImmutableList;
import com.google.protobuf.BytesValue;
import java.io.IOException;
import java.util.Collection;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
-import java.util.function.Consumer;
import java.util.function.Supplier;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
import org.apache.beam.fn.harness.fn.ThrowingConsumer;
-import org.apache.beam.fn.harness.fn.ThrowingRunnable;
import org.apache.beam.fn.v1.BeamFnApi;
import org.apache.beam.runners.dataflow.util.CloudObject;
import org.apache.beam.runners.dataflow.util.CloudObjects;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.common.runner.v1.RunnerApi;
-import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.KV;
import org.slf4j.Logger;
@@ -54,61 +48,9 @@ import org.slf4j.LoggerFactory;
* {@link #blockTillReadFinishes()} to finish.
*/
public class BeamFnDataReadRunner<OutputT> {
-
private static final Logger LOG = LoggerFactory.getLogger(BeamFnDataReadRunner.class);
- private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
- private static final String URN = "urn:org.apache.beam:source:runner:0.1";
-
- /** A registrar which provides a factory to handle reading from the Fn Api Data Plane. */
- @AutoService(PTransformRunnerFactory.Registrar.class)
- public static class Registrar implements
- PTransformRunnerFactory.Registrar {
-
- @Override
- public Map<String, PTransformRunnerFactory> getPTransformRunnerFactories() {
- return ImmutableMap.of(URN, new Factory());
- }
- }
- /** A factory for {@link BeamFnDataReadRunner}s. */
- static class Factory<OutputT>
- implements PTransformRunnerFactory<BeamFnDataReadRunner<OutputT>> {
-
- @Override
- public BeamFnDataReadRunner<OutputT> createRunnerForPTransform(
- PipelineOptions pipelineOptions,
- BeamFnDataClient beamFnDataClient,
- String pTransformId,
- RunnerApi.PTransform pTransform,
- Supplier<String> processBundleInstructionId,
- Map<String, RunnerApi.PCollection> pCollections,
- Map<String, RunnerApi.Coder> coders,
- Multimap<String, ThrowingConsumer<WindowedValue<?>>> pCollectionIdsToConsumers,
- Consumer<ThrowingRunnable> addStartFunction,
- Consumer<ThrowingRunnable> addFinishFunction) throws IOException {
-
- BeamFnApi.Target target = BeamFnApi.Target.newBuilder()
- .setPrimitiveTransformReference(pTransformId)
- .setName(getOnlyElement(pTransform.getOutputsMap().keySet()))
- .build();
- RunnerApi.Coder coderSpec = coders.get(pCollections.get(
- getOnlyElement(pTransform.getOutputsMap().values())).getCoderId());
- Collection<ThrowingConsumer<WindowedValue<OutputT>>> consumers =
- (Collection) pCollectionIdsToConsumers.get(
- getOnlyElement(pTransform.getOutputsMap().values()));
-
- BeamFnDataReadRunner<OutputT> runner = new BeamFnDataReadRunner<>(
- pTransform.getSpec(),
- processBundleInstructionId,
- target,
- coderSpec,
- beamFnDataClient,
- consumers);
- addStartFunction.accept(runner::registerInputLocation);
- addFinishFunction.accept(runner::blockTillReadFinishes);
- return runner;
- }
- }
+ private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
private final BeamFnApi.ApiServiceDescriptor apiServiceDescriptor;
private final Collection<ThrowingConsumer<WindowedValue<OutputT>>> consumers;
@@ -119,20 +61,20 @@ public class BeamFnDataReadRunner<OutputT> {
private CompletableFuture<Void> readFuture;
- BeamFnDataReadRunner(
+ public BeamFnDataReadRunner(
RunnerApi.FunctionSpec functionSpec,
Supplier<String> processBundleInstructionIdSupplier,
BeamFnApi.Target inputTarget,
RunnerApi.Coder coderSpec,
BeamFnDataClient beamFnDataClientFactory,
- Collection<ThrowingConsumer<WindowedValue<OutputT>>> consumers)
+ Map<String, Collection<ThrowingConsumer<WindowedValue<OutputT>>>> outputMap)
throws IOException {
this.apiServiceDescriptor = functionSpec.getParameter().unpack(BeamFnApi.RemoteGrpcPort.class)
.getApiServiceDescriptor();
this.inputTarget = inputTarget;
this.processBundleInstructionIdSupplier = processBundleInstructionIdSupplier;
this.beamFnDataClientFactory = beamFnDataClientFactory;
- this.consumers = consumers;
+ this.consumers = ImmutableList.copyOf(FluentIterable.concat(outputMap.values()));
@SuppressWarnings("unchecked")
Coder<WindowedValue<OutputT>> coder =
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/harness/src/main/java/org/apache/beam/runners/core/BeamFnDataWriteRunner.java
----------------------------------------------------------------------
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/runners/core/BeamFnDataWriteRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/runners/core/BeamFnDataWriteRunner.java
index c2a996b..a48df12 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/runners/core/BeamFnDataWriteRunner.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/runners/core/BeamFnDataWriteRunner.java
@@ -18,91 +18,30 @@
package org.apache.beam.runners.core;
-import static com.google.common.collect.Iterables.getOnlyElement;
-
import com.fasterxml.jackson.databind.ObjectMapper;
-import com.google.auto.service.AutoService;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Multimap;
import com.google.protobuf.BytesValue;
import java.io.IOException;
import java.util.Map;
-import java.util.function.Consumer;
import java.util.function.Supplier;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
import org.apache.beam.fn.harness.fn.CloseableThrowingConsumer;
-import org.apache.beam.fn.harness.fn.ThrowingConsumer;
-import org.apache.beam.fn.harness.fn.ThrowingRunnable;
import org.apache.beam.fn.v1.BeamFnApi;
import org.apache.beam.runners.dataflow.util.CloudObject;
import org.apache.beam.runners.dataflow.util.CloudObjects;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.common.runner.v1.RunnerApi;
-import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.KV;
/**
- * Registers as a consumer with the Beam Fn Data Api. Consumes elements and encodes them for
- * transmission.
+ * Registers as a consumer with the Beam Fn Data API. Propagates and elements consumed to
+ * the the registered consumer.
*
* <p>Can be re-used serially across {@link org.apache.beam.fn.v1.BeamFnApi.ProcessBundleRequest}s.
* For each request, call {@link #registerForOutput()} to start and call {@link #close()} to finish.
*/
public class BeamFnDataWriteRunner<InputT> {
-
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
- private static final String URN = "urn:org.apache.beam:sink:runner:0.1";
-
- /** A registrar which provides a factory to handle writing to the Fn Api Data Plane. */
- @AutoService(PTransformRunnerFactory.Registrar.class)
- public static class Registrar implements
- PTransformRunnerFactory.Registrar {
-
- @Override
- public Map<String, PTransformRunnerFactory> getPTransformRunnerFactories() {
- return ImmutableMap.of(URN, new Factory());
- }
- }
-
- /** A factory for {@link BeamFnDataWriteRunner}s. */
- static class Factory<InputT>
- implements PTransformRunnerFactory<BeamFnDataWriteRunner<InputT>> {
-
- @Override
- public BeamFnDataWriteRunner<InputT> createRunnerForPTransform(
- PipelineOptions pipelineOptions,
- BeamFnDataClient beamFnDataClient,
- String pTransformId,
- RunnerApi.PTransform pTransform,
- Supplier<String> processBundleInstructionId,
- Map<String, RunnerApi.PCollection> pCollections,
- Map<String, RunnerApi.Coder> coders,
- Multimap<String, ThrowingConsumer<WindowedValue<?>>> pCollectionIdsToConsumers,
- Consumer<ThrowingRunnable> addStartFunction,
- Consumer<ThrowingRunnable> addFinishFunction) throws IOException {
- BeamFnApi.Target target = BeamFnApi.Target.newBuilder()
- .setPrimitiveTransformReference(pTransformId)
- .setName(getOnlyElement(pTransform.getInputsMap().keySet()))
- .build();
- RunnerApi.Coder coderSpec = coders.get(
- pCollections.get(getOnlyElement(pTransform.getInputsMap().values())).getCoderId());
- BeamFnDataWriteRunner<InputT> runner =
- new BeamFnDataWriteRunner<>(
- pTransform.getSpec(),
- processBundleInstructionId,
- target,
- coderSpec,
- beamFnDataClient);
- addStartFunction.accept(runner::registerForOutput);
- pCollectionIdsToConsumers.put(
- getOnlyElement(pTransform.getInputsMap().values()),
- (ThrowingConsumer)
- (ThrowingConsumer<WindowedValue<InputT>>) runner::consume);
- addFinishFunction.accept(runner::close);
- return runner;
- }
- }
private final BeamFnApi.ApiServiceDescriptor apiServiceDescriptor;
private final BeamFnApi.Target outputTarget;
@@ -112,7 +51,7 @@ public class BeamFnDataWriteRunner<InputT> {
private CloseableThrowingConsumer<WindowedValue<InputT>> consumer;
- BeamFnDataWriteRunner(
+ public BeamFnDataWriteRunner(
RunnerApi.FunctionSpec functionSpec,
Supplier<String> processBundleInstructionIdSupplier,
BeamFnApi.Target outputTarget,
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/harness/src/main/java/org/apache/beam/runners/core/BoundedSourceRunner.java
----------------------------------------------------------------------
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/runners/core/BoundedSourceRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/runners/core/BoundedSourceRunner.java
index 3338c3a..4d530b8 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/runners/core/BoundedSourceRunner.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/runners/core/BoundedSourceRunner.java
@@ -18,20 +18,14 @@
package org.apache.beam.runners.core;
-import com.google.auto.service.AutoService;
+import com.google.common.collect.FluentIterable;
import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Multimap;
import com.google.protobuf.BytesValue;
import com.google.protobuf.InvalidProtocolBufferException;
import java.io.IOException;
import java.util.Collection;
import java.util.Map;
-import java.util.function.Consumer;
-import java.util.function.Supplier;
-import org.apache.beam.fn.harness.data.BeamFnDataClient;
import org.apache.beam.fn.harness.fn.ThrowingConsumer;
-import org.apache.beam.fn.harness.fn.ThrowingRunnable;
import org.apache.beam.sdk.common.runner.v1.RunnerApi;
import org.apache.beam.sdk.io.BoundedSource;
import org.apache.beam.sdk.io.Source.Reader;
@@ -40,77 +34,21 @@ import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.util.WindowedValue;
/**
- * A runner which creates {@link Reader}s for each {@link BoundedSource} sent as an input and
- * executes the {@link Reader}s read loop.
+ * A runner which creates {@link Reader}s for each {@link BoundedSource} and executes
+ * the {@link Reader}s read loop.
*/
public class BoundedSourceRunner<InputT extends BoundedSource<OutputT>, OutputT> {
-
- private static final String URN = "urn:org.apache.beam:source:java:0.1";
-
- /** A registrar which provides a factory to handle Java {@link BoundedSource}s. */
- @AutoService(PTransformRunnerFactory.Registrar.class)
- public static class Registrar implements
- PTransformRunnerFactory.Registrar {
-
- @Override
- public Map<String, PTransformRunnerFactory> getPTransformRunnerFactories() {
- return ImmutableMap.of(URN, new Factory());
- }
- }
-
- /** A factory for {@link BoundedSourceRunner}. */
- static class Factory<InputT extends BoundedSource<OutputT>, OutputT>
- implements PTransformRunnerFactory<BoundedSourceRunner<InputT, OutputT>> {
- @Override
- public BoundedSourceRunner<InputT, OutputT> createRunnerForPTransform(
- PipelineOptions pipelineOptions,
- BeamFnDataClient beamFnDataClient,
- String pTransformId,
- RunnerApi.PTransform pTransform,
- Supplier<String> processBundleInstructionId,
- Map<String, RunnerApi.PCollection> pCollections,
- Map<String, RunnerApi.Coder> coders,
- Multimap<String, ThrowingConsumer<WindowedValue<?>>> pCollectionIdsToConsumers,
- Consumer<ThrowingRunnable> addStartFunction,
- Consumer<ThrowingRunnable> addFinishFunction) {
-
- ImmutableList.Builder<ThrowingConsumer<WindowedValue<?>>> consumers = ImmutableList.builder();
- for (String pCollectionId : pTransform.getOutputsMap().values()) {
- consumers.addAll(pCollectionIdsToConsumers.get(pCollectionId));
- }
-
- @SuppressWarnings({"rawtypes", "unchecked"})
- BoundedSourceRunner<InputT, OutputT> runner = new BoundedSourceRunner(
- pipelineOptions,
- pTransform.getSpec(),
- consumers.build());
-
- // TODO: Remove and replace with source being sent across gRPC port
- addStartFunction.accept(runner::start);
-
- ThrowingConsumer runReadLoop =
- (ThrowingConsumer<WindowedValue<InputT>>) runner::runReadLoop;
- for (String pCollectionId : pTransform.getInputsMap().values()) {
- pCollectionIdsToConsumers.put(
- pCollectionId,
- runReadLoop);
- }
-
- return runner;
- }
- }
-
private final PipelineOptions pipelineOptions;
private final RunnerApi.FunctionSpec definition;
private final Collection<ThrowingConsumer<WindowedValue<OutputT>>> consumers;
- BoundedSourceRunner(
+ public BoundedSourceRunner(
PipelineOptions pipelineOptions,
RunnerApi.FunctionSpec definition,
- Collection<ThrowingConsumer<WindowedValue<OutputT>>> consumers) {
+ Map<String, Collection<ThrowingConsumer<WindowedValue<OutputT>>>> outputMap) {
this.pipelineOptions = pipelineOptions;
this.definition = definition;
- this.consumers = consumers;
+ this.consumers = ImmutableList.copyOf(FluentIterable.concat(outputMap.values()));
}
/**
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/harness/src/main/java/org/apache/beam/runners/core/FnApiDoFnRunner.java
----------------------------------------------------------------------
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/runners/core/FnApiDoFnRunner.java b/sdks/java/harness/src/main/java/org/apache/beam/runners/core/FnApiDoFnRunner.java
deleted file mode 100644
index b3cf3a7..0000000
--- a/sdks/java/harness/src/main/java/org/apache/beam/runners/core/FnApiDoFnRunner.java
+++ /dev/null
@@ -1,547 +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 com.google.auto.service.AutoService;
-import com.google.common.collect.Collections2;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.ImmutableMultimap;
-import com.google.common.collect.Multimap;
-import com.google.protobuf.ByteString;
-import com.google.protobuf.BytesValue;
-import com.google.protobuf.InvalidProtocolBufferException;
-import java.util.Collection;
-import java.util.HashSet;
-import java.util.Iterator;
-import java.util.Map;
-import java.util.Objects;
-import java.util.function.Consumer;
-import java.util.function.Supplier;
-import org.apache.beam.fn.harness.data.BeamFnDataClient;
-import org.apache.beam.fn.harness.fn.ThrowingConsumer;
-import org.apache.beam.fn.harness.fn.ThrowingRunnable;
-import org.apache.beam.runners.core.construction.ParDoTranslation;
-import org.apache.beam.runners.dataflow.util.DoFnInfo;
-import org.apache.beam.sdk.common.runner.v1.RunnerApi;
-import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.state.State;
-import org.apache.beam.sdk.state.TimeDomain;
-import org.apache.beam.sdk.state.Timer;
-import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.DoFn.OnTimerContext;
-import org.apache.beam.sdk.transforms.DoFn.ProcessContext;
-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.PaneInfo;
-import org.apache.beam.sdk.util.SerializableUtils;
-import org.apache.beam.sdk.util.UserCodeException;
-import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.sdk.values.WindowingStrategy;
-import org.joda.time.Instant;
-
-/**
- * A {@link DoFnRunner} specific to integrating with the Fn Api. This is to remove the layers
- * of abstraction caused by StateInternals/TimerInternals since they model state and timer
- * concepts differently.
- */
-public class FnApiDoFnRunner<InputT, OutputT> implements DoFnRunner<InputT, OutputT> {
- /**
- * A registrar which provides a factory to handle Java {@link DoFn}s.
- */
- @AutoService(PTransformRunnerFactory.Registrar.class)
- public static class Registrar implements
- PTransformRunnerFactory.Registrar {
-
- @Override
- public Map<String, PTransformRunnerFactory> getPTransformRunnerFactories() {
- return ImmutableMap.of(ParDoTranslation.CUSTOM_JAVA_DO_FN_URN, new Factory());
- }
- }
-
- /** A factory for {@link FnApiDoFnRunner}. */
- static class Factory<InputT, OutputT>
- implements PTransformRunnerFactory<DoFnRunner<InputT, OutputT>> {
-
- @Override
- public DoFnRunner<InputT, OutputT> createRunnerForPTransform(
- PipelineOptions pipelineOptions,
- BeamFnDataClient beamFnDataClient,
- String pTransformId,
- RunnerApi.PTransform pTransform,
- Supplier<String> processBundleInstructionId,
- Map<String, RunnerApi.PCollection> pCollections,
- Map<String, RunnerApi.Coder> coders,
- Multimap<String, ThrowingConsumer<WindowedValue<?>>> pCollectionIdsToConsumers,
- Consumer<ThrowingRunnable> addStartFunction,
- Consumer<ThrowingRunnable> addFinishFunction) {
-
- // For every output PCollection, create a map from output name to Consumer
- ImmutableMap.Builder<String, Collection<ThrowingConsumer<WindowedValue<?>>>>
- outputMapBuilder = ImmutableMap.builder();
- for (Map.Entry<String, String> entry : pTransform.getOutputsMap().entrySet()) {
- outputMapBuilder.put(
- entry.getKey(),
- pCollectionIdsToConsumers.get(entry.getValue()));
- }
- ImmutableMap<String, Collection<ThrowingConsumer<WindowedValue<?>>>> outputMap =
- outputMapBuilder.build();
-
- // Get the DoFnInfo from the serialized blob.
- ByteString serializedFn;
- try {
- serializedFn = pTransform.getSpec().getParameter().unpack(BytesValue.class).getValue();
- } catch (InvalidProtocolBufferException e) {
- throw new IllegalArgumentException(
- String.format("Unable to unwrap DoFn %s", pTransform.getSpec()), e);
- }
- @SuppressWarnings({"unchecked", "rawtypes"})
- DoFnInfo<InputT, OutputT> doFnInfo = (DoFnInfo) SerializableUtils.deserializeFromByteArray(
- serializedFn.toByteArray(), "DoFnInfo");
-
- // Verify that the DoFnInfo tag to output map matches the output map on the PTransform.
- checkArgument(
- Objects.equals(
- new HashSet<>(Collections2.transform(outputMap.keySet(), Long::parseLong)),
- doFnInfo.getOutputMap().keySet()),
- "Unexpected mismatch between transform output map %s and DoFnInfo output map %s.",
- outputMap.keySet(),
- doFnInfo.getOutputMap());
-
- ImmutableMultimap.Builder<TupleTag<?>,
- ThrowingConsumer<WindowedValue<?>>> tagToOutputMapBuilder =
- ImmutableMultimap.builder();
- for (Map.Entry<Long, TupleTag<?>> entry : doFnInfo.getOutputMap().entrySet()) {
- @SuppressWarnings({"unchecked", "rawtypes"})
- Collection<ThrowingConsumer<WindowedValue<?>>> consumers =
- outputMap.get(Long.toString(entry.getKey()));
- tagToOutputMapBuilder.putAll(entry.getValue(), consumers);
- }
-
- ImmutableMultimap<TupleTag<?>, ThrowingConsumer<WindowedValue<?>>> tagToOutputMap =
- tagToOutputMapBuilder.build();
-
- @SuppressWarnings({"unchecked", "rawtypes"})
- DoFnRunner<InputT, OutputT> runner = new FnApiDoFnRunner<>(
- pipelineOptions,
- doFnInfo.getDoFn(),
- (Collection<ThrowingConsumer<WindowedValue<OutputT>>>) (Collection)
- tagToOutputMap.get(doFnInfo.getOutputMap().get(doFnInfo.getMainOutput())),
- tagToOutputMap,
- doFnInfo.getWindowingStrategy());
-
- // Register the appropriate handlers.
- addStartFunction.accept(runner::startBundle);
- for (String pcollectionId : pTransform.getInputsMap().values()) {
- pCollectionIdsToConsumers.put(
- pcollectionId,
- (ThrowingConsumer) (ThrowingConsumer<WindowedValue<InputT>>) runner::processElement);
- }
- addFinishFunction.accept(runner::finishBundle);
- return runner;
- }
- }
-
- //////////////////////////////////////////////////////////////////////////////////////////////////
-
- private final PipelineOptions pipelineOptions;
- private final DoFn<InputT, OutputT> doFn;
- private final Collection<ThrowingConsumer<WindowedValue<OutputT>>> mainOutputConsumers;
- private final Multimap<TupleTag<?>, ThrowingConsumer<WindowedValue<?>>> outputMap;
- private final DoFnInvoker<InputT, OutputT> doFnInvoker;
- private final StartBundleContext startBundleContext;
- private final ProcessBundleContext processBundleContext;
- private final FinishBundleContext finishBundleContext;
-
- /**
- * The lifetime of this member is only valid during {@link #processElement(WindowedValue)}.
- */
- private WindowedValue<InputT> currentElement;
-
- /**
- * The lifetime of this member is only valid during {@link #processElement(WindowedValue)}.
- */
- private BoundedWindow currentWindow;
-
- FnApiDoFnRunner(
- PipelineOptions pipelineOptions,
- DoFn<InputT, OutputT> doFn,
- Collection<ThrowingConsumer<WindowedValue<OutputT>>> mainOutputConsumers,
- Multimap<TupleTag<?>, ThrowingConsumer<WindowedValue<?>>> outputMap,
- WindowingStrategy windowingStrategy) {
- this.pipelineOptions = pipelineOptions;
- this.doFn = doFn;
- this.mainOutputConsumers = mainOutputConsumers;
- this.outputMap = outputMap;
- this.doFnInvoker = DoFnInvokers.invokerFor(doFn);
- this.startBundleContext = new StartBundleContext();
- this.processBundleContext = new ProcessBundleContext();
- this.finishBundleContext = new FinishBundleContext();
- }
-
- @Override
- public void startBundle() {
- doFnInvoker.invokeStartBundle(startBundleContext);
- }
-
- @Override
- public void processElement(WindowedValue<InputT> elem) {
- currentElement = elem;
- try {
- Iterator<BoundedWindow> windowIterator =
- (Iterator<BoundedWindow>) elem.getWindows().iterator();
- while (windowIterator.hasNext()) {
- currentWindow = windowIterator.next();
- doFnInvoker.invokeProcessElement(processBundleContext);
- }
- } finally {
- currentElement = null;
- currentWindow = null;
- }
- }
-
- @Override
- public void onTimer(
- String timerId,
- BoundedWindow window,
- Instant timestamp,
- TimeDomain timeDomain) {
- throw new UnsupportedOperationException("TODO: Add support for timers");
- }
-
- @Override
- public void finishBundle() {
- doFnInvoker.invokeFinishBundle(finishBundleContext);
- }
-
- /**
- * Outputs the given element to the specified set of consumers wrapping any exceptions.
- */
- private <T> void outputTo(
- Collection<ThrowingConsumer<WindowedValue<T>>> consumers,
- WindowedValue<T> output) {
- Iterator<ThrowingConsumer<WindowedValue<T>>> consumerIterator;
- try {
- for (ThrowingConsumer<WindowedValue<T>> consumer : consumers) {
- consumer.accept(output);
- }
- } catch (Throwable t) {
- throw UserCodeException.wrap(t);
- }
- }
-
- /**
- * Provides arguments for a {@link DoFnInvoker} for {@link DoFn.StartBundle @StartBundle}.
- */
- private class StartBundleContext
- extends DoFn<InputT, OutputT>.StartBundleContext
- implements DoFnInvoker.ArgumentProvider<InputT, OutputT> {
-
- private StartBundleContext() {
- doFn.super();
- }
-
- @Override
- public PipelineOptions getPipelineOptions() {
- return pipelineOptions;
- }
-
- @Override
- public PipelineOptions pipelineOptions() {
- return pipelineOptions;
- }
-
- @Override
- public BoundedWindow window() {
- throw new UnsupportedOperationException(
- "Cannot access window outside of @ProcessElement and @OnTimer methods.");
- }
-
- @Override
- public DoFn<InputT, OutputT>.StartBundleContext startBundleContext(
- DoFn<InputT, OutputT> doFn) {
- return this;
- }
-
- @Override
- public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
- DoFn<InputT, OutputT> doFn) {
- throw new UnsupportedOperationException(
- "Cannot access FinishBundleContext outside of @FinishBundle method.");
- }
-
- @Override
- public DoFn<InputT, OutputT>.ProcessContext processContext(DoFn<InputT, OutputT> doFn) {
- throw new UnsupportedOperationException(
- "Cannot access ProcessContext outside of @ProcessElement method.");
- }
-
- @Override
- public DoFn<InputT, OutputT>.OnTimerContext onTimerContext(DoFn<InputT, OutputT> doFn) {
- throw new UnsupportedOperationException(
- "Cannot access OnTimerContext outside of @OnTimer methods.");
- }
-
- @Override
- public RestrictionTracker<?> restrictionTracker() {
- throw new UnsupportedOperationException(
- "Cannot access RestrictionTracker outside of @ProcessElement method.");
- }
-
- @Override
- public State state(String stateId) {
- throw new UnsupportedOperationException(
- "Cannot access state outside of @ProcessElement and @OnTimer methods.");
- }
-
- @Override
- public Timer timer(String timerId) {
- throw new UnsupportedOperationException(
- "Cannot access timers outside of @ProcessElement and @OnTimer methods.");
- }
- }
-
- /**
- * Provides arguments for a {@link DoFnInvoker} for {@link DoFn.ProcessElement @ProcessElement}.
- */
- private class ProcessBundleContext
- extends DoFn<InputT, OutputT>.ProcessContext
- implements DoFnInvoker.ArgumentProvider<InputT, OutputT> {
-
- private ProcessBundleContext() {
- doFn.super();
- }
-
- @Override
- public BoundedWindow window() {
- return currentWindow;
- }
-
- @Override
- public DoFn.StartBundleContext startBundleContext(DoFn<InputT, OutputT> doFn) {
- throw new UnsupportedOperationException(
- "Cannot access StartBundleContext outside of @StartBundle method.");
- }
-
- @Override
- public DoFn.FinishBundleContext finishBundleContext(DoFn<InputT, OutputT> doFn) {
- throw new UnsupportedOperationException(
- "Cannot access FinishBundleContext outside of @FinishBundle method.");
- }
-
- @Override
- public ProcessContext processContext(DoFn<InputT, OutputT> doFn) {
- return this;
- }
-
- @Override
- public OnTimerContext onTimerContext(DoFn<InputT, OutputT> doFn) {
- throw new UnsupportedOperationException("TODO: Add support for timers");
- }
-
- @Override
- public RestrictionTracker<?> restrictionTracker() {
- throw new UnsupportedOperationException("TODO: Add support for SplittableDoFn");
- }
-
- @Override
- public State state(String stateId) {
- throw new UnsupportedOperationException("TODO: Add support for state");
- }
-
- @Override
- public Timer timer(String timerId) {
- throw new UnsupportedOperationException("TODO: Add support for timers");
- }
-
- @Override
- public PipelineOptions getPipelineOptions() {
- return pipelineOptions;
- }
-
- @Override
- public PipelineOptions pipelineOptions() {
- return pipelineOptions;
- }
-
- @Override
- public void output(OutputT output) {
- outputTo(mainOutputConsumers,
- WindowedValue.of(
- output,
- currentElement.getTimestamp(),
- currentWindow,
- currentElement.getPane()));
- }
-
- @Override
- public void outputWithTimestamp(OutputT output, Instant timestamp) {
- outputTo(mainOutputConsumers,
- WindowedValue.of(
- output,
- timestamp,
- currentWindow,
- currentElement.getPane()));
- }
-
- @Override
- public <T> void output(TupleTag<T> tag, T output) {
- Collection<ThrowingConsumer<WindowedValue<T>>> consumers = (Collection) outputMap.get(tag);
- if (consumers == null) {
- throw new IllegalArgumentException(String.format("Unknown output tag %s", tag));
- }
- outputTo(consumers,
- WindowedValue.of(
- output,
- currentElement.getTimestamp(),
- currentWindow,
- currentElement.getPane()));
- }
-
- @Override
- public <T> void outputWithTimestamp(TupleTag<T> tag, T output, Instant timestamp) {
- Collection<ThrowingConsumer<WindowedValue<T>>> consumers = (Collection) outputMap.get(tag);
- if (consumers == null) {
- throw new IllegalArgumentException(String.format("Unknown output tag %s", tag));
- }
- outputTo(consumers,
- WindowedValue.of(
- output,
- timestamp,
- currentWindow,
- currentElement.getPane()));
- }
-
- @Override
- public InputT element() {
- return currentElement.getValue();
- }
-
- @Override
- public <T> T sideInput(PCollectionView<T> view) {
- throw new UnsupportedOperationException("TODO: Support side inputs");
- }
-
- @Override
- public Instant timestamp() {
- return currentElement.getTimestamp();
- }
-
- @Override
- public PaneInfo pane() {
- return currentElement.getPane();
- }
-
- @Override
- public void updateWatermark(Instant watermark) {
- throw new UnsupportedOperationException("TODO: Add support for SplittableDoFn");
- }
- }
-
- /**
- * Provides arguments for a {@link DoFnInvoker} for {@link DoFn.FinishBundle @FinishBundle}.
- */
- private class FinishBundleContext
- extends DoFn<InputT, OutputT>.FinishBundleContext
- implements DoFnInvoker.ArgumentProvider<InputT, OutputT> {
-
- private FinishBundleContext() {
- doFn.super();
- }
-
- @Override
- public PipelineOptions getPipelineOptions() {
- return pipelineOptions;
- }
-
- @Override
- public PipelineOptions pipelineOptions() {
- return pipelineOptions;
- }
-
- @Override
- public BoundedWindow window() {
- throw new UnsupportedOperationException(
- "Cannot access window outside of @ProcessElement and @OnTimer methods.");
- }
-
- @Override
- public DoFn<InputT, OutputT>.StartBundleContext startBundleContext(
- DoFn<InputT, OutputT> doFn) {
- throw new UnsupportedOperationException(
- "Cannot access StartBundleContext outside of @StartBundle method.");
- }
-
- @Override
- public DoFn<InputT, OutputT>.FinishBundleContext finishBundleContext(
- DoFn<InputT, OutputT> doFn) {
- return this;
- }
-
- @Override
- public DoFn<InputT, OutputT>.ProcessContext processContext(DoFn<InputT, OutputT> doFn) {
- throw new UnsupportedOperationException(
- "Cannot access ProcessContext outside of @ProcessElement method.");
- }
-
- @Override
- public DoFn<InputT, OutputT>.OnTimerContext onTimerContext(DoFn<InputT, OutputT> doFn) {
- throw new UnsupportedOperationException(
- "Cannot access OnTimerContext outside of @OnTimer methods.");
- }
-
- @Override
- public RestrictionTracker<?> restrictionTracker() {
- throw new UnsupportedOperationException(
- "Cannot access RestrictionTracker outside of @ProcessElement method.");
- }
-
- @Override
- public State state(String stateId) {
- throw new UnsupportedOperationException(
- "Cannot access state outside of @ProcessElement and @OnTimer methods.");
- }
-
- @Override
- public Timer timer(String timerId) {
- throw new UnsupportedOperationException(
- "Cannot access timers outside of @ProcessElement and @OnTimer methods.");
- }
-
- @Override
- public void output(OutputT output, Instant timestamp, BoundedWindow window) {
- outputTo(mainOutputConsumers,
- WindowedValue.of(output, timestamp, window, PaneInfo.NO_FIRING));
- }
-
- @Override
- public <T> void output(TupleTag<T> tag, T output, Instant timestamp, BoundedWindow window) {
- Collection<ThrowingConsumer<WindowedValue<T>>> consumers = (Collection) outputMap.get(tag);
- if (consumers == null) {
- throw new IllegalArgumentException(String.format("Unknown output tag %s", tag));
- }
- outputTo(consumers,
- WindowedValue.of(output, timestamp, window, PaneInfo.NO_FIRING));
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/harness/src/main/java/org/apache/beam/runners/core/PTransformRunnerFactory.java
----------------------------------------------------------------------
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/runners/core/PTransformRunnerFactory.java b/sdks/java/harness/src/main/java/org/apache/beam/runners/core/PTransformRunnerFactory.java
deleted file mode 100644
index b325db4..0000000
--- a/sdks/java/harness/src/main/java/org/apache/beam/runners/core/PTransformRunnerFactory.java
+++ /dev/null
@@ -1,81 +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.Multimap;
-import java.io.IOException;
-import java.util.Map;
-import java.util.function.Consumer;
-import java.util.function.Supplier;
-import org.apache.beam.fn.harness.data.BeamFnDataClient;
-import org.apache.beam.fn.harness.fn.ThrowingConsumer;
-import org.apache.beam.fn.harness.fn.ThrowingRunnable;
-import org.apache.beam.sdk.common.runner.v1.RunnerApi;
-import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.util.WindowedValue;
-
-/**
- * A factory able to instantiate an appropriate handler for a given PTransform.
- */
-public interface PTransformRunnerFactory<T> {
-
- /**
- * Creates and returns a handler for a given PTransform. Note that the handler must support
- * processing multiple bundles. The handler will be discarded if an error is thrown during
- * element processing, or during execution of start/finish.
- *
- * @param pipelineOptions Pipeline options
- * @param beamFnDataClient
- * @param pTransformId The id of the PTransform.
- * @param pTransform The PTransform definition.
- * @param processBundleInstructionId A supplier containing the active process bundle instruction
- * id.
- * @param pCollections A mapping from PCollection id to PCollection definition.
- * @param coders A mapping from coder id to coder definition.
- * @param pCollectionIdsToConsumers A mapping from PCollection id to a collection of consumers.
- * Note that if this handler is a consumer, it should register itself within this multimap under
- * the appropriate PCollection ids. Also note that all output consumers needed by this PTransform
- * (based on the values of the {@link RunnerApi.PTransform#getOutputsMap()} will have already
- * registered within this multimap.
- * @param addStartFunction A consumer to register a start bundle handler with.
- * @param addFinishFunction A consumer to register a finish bundle handler with.
- */
- T createRunnerForPTransform(
- PipelineOptions pipelineOptions,
- BeamFnDataClient beamFnDataClient,
- String pTransformId,
- RunnerApi.PTransform pTransform,
- Supplier<String> processBundleInstructionId,
- Map<String, RunnerApi.PCollection> pCollections,
- Map<String, RunnerApi.Coder> coders,
- Multimap<String, ThrowingConsumer<WindowedValue<?>>> pCollectionIdsToConsumers,
- Consumer<ThrowingRunnable> addStartFunction,
- Consumer<ThrowingRunnable> addFinishFunction) throws IOException;
-
- /**
- * A registrar which can return a mapping from {@link RunnerApi.FunctionSpec#getUrn()} to
- * a factory capable of instantiating an appropriate handler.
- */
- interface Registrar {
- /**
- * Returns a mapping from {@link RunnerApi.FunctionSpec#getUrn()} to a factory capable of
- * instantiating an appropriate handler.
- */
- Map<String, PTransformRunnerFactory> getPTransformRunnerFactories();
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java
index a616b2c..562f91f 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/ProcessBundleHandlerTest.java
@@ -18,28 +18,62 @@
package org.apache.beam.fn.harness.control;
+import static org.apache.beam.sdk.util.WindowedValue.timestampedValueInGlobalWindow;
+import static org.apache.beam.sdk.util.WindowedValue.valueInGlobalWindow;
import static org.hamcrest.Matchers.contains;
+import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.hamcrest.Matchers.empty;
import static org.hamcrest.Matchers.equalTo;
+import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.verifyZeroInteractions;
+import static org.mockito.Mockito.when;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.base.Suppliers;
+import com.google.common.collect.HashMultimap;
+import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Iterables;
import com.google.common.collect.Multimap;
+import com.google.protobuf.Any;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.BytesValue;
import com.google.protobuf.Message;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;
import java.util.function.Supplier;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
+import org.apache.beam.fn.harness.fn.CloseableThrowingConsumer;
import org.apache.beam.fn.harness.fn.ThrowingConsumer;
import org.apache.beam.fn.harness.fn.ThrowingRunnable;
import org.apache.beam.fn.v1.BeamFnApi;
-import org.apache.beam.runners.core.PTransformRunnerFactory;
+import org.apache.beam.runners.dataflow.util.CloudObjects;
+import org.apache.beam.runners.dataflow.util.DoFnInfo;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.coders.VarLongCoder;
import org.apache.beam.sdk.common.runner.v1.RunnerApi;
-import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.io.CountingSource;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.WindowingStrategy;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
@@ -48,14 +82,55 @@ import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.ArgumentCaptor;
import org.mockito.Captor;
+import org.mockito.Matchers;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
/** Tests for {@link ProcessBundleHandler}. */
@RunWith(JUnit4.class)
public class ProcessBundleHandlerTest {
+ private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+ private static final Coder<WindowedValue<String>> STRING_CODER =
+ WindowedValue.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE);
+ private static final String LONG_CODER_SPEC_ID = "998L";
+ private static final String STRING_CODER_SPEC_ID = "999L";
+ private static final BeamFnApi.RemoteGrpcPort REMOTE_PORT = BeamFnApi.RemoteGrpcPort.newBuilder()
+ .setApiServiceDescriptor(BeamFnApi.ApiServiceDescriptor.newBuilder()
+ .setId("58L")
+ .setUrl("TestUrl"))
+ .build();
+ private static final RunnerApi.Coder LONG_CODER_SPEC;
+ private static final RunnerApi.Coder STRING_CODER_SPEC;
+ static {
+ try {
+ STRING_CODER_SPEC = RunnerApi.Coder.newBuilder()
+ .setSpec(RunnerApi.SdkFunctionSpec.newBuilder()
+ .setSpec(RunnerApi.FunctionSpec.newBuilder()
+ .setParameter(Any.pack(BytesValue.newBuilder().setValue(ByteString.copyFrom(
+ OBJECT_MAPPER.writeValueAsBytes(CloudObjects.asCloudObject(STRING_CODER))))
+ .build())))
+ .build())
+ .build();
+ LONG_CODER_SPEC = RunnerApi.Coder.newBuilder()
+ .setSpec(RunnerApi.SdkFunctionSpec.newBuilder()
+ .setSpec(RunnerApi.FunctionSpec.newBuilder()
+ .setParameter(Any.pack(BytesValue.newBuilder().setValue(ByteString.copyFrom(
+ OBJECT_MAPPER.writeValueAsBytes(
+ CloudObjects.asCloudObject(WindowedValue.getFullCoder(VarLongCoder.of(),
+ GlobalWindow.Coder.INSTANCE)))))
+ .build())))
+ .build())
+ .build();
+ } catch (IOException e) {
+ throw new ExceptionInInitializerError(e);
+ }
+ }
+
private static final String DATA_INPUT_URN = "urn:org.apache.beam:source:runner:0.1";
private static final String DATA_OUTPUT_URN = "urn:org.apache.beam:sink:runner:0.1";
+ private static final String JAVA_DO_FN_URN = "urn:org.apache.beam:dofn:java:0.1";
+ private static final String JAVA_SOURCE_URN = "urn:org.apache.beam:source:java:0.1";
@Rule public ExpectedException thrown = ExpectedException.none();
@@ -86,16 +161,16 @@ public class ProcessBundleHandlerTest {
List<RunnerApi.PTransform> transformsProcessed = new ArrayList<>();
List<String> orderOfOperations = new ArrayList<>();
- PTransformRunnerFactory<Object> startFinishRecorder = new PTransformRunnerFactory<Object>() {
+ ProcessBundleHandler handler = new ProcessBundleHandler(
+ PipelineOptionsFactory.create(),
+ fnApiRegistry::get,
+ beamFnDataClient) {
@Override
- public Object createRunnerForPTransform(
- PipelineOptions pipelineOptions,
- BeamFnDataClient beamFnDataClient,
+ protected void createRunnerForPTransform(
String pTransformId,
RunnerApi.PTransform pTransform,
Supplier<String> processBundleInstructionId,
Map<String, RunnerApi.PCollection> pCollections,
- Map<String, RunnerApi.Coder> coders,
Multimap<String, ThrowingConsumer<WindowedValue<?>>> pCollectionIdsToConsumers,
Consumer<ThrowingRunnable> addStartFunction,
Consumer<ThrowingRunnable> addFinishFunction) throws IOException {
@@ -107,18 +182,8 @@ public class ProcessBundleHandlerTest {
() -> orderOfOperations.add("Start" + pTransformId));
addFinishFunction.accept(
() -> orderOfOperations.add("Finish" + pTransformId));
- return null;
}
};
-
- ProcessBundleHandler handler = new ProcessBundleHandler(
- PipelineOptionsFactory.create(),
- fnApiRegistry::get,
- beamFnDataClient,
- ImmutableMap.of(
- DATA_INPUT_URN, startFinishRecorder,
- DATA_OUTPUT_URN, startFinishRecorder));
-
handler.processBundle(BeamFnApi.InstructionRequest.newBuilder()
.setInstructionId("999L")
.setProcessBundle(
@@ -146,25 +211,21 @@ public class ProcessBundleHandlerTest {
ProcessBundleHandler handler = new ProcessBundleHandler(
PipelineOptionsFactory.create(),
fnApiRegistry::get,
- beamFnDataClient,
- ImmutableMap.of(DATA_INPUT_URN, new PTransformRunnerFactory<Object>() {
- @Override
- public Object createRunnerForPTransform(
- PipelineOptions pipelineOptions,
- BeamFnDataClient beamFnDataClient,
- String pTransformId,
- RunnerApi.PTransform pTransform,
- Supplier<String> processBundleInstructionId,
- Map<String, RunnerApi.PCollection> pCollections,
- Map<String, RunnerApi.Coder> coders,
- Multimap<String, ThrowingConsumer<WindowedValue<?>>> pCollectionIdsToConsumers,
- Consumer<ThrowingRunnable> addStartFunction,
- Consumer<ThrowingRunnable> addFinishFunction) throws IOException {
- thrown.expect(IllegalStateException.class);
- thrown.expectMessage("TestException");
- throw new IllegalStateException("TestException");
- }
- }));
+ beamFnDataClient) {
+ @Override
+ protected void createRunnerForPTransform(
+ String pTransformId,
+ RunnerApi.PTransform pTransform,
+ Supplier<String> processBundleInstructionId,
+ Map<String, RunnerApi.PCollection> pCollections,
+ Multimap<String, ThrowingConsumer<WindowedValue<?>>> pCollectionIdsToConsumers,
+ Consumer<ThrowingRunnable> addStartFunction,
+ Consumer<ThrowingRunnable> addFinishFunction) throws IOException {
+ thrown.expect(IllegalStateException.class);
+ thrown.expectMessage("TestException");
+ throw new IllegalStateException("TestException");
+ }
+ };
handler.processBundle(
BeamFnApi.InstructionRequest.newBuilder().setProcessBundle(
BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorReference("1L"))
@@ -184,26 +245,25 @@ public class ProcessBundleHandlerTest {
ProcessBundleHandler handler = new ProcessBundleHandler(
PipelineOptionsFactory.create(),
fnApiRegistry::get,
- beamFnDataClient,
- ImmutableMap.of(DATA_INPUT_URN, new PTransformRunnerFactory<Object>() {
- @Override
- public Object createRunnerForPTransform(
- PipelineOptions pipelineOptions,
- BeamFnDataClient beamFnDataClient,
- String pTransformId,
- RunnerApi.PTransform pTransform,
- Supplier<String> processBundleInstructionId,
- Map<String, RunnerApi.PCollection> pCollections,
- Map<String, RunnerApi.Coder> coders,
- Multimap<String, ThrowingConsumer<WindowedValue<?>>> pCollectionIdsToConsumers,
- Consumer<ThrowingRunnable> addStartFunction,
- Consumer<ThrowingRunnable> addFinishFunction) throws IOException {
- thrown.expect(IllegalStateException.class);
- thrown.expectMessage("TestException");
- addStartFunction.accept(ProcessBundleHandlerTest::throwException);
- return null;
- }
- }));
+ beamFnDataClient) {
+ @Override
+ protected void createRunnerForPTransform(
+ String pTransformId,
+ RunnerApi.PTransform pTransform,
+ Supplier<String> processBundleInstructionId,
+ Map<String, RunnerApi.PCollection> pCollections,
+ Multimap<String, ThrowingConsumer<WindowedValue<?>>> pCollectionIdsToConsumers,
+ Consumer<ThrowingRunnable> addStartFunction,
+ Consumer<ThrowingRunnable> addFinishFunction) throws IOException {
+ thrown.expect(IllegalStateException.class);
+ thrown.expectMessage("TestException");
+ addStartFunction.accept(this::throwException);
+ }
+
+ private void throwException() {
+ throw new IllegalStateException("TestException");
+ }
+ };
handler.processBundle(
BeamFnApi.InstructionRequest.newBuilder().setProcessBundle(
BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorReference("1L"))
@@ -223,33 +283,338 @@ public class ProcessBundleHandlerTest {
ProcessBundleHandler handler = new ProcessBundleHandler(
PipelineOptionsFactory.create(),
fnApiRegistry::get,
- beamFnDataClient,
- ImmutableMap.of(DATA_INPUT_URN, new PTransformRunnerFactory<Object>() {
- @Override
- public Object createRunnerForPTransform(
- PipelineOptions pipelineOptions,
- BeamFnDataClient beamFnDataClient,
- String pTransformId,
- RunnerApi.PTransform pTransform,
- Supplier<String> processBundleInstructionId,
- Map<String, RunnerApi.PCollection> pCollections,
- Map<String, RunnerApi.Coder> coders,
- Multimap<String, ThrowingConsumer<WindowedValue<?>>> pCollectionIdsToConsumers,
- Consumer<ThrowingRunnable> addStartFunction,
- Consumer<ThrowingRunnable> addFinishFunction) throws IOException {
- thrown.expect(IllegalStateException.class);
- thrown.expectMessage("TestException");
- addFinishFunction.accept(ProcessBundleHandlerTest::throwException);
- return null;
- }
- }));
+ beamFnDataClient) {
+ @Override
+ protected void createRunnerForPTransform(
+ String pTransformId,
+ RunnerApi.PTransform pTransform,
+ Supplier<String> processBundleInstructionId,
+ Map<String, RunnerApi.PCollection> pCollections,
+ Multimap<String, ThrowingConsumer<WindowedValue<?>>> pCollectionIdsToConsumers,
+ Consumer<ThrowingRunnable> addStartFunction,
+ Consumer<ThrowingRunnable> addFinishFunction) throws IOException {
+ thrown.expect(IllegalStateException.class);
+ thrown.expectMessage("TestException");
+ addFinishFunction.accept(this::throwException);
+ }
+
+ private void throwException() {
+ throw new IllegalStateException("TestException");
+ }
+ };
handler.processBundle(
BeamFnApi.InstructionRequest.newBuilder().setProcessBundle(
BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorReference("1L"))
.build());
}
- private static void throwException() {
- throw new IllegalStateException("TestException");
+ private static class TestDoFn extends DoFn<String, String> {
+ private static final TupleTag<String> mainOutput = new TupleTag<>("mainOutput");
+ private static final TupleTag<String> additionalOutput = new TupleTag<>("output");
+
+ private BoundedWindow window;
+
+ @ProcessElement
+ public void processElement(ProcessContext context, BoundedWindow window) {
+ context.output("MainOutput" + context.element());
+ context.output(additionalOutput, "AdditionalOutput" + context.element());
+ this.window = window;
+ }
+
+ @FinishBundle
+ public void finishBundle(FinishBundleContext context) {
+ if (window != null) {
+ context.output("FinishBundle", window.maxTimestamp(), window);
+ window = null;
+ }
+ }
+ }
+
+ /**
+ * Create a DoFn that has 3 inputs (inputATarget1, inputATarget2, inputBTarget) and 2 outputs
+ * (mainOutput, output). Validate that inputs are fed to the {@link DoFn} and that outputs
+ * are directed to the correct consumers.
+ */
+ @Test
+ public void testCreatingAndProcessingDoFn() throws Exception {
+ Map<String, Message> fnApiRegistry = ImmutableMap.of(STRING_CODER_SPEC_ID, STRING_CODER_SPEC);
+ String pTransformId = "100L";
+ String mainOutputId = "101";
+ String additionalOutputId = "102";
+
+ DoFnInfo<?, ?> doFnInfo = DoFnInfo.forFn(
+ new TestDoFn(),
+ WindowingStrategy.globalDefault(),
+ ImmutableList.of(),
+ StringUtf8Coder.of(),
+ Long.parseLong(mainOutputId),
+ ImmutableMap.of(
+ Long.parseLong(mainOutputId), TestDoFn.mainOutput,
+ Long.parseLong(additionalOutputId), TestDoFn.additionalOutput));
+ RunnerApi.FunctionSpec functionSpec = RunnerApi.FunctionSpec.newBuilder()
+ .setUrn(JAVA_DO_FN_URN)
+ .setParameter(Any.pack(BytesValue.newBuilder()
+ .setValue(ByteString.copyFrom(SerializableUtils.serializeToByteArray(doFnInfo)))
+ .build()))
+ .build();
+ RunnerApi.PTransform pTransform = RunnerApi.PTransform.newBuilder()
+ .setSpec(functionSpec)
+ .putInputs("inputA", "inputATarget")
+ .putInputs("inputB", "inputBTarget")
+ .putOutputs(mainOutputId, "mainOutputTarget")
+ .putOutputs(additionalOutputId, "additionalOutputTarget")
+ .build();
+
+ List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
+ List<WindowedValue<String>> additionalOutputValues = new ArrayList<>();
+ Multimap<String, ThrowingConsumer<WindowedValue<?>>> consumers = HashMultimap.create();
+ consumers.put("mainOutputTarget",
+ (ThrowingConsumer) (ThrowingConsumer<WindowedValue<String>>) mainOutputValues::add);
+ consumers.put("additionalOutputTarget",
+ (ThrowingConsumer) (ThrowingConsumer<WindowedValue<String>>) additionalOutputValues::add);
+ List<ThrowingRunnable> startFunctions = new ArrayList<>();
+ List<ThrowingRunnable> finishFunctions = new ArrayList<>();
+
+ ProcessBundleHandler handler = new ProcessBundleHandler(
+ PipelineOptionsFactory.create(),
+ fnApiRegistry::get,
+ beamFnDataClient);
+ handler.createRunnerForPTransform(
+ pTransformId,
+ pTransform,
+ Suppliers.ofInstance("57L")::get,
+ ImmutableMap.of(),
+ consumers,
+ startFunctions::add,
+ finishFunctions::add);
+
+ Iterables.getOnlyElement(startFunctions).run();
+ mainOutputValues.clear();
+
+ assertThat(consumers.keySet(), containsInAnyOrder(
+ "inputATarget", "inputBTarget", "mainOutputTarget", "additionalOutputTarget"));
+
+ Iterables.getOnlyElement(consumers.get("inputATarget")).accept(valueInGlobalWindow("A1"));
+ Iterables.getOnlyElement(consumers.get("inputATarget")).accept(valueInGlobalWindow("A2"));
+ Iterables.getOnlyElement(consumers.get("inputATarget")).accept(valueInGlobalWindow("B"));
+ assertThat(mainOutputValues, contains(
+ valueInGlobalWindow("MainOutputA1"),
+ valueInGlobalWindow("MainOutputA2"),
+ valueInGlobalWindow("MainOutputB")));
+ assertThat(additionalOutputValues, contains(
+ valueInGlobalWindow("AdditionalOutputA1"),
+ valueInGlobalWindow("AdditionalOutputA2"),
+ valueInGlobalWindow("AdditionalOutputB")));
+ mainOutputValues.clear();
+ additionalOutputValues.clear();
+
+ Iterables.getOnlyElement(finishFunctions).run();
+ assertThat(
+ mainOutputValues,
+ contains(
+ timestampedValueInGlobalWindow("FinishBundle", GlobalWindow.INSTANCE.maxTimestamp())));
+ mainOutputValues.clear();
+ }
+
+ @Test
+ public void testCreatingAndProcessingSource() throws Exception {
+ Map<String, Message> fnApiRegistry = ImmutableMap.of(LONG_CODER_SPEC_ID, LONG_CODER_SPEC);
+ List<WindowedValue<String>> outputValues = new ArrayList<>();
+
+ Multimap<String, ThrowingConsumer<WindowedValue<?>>> consumers = HashMultimap.create();
+ consumers.put("outputPC",
+ (ThrowingConsumer) (ThrowingConsumer<WindowedValue<String>>) outputValues::add);
+ List<ThrowingRunnable> startFunctions = new ArrayList<>();
+ List<ThrowingRunnable> finishFunctions = new ArrayList<>();
+
+ RunnerApi.FunctionSpec functionSpec = RunnerApi.FunctionSpec.newBuilder()
+ .setUrn(JAVA_SOURCE_URN)
+ .setParameter(Any.pack(BytesValue.newBuilder()
+ .setValue(ByteString.copyFrom(
+ SerializableUtils.serializeToByteArray(CountingSource.upTo(3))))
+ .build()))
+ .build();
+
+ RunnerApi.PTransform pTransform = RunnerApi.PTransform.newBuilder()
+ .setSpec(functionSpec)
+ .putInputs("input", "inputPC")
+ .putOutputs("output", "outputPC")
+ .build();
+
+ ProcessBundleHandler handler = new ProcessBundleHandler(
+ PipelineOptionsFactory.create(),
+ fnApiRegistry::get,
+ beamFnDataClient);
+
+ handler.createRunnerForPTransform(
+ "pTransformId",
+ pTransform,
+ Suppliers.ofInstance("57L")::get,
+ ImmutableMap.of(),
+ consumers,
+ startFunctions::add,
+ finishFunctions::add);
+
+ // This is testing a deprecated way of running sources and should be removed
+ // once all source definitions are instead propagated along the input edge.
+ Iterables.getOnlyElement(startFunctions).run();
+ assertThat(outputValues, contains(
+ valueInGlobalWindow(0L),
+ valueInGlobalWindow(1L),
+ valueInGlobalWindow(2L)));
+ outputValues.clear();
+
+ // Check that when passing a source along as an input, the source is processed.
+ assertThat(consumers.keySet(), containsInAnyOrder("inputPC", "outputPC"));
+ Iterables.getOnlyElement(consumers.get("inputPC")).accept(
+ valueInGlobalWindow(CountingSource.upTo(2)));
+ assertThat(outputValues, contains(
+ valueInGlobalWindow(0L),
+ valueInGlobalWindow(1L)));
+
+ assertThat(finishFunctions, empty());
+ }
+
+ @Test
+ public void testCreatingAndProcessingBeamFnDataReadRunner() throws Exception {
+ Map<String, Message> fnApiRegistry = ImmutableMap.of(STRING_CODER_SPEC_ID, STRING_CODER_SPEC);
+ String bundleId = "57";
+ String outputId = "101";
+
+ List<WindowedValue<String>> outputValues = new ArrayList<>();
+
+ Multimap<String, ThrowingConsumer<WindowedValue<?>>> consumers = HashMultimap.create();
+ consumers.put("outputPC",
+ (ThrowingConsumer) (ThrowingConsumer<WindowedValue<String>>) outputValues::add);
+ List<ThrowingRunnable> startFunctions = new ArrayList<>();
+ List<ThrowingRunnable> finishFunctions = new ArrayList<>();
+
+ RunnerApi.FunctionSpec functionSpec = RunnerApi.FunctionSpec.newBuilder()
+ .setUrn(DATA_INPUT_URN)
+ .setParameter(Any.pack(REMOTE_PORT))
+ .build();
+
+ RunnerApi.PTransform pTransform = RunnerApi.PTransform.newBuilder()
+ .setSpec(functionSpec)
+ .putOutputs(outputId, "outputPC")
+ .build();
+
+ ProcessBundleHandler handler = new ProcessBundleHandler(
+ PipelineOptionsFactory.create(),
+ fnApiRegistry::get,
+ beamFnDataClient);
+
+ handler.createRunnerForPTransform(
+ "pTransformId",
+ pTransform,
+ Suppliers.ofInstance(bundleId)::get,
+ ImmutableMap.of("outputPC",
+ RunnerApi.PCollection.newBuilder().setCoderId(STRING_CODER_SPEC_ID).build()),
+ consumers,
+ startFunctions::add,
+ finishFunctions::add);
+
+ verifyZeroInteractions(beamFnDataClient);
+
+ CompletableFuture<Void> completionFuture = new CompletableFuture<>();
+ when(beamFnDataClient.forInboundConsumer(any(), any(), any(), any()))
+ .thenReturn(completionFuture);
+ Iterables.getOnlyElement(startFunctions).run();
+ verify(beamFnDataClient).forInboundConsumer(
+ eq(REMOTE_PORT.getApiServiceDescriptor()),
+ eq(KV.of(bundleId, BeamFnApi.Target.newBuilder()
+ .setPrimitiveTransformReference("pTransformId")
+ .setName(outputId)
+ .build())),
+ eq(STRING_CODER),
+ consumerCaptor.capture());
+
+ consumerCaptor.getValue().accept(valueInGlobalWindow("TestValue"));
+ assertThat(outputValues, contains(valueInGlobalWindow("TestValue")));
+ outputValues.clear();
+
+ assertThat(consumers.keySet(), containsInAnyOrder("outputPC"));
+
+ completionFuture.complete(null);
+ Iterables.getOnlyElement(finishFunctions).run();
+
+ verifyNoMoreInteractions(beamFnDataClient);
+ }
+
+ @Test
+ public void testCreatingAndProcessingBeamFnDataWriteRunner() throws Exception {
+ Map<String, Message> fnApiRegistry = ImmutableMap.of(STRING_CODER_SPEC_ID, STRING_CODER_SPEC);
+ String bundleId = "57L";
+ String inputId = "100L";
+
+ Multimap<String, ThrowingConsumer<WindowedValue<?>>> consumers = HashMultimap.create();
+ List<ThrowingRunnable> startFunctions = new ArrayList<>();
+ List<ThrowingRunnable> finishFunctions = new ArrayList<>();
+
+ RunnerApi.FunctionSpec functionSpec = RunnerApi.FunctionSpec.newBuilder()
+ .setUrn(DATA_OUTPUT_URN)
+ .setParameter(Any.pack(REMOTE_PORT))
+ .build();
+
+ RunnerApi.PTransform pTransform = RunnerApi.PTransform.newBuilder()
+ .setSpec(functionSpec)
+ .putInputs(inputId, "inputPC")
+ .build();
+
+ ProcessBundleHandler handler = new ProcessBundleHandler(
+ PipelineOptionsFactory.create(),
+ fnApiRegistry::get,
+ beamFnDataClient);
+
+ handler.createRunnerForPTransform(
+ "ptransformId",
+ pTransform,
+ Suppliers.ofInstance(bundleId)::get,
+ ImmutableMap.of("inputPC",
+ RunnerApi.PCollection.newBuilder().setCoderId(STRING_CODER_SPEC_ID).build()),
+ consumers,
+ startFunctions::add,
+ finishFunctions::add);
+
+ verifyZeroInteractions(beamFnDataClient);
+
+ List<WindowedValue<String>> outputValues = new ArrayList<>();
+ AtomicBoolean wasCloseCalled = new AtomicBoolean();
+ CloseableThrowingConsumer<WindowedValue<String>> outputConsumer =
+ new CloseableThrowingConsumer<WindowedValue<String>>(){
+ @Override
+ public void close() throws Exception {
+ wasCloseCalled.set(true);
+ }
+
+ @Override
+ public void accept(WindowedValue<String> t) throws Exception {
+ outputValues.add(t);
+ }
+ };
+
+ when(beamFnDataClient.forOutboundConsumer(
+ any(),
+ any(),
+ Matchers.<Coder<WindowedValue<String>>>any())).thenReturn(outputConsumer);
+ Iterables.getOnlyElement(startFunctions).run();
+ verify(beamFnDataClient).forOutboundConsumer(
+ eq(REMOTE_PORT.getApiServiceDescriptor()),
+ eq(KV.of(bundleId, BeamFnApi.Target.newBuilder()
+ .setPrimitiveTransformReference("ptransformId")
+ .setName(inputId)
+ .build())),
+ eq(STRING_CODER));
+
+ assertThat(consumers.keySet(), containsInAnyOrder("inputPC"));
+ Iterables.getOnlyElement(consumers.get("inputPC")).accept(valueInGlobalWindow("TestValue"));
+ assertThat(outputValues, contains(valueInGlobalWindow("TestValue")));
+ outputValues.clear();
+
+ assertFalse(wasCloseCalled.get());
+ Iterables.getOnlyElement(finishFunctions).run();
+ assertTrue(wasCloseCalled.get());
+
+ verifyNoMoreInteractions(beamFnDataClient);
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/RegisterHandlerTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/RegisterHandlerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/RegisterHandlerTest.java
index 2b275af..b1f4410 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/RegisterHandlerTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/control/RegisterHandlerTest.java
@@ -44,14 +44,14 @@ public class RegisterHandlerTest {
.setRegister(BeamFnApi.RegisterRequest.newBuilder()
.addProcessBundleDescriptor(BeamFnApi.ProcessBundleDescriptor.newBuilder()
.setId("1L")
- .putCoders("10L", RunnerApi.Coder.newBuilder()
+ .putCodersyyy("10L", RunnerApi.Coder.newBuilder()
.setSpec(RunnerApi.SdkFunctionSpec.newBuilder()
.setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn("urn:10L").build())
.build())
.build())
.build())
.addProcessBundleDescriptor(BeamFnApi.ProcessBundleDescriptor.newBuilder().setId("2L")
- .putCoders("20L", RunnerApi.Coder.newBuilder()
+ .putCodersyyy("20L", RunnerApi.Coder.newBuilder()
.setSpec(RunnerApi.SdkFunctionSpec.newBuilder()
.setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn("urn:20L").build())
.build())
@@ -82,10 +82,10 @@ public class RegisterHandlerTest {
assertEquals(REGISTER_REQUEST.getRegister().getProcessBundleDescriptor(1),
handler.getById("2L"));
assertEquals(
- REGISTER_REQUEST.getRegister().getProcessBundleDescriptor(0).getCodersOrThrow("10L"),
+ REGISTER_REQUEST.getRegister().getProcessBundleDescriptor(0).getCodersyyyOrThrow("10L"),
handler.getById("10L"));
assertEquals(
- REGISTER_REQUEST.getRegister().getProcessBundleDescriptor(1).getCodersOrThrow("20L"),
+ REGISTER_REQUEST.getRegister().getProcessBundleDescriptor(1).getCodersyyyOrThrow("20L"),
handler.getById("20L"));
assertEquals(REGISTER_RESPONSE, responseFuture.get());
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BeamFnDataReadRunnerTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BeamFnDataReadRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BeamFnDataReadRunnerTest.java
index d6a476e..7e8ab1a 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BeamFnDataReadRunnerTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BeamFnDataReadRunnerTest.java
@@ -20,51 +20,41 @@ package org.apache.beam.runners.core;
import static org.apache.beam.sdk.util.WindowedValue.valueInGlobalWindow;
import static org.hamcrest.Matchers.contains;
-import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.junit.Assert.assertThat;
-import static org.junit.Assert.fail;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoMoreInteractions;
-import static org.mockito.Mockito.verifyZeroInteractions;
import static org.mockito.Mockito.when;
import com.fasterxml.jackson.databind.ObjectMapper;
-import com.google.common.base.Suppliers;
-import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Multimap;
import com.google.common.util.concurrent.Uninterruptibles;
import com.google.protobuf.Any;
import com.google.protobuf.ByteString;
import com.google.protobuf.BytesValue;
import java.io.IOException;
import java.util.ArrayList;
+import java.util.Collection;
import java.util.List;
-import java.util.ServiceLoader;
+import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
import org.apache.beam.fn.harness.fn.ThrowingConsumer;
-import org.apache.beam.fn.harness.fn.ThrowingRunnable;
import org.apache.beam.fn.harness.test.TestExecutors;
import org.apache.beam.fn.harness.test.TestExecutors.TestExecutorService;
import org.apache.beam.fn.v1.BeamFnApi;
-import org.apache.beam.runners.core.PTransformRunnerFactory.Registrar;
import org.apache.beam.runners.dataflow.util.CloudObjects;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.common.runner.v1.RunnerApi;
-import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.KV;
-import org.hamcrest.collection.IsMapContaining;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
@@ -78,18 +68,15 @@ import org.mockito.MockitoAnnotations;
/** Tests for {@link BeamFnDataReadRunner}. */
@RunWith(JUnit4.class)
public class BeamFnDataReadRunnerTest {
-
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
private static final BeamFnApi.RemoteGrpcPort PORT_SPEC = BeamFnApi.RemoteGrpcPort.newBuilder()
.setApiServiceDescriptor(BeamFnApi.ApiServiceDescriptor.getDefaultInstance()).build();
private static final RunnerApi.FunctionSpec FUNCTION_SPEC = RunnerApi.FunctionSpec.newBuilder()
.setParameter(Any.pack(PORT_SPEC)).build();
private static final Coder<WindowedValue<String>> CODER =
WindowedValue.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE);
- private static final String CODER_SPEC_ID = "string-coder-id";
private static final RunnerApi.Coder CODER_SPEC;
- private static final String URN = "urn:org.apache.beam:source:runner:0.1";
-
static {
try {
CODER_SPEC = RunnerApi.Coder.newBuilder().setSpec(
@@ -111,7 +98,7 @@ public class BeamFnDataReadRunnerTest {
.build();
@Rule public TestExecutorService executor = TestExecutors.from(Executors::newCachedThreadPool);
- @Mock private BeamFnDataClient mockBeamFnDataClient;
+ @Mock private BeamFnDataClient mockBeamFnDataClientFactory;
@Captor private ArgumentCaptor<ThrowingConsumer<WindowedValue<String>>> consumerCaptor;
@Before
@@ -120,93 +107,32 @@ public class BeamFnDataReadRunnerTest {
}
@Test
- public void testCreatingAndProcessingBeamFnDataReadRunner() throws Exception {
- String bundleId = "57";
- String outputId = "101";
-
- List<WindowedValue<String>> outputValues = new ArrayList<>();
-
- Multimap<String, ThrowingConsumer<WindowedValue<?>>> consumers = HashMultimap.create();
- consumers.put("outputPC",
- (ThrowingConsumer) (ThrowingConsumer<WindowedValue<String>>) outputValues::add);
- List<ThrowingRunnable> startFunctions = new ArrayList<>();
- List<ThrowingRunnable> finishFunctions = new ArrayList<>();
-
- RunnerApi.FunctionSpec functionSpec = RunnerApi.FunctionSpec.newBuilder()
- .setUrn("urn:org.apache.beam:source:runner:0.1")
- .setParameter(Any.pack(PORT_SPEC))
- .build();
-
- RunnerApi.PTransform pTransform = RunnerApi.PTransform.newBuilder()
- .setSpec(functionSpec)
- .putOutputs(outputId, "outputPC")
- .build();
-
- new BeamFnDataReadRunner.Factory<String>().createRunnerForPTransform(
- PipelineOptionsFactory.create(),
- mockBeamFnDataClient,
- "pTransformId",
- pTransform,
- Suppliers.ofInstance(bundleId)::get,
- ImmutableMap.of("outputPC",
- RunnerApi.PCollection.newBuilder().setCoderId(CODER_SPEC_ID).build()),
- ImmutableMap.of(CODER_SPEC_ID, CODER_SPEC),
- consumers,
- startFunctions::add,
- finishFunctions::add);
-
- verifyZeroInteractions(mockBeamFnDataClient);
-
- CompletableFuture<Void> completionFuture = new CompletableFuture<>();
- when(mockBeamFnDataClient.forInboundConsumer(any(), any(), any(), any()))
- .thenReturn(completionFuture);
- Iterables.getOnlyElement(startFunctions).run();
- verify(mockBeamFnDataClient).forInboundConsumer(
- eq(PORT_SPEC.getApiServiceDescriptor()),
- eq(KV.of(bundleId, BeamFnApi.Target.newBuilder()
- .setPrimitiveTransformReference("pTransformId")
- .setName(outputId)
- .build())),
- eq(CODER),
- consumerCaptor.capture());
-
- consumerCaptor.getValue().accept(valueInGlobalWindow("TestValue"));
- assertThat(outputValues, contains(valueInGlobalWindow("TestValue")));
- outputValues.clear();
-
- assertThat(consumers.keySet(), containsInAnyOrder("outputPC"));
-
- completionFuture.complete(null);
- Iterables.getOnlyElement(finishFunctions).run();
-
- verifyNoMoreInteractions(mockBeamFnDataClient);
- }
-
- @Test
public void testReuseForMultipleBundles() throws Exception {
CompletableFuture<Void> bundle1Future = new CompletableFuture<>();
CompletableFuture<Void> bundle2Future = new CompletableFuture<>();
- when(mockBeamFnDataClient.forInboundConsumer(
+ when(mockBeamFnDataClientFactory.forInboundConsumer(
any(),
any(),
any(),
any())).thenReturn(bundle1Future).thenReturn(bundle2Future);
List<WindowedValue<String>> valuesA = new ArrayList<>();
List<WindowedValue<String>> valuesB = new ArrayList<>();
-
+ Map<String, Collection<ThrowingConsumer<WindowedValue<String>>>> outputMap = ImmutableMap.of(
+ "outA", ImmutableList.of(valuesA::add),
+ "outB", ImmutableList.of(valuesB::add));
AtomicReference<String> bundleId = new AtomicReference<>("0");
BeamFnDataReadRunner<String> readRunner = new BeamFnDataReadRunner<>(
FUNCTION_SPEC,
bundleId::get,
INPUT_TARGET,
CODER_SPEC,
- mockBeamFnDataClient,
- ImmutableList.of(valuesA::add, valuesB::add));
+ mockBeamFnDataClientFactory,
+ outputMap);
// Process for bundle id 0
readRunner.registerInputLocation();
- verify(mockBeamFnDataClient).forInboundConsumer(
+ verify(mockBeamFnDataClientFactory).forInboundConsumer(
eq(PORT_SPEC.getApiServiceDescriptor()),
eq(KV.of(bundleId.get(), INPUT_TARGET)),
eq(CODER),
@@ -238,7 +164,7 @@ public class BeamFnDataReadRunnerTest {
valuesB.clear();
readRunner.registerInputLocation();
- verify(mockBeamFnDataClient).forInboundConsumer(
+ verify(mockBeamFnDataClientFactory).forInboundConsumer(
eq(PORT_SPEC.getApiServiceDescriptor()),
eq(KV.of(bundleId.get(), INPUT_TARGET)),
eq(CODER),
@@ -264,18 +190,6 @@ public class BeamFnDataReadRunnerTest {
assertThat(valuesA, contains(valueInGlobalWindow("GHI"), valueInGlobalWindow("JKL")));
assertThat(valuesB, contains(valueInGlobalWindow("GHI"), valueInGlobalWindow("JKL")));
- verifyNoMoreInteractions(mockBeamFnDataClient);
- }
-
- @Test
- public void testRegistration() {
- for (Registrar registrar :
- ServiceLoader.load(Registrar.class)) {
- if (registrar instanceof BeamFnDataReadRunner.Registrar) {
- assertThat(registrar.getPTransformRunnerFactories(), IsMapContaining.hasKey(URN));
- return;
- }
- }
- fail("Expected registrar not found.");
+ verifyNoMoreInteractions(mockBeamFnDataClientFactory);
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BeamFnDataWriteRunnerTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BeamFnDataWriteRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BeamFnDataWriteRunnerTest.java
index 64d9ea6..a3c874e 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BeamFnDataWriteRunnerTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BeamFnDataWriteRunnerTest.java
@@ -20,48 +20,31 @@ package org.apache.beam.runners.core;
import static org.apache.beam.sdk.util.WindowedValue.valueInGlobalWindow;
import static org.hamcrest.Matchers.contains;
-import static org.hamcrest.Matchers.containsInAnyOrder;
-import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoMoreInteractions;
-import static org.mockito.Mockito.verifyZeroInteractions;
import static org.mockito.Mockito.when;
import com.fasterxml.jackson.databind.ObjectMapper;
-import com.google.common.base.Suppliers;
-import com.google.common.collect.HashMultimap;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Multimap;
import com.google.protobuf.Any;
import com.google.protobuf.ByteString;
import com.google.protobuf.BytesValue;
import java.io.IOException;
import java.util.ArrayList;
-import java.util.List;
-import java.util.ServiceLoader;
-import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
import org.apache.beam.fn.harness.fn.CloseableThrowingConsumer;
-import org.apache.beam.fn.harness.fn.ThrowingConsumer;
-import org.apache.beam.fn.harness.fn.ThrowingRunnable;
import org.apache.beam.fn.v1.BeamFnApi;
-import org.apache.beam.runners.core.PTransformRunnerFactory.Registrar;
import org.apache.beam.runners.dataflow.util.CloudObjects;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.common.runner.v1.RunnerApi;
-import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.KV;
-import org.hamcrest.collection.IsMapContaining;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -73,18 +56,15 @@ import org.mockito.MockitoAnnotations;
/** Tests for {@link BeamFnDataWriteRunner}. */
@RunWith(JUnit4.class)
public class BeamFnDataWriteRunnerTest {
-
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
private static final BeamFnApi.RemoteGrpcPort PORT_SPEC = BeamFnApi.RemoteGrpcPort.newBuilder()
.setApiServiceDescriptor(BeamFnApi.ApiServiceDescriptor.getDefaultInstance()).build();
private static final RunnerApi.FunctionSpec FUNCTION_SPEC = RunnerApi.FunctionSpec.newBuilder()
.setParameter(Any.pack(PORT_SPEC)).build();
- private static final String CODER_ID = "string-coder-id";
private static final Coder<WindowedValue<String>> CODER =
WindowedValue.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE);
private static final RunnerApi.Coder CODER_SPEC;
- private static final String URN = "urn:org.apache.beam:sink:runner:0.1";
-
static {
try {
CODER_SPEC = RunnerApi.Coder.newBuilder().setSpec(
@@ -105,93 +85,18 @@ public class BeamFnDataWriteRunnerTest {
.setName("out")
.build();
- @Mock private BeamFnDataClient mockBeamFnDataClient;
+ @Mock private BeamFnDataClient mockBeamFnDataClientFactory;
@Before
public void setUp() {
MockitoAnnotations.initMocks(this);
}
-
- @Test
- public void testCreatingAndProcessingBeamFnDataWriteRunner() throws Exception {
- String bundleId = "57L";
- String inputId = "100L";
-
- Multimap<String, ThrowingConsumer<WindowedValue<?>>> consumers = HashMultimap.create();
- List<ThrowingRunnable> startFunctions = new ArrayList<>();
- List<ThrowingRunnable> finishFunctions = new ArrayList<>();
-
- RunnerApi.FunctionSpec functionSpec = RunnerApi.FunctionSpec.newBuilder()
- .setUrn("urn:org.apache.beam:sink:runner:0.1")
- .setParameter(Any.pack(PORT_SPEC))
- .build();
-
- RunnerApi.PTransform pTransform = RunnerApi.PTransform.newBuilder()
- .setSpec(functionSpec)
- .putInputs(inputId, "inputPC")
- .build();
-
- new BeamFnDataWriteRunner.Factory<String>().createRunnerForPTransform(
- PipelineOptionsFactory.create(),
- mockBeamFnDataClient,
- "ptransformId",
- pTransform,
- Suppliers.ofInstance(bundleId)::get,
- ImmutableMap.of("inputPC",
- RunnerApi.PCollection.newBuilder().setCoderId(CODER_ID).build()),
- ImmutableMap.of(CODER_ID, CODER_SPEC),
- consumers,
- startFunctions::add,
- finishFunctions::add);
-
- verifyZeroInteractions(mockBeamFnDataClient);
-
- List<WindowedValue<String>> outputValues = new ArrayList<>();
- AtomicBoolean wasCloseCalled = new AtomicBoolean();
- CloseableThrowingConsumer<WindowedValue<String>> outputConsumer =
- new CloseableThrowingConsumer<WindowedValue<String>>(){
- @Override
- public void close() throws Exception {
- wasCloseCalled.set(true);
- }
-
- @Override
- public void accept(WindowedValue<String> t) throws Exception {
- outputValues.add(t);
- }
- };
-
- when(mockBeamFnDataClient.forOutboundConsumer(
- any(),
- any(),
- Matchers.<Coder<WindowedValue<String>>>any())).thenReturn(outputConsumer);
- Iterables.getOnlyElement(startFunctions).run();
- verify(mockBeamFnDataClient).forOutboundConsumer(
- eq(PORT_SPEC.getApiServiceDescriptor()),
- eq(KV.of(bundleId, BeamFnApi.Target.newBuilder()
- .setPrimitiveTransformReference("ptransformId")
- .setName(inputId)
- .build())),
- eq(CODER));
-
- assertThat(consumers.keySet(), containsInAnyOrder("inputPC"));
- Iterables.getOnlyElement(consumers.get("inputPC")).accept(valueInGlobalWindow("TestValue"));
- assertThat(outputValues, contains(valueInGlobalWindow("TestValue")));
- outputValues.clear();
-
- assertFalse(wasCloseCalled.get());
- Iterables.getOnlyElement(finishFunctions).run();
- assertTrue(wasCloseCalled.get());
-
- verifyNoMoreInteractions(mockBeamFnDataClient);
- }
-
@Test
public void testReuseForMultipleBundles() throws Exception {
RecordingConsumer<WindowedValue<String>> valuesA = new RecordingConsumer<>();
RecordingConsumer<WindowedValue<String>> valuesB = new RecordingConsumer<>();
- when(mockBeamFnDataClient.forOutboundConsumer(
+ when(mockBeamFnDataClientFactory.forOutboundConsumer(
any(),
any(),
Matchers.<Coder<WindowedValue<String>>>any())).thenReturn(valuesA).thenReturn(valuesB);
@@ -201,12 +106,12 @@ public class BeamFnDataWriteRunnerTest {
bundleId::get,
OUTPUT_TARGET,
CODER_SPEC,
- mockBeamFnDataClient);
+ mockBeamFnDataClientFactory);
// Process for bundle id 0
writeRunner.registerForOutput();
- verify(mockBeamFnDataClient).forOutboundConsumer(
+ verify(mockBeamFnDataClientFactory).forOutboundConsumer(
eq(PORT_SPEC.getApiServiceDescriptor()),
eq(KV.of(bundleId.get(), OUTPUT_TARGET)),
eq(CODER));
@@ -224,7 +129,7 @@ public class BeamFnDataWriteRunnerTest {
valuesB.clear();
writeRunner.registerForOutput();
- verify(mockBeamFnDataClient).forOutboundConsumer(
+ verify(mockBeamFnDataClientFactory).forOutboundConsumer(
eq(PORT_SPEC.getApiServiceDescriptor()),
eq(KV.of(bundleId.get(), OUTPUT_TARGET)),
eq(CODER));
@@ -235,7 +140,7 @@ public class BeamFnDataWriteRunnerTest {
assertTrue(valuesB.closed);
assertThat(valuesB, contains(valueInGlobalWindow("GHI"), valueInGlobalWindow("JKL")));
- verifyNoMoreInteractions(mockBeamFnDataClient);
+ verifyNoMoreInteractions(mockBeamFnDataClientFactory);
}
private static class RecordingConsumer<T> extends ArrayList<T>
@@ -253,17 +158,6 @@ public class BeamFnDataWriteRunnerTest {
}
add(t);
}
- }
- @Test
- public void testRegistration() {
- for (Registrar registrar :
- ServiceLoader.load(Registrar.class)) {
- if (registrar instanceof BeamFnDataWriteRunner.Registrar) {
- assertThat(registrar.getPTransformRunnerFactories(), IsMapContaining.hasKey(URN));
- return;
- }
- }
- fail("Expected registrar not found.");
}
}
[05/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/ShardRecordsIteratorTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/ShardRecordsIteratorTest.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/ShardRecordsIteratorTest.java
index 4b2190f..49e806d 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/ShardRecordsIteratorTest.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/ShardRecordsIteratorTest.java
@@ -25,10 +25,8 @@ import static org.mockito.Matchers.anyListOf;
import static org.mockito.Mockito.when;
import com.amazonaws.services.kinesis.model.ExpiredIteratorException;
-
import java.io.IOException;
import java.util.Collections;
-
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -42,114 +40,112 @@ import org.mockito.stubbing.Answer;
*/
@RunWith(MockitoJUnitRunner.class)
public class ShardRecordsIteratorTest {
+ private static final String INITIAL_ITERATOR = "INITIAL_ITERATOR";
+ private static final String SECOND_ITERATOR = "SECOND_ITERATOR";
+ private static final String SECOND_REFRESHED_ITERATOR = "SECOND_REFRESHED_ITERATOR";
+ private static final String THIRD_ITERATOR = "THIRD_ITERATOR";
+ private static final String STREAM_NAME = "STREAM_NAME";
+ private static final String SHARD_ID = "SHARD_ID";
+
+ @Mock
+ private SimplifiedKinesisClient kinesisClient;
+ @Mock
+ private ShardCheckpoint firstCheckpoint, aCheckpoint, bCheckpoint, cCheckpoint, dCheckpoint;
+ @Mock
+ private GetKinesisRecordsResult firstResult, secondResult, thirdResult;
+ @Mock
+ private KinesisRecord a, b, c, d;
+ @Mock
+ private RecordFilter recordFilter;
+
+ private ShardRecordsIterator iterator;
+
+ @Before
+ public void setUp() throws IOException, TransientKinesisException {
+ when(firstCheckpoint.getShardIterator(kinesisClient)).thenReturn(INITIAL_ITERATOR);
+ when(firstCheckpoint.getStreamName()).thenReturn(STREAM_NAME);
+ when(firstCheckpoint.getShardId()).thenReturn(SHARD_ID);
+
+ when(firstCheckpoint.moveAfter(a)).thenReturn(aCheckpoint);
+ when(aCheckpoint.moveAfter(b)).thenReturn(bCheckpoint);
+ when(aCheckpoint.getStreamName()).thenReturn(STREAM_NAME);
+ when(aCheckpoint.getShardId()).thenReturn(SHARD_ID);
+ when(bCheckpoint.moveAfter(c)).thenReturn(cCheckpoint);
+ when(bCheckpoint.getStreamName()).thenReturn(STREAM_NAME);
+ when(bCheckpoint.getShardId()).thenReturn(SHARD_ID);
+ when(cCheckpoint.moveAfter(d)).thenReturn(dCheckpoint);
+ when(cCheckpoint.getStreamName()).thenReturn(STREAM_NAME);
+ when(cCheckpoint.getShardId()).thenReturn(SHARD_ID);
+ when(dCheckpoint.getStreamName()).thenReturn(STREAM_NAME);
+ when(dCheckpoint.getShardId()).thenReturn(SHARD_ID);
+
+ when(kinesisClient.getRecords(INITIAL_ITERATOR, STREAM_NAME, SHARD_ID))
+ .thenReturn(firstResult);
+ when(kinesisClient.getRecords(SECOND_ITERATOR, STREAM_NAME, SHARD_ID))
+ .thenReturn(secondResult);
+ when(kinesisClient.getRecords(THIRD_ITERATOR, STREAM_NAME, SHARD_ID))
+ .thenReturn(thirdResult);
+
+ when(firstResult.getNextShardIterator()).thenReturn(SECOND_ITERATOR);
+ when(secondResult.getNextShardIterator()).thenReturn(THIRD_ITERATOR);
+ when(thirdResult.getNextShardIterator()).thenReturn(THIRD_ITERATOR);
+
+ when(firstResult.getRecords()).thenReturn(Collections.<KinesisRecord>emptyList());
+ when(secondResult.getRecords()).thenReturn(Collections.<KinesisRecord>emptyList());
+ when(thirdResult.getRecords()).thenReturn(Collections.<KinesisRecord>emptyList());
+
+ when(recordFilter.apply(anyListOf(KinesisRecord.class), any(ShardCheckpoint
+ .class))).thenAnswer(new IdentityAnswer());
+
+ iterator = new ShardRecordsIterator(firstCheckpoint, kinesisClient, recordFilter);
+ }
+
+ @Test
+ public void returnsAbsentIfNoRecordsPresent() throws IOException, TransientKinesisException {
+ assertThat(iterator.next()).isEqualTo(CustomOptional.absent());
+ assertThat(iterator.next()).isEqualTo(CustomOptional.absent());
+ assertThat(iterator.next()).isEqualTo(CustomOptional.absent());
+ }
+
+ @Test
+ public void goesThroughAvailableRecords() throws IOException, TransientKinesisException {
+ when(firstResult.getRecords()).thenReturn(asList(a, b, c));
+ when(secondResult.getRecords()).thenReturn(singletonList(d));
+
+ assertThat(iterator.getCheckpoint()).isEqualTo(firstCheckpoint);
+ assertThat(iterator.next()).isEqualTo(CustomOptional.of(a));
+ assertThat(iterator.getCheckpoint()).isEqualTo(aCheckpoint);
+ assertThat(iterator.next()).isEqualTo(CustomOptional.of(b));
+ assertThat(iterator.getCheckpoint()).isEqualTo(bCheckpoint);
+ assertThat(iterator.next()).isEqualTo(CustomOptional.of(c));
+ assertThat(iterator.getCheckpoint()).isEqualTo(cCheckpoint);
+ assertThat(iterator.next()).isEqualTo(CustomOptional.of(d));
+ assertThat(iterator.getCheckpoint()).isEqualTo(dCheckpoint);
+ assertThat(iterator.next()).isEqualTo(CustomOptional.absent());
+ assertThat(iterator.getCheckpoint()).isEqualTo(dCheckpoint);
+ }
+
+ @Test
+ public void refreshesExpiredIterator() throws IOException, TransientKinesisException {
+ when(firstResult.getRecords()).thenReturn(singletonList(a));
+ when(secondResult.getRecords()).thenReturn(singletonList(b));
+
+ when(kinesisClient.getRecords(SECOND_ITERATOR, STREAM_NAME, SHARD_ID))
+ .thenThrow(ExpiredIteratorException.class);
+ when(aCheckpoint.getShardIterator(kinesisClient))
+ .thenReturn(SECOND_REFRESHED_ITERATOR);
+ when(kinesisClient.getRecords(SECOND_REFRESHED_ITERATOR, STREAM_NAME, SHARD_ID))
+ .thenReturn(secondResult);
+
+ assertThat(iterator.next()).isEqualTo(CustomOptional.of(a));
+ assertThat(iterator.next()).isEqualTo(CustomOptional.of(b));
+ assertThat(iterator.next()).isEqualTo(CustomOptional.absent());
+ }
- private static final String INITIAL_ITERATOR = "INITIAL_ITERATOR";
- private static final String SECOND_ITERATOR = "SECOND_ITERATOR";
- private static final String SECOND_REFRESHED_ITERATOR = "SECOND_REFRESHED_ITERATOR";
- private static final String THIRD_ITERATOR = "THIRD_ITERATOR";
- private static final String STREAM_NAME = "STREAM_NAME";
- private static final String SHARD_ID = "SHARD_ID";
-
- @Mock
- private SimplifiedKinesisClient kinesisClient;
- @Mock
- private ShardCheckpoint firstCheckpoint, aCheckpoint, bCheckpoint, cCheckpoint, dCheckpoint;
- @Mock
- private GetKinesisRecordsResult firstResult, secondResult, thirdResult;
- @Mock
- private KinesisRecord a, b, c, d;
- @Mock
- private RecordFilter recordFilter;
-
- private ShardRecordsIterator iterator;
-
- @Before
- public void setUp() throws IOException, TransientKinesisException {
- when(firstCheckpoint.getShardIterator(kinesisClient)).thenReturn(INITIAL_ITERATOR);
- when(firstCheckpoint.getStreamName()).thenReturn(STREAM_NAME);
- when(firstCheckpoint.getShardId()).thenReturn(SHARD_ID);
-
- when(firstCheckpoint.moveAfter(a)).thenReturn(aCheckpoint);
- when(aCheckpoint.moveAfter(b)).thenReturn(bCheckpoint);
- when(aCheckpoint.getStreamName()).thenReturn(STREAM_NAME);
- when(aCheckpoint.getShardId()).thenReturn(SHARD_ID);
- when(bCheckpoint.moveAfter(c)).thenReturn(cCheckpoint);
- when(bCheckpoint.getStreamName()).thenReturn(STREAM_NAME);
- when(bCheckpoint.getShardId()).thenReturn(SHARD_ID);
- when(cCheckpoint.moveAfter(d)).thenReturn(dCheckpoint);
- when(cCheckpoint.getStreamName()).thenReturn(STREAM_NAME);
- when(cCheckpoint.getShardId()).thenReturn(SHARD_ID);
- when(dCheckpoint.getStreamName()).thenReturn(STREAM_NAME);
- when(dCheckpoint.getShardId()).thenReturn(SHARD_ID);
-
- when(kinesisClient.getRecords(INITIAL_ITERATOR, STREAM_NAME, SHARD_ID))
- .thenReturn(firstResult);
- when(kinesisClient.getRecords(SECOND_ITERATOR, STREAM_NAME, SHARD_ID))
- .thenReturn(secondResult);
- when(kinesisClient.getRecords(THIRD_ITERATOR, STREAM_NAME, SHARD_ID))
- .thenReturn(thirdResult);
-
- when(firstResult.getNextShardIterator()).thenReturn(SECOND_ITERATOR);
- when(secondResult.getNextShardIterator()).thenReturn(THIRD_ITERATOR);
- when(thirdResult.getNextShardIterator()).thenReturn(THIRD_ITERATOR);
-
- when(firstResult.getRecords()).thenReturn(Collections.<KinesisRecord>emptyList());
- when(secondResult.getRecords()).thenReturn(Collections.<KinesisRecord>emptyList());
- when(thirdResult.getRecords()).thenReturn(Collections.<KinesisRecord>emptyList());
-
- when(recordFilter.apply(anyListOf(KinesisRecord.class), any(ShardCheckpoint
- .class))).thenAnswer(new IdentityAnswer());
-
- iterator = new ShardRecordsIterator(firstCheckpoint, kinesisClient, recordFilter);
- }
-
- @Test
- public void returnsAbsentIfNoRecordsPresent() throws IOException, TransientKinesisException {
- assertThat(iterator.next()).isEqualTo(CustomOptional.absent());
- assertThat(iterator.next()).isEqualTo(CustomOptional.absent());
- assertThat(iterator.next()).isEqualTo(CustomOptional.absent());
- }
-
- @Test
- public void goesThroughAvailableRecords() throws IOException, TransientKinesisException {
- when(firstResult.getRecords()).thenReturn(asList(a, b, c));
- when(secondResult.getRecords()).thenReturn(singletonList(d));
-
- assertThat(iterator.getCheckpoint()).isEqualTo(firstCheckpoint);
- assertThat(iterator.next()).isEqualTo(CustomOptional.of(a));
- assertThat(iterator.getCheckpoint()).isEqualTo(aCheckpoint);
- assertThat(iterator.next()).isEqualTo(CustomOptional.of(b));
- assertThat(iterator.getCheckpoint()).isEqualTo(bCheckpoint);
- assertThat(iterator.next()).isEqualTo(CustomOptional.of(c));
- assertThat(iterator.getCheckpoint()).isEqualTo(cCheckpoint);
- assertThat(iterator.next()).isEqualTo(CustomOptional.of(d));
- assertThat(iterator.getCheckpoint()).isEqualTo(dCheckpoint);
- assertThat(iterator.next()).isEqualTo(CustomOptional.absent());
- assertThat(iterator.getCheckpoint()).isEqualTo(dCheckpoint);
- }
-
- @Test
- public void refreshesExpiredIterator() throws IOException, TransientKinesisException {
- when(firstResult.getRecords()).thenReturn(singletonList(a));
- when(secondResult.getRecords()).thenReturn(singletonList(b));
-
- when(kinesisClient.getRecords(SECOND_ITERATOR, STREAM_NAME, SHARD_ID))
- .thenThrow(ExpiredIteratorException.class);
- when(aCheckpoint.getShardIterator(kinesisClient))
- .thenReturn(SECOND_REFRESHED_ITERATOR);
- when(kinesisClient.getRecords(SECOND_REFRESHED_ITERATOR, STREAM_NAME, SHARD_ID))
- .thenReturn(secondResult);
-
- assertThat(iterator.next()).isEqualTo(CustomOptional.of(a));
- assertThat(iterator.next()).isEqualTo(CustomOptional.of(b));
- assertThat(iterator.next()).isEqualTo(CustomOptional.absent());
- }
-
- private static class IdentityAnswer implements Answer<Object> {
-
- @Override
- public Object answer(InvocationOnMock invocation) throws Throwable {
- return invocation.getArguments()[0];
+ private static class IdentityAnswer implements Answer<Object> {
+ @Override
+ public Object answer(InvocationOnMock invocation) throws Throwable {
+ return invocation.getArguments()[0];
+ }
}
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/SimplifiedKinesisClientTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/SimplifiedKinesisClientTest.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/SimplifiedKinesisClientTest.java
index 2f8757c..96434fd 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/SimplifiedKinesisClientTest.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/SimplifiedKinesisClientTest.java
@@ -34,9 +34,7 @@ import com.amazonaws.services.kinesis.model.ProvisionedThroughputExceededExcepti
import com.amazonaws.services.kinesis.model.Shard;
import com.amazonaws.services.kinesis.model.ShardIteratorType;
import com.amazonaws.services.kinesis.model.StreamDescription;
-
import java.util.List;
-
import org.joda.time.Instant;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -48,180 +46,179 @@ import org.mockito.runners.MockitoJUnitRunner;
*/
@RunWith(MockitoJUnitRunner.class)
public class SimplifiedKinesisClientTest {
+ private static final String STREAM = "stream";
+ private static final String SHARD_1 = "shard-01";
+ private static final String SHARD_2 = "shard-02";
+ private static final String SHARD_3 = "shard-03";
+ private static final String SHARD_ITERATOR = "iterator";
+ private static final String SEQUENCE_NUMBER = "abc123";
+
+ @Mock
+ private AmazonKinesis kinesis;
+ @InjectMocks
+ private SimplifiedKinesisClient underTest;
+
+ @Test
+ public void shouldReturnIteratorStartingWithSequenceNumber() throws Exception {
+ given(kinesis.getShardIterator(new GetShardIteratorRequest()
+ .withStreamName(STREAM)
+ .withShardId(SHARD_1)
+ .withShardIteratorType(ShardIteratorType.AT_SEQUENCE_NUMBER)
+ .withStartingSequenceNumber(SEQUENCE_NUMBER)
+ )).willReturn(new GetShardIteratorResult()
+ .withShardIterator(SHARD_ITERATOR));
+
+ String stream = underTest.getShardIterator(STREAM, SHARD_1,
+ ShardIteratorType.AT_SEQUENCE_NUMBER, SEQUENCE_NUMBER, null);
+
+ assertThat(stream).isEqualTo(SHARD_ITERATOR);
+ }
+
+ @Test
+ public void shouldReturnIteratorStartingWithTimestamp() throws Exception {
+ Instant timestamp = Instant.now();
+ given(kinesis.getShardIterator(new GetShardIteratorRequest()
+ .withStreamName(STREAM)
+ .withShardId(SHARD_1)
+ .withShardIteratorType(ShardIteratorType.AT_SEQUENCE_NUMBER)
+ .withTimestamp(timestamp.toDate())
+ )).willReturn(new GetShardIteratorResult()
+ .withShardIterator(SHARD_ITERATOR));
+
+ String stream = underTest.getShardIterator(STREAM, SHARD_1,
+ ShardIteratorType.AT_SEQUENCE_NUMBER, null, timestamp);
+
+ assertThat(stream).isEqualTo(SHARD_ITERATOR);
+ }
+
+ @Test
+ public void shouldHandleExpiredIterationExceptionForGetShardIterator() {
+ shouldHandleGetShardIteratorError(new ExpiredIteratorException(""),
+ ExpiredIteratorException.class);
+ }
+
+ @Test
+ public void shouldHandleLimitExceededExceptionForGetShardIterator() {
+ shouldHandleGetShardIteratorError(new LimitExceededException(""),
+ TransientKinesisException.class);
+ }
+
+ @Test
+ public void shouldHandleProvisionedThroughputExceededExceptionForGetShardIterator() {
+ shouldHandleGetShardIteratorError(new ProvisionedThroughputExceededException(""),
+ TransientKinesisException.class);
+ }
+
+ @Test
+ public void shouldHandleServiceErrorForGetShardIterator() {
+ shouldHandleGetShardIteratorError(newAmazonServiceException(ErrorType.Service),
+ TransientKinesisException.class);
+ }
+
+ @Test
+ public void shouldHandleClientErrorForGetShardIterator() {
+ shouldHandleGetShardIteratorError(newAmazonServiceException(ErrorType.Client),
+ RuntimeException.class);
+ }
+
+ @Test
+ public void shouldHandleUnexpectedExceptionForGetShardIterator() {
+ shouldHandleGetShardIteratorError(new NullPointerException(),
+ RuntimeException.class);
+ }
+
+ private void shouldHandleGetShardIteratorError(
+ Exception thrownException,
+ Class<? extends Exception> expectedExceptionClass) {
+ GetShardIteratorRequest request = new GetShardIteratorRequest()
+ .withStreamName(STREAM)
+ .withShardId(SHARD_1)
+ .withShardIteratorType(ShardIteratorType.LATEST);
+
+ given(kinesis.getShardIterator(request)).willThrow(thrownException);
+
+ try {
+ underTest.getShardIterator(STREAM, SHARD_1, ShardIteratorType.LATEST, null, null);
+ failBecauseExceptionWasNotThrown(expectedExceptionClass);
+ } catch (Exception e) {
+ assertThat(e).isExactlyInstanceOf(expectedExceptionClass);
+ } finally {
+ reset(kinesis);
+ }
+ }
+
+ @Test
+ public void shouldListAllShards() throws Exception {
+ Shard shard1 = new Shard().withShardId(SHARD_1);
+ Shard shard2 = new Shard().withShardId(SHARD_2);
+ Shard shard3 = new Shard().withShardId(SHARD_3);
+ given(kinesis.describeStream(STREAM, null)).willReturn(new DescribeStreamResult()
+ .withStreamDescription(new StreamDescription()
+ .withShards(shard1, shard2)
+ .withHasMoreShards(true)));
+ given(kinesis.describeStream(STREAM, SHARD_2)).willReturn(new DescribeStreamResult()
+ .withStreamDescription(new StreamDescription()
+ .withShards(shard3)
+ .withHasMoreShards(false)));
+
+ List<Shard> shards = underTest.listShards(STREAM);
+
+ assertThat(shards).containsOnly(shard1, shard2, shard3);
+ }
+
+ @Test
+ public void shouldHandleExpiredIterationExceptionForShardListing() {
+ shouldHandleShardListingError(new ExpiredIteratorException(""),
+ ExpiredIteratorException.class);
+ }
+
+ @Test
+ public void shouldHandleLimitExceededExceptionForShardListing() {
+ shouldHandleShardListingError(new LimitExceededException(""),
+ TransientKinesisException.class);
+ }
+
+ @Test
+ public void shouldHandleProvisionedThroughputExceededExceptionForShardListing() {
+ shouldHandleShardListingError(new ProvisionedThroughputExceededException(""),
+ TransientKinesisException.class);
+ }
- private static final String STREAM = "stream";
- private static final String SHARD_1 = "shard-01";
- private static final String SHARD_2 = "shard-02";
- private static final String SHARD_3 = "shard-03";
- private static final String SHARD_ITERATOR = "iterator";
- private static final String SEQUENCE_NUMBER = "abc123";
-
- @Mock
- private AmazonKinesis kinesis;
- @InjectMocks
- private SimplifiedKinesisClient underTest;
-
- @Test
- public void shouldReturnIteratorStartingWithSequenceNumber() throws Exception {
- given(kinesis.getShardIterator(new GetShardIteratorRequest()
- .withStreamName(STREAM)
- .withShardId(SHARD_1)
- .withShardIteratorType(ShardIteratorType.AT_SEQUENCE_NUMBER)
- .withStartingSequenceNumber(SEQUENCE_NUMBER)
- )).willReturn(new GetShardIteratorResult()
- .withShardIterator(SHARD_ITERATOR));
-
- String stream = underTest.getShardIterator(STREAM, SHARD_1,
- ShardIteratorType.AT_SEQUENCE_NUMBER, SEQUENCE_NUMBER, null);
-
- assertThat(stream).isEqualTo(SHARD_ITERATOR);
- }
-
- @Test
- public void shouldReturnIteratorStartingWithTimestamp() throws Exception {
- Instant timestamp = Instant.now();
- given(kinesis.getShardIterator(new GetShardIteratorRequest()
- .withStreamName(STREAM)
- .withShardId(SHARD_1)
- .withShardIteratorType(ShardIteratorType.AT_SEQUENCE_NUMBER)
- .withTimestamp(timestamp.toDate())
- )).willReturn(new GetShardIteratorResult()
- .withShardIterator(SHARD_ITERATOR));
-
- String stream = underTest.getShardIterator(STREAM, SHARD_1,
- ShardIteratorType.AT_SEQUENCE_NUMBER, null, timestamp);
-
- assertThat(stream).isEqualTo(SHARD_ITERATOR);
- }
-
- @Test
- public void shouldHandleExpiredIterationExceptionForGetShardIterator() {
- shouldHandleGetShardIteratorError(new ExpiredIteratorException(""),
- ExpiredIteratorException.class);
- }
-
- @Test
- public void shouldHandleLimitExceededExceptionForGetShardIterator() {
- shouldHandleGetShardIteratorError(new LimitExceededException(""),
- TransientKinesisException.class);
- }
-
- @Test
- public void shouldHandleProvisionedThroughputExceededExceptionForGetShardIterator() {
- shouldHandleGetShardIteratorError(new ProvisionedThroughputExceededException(""),
- TransientKinesisException.class);
- }
-
- @Test
- public void shouldHandleServiceErrorForGetShardIterator() {
- shouldHandleGetShardIteratorError(newAmazonServiceException(ErrorType.Service),
- TransientKinesisException.class);
- }
-
- @Test
- public void shouldHandleClientErrorForGetShardIterator() {
- shouldHandleGetShardIteratorError(newAmazonServiceException(ErrorType.Client),
- RuntimeException.class);
- }
-
- @Test
- public void shouldHandleUnexpectedExceptionForGetShardIterator() {
- shouldHandleGetShardIteratorError(new NullPointerException(),
- RuntimeException.class);
- }
-
- private void shouldHandleGetShardIteratorError(
- Exception thrownException,
- Class<? extends Exception> expectedExceptionClass) {
- GetShardIteratorRequest request = new GetShardIteratorRequest()
- .withStreamName(STREAM)
- .withShardId(SHARD_1)
- .withShardIteratorType(ShardIteratorType.LATEST);
-
- given(kinesis.getShardIterator(request)).willThrow(thrownException);
-
- try {
- underTest.getShardIterator(STREAM, SHARD_1, ShardIteratorType.LATEST, null, null);
- failBecauseExceptionWasNotThrown(expectedExceptionClass);
- } catch (Exception e) {
- assertThat(e).isExactlyInstanceOf(expectedExceptionClass);
- } finally {
- reset(kinesis);
- }
- }
-
- @Test
- public void shouldListAllShards() throws Exception {
- Shard shard1 = new Shard().withShardId(SHARD_1);
- Shard shard2 = new Shard().withShardId(SHARD_2);
- Shard shard3 = new Shard().withShardId(SHARD_3);
- given(kinesis.describeStream(STREAM, null)).willReturn(new DescribeStreamResult()
- .withStreamDescription(new StreamDescription()
- .withShards(shard1, shard2)
- .withHasMoreShards(true)));
- given(kinesis.describeStream(STREAM, SHARD_2)).willReturn(new DescribeStreamResult()
- .withStreamDescription(new StreamDescription()
- .withShards(shard3)
- .withHasMoreShards(false)));
-
- List<Shard> shards = underTest.listShards(STREAM);
-
- assertThat(shards).containsOnly(shard1, shard2, shard3);
- }
-
- @Test
- public void shouldHandleExpiredIterationExceptionForShardListing() {
- shouldHandleShardListingError(new ExpiredIteratorException(""),
- ExpiredIteratorException.class);
- }
-
- @Test
- public void shouldHandleLimitExceededExceptionForShardListing() {
- shouldHandleShardListingError(new LimitExceededException(""),
- TransientKinesisException.class);
- }
-
- @Test
- public void shouldHandleProvisionedThroughputExceededExceptionForShardListing() {
- shouldHandleShardListingError(new ProvisionedThroughputExceededException(""),
- TransientKinesisException.class);
- }
-
- @Test
- public void shouldHandleServiceErrorForShardListing() {
- shouldHandleShardListingError(newAmazonServiceException(ErrorType.Service),
- TransientKinesisException.class);
- }
-
- @Test
- public void shouldHandleClientErrorForShardListing() {
- shouldHandleShardListingError(newAmazonServiceException(ErrorType.Client),
- RuntimeException.class);
- }
-
- @Test
- public void shouldHandleUnexpectedExceptionForShardListing() {
- shouldHandleShardListingError(new NullPointerException(),
- RuntimeException.class);
- }
-
- private void shouldHandleShardListingError(
- Exception thrownException,
- Class<? extends Exception> expectedExceptionClass) {
- given(kinesis.describeStream(STREAM, null)).willThrow(thrownException);
- try {
- underTest.listShards(STREAM);
- failBecauseExceptionWasNotThrown(expectedExceptionClass);
- } catch (Exception e) {
- assertThat(e).isExactlyInstanceOf(expectedExceptionClass);
- } finally {
- reset(kinesis);
- }
- }
-
- private AmazonServiceException newAmazonServiceException(ErrorType errorType) {
- AmazonServiceException exception = new AmazonServiceException("");
- exception.setErrorType(errorType);
- return exception;
- }
+ @Test
+ public void shouldHandleServiceErrorForShardListing() {
+ shouldHandleShardListingError(newAmazonServiceException(ErrorType.Service),
+ TransientKinesisException.class);
+ }
+
+ @Test
+ public void shouldHandleClientErrorForShardListing() {
+ shouldHandleShardListingError(newAmazonServiceException(ErrorType.Client),
+ RuntimeException.class);
+ }
+
+ @Test
+ public void shouldHandleUnexpectedExceptionForShardListing() {
+ shouldHandleShardListingError(new NullPointerException(),
+ RuntimeException.class);
+ }
+
+ private void shouldHandleShardListingError(
+ Exception thrownException,
+ Class<? extends Exception> expectedExceptionClass) {
+ given(kinesis.describeStream(STREAM, null)).willThrow(thrownException);
+ try {
+ underTest.listShards(STREAM);
+ failBecauseExceptionWasNotThrown(expectedExceptionClass);
+ } catch (Exception e) {
+ assertThat(e).isExactlyInstanceOf(expectedExceptionClass);
+ } finally {
+ reset(kinesis);
+ }
+ }
+
+ private AmazonServiceException newAmazonServiceException(ErrorType errorType) {
+ AmazonServiceException exception = new AmazonServiceException("");
+ exception.setErrorType(errorType);
+ return exception;
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/mongodb/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/mongodb/pom.xml b/sdks/java/io/mongodb/pom.xml
index d93cc41..912e20c 100644
--- a/sdks/java/io/mongodb/pom.xml
+++ b/sdks/java/io/mongodb/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbGridFSIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbGridFSIO.java b/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbGridFSIO.java
index 5b5412c..b63775d 100644
--- a/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbGridFSIO.java
+++ b/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbGridFSIO.java
@@ -117,7 +117,7 @@ import org.joda.time.Instant;
* to the file separated with line feeds.
* </p>
*/
-@Experimental(Experimental.Kind.SOURCE_SINK)
+@Experimental
public class MongoDbGridFSIO {
/**
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java b/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java
index 3b14182..620df74 100644
--- a/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java
+++ b/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java
@@ -18,13 +18,12 @@
package org.apache.beam.sdk.io.mongodb;
import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.base.Preconditions.checkNotNull;
import com.google.auto.value.AutoValue;
import com.google.common.annotations.VisibleForTesting;
import com.mongodb.BasicDBObject;
import com.mongodb.MongoClient;
-import com.mongodb.MongoClientOptions;
import com.mongodb.MongoClientURI;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoCursor;
@@ -94,27 +93,19 @@ import org.slf4j.LoggerFactory;
*
* }</pre>
*/
-@Experimental(Experimental.Kind.SOURCE_SINK)
+@Experimental
public class MongoDbIO {
private static final Logger LOG = LoggerFactory.getLogger(MongoDbIO.class);
/** Read data from MongoDB. */
public static Read read() {
- return new AutoValue_MongoDbIO_Read.Builder()
- .setKeepAlive(true)
- .setMaxConnectionIdleTime(60000)
- .setNumSplits(0)
- .build();
+ return new AutoValue_MongoDbIO_Read.Builder().setNumSplits(0).build();
}
/** Write data to MongoDB. */
public static Write write() {
- return new AutoValue_MongoDbIO_Write.Builder()
- .setKeepAlive(true)
- .setMaxConnectionIdleTime(60000)
- .setBatchSize(1024L)
- .build();
+ return new AutoValue_MongoDbIO_Write.Builder().setBatchSize(1024L).build();
}
private MongoDbIO() {
@@ -126,20 +117,16 @@ public class MongoDbIO {
@AutoValue
public abstract static class Read extends PTransform<PBegin, PCollection<Document>> {
@Nullable abstract String uri();
- abstract boolean keepAlive();
- abstract int maxConnectionIdleTime();
@Nullable abstract String database();
@Nullable abstract String collection();
@Nullable abstract String filter();
abstract int numSplits();
- abstract Builder builder();
+ abstract Builder toBuilder();
@AutoValue.Builder
abstract static class Builder {
abstract Builder setUri(String uri);
- abstract Builder setKeepAlive(boolean keepAlive);
- abstract Builder setMaxConnectionIdleTime(int maxConnectionIdleTime);
abstract Builder setDatabase(String database);
abstract Builder setCollection(String collection);
abstract Builder setFilter(String filter);
@@ -148,94 +135,31 @@ public class MongoDbIO {
}
/**
- * Define the location of the MongoDB instances using an URI. The URI describes the hosts to
- * be used and some options.
- *
- * <p>The format of the URI is:
- *
- * <pre>{@code
- * mongodb://[username:password@]host1[:port1]...[,hostN[:portN]]][/[database][?options]]
- * }</pre>
- *
- * <p>Where:
- * <ul>
- * <li>{@code mongodb://} is a required prefix to identify that this is a string in the
- * standard connection format.</li>
- * <li>{@code username:password@} are optional. If given, the driver will attempt to
- * login to a database after connecting to a database server. For some authentication
- * mechanisms, only the username is specified and the password is not, in which case
- * the ":" after the username is left off as well.</li>
- * <li>{@code host1} is the only required part of the URI. It identifies a server
- * address to connect to.</li>
- * <li>{@code :portX} is optional and defaults to {@code :27017} if not provided.</li>
- * <li>{@code /database} is the name of the database to login to and thus is only
- * relevant if the {@code username:password@} syntax is used. If not specified, the
- * "admin" database will be used by default. It has to be equivalent with the database
- * you specific with {@link Read#withDatabase(String)}.</li>
- * <li>{@code ?options} are connection options. Note that if {@code database} is absent
- * there is still a {@code /} required between the last {@code host} and the {@code ?}
- * introducing the options. Options are name=value pairs and the pairs are separated by
- * "{@code &}". The {@code KeepAlive} connection option can't be passed via the URI,
- * instead you have to use {@link Read#withKeepAlive(boolean)}. Same for the
- * {@code MaxConnectionIdleTime} connection option via
- * {@link Read#withMaxConnectionIdleTime(int)}.
- * </li>
- * </ul>
+ * Example documentation for withUri.
*/
public Read withUri(String uri) {
- checkArgument(uri != null, "MongoDbIO.read().withUri(uri) called with null uri");
- return builder().setUri(uri).build();
- }
-
- /**
- * Sets whether socket keep alive is enabled.
- */
- public Read withKeepAlive(boolean keepAlive) {
- return builder().setKeepAlive(keepAlive).build();
- }
-
- /**
- * Sets the maximum idle time for a pooled connection.
- */
- public Read withMaxConnectionIdleTime(int maxConnectionIdleTime) {
- return builder().setMaxConnectionIdleTime(maxConnectionIdleTime).build();
+ checkNotNull(uri);
+ return toBuilder().setUri(uri).build();
}
- /**
- * Sets the database to use.
- */
public Read withDatabase(String database) {
- checkArgument(database != null, "MongoDbIO.read().withDatabase(database) called with null"
- + " database");
- return builder().setDatabase(database).build();
+ checkNotNull(database);
+ return toBuilder().setDatabase(database).build();
}
- /**
- * Sets the collection to consider in the database.
- */
public Read withCollection(String collection) {
- checkArgument(collection != null, "MongoDbIO.read().withCollection(collection) called "
- + "with null collection");
- return builder().setCollection(collection).build();
+ checkNotNull(collection);
+ return toBuilder().setCollection(collection).build();
}
- /**
- * Sets a filter on the documents in a collection.
- */
public Read withFilter(String filter) {
- checkArgument(filter != null, "MongoDbIO.read().withFilter(filter) called with null "
- + "filter");
- return builder().setFilter(filter).build();
+ checkNotNull(filter);
+ return toBuilder().setFilter(filter).build();
}
- /**
- * Sets the user defined number of splits.
- */
public Read withNumSplits(int numSplits) {
- checkArgument(numSplits >= 0, "MongoDbIO.read().withNumSplits(numSplits) called with "
- + "invalid number. The number of splits has to be a positive value (currently %d)",
- numSplits);
- return builder().setNumSplits(numSplits).build();
+ checkArgument(numSplits >= 0);
+ return toBuilder().setNumSplits(numSplits).build();
}
@Override
@@ -245,19 +169,15 @@ public class MongoDbIO {
@Override
public void validate(PipelineOptions options) {
- checkState(uri() != null, "MongoDbIO.read() requires an URI to be set via withUri(uri)");
- checkState(database() != null, "MongoDbIO.read() requires a database to be set via "
- + "withDatabase(database)");
- checkState(collection() != null, "MongoDbIO.read() requires a collection to be set via "
- + "withCollection(collection)");
+ checkNotNull(uri(), "uri");
+ checkNotNull(database(), "database");
+ checkNotNull(collection(), "collection");
}
@Override
public void populateDisplayData(DisplayData.Builder builder) {
super.populateDisplayData(builder);
builder.add(DisplayData.item("uri", uri()));
- builder.add(DisplayData.item("keepAlive", keepAlive()));
- builder.add(DisplayData.item("maxConnectionIdleTime", maxConnectionIdleTime()));
builder.add(DisplayData.item("database", database()));
builder.add(DisplayData.item("collection", collection()));
builder.addIfNotNull(DisplayData.item("filter", filter()));
@@ -298,71 +218,61 @@ public class MongoDbIO {
@Override
public long getEstimatedSizeBytes(PipelineOptions pipelineOptions) {
- try (MongoClient mongoClient = new MongoClient(new MongoClientURI(spec.uri()))) {
- return getEstimatedSizeBytes(mongoClient, spec.database(), spec.collection());
- }
- }
-
- private long getEstimatedSizeBytes(MongoClient mongoClient,
- String database,
- String collection) {
- MongoDatabase mongoDatabase = mongoClient.getDatabase(database);
+ MongoClient mongoClient = new MongoClient(new MongoClientURI(spec.uri()));
+ MongoDatabase mongoDatabase = mongoClient.getDatabase(spec.database());
// get the Mongo collStats object
// it gives the size for the entire collection
BasicDBObject stat = new BasicDBObject();
- stat.append("collStats", collection);
+ stat.append("collStats", spec.collection());
Document stats = mongoDatabase.runCommand(stat);
-
return stats.get("size", Number.class).longValue();
}
@Override
public List<BoundedSource<Document>> split(long desiredBundleSizeBytes,
PipelineOptions options) {
- try (MongoClient mongoClient = new MongoClient(new MongoClientURI(spec.uri()))) {
- MongoDatabase mongoDatabase = mongoClient.getDatabase(spec.database());
-
- List<Document> splitKeys;
- if (spec.numSplits() > 0) {
- // the user defines his desired number of splits
- // calculate the batch size
- long estimatedSizeBytes = getEstimatedSizeBytes(mongoClient,
- spec.database(), spec.collection());
- desiredBundleSizeBytes = estimatedSizeBytes / spec.numSplits();
- }
-
- // the desired batch size is small, using default chunk size of 1MB
- if (desiredBundleSizeBytes < 1024 * 1024) {
- desiredBundleSizeBytes = 1 * 1024 * 1024;
- }
-
- // now we have the batch size (provided by user or provided by the runner)
- // we use Mongo splitVector command to get the split keys
- BasicDBObject splitVectorCommand = new BasicDBObject();
- splitVectorCommand.append("splitVector", spec.database() + "." + spec.collection());
- splitVectorCommand.append("keyPattern", new BasicDBObject().append("_id", 1));
- splitVectorCommand.append("force", false);
- // maxChunkSize is the Mongo partition size in MB
- LOG.debug("Splitting in chunk of {} MB", desiredBundleSizeBytes / 1024 / 1024);
- splitVectorCommand.append("maxChunkSize", desiredBundleSizeBytes / 1024 / 1024);
- Document splitVectorCommandResult = mongoDatabase.runCommand(splitVectorCommand);
- splitKeys = (List<Document>) splitVectorCommandResult.get("splitKeys");
-
- List<BoundedSource<Document>> sources = new ArrayList<>();
- if (splitKeys.size() < 1) {
- LOG.debug("Split keys is low, using an unique source");
- sources.add(this);
- return sources;
- }
+ MongoClient mongoClient = new MongoClient(new MongoClientURI(spec.uri()));
+ MongoDatabase mongoDatabase = mongoClient.getDatabase(spec.database());
+
+ List<Document> splitKeys;
+ if (spec.numSplits() > 0) {
+ // the user defines his desired number of splits
+ // calculate the batch size
+ long estimatedSizeBytes = getEstimatedSizeBytes(options);
+ desiredBundleSizeBytes = estimatedSizeBytes / spec.numSplits();
+ }
- LOG.debug("Number of splits is {}", splitKeys.size());
- for (String shardFilter : splitKeysToFilters(splitKeys, spec.filter())) {
- sources.add(new BoundedMongoDbSource(spec.withFilter(shardFilter)));
- }
+ // the desired batch size is small, using default chunk size of 1MB
+ if (desiredBundleSizeBytes < 1024 * 1024) {
+ desiredBundleSizeBytes = 1 * 1024 * 1024;
+ }
+ // now we have the batch size (provided by user or provided by the runner)
+ // we use Mongo splitVector command to get the split keys
+ BasicDBObject splitVectorCommand = new BasicDBObject();
+ splitVectorCommand.append("splitVector", spec.database() + "." + spec.collection());
+ splitVectorCommand.append("keyPattern", new BasicDBObject().append("_id", 1));
+ splitVectorCommand.append("force", false);
+ // maxChunkSize is the Mongo partition size in MB
+ LOG.debug("Splitting in chunk of {} MB", desiredBundleSizeBytes / 1024 / 1024);
+ splitVectorCommand.append("maxChunkSize", desiredBundleSizeBytes / 1024 / 1024);
+ Document splitVectorCommandResult = mongoDatabase.runCommand(splitVectorCommand);
+ splitKeys = (List<Document>) splitVectorCommandResult.get("splitKeys");
+
+ List<BoundedSource<Document>> sources = new ArrayList<>();
+ if (splitKeys.size() < 1) {
+ LOG.debug("Split keys is low, using an unique source");
+ sources.add(this);
return sources;
}
+
+ LOG.debug("Number of splits is {}", splitKeys.size());
+ for (String shardFilter : splitKeysToFilters(splitKeys, spec.filter())) {
+ sources.add(new BoundedMongoDbSource(spec.withFilter(shardFilter)));
+ }
+
+ return sources;
}
/**
@@ -457,10 +367,7 @@ public class MongoDbIO {
@Override
public boolean start() {
Read spec = source.spec;
- MongoClientOptions.Builder optionsBuilder = new MongoClientOptions.Builder();
- optionsBuilder.maxConnectionIdleTime(spec.maxConnectionIdleTime());
- optionsBuilder.socketKeepAlive(spec.keepAlive());
- client = new MongoClient(new MongoClientURI(spec.uri(), optionsBuilder));
+ client = new MongoClient(new MongoClientURI(spec.uri()));
MongoDatabase mongoDatabase = client.getDatabase(spec.database());
@@ -519,106 +426,36 @@ public class MongoDbIO {
*/
@AutoValue
public abstract static class Write extends PTransform<PCollection<Document>, PDone> {
-
@Nullable abstract String uri();
- abstract boolean keepAlive();
- abstract int maxConnectionIdleTime();
@Nullable abstract String database();
@Nullable abstract String collection();
abstract long batchSize();
- abstract Builder builder();
+ abstract Builder toBuilder();
@AutoValue.Builder
abstract static class Builder {
abstract Builder setUri(String uri);
- abstract Builder setKeepAlive(boolean keepAlive);
- abstract Builder setMaxConnectionIdleTime(int maxConnectionIdleTime);
abstract Builder setDatabase(String database);
abstract Builder setCollection(String collection);
abstract Builder setBatchSize(long batchSize);
abstract Write build();
}
- /**
- * Define the location of the MongoDB instances using an URI. The URI describes the hosts to
- * be used and some options.
- *
- * <p>The format of the URI is:
- *
- * <pre>{@code
- * mongodb://[username:password@]host1[:port1],...[,hostN[:portN]]][/[database][?options]]
- * }</pre>
- *
- * <p>Where:
- * <ul>
- * <li>{@code mongodb://} is a required prefix to identify that this is a string in the
- * standard connection format.</li>
- * <li>{@code username:password@} are optional. If given, the driver will attempt to
- * login to a database after connecting to a database server. For some authentication
- * mechanisms, only the username is specified and the password is not, in which case
- * the ":" after the username is left off as well.</li>
- * <li>{@code host1} is the only required part of the URI. It identifies a server
- * address to connect to.</li>
- * <li>{@code :portX} is optional and defaults to {@code :27017} if not provided.</li>
- * <li>{@code /database} is the name of the database to login to and thus is only
- * relevant if the {@code username:password@} syntax is used. If not specified, the
- * "admin" database will be used by default. It has to be equivalent with the database
- * you specific with {@link Write#withDatabase(String)}.</li>
- * <li>{@code ?options} are connection options. Note that if {@code database} is absent
- * there is still a {@code /} required between the last {@code host} and the {@code ?}
- * introducing the options. Options are name=value pairs and the pairs are separated by
- * "{@code &}". The {@code KeepAlive} connection option can't be passed via the URI, instead
- * you have to use {@link Write#withKeepAlive(boolean)}. Same for the
- * {@code MaxConnectionIdleTime} connection option via
- * {@link Write#withMaxConnectionIdleTime(int)}.
- * </li>
- * </ul>
- */
public Write withUri(String uri) {
- checkArgument(uri != null, "MongoDbIO.write().withUri(uri) called with null uri");
- return builder().setUri(uri).build();
- }
-
- /**
- * Sets whether socket keep alive is enabled.
- */
- public Write withKeepAlive(boolean keepAlive) {
- return builder().setKeepAlive(keepAlive).build();
- }
-
- /**
- * Sets the maximum idle time for a pooled connection.
- */
- public Write withMaxConnectionIdleTime(int maxConnectionIdleTime) {
- return builder().setMaxConnectionIdleTime(maxConnectionIdleTime).build();
+ return toBuilder().setUri(uri).build();
}
- /**
- * Sets the database to use.
- */
public Write withDatabase(String database) {
- checkArgument(database != null, "MongoDbIO.write().withDatabase(database) called with "
- + "null database");
- return builder().setDatabase(database).build();
+ return toBuilder().setDatabase(database).build();
}
- /**
- * Sets the collection where to write data in the database.
- */
public Write withCollection(String collection) {
- checkArgument(collection != null, "MongoDbIO.write().withCollection(collection) called "
- + "with null collection");
- return builder().setCollection(collection).build();
+ return toBuilder().setCollection(collection).build();
}
- /**
- * Define the size of the batch to group write operations.
- */
public Write withBatchSize(long batchSize) {
- checkArgument(batchSize >= 0, "MongoDbIO.write().withBatchSize(batchSize) called with "
- + "invalid batch size. Batch size has to be >= 0 (currently %d)", batchSize);
- return builder().setBatchSize(batchSize).build();
+ return toBuilder().setBatchSize(batchSize).build();
}
@Override
@@ -629,21 +466,10 @@ public class MongoDbIO {
@Override
public void validate(PipelineOptions options) {
- checkState(uri() != null, "MongoDbIO.write() requires an URI to be set via withUri(uri)");
- checkState(database() != null, "MongoDbIO.write() requires a database to be set via "
- + "withDatabase(database)");
- checkState(collection() != null, "MongoDbIO.write() requires a collection to be set via "
- + "withCollection(collection)");
- }
-
- @Override
- public void populateDisplayData(DisplayData.Builder builder) {
- builder.add(DisplayData.item("uri", uri()));
- builder.add(DisplayData.item("keepAlive", keepAlive()));
- builder.add(DisplayData.item("maxConnectionIdleTime", maxConnectionIdleTime()));
- builder.add(DisplayData.item("database", database()));
- builder.add(DisplayData.item("collection", collection()));
- builder.add(DisplayData.item("batchSize", batchSize()));
+ checkNotNull(uri(), "uri");
+ checkNotNull(database(), "database");
+ checkNotNull(collection(), "collection");
+ checkNotNull(batchSize(), "batchSize");
}
private static class WriteFn extends DoFn<Document, Void> {
@@ -657,10 +483,7 @@ public class MongoDbIO {
@Setup
public void createMongoClient() throws Exception {
- MongoClientOptions.Builder builder = new MongoClientOptions.Builder();
- builder.socketKeepAlive(spec.keepAlive());
- builder.maxConnectionIdleTime(spec.maxConnectionIdleTime());
- client = new MongoClient(new MongoClientURI(spec.uri(), builder));
+ client = new MongoClient(new MongoClientURI(spec.uri()));
}
@StartBundle
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbIOTest.java b/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbIOTest.java
index 67dbca4..cd26b48 100644
--- a/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbIOTest.java
+++ b/sdks/java/io/mongodb/src/test/java/org/apache/beam/sdk/io/mongodb/MongoDbIOTest.java
@@ -18,7 +18,6 @@
package org.apache.beam.sdk.io.mongodb;
import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
import com.mongodb.MongoClient;
import com.mongodb.client.MongoCollection;
@@ -190,42 +189,6 @@ public class MongoDbIOTest implements Serializable {
}
@Test
- public void testReadWithCustomConnectionOptions() throws Exception {
- MongoDbIO.Read read = MongoDbIO.read()
- .withUri("mongodb://localhost:" + port)
- .withKeepAlive(false)
- .withMaxConnectionIdleTime(10)
- .withDatabase(DATABASE)
- .withCollection(COLLECTION);
- assertFalse(read.keepAlive());
- assertEquals(10, read.maxConnectionIdleTime());
-
- PCollection<Document> documents = pipeline.apply(read);
-
- PAssert.thatSingleton(documents.apply("Count All", Count.<Document>globally()))
- .isEqualTo(1000L);
-
- PAssert.that(documents
- .apply("Map Scientist", MapElements.via(new SimpleFunction<Document, KV<String, Void>>() {
- public KV<String, Void> apply(Document input) {
- return KV.of(input.getString("scientist"), null);
- }
- }))
- .apply("Count Scientist", Count.<String, Void>perKey())
- ).satisfies(new SerializableFunction<Iterable<KV<String, Long>>, Void>() {
- @Override
- public Void apply(Iterable<KV<String, Long>> input) {
- for (KV<String, Long> element : input) {
- assertEquals(100L, element.getValue().longValue());
- }
- return null;
- }
- });
-
- pipeline.run();
- }
-
- @Test
public void testReadWithFilter() throws Exception {
PCollection<Document> output = pipeline.apply(
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/mqtt/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/mqtt/pom.xml b/sdks/java/io/mqtt/pom.xml
index 9fa1dc0..baaf771 100644
--- a/sdks/java/io/mqtt/pom.xml
+++ b/sdks/java/io/mqtt/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/mqtt/src/main/java/org/apache/beam/sdk/io/mqtt/MqttIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/mqtt/src/main/java/org/apache/beam/sdk/io/mqtt/MqttIO.java b/sdks/java/io/mqtt/src/main/java/org/apache/beam/sdk/io/mqtt/MqttIO.java
index add5cb5..228a85d 100644
--- a/sdks/java/io/mqtt/src/main/java/org/apache/beam/sdk/io/mqtt/MqttIO.java
+++ b/sdks/java/io/mqtt/src/main/java/org/apache/beam/sdk/io/mqtt/MqttIO.java
@@ -97,7 +97,7 @@ import org.slf4j.LoggerFactory;
*
* }</pre>
*/
-@Experimental(Experimental.Kind.SOURCE_SINK)
+@Experimental
public class MqttIO {
private static final Logger LOG = LoggerFactory.getLogger(MqttIO.class);
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/pom.xml b/sdks/java/io/pom.xml
index b7909fa..44f3baa 100644
--- a/sdks/java/io/pom.xml
+++ b/sdks/java/io/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
@@ -32,8 +32,38 @@
<description>Beam SDK Java IO provides different connectivity components
(sources and sinks) to consume and produce data from systems.</description>
+ <properties>
+ <!--
+ This is the version of Hadoop used to compile the hadoop-common module.
+ This dependency is defined with a provided scope.
+ Users must supply their own Hadoop version at runtime.
+ -->
+ <hadoop.version>2.7.3</hadoop.version>
+ </properties>
+
+ <dependencyManagement>
+ <dependencies>
+ <dependency>
+ <groupId>org.apache.hadoop</groupId>
+ <artifactId>hadoop-client</artifactId>
+ <version>${hadoop.version}</version>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.hadoop</groupId>
+ <artifactId>hadoop-common</artifactId>
+ <version>${hadoop.version}</version>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.hadoop</groupId>
+ <artifactId>hadoop-mapreduce-client-core</artifactId>
+ <version>${hadoop.version}</version>
+ </dependency>
+ </dependencies>
+ </dependencyManagement>
+
<modules>
- <module>amqp</module>
<module>cassandra</module>
<module>common</module>
<module>elasticsearch</module>
@@ -42,7 +72,6 @@
<module>hadoop-file-system</module>
<module>hadoop</module>
<module>hbase</module>
- <module>hcatalog</module>
<module>jdbc</module>
<module>jms</module>
<module>kafka</module>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/xml/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/xml/pom.xml b/sdks/java/io/xml/pom.xml
index 7b5804e..cf7dd33 100644
--- a/sdks/java/io/xml/pom.xml
+++ b/sdks/java/io/xml/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/xml/src/main/java/org/apache/beam/sdk/io/xml/XmlIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/xml/src/main/java/org/apache/beam/sdk/io/xml/XmlIO.java b/sdks/java/io/xml/src/main/java/org/apache/beam/sdk/io/xml/XmlIO.java
index 442fba5..7255a94 100644
--- a/sdks/java/io/xml/src/main/java/org/apache/beam/sdk/io/xml/XmlIO.java
+++ b/sdks/java/io/xml/src/main/java/org/apache/beam/sdk/io/xml/XmlIO.java
@@ -36,7 +36,6 @@ import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.transforms.SerializableFunctions;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
@@ -522,8 +521,7 @@ public class XmlIO {
@Override
public PDone expand(PCollection<T> input) {
- return input.apply(
- org.apache.beam.sdk.io.WriteFiles.to(createSink(), SerializableFunctions.<T>identity()));
+ return input.apply(org.apache.beam.sdk.io.WriteFiles.to(createSink()));
}
@VisibleForTesting
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/xml/src/main/java/org/apache/beam/sdk/io/xml/XmlSink.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/xml/src/main/java/org/apache/beam/sdk/io/xml/XmlSink.java b/sdks/java/io/xml/src/main/java/org/apache/beam/sdk/io/xml/XmlSink.java
index 74e0bda..6ae83f2 100644
--- a/sdks/java/io/xml/src/main/java/org/apache/beam/sdk/io/xml/XmlSink.java
+++ b/sdks/java/io/xml/src/main/java/org/apache/beam/sdk/io/xml/XmlSink.java
@@ -25,7 +25,6 @@ import javax.xml.bind.JAXBContext;
import javax.xml.bind.Marshaller;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.io.DefaultFilenamePolicy;
-import org.apache.beam.sdk.io.DynamicFileDestinations;
import org.apache.beam.sdk.io.FileBasedSink;
import org.apache.beam.sdk.io.ShardNameTemplate;
import org.apache.beam.sdk.io.fs.ResourceId;
@@ -35,18 +34,18 @@ import org.apache.beam.sdk.util.CoderUtils;
import org.apache.beam.sdk.util.MimeTypes;
/** Implementation of {@link XmlIO#write}. */
-class XmlSink<T> extends FileBasedSink<T, Void> {
+class XmlSink<T> extends FileBasedSink<T> {
private static final String XML_EXTENSION = ".xml";
private final XmlIO.Write<T> spec;
- private static <T> DefaultFilenamePolicy makeFilenamePolicy(XmlIO.Write<T> spec) {
- return DefaultFilenamePolicy.fromStandardParameters(
+ private static DefaultFilenamePolicy makeFilenamePolicy(XmlIO.Write<?> spec) {
+ return DefaultFilenamePolicy.constructUsingStandardParameters(
spec.getFilenamePrefix(), ShardNameTemplate.INDEX_OF_MAX, XML_EXTENSION, false);
}
XmlSink(XmlIO.Write<T> spec) {
- super(spec.getFilenamePrefix(), DynamicFileDestinations.constant(makeFilenamePolicy(spec)));
+ super(spec.getFilenamePrefix(), makeFilenamePolicy(spec));
this.spec = spec;
}
@@ -76,8 +75,10 @@ class XmlSink<T> extends FileBasedSink<T, Void> {
super.populateDisplayData(builder);
}
- /** {@link WriteOperation} for XML {@link FileBasedSink}s. */
- protected static final class XmlWriteOperation<T> extends WriteOperation<T, Void> {
+ /**
+ * {@link WriteOperation} for XML {@link FileBasedSink}s.
+ */
+ protected static final class XmlWriteOperation<T> extends WriteOperation<T> {
public XmlWriteOperation(XmlSink<T> sink) {
super(sink);
}
@@ -111,8 +112,10 @@ class XmlSink<T> extends FileBasedSink<T, Void> {
}
}
- /** A {@link Writer} that can write objects as XML elements. */
- protected static final class XmlWriter<T> extends Writer<T, Void> {
+ /**
+ * A {@link Writer} that can write objects as XML elements.
+ */
+ protected static final class XmlWriter<T> extends Writer<T> {
final Marshaller marshaller;
private OutputStream os = null;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/xml/src/test/java/org/apache/beam/sdk/io/xml/XmlSinkTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/xml/src/test/java/org/apache/beam/sdk/io/xml/XmlSinkTest.java b/sdks/java/io/xml/src/test/java/org/apache/beam/sdk/io/xml/XmlSinkTest.java
index d1584dc..aa0c1c3 100644
--- a/sdks/java/io/xml/src/test/java/org/apache/beam/sdk/io/xml/XmlSinkTest.java
+++ b/sdks/java/io/xml/src/test/java/org/apache/beam/sdk/io/xml/XmlSinkTest.java
@@ -197,8 +197,8 @@ public class XmlSinkTest {
.withRecordClass(Integer.class);
DisplayData displayData = DisplayData.from(write);
- assertThat(
- displayData, hasDisplayItem("filenamePattern", "/path/to/file-SSSSS-of-NNNNN" + ".xml"));
+
+ assertThat(displayData, hasDisplayItem("filenamePattern", "file-SSSSS-of-NNNNN.xml"));
assertThat(displayData, hasDisplayItem("rootElement", "bird"));
assertThat(displayData, hasDisplayItem("recordClass", Integer.class));
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/java8tests/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/java8tests/pom.xml b/sdks/java/java8tests/pom.xml
index 2378014..b90a757 100644
--- a/sdks/java/java8tests/pom.xml
+++ b/sdks/java/java8tests/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/javadoc/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/javadoc/pom.xml b/sdks/java/javadoc/pom.xml
index 51109fb..54dae3a 100644
--- a/sdks/java/javadoc/pom.xml
+++ b/sdks/java/javadoc/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../../../pom.xml</relativePath>
</parent>
@@ -99,16 +99,6 @@
<dependency>
<groupId>org.apache.beam</groupId>
- <artifactId>beam-sdks-java-io-amqp</artifactId>
- </dependency>
-
- <dependency>
- <groupId>org.apache.beam</groupId>
- <artifactId>beam-sdks-java-io-cassandra</artifactId>
- </dependency>
-
- <dependency>
- <groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-elasticsearch</artifactId>
</dependency>
@@ -134,11 +124,6 @@
<dependency>
<groupId>org.apache.beam</groupId>
- <artifactId>beam-sdks-java-io-hcatalog</artifactId>
- </dependency>
-
- <dependency>
- <groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-jdbc</artifactId>
</dependency>
@@ -211,11 +196,13 @@
<dependency>
<groupId>org.apache.spark</groupId>
<artifactId>spark-core_2.10</artifactId>
+ <version>${spark.version}</version>
</dependency>
<dependency>
<groupId>org.apache.spark</groupId>
<artifactId>spark-streaming_2.10</artifactId>
+ <version>${spark.version}</version>
</dependency>
</dependencies>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/maven-archetypes/examples-java8/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/maven-archetypes/examples-java8/pom.xml b/sdks/java/maven-archetypes/examples-java8/pom.xml
index b60a695..b57644d 100644
--- a/sdks/java/maven-archetypes/examples-java8/pom.xml
+++ b/sdks/java/maven-archetypes/examples-java8/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-maven-archetypes-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/maven-archetypes/examples-java8/src/main/resources/archetype-resources/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/maven-archetypes/examples-java8/src/main/resources/archetype-resources/pom.xml b/sdks/java/maven-archetypes/examples-java8/src/main/resources/archetype-resources/pom.xml
index 4517861..af4fbd3 100644
--- a/sdks/java/maven-archetypes/examples-java8/src/main/resources/archetype-resources/pom.xml
+++ b/sdks/java/maven-archetypes/examples-java8/src/main/resources/archetype-resources/pom.xml
@@ -242,6 +242,7 @@
<dependency>
<groupId>org.apache.spark</groupId>
<artifactId>spark-streaming_2.10</artifactId>
+ <version>${spark.version}</version>
<scope>runtime</scope>
<exclusions>
<exclusion>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/maven-archetypes/examples/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/maven-archetypes/examples/pom.xml b/sdks/java/maven-archetypes/examples/pom.xml
index 2a02039..c1378cb 100644
--- a/sdks/java/maven-archetypes/examples/pom.xml
+++ b/sdks/java/maven-archetypes/examples/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-maven-archetypes-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/maven-archetypes/examples/src/main/resources/archetype-resources/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/maven-archetypes/examples/src/main/resources/archetype-resources/pom.xml b/sdks/java/maven-archetypes/examples/src/main/resources/archetype-resources/pom.xml
index d039ddb..b8b9c9f 100644
--- a/sdks/java/maven-archetypes/examples/src/main/resources/archetype-resources/pom.xml
+++ b/sdks/java/maven-archetypes/examples/src/main/resources/archetype-resources/pom.xml
@@ -241,6 +241,7 @@
<dependency>
<groupId>org.apache.spark</groupId>
<artifactId>spark-streaming_2.10</artifactId>
+ <version>${spark.version}</version>
<scope>runtime</scope>
<exclusions>
<exclusion>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/maven-archetypes/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/maven-archetypes/pom.xml b/sdks/java/maven-archetypes/pom.xml
index d676b31..b7fe274 100644
--- a/sdks/java/maven-archetypes/pom.xml
+++ b/sdks/java/maven-archetypes/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/maven-archetypes/starter/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/maven-archetypes/starter/pom.xml b/sdks/java/maven-archetypes/starter/pom.xml
index 8024b52..06b41c8 100644
--- a/sdks/java/maven-archetypes/starter/pom.xml
+++ b/sdks/java/maven-archetypes/starter/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-maven-archetypes-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/maven-archetypes/starter/src/test/resources/projects/basic/reference/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/maven-archetypes/starter/src/test/resources/projects/basic/reference/pom.xml b/sdks/java/maven-archetypes/starter/src/test/resources/projects/basic/reference/pom.xml
index 6056fb0..60405e6 100644
--- a/sdks/java/maven-archetypes/starter/src/test/resources/projects/basic/reference/pom.xml
+++ b/sdks/java/maven-archetypes/starter/src/test/resources/projects/basic/reference/pom.xml
@@ -28,7 +28,7 @@
<beam.version>@project.version@</beam.version>
<maven-compiler-plugin.version>3.6.1</maven-compiler-plugin.version>
- <maven-exec-plugin.version>1.6.0</maven-exec-plugin.version>
+ <maven-exec-plugin.version>1.4.0</maven-exec-plugin.version>
<slf4j.version>1.7.14</slf4j.version>
</properties>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/pom.xml b/sdks/java/pom.xml
index 3144193..250c85a 100644
--- a/sdks/java/pom.xml
+++ b/sdks/java/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/pom.xml b/sdks/pom.xml
index aec8762..27b9610 100644
--- a/sdks/pom.xml
+++ b/sdks/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/coders/coder_impl.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/coders/coder_impl.py b/sdks/python/apache_beam/coders/coder_impl.py
index 2670250..10298bf 100644
--- a/sdks/python/apache_beam/coders/coder_impl.py
+++ b/sdks/python/apache_beam/coders/coder_impl.py
@@ -710,10 +710,6 @@ class WindowedValueCoderImpl(StreamCoderImpl):
timestamp = MAX_TIMESTAMP.micros
else:
timestamp *= 1000
- if timestamp > MAX_TIMESTAMP.micros:
- timestamp = MAX_TIMESTAMP.micros
- if timestamp < MIN_TIMESTAMP.micros:
- timestamp = MIN_TIMESTAMP.micros
windows = self._windows_coder.decode_from_stream(in_stream, True)
# Read PaneInfo encoded byte.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/coders/coders.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/coders/coders.py b/sdks/python/apache_beam/coders/coders.py
index c56ef52..f40045d 100644
--- a/sdks/python/apache_beam/coders/coders.py
+++ b/sdks/python/apache_beam/coders/coders.py
@@ -25,7 +25,6 @@ import cPickle as pickle
import google.protobuf
from apache_beam.coders import coder_impl
-from apache_beam.portability.api import beam_runner_api_pb2
from apache_beam.utils import urns
from apache_beam.utils import proto_utils
@@ -206,6 +205,7 @@ class Coder(object):
"""For internal use only; no backwards-compatibility guarantees.
"""
# TODO(BEAM-115): Use specialized URNs and components.
+ from apache_beam.runners.api import beam_runner_api_pb2
return beam_runner_api_pb2.Coder(
spec=beam_runner_api_pb2.SdkFunctionSpec(
spec=beam_runner_api_pb2.FunctionSpec(
@@ -286,11 +286,6 @@ class BytesCoder(FastCoder):
def is_deterministic(self):
return True
- def as_cloud_object(self):
- return {
- '@type': 'kind:bytes',
- }
-
def __eq__(self, other):
return type(self) == type(other)
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/coders/coders_test_common.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/coders/coders_test_common.py b/sdks/python/apache_beam/coders/coders_test_common.py
index 577c53a..c9b67b3 100644
--- a/sdks/python/apache_beam/coders/coders_test_common.py
+++ b/sdks/python/apache_beam/coders/coders_test_common.py
@@ -23,8 +23,6 @@ import unittest
import dill
-from apache_beam.transforms.window import GlobalWindow
-from apache_beam.utils.timestamp import MIN_TIMESTAMP
import observable
from apache_beam.transforms import window
from apache_beam.utils import timestamp
@@ -289,12 +287,6 @@ class CodersTest(unittest.TestCase):
# Test binary representation
self.assertEqual('\x7f\xdf;dZ\x1c\xac\t\x00\x00\x00\x01\x0f\x01',
coder.encode(window.GlobalWindows.windowed_value(1)))
-
- # Test decoding large timestamp
- self.assertEqual(
- coder.decode('\x7f\xdf;dZ\x1c\xac\x08\x00\x00\x00\x01\x0f\x00'),
- windowed_value.create(0, MIN_TIMESTAMP.micros, (GlobalWindow(),)))
-
# Test unnested
self.check_coder(
coders.WindowedValueCoder(coders.VarIntCoder()),
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/examples/snippets/snippets_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/examples/snippets/snippets_test.py b/sdks/python/apache_beam/examples/snippets/snippets_test.py
index 31f71b3..9183d0d 100644
--- a/sdks/python/apache_beam/examples/snippets/snippets_test.py
+++ b/sdks/python/apache_beam/examples/snippets/snippets_test.py
@@ -589,22 +589,6 @@ class SnippetsTest(unittest.TestCase):
snippets.model_textio_compressed(
{'read': gzip_file_name}, ['aa', 'bb', 'cc'])
- def test_model_textio_gzip_concatenated(self):
- temp_path_1 = self.create_temp_file('a\nb\nc\n')
- temp_path_2 = self.create_temp_file('p\nq\nr\n')
- temp_path_3 = self.create_temp_file('x\ny\nz')
- gzip_file_name = temp_path_1 + '.gz'
- with open(temp_path_1) as src, gzip.open(gzip_file_name, 'wb') as dst:
- dst.writelines(src)
- with open(temp_path_2) as src, gzip.open(gzip_file_name, 'ab') as dst:
- dst.writelines(src)
- with open(temp_path_3) as src, gzip.open(gzip_file_name, 'ab') as dst:
- dst.writelines(src)
- # Add the temporary gzip file to be cleaned up as well.
- self.temp_files.append(gzip_file_name)
- snippets.model_textio_compressed(
- {'read': gzip_file_name}, ['a', 'b', 'c', 'p', 'q', 'r', 'x', 'y', 'z'])
-
@unittest.skipIf(datastore_pb2 is None, 'GCP dependencies are not installed')
def test_model_datastoreio(self):
# We cannot test datastoreio functionality in unit tests therefore we limit
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/examples/streaming_wordcount.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/examples/streaming_wordcount.py b/sdks/python/apache_beam/examples/streaming_wordcount.py
index 7696d77..ed8b5d0 100644
--- a/sdks/python/apache_beam/examples/streaming_wordcount.py
+++ b/sdks/python/apache_beam/examples/streaming_wordcount.py
@@ -25,44 +25,35 @@ from __future__ import absolute_import
import argparse
import logging
+import re
import apache_beam as beam
-from apache_beam.options.pipeline_options import PipelineOptions
-from apache_beam.options.pipeline_options import StandardOptions
import apache_beam.transforms.window as window
-def split_fn(lines):
- import re
- return re.findall(r'[A-Za-z\']+', lines)
-
-
def run(argv=None):
"""Build and run the pipeline."""
+
parser = argparse.ArgumentParser()
parser.add_argument(
'--input_topic', required=True,
- help=('Input PubSub topic of the form '
- '"projects/<PROJECT>/topics/<TOPIC>".'))
+ help='Input PubSub topic of the form "/topics/<PROJECT>/<TOPIC>".')
parser.add_argument(
'--output_topic', required=True,
- help=('Output PubSub topic of the form '
- '"projects/<PROJECT>/topic/<TOPIC>".'))
+ help='Output PubSub topic of the form "/topics/<PROJECT>/<TOPIC>".')
known_args, pipeline_args = parser.parse_known_args(argv)
- options = PipelineOptions(pipeline_args)
- options.view_as(StandardOptions).streaming = True
- with beam.Pipeline(options=options) as p:
+ with beam.Pipeline(argv=pipeline_args) as p:
- # Read from PubSub into a PCollection.
+ # Read the text file[pattern] into a PCollection.
lines = p | beam.io.ReadStringsFromPubSub(known_args.input_topic)
# Capitalize the characters in each line.
transformed = (lines
- # Use a pre-defined function that imports the re package.
| 'Split' >> (
- beam.FlatMap(split_fn).with_output_types(unicode))
+ beam.FlatMap(lambda x: re.findall(r'[A-Za-z\']+', x))
+ .with_output_types(unicode))
| 'PairWithOne' >> beam.Map(lambda x: (x, 1))
| beam.WindowInto(window.FixedWindows(15, 0))
| 'Group' >> beam.GroupByKey()
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/examples/windowed_wordcount.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/examples/windowed_wordcount.py b/sdks/python/apache_beam/examples/windowed_wordcount.py
deleted file mode 100644
index bd57847..0000000
--- a/sdks/python/apache_beam/examples/windowed_wordcount.py
+++ /dev/null
@@ -1,93 +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.
-#
-
-"""A streaming word-counting workflow.
-
-Important: streaming pipeline support in Python Dataflow is in development
-and is not yet available for use.
-"""
-
-from __future__ import absolute_import
-
-import argparse
-import logging
-
-
-import apache_beam as beam
-import apache_beam.transforms.window as window
-
-TABLE_SCHEMA = ('word:STRING, count:INTEGER, '
- 'window_start:TIMESTAMP, window_end:TIMESTAMP')
-
-
-def find_words(element):
- import re
- return re.findall(r'[A-Za-z\']+', element)
-
-
-class FormatDoFn(beam.DoFn):
- def process(self, element, window=beam.DoFn.WindowParam):
- ts_format = '%Y-%m-%d %H:%M:%S.%f UTC'
- window_start = window.start.to_utc_datetime().strftime(ts_format)
- window_end = window.end.to_utc_datetime().strftime(ts_format)
- return [{'word': element[0],
- 'count': element[1],
- 'window_start':window_start,
- 'window_end':window_end}]
-
-
-def run(argv=None):
- """Build and run the pipeline."""
-
- parser = argparse.ArgumentParser()
- parser.add_argument(
- '--input_topic', required=True,
- help='Input PubSub topic of the form "/topics/<PROJECT>/<TOPIC>".')
- parser.add_argument(
- '--output_table', required=True,
- help=
- ('Output BigQuery table for results specified as: PROJECT:DATASET.TABLE '
- 'or DATASET.TABLE.'))
- known_args, pipeline_args = parser.parse_known_args(argv)
-
- with beam.Pipeline(argv=pipeline_args) as p:
-
- # Read the text from PubSub messages
- lines = p | beam.io.ReadStringsFromPubSub(known_args.input_topic)
-
- # Capitalize the characters in each line.
- transformed = (lines
- | 'Split' >> (beam.FlatMap(find_words)
- .with_output_types(unicode))
- | 'PairWithOne' >> beam.Map(lambda x: (x, 1))
- | beam.WindowInto(window.FixedWindows(2*60, 0))
- | 'Group' >> beam.GroupByKey()
- | 'Count' >> beam.Map(lambda (word, ones): (word, sum(ones)))
- | 'Format' >> beam.ParDo(FormatDoFn()))
-
- # Write to BigQuery.
- # pylint: disable=expression-not-assigned
- transformed | 'Write' >> beam.io.WriteToBigQuery(
- known_args.output_table,
- schema=TABLE_SCHEMA,
- create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
- write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND)
-
-
-if __name__ == '__main__':
- logging.getLogger().setLevel(logging.INFO)
- run()
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/io/filesystem.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/io/filesystem.py b/sdks/python/apache_beam/io/filesystem.py
index 1f65d0a..db6a1d0 100644
--- a/sdks/python/apache_beam/io/filesystem.py
+++ b/sdks/python/apache_beam/io/filesystem.py
@@ -26,8 +26,6 @@ import zlib
import logging
import time
-from apache_beam.utils.plugin import BeamPlugin
-
logger = logging.getLogger(__name__)
DEFAULT_READ_BUFFER_SIZE = 16 * 1024 * 1024
@@ -201,14 +199,6 @@ class CompressedFile(object):
assert False, 'Possible file corruption.'
except EOFError:
pass # All is as expected!
- elif self._compression_type == CompressionTypes.GZIP:
- # If Gzip file check if there is unused data generated by gzip concat
- if self._decompressor.unused_data != '':
- buf = self._decompressor.unused_data
- self._decompressor = zlib.decompressobj(self._gzip_mask)
- decompressed = self._decompressor.decompress(buf)
- self._read_buffer.write(decompressed)
- continue
else:
self._read_buffer.write(self._decompressor.flush())
@@ -419,7 +409,7 @@ class BeamIOError(IOError):
self.exception_details = exception_details
-class FileSystem(BeamPlugin):
+class FileSystem(object):
"""A class that defines the functions that can be performed on a filesystem.
All methods are abstract and they are for file system providers to
@@ -439,6 +429,16 @@ class FileSystem(BeamPlugin):
return compression_type
@classmethod
+ def get_all_subclasses(cls):
+ """Get all the subclasses of the FileSystem class
+ """
+ all_subclasses = []
+ for subclass in cls.__subclasses__():
+ all_subclasses.append(subclass)
+ all_subclasses.extend(subclass.get_all_subclasses())
+ return all_subclasses
+
+ @classmethod
def scheme(cls):
"""URI scheme for the FileSystem
"""
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/io/gcp/gcsio.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/io/gcp/gcsio.py b/sdks/python/apache_beam/io/gcp/gcsio.py
index 643fbc7..d43c8ba 100644
--- a/sdks/python/apache_beam/io/gcp/gcsio.py
+++ b/sdks/python/apache_beam/io/gcp/gcsio.py
@@ -31,7 +31,6 @@ import re
import threading
import time
import traceback
-import httplib2
from apache_beam.utils import retry
@@ -69,10 +68,6 @@ except ImportError:
# +---------------+------------+-------------+-------------+-------------+
DEFAULT_READ_BUFFER_SIZE = 16 * 1024 * 1024
-# This is the number of seconds the library will wait for GCS operations to
-# complete.
-DEFAULT_HTTP_TIMEOUT_SECONDS = 60
-
# This is the number of seconds the library will wait for a partial-file read
# operation from GCS to complete before retrying.
DEFAULT_READ_SEGMENT_TIMEOUT_SECONDS = 60
@@ -104,7 +99,6 @@ class GcsIO(object):
def __new__(cls, storage_client=None):
if storage_client:
- # This path is only used for testing.
return super(GcsIO, cls).__new__(cls, storage_client)
else:
# Create a single storage client for each thread. We would like to avoid
@@ -114,9 +108,7 @@ class GcsIO(object):
local_state = threading.local()
if getattr(local_state, 'gcsio_instance', None) is None:
credentials = auth.get_service_credentials()
- storage_client = storage.StorageV1(
- credentials=credentials,
- http=httplib2.Http(timeout=DEFAULT_HTTP_TIMEOUT_SECONDS))
+ storage_client = storage.StorageV1(credentials=credentials)
local_state.gcsio_instance = (
super(GcsIO, cls).__new__(cls, storage_client))
local_state.gcsio_instance.client = storage_client
[24/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/test/java/org/apache/beam/runners/core/InMemoryStateInternalsTest.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/InMemoryStateInternalsTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/InMemoryStateInternalsTest.java
index 1c6cd30..b526305 100644
--- a/runners/core-java/src/test/java/org/apache/beam/runners/core/InMemoryStateInternalsTest.java
+++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/InMemoryStateInternalsTest.java
@@ -17,88 +17,545 @@
*/
package org.apache.beam.runners.core;
+import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.hamcrest.Matchers.equalTo;
+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.assertNull;
import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
+import java.util.Arrays;
+import java.util.Map;
+import java.util.Objects;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.state.BagState;
import org.apache.beam.sdk.state.CombiningState;
+import org.apache.beam.sdk.state.GroupingState;
import org.apache.beam.sdk.state.MapState;
+import org.apache.beam.sdk.state.ReadableState;
import org.apache.beam.sdk.state.SetState;
-import org.apache.beam.sdk.state.State;
import org.apache.beam.sdk.state.ValueState;
import org.apache.beam.sdk.state.WatermarkHoldState;
import org.apache.beam.sdk.transforms.Sum;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
import org.hamcrest.Matchers;
+import org.joda.time.Instant;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
-import org.junit.runners.Suite;
/**
- * Tests for {@link InMemoryStateInternals}. This is based on {@link StateInternalsTest}.
+ * Tests for {@link InMemoryStateInternals}.
*/
-@RunWith(Suite.class)
-@Suite.SuiteClasses({
- InMemoryStateInternalsTest.StandardStateInternalsTests.class,
- InMemoryStateInternalsTest.OtherTests.class
-})
+@RunWith(JUnit4.class)
public class InMemoryStateInternalsTest {
+ private static final BoundedWindow WINDOW_1 = new IntervalWindow(new Instant(0), new Instant(10));
+ private static final StateNamespace NAMESPACE_1 = new StateNamespaceForTest("ns1");
+ private static final StateNamespace NAMESPACE_2 = new StateNamespaceForTest("ns2");
+ private static final StateNamespace NAMESPACE_3 = new StateNamespaceForTest("ns3");
- /**
- * A standard StateInternals test.
- */
- @RunWith(JUnit4.class)
- public static class StandardStateInternalsTests extends StateInternalsTest {
- @Override
- protected StateInternals createStateInternals() {
- return new InMemoryStateInternals<>("dummyKey");
- }
+ private static final StateTag<ValueState<String>> STRING_VALUE_ADDR =
+ StateTags.value("stringValue", StringUtf8Coder.of());
+ private static final StateTag<CombiningState<Integer, int[], Integer>>
+ SUM_INTEGER_ADDR = StateTags.combiningValueFromInputInternal(
+ "sumInteger", VarIntCoder.of(), Sum.ofIntegers());
+ private static final StateTag<BagState<String>> STRING_BAG_ADDR =
+ StateTags.bag("stringBag", StringUtf8Coder.of());
+ private static final StateTag<SetState<String>> STRING_SET_ADDR =
+ StateTags.set("stringSet", StringUtf8Coder.of());
+ private static final StateTag<MapState<String, Integer>> STRING_MAP_ADDR =
+ StateTags.map("stringMap", StringUtf8Coder.of(), VarIntCoder.of());
+ private static final StateTag<WatermarkHoldState> WATERMARK_EARLIEST_ADDR =
+ StateTags.watermarkStateInternal("watermark", TimestampCombiner.EARLIEST);
+ private static final StateTag<WatermarkHoldState> WATERMARK_LATEST_ADDR =
+ StateTags.watermarkStateInternal("watermark", TimestampCombiner.LATEST);
+ private static final StateTag<WatermarkHoldState> WATERMARK_EOW_ADDR =
+ StateTags.watermarkStateInternal("watermark", TimestampCombiner.END_OF_WINDOW);
+
+ InMemoryStateInternals<String> underTest = InMemoryStateInternals.forKey("dummyKey");
+
+ @Test
+ public void testValue() throws Exception {
+ ValueState<String> value = underTest.state(NAMESPACE_1, STRING_VALUE_ADDR);
+
+ // State instances are cached, but depend on the namespace.
+ assertThat(underTest.state(NAMESPACE_1, STRING_VALUE_ADDR), Matchers.sameInstance(value));
+ assertThat(
+ underTest.state(NAMESPACE_2, STRING_VALUE_ADDR),
+ Matchers.not(Matchers.sameInstance(value)));
+
+ assertThat(value.read(), Matchers.nullValue());
+ value.write("hello");
+ assertThat(value.read(), equalTo("hello"));
+ value.write("world");
+ assertThat(value.read(), equalTo("world"));
+
+ value.clear();
+ assertThat(value.read(), Matchers.nullValue());
+ assertThat(underTest.state(NAMESPACE_1, STRING_VALUE_ADDR), Matchers.sameInstance(value));
+ }
+
+ @Test
+ public void testBag() throws Exception {
+ BagState<String> value = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
+
+ // State instances are cached, but depend on the namespace.
+ assertThat(value, equalTo(underTest.state(NAMESPACE_1, STRING_BAG_ADDR)));
+ assertThat(value, not(equalTo(underTest.state(NAMESPACE_2, STRING_BAG_ADDR))));
+
+ assertThat(value.read(), Matchers.emptyIterable());
+ value.add("hello");
+ assertThat(value.read(), containsInAnyOrder("hello"));
+
+ value.add("world");
+ assertThat(value.read(), containsInAnyOrder("hello", "world"));
+
+ value.clear();
+ assertThat(value.read(), Matchers.emptyIterable());
+ assertThat(underTest.state(NAMESPACE_1, STRING_BAG_ADDR), Matchers.sameInstance(value));
+ }
+
+ @Test
+ public void testBagIsEmpty() throws Exception {
+ BagState<String> value = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
+
+ assertThat(value.isEmpty().read(), Matchers.is(true));
+ ReadableState<Boolean> readFuture = value.isEmpty();
+ value.add("hello");
+ assertThat(readFuture.read(), Matchers.is(false));
+
+ value.clear();
+ assertThat(readFuture.read(), Matchers.is(true));
+ }
+
+ @Test
+ public void testMergeBagIntoSource() throws Exception {
+ BagState<String> bag1 = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
+ BagState<String> bag2 = underTest.state(NAMESPACE_2, STRING_BAG_ADDR);
+
+ bag1.add("Hello");
+ bag2.add("World");
+ bag1.add("!");
+
+ StateMerging.mergeBags(Arrays.asList(bag1, bag2), bag1);
+
+ // Reading the merged bag gets both the contents
+ assertThat(bag1.read(), containsInAnyOrder("Hello", "World", "!"));
+ assertThat(bag2.read(), Matchers.emptyIterable());
+ }
+
+ @Test
+ public void testMergeBagIntoNewNamespace() throws Exception {
+ BagState<String> bag1 = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
+ BagState<String> bag2 = underTest.state(NAMESPACE_2, STRING_BAG_ADDR);
+ BagState<String> bag3 = underTest.state(NAMESPACE_3, STRING_BAG_ADDR);
+
+ bag1.add("Hello");
+ bag2.add("World");
+ bag1.add("!");
+
+ StateMerging.mergeBags(Arrays.asList(bag1, bag2, bag3), bag3);
+
+ // Reading the merged bag gets both the contents
+ assertThat(bag3.read(), containsInAnyOrder("Hello", "World", "!"));
+ assertThat(bag1.read(), Matchers.emptyIterable());
+ assertThat(bag2.read(), Matchers.emptyIterable());
+ }
+
+ @Test
+ public void testSet() throws Exception {
+ SetState<String> value = underTest.state(NAMESPACE_1, STRING_SET_ADDR);
+
+ // State instances are cached, but depend on the namespace.
+ assertThat(value, equalTo(underTest.state(NAMESPACE_1, STRING_SET_ADDR)));
+ assertThat(value, not(equalTo(underTest.state(NAMESPACE_2, STRING_SET_ADDR))));
+
+ // empty
+ assertThat(value.read(), Matchers.emptyIterable());
+ assertFalse(value.contains("A").read());
+
+ // add
+ value.add("A");
+ value.add("B");
+ value.add("A");
+ assertFalse(value.addIfAbsent("B").read());
+ assertThat(value.read(), containsInAnyOrder("A", "B"));
+
+ // remove
+ value.remove("A");
+ assertThat(value.read(), containsInAnyOrder("B"));
+ value.remove("C");
+ assertThat(value.read(), containsInAnyOrder("B"));
+
+ // contains
+ assertFalse(value.contains("A").read());
+ assertTrue(value.contains("B").read());
+ value.add("C");
+ value.add("D");
+
+ // readLater
+ assertThat(value.readLater().read(), containsInAnyOrder("B", "C", "D"));
+ SetState<String> later = value.readLater();
+ assertThat(later.read(), hasItems("C", "D"));
+ assertFalse(later.contains("A").read());
+
+ // clear
+ value.clear();
+ assertThat(value.read(), Matchers.emptyIterable());
+ assertThat(underTest.state(NAMESPACE_1, STRING_SET_ADDR), Matchers.sameInstance(value));
+
+ }
+
+ @Test
+ public void testSetIsEmpty() throws Exception {
+ SetState<String> value = underTest.state(NAMESPACE_1, STRING_SET_ADDR);
+
+ assertThat(value.isEmpty().read(), Matchers.is(true));
+ ReadableState<Boolean> readFuture = value.isEmpty();
+ value.add("hello");
+ assertThat(readFuture.read(), Matchers.is(false));
+
+ value.clear();
+ assertThat(readFuture.read(), Matchers.is(true));
+ }
+
+ @Test
+ public void testMergeSetIntoSource() throws Exception {
+ SetState<String> set1 = underTest.state(NAMESPACE_1, STRING_SET_ADDR);
+ SetState<String> set2 = underTest.state(NAMESPACE_2, STRING_SET_ADDR);
+
+ set1.add("Hello");
+ set2.add("Hello");
+ set2.add("World");
+ set1.add("!");
+
+ StateMerging.mergeSets(Arrays.asList(set1, set2), set1);
+
+ // Reading the merged set gets both the contents
+ assertThat(set1.read(), containsInAnyOrder("Hello", "World", "!"));
+ assertThat(set2.read(), Matchers.emptyIterable());
}
- /**
- * A specific test of InMemoryStateInternals.
- */
- @RunWith(JUnit4.class)
- public static class OtherTests {
-
- private static final StateNamespace NAMESPACE = new StateNamespaceForTest("ns");
-
- private static final StateTag<ValueState<String>> STRING_VALUE_ADDR =
- StateTags.value("stringValue", StringUtf8Coder.of());
- private static final StateTag<CombiningState<Integer, int[], Integer>>
- SUM_INTEGER_ADDR = StateTags.combiningValueFromInputInternal(
- "sumInteger", VarIntCoder.of(), Sum.ofIntegers());
- private static final StateTag<BagState<String>> STRING_BAG_ADDR =
- StateTags.bag("stringBag", StringUtf8Coder.of());
- private static final StateTag<SetState<String>> STRING_SET_ADDR =
- StateTags.set("stringSet", StringUtf8Coder.of());
- private static final StateTag<MapState<String, Integer>> STRING_MAP_ADDR =
- StateTags.map("stringMap", StringUtf8Coder.of(), VarIntCoder.of());
- private static final StateTag<WatermarkHoldState> WATERMARK_EARLIEST_ADDR =
- StateTags.watermarkStateInternal("watermark", TimestampCombiner.EARLIEST);
- private static final StateTag<WatermarkHoldState> WATERMARK_LATEST_ADDR =
- StateTags.watermarkStateInternal("watermark", TimestampCombiner.LATEST);
- private static final StateTag<WatermarkHoldState> WATERMARK_EOW_ADDR =
- StateTags.watermarkStateInternal("watermark", TimestampCombiner.END_OF_WINDOW);
-
- StateInternals underTest = new InMemoryStateInternals<>("dummyKey");
-
- @Test
- public void testSameInstance() {
- assertSameInstance(STRING_VALUE_ADDR);
- assertSameInstance(SUM_INTEGER_ADDR);
- assertSameInstance(STRING_BAG_ADDR);
- assertSameInstance(STRING_SET_ADDR);
- assertSameInstance(STRING_MAP_ADDR);
- assertSameInstance(WATERMARK_EARLIEST_ADDR);
+ @Test
+ public void testMergeSetIntoNewNamespace() throws Exception {
+ SetState<String> set1 = underTest.state(NAMESPACE_1, STRING_SET_ADDR);
+ SetState<String> set2 = underTest.state(NAMESPACE_2, STRING_SET_ADDR);
+ SetState<String> set3 = underTest.state(NAMESPACE_3, STRING_SET_ADDR);
+
+ set1.add("Hello");
+ set2.add("Hello");
+ set2.add("World");
+ set1.add("!");
+
+ StateMerging.mergeSets(Arrays.asList(set1, set2, set3), set3);
+
+ // Reading the merged set gets both the contents
+ assertThat(set3.read(), containsInAnyOrder("Hello", "World", "!"));
+ assertThat(set1.read(), Matchers.emptyIterable());
+ assertThat(set2.read(), Matchers.emptyIterable());
+ }
+
+ // for testMap
+ private static class MapEntry<K, V> implements Map.Entry<K, V> {
+ private K key;
+ private V value;
+
+ private MapEntry(K key, V value) {
+ this.key = key;
+ this.value = value;
+ }
+
+ static <K, V> Map.Entry<K, V> of(K k, V v) {
+ return new MapEntry<>(k, v);
+ }
+
+ public final K getKey() {
+ return key;
+ }
+ public final V getValue() {
+ return value;
+ }
+
+ public final String toString() {
+ return key + "=" + value;
+ }
+
+ public final int hashCode() {
+ return Objects.hashCode(key) ^ Objects.hashCode(value);
+ }
+
+ public final V setValue(V newValue) {
+ V oldValue = value;
+ value = newValue;
+ return oldValue;
}
- private <T extends State> void assertSameInstance(StateTag<T> address) {
- assertThat(underTest.state(NAMESPACE, address),
- Matchers.sameInstance(underTest.state(NAMESPACE, address)));
+ public final boolean equals(Object o) {
+ if (o == this) {
+ return true;
+ }
+ if (o instanceof Map.Entry) {
+ Map.Entry<?, ?> e = (Map.Entry<?, ?>) o;
+ if (Objects.equals(key, e.getKey())
+ && Objects.equals(value, e.getValue())) {
+ return true;
+ }
+ }
+ return false;
}
}
+ @Test
+ public void testMap() throws Exception {
+ MapState<String, Integer> value = underTest.state(NAMESPACE_1, STRING_MAP_ADDR);
+
+ // State instances are cached, but depend on the namespace.
+ assertThat(value, equalTo(underTest.state(NAMESPACE_1, STRING_MAP_ADDR)));
+ assertThat(value, not(equalTo(underTest.state(NAMESPACE_2, STRING_MAP_ADDR))));
+
+ // put
+ assertThat(value.entries().read(), Matchers.emptyIterable());
+ value.put("A", 1);
+ value.put("B", 2);
+ value.put("A", 11);
+ assertThat(value.putIfAbsent("B", 22).read(), equalTo(2));
+ assertThat(value.entries().read(), containsInAnyOrder(MapEntry.of("A", 11),
+ MapEntry.of("B", 2)));
+
+ // remove
+ value.remove("A");
+ assertThat(value.entries().read(), containsInAnyOrder(MapEntry.of("B", 2)));
+ value.remove("C");
+ assertThat(value.entries().read(), containsInAnyOrder(MapEntry.of("B", 2)));
+
+ // get
+ assertNull(value.get("A").read());
+ assertThat(value.get("B").read(), equalTo(2));
+ value.put("C", 3);
+ value.put("D", 4);
+ assertThat(value.get("C").read(), equalTo(3));
+
+ // iterate
+ value.put("E", 5);
+ value.remove("C");
+ assertThat(value.keys().read(), containsInAnyOrder("B", "D", "E"));
+ assertThat(value.values().read(), containsInAnyOrder(2, 4, 5));
+ assertThat(
+ value.entries().read(),
+ containsInAnyOrder(MapEntry.of("B", 2), MapEntry.of("D", 4), MapEntry.of("E", 5)));
+
+ // readLater
+ assertThat(value.get("B").readLater().read(), equalTo(2));
+ assertNull(value.get("A").readLater().read());
+ assertThat(
+ value.entries().readLater().read(),
+ containsInAnyOrder(MapEntry.of("B", 2), MapEntry.of("D", 4), MapEntry.of("E", 5)));
+
+ // clear
+ value.clear();
+ assertThat(value.entries().read(), Matchers.emptyIterable());
+ assertThat(underTest.state(NAMESPACE_1, STRING_MAP_ADDR), Matchers.sameInstance(value));
+ }
+
+ @Test
+ public void testCombiningValue() throws Exception {
+ GroupingState<Integer, Integer> value = underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
+
+ // State instances are cached, but depend on the namespace.
+ assertEquals(value, underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, SUM_INTEGER_ADDR)));
+
+ assertThat(value.read(), equalTo(0));
+ value.add(2);
+ assertThat(value.read(), equalTo(2));
+
+ value.add(3);
+ assertThat(value.read(), equalTo(5));
+
+ value.clear();
+ assertThat(value.read(), equalTo(0));
+ assertThat(underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR), Matchers.sameInstance(value));
+ }
+
+ @Test
+ public void testCombiningIsEmpty() throws Exception {
+ GroupingState<Integer, Integer> value = underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
+
+ assertThat(value.isEmpty().read(), Matchers.is(true));
+ ReadableState<Boolean> readFuture = value.isEmpty();
+ value.add(5);
+ assertThat(readFuture.read(), Matchers.is(false));
+
+ value.clear();
+ assertThat(readFuture.read(), Matchers.is(true));
+ }
+
+ @Test
+ public void testMergeCombiningValueIntoSource() throws Exception {
+ CombiningState<Integer, int[], Integer> value1 =
+ underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
+ CombiningState<Integer, int[], Integer> value2 =
+ underTest.state(NAMESPACE_2, SUM_INTEGER_ADDR);
+
+ value1.add(5);
+ value2.add(10);
+ value1.add(6);
+
+ assertThat(value1.read(), equalTo(11));
+ assertThat(value2.read(), equalTo(10));
+
+ // Merging clears the old values and updates the result value.
+ StateMerging.mergeCombiningValues(Arrays.asList(value1, value2), value1);
+
+ assertThat(value1.read(), equalTo(21));
+ assertThat(value2.read(), equalTo(0));
+ }
+
+ @Test
+ public void testMergeCombiningValueIntoNewNamespace() throws Exception {
+ CombiningState<Integer, int[], Integer> value1 =
+ underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
+ CombiningState<Integer, int[], Integer> value2 =
+ underTest.state(NAMESPACE_2, SUM_INTEGER_ADDR);
+ CombiningState<Integer, int[], Integer> value3 =
+ underTest.state(NAMESPACE_3, SUM_INTEGER_ADDR);
+
+ value1.add(5);
+ value2.add(10);
+ value1.add(6);
+
+ StateMerging.mergeCombiningValues(Arrays.asList(value1, value2), value3);
+
+ // Merging clears the old values and updates the result value.
+ assertThat(value1.read(), equalTo(0));
+ assertThat(value2.read(), equalTo(0));
+ assertThat(value3.read(), equalTo(21));
+ }
+
+ @Test
+ public void testWatermarkEarliestState() throws Exception {
+ WatermarkHoldState value =
+ underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR);
+
+ // State instances are cached, but depend on the namespace.
+ assertEquals(value, underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, WATERMARK_EARLIEST_ADDR)));
+
+ assertThat(value.read(), Matchers.nullValue());
+ value.add(new Instant(2000));
+ assertThat(value.read(), equalTo(new Instant(2000)));
+
+ value.add(new Instant(3000));
+ assertThat(value.read(), equalTo(new Instant(2000)));
+
+ value.add(new Instant(1000));
+ assertThat(value.read(), equalTo(new Instant(1000)));
+
+ value.clear();
+ assertThat(value.read(), equalTo(null));
+ assertThat(underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR), Matchers.sameInstance(value));
+ }
+
+ @Test
+ public void testWatermarkLatestState() throws Exception {
+ WatermarkHoldState value =
+ underTest.state(NAMESPACE_1, WATERMARK_LATEST_ADDR);
+
+ // State instances are cached, but depend on the namespace.
+ assertEquals(value, underTest.state(NAMESPACE_1, WATERMARK_LATEST_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, WATERMARK_LATEST_ADDR)));
+
+ assertThat(value.read(), Matchers.nullValue());
+ value.add(new Instant(2000));
+ assertThat(value.read(), equalTo(new Instant(2000)));
+
+ value.add(new Instant(3000));
+ assertThat(value.read(), equalTo(new Instant(3000)));
+
+ value.add(new Instant(1000));
+ assertThat(value.read(), equalTo(new Instant(3000)));
+
+ value.clear();
+ assertThat(value.read(), equalTo(null));
+ assertThat(underTest.state(NAMESPACE_1, WATERMARK_LATEST_ADDR), Matchers.sameInstance(value));
+ }
+
+ @Test
+ public void testWatermarkEndOfWindowState() throws Exception {
+ WatermarkHoldState value = underTest.state(NAMESPACE_1, WATERMARK_EOW_ADDR);
+
+ // State instances are cached, but depend on the namespace.
+ assertEquals(value, underTest.state(NAMESPACE_1, WATERMARK_EOW_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, WATERMARK_EOW_ADDR)));
+
+ assertThat(value.read(), Matchers.nullValue());
+ value.add(new Instant(2000));
+ assertThat(value.read(), equalTo(new Instant(2000)));
+
+ value.clear();
+ assertThat(value.read(), equalTo(null));
+ assertThat(underTest.state(NAMESPACE_1, WATERMARK_EOW_ADDR), Matchers.sameInstance(value));
+ }
+
+ @Test
+ public void testWatermarkStateIsEmpty() throws Exception {
+ WatermarkHoldState value =
+ underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR);
+
+ assertThat(value.isEmpty().read(), Matchers.is(true));
+ ReadableState<Boolean> readFuture = value.isEmpty();
+ value.add(new Instant(1000));
+ assertThat(readFuture.read(), Matchers.is(false));
+
+ value.clear();
+ assertThat(readFuture.read(), Matchers.is(true));
+ }
+
+ @Test
+ public void testMergeEarliestWatermarkIntoSource() throws Exception {
+ WatermarkHoldState value1 =
+ underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR);
+ WatermarkHoldState value2 =
+ underTest.state(NAMESPACE_2, WATERMARK_EARLIEST_ADDR);
+
+ value1.add(new Instant(3000));
+ value2.add(new Instant(5000));
+ value1.add(new Instant(4000));
+ value2.add(new Instant(2000));
+
+ // Merging clears the old values and updates the merged value.
+ StateMerging.mergeWatermarks(Arrays.asList(value1, value2), value1, WINDOW_1);
+
+ assertThat(value1.read(), equalTo(new Instant(2000)));
+ assertThat(value2.read(), equalTo(null));
+ }
+
+ @Test
+ public void testMergeLatestWatermarkIntoSource() throws Exception {
+ WatermarkHoldState value1 =
+ underTest.state(NAMESPACE_1, WATERMARK_LATEST_ADDR);
+ WatermarkHoldState value2 =
+ underTest.state(NAMESPACE_2, WATERMARK_LATEST_ADDR);
+ WatermarkHoldState value3 =
+ underTest.state(NAMESPACE_3, WATERMARK_LATEST_ADDR);
+
+ value1.add(new Instant(3000));
+ value2.add(new Instant(5000));
+ value1.add(new Instant(4000));
+ value2.add(new Instant(2000));
+
+ // Merging clears the old values and updates the result value.
+ StateMerging.mergeWatermarks(Arrays.asList(value1, value2), value3, WINDOW_1);
+
+ // Merging clears the old values and updates the result value.
+ assertThat(value3.read(), equalTo(new Instant(5000)));
+ assertThat(value1.read(), equalTo(null));
+ assertThat(value2.read(), equalTo(null));
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/test/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvokerTest.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvokerTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvokerTest.java
index 959909e..a2f6acc 100644
--- a/runners/core-java/src/test/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvokerTest.java
+++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvokerTest.java
@@ -17,22 +17,18 @@
*/
package org.apache.beam.runners.core;
-import static org.apache.beam.sdk.transforms.DoFn.ProcessContinuation.resume;
-import static org.apache.beam.sdk.transforms.DoFn.ProcessContinuation.stop;
import static org.hamcrest.Matchers.greaterThan;
import static org.hamcrest.Matchers.lessThan;
import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertThat;
-import static org.junit.Assert.assertTrue;
import java.util.Collection;
import java.util.concurrent.Executors;
-import org.apache.beam.sdk.io.range.OffsetRange;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.reflect.DoFnInvokers;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRange;
import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
@@ -46,27 +42,19 @@ import org.junit.Test;
/** Tests for {@link OutputAndTimeBoundedSplittableProcessElementInvoker}. */
public class OutputAndTimeBoundedSplittableProcessElementInvokerTest {
private static class SomeFn extends DoFn<Integer, String> {
- private final int numOutputsPerProcessCall;
private final Duration sleepBeforeEachOutput;
- private SomeFn(int numOutputsPerProcessCall, Duration sleepBeforeEachOutput) {
- this.numOutputsPerProcessCall = numOutputsPerProcessCall;
+ private SomeFn(Duration sleepBeforeEachOutput) {
this.sleepBeforeEachOutput = sleepBeforeEachOutput;
}
@ProcessElement
- public ProcessContinuation process(ProcessContext context, OffsetRangeTracker tracker)
+ public void process(ProcessContext context, OffsetRangeTracker tracker)
throws Exception {
- for (long i = tracker.currentRestriction().getFrom(), numIterations = 1;
- tracker.tryClaim(i);
- ++i, ++numIterations) {
+ for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) {
Thread.sleep(sleepBeforeEachOutput.getMillis());
context.output("" + i);
- if (numIterations == numOutputsPerProcessCall) {
- return resume();
- }
}
- return stop();
}
@GetInitialRestriction
@@ -76,8 +64,8 @@ public class OutputAndTimeBoundedSplittableProcessElementInvokerTest {
}
private SplittableProcessElementInvoker<Integer, String, OffsetRange, OffsetRangeTracker>.Result
- runTest(int totalNumOutputs, int numOutputsPerProcessCall, Duration sleepPerElement) {
- SomeFn fn = new SomeFn(numOutputsPerProcessCall, sleepPerElement);
+ runTest(int count, Duration sleepPerElement) {
+ SomeFn fn = new SomeFn(sleepPerElement);
SplittableProcessElementInvoker<Integer, String, OffsetRange, OffsetRangeTracker> invoker =
new OutputAndTimeBoundedSplittableProcessElementInvoker<>(
fn,
@@ -105,15 +93,14 @@ public class OutputAndTimeBoundedSplittableProcessElementInvokerTest {
return invoker.invokeProcessElement(
DoFnInvokers.invokerFor(fn),
- WindowedValue.of(totalNumOutputs, Instant.now(), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING),
- new OffsetRangeTracker(new OffsetRange(0, totalNumOutputs)));
+ WindowedValue.of(count, Instant.now(), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING),
+ new OffsetRangeTracker(new OffsetRange(0, count)));
}
@Test
public void testInvokeProcessElementOutputBounded() throws Exception {
SplittableProcessElementInvoker<Integer, String, OffsetRange, OffsetRangeTracker>.Result res =
- runTest(10000, Integer.MAX_VALUE, Duration.ZERO);
- assertFalse(res.getContinuation().shouldResume());
+ runTest(10000, Duration.ZERO);
OffsetRange residualRange = res.getResidualRestriction();
// Should process the first 100 elements.
assertEquals(1000, residualRange.getFrom());
@@ -123,8 +110,7 @@ public class OutputAndTimeBoundedSplittableProcessElementInvokerTest {
@Test
public void testInvokeProcessElementTimeBounded() throws Exception {
SplittableProcessElementInvoker<Integer, String, OffsetRange, OffsetRangeTracker>.Result res =
- runTest(10000, Integer.MAX_VALUE, Duration.millis(100));
- assertFalse(res.getContinuation().shouldResume());
+ runTest(10000, Duration.millis(100));
OffsetRange residualRange = res.getResidualRestriction();
// Should process ideally around 30 elements - but due to timing flakiness, we can't enforce
// that precisely. Just test that it's not egregiously off.
@@ -134,18 +120,9 @@ public class OutputAndTimeBoundedSplittableProcessElementInvokerTest {
}
@Test
- public void testInvokeProcessElementVoluntaryReturnStop() throws Exception {
+ public void testInvokeProcessElementVoluntaryReturn() throws Exception {
SplittableProcessElementInvoker<Integer, String, OffsetRange, OffsetRangeTracker>.Result res =
- runTest(5, Integer.MAX_VALUE, Duration.millis(100));
- assertFalse(res.getContinuation().shouldResume());
+ runTest(5, Duration.millis(100));
assertNull(res.getResidualRestriction());
}
-
- @Test
- public void testInvokeProcessElementVoluntaryReturnResume() throws Exception {
- SplittableProcessElementInvoker<Integer, String, OffsetRange, OffsetRangeTracker>.Result res =
- runTest(10, 5, Duration.millis(100));
- assertTrue(res.getContinuation().shouldResume());
- assertEquals(new OffsetRange(5, 10), res.getResidualRestriction());
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnRunnerTest.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnRunnerTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnRunnerTest.java
index 4f13af1..9e71300 100644
--- a/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnRunnerTest.java
+++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnRunnerTest.java
@@ -55,7 +55,6 @@ import org.apache.beam.sdk.transforms.windowing.AfterPane;
import org.apache.beam.sdk.transforms.windowing.AfterProcessingTime;
import org.apache.beam.sdk.transforms.windowing.AfterWatermark;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.transforms.windowing.DefaultTrigger;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
@@ -68,7 +67,6 @@ import org.apache.beam.sdk.transforms.windowing.Sessions;
import org.apache.beam.sdk.transforms.windowing.SlidingWindows;
import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
import org.apache.beam.sdk.transforms.windowing.Trigger;
-import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.transforms.windowing.Window.ClosingBehavior;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.transforms.windowing.WindowMappingFn;
@@ -142,40 +140,7 @@ public class ReduceFnRunnerTest {
}
})
.when(mockTrigger).onFire(anyTriggerContext());
- }
-
- /**
- * Tests that a processing time timer does not cause window GC.
- */
- @Test
- public void testProcessingTimeTimerDoesNotGc() throws Exception {
- WindowingStrategy<?, IntervalWindow> strategy =
- WindowingStrategy.of((WindowFn<?, IntervalWindow>) FixedWindows.of(Duration.millis(100)))
- .withTimestampCombiner(TimestampCombiner.EARLIEST)
- .withMode(AccumulationMode.ACCUMULATING_FIRED_PANES)
- .withAllowedLateness(Duration.ZERO)
- .withTrigger(
- Repeatedly.forever(
- AfterProcessingTime.pastFirstElementInPane().plusDelayOf(Duration.millis(10))));
-
- ReduceFnTester<Integer, Integer, IntervalWindow> tester =
- ReduceFnTester.combining(strategy, Sum.ofIntegers(), VarIntCoder.of());
-
- tester.advanceProcessingTime(new Instant(5000));
- injectElement(tester, 2); // processing timer @ 5000 + 10; EOW timer @ 100
- injectElement(tester, 5);
-
- tester.advanceProcessingTime(new Instant(10000));
-
- tester.assertHasOnlyGlobalAndStateFor(
- new IntervalWindow(new Instant(0), new Instant(100)));
-
- assertThat(
- tester.extractOutput(),
- contains(
- isSingleWindowedValue(
- equalTo(7), 2, 0, 100, PaneInfo.createPane(true, false, Timing.EARLY, 0, 0))));
- }
+ }
@Test
public void testOnElementBufferingDiscarding() throws Exception {
@@ -246,52 +211,6 @@ public class ReduceFnRunnerTest {
tester.assertHasOnlyGlobalAndFinishedSetsFor(firstWindow);
}
- /**
- * Tests that with the default trigger we will not produce two ON_TIME panes, even
- * if there are two outputs that are both candidates.
- */
- @Test
- public void testOnlyOneOnTimePane() throws Exception {
- WindowingStrategy<?, IntervalWindow> strategy =
- WindowingStrategy.of((WindowFn<?, IntervalWindow>) FixedWindows.of(Duration.millis(10)))
- .withTrigger(DefaultTrigger.of())
- .withMode(AccumulationMode.ACCUMULATING_FIRED_PANES)
- .withAllowedLateness(Duration.millis(100));
-
- ReduceFnTester<Integer, Integer, IntervalWindow> tester =
- ReduceFnTester.combining(strategy, Sum.ofIntegers(), VarIntCoder.of());
-
- tester.advanceInputWatermark(new Instant(0));
-
- int value1 = 1;
- int value2 = 3;
-
- // A single element that should be in the ON_TIME output
- tester.injectElements(
- TimestampedValue.of(value1, new Instant(1)));
-
- // Should fire ON_TIME
- tester.advanceInputWatermark(new Instant(10));
-
- // The DefaultTrigger should cause output labeled LATE, even though it does not have to be
- // labeled as such.
- tester.injectElements(
- TimestampedValue.of(value2, new Instant(3)));
-
- List<WindowedValue<Integer>> output = tester.extractOutput();
- assertEquals(2, output.size());
-
- assertThat(output.get(0), WindowMatchers.isWindowedValue(equalTo(value1)));
- assertThat(output.get(1), WindowMatchers.isWindowedValue(equalTo(value1 + value2)));
-
- assertThat(
- output.get(0),
- WindowMatchers.valueWithPaneInfo(PaneInfo.createPane(true, false, Timing.ON_TIME, 0, 0)));
- assertThat(
- output.get(1),
- WindowMatchers.valueWithPaneInfo(PaneInfo.createPane(false, false, Timing.LATE, 1, 1)));
- }
-
@Test
public void testOnElementCombiningDiscarding() throws Exception {
// Test basic execution of a trigger using a non-combining window set and discarding mode.
@@ -331,76 +250,6 @@ public class ReduceFnRunnerTest {
}
/**
- * Tests that when a processing time timer comes in after a window is expired
- * it is just ignored.
- */
- @Test
- public void testLateProcessingTimeTimer() throws Exception {
- WindowingStrategy<?, IntervalWindow> strategy =
- WindowingStrategy.of((WindowFn<?, IntervalWindow>) FixedWindows.of(Duration.millis(100)))
- .withTimestampCombiner(TimestampCombiner.EARLIEST)
- .withMode(AccumulationMode.ACCUMULATING_FIRED_PANES)
- .withAllowedLateness(Duration.ZERO)
- .withTrigger(
- Repeatedly.forever(
- AfterProcessingTime.pastFirstElementInPane().plusDelayOf(Duration.millis(10))));
-
- ReduceFnTester<Integer, Integer, IntervalWindow> tester =
- ReduceFnTester.combining(strategy, Sum.ofIntegers(), VarIntCoder.of());
-
- tester.advanceProcessingTime(new Instant(5000));
- injectElement(tester, 2); // processing timer @ 5000 + 10; EOW timer @ 100
- injectElement(tester, 5);
-
- // After this advancement, the window is expired and only the GC process
- // should be allowed to touch it
- tester.advanceInputWatermarkNoTimers(new Instant(100));
-
- // This should not output
- tester.advanceProcessingTime(new Instant(6000));
-
- assertThat(tester.extractOutput(), emptyIterable());
- }
-
- /**
- * Tests that when a processing time timer comes in after a window is expired
- * but in the same bundle it does not cause a spurious output.
- */
- @Test
- public void testCombiningAccumulatingProcessingTime() throws Exception {
- WindowingStrategy<?, IntervalWindow> strategy =
- WindowingStrategy.of((WindowFn<?, IntervalWindow>) FixedWindows.of(Duration.millis(100)))
- .withTimestampCombiner(TimestampCombiner.EARLIEST)
- .withMode(AccumulationMode.ACCUMULATING_FIRED_PANES)
- .withAllowedLateness(Duration.ZERO)
- .withTrigger(
- Repeatedly.forever(
- AfterProcessingTime.pastFirstElementInPane().plusDelayOf(Duration.millis(10))));
-
- ReduceFnTester<Integer, Integer, IntervalWindow> tester =
- ReduceFnTester.combining(strategy, Sum.ofIntegers(), VarIntCoder.of());
-
- tester.advanceProcessingTime(new Instant(5000));
- injectElement(tester, 2); // processing timer @ 5000 + 10; EOW timer @ 100
- injectElement(tester, 5);
-
- tester.advanceInputWatermarkNoTimers(new Instant(100));
- tester.advanceProcessingTimeNoTimers(new Instant(5010));
-
- // Fires the GC/EOW timer at the same time as the processing time timer.
- tester.fireTimers(
- new IntervalWindow(new Instant(0), new Instant(100)),
- TimestampedValue.of(TimeDomain.EVENT_TIME, new Instant(100)),
- TimestampedValue.of(TimeDomain.PROCESSING_TIME, new Instant(5010)));
-
- assertThat(
- tester.extractOutput(),
- contains(
- isSingleWindowedValue(
- equalTo(7), 2, 0, 100, PaneInfo.createPane(true, true, Timing.ON_TIME, 0, 0))));
- }
-
- /**
* Tests that the garbage collection time for a fixed window does not overflow the end of time.
*/
@Test
@@ -467,67 +316,6 @@ public class ReduceFnRunnerTest {
assertThat(tester.extractOutput(), contains(isWindowedValue(equalTo(55))));
}
- /**
- * Tests that when a processing time timers comes in after a window is expired
- * and GC'd it does not cause a spurious output.
- */
- @Test
- public void testCombiningAccumulatingProcessingTimeSeparateBundles() throws Exception {
- WindowingStrategy<?, IntervalWindow> strategy =
- WindowingStrategy.of((WindowFn<?, IntervalWindow>) FixedWindows.of(Duration.millis(100)))
- .withTimestampCombiner(TimestampCombiner.EARLIEST)
- .withMode(AccumulationMode.ACCUMULATING_FIRED_PANES)
- .withAllowedLateness(Duration.ZERO)
- .withTrigger(
- Repeatedly.forever(
- AfterProcessingTime.pastFirstElementInPane().plusDelayOf(Duration.millis(10))));
-
- ReduceFnTester<Integer, Integer, IntervalWindow> tester =
- ReduceFnTester.combining(strategy, Sum.ofIntegers(), VarIntCoder.of());
-
- tester.advanceProcessingTime(new Instant(5000));
- injectElement(tester, 2); // processing timer @ 5000 + 10; EOW timer @ 100
- injectElement(tester, 5);
-
- tester.advanceInputWatermark(new Instant(100));
- tester.advanceProcessingTime(new Instant(5011));
-
- assertThat(
- tester.extractOutput(),
- contains(
- isSingleWindowedValue(
- equalTo(7), 2, 0, 100, PaneInfo.createPane(true, true, Timing.ON_TIME, 0, 0))));
- }
-
- /**
- * Tests that if end-of-window and GC timers come in together, that the pane is correctly
- * marked as final.
- */
- @Test
- public void testCombiningAccumulatingEventTime() throws Exception {
- WindowingStrategy<?, IntervalWindow> strategy =
- WindowingStrategy.of((WindowFn<?, IntervalWindow>) FixedWindows.of(Duration.millis(100)))
- .withTimestampCombiner(TimestampCombiner.EARLIEST)
- .withMode(AccumulationMode.ACCUMULATING_FIRED_PANES)
- .withAllowedLateness(Duration.millis(1))
- .withTrigger(Repeatedly.forever(AfterWatermark.pastEndOfWindow()));
-
- ReduceFnTester<Integer, Integer, IntervalWindow> tester =
- ReduceFnTester.combining(strategy, Sum.ofIntegers(), VarIntCoder.of());
-
- injectElement(tester, 2); // processing timer @ 5000 + 10; EOW timer @ 100
- injectElement(tester, 5);
-
- tester.advanceInputWatermark(new Instant(1000));
-
- assertThat(
- tester.extractOutput(),
- contains(
- isSingleWindowedValue(
- equalTo(7), 2, 0, 100, PaneInfo.createPane(true, true, Timing.ON_TIME, 0, 0))));
- }
-
-
@Test
public void testOnElementCombiningAccumulating() throws Exception {
// Test basic execution of a trigger using a non-combining window set and accumulating mode.
@@ -1501,166 +1289,6 @@ public class ReduceFnRunnerTest {
}
/**
- * Test that it won't fire an empty on-time pane when OnTimeBehavior is FIRE_IF_NON_EMPTY.
- */
- @Test
- public void testEmptyOnTimeWithOnTimeBehaviorFireIfNonEmpty() throws Exception {
-
- WindowingStrategy<?, IntervalWindow> strategy =
- WindowingStrategy.of((WindowFn<?, IntervalWindow>) FixedWindows.of(Duration.millis(10)))
- .withTimestampCombiner(TimestampCombiner.EARLIEST)
- .withTrigger(
- AfterEach.<IntervalWindow>inOrder(
- Repeatedly.forever(
- AfterProcessingTime.pastFirstElementInPane()
- .plusDelayOf(new Duration(5)))
- .orFinally(AfterWatermark.pastEndOfWindow()),
- Repeatedly.forever(
- AfterProcessingTime.pastFirstElementInPane()
- .plusDelayOf(new Duration(25)))))
- .withMode(AccumulationMode.ACCUMULATING_FIRED_PANES)
- .withAllowedLateness(Duration.millis(100))
- .withClosingBehavior(ClosingBehavior.FIRE_ALWAYS)
- .withOnTimeBehavior(Window.OnTimeBehavior.FIRE_IF_NON_EMPTY);
-
- ReduceFnTester<Integer, Integer, IntervalWindow> tester =
- ReduceFnTester.combining(strategy, Sum.ofIntegers(), VarIntCoder.of());
-
- tester.advanceInputWatermark(new Instant(0));
- tester.advanceProcessingTime(new Instant(0));
-
- // Processing time timer for 5
- tester.injectElements(
- TimestampedValue.of(1, new Instant(1)),
- TimestampedValue.of(1, new Instant(3)),
- TimestampedValue.of(1, new Instant(7)),
- TimestampedValue.of(1, new Instant(5)));
-
- // Should fire early pane
- tester.advanceProcessingTime(new Instant(6));
-
- // Should not fire empty on time pane
- tester.advanceInputWatermark(new Instant(11));
-
- // Should fire final GC pane
- tester.advanceInputWatermark(new Instant(10 + 100));
- List<WindowedValue<Integer>> output = tester.extractOutput();
- assertEquals(2, output.size());
-
- assertThat(output.get(0), WindowMatchers.isSingleWindowedValue(4, 1, 0, 10));
- assertThat(output.get(1), WindowMatchers.isSingleWindowedValue(4, 9, 0, 10));
-
- assertThat(
- output.get(0),
- WindowMatchers.valueWithPaneInfo(PaneInfo.createPane(true, false, Timing.EARLY, 0, -1)));
- assertThat(
- output.get(1),
- WindowMatchers.valueWithPaneInfo(PaneInfo.createPane(false, true, Timing.LATE, 1, 0)));
- }
-
- /**
- * Test that it fires an empty on-time isFinished pane when OnTimeBehavior is FIRE_ALWAYS
- * and ClosingBehavior is FIRE_IF_NON_EMPTY.
- *
- * <p>This is a test just for backward compatibility.
- */
- @Test
- public void testEmptyOnTimeWithOnTimeBehaviorBackwardCompatibility() throws Exception {
-
- WindowingStrategy<?, IntervalWindow> strategy =
- WindowingStrategy.of((WindowFn<?, IntervalWindow>) FixedWindows.of(Duration.millis(10)))
- .withTimestampCombiner(TimestampCombiner.EARLIEST)
- .withTrigger(AfterWatermark.pastEndOfWindow()
- .withEarlyFirings(AfterPane.elementCountAtLeast(1)))
- .withMode(AccumulationMode.ACCUMULATING_FIRED_PANES)
- .withAllowedLateness(Duration.millis(0))
- .withClosingBehavior(ClosingBehavior.FIRE_IF_NON_EMPTY);
-
- ReduceFnTester<Integer, Integer, IntervalWindow> tester =
- ReduceFnTester.combining(strategy, Sum.ofIntegers(), VarIntCoder.of());
-
- tester.advanceInputWatermark(new Instant(0));
- tester.advanceProcessingTime(new Instant(0));
-
- tester.injectElements(
- TimestampedValue.of(1, new Instant(1)));
-
- // Should fire empty on time isFinished pane
- tester.advanceInputWatermark(new Instant(11));
-
- List<WindowedValue<Integer>> output = tester.extractOutput();
- assertEquals(2, output.size());
-
- assertThat(
- output.get(0),
- WindowMatchers.valueWithPaneInfo(PaneInfo.createPane(true, false, Timing.EARLY, 0, -1)));
- assertThat(
- output.get(1),
- WindowMatchers.valueWithPaneInfo(PaneInfo.createPane(false, true, Timing.ON_TIME, 1, 0)));
- }
-
- /**
- * Test that it won't fire an empty on-time pane when OnTimeBehavior is FIRE_IF_NON_EMPTY
- * and when receiving late data.
- */
- @Test
- public void testEmptyOnTimeWithOnTimeBehaviorFireIfNonEmptyAndLateData() throws Exception {
-
- WindowingStrategy<?, IntervalWindow> strategy =
- WindowingStrategy.of((WindowFn<?, IntervalWindow>) FixedWindows.of(Duration.millis(10)))
- .withTimestampCombiner(TimestampCombiner.EARLIEST)
- .withTrigger(
- AfterEach.<IntervalWindow>inOrder(
- Repeatedly.forever(
- AfterProcessingTime.pastFirstElementInPane()
- .plusDelayOf(new Duration(5)))
- .orFinally(AfterWatermark.pastEndOfWindow()),
- Repeatedly.forever(
- AfterProcessingTime.pastFirstElementInPane()
- .plusDelayOf(new Duration(25)))))
- .withMode(AccumulationMode.ACCUMULATING_FIRED_PANES)
- .withAllowedLateness(Duration.millis(100))
- .withOnTimeBehavior(Window.OnTimeBehavior.FIRE_IF_NON_EMPTY);
-
- ReduceFnTester<Integer, Integer, IntervalWindow> tester =
- ReduceFnTester.combining(strategy, Sum.ofIntegers(), VarIntCoder.of());
-
- tester.advanceInputWatermark(new Instant(0));
- tester.advanceProcessingTime(new Instant(0));
-
- // Processing time timer for 5
- tester.injectElements(
- TimestampedValue.of(1, new Instant(1)),
- TimestampedValue.of(1, new Instant(3)),
- TimestampedValue.of(1, new Instant(7)),
- TimestampedValue.of(1, new Instant(5)));
-
- // Should fire early pane
- tester.advanceProcessingTime(new Instant(6));
-
- // Should not fire empty on time pane
- tester.advanceInputWatermark(new Instant(11));
-
- // Processing late data, and should fire late pane
- tester.injectElements(
- TimestampedValue.of(1, new Instant(9)));
- tester.advanceProcessingTime(new Instant(6 + 25 + 1));
-
- List<WindowedValue<Integer>> output = tester.extractOutput();
- assertEquals(2, output.size());
-
- assertThat(output.get(0), WindowMatchers.isSingleWindowedValue(4, 1, 0, 10));
- assertThat(output.get(1), WindowMatchers.isSingleWindowedValue(5, 9, 0, 10));
-
- assertThat(
- output.get(0),
- WindowMatchers.valueWithPaneInfo(PaneInfo.createPane(true, false, Timing.EARLY, 0, -1)));
- assertThat(
- output.get(1),
- WindowMatchers.valueWithPaneInfo(PaneInfo.createPane(false, false, Timing.LATE, 1, 0)));
- }
-
- /**
* Tests for processing time firings after the watermark passes the end of the window.
* Specifically, verify the proper triggerings and pane-info of a typical speculative/on-time/late
* when the on-time pane is non-empty.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java
index 7ca96b9..7f83eae 100644
--- a/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java
+++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java
@@ -318,19 +318,6 @@ public class ReduceFnTester<InputT, OutputT, W extends BoundedWindow> {
}
@SafeVarargs
- public final void assertHasOnlyGlobalAndStateFor(W... expectedWindows) {
- assertHasOnlyGlobalAndAllowedTags(
- ImmutableSet.copyOf(expectedWindows),
- ImmutableSet.<StateTag<?>>of(
- ((SystemReduceFn<?, ?, ?, ?, ?>) reduceFn).getBufferTag(),
- TriggerStateMachineRunner.FINISHED_BITS_TAG,
- PaneInfoTracker.PANE_INFO_TAG,
- WatermarkHold.watermarkHoldTagForTimestampCombiner(
- objectStrategy.getTimestampCombiner()),
- WatermarkHold.EXTRA_HOLD_TAG));
- }
-
- @SafeVarargs
public final void assertHasOnlyGlobalAndFinishedSetsAndPaneInfoFor(W... expectedWindows) {
assertHasOnlyGlobalAndAllowedTags(
ImmutableSet.copyOf(expectedWindows),
@@ -433,10 +420,6 @@ public class ReduceFnTester<InputT, OutputT, W extends BoundedWindow> {
return result;
}
- public void advanceInputWatermarkNoTimers(Instant newInputWatermark) throws Exception {
- timerInternals.advanceInputWatermark(newInputWatermark);
- }
-
/**
* Advance the input watermark to the specified time, firing any timers that should
* fire. Then advance the output watermark as far as possible.
@@ -468,10 +451,6 @@ public class ReduceFnTester<InputT, OutputT, W extends BoundedWindow> {
runner.persist();
}
- public void advanceProcessingTimeNoTimers(Instant newProcessingTime) throws Exception {
- timerInternals.advanceProcessingTime(newProcessingTime);
- }
-
/**
* If {@link #autoAdvanceOutputWatermark} is {@literal false}, advance the output watermark
* to the given value. Otherwise throw.
@@ -529,8 +508,8 @@ public class ReduceFnTester<InputT, OutputT, W extends BoundedWindow> {
for (TimestampedValue<InputT> value : values) {
WindowTracing.trace("TriggerTester.injectElements: {}", value);
}
-
- Iterable<WindowedValue<InputT>> inputs =
+ ReduceFnRunner<String, InputT, OutputT, W> runner = createRunner();
+ runner.processElements(
Iterables.transform(
Arrays.asList(values),
new Function<TimestampedValue<InputT>, WindowedValue<InputT>>() {
@@ -548,12 +527,7 @@ public class ReduceFnTester<InputT, OutputT, W extends BoundedWindow> {
throw new RuntimeException(e);
}
}
- });
-
- ReduceFnRunner<String, InputT, OutputT, W> runner = createRunner();
- runner.processElements(
- new LateDataDroppingDoFnRunner.LateDataFilter(objectStrategy, timerInternals)
- .filter(KEY, inputs));
+ }));
// Persist after each bundle.
runner.persist();
@@ -561,27 +535,13 @@ public class ReduceFnTester<InputT, OutputT, W extends BoundedWindow> {
public void fireTimer(W window, Instant timestamp, TimeDomain domain) throws Exception {
ReduceFnRunner<String, InputT, OutputT, W> runner = createRunner();
- ArrayList<TimerData> timers = new ArrayList<>(1);
+ ArrayList timers = new ArrayList(1);
timers.add(
TimerData.of(StateNamespaces.window(windowFn.windowCoder(), window), timestamp, domain));
runner.onTimers(timers);
runner.persist();
}
- public void fireTimers(W window, TimestampedValue<TimeDomain>... timers) throws Exception {
- ReduceFnRunner<String, InputT, OutputT, W> runner = createRunner();
- ArrayList<TimerData> timerData = new ArrayList<>(timers.length);
- for (TimestampedValue<TimeDomain> timer : timers) {
- timerData.add(
- TimerData.of(
- StateNamespaces.window(windowFn.windowCoder(), window),
- timer.getTimestamp(),
- timer.getValue()));
- }
- runner.onTimers(timerData);
- runner.persist();
- }
-
/**
* Convey the simulated state and implement {@link #outputWindowedValue} to capture all output
* elements.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/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
index 7449af3..d242431 100644
--- 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
@@ -17,9 +17,6 @@
*/
package org.apache.beam.runners.core;
-import static org.apache.beam.sdk.transforms.DoFn.ProcessContinuation.resume;
-import static org.apache.beam.sdk.transforms.DoFn.ProcessContinuation.stop;
-import static org.hamcrest.Matchers.contains;
import static org.hamcrest.Matchers.greaterThanOrEqualTo;
import static org.hamcrest.Matchers.hasItem;
import static org.hamcrest.Matchers.hasItems;
@@ -38,15 +35,16 @@ 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.io.range.OffsetRange;
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;
@@ -55,7 +53,6 @@ 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.KV;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TimestampedValue;
import org.apache.beam.sdk.values.TupleTag;
@@ -114,7 +111,9 @@ public class SplittableParDoProcessFnTest {
private static class ProcessFnTester<
InputT, OutputT, RestrictionT, TrackerT extends RestrictionTracker<RestrictionT>>
implements AutoCloseable {
- private final DoFnTester<KeyedWorkItem<String, KV<InputT, RestrictionT>>, OutputT> tester;
+ private final DoFnTester<
+ KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT>
+ tester;
private Instant currentProcessingTime;
private InMemoryTimerInternals timerInternals;
@@ -195,13 +194,14 @@ public class SplittableParDoProcessFnTest {
void startElement(InputT element, RestrictionT restriction) throws Exception {
startElement(
WindowedValue.of(
- KV.of(element, restriction),
+ ElementAndRestriction.of(element, restriction),
currentProcessingTime,
GlobalWindow.INSTANCE,
PaneInfo.ON_TIME_AND_ONLY_FIRING));
}
- void startElement(WindowedValue<KV<InputT, RestrictionT>> windowedValue) throws Exception {
+ void startElement(WindowedValue<ElementAndRestriction<InputT, RestrictionT>> windowedValue)
+ throws Exception {
tester.processElement(
KeyedWorkItems.elementsWorkItem("key", Collections.singletonList(windowedValue)));
}
@@ -223,7 +223,8 @@ public class SplittableParDoProcessFnTest {
return false;
}
tester.processElement(
- KeyedWorkItems.<String, KV<InputT, RestrictionT>>timersWorkItem("key", timers));
+ KeyedWorkItems.<String, ElementAndRestriction<InputT, RestrictionT>>timersWorkItem(
+ "key", timers));
return true;
}
@@ -308,7 +309,7 @@ public class SplittableParDoProcessFnTest {
MAX_BUNDLE_DURATION);
tester.startElement(
WindowedValue.of(
- KV.of(42, new SomeRestriction()),
+ ElementAndRestriction.of(42, new SomeRestriction()),
base,
Collections.singletonList(w),
PaneInfo.ON_TIME_AND_ONLY_FIRING));
@@ -368,71 +369,16 @@ public class SplittableParDoProcessFnTest {
assertEquals(null, tester.getWatermarkHold());
}
- /** A simple splittable {@link DoFn} that outputs the given element every 5 seconds forever. */
- private static class SelfInitiatedResumeFn extends DoFn<Integer, String> {
- @ProcessElement
- public ProcessContinuation process(ProcessContext c, SomeRestrictionTracker tracker) {
- c.output(c.element().toString());
- return resume().withResumeDelay(Duration.standardSeconds(5));
- }
-
- @GetInitialRestriction
- public SomeRestriction getInitialRestriction(Integer elem) {
- return new SomeRestriction();
- }
- }
-
- @Test
- public void testResumeSetsTimer() throws Exception {
- DoFn<Integer, String> fn = new SelfInitiatedResumeFn();
- Instant base = Instant.now();
- 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(42, new SomeRestriction());
- assertThat(tester.takeOutputElements(), contains("42"));
-
- // Should resume after 5 seconds: advancing by 3 seconds should have no effect.
- assertFalse(tester.advanceProcessingTimeBy(Duration.standardSeconds(3)));
- assertTrue(tester.takeOutputElements().isEmpty());
-
- // 6 seconds should be enough should invoke the fn again.
- assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(3)));
- assertThat(tester.takeOutputElements(), contains("42"));
-
- // Should again resume after 5 seconds: advancing by 3 seconds should again have no effect.
- assertFalse(tester.advanceProcessingTimeBy(Duration.standardSeconds(3)));
- assertTrue(tester.takeOutputElements().isEmpty());
-
- // 6 seconds should again be enough.
- assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(3)));
- assertThat(tester.takeOutputElements(), contains("42"));
- }
-
- /** A splittable {@link DoFn} that generates the sequence [init, init + total). */
+ /**
+ * A splittable {@link DoFn} that generates the sequence [init, init + total).
+ */
private static class CounterFn extends DoFn<Integer, String> {
- private final int numOutputsPerCall;
-
- public CounterFn(int numOutputsPerCall) {
- this.numOutputsPerCall = numOutputsPerCall;
- }
-
@ProcessElement
- public ProcessContinuation process(ProcessContext c, OffsetRangeTracker tracker) {
- for (long i = tracker.currentRestriction().getFrom(), numIterations = 0;
- tracker.tryClaim(i); ++i, ++numIterations) {
+ public void process(ProcessContext c, OffsetRangeTracker tracker) {
+ for (long i = tracker.currentRestriction().getFrom();
+ tracker.tryClaim(i); ++i) {
c.output(String.valueOf(c.element() + i));
- if (numIterations == numOutputsPerCall) {
- return resume();
- }
}
- return stop();
}
@GetInitialRestriction
@@ -441,35 +387,10 @@ public class SplittableParDoProcessFnTest {
}
}
- public void testResumeCarriesOverState() throws Exception {
- DoFn<Integer, String> fn = new CounterFn(1);
- Instant base = Instant.now();
- ProcessFnTester<Integer, String, OffsetRange, OffsetRangeTracker> tester =
- new ProcessFnTester<>(
- base,
- fn,
- BigEndianIntegerCoder.of(),
- SerializableCoder.of(OffsetRange.class),
- MAX_OUTPUTS_PER_BUNDLE,
- MAX_BUNDLE_DURATION);
-
- tester.startElement(42, new OffsetRange(0, 3));
- assertThat(tester.takeOutputElements(), contains("42"));
- assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
- assertThat(tester.takeOutputElements(), contains("43"));
- assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
- assertThat(tester.takeOutputElements(), contains("44"));
- assertTrue(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
- // After outputting all 3 items, should not output anything more.
- assertEquals(0, tester.takeOutputElements().size());
- // Should also not ask to resume.
- assertFalse(tester.advanceProcessingTimeBy(Duration.standardSeconds(1)));
- }
-
@Test
public void testCheckpointsAfterNumOutputs() throws Exception {
int max = 100;
- DoFn<Integer, String> fn = new CounterFn(Integer.MAX_VALUE);
+ DoFn<Integer, String> fn = new CounterFn();
Instant base = Instant.now();
int baseIndex = 42;
@@ -511,7 +432,7 @@ public class SplittableParDoProcessFnTest {
// 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(Integer.MAX_VALUE);
+ DoFn<Integer, String> fn = new CounterFn();
Instant base = Instant.now();
int baseIndex = 42;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/test/java/org/apache/beam/runners/core/StateInternalsTest.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/StateInternalsTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/StateInternalsTest.java
deleted file mode 100644
index ae07fe6..0000000
--- a/runners/core-java/src/test/java/org/apache/beam/runners/core/StateInternalsTest.java
+++ /dev/null
@@ -1,613 +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.containsInAnyOrder;
-import static org.hamcrest.Matchers.equalTo;
-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.assertNull;
-import static org.junit.Assert.assertThat;
-import static org.junit.Assert.assertTrue;
-
-import com.google.common.collect.Iterables;
-import java.util.Arrays;
-import java.util.Map;
-import java.util.Objects;
-import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.sdk.coders.VarIntCoder;
-import org.apache.beam.sdk.state.BagState;
-import org.apache.beam.sdk.state.CombiningState;
-import org.apache.beam.sdk.state.GroupingState;
-import org.apache.beam.sdk.state.MapState;
-import org.apache.beam.sdk.state.ReadableState;
-import org.apache.beam.sdk.state.SetState;
-import org.apache.beam.sdk.state.ValueState;
-import org.apache.beam.sdk.state.WatermarkHoldState;
-import org.apache.beam.sdk.transforms.Sum;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
-import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
-import org.hamcrest.Matchers;
-import org.joda.time.Instant;
-import org.junit.Before;
-import org.junit.Test;
-
-/**
- * Tests for {@link StateInternals}.
- */
-public abstract class StateInternalsTest {
-
- private static final BoundedWindow WINDOW_1 = new IntervalWindow(new Instant(0), new Instant(10));
- private static final StateNamespace NAMESPACE_1 = new StateNamespaceForTest("ns1");
- private static final StateNamespace NAMESPACE_2 = new StateNamespaceForTest("ns2");
- private static final StateNamespace NAMESPACE_3 = new StateNamespaceForTest("ns3");
-
- private static final StateTag<ValueState<String>> STRING_VALUE_ADDR =
- StateTags.value("stringValue", StringUtf8Coder.of());
- private static final StateTag<CombiningState<Integer, int[], Integer>>
- SUM_INTEGER_ADDR = StateTags.combiningValueFromInputInternal(
- "sumInteger", VarIntCoder.of(), Sum.ofIntegers());
- private static final StateTag<BagState<String>> STRING_BAG_ADDR =
- StateTags.bag("stringBag", StringUtf8Coder.of());
- private static final StateTag<SetState<String>> STRING_SET_ADDR =
- StateTags.set("stringSet", StringUtf8Coder.of());
- private static final StateTag<MapState<String, Integer>> STRING_MAP_ADDR =
- StateTags.map("stringMap", StringUtf8Coder.of(), VarIntCoder.of());
- private static final StateTag<WatermarkHoldState> WATERMARK_EARLIEST_ADDR =
- StateTags.watermarkStateInternal("watermark", TimestampCombiner.EARLIEST);
- private static final StateTag<WatermarkHoldState> WATERMARK_LATEST_ADDR =
- StateTags.watermarkStateInternal("watermark", TimestampCombiner.LATEST);
- private static final StateTag<WatermarkHoldState> WATERMARK_EOW_ADDR =
- StateTags.watermarkStateInternal("watermark", TimestampCombiner.END_OF_WINDOW);
-
- private StateInternals underTest;
-
- @Before
- public void setUp() {
- this.underTest = createStateInternals();
- }
-
- protected abstract StateInternals createStateInternals();
-
- @Test
- public void testValue() throws Exception {
- ValueState<String> value = underTest.state(NAMESPACE_1, STRING_VALUE_ADDR);
-
- // State instances are cached, but depend on the namespace.
- assertThat(underTest.state(NAMESPACE_1, STRING_VALUE_ADDR), equalTo(value));
- assertThat(
- underTest.state(NAMESPACE_2, STRING_VALUE_ADDR),
- Matchers.not(equalTo(value)));
-
- assertThat(value.read(), Matchers.nullValue());
- value.write("hello");
- assertThat(value.read(), equalTo("hello"));
- value.write("world");
- assertThat(value.read(), equalTo("world"));
-
- value.clear();
- assertThat(value.read(), Matchers.nullValue());
- assertThat(underTest.state(NAMESPACE_1, STRING_VALUE_ADDR), equalTo(value));
- }
-
- @Test
- public void testBag() throws Exception {
- BagState<String> value = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
-
- // State instances are cached, but depend on the namespace.
- assertThat(value, equalTo(underTest.state(NAMESPACE_1, STRING_BAG_ADDR)));
- assertThat(value, not(equalTo(underTest.state(NAMESPACE_2, STRING_BAG_ADDR))));
-
- assertThat(value.read(), Matchers.emptyIterable());
- value.add("hello");
- assertThat(value.read(), containsInAnyOrder("hello"));
-
- value.add("world");
- assertThat(value.read(), containsInAnyOrder("hello", "world"));
-
- value.clear();
- assertThat(value.read(), Matchers.emptyIterable());
- assertThat(underTest.state(NAMESPACE_1, STRING_BAG_ADDR), equalTo(value));
- }
-
- @Test
- public void testBagIsEmpty() throws Exception {
- BagState<String> value = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
-
- assertThat(value.isEmpty().read(), Matchers.is(true));
- ReadableState<Boolean> readFuture = value.isEmpty();
- value.add("hello");
- assertThat(readFuture.read(), Matchers.is(false));
-
- value.clear();
- assertThat(readFuture.read(), Matchers.is(true));
- }
-
- @Test
- public void testMergeBagIntoSource() throws Exception {
- BagState<String> bag1 = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
- BagState<String> bag2 = underTest.state(NAMESPACE_2, STRING_BAG_ADDR);
-
- bag1.add("Hello");
- bag2.add("World");
- bag1.add("!");
-
- StateMerging.mergeBags(Arrays.asList(bag1, bag2), bag1);
-
- // Reading the merged bag gets both the contents
- assertThat(bag1.read(), containsInAnyOrder("Hello", "World", "!"));
- assertThat(bag2.read(), Matchers.emptyIterable());
- }
-
- @Test
- public void testMergeBagIntoNewNamespace() throws Exception {
- BagState<String> bag1 = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
- BagState<String> bag2 = underTest.state(NAMESPACE_2, STRING_BAG_ADDR);
- BagState<String> bag3 = underTest.state(NAMESPACE_3, STRING_BAG_ADDR);
-
- bag1.add("Hello");
- bag2.add("World");
- bag1.add("!");
-
- StateMerging.mergeBags(Arrays.asList(bag1, bag2, bag3), bag3);
-
- // Reading the merged bag gets both the contents
- assertThat(bag3.read(), containsInAnyOrder("Hello", "World", "!"));
- assertThat(bag1.read(), Matchers.emptyIterable());
- assertThat(bag2.read(), Matchers.emptyIterable());
- }
-
- @Test
- public void testSet() throws Exception {
-
- SetState<String> value = underTest.state(NAMESPACE_1, STRING_SET_ADDR);
-
- // State instances are cached, but depend on the namespace.
- assertThat(value, equalTo(underTest.state(NAMESPACE_1, STRING_SET_ADDR)));
- assertThat(value, not(equalTo(underTest.state(NAMESPACE_2, STRING_SET_ADDR))));
-
- // empty
- assertThat(value.read(), Matchers.emptyIterable());
- assertFalse(value.contains("A").read());
-
- // add
- value.add("A");
- value.add("B");
- value.add("A");
- assertFalse(value.addIfAbsent("B").read());
- assertThat(value.read(), containsInAnyOrder("A", "B"));
-
- // remove
- value.remove("A");
- assertThat(value.read(), containsInAnyOrder("B"));
- value.remove("C");
- assertThat(value.read(), containsInAnyOrder("B"));
-
- // contains
- assertFalse(value.contains("A").read());
- assertTrue(value.contains("B").read());
- value.add("C");
- value.add("D");
-
- // readLater
- assertThat(value.readLater().read(), containsInAnyOrder("B", "C", "D"));
- SetState<String> later = value.readLater();
- assertThat(later.read(), hasItems("C", "D"));
- assertFalse(later.contains("A").read());
-
- // clear
- value.clear();
- assertThat(value.read(), Matchers.emptyIterable());
- assertThat(underTest.state(NAMESPACE_1, STRING_SET_ADDR), equalTo(value));
-
- }
-
- @Test
- public void testSetIsEmpty() throws Exception {
-
- SetState<String> value = underTest.state(NAMESPACE_1, STRING_SET_ADDR);
-
- assertThat(value.isEmpty().read(), Matchers.is(true));
- ReadableState<Boolean> readFuture = value.isEmpty();
- value.add("hello");
- assertThat(readFuture.read(), Matchers.is(false));
-
- value.clear();
- assertThat(readFuture.read(), Matchers.is(true));
- }
-
- @Test
- public void testMergeSetIntoSource() throws Exception {
-
- SetState<String> set1 = underTest.state(NAMESPACE_1, STRING_SET_ADDR);
- SetState<String> set2 = underTest.state(NAMESPACE_2, STRING_SET_ADDR);
-
- set1.add("Hello");
- set2.add("Hello");
- set2.add("World");
- set1.add("!");
-
- StateMerging.mergeSets(Arrays.asList(set1, set2), set1);
-
- // Reading the merged set gets both the contents
- assertThat(set1.read(), containsInAnyOrder("Hello", "World", "!"));
- assertThat(set2.read(), Matchers.emptyIterable());
- }
-
- @Test
- public void testMergeSetIntoNewNamespace() throws Exception {
-
- SetState<String> set1 = underTest.state(NAMESPACE_1, STRING_SET_ADDR);
- SetState<String> set2 = underTest.state(NAMESPACE_2, STRING_SET_ADDR);
- SetState<String> set3 = underTest.state(NAMESPACE_3, STRING_SET_ADDR);
-
- set1.add("Hello");
- set2.add("Hello");
- set2.add("World");
- set1.add("!");
-
- StateMerging.mergeSets(Arrays.asList(set1, set2, set3), set3);
-
- // Reading the merged set gets both the contents
- assertThat(set3.read(), containsInAnyOrder("Hello", "World", "!"));
- assertThat(set1.read(), Matchers.emptyIterable());
- assertThat(set2.read(), Matchers.emptyIterable());
- }
-
- // for testMap
- private static class MapEntry<K, V> implements Map.Entry<K, V> {
- private K key;
- private V value;
-
- private MapEntry(K key, V value) {
- this.key = key;
- this.value = value;
- }
-
- static <K, V> Map.Entry<K, V> of(K k, V v) {
- return new MapEntry<>(k, v);
- }
-
- public final K getKey() {
- return key;
- }
- public final V getValue() {
- return value;
- }
-
- public final String toString() {
- return key + "=" + value;
- }
-
- public final int hashCode() {
- return Objects.hashCode(key) ^ Objects.hashCode(value);
- }
-
- public final V setValue(V newValue) {
- V oldValue = value;
- value = newValue;
- return oldValue;
- }
-
- public final boolean equals(Object o) {
- if (o == this) {
- return true;
- }
- if (o instanceof Map.Entry) {
- Map.Entry<?, ?> e = (Map.Entry<?, ?>) o;
- if (Objects.equals(key, e.getKey())
- && Objects.equals(value, e.getValue())) {
- return true;
- }
- }
- return false;
- }
- }
-
- @Test
- public void testMap() throws Exception {
-
- MapState<String, Integer> value = underTest.state(NAMESPACE_1, STRING_MAP_ADDR);
-
- // State instances are cached, but depend on the namespace.
- assertThat(value, equalTo(underTest.state(NAMESPACE_1, STRING_MAP_ADDR)));
- assertThat(value, not(equalTo(underTest.state(NAMESPACE_2, STRING_MAP_ADDR))));
-
- // put
- assertThat(value.entries().read(), Matchers.emptyIterable());
- value.put("A", 1);
- value.put("B", 2);
- value.put("A", 11);
- assertThat(value.putIfAbsent("B", 22).read(), equalTo(2));
- assertThat(value.entries().read(), containsInAnyOrder(MapEntry.of("A", 11),
- MapEntry.of("B", 2)));
-
- // remove
- value.remove("A");
- assertThat(value.entries().read(), containsInAnyOrder(MapEntry.of("B", 2)));
- value.remove("C");
- assertThat(value.entries().read(), containsInAnyOrder(MapEntry.of("B", 2)));
-
- // get
- assertNull(value.get("A").read());
- assertThat(value.get("B").read(), equalTo(2));
- value.put("C", 3);
- value.put("D", 4);
- assertThat(value.get("C").read(), equalTo(3));
-
- // iterate
- value.put("E", 5);
- value.remove("C");
- assertThat(value.keys().read(), containsInAnyOrder("B", "D", "E"));
- assertThat(value.values().read(), containsInAnyOrder(2, 4, 5));
- assertThat(
- value.entries().read(),
- containsInAnyOrder(MapEntry.of("B", 2), MapEntry.of("D", 4), MapEntry.of("E", 5)));
-
- // readLater
- assertThat(value.get("B").readLater().read(), equalTo(2));
- assertNull(value.get("A").readLater().read());
- assertThat(
- value.entries().readLater().read(),
- containsInAnyOrder(MapEntry.of("B", 2), MapEntry.of("D", 4), MapEntry.of("E", 5)));
-
- // clear
- value.clear();
- assertThat(value.entries().read(), Matchers.emptyIterable());
- assertThat(underTest.state(NAMESPACE_1, STRING_MAP_ADDR), equalTo(value));
- }
-
- @Test
- public void testCombiningValue() throws Exception {
-
- GroupingState<Integer, Integer> value = underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
-
- // State instances are cached, but depend on the namespace.
- assertEquals(value, underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR));
- assertFalse(value.equals(underTest.state(NAMESPACE_2, SUM_INTEGER_ADDR)));
-
- assertThat(value.read(), equalTo(0));
- value.add(2);
- assertThat(value.read(), equalTo(2));
-
- value.add(3);
- assertThat(value.read(), equalTo(5));
-
- value.clear();
- assertThat(value.read(), equalTo(0));
- assertThat(underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR), equalTo(value));
- }
-
- @Test
- public void testCombiningIsEmpty() throws Exception {
- GroupingState<Integer, Integer> value = underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
-
- assertThat(value.isEmpty().read(), Matchers.is(true));
- ReadableState<Boolean> readFuture = value.isEmpty();
- value.add(5);
- assertThat(readFuture.read(), Matchers.is(false));
-
- value.clear();
- assertThat(readFuture.read(), Matchers.is(true));
- }
-
- @Test
- public void testMergeCombiningValueIntoSource() throws Exception {
- CombiningState<Integer, int[], Integer> value1 =
- underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
- CombiningState<Integer, int[], Integer> value2 =
- underTest.state(NAMESPACE_2, SUM_INTEGER_ADDR);
-
- value1.add(5);
- value2.add(10);
- value1.add(6);
-
- assertThat(value1.read(), equalTo(11));
- assertThat(value2.read(), equalTo(10));
-
- // Merging clears the old values and updates the result value.
- StateMerging.mergeCombiningValues(Arrays.asList(value1, value2), value1);
-
- assertThat(value1.read(), equalTo(21));
- assertThat(value2.read(), equalTo(0));
- }
-
- @Test
- public void testMergeCombiningValueIntoNewNamespace() throws Exception {
- CombiningState<Integer, int[], Integer> value1 =
- underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
- CombiningState<Integer, int[], Integer> value2 =
- underTest.state(NAMESPACE_2, SUM_INTEGER_ADDR);
- CombiningState<Integer, int[], Integer> value3 =
- underTest.state(NAMESPACE_3, SUM_INTEGER_ADDR);
-
- value1.add(5);
- value2.add(10);
- value1.add(6);
-
- StateMerging.mergeCombiningValues(Arrays.asList(value1, value2), value3);
-
- // Merging clears the old values and updates the result value.
- assertThat(value1.read(), equalTo(0));
- assertThat(value2.read(), equalTo(0));
- assertThat(value3.read(), equalTo(21));
- }
-
- @Test
- public void testWatermarkEarliestState() throws Exception {
- WatermarkHoldState value =
- underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR);
-
- // State instances are cached, but depend on the namespace.
- assertEquals(value, underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR));
- assertFalse(value.equals(underTest.state(NAMESPACE_2, WATERMARK_EARLIEST_ADDR)));
-
- assertThat(value.read(), Matchers.nullValue());
- value.add(new Instant(2000));
- assertThat(value.read(), equalTo(new Instant(2000)));
-
- value.add(new Instant(3000));
- assertThat(value.read(), equalTo(new Instant(2000)));
-
- value.add(new Instant(1000));
- assertThat(value.read(), equalTo(new Instant(1000)));
-
- value.clear();
- assertThat(value.read(), equalTo(null));
- assertThat(underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR), equalTo(value));
- }
-
- @Test
- public void testWatermarkLatestState() throws Exception {
- WatermarkHoldState value =
- underTest.state(NAMESPACE_1, WATERMARK_LATEST_ADDR);
-
- // State instances are cached, but depend on the namespace.
- assertEquals(value, underTest.state(NAMESPACE_1, WATERMARK_LATEST_ADDR));
- assertFalse(value.equals(underTest.state(NAMESPACE_2, WATERMARK_LATEST_ADDR)));
-
- assertThat(value.read(), Matchers.nullValue());
- value.add(new Instant(2000));
- assertThat(value.read(), equalTo(new Instant(2000)));
-
- value.add(new Instant(3000));
- assertThat(value.read(), equalTo(new Instant(3000)));
-
- value.add(new Instant(1000));
- assertThat(value.read(), equalTo(new Instant(3000)));
-
- value.clear();
- assertThat(value.read(), equalTo(null));
- assertThat(underTest.state(NAMESPACE_1, WATERMARK_LATEST_ADDR), equalTo(value));
- }
-
- @Test
- public void testWatermarkEndOfWindowState() throws Exception {
- WatermarkHoldState value = underTest.state(NAMESPACE_1, WATERMARK_EOW_ADDR);
-
- // State instances are cached, but depend on the namespace.
- assertEquals(value, underTest.state(NAMESPACE_1, WATERMARK_EOW_ADDR));
- assertFalse(value.equals(underTest.state(NAMESPACE_2, WATERMARK_EOW_ADDR)));
-
- assertThat(value.read(), Matchers.nullValue());
- value.add(new Instant(2000));
- assertThat(value.read(), equalTo(new Instant(2000)));
-
- value.clear();
- assertThat(value.read(), equalTo(null));
- assertThat(underTest.state(NAMESPACE_1, WATERMARK_EOW_ADDR), equalTo(value));
- }
-
- @Test
- public void testWatermarkStateIsEmpty() throws Exception {
- WatermarkHoldState value =
- underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR);
-
- assertThat(value.isEmpty().read(), Matchers.is(true));
- ReadableState<Boolean> readFuture = value.isEmpty();
- value.add(new Instant(1000));
- assertThat(readFuture.read(), Matchers.is(false));
-
- value.clear();
- assertThat(readFuture.read(), Matchers.is(true));
- }
-
- @Test
- public void testMergeEarliestWatermarkIntoSource() throws Exception {
- WatermarkHoldState value1 =
- underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR);
- WatermarkHoldState value2 =
- underTest.state(NAMESPACE_2, WATERMARK_EARLIEST_ADDR);
-
- value1.add(new Instant(3000));
- value2.add(new Instant(5000));
- value1.add(new Instant(4000));
- value2.add(new Instant(2000));
-
- // Merging clears the old values and updates the merged value.
- StateMerging.mergeWatermarks(Arrays.asList(value1, value2), value1, WINDOW_1);
-
- assertThat(value1.read(), equalTo(new Instant(2000)));
- assertThat(value2.read(), equalTo(null));
- }
-
- @Test
- public void testMergeLatestWatermarkIntoSource() throws Exception {
- WatermarkHoldState value1 =
- underTest.state(NAMESPACE_1, WATERMARK_LATEST_ADDR);
- WatermarkHoldState value2 =
- underTest.state(NAMESPACE_2, WATERMARK_LATEST_ADDR);
- WatermarkHoldState value3 =
- underTest.state(NAMESPACE_3, WATERMARK_LATEST_ADDR);
-
- value1.add(new Instant(3000));
- value2.add(new Instant(5000));
- value1.add(new Instant(4000));
- value2.add(new Instant(2000));
-
- // Merging clears the old values and updates the result value.
- StateMerging.mergeWatermarks(Arrays.asList(value1, value2), value3, WINDOW_1);
-
- // Merging clears the old values and updates the result value.
- assertThat(value3.read(), equalTo(new Instant(5000)));
- assertThat(value1.read(), equalTo(null));
- assertThat(value2.read(), equalTo(null));
- }
-
- @Test
- public void testSetReadable() throws Exception {
- SetState<String> value = underTest.state(NAMESPACE_1, STRING_SET_ADDR);
-
- // test contains
- ReadableState<Boolean> readable = value.contains("A");
- value.add("A");
- assertFalse(readable.read());
-
- // test addIfAbsent
- value.addIfAbsent("B");
- assertTrue(value.contains("B").read());
- }
-
- @Test
- public void testMapReadable() throws Exception {
- MapState<String, Integer> value = underTest.state(NAMESPACE_1, STRING_MAP_ADDR);
-
- // test iterable, should just return a iterable view of the values contained in this map.
- // The iterable is backed by the map, so changes to the map are reflected in the iterable.
- ReadableState<Iterable<String>> keys = value.keys();
- ReadableState<Iterable<Integer>> values = value.values();
- ReadableState<Iterable<Map.Entry<String, Integer>>> entries = value.entries();
- value.put("A", 1);
- assertFalse(Iterables.isEmpty(keys.read()));
- assertFalse(Iterables.isEmpty(values.read()));
- assertFalse(Iterables.isEmpty(entries.read()));
-
- // test get
- ReadableState<Integer> get = value.get("B");
- value.put("B", 2);
- assertNull(get.read());
-
- // test addIfAbsent
- value.putIfAbsent("C", 3);
- assertThat(value.get("C").read(), equalTo(3));
- }
-
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/test/java/org/apache/beam/runners/core/WindowMatchers.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/WindowMatchers.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/WindowMatchers.java
index 26cbfee..9769d10 100644
--- a/runners/core-java/src/test/java/org/apache/beam/runners/core/WindowMatchers.java
+++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/WindowMatchers.java
@@ -116,21 +116,6 @@ public class WindowMatchers {
}
public static <T> Matcher<WindowedValue<? extends T>> isSingleWindowedValue(
- Matcher<T> valueMatcher,
- long timestamp,
- long windowStart,
- long windowEnd,
- PaneInfo paneInfo) {
- IntervalWindow intervalWindow =
- new IntervalWindow(new Instant(windowStart), new Instant(windowEnd));
- return WindowMatchers.<T>isSingleWindowedValue(
- valueMatcher,
- Matchers.describedAs("%0", Matchers.equalTo(new Instant(timestamp)), timestamp),
- Matchers.<BoundedWindow>equalTo(intervalWindow),
- Matchers.equalTo(paneInfo));
- }
-
- public static <T> Matcher<WindowedValue<? extends T>> isSingleWindowedValue(
Matcher<? super T> valueMatcher,
Matcher<? super Instant> timestampMatcher,
Matcher<? super BoundedWindow> windowMatcher) {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/AfterFirstStateMachineTest.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/AfterFirstStateMachineTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/AfterFirstStateMachineTest.java
index 2be90de..453c8ff 100644
--- a/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/AfterFirstStateMachineTest.java
+++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/AfterFirstStateMachineTest.java
@@ -21,6 +21,7 @@ import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.when;
+import org.apache.beam.runners.core.triggers.TriggerStateMachine.OnceTriggerStateMachine;
import org.apache.beam.runners.core.triggers.TriggerStateMachineTester.SimpleTriggerStateMachineTester;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
@@ -41,8 +42,8 @@ import org.mockito.MockitoAnnotations;
@RunWith(JUnit4.class)
public class AfterFirstStateMachineTest {
- @Mock private TriggerStateMachine mockTrigger1;
- @Mock private TriggerStateMachine mockTrigger2;
+ @Mock private OnceTriggerStateMachine mockTrigger1;
+ @Mock private OnceTriggerStateMachine mockTrigger2;
private SimpleTriggerStateMachineTester<IntervalWindow> tester;
private static TriggerStateMachine.TriggerContext anyTriggerContext() {
return Mockito.<TriggerStateMachine.TriggerContext>any();
[11/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BoundedSourceRunnerTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BoundedSourceRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BoundedSourceRunnerTest.java
index 6c9a4cb..d8ed121 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BoundedSourceRunnerTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/runners/core/BoundedSourceRunnerTest.java
@@ -20,35 +20,25 @@ package org.apache.beam.runners.core;
import static org.apache.beam.sdk.util.WindowedValue.valueInGlobalWindow;
import static org.hamcrest.Matchers.contains;
-import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.collection.IsEmptyCollection.empty;
import static org.junit.Assert.assertThat;
-import static org.junit.Assert.fail;
-import com.google.common.base.Suppliers;
-import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Multimap;
import com.google.protobuf.Any;
import com.google.protobuf.ByteString;
import com.google.protobuf.BytesValue;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
-import java.util.ServiceLoader;
+import java.util.Map;
import org.apache.beam.fn.harness.fn.ThrowingConsumer;
-import org.apache.beam.fn.harness.fn.ThrowingRunnable;
-import org.apache.beam.runners.core.PTransformRunnerFactory.Registrar;
import org.apache.beam.sdk.common.runner.v1.RunnerApi;
import org.apache.beam.sdk.io.BoundedSource;
import org.apache.beam.sdk.io.CountingSource;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.util.WindowedValue;
-import org.hamcrest.Matchers;
-import org.hamcrest.collection.IsMapContaining;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
@@ -56,25 +46,27 @@ import org.junit.runners.JUnit4;
/** Tests for {@link BoundedSourceRunner}. */
@RunWith(JUnit4.class)
public class BoundedSourceRunnerTest {
-
- public static final String URN = "urn:org.apache.beam:source:java:0.1";
-
@Test
public void testRunReadLoopWithMultipleSources() throws Exception {
- List<WindowedValue<Long>> out1Values = new ArrayList<>();
+ List<WindowedValue<Long>> out1ValuesA = new ArrayList<>();
+ List<WindowedValue<Long>> out1ValuesB = new ArrayList<>();
List<WindowedValue<Long>> out2Values = new ArrayList<>();
- Collection<ThrowingConsumer<WindowedValue<Long>>> consumers =
- ImmutableList.of(out1Values::add, out2Values::add);
+ Map<String, Collection<ThrowingConsumer<WindowedValue<Long>>>> outputMap = ImmutableMap.of(
+ "out1", ImmutableList.of(out1ValuesA::add, out1ValuesB::add),
+ "out2", ImmutableList.of(out2Values::add));
- BoundedSourceRunner<BoundedSource<Long>, Long> runner = new BoundedSourceRunner<>(
+ BoundedSourceRunner<BoundedSource<Long>, Long> runner =
+ new BoundedSourceRunner<>(
PipelineOptionsFactory.create(),
RunnerApi.FunctionSpec.getDefaultInstance(),
- consumers);
+ outputMap);
runner.runReadLoop(valueInGlobalWindow(CountingSource.upTo(2)));
runner.runReadLoop(valueInGlobalWindow(CountingSource.upTo(1)));
- assertThat(out1Values,
+ assertThat(out1ValuesA,
+ contains(valueInGlobalWindow(0L), valueInGlobalWindow(1L), valueInGlobalWindow(0L)));
+ assertThat(out1ValuesB,
contains(valueInGlobalWindow(0L), valueInGlobalWindow(1L), valueInGlobalWindow(0L)));
assertThat(out2Values,
contains(valueInGlobalWindow(0L), valueInGlobalWindow(1L), valueInGlobalWindow(0L)));
@@ -82,106 +74,40 @@ public class BoundedSourceRunnerTest {
@Test
public void testRunReadLoopWithEmptySource() throws Exception {
- List<WindowedValue<Long>> outValues = new ArrayList<>();
- Collection<ThrowingConsumer<WindowedValue<Long>>> consumers =
- ImmutableList.of(outValues::add);
+ List<WindowedValue<Long>> out1Values = new ArrayList<>();
+ Map<String, Collection<ThrowingConsumer<WindowedValue<Long>>>> outputMap = ImmutableMap.of(
+ "out1", ImmutableList.of(out1Values::add));
- BoundedSourceRunner<BoundedSource<Long>, Long> runner = new BoundedSourceRunner<>(
+ BoundedSourceRunner<BoundedSource<Long>, Long> runner =
+ new BoundedSourceRunner<>(
PipelineOptionsFactory.create(),
RunnerApi.FunctionSpec.getDefaultInstance(),
- consumers);
+ outputMap);
runner.runReadLoop(valueInGlobalWindow(CountingSource.upTo(0)));
- assertThat(outValues, empty());
+ assertThat(out1Values, empty());
}
@Test
public void testStart() throws Exception {
List<WindowedValue<Long>> outValues = new ArrayList<>();
- Collection<ThrowingConsumer<WindowedValue<Long>>> consumers =
- ImmutableList.of(outValues::add);
+ Map<String, Collection<ThrowingConsumer<WindowedValue<Long>>>> outputMap = ImmutableMap.of(
+ "out", ImmutableList.of(outValues::add));
ByteString encodedSource =
ByteString.copyFrom(SerializableUtils.serializeToByteArray(CountingSource.upTo(3)));
- BoundedSourceRunner<BoundedSource<Long>, Long> runner = new BoundedSourceRunner<>(
+ BoundedSourceRunner<BoundedSource<Long>, Long> runner =
+ new BoundedSourceRunner<>(
PipelineOptionsFactory.create(),
- RunnerApi.FunctionSpec.newBuilder().setParameter(
+ RunnerApi.FunctionSpec.newBuilder().setParameter(
Any.pack(BytesValue.newBuilder().setValue(encodedSource).build())).build(),
- consumers);
+ outputMap);
runner.start();
assertThat(outValues,
contains(valueInGlobalWindow(0L), valueInGlobalWindow(1L), valueInGlobalWindow(2L)));
}
-
- @Test
- public void testCreatingAndProcessingSourceFromFactory() throws Exception {
- List<WindowedValue<String>> outputValues = new ArrayList<>();
-
- Multimap<String, ThrowingConsumer<WindowedValue<?>>> consumers = HashMultimap.create();
- consumers.put("outputPC",
- (ThrowingConsumer) (ThrowingConsumer<WindowedValue<String>>) outputValues::add);
- List<ThrowingRunnable> startFunctions = new ArrayList<>();
- List<ThrowingRunnable> finishFunctions = new ArrayList<>();
-
- RunnerApi.FunctionSpec functionSpec = RunnerApi.FunctionSpec.newBuilder()
- .setUrn("urn:org.apache.beam:source:java:0.1")
- .setParameter(Any.pack(BytesValue.newBuilder()
- .setValue(ByteString.copyFrom(
- SerializableUtils.serializeToByteArray(CountingSource.upTo(3))))
- .build()))
- .build();
-
- RunnerApi.PTransform pTransform = RunnerApi.PTransform.newBuilder()
- .setSpec(functionSpec)
- .putInputs("input", "inputPC")
- .putOutputs("output", "outputPC")
- .build();
-
- new BoundedSourceRunner.Factory<>().createRunnerForPTransform(
- PipelineOptionsFactory.create(),
- null /* beamFnDataClient */,
- "pTransformId",
- pTransform,
- Suppliers.ofInstance("57L")::get,
- ImmutableMap.of(),
- ImmutableMap.of(),
- consumers,
- startFunctions::add,
- finishFunctions::add);
-
- // This is testing a deprecated way of running sources and should be removed
- // once all source definitions are instead propagated along the input edge.
- Iterables.getOnlyElement(startFunctions).run();
- assertThat(outputValues, contains(
- valueInGlobalWindow(0L),
- valueInGlobalWindow(1L),
- valueInGlobalWindow(2L)));
- outputValues.clear();
-
- // Check that when passing a source along as an input, the source is processed.
- assertThat(consumers.keySet(), containsInAnyOrder("inputPC", "outputPC"));
- Iterables.getOnlyElement(consumers.get("inputPC")).accept(
- valueInGlobalWindow(CountingSource.upTo(2)));
- assertThat(outputValues, contains(
- valueInGlobalWindow(0L),
- valueInGlobalWindow(1L)));
-
- assertThat(finishFunctions, Matchers.empty());
- }
-
- @Test
- public void testRegistration() {
- for (Registrar registrar :
- ServiceLoader.load(Registrar.class)) {
- if (registrar instanceof BoundedSourceRunner.Registrar) {
- assertThat(registrar.getPTransformRunnerFactories(), IsMapContaining.hasKey(URN));
- return;
- }
- }
- fail("Expected registrar not found.");
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/FnApiDoFnRunnerTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/FnApiDoFnRunnerTest.java b/sdks/java/harness/src/test/java/org/apache/beam/runners/core/FnApiDoFnRunnerTest.java
deleted file mode 100644
index c4df77a..0000000
--- a/sdks/java/harness/src/test/java/org/apache/beam/runners/core/FnApiDoFnRunnerTest.java
+++ /dev/null
@@ -1,210 +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.apache.beam.sdk.util.WindowedValue.timestampedValueInGlobalWindow;
-import static org.apache.beam.sdk.util.WindowedValue.valueInGlobalWindow;
-import static org.hamcrest.Matchers.contains;
-import static org.hamcrest.Matchers.containsInAnyOrder;
-import static org.junit.Assert.assertThat;
-import static org.junit.Assert.fail;
-
-import com.fasterxml.jackson.databind.ObjectMapper;
-import com.google.common.base.Suppliers;
-import com.google.common.collect.HashMultimap;
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Multimap;
-import com.google.protobuf.Any;
-import com.google.protobuf.ByteString;
-import com.google.protobuf.BytesValue;
-import com.google.protobuf.Message;
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-import java.util.ServiceLoader;
-import org.apache.beam.fn.harness.fn.ThrowingConsumer;
-import org.apache.beam.fn.harness.fn.ThrowingRunnable;
-import org.apache.beam.runners.core.PTransformRunnerFactory.Registrar;
-import org.apache.beam.runners.core.construction.ParDoTranslation;
-import org.apache.beam.runners.dataflow.util.CloudObjects;
-import org.apache.beam.runners.dataflow.util.DoFnInfo;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.sdk.common.runner.v1.RunnerApi;
-import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
-import org.apache.beam.sdk.util.SerializableUtils;
-import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.sdk.values.WindowingStrategy;
-import org.hamcrest.collection.IsMapContaining;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-/** Tests for {@link FnApiDoFnRunner}. */
-@RunWith(JUnit4.class)
-public class FnApiDoFnRunnerTest {
-
- private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
- private static final Coder<WindowedValue<String>> STRING_CODER =
- WindowedValue.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE);
- private static final String STRING_CODER_SPEC_ID = "999L";
- private static final RunnerApi.Coder STRING_CODER_SPEC;
-
- static {
- try {
- STRING_CODER_SPEC = RunnerApi.Coder.newBuilder()
- .setSpec(RunnerApi.SdkFunctionSpec.newBuilder()
- .setSpec(RunnerApi.FunctionSpec.newBuilder()
- .setParameter(Any.pack(BytesValue.newBuilder().setValue(ByteString.copyFrom(
- OBJECT_MAPPER.writeValueAsBytes(CloudObjects.asCloudObject(STRING_CODER))))
- .build())))
- .build())
- .build();
- } catch (IOException e) {
- throw new ExceptionInInitializerError(e);
- }
- }
-
- private static class TestDoFn extends DoFn<String, String> {
- private static final TupleTag<String> mainOutput = new TupleTag<>("mainOutput");
- private static final TupleTag<String> additionalOutput = new TupleTag<>("output");
-
- private BoundedWindow window;
-
- @ProcessElement
- public void processElement(ProcessContext context, BoundedWindow window) {
- context.output("MainOutput" + context.element());
- context.output(additionalOutput, "AdditionalOutput" + context.element());
- this.window = window;
- }
-
- @FinishBundle
- public void finishBundle(FinishBundleContext context) {
- if (window != null) {
- context.output("FinishBundle", window.maxTimestamp(), window);
- window = null;
- }
- }
- }
-
- /**
- * Create a DoFn that has 3 inputs (inputATarget1, inputATarget2, inputBTarget) and 2 outputs
- * (mainOutput, output). Validate that inputs are fed to the {@link DoFn} and that outputs
- * are directed to the correct consumers.
- */
- @Test
- public void testCreatingAndProcessingDoFn() throws Exception {
- Map<String, Message> fnApiRegistry = ImmutableMap.of(STRING_CODER_SPEC_ID, STRING_CODER_SPEC);
- String pTransformId = "pTransformId";
- String mainOutputId = "101";
- String additionalOutputId = "102";
-
- DoFnInfo<?, ?> doFnInfo = DoFnInfo.forFn(
- new TestDoFn(),
- WindowingStrategy.globalDefault(),
- ImmutableList.of(),
- StringUtf8Coder.of(),
- Long.parseLong(mainOutputId),
- ImmutableMap.of(
- Long.parseLong(mainOutputId), TestDoFn.mainOutput,
- Long.parseLong(additionalOutputId), TestDoFn.additionalOutput));
- RunnerApi.FunctionSpec functionSpec = RunnerApi.FunctionSpec.newBuilder()
- .setUrn(ParDoTranslation.CUSTOM_JAVA_DO_FN_URN)
- .setParameter(Any.pack(BytesValue.newBuilder()
- .setValue(ByteString.copyFrom(SerializableUtils.serializeToByteArray(doFnInfo)))
- .build()))
- .build();
- RunnerApi.PTransform pTransform = RunnerApi.PTransform.newBuilder()
- .setSpec(functionSpec)
- .putInputs("inputA", "inputATarget")
- .putInputs("inputB", "inputBTarget")
- .putOutputs(mainOutputId, "mainOutputTarget")
- .putOutputs(additionalOutputId, "additionalOutputTarget")
- .build();
-
- List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
- List<WindowedValue<String>> additionalOutputValues = new ArrayList<>();
- Multimap<String, ThrowingConsumer<WindowedValue<?>>> consumers = HashMultimap.create();
- consumers.put("mainOutputTarget",
- (ThrowingConsumer) (ThrowingConsumer<WindowedValue<String>>) mainOutputValues::add);
- consumers.put("additionalOutputTarget",
- (ThrowingConsumer) (ThrowingConsumer<WindowedValue<String>>) additionalOutputValues::add);
- List<ThrowingRunnable> startFunctions = new ArrayList<>();
- List<ThrowingRunnable> finishFunctions = new ArrayList<>();
-
- new FnApiDoFnRunner.Factory<>().createRunnerForPTransform(
- PipelineOptionsFactory.create(),
- null /* beamFnDataClient */,
- pTransformId,
- pTransform,
- Suppliers.ofInstance("57L")::get,
- ImmutableMap.of(),
- ImmutableMap.of(),
- consumers,
- startFunctions::add,
- finishFunctions::add);
-
- Iterables.getOnlyElement(startFunctions).run();
- mainOutputValues.clear();
-
- assertThat(consumers.keySet(), containsInAnyOrder(
- "inputATarget", "inputBTarget", "mainOutputTarget", "additionalOutputTarget"));
-
- Iterables.getOnlyElement(consumers.get("inputATarget")).accept(valueInGlobalWindow("A1"));
- Iterables.getOnlyElement(consumers.get("inputATarget")).accept(valueInGlobalWindow("A2"));
- Iterables.getOnlyElement(consumers.get("inputATarget")).accept(valueInGlobalWindow("B"));
- assertThat(mainOutputValues, contains(
- valueInGlobalWindow("MainOutputA1"),
- valueInGlobalWindow("MainOutputA2"),
- valueInGlobalWindow("MainOutputB")));
- assertThat(additionalOutputValues, contains(
- valueInGlobalWindow("AdditionalOutputA1"),
- valueInGlobalWindow("AdditionalOutputA2"),
- valueInGlobalWindow("AdditionalOutputB")));
- mainOutputValues.clear();
- additionalOutputValues.clear();
-
- Iterables.getOnlyElement(finishFunctions).run();
- assertThat(
- mainOutputValues,
- contains(
- timestampedValueInGlobalWindow("FinishBundle", GlobalWindow.INSTANCE.maxTimestamp())));
- mainOutputValues.clear();
- }
-
- @Test
- public void testRegistration() {
- for (Registrar registrar :
- ServiceLoader.load(Registrar.class)) {
- if (registrar instanceof FnApiDoFnRunner.Registrar) {
- assertThat(registrar.getPTransformRunnerFactories(),
- IsMapContaining.hasKey(ParDoTranslation.CUSTOM_JAVA_DO_FN_URN));
- return;
- }
- }
- fail("Expected registrar not found.");
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/amqp/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/amqp/pom.xml b/sdks/java/io/amqp/pom.xml
deleted file mode 100644
index 8da9448..0000000
--- a/sdks/java/io/amqp/pom.xml
+++ /dev/null
@@ -1,100 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<!--
- 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.
--->
-<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
-
- <modelVersion>4.0.0</modelVersion>
-
- <parent>
- <groupId>org.apache.beam</groupId>
- <artifactId>beam-sdks-java-io-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
- <relativePath>../pom.xml</relativePath>
- </parent>
-
- <artifactId>beam-sdks-java-io-amqp</artifactId>
- <name>Apache Beam :: SDKs :: Java :: IO :: AMQP</name>
- <description>IO to read and write using AMQP 1.0 protocol (http://www.amqp.org).</description>
-
- <dependencies>
- <dependency>
- <groupId>org.apache.beam</groupId>
- <artifactId>beam-sdks-java-core</artifactId>
- </dependency>
-
- <dependency>
- <groupId>org.slf4j</groupId>
- <artifactId>slf4j-api</artifactId>
- </dependency>
-
- <dependency>
- <groupId>joda-time</groupId>
- <artifactId>joda-time</artifactId>
- </dependency>
-
- <dependency>
- <groupId>com.google.guava</groupId>
- <artifactId>guava</artifactId>
- </dependency>
-
- <dependency>
- <groupId>com.google.code.findbugs</groupId>
- <artifactId>jsr305</artifactId>
- </dependency>
-
- <dependency>
- <groupId>org.apache.qpid</groupId>
- <artifactId>proton-j</artifactId>
- <version>0.13.1</version>
- </dependency>
-
- <!-- compile dependencies -->
- <dependency>
- <groupId>com.google.auto.value</groupId>
- <artifactId>auto-value</artifactId>
- <scope>provided</scope>
- </dependency>
- <dependency>
- <groupId>com.google.auto.service</groupId>
- <artifactId>auto-service</artifactId>
- <optional>true</optional>
- </dependency>
-
- <!-- test dependencies -->
- <dependency>
- <groupId>org.slf4j</groupId>
- <artifactId>slf4j-jdk14</artifactId>
- <scope>test</scope>
- </dependency>
- <dependency>
- <groupId>junit</groupId>
- <artifactId>junit</artifactId>
- <scope>test</scope>
- </dependency>
- <dependency>
- <groupId>org.hamcrest</groupId>
- <artifactId>hamcrest-all</artifactId>
- <scope>test</scope>
- </dependency>
- <dependency>
- <groupId>org.apache.beam</groupId>
- <artifactId>beam-runners-direct-java</artifactId>
- <scope>test</scope>
- </dependency>
- </dependencies>
-
-</project>
\ No newline at end of file
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpIO.java b/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpIO.java
deleted file mode 100644
index 1f307b2..0000000
--- a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpIO.java
+++ /dev/null
@@ -1,399 +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.sdk.io.amqp;
-
-import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Preconditions.checkState;
-
-import com.google.auto.value.AutoValue;
-import com.google.common.base.Joiner;
-
-import java.io.IOException;
-import java.io.Serializable;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.NoSuchElementException;
-
-import javax.annotation.Nullable;
-
-import org.apache.beam.sdk.annotations.Experimental;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.SerializableCoder;
-import org.apache.beam.sdk.io.UnboundedSource;
-import org.apache.beam.sdk.options.PipelineOptions;
-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.display.DisplayData;
-import org.apache.beam.sdk.values.PBegin;
-import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PDone;
-import org.apache.qpid.proton.message.Message;
-import org.apache.qpid.proton.messenger.Messenger;
-import org.apache.qpid.proton.messenger.Tracker;
-import org.joda.time.Duration;
-import org.joda.time.Instant;
-
-/**
- * AmqpIO supports AMQP 1.0 protocol using the Apache QPid Proton-J library.
- *
- * <p>It's also possible to use AMQP 1.0 protocol via Apache Qpid JMS connection factory and the
- * Apache Beam JmsIO.
- *
- * <h3>Binding AMQP and receive messages</h3>
- *
- * <p>The {@link AmqpIO} {@link Read} can bind a AMQP listener endpoint and receive messages. It can
- * also connect to a AMPQ broker (such as Apache Qpid or Apache ActiveMQ).
- *
- * <p>{@link AmqpIO} {@link Read} returns an unbounded {@link PCollection} of {@link Message}
- * containing the received messages.
- *
- * <p>To configure a AMQP source, you have to provide a list of addresses where it will receive
- * messages. An address has the following form: {@code
- * [amqp[s]://][user[:password]@]domain[/[name]]} where {@code domain} can be one of {@code
- * host | host:port | ip | ip:port | name}. NB: the {@code ~} character allows to bind a AMQP
- * listener instead of connecting to a remote broker. For instance {@code amqp://~0.0.0.0:1234}
- * will bind a AMQP listener on any network interface on the 1234 port number.
- *
- * <p>The following example illustrates how to configure a AMQP source:
- *
- * <pre>{@code
- *
- * pipeline.apply(AmqpIO.read()
- * .withAddresses(Collections.singletonList("amqp://host:1234")))
- *
- * }</pre>
- *
- * <h3>Sending messages to a AMQP endpoint</h3>
- *
- * <p>{@link AmqpIO} provides a sink to send {@link PCollection} elements as messages.
- *
- * <p>As for the {@link Read}, {@link AmqpIO} {@link Write} requires a list of addresses where to
- * send messages. The following example illustrates how to configure the {@link AmqpIO}
- * {@link Write}:
- *
- * <pre>{@code
- *
- * pipeline
- * .apply(...) // provide PCollection<Message>
- * .apply(AmqpIO.write());
- *
- * }</pre>
- */
-@Experimental(Experimental.Kind.SOURCE_SINK)
-public class AmqpIO {
-
- public static Read read() {
- return new AutoValue_AmqpIO_Read.Builder().setMaxNumRecords(Long.MAX_VALUE).build();
- }
-
- public static Write write() {
- return new AutoValue_AmqpIO_Write();
- }
-
- private AmqpIO() {
- }
-
- /**
- * A {@link PTransform} to read/receive messages using AMQP 1.0 protocol.
- */
- @AutoValue
- public abstract static class Read extends PTransform<PBegin, PCollection<Message>> {
-
- @Nullable abstract List<String> addresses();
- abstract long maxNumRecords();
- @Nullable abstract Duration maxReadTime();
-
- abstract Builder builder();
-
- @AutoValue.Builder
- abstract static class Builder {
- abstract Builder setAddresses(List<String> addresses);
- abstract Builder setMaxNumRecords(long maxNumRecords);
- abstract Builder setMaxReadTime(Duration maxReadTime);
- abstract Read build();
- }
-
- /**
- * Define the AMQP addresses where to receive messages.
- */
- public Read withAddresses(List<String> addresses) {
- checkArgument(addresses != null, "AmqpIO.read().withAddresses(addresses) called with null"
- + " addresses");
- checkArgument(!addresses.isEmpty(), "AmqpIO.read().withAddresses(addresses) called with "
- + "empty addresses list");
- return builder().setAddresses(addresses).build();
- }
-
- /**
- * Define the max number of records received by the {@link Read}.
- * When the max number of records is lower than {@code Long.MAX_VALUE}, the {@link Read} will
- * provide a bounded {@link PCollection}.
- */
- public Read withMaxNumRecords(long maxNumRecords) {
- checkArgument(maxReadTime() == null,
- "maxNumRecord and maxReadTime are exclusive");
- return builder().setMaxNumRecords(maxNumRecords).build();
- }
-
- /**
- * Define the max read time (duration) while the {@link Read} will receive messages.
- * When this max read time is not null, the {@link Read} will provide a bounded
- * {@link PCollection}.
- */
- public Read withMaxReadTime(Duration maxReadTime) {
- checkArgument(maxNumRecords() == Long.MAX_VALUE,
- "maxNumRecord and maxReadTime are exclusive");
- return builder().setMaxReadTime(maxReadTime).build();
- }
-
- @Override
- public void validate(PipelineOptions pipelineOptions) {
- checkState(addresses() != null, "AmqIO.read() requires addresses list to be set via "
- + "withAddresses(addresses)");
- checkState(!addresses().isEmpty(), "AmqIO.read() requires a non-empty addresses list to be"
- + " set via withAddresses(addresses)");
- }
-
- @Override
- public void populateDisplayData(DisplayData.Builder builder) {
- builder.add(DisplayData.item("addresses", Joiner.on(" ").join(addresses())));
- }
-
- @Override
- public PCollection<Message> expand(PBegin input) {
- org.apache.beam.sdk.io.Read.Unbounded<Message> unbounded =
- org.apache.beam.sdk.io.Read.from(new UnboundedAmqpSource(this));
-
- PTransform<PBegin, PCollection<Message>> transform = unbounded;
-
- if (maxNumRecords() != Long.MAX_VALUE) {
- transform = unbounded.withMaxNumRecords(maxNumRecords());
- } else if (maxReadTime() != null) {
- transform = unbounded.withMaxReadTime(maxReadTime());
- }
-
- return input.getPipeline().apply(transform);
- }
-
- }
-
- private static class AmqpCheckpointMark implements UnboundedSource.CheckpointMark, Serializable {
-
- private transient Messenger messenger;
- private transient List<Tracker> trackers = new ArrayList<>();
-
- public AmqpCheckpointMark() {
- }
-
- @Override
- public void finalizeCheckpoint() {
- for (Tracker tracker : trackers) {
- // flag as not cumulative
- messenger.accept(tracker, 0);
- }
- trackers.clear();
- }
-
- // set an empty list to messages when deserialize
- private void readObject(java.io.ObjectInputStream stream)
- throws java.io.IOException, ClassNotFoundException {
- trackers = new ArrayList<>();
- }
-
- }
-
- private static class UnboundedAmqpSource
- extends UnboundedSource<Message, AmqpCheckpointMark> {
-
- private final Read spec;
-
- public UnboundedAmqpSource(Read spec) {
- this.spec = spec;
- }
-
- @Override
- public List<UnboundedAmqpSource> split(int desiredNumSplits,
- PipelineOptions pipelineOptions) {
- // amqp is a queue system, so, it's possible to have multiple concurrent sources, even if
- // they bind the listener
- List<UnboundedAmqpSource> sources = new ArrayList<>();
- for (int i = 0; i < Math.max(1, desiredNumSplits); ++i) {
- sources.add(new UnboundedAmqpSource(spec));
- }
- return sources;
- }
-
- @Override
- public UnboundedReader<Message> createReader(PipelineOptions pipelineOptions,
- AmqpCheckpointMark checkpointMark) {
- return new UnboundedAmqpReader(this, checkpointMark);
- }
-
- @Override
- public Coder<Message> getDefaultOutputCoder() {
- return new AmqpMessageCoder();
- }
-
- @Override
- public Coder<AmqpCheckpointMark> getCheckpointMarkCoder() {
- return SerializableCoder.of(AmqpCheckpointMark.class);
- }
-
- @Override
- public void validate() {
- spec.validate(null);
- }
-
- }
-
- private static class UnboundedAmqpReader extends UnboundedSource.UnboundedReader<Message> {
-
- private final UnboundedAmqpSource source;
-
- private Messenger messenger;
- private Message current;
- private Instant currentTimestamp;
- private Instant watermark = new Instant(Long.MIN_VALUE);
- private AmqpCheckpointMark checkpointMark;
-
- public UnboundedAmqpReader(UnboundedAmqpSource source, AmqpCheckpointMark checkpointMark) {
- this.source = source;
- this.current = null;
- if (checkpointMark != null) {
- this.checkpointMark = checkpointMark;
- } else {
- this.checkpointMark = new AmqpCheckpointMark();
- }
- }
-
- @Override
- public Instant getWatermark() {
- return watermark;
- }
-
- @Override
- public Instant getCurrentTimestamp() {
- if (current == null) {
- throw new NoSuchElementException();
- }
- return currentTimestamp;
- }
-
- @Override
- public Message getCurrent() {
- if (current == null) {
- throw new NoSuchElementException();
- }
- return current;
- }
-
- @Override
- public UnboundedSource.CheckpointMark getCheckpointMark() {
- return checkpointMark;
- }
-
- @Override
- public UnboundedAmqpSource getCurrentSource() {
- return source;
- }
-
- @Override
- public boolean start() throws IOException {
- Read spec = source.spec;
- messenger = Messenger.Factory.create();
- messenger.start();
- for (String address : spec.addresses()) {
- messenger.subscribe(address);
- }
- checkpointMark.messenger = messenger;
- return advance();
- }
-
- @Override
- public boolean advance() {
- messenger.recv();
- if (messenger.incoming() <= 0) {
- current = null;
- return false;
- }
- Message message = messenger.get();
- Tracker tracker = messenger.incomingTracker();
- checkpointMark.trackers.add(tracker);
- currentTimestamp = new Instant(message.getCreationTime());
- watermark = currentTimestamp;
- current = message;
- return true;
- }
-
- @Override
- public void close() {
- if (messenger != null) {
- messenger.stop();
- }
- }
-
- }
-
- /**
- * A {@link PTransform} to send messages using AMQP 1.0 protocol.
- */
- @AutoValue
- public abstract static class Write extends PTransform<PCollection<Message>, PDone> {
-
- @Override
- public PDone expand(PCollection<Message> input) {
- input.apply(ParDo.of(new WriteFn(this)));
- return PDone.in(input.getPipeline());
- }
-
- private static class WriteFn extends DoFn<Message, Void> {
-
- private final Write spec;
-
- private transient Messenger messenger;
-
- public WriteFn(Write spec) {
- this.spec = spec;
- }
-
- @Setup
- public void setup() throws Exception {
- messenger = Messenger.Factory.create();
- messenger.start();
- }
-
- @ProcessElement
- public void processElement(ProcessContext processContext) throws Exception {
- Message message = processContext.element();
- messenger.put(message);
- messenger.send();
- }
-
- @Teardown
- public void teardown() throws Exception {
- if (messenger != null) {
- messenger.stop();
- }
- }
-
- }
-
- }
-
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoder.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoder.java b/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoder.java
deleted file mode 100644
index 5a55260..0000000
--- a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoder.java
+++ /dev/null
@@ -1,79 +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.sdk.io.amqp;
-
-import com.google.common.io.ByteStreams;
-
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
-import java.nio.BufferOverflowException;
-
-import org.apache.beam.sdk.coders.CoderException;
-import org.apache.beam.sdk.coders.CustomCoder;
-import org.apache.beam.sdk.util.VarInt;
-import org.apache.qpid.proton.message.Message;
-
-/**
- * A coder for AMQP message.
- */
-public class AmqpMessageCoder extends CustomCoder<Message> {
-
- private static final int[] MESSAGE_SIZES = new int[]{
- 8 * 1024,
- 64 * 1024,
- 1 * 1024 * 1024,
- 64 * 1024 * 1024
- };
-
- static AmqpMessageCoder of() {
- return new AmqpMessageCoder();
- }
-
- @Override
- public void encode(Message value, OutputStream outStream) throws CoderException, IOException {
- for (int maxMessageSize : MESSAGE_SIZES) {
- try {
- encode(value, outStream, maxMessageSize);
- return;
- } catch (Exception e) {
- continue;
- }
- }
- throw new CoderException("Message is larger than the max size supported by the coder");
- }
-
- private void encode(Message value, OutputStream outStream, int messageSize) throws
- IOException, BufferOverflowException {
- byte[] data = new byte[messageSize];
- int bytesWritten = value.encode(data, 0, data.length);
- VarInt.encode(bytesWritten, outStream);
- outStream.write(data, 0, bytesWritten);
- }
-
- @Override
- public Message decode(InputStream inStream) throws CoderException, IOException {
- Message message = Message.Factory.create();
- int bytesToRead = VarInt.decodeInt(inStream);
- byte[] encodedMessage = new byte[bytesToRead];
- ByteStreams.readFully(inStream, encodedMessage);
- message.decode(encodedMessage, 0, encodedMessage.length);
- return message;
- }
-
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoderProviderRegistrar.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoderProviderRegistrar.java b/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoderProviderRegistrar.java
deleted file mode 100644
index bc3445c..0000000
--- a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoderProviderRegistrar.java
+++ /dev/null
@@ -1,44 +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.sdk.io.amqp;
-
-import com.google.auto.service.AutoService;
-import com.google.common.collect.ImmutableList;
-
-import java.util.List;
-
-import org.apache.beam.sdk.coders.CoderProvider;
-import org.apache.beam.sdk.coders.CoderProviderRegistrar;
-import org.apache.beam.sdk.coders.CoderProviders;
-import org.apache.beam.sdk.values.TypeDescriptor;
-import org.apache.qpid.proton.message.Message;
-
-/**
- * A {@link CoderProviderRegistrar} for standard types used with {@link AmqpIO}.
- */
-@AutoService(CoderProviderRegistrar.class)
-public class AmqpMessageCoderProviderRegistrar implements CoderProviderRegistrar {
-
- @Override
- public List<CoderProvider> getCoderProviders() {
- return ImmutableList.of(
- CoderProviders.forCoder(TypeDescriptor.of(Message.class),
- AmqpMessageCoder.of()));
- }
-
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/package-info.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/package-info.java b/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/package-info.java
deleted file mode 100644
index 091f234..0000000
--- a/sdks/java/io/amqp/src/main/java/org/apache/beam/sdk/io/amqp/package-info.java
+++ /dev/null
@@ -1,22 +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.
- */
-
-/**
- * Transforms for reading and writing using AMQP 1.0 protocol.
- */
-package org.apache.beam.sdk.io.amqp;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/amqp/src/test/java/org/apache/beam/sdk/io/amqp/AmqpIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/amqp/src/test/java/org/apache/beam/sdk/io/amqp/AmqpIOTest.java b/sdks/java/io/amqp/src/test/java/org/apache/beam/sdk/io/amqp/AmqpIOTest.java
deleted file mode 100644
index c8fe4e8..0000000
--- a/sdks/java/io/amqp/src/test/java/org/apache/beam/sdk/io/amqp/AmqpIOTest.java
+++ /dev/null
@@ -1,148 +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.sdk.io.amqp;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
-
-import java.net.ServerSocket;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
-
-import org.apache.beam.sdk.testing.PAssert;
-import org.apache.beam.sdk.testing.TestPipeline;
-import org.apache.beam.sdk.transforms.Count;
-import org.apache.beam.sdk.transforms.Create;
-import org.apache.beam.sdk.values.PCollection;
-import org.apache.qpid.proton.amqp.messaging.AmqpValue;
-import org.apache.qpid.proton.message.Message;
-import org.apache.qpid.proton.messenger.Messenger;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/**
- * Tests on {@link AmqpIO}.
- */
-@RunWith(JUnit4.class)
-public class AmqpIOTest {
-
- private static final Logger LOG = LoggerFactory.getLogger(AmqpIOTest.class);
-
- private int port;
-
- @Rule public TestPipeline pipeline = TestPipeline.create();
-
- @Before
- public void findFreeNetworkPort() throws Exception {
- LOG.info("Finding free network port");
- ServerSocket socket = new ServerSocket(0);
- port = socket.getLocalPort();
- socket.close();
- }
-
- @Test
- public void testRead() throws Exception {
- PCollection<Message> output = pipeline.apply(AmqpIO.read()
- .withMaxNumRecords(100)
- .withAddresses(Collections.singletonList("amqp://~localhost:" + port)));
- PAssert.thatSingleton(output.apply(Count.<Message>globally())).isEqualTo(100L);
-
- Thread sender = new Thread() {
- public void run() {
- try {
- Thread.sleep(500);
- Messenger sender = Messenger.Factory.create();
- sender.start();
- for (int i = 0; i < 100; i++) {
- Message message = Message.Factory.create();
- message.setAddress("amqp://localhost:" + port);
- message.setBody(new AmqpValue("Test " + i));
- sender.put(message);
- sender.send();
- }
- sender.stop();
- } catch (Exception e) {
- LOG.error("Sender error", e);
- }
- }
- };
- try {
- sender.start();
- pipeline.run();
- } finally {
- sender.join();
- }
- }
-
- @Test
- public void testWrite() throws Exception {
- final List<String> received = new ArrayList<>();
- Thread receiver = new Thread() {
- @Override
- public void run() {
- try {
- Messenger messenger = Messenger.Factory.create();
- messenger.start();
- messenger.subscribe("amqp://~localhost:" + port);
- while (received.size() < 100) {
- messenger.recv();
- while (messenger.incoming() > 0) {
- Message message = messenger.get();
- LOG.info("Received: " + message.getBody().toString());
- received.add(message.getBody().toString());
- }
- }
- messenger.stop();
- } catch (Exception e) {
- LOG.error("Receiver error", e);
- }
- }
- };
- LOG.info("Starting AMQP receiver");
- receiver.start();
-
- List<Message> data = new ArrayList<>();
- for (int i = 0; i < 100; i++) {
- Message message = Message.Factory.create();
- message.setBody(new AmqpValue("Test " + i));
- message.setAddress("amqp://localhost:" + port);
- message.setSubject("test");
- data.add(message);
- }
- pipeline.apply(Create.of(data).withCoder(AmqpMessageCoder.of())).apply(AmqpIO.write());
- LOG.info("Starting pipeline");
- try {
- pipeline.run();
- } finally {
- LOG.info("Join receiver thread");
- receiver.join();
- }
-
- assertEquals(100, received.size());
- for (int i = 0; i < 100; i++) {
- assertTrue(received.contains("AmqpValue{Test " + i + "}"));
- }
- }
-
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/amqp/src/test/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoderTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/amqp/src/test/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoderTest.java b/sdks/java/io/amqp/src/test/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoderTest.java
deleted file mode 100644
index 7a8efeb..0000000
--- a/sdks/java/io/amqp/src/test/java/org/apache/beam/sdk/io/amqp/AmqpMessageCoderTest.java
+++ /dev/null
@@ -1,89 +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.sdk.io.amqp;
-
-import static org.junit.Assert.assertEquals;
-
-import com.google.common.base.Joiner;
-
-import java.util.Collections;
-
-import org.apache.beam.sdk.coders.CoderException;
-import org.apache.beam.sdk.util.CoderUtils;
-import org.apache.qpid.proton.amqp.messaging.AmqpValue;
-import org.apache.qpid.proton.message.Message;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.rules.ExpectedException;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-/**
- * Test on {@link AmqpMessageCoder}.
- */
-@RunWith(JUnit4.class)
-public class AmqpMessageCoderTest {
-
- @Rule
- public ExpectedException thrown = ExpectedException.none();
-
- @Test
- public void encodeDecode() throws Exception {
- Message message = Message.Factory.create();
- message.setBody(new AmqpValue("body"));
- message.setAddress("address");
- message.setSubject("test");
- AmqpMessageCoder coder = AmqpMessageCoder.of();
-
- Message clone = CoderUtils.clone(coder, message);
-
- assertEquals("AmqpValue{body}", clone.getBody().toString());
- assertEquals("address", clone.getAddress());
- assertEquals("test", clone.getSubject());
- }
-
- @Test
- public void encodeDecodeTooMuchLargerMessage() throws Exception {
- thrown.expect(CoderException.class);
- Message message = Message.Factory.create();
- message.setAddress("address");
- message.setSubject("subject");
- String body = Joiner.on("").join(Collections.nCopies(64 * 1024 * 1024, " "));
- message.setBody(new AmqpValue(body));
-
- AmqpMessageCoder coder = AmqpMessageCoder.of();
-
- byte[] encoded = CoderUtils.encodeToByteArray(coder, message);
- }
-
- @Test
- public void encodeDecodeLargeMessage() throws Exception {
- Message message = Message.Factory.create();
- message.setAddress("address");
- message.setSubject("subject");
- String body = Joiner.on("").join(Collections.nCopies(32 * 1024 * 1024, " "));
- message.setBody(new AmqpValue(body));
-
- AmqpMessageCoder coder = AmqpMessageCoder.of();
-
- Message clone = CoderUtils.clone(coder, message);
-
- clone.getBody().toString().equals(message.getBody().toString());
- }
-
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/cassandra/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/cassandra/pom.xml b/sdks/java/io/cassandra/pom.xml
index c74477e..8249f57 100644
--- a/sdks/java/io/cassandra/pom.xml
+++ b/sdks/java/io/cassandra/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/cassandra/src/main/java/org/apache/beam/sdk/io/cassandra/CassandraIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/cassandra/src/main/java/org/apache/beam/sdk/io/cassandra/CassandraIO.java b/sdks/java/io/cassandra/src/main/java/org/apache/beam/sdk/io/cassandra/CassandraIO.java
index 32905b7..b6f4ef6 100644
--- a/sdks/java/io/cassandra/src/main/java/org/apache/beam/sdk/io/cassandra/CassandraIO.java
+++ b/sdks/java/io/cassandra/src/main/java/org/apache/beam/sdk/io/cassandra/CassandraIO.java
@@ -82,7 +82,7 @@ import org.slf4j.LoggerFactory;
* .withEntity(Person.class));
* }</pre>
*/
-@Experimental(Experimental.Kind.SOURCE_SINK)
+@Experimental
public class CassandraIO {
private static final Logger LOG = LoggerFactory.getLogger(CassandraIO.class);
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/common/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/common/pom.xml b/sdks/java/io/common/pom.xml
index df0d94b..f7525fd 100644
--- a/sdks/java/io/common/pom.xml
+++ b/sdks/java/io/common/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/common/src/test/java/org/apache/beam/sdk/io/common/IOTestPipelineOptions.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/common/src/test/java/org/apache/beam/sdk/io/common/IOTestPipelineOptions.java b/sdks/java/io/common/src/test/java/org/apache/beam/sdk/io/common/IOTestPipelineOptions.java
index 25ab929..387fd22 100644
--- a/sdks/java/io/common/src/test/java/org/apache/beam/sdk/io/common/IOTestPipelineOptions.java
+++ b/sdks/java/io/common/src/test/java/org/apache/beam/sdk/io/common/IOTestPipelineOptions.java
@@ -71,7 +71,11 @@ public interface IOTestPipelineOptions extends TestPipelineOptions {
Integer getElasticsearchHttpPort();
void setElasticsearchHttpPort(Integer value);
- /* Cassandra */
+ @Description("Tcp port for elasticsearch server")
+ @Default.Integer(9300)
+ Integer getElasticsearchTcpPort();
+ void setElasticsearchTcpPort(Integer value);
+
@Description("Host for Cassandra server (host name/ip address)")
@Default.String("cassandra-host")
String getCassandraHost();
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/elasticsearch/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/elasticsearch/pom.xml b/sdks/java/io/elasticsearch/pom.xml
index e0a7f21..03632ce 100644
--- a/sdks/java/io/elasticsearch/pom.xml
+++ b/sdks/java/io/elasticsearch/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
@@ -137,14 +137,6 @@
<scope>test</scope>
</dependency>
- <!-- This optional dependency is used by the test framework. Avoids a warning -->
- <dependency>
- <groupId>net.java.dev.jna</groupId>
- <artifactId>jna</artifactId>
- <version>4.1.0</version>
- <scope>test</scope>
- </dependency>
-
<dependency>
<groupId>org.apache.beam</groupId>
<artifactId>beam-runners-direct-java</artifactId>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/elasticsearch/src/main/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/elasticsearch/src/main/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIO.java b/sdks/java/io/elasticsearch/src/main/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIO.java
index 4d76887..f6ceef2 100644
--- a/sdks/java/io/elasticsearch/src/main/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIO.java
+++ b/sdks/java/io/elasticsearch/src/main/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIO.java
@@ -113,7 +113,7 @@ import org.elasticsearch.client.RestClientBuilder;
* <p>Optionally, you can provide {@code withBatchSize()} and {@code withBatchSizeBytes()}
* to specify the size of the write batch in number of documents or in bytes.
*/
-@Experimental(Experimental.Kind.SOURCE_SINK)
+@Experimental
public class ElasticsearchIO {
public static Read read() {
@@ -139,7 +139,7 @@ public class ElasticsearchIO {
private static final ObjectMapper mapper = new ObjectMapper();
- static JsonNode parseResponse(Response response) throws IOException {
+ private static JsonNode parseResponse(Response response) throws IOException {
return mapper.readValue(response.getEntity().getContent(), JsonNode.class);
}
@@ -264,7 +264,7 @@ public class ElasticsearchIO {
builder.addIfNotNull(DisplayData.item("username", getUsername()));
}
- RestClient createClient() throws MalformedURLException {
+ private RestClient createClient() throws MalformedURLException {
HttpHost[] hosts = new HttpHost[getAddresses().size()];
int i = 0;
for (String address : getAddresses()) {
@@ -455,7 +455,16 @@ public class ElasticsearchIO {
while (shards.hasNext()) {
Map.Entry<String, JsonNode> shardJson = shards.next();
String shardId = shardJson.getKey();
- sources.add(new BoundedElasticsearchSource(spec, shardId));
+ JsonNode value = (JsonNode) shardJson.getValue();
+ boolean isPrimaryShard =
+ value
+ .path(0)
+ .path("routing")
+ .path("primary")
+ .asBoolean();
+ if (isPrimaryShard) {
+ sources.add(new BoundedElasticsearchSource(spec, shardId));
+ }
}
checkArgument(!sources.isEmpty(), "No primary shard found");
return sources;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticSearchIOTestUtils.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticSearchIOTestUtils.java b/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticSearchIOTestUtils.java
index 203963d..b0d161f 100644
--- a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticSearchIOTestUtils.java
+++ b/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticSearchIOTestUtils.java
@@ -17,17 +17,19 @@
*/
package org.apache.beam.sdk.io.elasticsearch;
-import com.fasterxml.jackson.databind.JsonNode;
import java.io.IOException;
import java.util.ArrayList;
-import java.util.Collections;
import java.util.List;
-import org.apache.http.HttpEntity;
-import org.apache.http.entity.ContentType;
-import org.apache.http.message.BasicHeader;
-import org.apache.http.nio.entity.NStringEntity;
-import org.elasticsearch.client.Response;
-import org.elasticsearch.client.RestClient;
+import org.elasticsearch.action.admin.indices.exists.indices.IndicesExistsRequest;
+import org.elasticsearch.action.admin.indices.exists.indices.IndicesExistsResponse;
+import org.elasticsearch.action.admin.indices.upgrade.post.UpgradeRequest;
+import org.elasticsearch.action.bulk.BulkRequestBuilder;
+import org.elasticsearch.action.bulk.BulkResponse;
+import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.client.Client;
+import org.elasticsearch.client.IndicesAdminClient;
+import org.elasticsearch.client.Requests;
+import org.elasticsearch.index.IndexNotFoundException;
/** Test utilities to use with {@link ElasticsearchIO}. */
class ElasticSearchIOTestUtils {
@@ -39,68 +41,57 @@ class ElasticSearchIOTestUtils {
}
/** Deletes the given index synchronously. */
- static void deleteIndex(String index, RestClient restClient) throws IOException {
- try {
- restClient.performRequest("DELETE", String.format("/%s", index), new BasicHeader("", ""));
- } catch (IOException e) {
- // it is fine to ignore this expression as deleteIndex occurs in @before,
- // so when the first tests is run, the index does not exist yet
- if (!e.getMessage().contains("index_not_found_exception")){
- throw e;
- }
+ static void deleteIndex(String index, Client client) throws Exception {
+ IndicesAdminClient indices = client.admin().indices();
+ IndicesExistsResponse indicesExistsResponse =
+ indices.exists(new IndicesExistsRequest(index)).get();
+ if (indicesExistsResponse.isExists()) {
+ indices.prepareClose(index).get();
+ indices.delete(Requests.deleteIndexRequest(index)).get();
}
}
/** Inserts the given number of test documents into Elasticsearch. */
- static void insertTestDocuments(String index, String type, long numDocs, RestClient restClient)
- throws IOException {
+ static void insertTestDocuments(String index, String type, long numDocs, Client client)
+ throws Exception {
+ final BulkRequestBuilder bulkRequestBuilder = client.prepareBulk().setRefresh(true);
List<String> data =
ElasticSearchIOTestUtils.createDocuments(
numDocs, ElasticSearchIOTestUtils.InjectionMode.DO_NOT_INJECT_INVALID_DOCS);
- StringBuilder bulkRequest = new StringBuilder();
for (String document : data) {
- bulkRequest.append(String.format("{ \"index\" : {} }%n%s%n", document));
+ bulkRequestBuilder.add(client.prepareIndex(index, type, null).setSource(document));
}
- String endPoint = String.format("/%s/%s/_bulk", index, type);
- HttpEntity requestBody =
- new NStringEntity(bulkRequest.toString(), ContentType.APPLICATION_JSON);
- Response response = restClient.performRequest("POST", endPoint,
- Collections.singletonMap("refresh", "true"), requestBody,
- new BasicHeader("", ""));
- JsonNode searchResult = ElasticsearchIO.parseResponse(response);
- boolean errors = searchResult.path("errors").asBoolean();
- if (errors){
+ final BulkResponse bulkResponse = bulkRequestBuilder.execute().actionGet();
+ if (bulkResponse.hasFailures()) {
throw new IOException(
- String.format("Failed to insert test documents in index %s", index));
+ String.format(
+ "Cannot insert test documents in index %s : %s",
+ index, bulkResponse.buildFailureMessage()));
}
}
/**
- * Forces a refresh of the given index to make recently inserted documents available for search.
+ * Forces an upgrade of the given index to make recently inserted documents available for search.
*
* @return The number of docs in the index
*/
- static long refreshIndexAndGetCurrentNumDocs(String index, String type, RestClient restClient)
- throws IOException {
- long result = 0;
+ static long upgradeIndexAndGetCurrentNumDocs(String index, String type, Client client) {
try {
- String endPoint = String.format("/%s/_refresh", index);
- restClient.performRequest("POST", endPoint, new BasicHeader("", ""));
-
- endPoint = String.format("/%s/%s/_search", index, type);
- Response response = restClient.performRequest("GET", endPoint, new BasicHeader("", ""));
- JsonNode searchResult = ElasticsearchIO.parseResponse(response);
- result = searchResult.path("hits").path("total").asLong();
- } catch (IOException e) {
+ client.admin().indices().upgrade(new UpgradeRequest(index)).actionGet();
+ SearchResponse response =
+ client.prepareSearch(index).setTypes(type).execute().actionGet(5000);
+ return response.getHits().getTotalHits();
// it is fine to ignore bellow exceptions because in testWriteWithBatchSize* sometimes,
// we call upgrade before any doc have been written
// (when there are fewer docs processed than batchSize).
// In that cases index/type has not been created (created upon first doc insertion)
- if (!e.getMessage().contains("index_not_found_exception")){
+ } catch (IndexNotFoundException e) {
+ } catch (java.lang.IllegalArgumentException e) {
+ if (!e.getMessage().contains("No search type")) {
throw e;
}
}
- return result;
+ return 0;
}
/**
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java b/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java
index 7c37e87..2d6393a 100644
--- a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java
+++ b/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOIT.java
@@ -32,7 +32,7 @@ import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Count;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.values.PCollection;
-import org.elasticsearch.client.RestClient;
+import org.elasticsearch.client.transport.TransportClient;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Rule;
@@ -57,7 +57,7 @@ import org.slf4j.LoggerFactory;
*/
public class ElasticsearchIOIT {
private static final Logger LOG = LoggerFactory.getLogger(ElasticsearchIOIT.class);
- private static RestClient restClient;
+ private static TransportClient client;
private static IOTestPipelineOptions options;
private static ElasticsearchIO.ConnectionConfiguration readConnectionConfiguration;
@Rule public TestPipeline pipeline = TestPipeline.create();
@@ -66,16 +66,16 @@ public class ElasticsearchIOIT {
public static void beforeClass() throws Exception {
PipelineOptionsFactory.register(IOTestPipelineOptions.class);
options = TestPipeline.testingPipelineOptions().as(IOTestPipelineOptions.class);
+ client = ElasticsearchTestDataSet.getClient(options);
readConnectionConfiguration =
ElasticsearchTestDataSet.getConnectionConfiguration(
options, ElasticsearchTestDataSet.ReadOrWrite.READ);
- restClient = readConnectionConfiguration.createClient();
}
@AfterClass
public static void afterClass() throws Exception {
- ElasticsearchTestDataSet.deleteIndex(restClient, ElasticsearchTestDataSet.ReadOrWrite.WRITE);
- restClient.close();
+ ElasticsearchTestDataSet.deleteIndex(client, ElasticsearchTestDataSet.ReadOrWrite.WRITE);
+ client.close();
}
@Test
@@ -128,8 +128,8 @@ public class ElasticsearchIOIT {
pipeline.run();
long currentNumDocs =
- ElasticSearchIOTestUtils.refreshIndexAndGetCurrentNumDocs(
- ElasticsearchTestDataSet.ES_INDEX, ElasticsearchTestDataSet.ES_TYPE, restClient);
+ ElasticSearchIOTestUtils.upgradeIndexAndGetCurrentNumDocs(
+ ElasticsearchTestDataSet.ES_INDEX, ElasticsearchTestDataSet.ES_TYPE, client);
assertEquals(ElasticsearchTestDataSet.NUM_DOCS, currentNumDocs);
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java b/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java
index b349a29..260af79 100644
--- a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java
+++ b/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIOTest.java
@@ -39,11 +39,11 @@ import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFnTester;
import org.apache.beam.sdk.values.PCollection;
import org.elasticsearch.action.search.SearchResponse;
-import org.elasticsearch.client.RestClient;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.node.Node;
+import org.elasticsearch.node.NodeBuilder;
import org.hamcrest.CustomMatcher;
import org.junit.AfterClass;
import org.junit.Before;
@@ -74,10 +74,9 @@ public class ElasticsearchIOTest implements Serializable {
private static final long BATCH_SIZE_BYTES = 2048L;
private static Node node;
- private static RestClient restClient;
private static ElasticsearchIO.ConnectionConfiguration connectionConfiguration;
- @ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder();
+ @ClassRule public static TemporaryFolder folder = new TemporaryFolder();
@Rule
public TestPipeline pipeline = TestPipeline.create();
@@ -92,8 +91,8 @@ public class ElasticsearchIOTest implements Serializable {
.put("cluster.name", "beam")
.put("http.enabled", "true")
.put("node.data", "true")
- .put("path.data", TEMPORARY_FOLDER.getRoot().getPath())
- .put("path.home", TEMPORARY_FOLDER.getRoot().getPath())
+ .put("path.data", folder.getRoot().getPath())
+ .put("path.home", folder.getRoot().getPath())
.put("node.name", "beam")
.put("network.host", ES_IP)
.put("http.port", esHttpPort)
@@ -101,29 +100,27 @@ public class ElasticsearchIOTest implements Serializable {
// had problems with some jdk, embedded ES was too slow for bulk insertion,
// and queue of 50 was full. No pb with real ES instance (cf testWrite integration test)
.put("threadpool.bulk.queue_size", 100);
- node = new Node(settingsBuilder.build());
+ node = NodeBuilder.nodeBuilder().settings(settingsBuilder).build();
LOG.info("Elasticsearch node created");
node.start();
connectionConfiguration =
ElasticsearchIO.ConnectionConfiguration.create(
new String[] {"http://" + ES_IP + ":" + esHttpPort}, ES_INDEX, ES_TYPE);
- restClient = connectionConfiguration.createClient();
}
@AfterClass
- public static void afterClass() throws IOException{
- restClient.close();
+ public static void afterClass() {
node.close();
}
@Before
public void before() throws Exception {
- ElasticSearchIOTestUtils.deleteIndex(ES_INDEX, restClient);
+ ElasticSearchIOTestUtils.deleteIndex(ES_INDEX, node.client());
}
@Test
public void testSizes() throws Exception {
- ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, restClient);
+ ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, node.client());
PipelineOptions options = PipelineOptionsFactory.create();
ElasticsearchIO.Read read =
ElasticsearchIO.read().withConnectionConfiguration(connectionConfiguration);
@@ -137,7 +134,7 @@ public class ElasticsearchIOTest implements Serializable {
@Test
public void testRead() throws Exception {
- ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, restClient);
+ ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, node.client());
PCollection<String> output =
pipeline.apply(
@@ -153,7 +150,7 @@ public class ElasticsearchIOTest implements Serializable {
@Test
public void testReadWithQuery() throws Exception {
- ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, restClient);
+ ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, node.client());
String query =
"{\n"
@@ -188,7 +185,7 @@ public class ElasticsearchIOTest implements Serializable {
pipeline.run();
long currentNumDocs =
- ElasticSearchIOTestUtils.refreshIndexAndGetCurrentNumDocs(ES_INDEX, ES_TYPE, restClient);
+ ElasticSearchIOTestUtils.upgradeIndexAndGetCurrentNumDocs(ES_INDEX, ES_TYPE, node.client());
assertEquals(NUM_DOCS, currentNumDocs);
QueryBuilder queryBuilder = QueryBuilders.queryStringQuery("Einstein").field("scientist");
@@ -261,8 +258,9 @@ public class ElasticsearchIOTest implements Serializable {
if ((numDocsProcessed % 100) == 0) {
// force the index to upgrade after inserting for the inserted docs
// to be searchable immediately
- long currentNumDocs = ElasticSearchIOTestUtils
- .refreshIndexAndGetCurrentNumDocs(ES_INDEX, ES_TYPE, restClient);
+ long currentNumDocs =
+ ElasticSearchIOTestUtils.upgradeIndexAndGetCurrentNumDocs(
+ ES_INDEX, ES_TYPE, node.client());
if ((numDocsProcessed % BATCH_SIZE) == 0) {
/* bundle end */
assertEquals(
@@ -306,8 +304,8 @@ public class ElasticsearchIOTest implements Serializable {
// force the index to upgrade after inserting for the inserted docs
// to be searchable immediately
long currentNumDocs =
- ElasticSearchIOTestUtils.refreshIndexAndGetCurrentNumDocs(
- ES_INDEX, ES_TYPE, restClient);
+ ElasticSearchIOTestUtils.upgradeIndexAndGetCurrentNumDocs(
+ ES_INDEX, ES_TYPE, node.client());
if (sizeProcessed / BATCH_SIZE_BYTES > batchInserted) {
/* bundle end */
assertThat(
@@ -329,7 +327,7 @@ public class ElasticsearchIOTest implements Serializable {
@Test
public void testSplit() throws Exception {
- ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, restClient);
+ ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, node.client());
PipelineOptions options = PipelineOptionsFactory.create();
ElasticsearchIO.Read read =
ElasticsearchIO.read().withConnectionConfiguration(connectionConfiguration);
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchTestDataSet.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchTestDataSet.java b/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchTestDataSet.java
index 2a2dbe9..3a9aae6 100644
--- a/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchTestDataSet.java
+++ b/sdks/java/io/elasticsearch/src/test/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchTestDataSet.java
@@ -17,11 +17,13 @@
*/
package org.apache.beam.sdk.io.elasticsearch;
+import static java.net.InetAddress.getByName;
import java.io.IOException;
import org.apache.beam.sdk.io.common.IOTestPipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.elasticsearch.client.RestClient;
+import org.elasticsearch.client.transport.TransportClient;
+import org.elasticsearch.common.transport.InetSocketTransportAddress;
/**
* Manipulates test data used by the {@link ElasticsearchIO}
@@ -49,6 +51,7 @@ public class ElasticsearchTestDataSet {
* -Dexec.mainClass=org.apache.beam.sdk.io.elasticsearch.ElasticsearchTestDataSet \
* -Dexec.args="--elasticsearchServer=1.2.3.4 \
* --elasticsearchHttpPort=9200 \
+ * --elasticsearchTcpPort=9300" \
* -Dexec.classpathScope=test
* </pre>
*
@@ -59,20 +62,29 @@ public class ElasticsearchTestDataSet {
PipelineOptionsFactory.register(IOTestPipelineOptions.class);
IOTestPipelineOptions options =
PipelineOptionsFactory.fromArgs(args).as(IOTestPipelineOptions.class);
- createAndPopulateReadIndex(options);
+
+ createAndPopulateIndex(getClient(options), ReadOrWrite.READ);
}
- private static void createAndPopulateReadIndex(IOTestPipelineOptions options) throws Exception {
- RestClient restClient = getConnectionConfiguration(options, ReadOrWrite.READ).createClient();
+ private static void createAndPopulateIndex(TransportClient client, ReadOrWrite rOw)
+ throws Exception {
// automatically creates the index and insert docs
- try {
- ElasticSearchIOTestUtils.insertTestDocuments(ES_INDEX, ES_TYPE, NUM_DOCS, restClient);
- } finally {
- restClient.close();
- }
+ ElasticSearchIOTestUtils.insertTestDocuments(
+ (rOw == ReadOrWrite.READ) ? ES_INDEX : writeIndex, ES_TYPE, NUM_DOCS, client);
+ }
+
+ public static TransportClient getClient(IOTestPipelineOptions options) throws Exception {
+ TransportClient client =
+ TransportClient.builder()
+ .build()
+ .addTransportAddress(
+ new InetSocketTransportAddress(
+ getByName(options.getElasticsearchServer()),
+ options.getElasticsearchTcpPort()));
+ return client;
}
- static ElasticsearchIO.ConnectionConfiguration getConnectionConfiguration(
+ public static ElasticsearchIO.ConnectionConfiguration getConnectionConfiguration(
IOTestPipelineOptions options, ReadOrWrite rOw) throws IOException {
ElasticsearchIO.ConnectionConfiguration connectionConfiguration =
ElasticsearchIO.ConnectionConfiguration.create(
@@ -87,9 +99,8 @@ public class ElasticsearchTestDataSet {
return connectionConfiguration;
}
- static void deleteIndex(RestClient restClient, ReadOrWrite rOw) throws Exception {
- ElasticSearchIOTestUtils
- .deleteIndex((rOw == ReadOrWrite.READ) ? ES_INDEX : writeIndex, restClient);
+ public static void deleteIndex(TransportClient client, ReadOrWrite rOw) throws Exception {
+ ElasticSearchIOTestUtils.deleteIndex((rOw == ReadOrWrite.READ) ? ES_INDEX : writeIndex, client);
}
/** Enum that tells whether we use the index for reading or for writing. */
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/pom.xml b/sdks/java/io/google-cloud-platform/pom.xml
index a1495f2..8b53820 100644
--- a/sdks/java/io/google-cloud-platform/pom.xml
+++ b/sdks/java/io/google-cloud-platform/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
@@ -93,12 +93,7 @@
<dependency>
<groupId>com.google.api</groupId>
- <artifactId>gax-grpc</artifactId>
- </dependency>
-
- <dependency>
- <groupId>com.google.cloud</groupId>
- <artifactId>google-cloud-core-grpc</artifactId>
+ <artifactId>api-common</artifactId>
</dependency>
<dependency>
@@ -258,6 +253,11 @@
<artifactId>proto-google-common-protos</artifactId>
</dependency>
+ <dependency>
+ <groupId>org.apache.commons</groupId>
+ <artifactId>commons-lang3</artifactId>
+ </dependency>
+
<!-- Test dependencies -->
<dependency>
<groupId>org.apache.beam</groupId>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java
index e46b1d3..4393a63 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java
@@ -32,7 +32,6 @@ import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.ListCoder;
import org.apache.beam.sdk.coders.NullableCoder;
-import org.apache.beam.sdk.coders.ShardedKeyCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.CreateDisposition;
@@ -58,7 +57,6 @@ import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.ShardedKey;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
import org.apache.beam.sdk.values.TypeDescriptor;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java
index c5c2462..edb1e0d 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java
@@ -23,7 +23,8 @@ import static com.google.common.base.Preconditions.checkArgument;
import com.google.api.services.bigquery.model.TableSchema;
import com.google.common.collect.Lists;
import java.io.Serializable;
-import java.lang.reflect.TypeVariable;
+import java.lang.reflect.ParameterizedType;
+import java.lang.reflect.Type;
import java.util.List;
import javax.annotation.Nullable;
import org.apache.beam.sdk.coders.CannotProvideCoderException;
@@ -31,7 +32,6 @@ import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderRegistry;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.sdk.values.ValueInSingleWindow;
/**
@@ -158,16 +158,21 @@ public abstract class DynamicDestinations<T, DestinationT> implements Serializab
}
// If dynamicDestinations doesn't provide a coder, try to find it in the coder registry.
// We must first use reflection to figure out what the type parameter is.
- TypeDescriptor<?> superDescriptor =
- TypeDescriptor.of(getClass()).getSupertype(DynamicDestinations.class);
- if (!superDescriptor.getRawType().equals(DynamicDestinations.class)) {
- throw new AssertionError(
- "Couldn't find the DynamicDestinations superclass of " + this.getClass());
+ for (Type superclass = getClass().getGenericSuperclass();
+ superclass != null;
+ superclass = ((Class) superclass).getGenericSuperclass()) {
+ if (superclass instanceof ParameterizedType) {
+ ParameterizedType parameterized = (ParameterizedType) superclass;
+ if (parameterized.getRawType() == DynamicDestinations.class) {
+ // DestinationT is the second parameter.
+ Type parameter = parameterized.getActualTypeArguments()[1];
+ @SuppressWarnings("unchecked")
+ Class<DestinationT> parameterClass = (Class<DestinationT>) parameter;
+ return registry.getCoder(parameterClass);
+ }
+ }
}
- TypeVariable typeVariable = superDescriptor.getTypeParameter("DestinationT");
- @SuppressWarnings("unchecked")
- TypeDescriptor<DestinationT> descriptor =
- (TypeDescriptor<DestinationT>) superDescriptor.resolveType(typeVariable);
- return registry.getCoder(descriptor);
+ throw new AssertionError(
+ "Couldn't find the DynamicDestinations superclass of " + this.getClass());
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/GenerateShardedTable.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/GenerateShardedTable.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/GenerateShardedTable.java
index 55672ff..90d41a0 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/GenerateShardedTable.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/GenerateShardedTable.java
@@ -23,7 +23,6 @@ import java.util.concurrent.ThreadLocalRandom;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.values.KV;
-import org.apache.beam.sdk.values.ShardedKey;
/**
* Given a write to a specific table, assign that to one of the
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/ShardedKey.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/ShardedKey.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/ShardedKey.java
new file mode 100644
index 0000000..c2b739f
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/ShardedKey.java
@@ -0,0 +1,67 @@
+/*
+ * 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.sdk.io.gcp.bigquery;
+
+import java.io.Serializable;
+import java.util.Objects;
+
+/**
+ * A key and a shard number.
+ */
+class ShardedKey<K> implements Serializable {
+ private static final long serialVersionUID = 1L;
+ private final K key;
+ private final int shardNumber;
+
+ public static <K> ShardedKey<K> of(K key, int shardNumber) {
+ return new ShardedKey<>(key, shardNumber);
+ }
+
+ ShardedKey(K key, int shardNumber) {
+ this.key = key;
+ this.shardNumber = shardNumber;
+ }
+
+ public K getKey() {
+ return key;
+ }
+
+ public int getShardNumber() {
+ return shardNumber;
+ }
+
+ @Override
+ public String toString() {
+ return "key: " + key + " shard: " + shardNumber;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (!(o instanceof ShardedKey)) {
+ return false;
+ }
+ ShardedKey<K> other = (ShardedKey<K>) o;
+ return Objects.equals(key, other.key) && Objects.equals(shardNumber, other.shardNumber);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(key, shardNumber);
+ }
+}
[23/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/AfterWatermarkStateMachineTest.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/AfterWatermarkStateMachineTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/AfterWatermarkStateMachineTest.java
index 45a5cfb..e4d10a0 100644
--- a/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/AfterWatermarkStateMachineTest.java
+++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/AfterWatermarkStateMachineTest.java
@@ -25,6 +25,7 @@ import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import org.apache.beam.runners.core.triggers.TriggerStateMachine.OnMergeContext;
+import org.apache.beam.runners.core.triggers.TriggerStateMachine.OnceTriggerStateMachine;
import org.apache.beam.runners.core.triggers.TriggerStateMachineTester.SimpleTriggerStateMachineTester;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
@@ -45,8 +46,8 @@ import org.mockito.MockitoAnnotations;
@RunWith(JUnit4.class)
public class AfterWatermarkStateMachineTest {
- @Mock private TriggerStateMachine mockEarly;
- @Mock private TriggerStateMachine mockLate;
+ @Mock private OnceTriggerStateMachine mockEarly;
+ @Mock private OnceTriggerStateMachine mockLate;
private SimpleTriggerStateMachineTester<IntervalWindow> tester;
private static TriggerStateMachine.TriggerContext anyTriggerContext() {
@@ -69,7 +70,7 @@ public class AfterWatermarkStateMachineTest {
MockitoAnnotations.initMocks(this);
}
- public void testRunningAsTrigger(TriggerStateMachine mockTrigger, IntervalWindow window)
+ public void testRunningAsTrigger(OnceTriggerStateMachine mockTrigger, IntervalWindow window)
throws Exception {
// Don't fire due to mock saying no
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/StubTriggerStateMachine.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/StubTriggerStateMachine.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/StubTriggerStateMachine.java
index 1bc757e..4512848 100644
--- a/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/StubTriggerStateMachine.java
+++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/triggers/StubTriggerStateMachine.java
@@ -18,11 +18,12 @@
package org.apache.beam.runners.core.triggers;
import com.google.common.collect.Lists;
+import org.apache.beam.runners.core.triggers.TriggerStateMachine.OnceTriggerStateMachine;
/**
- * No-op {@link TriggerStateMachine} implementation for testing.
+ * No-op {@link OnceTriggerStateMachine} implementation for testing.
*/
-abstract class StubTriggerStateMachine extends TriggerStateMachine {
+abstract class StubTriggerStateMachine extends OnceTriggerStateMachine {
/**
* Create a stub {@link TriggerStateMachine} instance which returns the specified name on {@link
* #toString()}.
@@ -41,7 +42,7 @@ abstract class StubTriggerStateMachine extends TriggerStateMachine {
}
@Override
- public void onFire(TriggerContext context) throws Exception {
+ protected void onOnlyFiring(TriggerContext context) throws Exception {
}
@Override
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/pom.xml
----------------------------------------------------------------------
diff --git a/runners/direct-java/pom.xml b/runners/direct-java/pom.xml
index e14e813..bec2113 100644
--- a/runners/direct-java/pom.xml
+++ b/runners/direct-java/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-runners-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
@@ -117,7 +117,7 @@
</relocation>
</relocations>
<transformers>
- <transformer implementation="org.apache.maven.plugins.shade.resource.ServicesResourceTransformer" />
+ <transformer implementation="org.apache.maven.plugins.shade.resource.ServicesResourceTransformer"/>
</transformers>
</configuration>
</execution>
@@ -155,8 +155,7 @@
<systemPropertyVariables>
<beamTestPipelineOptions>
[
- "--runner=DirectRunner",
- "--runnerDeterminedSharding=false"
+ "--runner=DirectRunner"
]
</beamTestPipelineOptions>
</systemPropertyVariables>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/CommittedResult.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/CommittedResult.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/CommittedResult.java
index 70e3ac3..8c45449 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/CommittedResult.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/CommittedResult.java
@@ -19,8 +19,8 @@
package org.apache.beam.runners.direct;
import com.google.auto.value.AutoValue;
-import com.google.common.base.Optional;
import java.util.Set;
+import javax.annotation.Nullable;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.transforms.View.CreatePCollectionView;
@@ -36,10 +36,12 @@ abstract class CommittedResult {
/**
* Returns the {@link CommittedBundle} that contains the input elements that could not be
- * processed by the evaluation. The returned optional is present if there were any unprocessed
- * input elements, and absent otherwise.
+ * processed by the evaluation.
+ *
+ * <p>{@code null} if the input bundle was null.
*/
- public abstract Optional<? extends CommittedBundle<?>> getUnprocessedInputs();
+ @Nullable
+ public abstract CommittedBundle<?> getUnprocessedInputs();
/**
* Returns the outputs produced by the transform.
@@ -57,7 +59,7 @@ abstract class CommittedResult {
public static CommittedResult create(
TransformResult<?> original,
- Optional<? extends CommittedBundle<?>> unprocessedElements,
+ CommittedBundle<?> unprocessedElements,
Iterable<? extends CommittedBundle<?>> outputs,
Set<OutputType> producedOutputs) {
return new AutoValue_CommittedResult(original.getTransform(),
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraph.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraph.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraph.java
index ad17b2b..c2c0afa 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraph.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraph.java
@@ -17,8 +17,6 @@
*/
package org.apache.beam.runners.direct;
-import static com.google.common.base.Preconditions.checkArgument;
-
import com.google.common.collect.ListMultimap;
import java.util.Collection;
import java.util.List;
@@ -38,8 +36,7 @@ import org.apache.beam.sdk.values.PValue;
class DirectGraph {
private final Map<PCollection<?>, AppliedPTransform<?, ?, ?>> producers;
private final Map<PCollectionView<?>, AppliedPTransform<?, ?, ?>> viewWriters;
- private final ListMultimap<PInput, AppliedPTransform<?, ?, ?>> perElementConsumers;
- private final ListMultimap<PValue, AppliedPTransform<?, ?, ?>> allConsumers;
+ private final ListMultimap<PInput, AppliedPTransform<?, ?, ?>> primitiveConsumers;
private final Set<AppliedPTransform<?, ?, ?>> rootTransforms;
private final Map<AppliedPTransform<?, ?, ?>, String> stepNames;
@@ -47,36 +44,23 @@ class DirectGraph {
public static DirectGraph create(
Map<PCollection<?>, AppliedPTransform<?, ?, ?>> producers,
Map<PCollectionView<?>, AppliedPTransform<?, ?, ?>> viewWriters,
- ListMultimap<PInput, AppliedPTransform<?, ?, ?>> perElementConsumers,
- ListMultimap<PValue, AppliedPTransform<?, ?, ?>> allConsumers,
+ ListMultimap<PInput, AppliedPTransform<?, ?, ?>> primitiveConsumers,
Set<AppliedPTransform<?, ?, ?>> rootTransforms,
Map<AppliedPTransform<?, ?, ?>, String> stepNames) {
- return new DirectGraph(
- producers, viewWriters, perElementConsumers, allConsumers, rootTransforms, stepNames);
+ return new DirectGraph(producers, viewWriters, primitiveConsumers, rootTransforms, stepNames);
}
private DirectGraph(
Map<PCollection<?>, AppliedPTransform<?, ?, ?>> producers,
Map<PCollectionView<?>, AppliedPTransform<?, ?, ?>> viewWriters,
- ListMultimap<PInput, AppliedPTransform<?, ?, ?>> perElementConsumers,
- ListMultimap<PValue, AppliedPTransform<?, ?, ?>> allConsumers,
+ ListMultimap<PInput, AppliedPTransform<?, ?, ?>> primitiveConsumers,
Set<AppliedPTransform<?, ?, ?>> rootTransforms,
Map<AppliedPTransform<?, ?, ?>, String> stepNames) {
this.producers = producers;
this.viewWriters = viewWriters;
- this.perElementConsumers = perElementConsumers;
- this.allConsumers = allConsumers;
+ this.primitiveConsumers = primitiveConsumers;
this.rootTransforms = rootTransforms;
this.stepNames = stepNames;
- for (AppliedPTransform<?, ?, ?> step : stepNames.keySet()) {
- for (PValue input : step.getInputs().values()) {
- checkArgument(
- allConsumers.get(input).contains(step),
- "Step %s lists value %s as input, but it is not in the graph of consumers",
- step.getFullName(),
- input);
- }
- }
}
AppliedPTransform<?, ?, ?> getProducer(PCollection<?> produced) {
@@ -87,22 +71,14 @@ class DirectGraph {
return viewWriters.get(view);
}
- List<AppliedPTransform<?, ?, ?>> getPerElementConsumers(PValue consumed) {
- return perElementConsumers.get(consumed);
- }
-
- List<AppliedPTransform<?, ?, ?>> getAllConsumers(PValue consumed) {
- return allConsumers.get(consumed);
+ List<AppliedPTransform<?, ?, ?>> getPrimitiveConsumers(PValue consumed) {
+ return primitiveConsumers.get(consumed);
}
Set<AppliedPTransform<?, ?, ?>> getRootTransforms() {
return rootTransforms;
}
- Set<PCollection<?>> getPCollections() {
- return producers.keySet();
- }
-
Set<PCollectionView<?>> getViews() {
return viewWriters.keySet();
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraphVisitor.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraphVisitor.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraphVisitor.java
index 675de2c..d54de5d 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraphVisitor.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraphVisitor.java
@@ -21,26 +21,19 @@ import static com.google.common.base.Preconditions.checkState;
import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.ListMultimap;
-import com.google.common.collect.Sets;
-import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
-import org.apache.beam.runners.core.construction.TransformInputs;
-import org.apache.beam.runners.direct.ViewOverrideFactory.WriteView;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.Pipeline.PipelineVisitor;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.runners.TransformHierarchy;
import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PInput;
import org.apache.beam.sdk.values.PValue;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
/**
* Tracks the {@link AppliedPTransform AppliedPTransforms} that consume each {@link PValue} in the
@@ -48,15 +41,11 @@ import org.slf4j.LoggerFactory;
* input after the upstream transform has produced and committed output.
*/
class DirectGraphVisitor extends PipelineVisitor.Defaults {
- private static final Logger LOG = LoggerFactory.getLogger(DirectGraphVisitor.class);
private Map<PCollection<?>, AppliedPTransform<?, ?, ?>> producers = new HashMap<>();
private Map<PCollectionView<?>, AppliedPTransform<?, ?, ?>> viewWriters = new HashMap<>();
- private Set<PCollectionView<?>> consumedViews = new HashSet<>();
- private ListMultimap<PInput, AppliedPTransform<?, ?, ?>> perElementConsumers =
- ArrayListMultimap.create();
- private ListMultimap<PValue, AppliedPTransform<?, ?, ?>> allConsumers =
+ private ListMultimap<PInput, AppliedPTransform<?, ?, ?>> primitiveConsumers =
ArrayListMultimap.create();
private Set<AppliedPTransform<?, ?, ?>> rootTransforms = new HashSet<>();
@@ -84,13 +73,6 @@ class DirectGraphVisitor extends PipelineVisitor.Defaults {
getClass().getSimpleName());
if (node.isRootNode()) {
finalized = true;
- checkState(
- viewWriters.keySet().containsAll(consumedViews),
- "All %ss that are consumed must be written by some %s %s: Missing %s",
- PCollectionView.class.getSimpleName(),
- WriteView.class.getSimpleName(),
- PTransform.class.getSimpleName(),
- Sets.difference(consumedViews, viewWriters.keySet()));
}
}
@@ -101,30 +83,18 @@ class DirectGraphVisitor extends PipelineVisitor.Defaults {
if (node.getInputs().isEmpty()) {
rootTransforms.add(appliedTransform);
} else {
- Collection<PValue> mainInputs =
- TransformInputs.nonAdditionalInputs(node.toAppliedPTransform(getPipeline()));
- if (!mainInputs.containsAll(node.getInputs().values())) {
- LOG.debug(
- "Inputs reduced to {} from {} by removing additional inputs",
- mainInputs,
- node.getInputs().values());
- }
- for (PValue value : mainInputs) {
- perElementConsumers.put(value, appliedTransform);
- }
for (PValue value : node.getInputs().values()) {
- allConsumers.put(value, appliedTransform);
+ primitiveConsumers.put(value, appliedTransform);
+ }
+ if (node.getTransform() instanceof ViewOverrideFactory.WriteView) {
+ viewWriters.put(
+ ((ViewOverrideFactory.WriteView<?, ?>) node.getTransform()).getView(),
+ node.toAppliedPTransform(getPipeline()));
}
- }
- if (node.getTransform() instanceof ParDo.MultiOutput) {
- consumedViews.addAll(((ParDo.MultiOutput<?, ?>) node.getTransform()).getSideInputs());
- } else if (node.getTransform() instanceof ViewOverrideFactory.WriteView) {
- viewWriters.put(
- ((WriteView) node.getTransform()).getView(), node.toAppliedPTransform(getPipeline()));
}
}
- @Override
+ @Override
public void visitValue(PValue value, TransformHierarchy.Node producer) {
AppliedPTransform<?, ?, ?> appliedTransform = getAppliedTransform(producer);
if (value instanceof PCollection && !producers.containsKey(value)) {
@@ -149,6 +119,6 @@ class DirectGraphVisitor extends PipelineVisitor.Defaults {
public DirectGraph getGraph() {
checkState(finalized, "Can't get a graph before the Pipeline has been completely traversed");
return DirectGraph.create(
- producers, viewWriters, perElementConsumers, allConsumers, rootTransforms, stepNames);
+ producers, viewWriters, primitiveConsumers, rootTransforms, stepNames);
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGroupByKey.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGroupByKey.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGroupByKey.java
index 06b8e29..2fc0dd4 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGroupByKey.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGroupByKey.java
@@ -36,17 +36,13 @@ import org.apache.beam.sdk.values.WindowingStrategy;
class DirectGroupByKey<K, V>
extends ForwardingPTransform<PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>> {
- private final PTransform<PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>> original;
+ private final GroupByKey<K, V> original;
static final String DIRECT_GBKO_URN = "urn:beam:directrunner:transforms:gbko:v1";
static final String DIRECT_GABW_URN = "urn:beam:directrunner:transforms:gabw:v1";
- private final WindowingStrategy<?, ?> outputWindowingStrategy;
- DirectGroupByKey(
- PTransform<PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>> original,
- WindowingStrategy<?, ?> outputWindowingStrategy) {
- this.original = original;
- this.outputWindowingStrategy = outputWindowingStrategy;
+ DirectGroupByKey(GroupByKey<K, V> from) {
+ this.original = from;
}
@Override
@@ -61,6 +57,9 @@ class DirectGroupByKey<K, V>
// key/value input elements and the window merge operation of the
// window function associated with the input PCollection.
WindowingStrategy<?, ?> inputWindowingStrategy = input.getWindowingStrategy();
+ // Update the windowing strategy as appropriate.
+ WindowingStrategy<?, ?> outputWindowingStrategy =
+ original.updateWindowingStrategy(inputWindowingStrategy);
// By default, implement GroupByKey via a series of lower-level operations.
return input
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGroupByKeyOverrideFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGroupByKeyOverrideFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGroupByKeyOverrideFactory.java
index 9c2de3d..c2eb5e7 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGroupByKeyOverrideFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGroupByKeyOverrideFactory.java
@@ -17,34 +17,26 @@
*/
package org.apache.beam.runners.direct;
-import com.google.common.collect.Iterables;
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.runners.PTransformOverrideFactory;
import org.apache.beam.sdk.transforms.GroupByKey;
-import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
/** A {@link PTransformOverrideFactory} for {@link GroupByKey} PTransforms. */
final class DirectGroupByKeyOverrideFactory<K, V>
extends SingleInputOutputOverrideFactory<
- PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>,
- PTransform<PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>>> {
+ PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>, GroupByKey<K, V>> {
@Override
public PTransformReplacement<PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>>
getReplacementTransform(
AppliedPTransform<
- PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>,
- PTransform<PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>>>
+ PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>, GroupByKey<K, V>>
transform) {
-
- PCollection<KV<K, Iterable<V>>> output =
- (PCollection<KV<K, Iterable<V>>>) Iterables.getOnlyElement(transform.getOutputs().values());
-
return PTransformReplacement.of(
PTransformReplacements.getSingletonMainInput(transform),
- new DirectGroupByKey<>(transform.getTransform(), output.getWindowingStrategy()));
+ new DirectGroupByKey<>(transform.getTransform()));
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRegistrar.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRegistrar.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRegistrar.java
index 53fb2f2..0e6fbab 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRegistrar.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRegistrar.java
@@ -50,7 +50,7 @@ public class DirectRegistrar {
@Override
public Iterable<Class<? extends PipelineOptions>> getPipelineOptions() {
return ImmutableList.<Class<? extends PipelineOptions>>of(
- DirectOptions.class, DirectTestOptions.class);
+ DirectOptions.class);
}
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/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 4621224..dbd1ec4 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
@@ -30,7 +30,6 @@ import java.util.Map;
import java.util.Set;
import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems;
import org.apache.beam.runners.core.construction.PTransformMatchers;
-import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.runners.core.construction.SplittableParDo;
import org.apache.beam.runners.direct.DirectRunner.DirectPipelineResult;
import org.apache.beam.runners.direct.TestStreamEvaluatorFactory.DirectTestStreamFactory;
@@ -38,11 +37,17 @@ import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.Pipeline.PipelineExecutionException;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.PipelineRunner;
+import org.apache.beam.sdk.io.Read;
import org.apache.beam.sdk.metrics.MetricResults;
import org.apache.beam.sdk.metrics.MetricsEnvironment;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.runners.PTransformOverride;
+import org.apache.beam.sdk.testing.TestStream;
+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.ParDo.MultiOutput;
+import org.apache.beam.sdk.transforms.View.CreatePCollectionView;
import org.apache.beam.sdk.util.UserCodeException;
import org.apache.beam.sdk.values.PCollection;
import org.joda.time.Duration;
@@ -69,17 +74,16 @@ public class DirectRunner extends PipelineRunner<DirectPipelineResult> {
IMMUTABILITY {
@Override
public boolean appliesTo(PCollection<?> collection, DirectGraph graph) {
- return CONTAINS_UDF.contains(
- PTransformTranslation.urnForTransform(graph.getProducer(collection).getTransform()));
+ return CONTAINS_UDF.contains(graph.getProducer(collection).getTransform().getClass());
}
};
/**
* The set of {@link PTransform PTransforms} that execute a UDF. Useful for some enforcements.
*/
- private static final Set<String> CONTAINS_UDF =
+ private static final Set<Class<? extends PTransform>> CONTAINS_UDF =
ImmutableSet.of(
- PTransformTranslation.READ_TRANSFORM_URN, PTransformTranslation.PAR_DO_TRANSFORM_URN);
+ Read.Bounded.class, Read.Unbounded.class, ParDo.SingleOutput.class, MultiOutput.class);
public abstract boolean appliesTo(PCollection<?> collection, DirectGraph graph);
@@ -108,19 +112,22 @@ public class DirectRunner extends PipelineRunner<DirectPipelineResult> {
return bundleFactory;
}
- private static Map<String, Collection<ModelEnforcementFactory>>
+ @SuppressWarnings("rawtypes")
+ private static Map<Class<? extends PTransform>, Collection<ModelEnforcementFactory>>
defaultModelEnforcements(Set<Enforcement> enabledEnforcements) {
- ImmutableMap.Builder<String, Collection<ModelEnforcementFactory>> enforcements =
- ImmutableMap.builder();
+ ImmutableMap.Builder<Class<? extends PTransform>, Collection<ModelEnforcementFactory>>
+ enforcements = ImmutableMap.builder();
ImmutableList.Builder<ModelEnforcementFactory> enabledParDoEnforcements =
ImmutableList.builder();
if (enabledEnforcements.contains(Enforcement.IMMUTABILITY)) {
enabledParDoEnforcements.add(ImmutabilityEnforcementFactory.create());
}
Collection<ModelEnforcementFactory> parDoEnforcements = enabledParDoEnforcements.build();
- enforcements.put(PTransformTranslation.PAR_DO_TRANSFORM_URN, parDoEnforcements);
+ enforcements.put(ParDo.SingleOutput.class, parDoEnforcements);
+ enforcements.put(MultiOutput.class, parDoEnforcements);
return enforcements.build();
}
+
}
////////////////////////////////////////////////////////////////////////////////////////////////
@@ -216,45 +223,42 @@ public class DirectRunner extends PipelineRunner<DirectPipelineResult> {
@SuppressWarnings("rawtypes")
@VisibleForTesting
List<PTransformOverride> defaultTransformOverrides() {
- DirectTestOptions testOptions = options.as(DirectTestOptions.class);
- ImmutableList.Builder<PTransformOverride> builder = ImmutableList.builder();
- if (testOptions.isRunnerDeterminedSharding()) {
- builder.add(
- PTransformOverride.of(
- PTransformMatchers.writeWithRunnerDeterminedSharding(),
- new WriteWithShardingFactory())); /* Uses a view internally. */
- }
- builder = builder.add(
- PTransformOverride.of(
- PTransformMatchers.urnEqualTo(PTransformTranslation.CREATE_VIEW_TRANSFORM_URN),
- new ViewOverrideFactory())) /* Uses pardos and GBKs */
+ return ImmutableList.<PTransformOverride>builder()
+ .add(
+ PTransformOverride.of(
+ PTransformMatchers.writeWithRunnerDeterminedSharding(),
+ new WriteWithShardingFactory())) /* Uses a view internally. */
+ .add(
+ PTransformOverride.of(
+ PTransformMatchers.classEqualTo(CreatePCollectionView.class),
+ new ViewOverrideFactory())) /* Uses pardos and GBKs */
.add(
PTransformOverride.of(
- PTransformMatchers.urnEqualTo(PTransformTranslation.TEST_STREAM_TRANSFORM_URN),
+ PTransformMatchers.classEqualTo(TestStream.class),
new DirectTestStreamFactory(this))) /* primitive */
// SplittableParMultiDo is implemented in terms of nonsplittable simple ParDos and extra
// primitives
.add(
PTransformOverride.of(
- PTransformMatchers.splittableParDo(), new ParDoMultiOverrideFactory()))
+ PTransformMatchers.splittableParDoMulti(), new ParDoMultiOverrideFactory()))
// state and timer pardos are implemented in terms of simple ParDos and extra primitives
.add(
PTransformOverride.of(
- PTransformMatchers.stateOrTimerParDo(), new ParDoMultiOverrideFactory()))
+ PTransformMatchers.stateOrTimerParDoMulti(), new ParDoMultiOverrideFactory()))
.add(
PTransformOverride.of(
- PTransformMatchers.urnEqualTo(
- SplittableParDo.SPLITTABLE_PROCESS_KEYED_ELEMENTS_URN),
+ PTransformMatchers.classEqualTo(SplittableParDo.ProcessKeyedElements.class),
new SplittableParDoViaKeyedWorkItems.OverrideFactory()))
.add(
PTransformOverride.of(
- PTransformMatchers.urnEqualTo(SplittableParDo.SPLITTABLE_GBKIKWI_URN),
+ PTransformMatchers.classEqualTo(
+ SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems.class),
new DirectGBKIntoKeyedWorkItemsOverrideFactory())) /* Returns a GBKO */
.add(
PTransformOverride.of(
- PTransformMatchers.urnEqualTo(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN),
- new DirectGroupByKeyOverrideFactory())); /* returns two chained primitives. */
- return builder.build();
+ PTransformMatchers.classEqualTo(GroupByKey.class),
+ new DirectGroupByKeyOverrideFactory())) /* returns two chained primitives. */
+ .build();
}
/**
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectTestOptions.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectTestOptions.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectTestOptions.java
deleted file mode 100644
index a426443..0000000
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectTestOptions.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.direct;
-
-import org.apache.beam.sdk.annotations.Internal;
-import org.apache.beam.sdk.options.ApplicationNameOptions;
-import org.apache.beam.sdk.options.Default;
-import org.apache.beam.sdk.options.Description;
-import org.apache.beam.sdk.options.Hidden;
-import org.apache.beam.sdk.options.PipelineOptions;
-
-/**
- * Internal-only options for tweaking the behavior of the {@link DirectRunner} in ways that users
- * should never do.
- *
- * <p>Currently, the only use is to disable user-friendly overrides that prevent fully testing
- * certain composite transforms.
- */
-@Internal
-@Hidden
-public interface DirectTestOptions extends PipelineOptions, ApplicationNameOptions {
- @Default.Boolean(true)
- @Description(
- "Indicates whether this is an automatically-run unit test.")
- boolean isRunnerDeterminedSharding();
- void setRunnerDeterminedSharding(boolean goAheadAndDetermineSharding);
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/EvaluationContext.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/EvaluationContext.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/EvaluationContext.java
index d192785..e215070 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/EvaluationContext.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/EvaluationContext.java
@@ -20,7 +20,6 @@ package org.apache.beam.runners.direct;
import static com.google.common.base.Preconditions.checkNotNull;
import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Optional;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.util.concurrent.MoreExecutors;
@@ -159,9 +158,12 @@ class EvaluationContext {
} else {
outputTypes.add(OutputType.BUNDLE);
}
- CommittedResult committedResult =
- CommittedResult.create(
- result, getUnprocessedInput(completedBundle, result), committedBundles, outputTypes);
+ CommittedResult committedResult = CommittedResult.create(result,
+ completedBundle == null
+ ? null
+ : completedBundle.withElements((Iterable) result.getUnprocessedElements()),
+ committedBundles,
+ outputTypes);
// Update state internals
CopyOnAccessInMemoryStateInternals theirState = result.getState();
if (theirState != null) {
@@ -185,22 +187,6 @@ class EvaluationContext {
return committedResult;
}
- /**
- * Returns an {@link Optional} containing a bundle which contains all of the unprocessed elements
- * that were not processed from the {@code completedBundle}. If all of the elements of the {@code
- * completedBundle} were processed, or if {@code completedBundle} is null, returns an absent
- * {@link Optional}.
- */
- private Optional<? extends CommittedBundle<?>> getUnprocessedInput(
- @Nullable CommittedBundle<?> completedBundle, TransformResult<?> result) {
- if (completedBundle == null || Iterables.isEmpty(result.getUnprocessedElements())) {
- return Optional.absent();
- }
- CommittedBundle<?> residual =
- completedBundle.withElements((Iterable) result.getUnprocessedElements());
- return Optional.of(residual);
- }
-
private Iterable<? extends CommittedBundle<?>> commitBundles(
Iterable<? extends UncommittedBundle<?>> bundles) {
ImmutableList.Builder<CommittedBundle<?>> completed = ImmutableList.builder();
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ExecutorServiceParallelExecutor.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ExecutorServiceParallelExecutor.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ExecutorServiceParallelExecutor.java
index 75e2562..71ab4cc 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ExecutorServiceParallelExecutor.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ExecutorServiceParallelExecutor.java
@@ -49,11 +49,11 @@ import javax.annotation.Nullable;
import org.apache.beam.runners.core.KeyedWorkItem;
import org.apache.beam.runners.core.KeyedWorkItems;
import org.apache.beam.runners.core.TimerInternals.TimerData;
-import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.runners.direct.WatermarkManager.FiredTimers;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineResult.State;
import org.apache.beam.sdk.runners.AppliedPTransform;
+import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.util.UserCodeException;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.PCollection;
@@ -77,7 +77,9 @@ final class ExecutorServiceParallelExecutor implements PipelineExecutor {
private final DirectGraph graph;
private final RootProviderRegistry rootProviderRegistry;
private final TransformEvaluatorRegistry registry;
- private final Map<String, Collection<ModelEnforcementFactory>> transformEnforcements;
+ @SuppressWarnings("rawtypes")
+ private final Map<Class<? extends PTransform>, Collection<ModelEnforcementFactory>>
+ transformEnforcements;
private final EvaluationContext evaluationContext;
@@ -110,7 +112,9 @@ final class ExecutorServiceParallelExecutor implements PipelineExecutor {
DirectGraph graph,
RootProviderRegistry rootProviderRegistry,
TransformEvaluatorRegistry registry,
- Map<String, Collection<ModelEnforcementFactory>> transformEnforcements,
+ @SuppressWarnings("rawtypes")
+ Map<Class<? extends PTransform>, Collection<ModelEnforcementFactory>>
+ transformEnforcements,
EvaluationContext context) {
return new ExecutorServiceParallelExecutor(
targetParallelism,
@@ -126,7 +130,8 @@ final class ExecutorServiceParallelExecutor implements PipelineExecutor {
DirectGraph graph,
RootProviderRegistry rootProviderRegistry,
TransformEvaluatorRegistry registry,
- Map<String, Collection<ModelEnforcementFactory>> transformEnforcements,
+ @SuppressWarnings("rawtypes")
+ Map<Class<? extends PTransform>, Collection<ModelEnforcementFactory>> transformEnforcements,
EvaluationContext context) {
this.targetParallelism = targetParallelism;
// Don't use Daemon threads for workers. The Pipeline should continue to execute even if there
@@ -232,8 +237,7 @@ final class ExecutorServiceParallelExecutor implements PipelineExecutor {
Collection<ModelEnforcementFactory> enforcements =
MoreObjects.firstNonNull(
- transformEnforcements.get(
- PTransformTranslation.urnForTransform(transform.getTransform())),
+ transformEnforcements.get(transform.getTransform().getClass()),
Collections.<ModelEnforcementFactory>emptyList());
TransformExecutor<T> callable =
@@ -351,18 +355,17 @@ final class ExecutorServiceParallelExecutor implements PipelineExecutor {
for (CommittedBundle<?> outputBundle : committedResult.getOutputs()) {
allUpdates.offer(
ExecutorUpdate.fromBundle(
- outputBundle, graph.getPerElementConsumers(outputBundle.getPCollection())));
+ outputBundle, graph.getPrimitiveConsumers(outputBundle.getPCollection())));
}
- Optional<? extends CommittedBundle<?>> unprocessedInputs =
- committedResult.getUnprocessedInputs();
- if (unprocessedInputs.isPresent()) {
+ CommittedBundle<?> unprocessedInputs = committedResult.getUnprocessedInputs();
+ if (unprocessedInputs != null && !Iterables.isEmpty(unprocessedInputs.getElements())) {
if (inputBundle.getPCollection() == null) {
// TODO: Split this logic out of an if statement
- pendingRootBundles.get(result.getTransform()).offer(unprocessedInputs.get());
+ pendingRootBundles.get(result.getTransform()).offer(unprocessedInputs);
} else {
allUpdates.offer(
ExecutorUpdate.fromBundle(
- unprocessedInputs.get(),
+ unprocessedInputs,
Collections.<AppliedPTransform<?, ?, ?>>singleton(
committedResult.getTransform())));
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoEvaluatorFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoEvaluatorFactory.java
index 516f798..8aa75cf 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoEvaluatorFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoEvaluatorFactory.java
@@ -20,6 +20,7 @@ package org.apache.beam.runners.direct;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
+import com.google.common.collect.Iterables;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@@ -78,7 +79,6 @@ final class ParDoEvaluatorFactory<InputT, OutputT> implements TransformEvaluator
(TransformEvaluator<T>)
createEvaluator(
(AppliedPTransform) application,
- (PCollection<InputT>) inputBundle.getPCollection(),
inputBundle.getKey(),
doFn,
transform.getSideInputs(),
@@ -102,7 +102,6 @@ final class ParDoEvaluatorFactory<InputT, OutputT> implements TransformEvaluator
@SuppressWarnings({"unchecked", "rawtypes"})
DoFnLifecycleManagerRemovingTransformEvaluator<InputT> createEvaluator(
AppliedPTransform<PCollection<InputT>, PCollectionTuple, ?> application,
- PCollection<InputT> mainInput,
StructuralKey<?> inputBundleKey,
DoFn<InputT, OutputT> doFn,
List<PCollectionView<?>> sideInputs,
@@ -121,7 +120,6 @@ final class ParDoEvaluatorFactory<InputT, OutputT> implements TransformEvaluator
createParDoEvaluator(
application,
inputBundleKey,
- mainInput,
sideInputs,
mainOutputTag,
additionalOutputTags,
@@ -134,7 +132,6 @@ final class ParDoEvaluatorFactory<InputT, OutputT> implements TransformEvaluator
ParDoEvaluator<InputT> createParDoEvaluator(
AppliedPTransform<PCollection<InputT>, PCollectionTuple, ?> application,
StructuralKey<?> key,
- PCollection<InputT> mainInput,
List<PCollectionView<?>> sideInputs,
TupleTag<OutputT> mainOutputTag,
List<TupleTag<?>> additionalOutputTags,
@@ -147,7 +144,8 @@ final class ParDoEvaluatorFactory<InputT, OutputT> implements TransformEvaluator
evaluationContext,
stepContext,
application,
- mainInput.getWindowingStrategy(),
+ ((PCollection<InputT>) Iterables.getOnlyElement(application.getInputs().values()))
+ .getWindowingStrategy(),
fn,
key,
sideInputs,
@@ -175,4 +173,5 @@ final class ParDoEvaluatorFactory<InputT, OutputT> implements TransformEvaluator
}
return pcs;
}
+
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/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 891d102..858ea34 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
@@ -19,17 +19,15 @@ package org.apache.beam.runners.direct;
import static com.google.common.base.Preconditions.checkState;
-import java.io.IOException;
-import java.util.List;
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.construction.PTransformReplacements;
import org.apache.beam.runners.core.construction.PTransformTranslation;
-import org.apache.beam.runners.core.construction.ParDoTranslation;
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;
import org.apache.beam.sdk.runners.AppliedPTransform;
@@ -38,6 +36,7 @@ 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.ParDo.MultiOutput;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature;
import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
import org.apache.beam.sdk.transforms.windowing.AfterPane;
@@ -49,8 +48,6 @@ 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.PCollectionViews;
import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
@@ -63,48 +60,36 @@ import org.apache.beam.sdk.values.WindowingStrategy;
*/
class ParDoMultiOverrideFactory<InputT, OutputT>
implements PTransformOverrideFactory<
- PCollection<? extends InputT>, PCollectionTuple,
- PTransform<PCollection<? extends InputT>, PCollectionTuple>> {
+ PCollection<? extends InputT>, PCollectionTuple, MultiOutput<InputT, OutputT>> {
@Override
public PTransformReplacement<PCollection<? extends InputT>, PCollectionTuple>
getReplacementTransform(
AppliedPTransform<
- PCollection<? extends InputT>, PCollectionTuple,
- PTransform<PCollection<? extends InputT>, PCollectionTuple>>
- application) {
-
- try {
- return PTransformReplacement.of(
- PTransformReplacements.getSingletonMainInput(application),
- getReplacementForApplication(application));
- } catch (IOException exc) {
- throw new RuntimeException(exc);
- }
+ PCollection<? extends InputT>, PCollectionTuple, MultiOutput<InputT, OutputT>>
+ transform) {
+ return PTransformReplacement.of(
+ PTransformReplacements.getSingletonMainInput(transform),
+ getReplacementTransform(transform.getTransform()));
}
@SuppressWarnings("unchecked")
- private PTransform<PCollection<? extends InputT>, PCollectionTuple> getReplacementForApplication(
- AppliedPTransform<
- PCollection<? extends InputT>, PCollectionTuple,
- PTransform<PCollection<? extends InputT>, PCollectionTuple>>
- application)
- throws IOException {
-
- DoFn<InputT, OutputT> fn = (DoFn<InputT, OutputT>) ParDoTranslation.getDoFn(application);
+ private PTransform<PCollection<? extends InputT>, PCollectionTuple> getReplacementTransform(
+ MultiOutput<InputT, OutputT> transform) {
+ DoFn<InputT, OutputT> fn = transform.getFn();
DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());
-
if (signature.processElement().isSplittable()) {
- return (PTransform) SplittableParDo.forAppliedParDo(application);
+ return new SplittableParDo(transform);
} else if (signature.stateDeclarations().size() > 0
|| signature.timerDeclarations().size() > 0) {
- return new GbkThenStatefulParDo(
- fn,
- ParDoTranslation.getMainOutputTag(application),
- ParDoTranslation.getAdditionalOutputTags(application),
- ParDoTranslation.getSideInputs(application));
+ // Based on the fact that the signature is stateful, DoFnSignatures ensures
+ // that it is also keyed
+ MultiOutput<KV<?, ?>, OutputT> keyedTransform =
+ (MultiOutput<KV<?, ?>, OutputT>) transform;
+
+ return new GbkThenStatefulParDo(keyedTransform);
} else {
- return application.getTransform();
+ return transform;
}
}
@@ -116,25 +101,10 @@ class ParDoMultiOverrideFactory<InputT, OutputT>
static class GbkThenStatefulParDo<K, InputT, OutputT>
extends PTransform<PCollection<KV<K, InputT>>, PCollectionTuple> {
- private final transient DoFn<KV<K, InputT>, OutputT> doFn;
- private final TupleTagList additionalOutputTags;
- private final TupleTag<OutputT> mainOutputTag;
- private final List<PCollectionView<?>> sideInputs;
-
- public GbkThenStatefulParDo(
- DoFn<KV<K, InputT>, OutputT> doFn,
- TupleTag<OutputT> mainOutputTag,
- TupleTagList additionalOutputTags,
- List<PCollectionView<?>> sideInputs) {
- this.doFn = doFn;
- this.additionalOutputTags = additionalOutputTags;
- this.mainOutputTag = mainOutputTag;
- this.sideInputs = sideInputs;
- }
+ private final MultiOutput<KV<K, InputT>, OutputT> underlyingParDo;
- @Override
- public Map<TupleTag<?>, PValue> getAdditionalInputs() {
- return PCollectionViews.toAdditionalInputs(sideInputs);
+ public GbkThenStatefulParDo(MultiOutput<KV<K, InputT>, OutputT> underlyingParDo) {
+ this.underlyingParDo = underlyingParDo;
}
@Override
@@ -190,9 +160,7 @@ class ParDoMultiOverrideFactory<InputT, OutputT>
adjustedInput
// Explode the resulting iterable into elements that are exactly the ones from
// the input
- .apply(
- "Stateful ParDo",
- new StatefulParDo<>(doFn, mainOutputTag, additionalOutputTags, sideInputs));
+ .apply("Stateful ParDo", new StatefulParDo<>(underlyingParDo, input));
return outputs;
}
@@ -204,41 +172,25 @@ class ParDoMultiOverrideFactory<InputT, OutputT>
static class StatefulParDo<K, InputT, OutputT>
extends PTransformTranslation.RawPTransform<
PCollection<? extends KeyedWorkItem<K, KV<K, InputT>>>, PCollectionTuple> {
- private final transient DoFn<KV<K, InputT>, OutputT> doFn;
- private final TupleTagList additionalOutputTags;
- private final TupleTag<OutputT> mainOutputTag;
- private final List<PCollectionView<?>> sideInputs;
+ private final transient MultiOutput<KV<K, InputT>, OutputT> underlyingParDo;
+ private final transient PCollection<KV<K, InputT>> originalInput;
public StatefulParDo(
- DoFn<KV<K, InputT>, OutputT> doFn,
- TupleTag<OutputT> mainOutputTag,
- TupleTagList additionalOutputTags,
- List<PCollectionView<?>> sideInputs) {
- this.doFn = doFn;
- this.mainOutputTag = mainOutputTag;
- this.additionalOutputTags = additionalOutputTags;
- this.sideInputs = sideInputs;
- }
-
- public DoFn<KV<K, InputT>, OutputT> getDoFn() {
- return doFn;
- }
-
- public TupleTag<OutputT> getMainOutputTag() {
- return mainOutputTag;
- }
-
- public List<PCollectionView<?>> getSideInputs() {
- return sideInputs;
+ MultiOutput<KV<K, InputT>, OutputT> underlyingParDo,
+ PCollection<KV<K, InputT>> originalInput) {
+ this.underlyingParDo = underlyingParDo;
+ this.originalInput = originalInput;
}
- public TupleTagList getAdditionalOutputTags() {
- return additionalOutputTags;
+ public MultiOutput<KV<K, InputT>, OutputT> getUnderlyingParDo() {
+ return underlyingParDo;
}
@Override
- public Map<TupleTag<?>, PValue> getAdditionalInputs() {
- return PCollectionViews.toAdditionalInputs(sideInputs);
+ public <T> Coder<T> getDefaultOutputCoder(
+ PCollection<? extends KeyedWorkItem<K, KV<K, InputT>>> input, PCollection<T> output)
+ throws CannotProvideCoderException {
+ return underlyingParDo.getDefaultOutputCoder(originalInput, output);
}
@Override
@@ -247,7 +199,8 @@ class ParDoMultiOverrideFactory<InputT, OutputT>
PCollectionTuple outputs =
PCollectionTuple.ofPrimitiveOutputsInternal(
input.getPipeline(),
- TupleTagList.of(getMainOutputTag()).and(getAdditionalOutputTags().getAll()),
+ TupleTagList.of(underlyingParDo.getMainOutputTag())
+ .and(underlyingParDo.getAdditionalOutputTags().getAll()),
input.getWindowingStrategy(),
input.isBounded());
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/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 e6b51b7..b85f481c 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
@@ -35,6 +35,7 @@ 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;
@@ -42,7 +43,6 @@ import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
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;
@@ -54,7 +54,8 @@ import org.joda.time.Instant;
class SplittableProcessElementsEvaluatorFactory<
InputT, OutputT, RestrictionT, TrackerT extends RestrictionTracker<RestrictionT>>
implements TransformEvaluatorFactory {
- private final ParDoEvaluatorFactory<KeyedWorkItem<String, KV<InputT, RestrictionT>>, OutputT>
+ private final ParDoEvaluatorFactory<
+ KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT>
delegateFactory;
private final EvaluationContext evaluationContext;
@@ -83,13 +84,14 @@ class SplittableProcessElementsEvaluatorFactory<
}
@SuppressWarnings({"unchecked", "rawtypes"})
- private TransformEvaluator<KeyedWorkItem<String, KV<InputT, RestrictionT>>> createEvaluator(
- AppliedPTransform<
- PCollection<KeyedWorkItem<String, KV<InputT, RestrictionT>>>, PCollectionTuple,
- ProcessElements<InputT, OutputT, RestrictionT, TrackerT>>
- application,
- CommittedBundle<InputT> inputBundle)
- throws Exception {
+ private TransformEvaluator<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>>
+ createEvaluator(
+ AppliedPTransform<
+ PCollection<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>>,
+ PCollectionTuple, ProcessElements<InputT, OutputT, RestrictionT, TrackerT>>
+ application,
+ CommittedBundle<InputT> inputBundle)
+ throws Exception {
final ProcessElements<InputT, OutputT, RestrictionT, TrackerT> transform =
application.getTransform();
@@ -99,7 +101,9 @@ class SplittableProcessElementsEvaluatorFactory<
DoFnLifecycleManager fnManager = DoFnLifecycleManager.of(processFn);
processFn =
((ProcessFn<InputT, OutputT, RestrictionT, TrackerT>)
- fnManager.<KeyedWorkItem<String, KV<InputT, RestrictionT>>, OutputT>get());
+ fnManager
+ .<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT>
+ get());
String stepName = evaluationContext.getStepName(application);
final DirectExecutionContext.DirectStepContext stepContext =
@@ -107,13 +111,11 @@ class SplittableProcessElementsEvaluatorFactory<
.getExecutionContext(application, inputBundle.getKey())
.getStepContext(stepName);
- final ParDoEvaluator<KeyedWorkItem<String, KV<InputT, RestrictionT>>>
+ final ParDoEvaluator<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>>
parDoEvaluator =
delegateFactory.createParDoEvaluator(
application,
inputBundle.getKey(),
- (PCollection<KeyedWorkItem<String, KV<InputT, RestrictionT>>>)
- inputBundle.getPCollection(),
transform.getSideInputs(),
transform.getMainOutputTag(),
transform.getAdditionalOutputTags().getAll(),
@@ -185,16 +187,17 @@ class SplittableProcessElementsEvaluatorFactory<
}
private static <InputT, OutputT, RestrictionT>
- ParDoEvaluator.DoFnRunnerFactory<KeyedWorkItem<String, KV<InputT, RestrictionT>>, OutputT>
+ ParDoEvaluator.DoFnRunnerFactory<
+ KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT>
processFnRunnerFactory() {
return new ParDoEvaluator.DoFnRunnerFactory<
- KeyedWorkItem<String, KV<InputT, RestrictionT>>, OutputT>() {
+ KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT>() {
@Override
public PushbackSideInputDoFnRunner<
- KeyedWorkItem<String, KV<InputT, RestrictionT>>, OutputT>
+ KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT>
createRunner(
PipelineOptions options,
- DoFn<KeyedWorkItem<String, KV<InputT, RestrictionT>>, OutputT> fn,
+ DoFn<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT> fn,
List<PCollectionView<?>> sideInputs,
ReadyCheckingSideInputReader sideInputReader,
OutputManager outputManager,
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactory.java
index bdec9c8..506c84c 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactory.java
@@ -98,7 +98,7 @@ final class StatefulParDoEvaluatorFactory<K, InputT, OutputT> implements Transfo
throws Exception {
final DoFn<KV<K, InputT>, OutputT> doFn =
- application.getTransform().getDoFn();
+ application.getTransform().getUnderlyingParDo().getFn();
final DoFnSignature signature = DoFnSignatures.getSignature(doFn.getClass());
// If the DoFn is stateful, schedule state clearing.
@@ -117,12 +117,11 @@ final class StatefulParDoEvaluatorFactory<K, InputT, OutputT> implements Transfo
DoFnLifecycleManagerRemovingTransformEvaluator<KV<K, InputT>> delegateEvaluator =
delegateFactory.createEvaluator(
(AppliedPTransform) application,
- (PCollection) inputBundle.getPCollection(),
inputBundle.getKey(),
doFn,
- application.getTransform().getSideInputs(),
- application.getTransform().getMainOutputTag(),
- application.getTransform().getAdditionalOutputTags().getAll());
+ application.getTransform().getUnderlyingParDo().getSideInputs(),
+ application.getTransform().getUnderlyingParDo().getMainOutputTag(),
+ application.getTransform().getUnderlyingParDo().getAdditionalOutputTags().getAll());
return new StatefulParDoEvaluator<>(delegateEvaluator);
}
@@ -152,11 +151,12 @@ final class StatefulParDoEvaluatorFactory<K, InputT, OutputT> implements Transfo
transformOutputWindow
.getTransform()
.getTransform()
+ .getUnderlyingParDo()
.getMainOutputTag());
WindowingStrategy<?, ?> windowingStrategy = pc.getWindowingStrategy();
BoundedWindow window = transformOutputWindow.getWindow();
final DoFn<?, ?> doFn =
- transformOutputWindow.getTransform().getTransform().getDoFn();
+ transformOutputWindow.getTransform().getTransform().getUnderlyingParDo().getFn();
final DoFnSignature signature = DoFnSignatures.getSignature(doFn.getClass());
final DirectStepContext stepContext =
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactory.java
index 16c8589..2da7a71 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactory.java
@@ -22,7 +22,6 @@ import com.google.auto.value.AutoValue;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Supplier;
import com.google.common.collect.Iterables;
-import java.io.IOException;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
@@ -31,7 +30,6 @@ import java.util.concurrent.atomic.AtomicReference;
import javax.annotation.Nullable;
import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.runners.core.construction.ReplacementOutputs;
-import org.apache.beam.runners.core.construction.TestStreamTranslation;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.runners.PTransformOverrideFactory;
import org.apache.beam.sdk.testing.TestStream;
@@ -162,8 +160,7 @@ class TestStreamEvaluatorFactory implements TransformEvaluatorFactory {
}
static class DirectTestStreamFactory<T>
- implements PTransformOverrideFactory<
- PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>> {
+ implements PTransformOverrideFactory<PBegin, PCollection<T>, TestStream<T>> {
private final DirectRunner runner;
DirectTestStreamFactory(DirectRunner runner) {
@@ -172,17 +169,10 @@ class TestStreamEvaluatorFactory implements TransformEvaluatorFactory {
@Override
public PTransformReplacement<PBegin, PCollection<T>> getReplacementTransform(
- AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>> transform) {
- try {
- return PTransformReplacement.of(
- transform.getPipeline().begin(),
- new DirectTestStream<T>(runner, TestStreamTranslation.getTestStream(transform)));
- } catch (IOException exc) {
- throw new RuntimeException(
- String.format(
- "Transform could not be converted to %s", TestStream.class.getSimpleName()),
- exc);
- }
+ AppliedPTransform<PBegin, PCollection<T>, TestStream<T>> transform) {
+ return PTransformReplacement.of(
+ transform.getPipeline().begin(),
+ new DirectTestStream<T>(runner, transform.getTransform()));
}
@Override
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ViewEvaluatorFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ViewEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ViewEvaluatorFactory.java
index 8a281a7..057f4a1 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ViewEvaluatorFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ViewEvaluatorFactory.java
@@ -28,6 +28,7 @@ import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.View.CreatePCollectionView;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
/**
* The {@link DirectRunner} {@link TransformEvaluatorFactory} for the {@link CreatePCollectionView}
@@ -59,13 +60,12 @@ class ViewEvaluatorFactory implements TransformEvaluatorFactory {
public void cleanup() throws Exception {}
private <InT, OuT> TransformEvaluator<Iterable<InT>> createEvaluator(
- final AppliedPTransform<
- PCollection<Iterable<InT>>, PCollection<Iterable<InT>>, WriteView<InT, OuT>>
+ final AppliedPTransform<PCollection<Iterable<InT>>, PCollectionView<OuT>, WriteView<InT, OuT>>
application) {
PCollection<Iterable<InT>> input =
(PCollection<Iterable<InT>>) Iterables.getOnlyElement(application.getInputs().values());
- final PCollectionViewWriter<InT, OuT> writer =
- context.createPCollectionViewWriter(input, application.getTransform().getView());
+ final PCollectionViewWriter<InT, OuT> writer = context.createPCollectionViewWriter(input,
+ (PCollectionView<OuT>) Iterables.getOnlyElement(application.getOutputs().values()));
return new TransformEvaluator<Iterable<InT>>() {
private final List<WindowedValue<InT>> elements = new ArrayList<>();
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ViewOverrideFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ViewOverrideFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ViewOverrideFactory.java
index 5dcf016..fdff63d 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ViewOverrideFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ViewOverrideFactory.java
@@ -18,12 +18,11 @@
package org.apache.beam.runners.direct;
-import java.io.IOException;
+import java.util.Collections;
import java.util.Map;
-import org.apache.beam.runners.core.construction.CreatePCollectionViewTranslation;
+import org.apache.beam.runners.core.construction.ForwardingPTransform;
import org.apache.beam.runners.core.construction.PTransformReplacements;
import org.apache.beam.runners.core.construction.PTransformTranslation.RawPTransform;
-import org.apache.beam.runners.core.construction.ReplacementOutputs;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.runners.AppliedPTransform;
@@ -44,56 +43,46 @@ import org.apache.beam.sdk.values.TupleTag;
*/
class ViewOverrideFactory<ElemT, ViewT>
implements PTransformOverrideFactory<
- PCollection<ElemT>, PCollection<ElemT>,
- PTransform<PCollection<ElemT>, PCollection<ElemT>>> {
+ PCollection<ElemT>, PCollectionView<ViewT>, CreatePCollectionView<ElemT, ViewT>> {
@Override
- public PTransformReplacement<PCollection<ElemT>, PCollection<ElemT>> getReplacementTransform(
+ public PTransformReplacement<PCollection<ElemT>, PCollectionView<ViewT>> getReplacementTransform(
AppliedPTransform<
- PCollection<ElemT>, PCollection<ElemT>,
- PTransform<PCollection<ElemT>, PCollection<ElemT>>>
+ PCollection<ElemT>, PCollectionView<ViewT>, CreatePCollectionView<ElemT, ViewT>>
transform) {
-
- PCollectionView<ViewT> view;
- try {
- view = CreatePCollectionViewTranslation.getView(transform);
- } catch (IOException exc) {
- throw new RuntimeException(
- String.format(
- "Could not extract %s from transform %s",
- PCollectionView.class.getSimpleName(), transform),
- exc);
- }
-
- return PTransformReplacement.of(
+ return PTransformReplacement.of(
PTransformReplacements.getSingletonMainInput(transform),
- new GroupAndWriteView<ElemT, ViewT>(view));
+ new GroupAndWriteView<>(transform.getTransform()));
}
@Override
public Map<PValue, ReplacementOutput> mapOutputs(
- Map<TupleTag<?>, PValue> outputs, PCollection<ElemT> newOutput) {
- return ReplacementOutputs.singleton(outputs, newOutput);
+ Map<TupleTag<?>, PValue> outputs, PCollectionView<ViewT> newOutput) {
+ return Collections.emptyMap();
}
/** The {@link DirectRunner} composite override for {@link CreatePCollectionView}. */
static class GroupAndWriteView<ElemT, ViewT>
- extends PTransform<PCollection<ElemT>, PCollection<ElemT>> {
- private final PCollectionView<ViewT> view;
+ extends ForwardingPTransform<PCollection<ElemT>, PCollectionView<ViewT>> {
+ private final CreatePCollectionView<ElemT, ViewT> og;
- private GroupAndWriteView(PCollectionView<ViewT> view) {
- this.view = view;
+ private GroupAndWriteView(CreatePCollectionView<ElemT, ViewT> og) {
+ this.og = og;
}
@Override
- public PCollection<ElemT> expand(final PCollection<ElemT> input) {
- input
+ public PCollectionView<ViewT> expand(PCollection<ElemT> input) {
+ return input
.apply(WithKeys.<Void, ElemT>of((Void) null))
.setCoder(KvCoder.of(VoidCoder.of(), input.getCoder()))
.apply(GroupByKey.<Void, ElemT>create())
.apply(Values.<Iterable<ElemT>>create())
- .apply(new WriteView<ElemT, ViewT>(view));
- return input;
+ .apply(new WriteView<ElemT, ViewT>(og));
+ }
+
+ @Override
+ protected PTransform<PCollection<ElemT>, PCollectionView<ViewT>> delegate() {
+ return og;
}
}
@@ -105,24 +94,22 @@ class ViewOverrideFactory<ElemT, ViewT>
* to {@link ViewT}.
*/
static final class WriteView<ElemT, ViewT>
- extends RawPTransform<PCollection<Iterable<ElemT>>, PCollection<Iterable<ElemT>>> {
- private final PCollectionView<ViewT> view;
+ extends RawPTransform<PCollection<Iterable<ElemT>>, PCollectionView<ViewT>> {
+ private final CreatePCollectionView<ElemT, ViewT> og;
- WriteView(PCollectionView<ViewT> view) {
- this.view = view;
+ WriteView(CreatePCollectionView<ElemT, ViewT> og) {
+ this.og = og;
}
@Override
@SuppressWarnings("deprecation")
- public PCollection<Iterable<ElemT>> expand(PCollection<Iterable<ElemT>> input) {
- return PCollection.<Iterable<ElemT>>createPrimitiveOutputInternal(
- input.getPipeline(), input.getWindowingStrategy(), input.isBounded())
- .setCoder(input.getCoder());
+ public PCollectionView<ViewT> expand(PCollection<Iterable<ElemT>> input) {
+ return og.getView();
}
@SuppressWarnings("deprecation")
public PCollectionView<ViewT> getView() {
- return view;
+ return og.getView();
}
@Override
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java
index 599b74f..40ce163 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WatermarkManager.java
@@ -54,7 +54,6 @@ import javax.annotation.concurrent.GuardedBy;
import org.apache.beam.runners.core.StateNamespace;
import org.apache.beam.runners.core.TimerInternals;
import org.apache.beam.runners.core.TimerInternals.TimerData;
-import org.apache.beam.runners.core.construction.TransformInputs;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.state.TimeDomain;
@@ -63,6 +62,7 @@ import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.joda.time.Instant;
/**
@@ -831,11 +831,11 @@ class WatermarkManager {
private Collection<Watermark> getInputProcessingWatermarks(AppliedPTransform<?, ?, ?> transform) {
ImmutableList.Builder<Watermark> inputWmsBuilder = ImmutableList.builder();
- Collection<PValue> inputs = TransformInputs.nonAdditionalInputs(transform);
+ Map<TupleTag<?>, PValue> inputs = transform.getInputs();
if (inputs.isEmpty()) {
inputWmsBuilder.add(THE_END_OF_TIME);
}
- for (PValue pvalue : inputs) {
+ for (PValue pvalue : inputs.values()) {
Watermark producerOutputWatermark =
getValueWatermark(pvalue).synchronizedProcessingOutputWatermark;
inputWmsBuilder.add(producerOutputWatermark);
@@ -845,11 +845,11 @@ class WatermarkManager {
private List<Watermark> getInputWatermarks(AppliedPTransform<?, ?, ?> transform) {
ImmutableList.Builder<Watermark> inputWatermarksBuilder = ImmutableList.builder();
- Collection< PValue> inputs = TransformInputs.nonAdditionalInputs(transform);
+ Map<TupleTag<?>, PValue> inputs = transform.getInputs();
if (inputs.isEmpty()) {
inputWatermarksBuilder.add(THE_END_OF_TIME);
}
- for (PValue pvalue : inputs) {
+ for (PValue pvalue : inputs.values()) {
Watermark producerOutputWatermark = getValueWatermark(pvalue).outputWatermark;
inputWatermarksBuilder.add(producerOutputWatermark);
}
@@ -987,16 +987,16 @@ class WatermarkManager {
// refresh.
for (CommittedBundle<?> bundle : result.getOutputs()) {
for (AppliedPTransform<?, ?, ?> consumer :
- graph.getPerElementConsumers(bundle.getPCollection())) {
+ graph.getPrimitiveConsumers(bundle.getPCollection())) {
TransformWatermarks watermarks = transformToWatermarks.get(consumer);
watermarks.addPending(bundle);
}
}
TransformWatermarks completedTransform = transformToWatermarks.get(result.getTransform());
- if (result.getUnprocessedInputs().isPresent()) {
+ if (input != null) {
// Add the unprocessed inputs
- completedTransform.addPending(result.getUnprocessedInputs().get());
+ completedTransform.addPending(result.getUnprocessedInputs());
}
completedTransform.updateTimers(timerUpdate);
if (input != null) {
@@ -1035,7 +1035,7 @@ class WatermarkManager {
if (updateResult.isAdvanced()) {
Set<AppliedPTransform<?, ?, ?>> additionalRefreshes = new HashSet<>();
for (PValue outputPValue : toRefresh.getOutputs().values()) {
- additionalRefreshes.addAll(graph.getPerElementConsumers(outputPValue));
+ additionalRefreshes.addAll(graph.getPrimitiveConsumers(outputPValue));
}
return additionalRefreshes;
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WriteWithShardingFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WriteWithShardingFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WriteWithShardingFactory.java
index ba796ae..65a5a19 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WriteWithShardingFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WriteWithShardingFactory.java
@@ -21,13 +21,11 @@ package org.apache.beam.runners.direct;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Supplier;
import com.google.common.base.Suppliers;
-import java.io.IOException;
import java.io.Serializable;
import java.util.Collections;
import java.util.Map;
import java.util.concurrent.ThreadLocalRandom;
import org.apache.beam.runners.core.construction.PTransformReplacements;
-import org.apache.beam.runners.core.construction.WriteFilesTranslation;
import org.apache.beam.sdk.io.WriteFiles;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.runners.PTransformOverrideFactory;
@@ -45,35 +43,23 @@ import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TupleTag;
/**
- * A {@link PTransformOverrideFactory} that overrides {@link WriteFiles} {@link PTransform
- * PTransforms} with an unspecified number of shards with a write with a specified number of shards.
- * The number of shards is the log base 10 of the number of input records, with up to 2 additional
- * shards.
+ * A {@link PTransformOverrideFactory} that overrides {@link WriteFiles}
+ * {@link PTransform PTransforms} with an unspecified number of shards with a write with a
+ * specified number of shards. The number of shards is the log base 10 of the number of input
+ * records, with up to 2 additional shards.
*/
class WriteWithShardingFactory<InputT>
- implements PTransformOverrideFactory<
- PCollection<InputT>, PDone, PTransform<PCollection<InputT>, PDone>> {
+ implements PTransformOverrideFactory<PCollection<InputT>, PDone, WriteFiles<InputT>> {
static final int MAX_RANDOM_EXTRA_SHARDS = 3;
@VisibleForTesting static final int MIN_SHARDS_FOR_LOG = 3;
@Override
public PTransformReplacement<PCollection<InputT>, PDone> getReplacementTransform(
- AppliedPTransform<PCollection<InputT>, PDone, PTransform<PCollection<InputT>, PDone>>
- transform) {
- try {
- WriteFiles<InputT, ?, ?> replacement =
- WriteFiles.to(
- WriteFilesTranslation.getSink(transform),
- WriteFilesTranslation.getFormatFunction(transform));
- if (WriteFilesTranslation.isWindowedWrites(transform)) {
- replacement = replacement.withWindowedWrites();
- }
- return PTransformReplacement.of(
- PTransformReplacements.getSingletonMainInput(transform),
- replacement.withSharding(new LogElementShardsWithDrift<InputT>()));
- } catch (IOException e) {
- throw new RuntimeException(e);
- }
+ AppliedPTransform<PCollection<InputT>, PDone, WriteFiles<InputT>> transform) {
+
+ return PTransformReplacement.of(
+ PTransformReplacements.getSingletonMainInput(transform),
+ transform.getTransform().withSharding(new LogElementShardsWithDrift<InputT>()));
}
@Override
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/CommittedResultTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/CommittedResultTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/CommittedResultTest.java
index 8b95b34..cf19dc2 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/CommittedResultTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/CommittedResultTest.java
@@ -18,9 +18,9 @@
package org.apache.beam.runners.direct;
+import static org.hamcrest.Matchers.nullValue;
import static org.junit.Assert.assertThat;
-import com.google.common.base.Optional;
import com.google.common.collect.ImmutableList;
import java.io.Serializable;
import java.util.Collections;
@@ -72,7 +72,7 @@ public class CommittedResultTest implements Serializable {
CommittedResult result =
CommittedResult.create(
StepTransformResult.withoutHold(transform).build(),
- Optional.<CommittedBundle<?>>absent(),
+ bundleFactory.createBundle(created).commit(Instant.now()),
Collections.<CommittedBundle<?>>emptyList(),
EnumSet.noneOf(OutputType.class));
@@ -88,11 +88,11 @@ public class CommittedResultTest implements Serializable {
CommittedResult result =
CommittedResult.create(
StepTransformResult.withoutHold(transform).build(),
- Optional.of(bundle),
+ bundle,
Collections.<CommittedBundle<?>>emptyList(),
EnumSet.noneOf(OutputType.class));
- assertThat(result.getUnprocessedInputs().get(),
+ assertThat(result.getUnprocessedInputs(),
Matchers.<CommittedBundle<?>>equalTo(bundle));
}
@@ -101,14 +101,11 @@ public class CommittedResultTest implements Serializable {
CommittedResult result =
CommittedResult.create(
StepTransformResult.withoutHold(transform).build(),
- Optional.<CommittedBundle<?>>absent(),
+ null,
Collections.<CommittedBundle<?>>emptyList(),
EnumSet.noneOf(OutputType.class));
- assertThat(
- result.getUnprocessedInputs(),
- Matchers.<Optional<? extends CommittedBundle<?>>>equalTo(
- Optional.<CommittedBundle<?>>absent()));
+ assertThat(result.getUnprocessedInputs(), nullValue());
}
@Test
@@ -123,7 +120,7 @@ public class CommittedResultTest implements Serializable {
CommittedResult result =
CommittedResult.create(
StepTransformResult.withoutHold(transform).build(),
- Optional.<CommittedBundle<?>>absent(),
+ bundleFactory.createBundle(created).commit(Instant.now()),
outputs,
EnumSet.of(OutputType.BUNDLE, OutputType.PCOLLECTION_VIEW));
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectGraphVisitorTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectGraphVisitorTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectGraphVisitorTest.java
index bf3e83e..576edf3 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectGraphVisitorTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectGraphVisitorTest.java
@@ -151,13 +151,13 @@ public class DirectGraphVisitorTest implements Serializable {
graph.getProducer(flattened);
assertThat(
- graph.getPerElementConsumers(created),
+ graph.getPrimitiveConsumers(created),
Matchers.<AppliedPTransform<?, ?, ?>>containsInAnyOrder(
transformedProducer, flattenedProducer));
assertThat(
- graph.getPerElementConsumers(transformed),
+ graph.getPrimitiveConsumers(transformed),
Matchers.<AppliedPTransform<?, ?, ?>>containsInAnyOrder(flattenedProducer));
- assertThat(graph.getPerElementConsumers(flattened), emptyIterable());
+ assertThat(graph.getPrimitiveConsumers(flattened), emptyIterable());
}
@Test
@@ -173,10 +173,10 @@ public class DirectGraphVisitorTest implements Serializable {
AppliedPTransform<?, ?, ?> flattenedProducer = graph.getProducer(flattened);
assertThat(
- graph.getPerElementConsumers(created),
+ graph.getPrimitiveConsumers(created),
Matchers.<AppliedPTransform<?, ?, ?>>containsInAnyOrder(flattenedProducer,
flattenedProducer));
- assertThat(graph.getPerElementConsumers(flattened), emptyIterable());
+ assertThat(graph.getPrimitiveConsumers(flattened), emptyIterable());
}
@Test
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectGraphs.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectGraphs.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectGraphs.java
index 7707f7f..43de091 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectGraphs.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectGraphs.java
@@ -18,7 +18,6 @@
package org.apache.beam.runners.direct;
import org.apache.beam.sdk.Pipeline;
-import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;
@@ -26,12 +25,6 @@ import org.apache.beam.sdk.values.PValue;
/** Test utilities for the {@link DirectRunner}. */
final class DirectGraphs {
- public static void performDirectOverrides(Pipeline p) {
- p.replaceAll(
- DirectRunner.fromOptions(PipelineOptionsFactory.create().as(DirectOptions.class))
- .defaultTransformOverrides());
- }
-
public static DirectGraph getGraph(Pipeline p) {
DirectGraphVisitor visitor = new DirectGraphVisitor();
p.traverseTopologically(visitor);
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectRegistrarTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectRegistrarTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectRegistrarTest.java
index 4b909bc..603e43e 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectRegistrarTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectRegistrarTest.java
@@ -37,7 +37,7 @@ public class DirectRegistrarTest {
@Test
public void testCorrectOptionsAreReturned() {
assertEquals(
- ImmutableList.of(DirectOptions.class, DirectTestOptions.class),
+ ImmutableList.of(DirectOptions.class),
new Options().getPipelineOptions());
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EvaluationContextTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EvaluationContextTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EvaluationContextTest.java
index 699a318..c0e43d6 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EvaluationContextTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/EvaluationContextTest.java
@@ -101,13 +101,14 @@ public class EvaluationContextTest {
view = created.apply(View.<Integer>asIterable());
unbounded = p.apply(GenerateSequence.from(0));
- p.replaceAll(runner.defaultTransformOverrides());
+ p.replaceAll(
+ DirectRunner.fromOptions(TestPipeline.testingPipelineOptions())
+ .defaultTransformOverrides());
KeyedPValueTrackingVisitor keyedPValueTrackingVisitor = KeyedPValueTrackingVisitor.create();
p.traverseTopologically(keyedPValueTrackingVisitor);
BundleFactory bundleFactory = ImmutableListBundleFactory.create();
- DirectGraphs.performDirectOverrides(p);
graph = DirectGraphs.getGraph(p);
context =
EvaluationContext.create(
@@ -414,7 +415,7 @@ public class EvaluationContextTest {
StepTransformResult.withoutHold(unboundedProducer).build());
assertThat(context.isDone(), is(false));
- for (AppliedPTransform<?, ?, ?> consumers : graph.getPerElementConsumers(created)) {
+ for (AppliedPTransform<?, ?, ?> consumers : graph.getPrimitiveConsumers(created)) {
context.handleResult(
committedBundle,
ImmutableList.<TimerData>of(),
[06/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/SimplifiedKinesisClient.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/SimplifiedKinesisClient.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/SimplifiedKinesisClient.java
index 80c950f..3e3984a 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/SimplifiedKinesisClient.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/SimplifiedKinesisClient.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.sdk.io.kinesis;
+
import com.amazonaws.AmazonServiceException;
import com.amazonaws.services.kinesis.AmazonKinesis;
import com.amazonaws.services.kinesis.clientlibrary.types.UserRecord;
@@ -30,11 +31,9 @@ import com.amazonaws.services.kinesis.model.Shard;
import com.amazonaws.services.kinesis.model.ShardIteratorType;
import com.amazonaws.services.kinesis.model.StreamDescription;
import com.google.common.collect.Lists;
-
import java.util.Date;
import java.util.List;
import java.util.concurrent.Callable;
-
import org.joda.time.Instant;
/**
@@ -42,121 +41,117 @@ import org.joda.time.Instant;
* proper error handling.
*/
class SimplifiedKinesisClient {
+ private final AmazonKinesis kinesis;
- private final AmazonKinesis kinesis;
-
- public SimplifiedKinesisClient(AmazonKinesis kinesis) {
- this.kinesis = kinesis;
- }
-
- public static SimplifiedKinesisClient from(KinesisClientProvider provider) {
- return new SimplifiedKinesisClient(provider.get());
- }
-
- public String getShardIterator(final String streamName, final String shardId,
- final ShardIteratorType shardIteratorType,
- final String startingSequenceNumber, final Instant timestamp)
- throws TransientKinesisException {
- final Date date = timestamp != null ? timestamp.toDate() : null;
- return wrapExceptions(new Callable<String>() {
-
- @Override
- public String call() throws Exception {
- return kinesis.getShardIterator(new GetShardIteratorRequest()
- .withStreamName(streamName)
- .withShardId(shardId)
- .withShardIteratorType(shardIteratorType)
- .withStartingSequenceNumber(startingSequenceNumber)
- .withTimestamp(date)
- ).getShardIterator();
- }
- });
- }
-
- public List<Shard> listShards(final String streamName) throws TransientKinesisException {
- return wrapExceptions(new Callable<List<Shard>>() {
-
- @Override
- public List<Shard> call() throws Exception {
- List<Shard> shards = Lists.newArrayList();
- String lastShardId = null;
-
- StreamDescription description;
- do {
- description = kinesis.describeStream(streamName, lastShardId)
- .getStreamDescription();
+ public SimplifiedKinesisClient(AmazonKinesis kinesis) {
+ this.kinesis = kinesis;
+ }
- shards.addAll(description.getShards());
- lastShardId = shards.get(shards.size() - 1).getShardId();
- } while (description.getHasMoreShards());
+ public static SimplifiedKinesisClient from(KinesisClientProvider provider) {
+ return new SimplifiedKinesisClient(provider.get());
+ }
- return shards;
- }
- });
- }
+ public String getShardIterator(final String streamName, final String shardId,
+ final ShardIteratorType shardIteratorType,
+ final String startingSequenceNumber, final Instant timestamp)
+ throws TransientKinesisException {
+ final Date date = timestamp != null ? timestamp.toDate() : null;
+ return wrapExceptions(new Callable<String>() {
+ @Override
+ public String call() throws Exception {
+ return kinesis.getShardIterator(new GetShardIteratorRequest()
+ .withStreamName(streamName)
+ .withShardId(shardId)
+ .withShardIteratorType(shardIteratorType)
+ .withStartingSequenceNumber(startingSequenceNumber)
+ .withTimestamp(date)
+ ).getShardIterator();
+ }
+ });
+ }
- /**
- * Gets records from Kinesis and deaggregates them if needed.
- *
- * @return list of deaggregated records
- * @throws TransientKinesisException - in case of recoverable situation
- */
- public GetKinesisRecordsResult getRecords(String shardIterator, String streamName,
- String shardId) throws TransientKinesisException {
- return getRecords(shardIterator, streamName, shardId, null);
- }
+ public List<Shard> listShards(final String streamName) throws TransientKinesisException {
+ return wrapExceptions(new Callable<List<Shard>>() {
+ @Override
+ public List<Shard> call() throws Exception {
+ List<Shard> shards = Lists.newArrayList();
+ String lastShardId = null;
+
+ StreamDescription description;
+ do {
+ description = kinesis.describeStream(streamName, lastShardId)
+ .getStreamDescription();
+
+ shards.addAll(description.getShards());
+ lastShardId = shards.get(shards.size() - 1).getShardId();
+ } while (description.getHasMoreShards());
+
+ return shards;
+ }
+ });
+ }
- /**
- * Gets records from Kinesis and deaggregates them if needed.
- *
- * @return list of deaggregated records
- * @throws TransientKinesisException - in case of recoverable situation
- */
- public GetKinesisRecordsResult getRecords(final String shardIterator, final String streamName,
- final String shardId, final Integer limit)
- throws
- TransientKinesisException {
- return wrapExceptions(new Callable<GetKinesisRecordsResult>() {
+ /**
+ * Gets records from Kinesis and deaggregates them if needed.
+ *
+ * @return list of deaggregated records
+ * @throws TransientKinesisException - in case of recoverable situation
+ */
+ public GetKinesisRecordsResult getRecords(String shardIterator, String streamName,
+ String shardId) throws TransientKinesisException {
+ return getRecords(shardIterator, streamName, shardId, null);
+ }
- @Override
- public GetKinesisRecordsResult call() throws Exception {
- GetRecordsResult response = kinesis.getRecords(new GetRecordsRequest()
- .withShardIterator(shardIterator)
- .withLimit(limit));
- return new GetKinesisRecordsResult(
- UserRecord.deaggregate(response.getRecords()),
- response.getNextShardIterator(),
- streamName, shardId);
- }
- });
- }
+ /**
+ * Gets records from Kinesis and deaggregates them if needed.
+ *
+ * @return list of deaggregated records
+ * @throws TransientKinesisException - in case of recoverable situation
+ */
+ public GetKinesisRecordsResult getRecords(final String shardIterator, final String streamName,
+ final String shardId, final Integer limit)
+ throws
+ TransientKinesisException {
+ return wrapExceptions(new Callable<GetKinesisRecordsResult>() {
+ @Override
+ public GetKinesisRecordsResult call() throws Exception {
+ GetRecordsResult response = kinesis.getRecords(new GetRecordsRequest()
+ .withShardIterator(shardIterator)
+ .withLimit(limit));
+ return new GetKinesisRecordsResult(
+ UserRecord.deaggregate(response.getRecords()),
+ response.getNextShardIterator(),
+ streamName, shardId);
+ }
+ });
+ }
- /**
- * Wraps Amazon specific exceptions into more friendly format.
- *
- * @throws TransientKinesisException - in case of recoverable situation, i.e.
- * the request rate is too high, Kinesis remote service
- * failed, network issue, etc.
- * @throws ExpiredIteratorException - if iterator needs to be refreshed
- * @throws RuntimeException - in all other cases
- */
- private <T> T wrapExceptions(Callable<T> callable) throws TransientKinesisException {
- try {
- return callable.call();
- } catch (ExpiredIteratorException e) {
- throw e;
- } catch (LimitExceededException | ProvisionedThroughputExceededException e) {
- throw new TransientKinesisException(
- "Too many requests to Kinesis. Wait some time and retry.", e);
- } catch (AmazonServiceException e) {
- if (e.getErrorType() == AmazonServiceException.ErrorType.Service) {
- throw new TransientKinesisException(
- "Kinesis backend failed. Wait some time and retry.", e);
- }
- throw new RuntimeException("Kinesis client side failure", e);
- } catch (Exception e) {
- throw new RuntimeException("Unknown kinesis failure, when trying to reach kinesis", e);
+ /**
+ * Wraps Amazon specific exceptions into more friendly format.
+ *
+ * @throws TransientKinesisException - in case of recoverable situation, i.e.
+ * the request rate is too high, Kinesis remote service
+ * failed, network issue, etc.
+ * @throws ExpiredIteratorException - if iterator needs to be refreshed
+ * @throws RuntimeException - in all other cases
+ */
+ private <T> T wrapExceptions(Callable<T> callable) throws TransientKinesisException {
+ try {
+ return callable.call();
+ } catch (ExpiredIteratorException e) {
+ throw e;
+ } catch (LimitExceededException | ProvisionedThroughputExceededException e) {
+ throw new TransientKinesisException(
+ "Too many requests to Kinesis. Wait some time and retry.", e);
+ } catch (AmazonServiceException e) {
+ if (e.getErrorType() == AmazonServiceException.ErrorType.Service) {
+ throw new TransientKinesisException(
+ "Kinesis backend failed. Wait some time and retry.", e);
+ }
+ throw new RuntimeException("Kinesis client side failure", e);
+ } catch (Exception e) {
+ throw new RuntimeException("Unknown kinesis failure, when trying to reach kinesis", e);
+ }
}
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/StartingPoint.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/StartingPoint.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/StartingPoint.java
index f9298fa..d8842c4 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/StartingPoint.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/StartingPoint.java
@@ -17,14 +17,13 @@
*/
package org.apache.beam.sdk.io.kinesis;
+
import static com.google.common.base.Preconditions.checkNotNull;
import com.amazonaws.services.kinesis.clientlibrary.lib.worker.InitialPositionInStream;
import com.amazonaws.services.kinesis.model.ShardIteratorType;
-
import java.io.Serializable;
import java.util.Objects;
-
import org.joda.time.Instant;
/**
@@ -33,55 +32,54 @@ import org.joda.time.Instant;
* in which case the reader will start reading at the specified point in time.
*/
class StartingPoint implements Serializable {
+ private final InitialPositionInStream position;
+ private final Instant timestamp;
- private final InitialPositionInStream position;
- private final Instant timestamp;
-
- public StartingPoint(InitialPositionInStream position) {
- this.position = checkNotNull(position, "position");
- this.timestamp = null;
- }
-
- public StartingPoint(Instant timestamp) {
- this.timestamp = checkNotNull(timestamp, "timestamp");
- this.position = null;
- }
+ public StartingPoint(InitialPositionInStream position) {
+ this.position = checkNotNull(position, "position");
+ this.timestamp = null;
+ }
- public InitialPositionInStream getPosition() {
- return position;
- }
+ public StartingPoint(Instant timestamp) {
+ this.timestamp = checkNotNull(timestamp, "timestamp");
+ this.position = null;
+ }
- public String getPositionName() {
- return position != null ? position.name() : ShardIteratorType.AT_TIMESTAMP.name();
- }
+ public InitialPositionInStream getPosition() {
+ return position;
+ }
- public Instant getTimestamp() {
- return timestamp != null ? timestamp : null;
- }
+ public String getPositionName() {
+ return position != null ? position.name() : ShardIteratorType.AT_TIMESTAMP.name();
+ }
- @Override
- public boolean equals(Object o) {
- if (this == o) {
- return true;
+ public Instant getTimestamp() {
+ return timestamp != null ? timestamp : null;
}
- if (o == null || getClass() != o.getClass()) {
- return false;
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ StartingPoint that = (StartingPoint) o;
+ return position == that.position && Objects.equals(timestamp, that.timestamp);
}
- StartingPoint that = (StartingPoint) o;
- return position == that.position && Objects.equals(timestamp, that.timestamp);
- }
- @Override
- public int hashCode() {
- return Objects.hash(position, timestamp);
- }
+ @Override
+ public int hashCode() {
+ return Objects.hash(position, timestamp);
+ }
- @Override
- public String toString() {
- if (timestamp == null) {
- return position.toString();
- } else {
- return "Starting at timestamp " + timestamp;
+ @Override
+ public String toString() {
+ if (timestamp == null) {
+ return position.toString();
+ } else {
+ return "Starting at timestamp " + timestamp;
+ }
}
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/StaticCheckpointGenerator.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/StaticCheckpointGenerator.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/StaticCheckpointGenerator.java
index 1ec865d..22dc973 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/StaticCheckpointGenerator.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/StaticCheckpointGenerator.java
@@ -23,21 +23,20 @@ import static com.google.common.base.Preconditions.checkNotNull;
* Always returns the same instance of checkpoint.
*/
class StaticCheckpointGenerator implements CheckpointGenerator {
+ private final KinesisReaderCheckpoint checkpoint;
- private final KinesisReaderCheckpoint checkpoint;
+ public StaticCheckpointGenerator(KinesisReaderCheckpoint checkpoint) {
+ checkNotNull(checkpoint, "checkpoint");
+ this.checkpoint = checkpoint;
+ }
- public StaticCheckpointGenerator(KinesisReaderCheckpoint checkpoint) {
- checkNotNull(checkpoint, "checkpoint");
- this.checkpoint = checkpoint;
- }
+ @Override
+ public KinesisReaderCheckpoint generate(SimplifiedKinesisClient client) {
+ return checkpoint;
+ }
- @Override
- public KinesisReaderCheckpoint generate(SimplifiedKinesisClient client) {
- return checkpoint;
- }
-
- @Override
- public String toString() {
- return checkpoint.toString();
- }
+ @Override
+ public String toString() {
+ return checkpoint.toString();
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/TransientKinesisException.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/TransientKinesisException.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/TransientKinesisException.java
index 68ca0d7..57ad8a8 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/TransientKinesisException.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/TransientKinesisException.java
@@ -23,8 +23,7 @@ import com.amazonaws.AmazonServiceException;
* A transient exception thrown by Kinesis.
*/
class TransientKinesisException extends Exception {
-
- public TransientKinesisException(String s, AmazonServiceException e) {
- super(s, e);
- }
+ public TransientKinesisException(String s, AmazonServiceException e) {
+ super(s, e);
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/AmazonKinesisMock.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/AmazonKinesisMock.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/AmazonKinesisMock.java
index 994d6e3..046c9d9 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/AmazonKinesisMock.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/AmazonKinesisMock.java
@@ -66,12 +66,10 @@ import com.amazonaws.services.kinesis.model.SplitShardRequest;
import com.amazonaws.services.kinesis.model.SplitShardResult;
import com.amazonaws.services.kinesis.model.StreamDescription;
import com.google.common.base.Function;
-
import java.io.Serializable;
import java.nio.ByteBuffer;
import java.util.List;
import javax.annotation.Nullable;
-
import org.apache.commons.lang.builder.EqualsBuilder;
import org.joda.time.Instant;
@@ -80,301 +78,298 @@ import org.joda.time.Instant;
*/
class AmazonKinesisMock implements AmazonKinesis {
- static class TestData implements Serializable {
+ static class TestData implements Serializable {
+ private final String data;
+ private final Instant arrivalTimestamp;
+ private final String sequenceNumber;
+
+ public TestData(KinesisRecord record) {
+ this(new String(record.getData().array()),
+ record.getApproximateArrivalTimestamp(),
+ record.getSequenceNumber());
+ }
+
+ public TestData(String data, Instant arrivalTimestamp, String sequenceNumber) {
+ this.data = data;
+ this.arrivalTimestamp = arrivalTimestamp;
+ this.sequenceNumber = sequenceNumber;
+ }
+
+ public Record convertToRecord() {
+ return new Record().
+ withApproximateArrivalTimestamp(arrivalTimestamp.toDate()).
+ withData(ByteBuffer.wrap(data.getBytes())).
+ withSequenceNumber(sequenceNumber).
+ withPartitionKey("");
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ return EqualsBuilder.reflectionEquals(this, obj);
+ }
+
+ @Override
+ public int hashCode() {
+ return reflectionHashCode(this);
+ }
+ }
+
+ static class Provider implements KinesisClientProvider {
+
+ private final List<List<TestData>> shardedData;
+ private final int numberOfRecordsPerGet;
+
+ public Provider(List<List<TestData>> shardedData, int numberOfRecordsPerGet) {
+ this.shardedData = shardedData;
+ this.numberOfRecordsPerGet = numberOfRecordsPerGet;
+ }
+
+ @Override
+ public AmazonKinesis get() {
+ return new AmazonKinesisMock(transform(shardedData,
+ new Function<List<TestData>, List<Record>>() {
+ @Override
+ public List<Record> apply(@Nullable List<TestData> testDatas) {
+ return transform(testDatas, new Function<TestData, Record>() {
+ @Override
+ public Record apply(@Nullable TestData testData) {
+ return testData.convertToRecord();
+ }
+ });
+ }
+ }), numberOfRecordsPerGet);
+ }
+ }
- private final String data;
- private final Instant arrivalTimestamp;
- private final String sequenceNumber;
+ private final List<List<Record>> shardedData;
+ private final int numberOfRecordsPerGet;
- public TestData(KinesisRecord record) {
- this(new String(record.getData().array()),
- record.getApproximateArrivalTimestamp(),
- record.getSequenceNumber());
+ public AmazonKinesisMock(List<List<Record>> shardedData, int numberOfRecordsPerGet) {
+ this.shardedData = shardedData;
+ this.numberOfRecordsPerGet = numberOfRecordsPerGet;
}
- public TestData(String data, Instant arrivalTimestamp, String sequenceNumber) {
- this.data = data;
- this.arrivalTimestamp = arrivalTimestamp;
- this.sequenceNumber = sequenceNumber;
+ @Override
+ public GetRecordsResult getRecords(GetRecordsRequest getRecordsRequest) {
+ String[] shardIteratorParts = getRecordsRequest.getShardIterator().split(":");
+ int shardId = parseInt(shardIteratorParts[0]);
+ int startingRecord = parseInt(shardIteratorParts[1]);
+ List<Record> shardData = shardedData.get(shardId);
+
+ int toIndex = min(startingRecord + numberOfRecordsPerGet, shardData.size());
+ int fromIndex = min(startingRecord, toIndex);
+ return new GetRecordsResult().
+ withRecords(shardData.subList(fromIndex, toIndex)).
+ withNextShardIterator(String.format("%s:%s", shardId, toIndex));
}
- public Record convertToRecord() {
- return new Record().
- withApproximateArrivalTimestamp(arrivalTimestamp.toDate()).
- withData(ByteBuffer.wrap(data.getBytes())).
- withSequenceNumber(sequenceNumber).
- withPartitionKey("");
+ @Override
+ public GetShardIteratorResult getShardIterator(
+ GetShardIteratorRequest getShardIteratorRequest) {
+ ShardIteratorType shardIteratorType = ShardIteratorType.fromValue(
+ getShardIteratorRequest.getShardIteratorType());
+
+ String shardIterator;
+ if (shardIteratorType == ShardIteratorType.TRIM_HORIZON) {
+ shardIterator = String.format("%s:%s", getShardIteratorRequest.getShardId(), 0);
+ } else {
+ throw new RuntimeException("Not implemented");
+ }
+
+ return new GetShardIteratorResult().withShardIterator(shardIterator);
}
@Override
- public boolean equals(Object obj) {
- return EqualsBuilder.reflectionEquals(this, obj);
+ public DescribeStreamResult describeStream(String streamName, String exclusiveStartShardId) {
+ int nextShardId = 0;
+ if (exclusiveStartShardId != null) {
+ nextShardId = parseInt(exclusiveStartShardId) + 1;
+ }
+ boolean hasMoreShards = nextShardId + 1 < shardedData.size();
+
+ List<Shard> shards = newArrayList();
+ if (nextShardId < shardedData.size()) {
+ shards.add(new Shard().withShardId(Integer.toString(nextShardId)));
+ }
+
+ return new DescribeStreamResult().withStreamDescription(
+ new StreamDescription().withHasMoreShards(hasMoreShards).withShards(shards)
+ );
}
@Override
- public int hashCode() {
- return reflectionHashCode(this);
+ public void setEndpoint(String endpoint) {
+
}
- }
- static class Provider implements KinesisClientProvider {
+ @Override
+ public void setRegion(Region region) {
- private final List<List<TestData>> shardedData;
- private final int numberOfRecordsPerGet;
+ }
+
+ @Override
+ public AddTagsToStreamResult addTagsToStream(AddTagsToStreamRequest addTagsToStreamRequest) {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public CreateStreamResult createStream(CreateStreamRequest createStreamRequest) {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public CreateStreamResult createStream(String streamName, Integer shardCount) {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public DecreaseStreamRetentionPeriodResult decreaseStreamRetentionPeriod(
+ DecreaseStreamRetentionPeriodRequest decreaseStreamRetentionPeriodRequest) {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public DeleteStreamResult deleteStream(DeleteStreamRequest deleteStreamRequest) {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public DeleteStreamResult deleteStream(String streamName) {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public DescribeStreamResult describeStream(DescribeStreamRequest describeStreamRequest) {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public DescribeStreamResult describeStream(String streamName) {
+
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public DescribeStreamResult describeStream(String streamName,
+ Integer limit, String exclusiveStartShardId) {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public DisableEnhancedMonitoringResult disableEnhancedMonitoring(
+ DisableEnhancedMonitoringRequest disableEnhancedMonitoringRequest) {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public EnableEnhancedMonitoringResult enableEnhancedMonitoring(
+ EnableEnhancedMonitoringRequest enableEnhancedMonitoringRequest) {
+ throw new RuntimeException("Not implemented");
+ }
- public Provider(List<List<TestData>> shardedData, int numberOfRecordsPerGet) {
- this.shardedData = shardedData;
- this.numberOfRecordsPerGet = numberOfRecordsPerGet;
+ @Override
+ public GetShardIteratorResult getShardIterator(String streamName,
+ String shardId,
+ String shardIteratorType) {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public GetShardIteratorResult getShardIterator(String streamName,
+ String shardId,
+ String shardIteratorType,
+ String startingSequenceNumber) {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public IncreaseStreamRetentionPeriodResult increaseStreamRetentionPeriod(
+ IncreaseStreamRetentionPeriodRequest increaseStreamRetentionPeriodRequest) {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public ListStreamsResult listStreams(ListStreamsRequest listStreamsRequest) {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public ListStreamsResult listStreams() {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public ListStreamsResult listStreams(String exclusiveStartStreamName) {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public ListStreamsResult listStreams(Integer limit, String exclusiveStartStreamName) {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public ListTagsForStreamResult listTagsForStream(
+ ListTagsForStreamRequest listTagsForStreamRequest) {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public MergeShardsResult mergeShards(MergeShardsRequest mergeShardsRequest) {
+ throw new RuntimeException("Not implemented");
}
@Override
- public AmazonKinesis get() {
- return new AmazonKinesisMock(transform(shardedData,
- new Function<List<TestData>, List<Record>>() {
+ public MergeShardsResult mergeShards(String streamName,
+ String shardToMerge, String adjacentShardToMerge) {
+ throw new RuntimeException("Not implemented");
+ }
- @Override
- public List<Record> apply(@Nullable List<TestData> testDatas) {
- return transform(testDatas, new Function<TestData, Record>() {
+ @Override
+ public PutRecordResult putRecord(PutRecordRequest putRecordRequest) {
+ throw new RuntimeException("Not implemented");
+ }
- @Override
- public Record apply(@Nullable TestData testData) {
- return testData.convertToRecord();
- }
- });
- }
- }), numberOfRecordsPerGet);
+ @Override
+ public PutRecordResult putRecord(String streamName, ByteBuffer data, String partitionKey) {
+ throw new RuntimeException("Not implemented");
}
- }
- private final List<List<Record>> shardedData;
- private final int numberOfRecordsPerGet;
+ @Override
+ public PutRecordResult putRecord(String streamName, ByteBuffer data,
+ String partitionKey, String sequenceNumberForOrdering) {
+ throw new RuntimeException("Not implemented");
+ }
- public AmazonKinesisMock(List<List<Record>> shardedData, int numberOfRecordsPerGet) {
- this.shardedData = shardedData;
- this.numberOfRecordsPerGet = numberOfRecordsPerGet;
- }
+ @Override
+ public PutRecordsResult putRecords(PutRecordsRequest putRecordsRequest) {
+ throw new RuntimeException("Not implemented");
+ }
- @Override
- public GetRecordsResult getRecords(GetRecordsRequest getRecordsRequest) {
- String[] shardIteratorParts = getRecordsRequest.getShardIterator().split(":");
- int shardId = parseInt(shardIteratorParts[0]);
- int startingRecord = parseInt(shardIteratorParts[1]);
- List<Record> shardData = shardedData.get(shardId);
+ @Override
+ public RemoveTagsFromStreamResult removeTagsFromStream(
+ RemoveTagsFromStreamRequest removeTagsFromStreamRequest) {
+ throw new RuntimeException("Not implemented");
+ }
- int toIndex = min(startingRecord + numberOfRecordsPerGet, shardData.size());
- int fromIndex = min(startingRecord, toIndex);
- return new GetRecordsResult().
- withRecords(shardData.subList(fromIndex, toIndex)).
- withNextShardIterator(String.format("%s:%s", shardId, toIndex));
- }
+ @Override
+ public SplitShardResult splitShard(SplitShardRequest splitShardRequest) {
+ throw new RuntimeException("Not implemented");
+ }
+
+ @Override
+ public SplitShardResult splitShard(String streamName,
+ String shardToSplit, String newStartingHashKey) {
+ throw new RuntimeException("Not implemented");
+ }
- @Override
- public GetShardIteratorResult getShardIterator(
- GetShardIteratorRequest getShardIteratorRequest) {
- ShardIteratorType shardIteratorType = ShardIteratorType.fromValue(
- getShardIteratorRequest.getShardIteratorType());
-
- String shardIterator;
- if (shardIteratorType == ShardIteratorType.TRIM_HORIZON) {
- shardIterator = String.format("%s:%s", getShardIteratorRequest.getShardId(), 0);
- } else {
- throw new RuntimeException("Not implemented");
- }
-
- return new GetShardIteratorResult().withShardIterator(shardIterator);
- }
-
- @Override
- public DescribeStreamResult describeStream(String streamName, String exclusiveStartShardId) {
- int nextShardId = 0;
- if (exclusiveStartShardId != null) {
- nextShardId = parseInt(exclusiveStartShardId) + 1;
- }
- boolean hasMoreShards = nextShardId + 1 < shardedData.size();
-
- List<Shard> shards = newArrayList();
- if (nextShardId < shardedData.size()) {
- shards.add(new Shard().withShardId(Integer.toString(nextShardId)));
- }
-
- return new DescribeStreamResult().withStreamDescription(
- new StreamDescription().withHasMoreShards(hasMoreShards).withShards(shards)
- );
- }
-
- @Override
- public void setEndpoint(String endpoint) {
-
- }
-
- @Override
- public void setRegion(Region region) {
-
- }
-
- @Override
- public AddTagsToStreamResult addTagsToStream(AddTagsToStreamRequest addTagsToStreamRequest) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public CreateStreamResult createStream(CreateStreamRequest createStreamRequest) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public CreateStreamResult createStream(String streamName, Integer shardCount) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public DecreaseStreamRetentionPeriodResult decreaseStreamRetentionPeriod(
- DecreaseStreamRetentionPeriodRequest decreaseStreamRetentionPeriodRequest) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public DeleteStreamResult deleteStream(DeleteStreamRequest deleteStreamRequest) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public DeleteStreamResult deleteStream(String streamName) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public DescribeStreamResult describeStream(DescribeStreamRequest describeStreamRequest) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public DescribeStreamResult describeStream(String streamName) {
-
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public DescribeStreamResult describeStream(String streamName,
- Integer limit, String exclusiveStartShardId) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public DisableEnhancedMonitoringResult disableEnhancedMonitoring(
- DisableEnhancedMonitoringRequest disableEnhancedMonitoringRequest) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public EnableEnhancedMonitoringResult enableEnhancedMonitoring(
- EnableEnhancedMonitoringRequest enableEnhancedMonitoringRequest) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public GetShardIteratorResult getShardIterator(String streamName,
- String shardId,
- String shardIteratorType) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public GetShardIteratorResult getShardIterator(String streamName,
- String shardId,
- String shardIteratorType,
- String startingSequenceNumber) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public IncreaseStreamRetentionPeriodResult increaseStreamRetentionPeriod(
- IncreaseStreamRetentionPeriodRequest increaseStreamRetentionPeriodRequest) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public ListStreamsResult listStreams(ListStreamsRequest listStreamsRequest) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public ListStreamsResult listStreams() {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public ListStreamsResult listStreams(String exclusiveStartStreamName) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public ListStreamsResult listStreams(Integer limit, String exclusiveStartStreamName) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public ListTagsForStreamResult listTagsForStream(
- ListTagsForStreamRequest listTagsForStreamRequest) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public MergeShardsResult mergeShards(MergeShardsRequest mergeShardsRequest) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public MergeShardsResult mergeShards(String streamName,
- String shardToMerge, String adjacentShardToMerge) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public PutRecordResult putRecord(PutRecordRequest putRecordRequest) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public PutRecordResult putRecord(String streamName, ByteBuffer data, String partitionKey) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public PutRecordResult putRecord(String streamName, ByteBuffer data,
- String partitionKey, String sequenceNumberForOrdering) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public PutRecordsResult putRecords(PutRecordsRequest putRecordsRequest) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public RemoveTagsFromStreamResult removeTagsFromStream(
- RemoveTagsFromStreamRequest removeTagsFromStreamRequest) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public SplitShardResult splitShard(SplitShardRequest splitShardRequest) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public SplitShardResult splitShard(String streamName,
- String shardToSplit, String newStartingHashKey) {
- throw new RuntimeException("Not implemented");
- }
-
- @Override
- public void shutdown() {
-
- }
-
- @Override
- public ResponseMetadata getCachedResponseMetadata(AmazonWebServiceRequest request) {
- throw new RuntimeException("Not implemented");
- }
+ @Override
+ public void shutdown() {
+
+ }
+
+ @Override
+ public ResponseMetadata getCachedResponseMetadata(AmazonWebServiceRequest request) {
+ throw new RuntimeException("Not implemented");
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/CustomOptionalTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/CustomOptionalTest.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/CustomOptionalTest.java
index 0b16bb7..00acffe 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/CustomOptionalTest.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/CustomOptionalTest.java
@@ -18,27 +18,24 @@
package org.apache.beam.sdk.io.kinesis;
import com.google.common.testing.EqualsTester;
-
import java.util.NoSuchElementException;
-
import org.junit.Test;
/**
* Tests {@link CustomOptional}.
*/
public class CustomOptionalTest {
+ @Test(expected = NoSuchElementException.class)
+ public void absentThrowsNoSuchElementExceptionOnGet() {
+ CustomOptional.absent().get();
+ }
- @Test(expected = NoSuchElementException.class)
- public void absentThrowsNoSuchElementExceptionOnGet() {
- CustomOptional.absent().get();
- }
-
- @Test
- public void testEqualsAndHashCode() {
- new EqualsTester()
- .addEqualityGroup(CustomOptional.absent(), CustomOptional.absent())
- .addEqualityGroup(CustomOptional.of(3), CustomOptional.of(3))
- .addEqualityGroup(CustomOptional.of(11))
- .addEqualityGroup(CustomOptional.of("3")).testEquals();
- }
+ @Test
+ public void testEqualsAndHashCode() {
+ new EqualsTester()
+ .addEqualityGroup(CustomOptional.absent(), CustomOptional.absent())
+ .addEqualityGroup(CustomOptional.of(3), CustomOptional.of(3))
+ .addEqualityGroup(CustomOptional.of(11))
+ .addEqualityGroup(CustomOptional.of("3")).testEquals();
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/DynamicCheckpointGeneratorTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/DynamicCheckpointGeneratorTest.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/DynamicCheckpointGeneratorTest.java
index 1bb9717..c92ac9a 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/DynamicCheckpointGeneratorTest.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/DynamicCheckpointGeneratorTest.java
@@ -28,29 +28,30 @@ import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.runners.MockitoJUnitRunner;
+
/***
*/
@RunWith(MockitoJUnitRunner.class)
public class DynamicCheckpointGeneratorTest {
- @Mock
- private SimplifiedKinesisClient kinesisClient;
- @Mock
- private Shard shard1, shard2, shard3;
+ @Mock
+ private SimplifiedKinesisClient kinesisClient;
+ @Mock
+ private Shard shard1, shard2, shard3;
- @Test
- public void shouldMapAllShardsToCheckpoints() throws Exception {
- given(shard1.getShardId()).willReturn("shard-01");
- given(shard2.getShardId()).willReturn("shard-02");
- given(shard3.getShardId()).willReturn("shard-03");
- given(kinesisClient.listShards("stream")).willReturn(asList(shard1, shard2, shard3));
+ @Test
+ public void shouldMapAllShardsToCheckpoints() throws Exception {
+ given(shard1.getShardId()).willReturn("shard-01");
+ given(shard2.getShardId()).willReturn("shard-02");
+ given(shard3.getShardId()).willReturn("shard-03");
+ given(kinesisClient.listShards("stream")).willReturn(asList(shard1, shard2, shard3));
- StartingPoint startingPoint = new StartingPoint(InitialPositionInStream.LATEST);
- DynamicCheckpointGenerator underTest = new DynamicCheckpointGenerator("stream",
- startingPoint);
+ StartingPoint startingPoint = new StartingPoint(InitialPositionInStream.LATEST);
+ DynamicCheckpointGenerator underTest = new DynamicCheckpointGenerator("stream",
+ startingPoint);
- KinesisReaderCheckpoint checkpoint = underTest.generate(kinesisClient);
+ KinesisReaderCheckpoint checkpoint = underTest.generate(kinesisClient);
- assertThat(checkpoint).hasSize(3);
- }
+ assertThat(checkpoint).hasSize(3);
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisMockReadTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisMockReadTest.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisMockReadTest.java
index 44ad67d..567e25f 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisMockReadTest.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisMockReadTest.java
@@ -21,9 +21,7 @@ import static com.google.common.collect.Lists.newArrayList;
import com.amazonaws.services.kinesis.clientlibrary.lib.worker.InitialPositionInStream;
import com.google.common.collect.Iterables;
-
import java.util.List;
-
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.DoFn;
@@ -38,60 +36,59 @@ import org.junit.Test;
*/
public class KinesisMockReadTest {
- @Rule
- public final transient TestPipeline p = TestPipeline.create();
-
- @Test
- public void readsDataFromMockKinesis() {
- int noOfShards = 3;
- int noOfEventsPerShard = 100;
- List<List<AmazonKinesisMock.TestData>> testData =
- provideTestData(noOfShards, noOfEventsPerShard);
-
- PCollection<AmazonKinesisMock.TestData> result = p
- .apply(
- KinesisIO.read()
- .from("stream", InitialPositionInStream.TRIM_HORIZON)
- .withClientProvider(new AmazonKinesisMock.Provider(testData, 10))
- .withMaxNumRecords(noOfShards * noOfEventsPerShard))
- .apply(ParDo.of(new KinesisRecordToTestData()));
- PAssert.that(result).containsInAnyOrder(Iterables.concat(testData));
- p.run();
- }
-
- private static class KinesisRecordToTestData extends
- DoFn<KinesisRecord, AmazonKinesisMock.TestData> {
+ @Rule
+ public final transient TestPipeline p = TestPipeline.create();
+
+ @Test
+ public void readsDataFromMockKinesis() {
+ int noOfShards = 3;
+ int noOfEventsPerShard = 100;
+ List<List<AmazonKinesisMock.TestData>> testData =
+ provideTestData(noOfShards, noOfEventsPerShard);
+
+ PCollection<AmazonKinesisMock.TestData> result = p
+ .apply(
+ KinesisIO.read()
+ .from("stream", InitialPositionInStream.TRIM_HORIZON)
+ .withClientProvider(new AmazonKinesisMock.Provider(testData, 10))
+ .withMaxNumRecords(noOfShards * noOfEventsPerShard))
+ .apply(ParDo.of(new KinesisRecordToTestData()));
+ PAssert.that(result).containsInAnyOrder(Iterables.concat(testData));
+ p.run();
+ }
- @ProcessElement
- public void processElement(ProcessContext c) throws Exception {
- c.output(new AmazonKinesisMock.TestData(c.element()));
+ private static class KinesisRecordToTestData extends
+ DoFn<KinesisRecord, AmazonKinesisMock.TestData> {
+ @ProcessElement
+ public void processElement(ProcessContext c) throws Exception {
+ c.output(new AmazonKinesisMock.TestData(c.element()));
+ }
}
- }
- private List<List<AmazonKinesisMock.TestData>> provideTestData(
- int noOfShards,
- int noOfEventsPerShard) {
+ private List<List<AmazonKinesisMock.TestData>> provideTestData(
+ int noOfShards,
+ int noOfEventsPerShard) {
- int seqNumber = 0;
+ int seqNumber = 0;
- List<List<AmazonKinesisMock.TestData>> shardedData = newArrayList();
- for (int i = 0; i < noOfShards; ++i) {
- List<AmazonKinesisMock.TestData> shardData = newArrayList();
- shardedData.add(shardData);
+ List<List<AmazonKinesisMock.TestData>> shardedData = newArrayList();
+ for (int i = 0; i < noOfShards; ++i) {
+ List<AmazonKinesisMock.TestData> shardData = newArrayList();
+ shardedData.add(shardData);
- DateTime arrival = DateTime.now();
- for (int j = 0; j < noOfEventsPerShard; ++j) {
- arrival = arrival.plusSeconds(1);
+ DateTime arrival = DateTime.now();
+ for (int j = 0; j < noOfEventsPerShard; ++j) {
+ arrival = arrival.plusSeconds(1);
- seqNumber++;
- shardData.add(new AmazonKinesisMock.TestData(
- Integer.toString(seqNumber),
- arrival.toInstant(),
- Integer.toString(seqNumber))
- );
- }
- }
+ seqNumber++;
+ shardData.add(new AmazonKinesisMock.TestData(
+ Integer.toString(seqNumber),
+ arrival.toInstant(),
+ Integer.toString(seqNumber))
+ );
+ }
+ }
- return shardedData;
- }
+ return shardedData;
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderCheckpointTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderCheckpointTest.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderCheckpointTest.java
index 1038a47..8c8da64 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderCheckpointTest.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderCheckpointTest.java
@@ -17,14 +17,13 @@
*/
package org.apache.beam.sdk.io.kinesis;
+
import static java.util.Arrays.asList;
import static org.assertj.core.api.Assertions.assertThat;
import com.google.common.collect.Iterables;
-
import java.util.Iterator;
import java.util.List;
-
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -36,34 +35,33 @@ import org.mockito.runners.MockitoJUnitRunner;
*/
@RunWith(MockitoJUnitRunner.class)
public class KinesisReaderCheckpointTest {
+ @Mock
+ private ShardCheckpoint a, b, c;
- @Mock
- private ShardCheckpoint a, b, c;
-
- private KinesisReaderCheckpoint checkpoint;
+ private KinesisReaderCheckpoint checkpoint;
- @Before
- public void setUp() {
- checkpoint = new KinesisReaderCheckpoint(asList(a, b, c));
- }
+ @Before
+ public void setUp() {
+ checkpoint = new KinesisReaderCheckpoint(asList(a, b, c));
+ }
- @Test
- public void splitsCheckpointAccordingly() {
- verifySplitInto(1);
- verifySplitInto(2);
- verifySplitInto(3);
- verifySplitInto(4);
- }
+ @Test
+ public void splitsCheckpointAccordingly() {
+ verifySplitInto(1);
+ verifySplitInto(2);
+ verifySplitInto(3);
+ verifySplitInto(4);
+ }
- @Test(expected = UnsupportedOperationException.class)
- public void isImmutable() {
- Iterator<ShardCheckpoint> iterator = checkpoint.iterator();
- iterator.remove();
- }
+ @Test(expected = UnsupportedOperationException.class)
+ public void isImmutable() {
+ Iterator<ShardCheckpoint> iterator = checkpoint.iterator();
+ iterator.remove();
+ }
- private void verifySplitInto(int size) {
- List<KinesisReaderCheckpoint> split = checkpoint.splitInto(size);
- assertThat(Iterables.concat(split)).containsOnly(a, b, c);
- assertThat(split).hasSize(Math.min(size, 3));
- }
+ private void verifySplitInto(int size) {
+ List<KinesisReaderCheckpoint> split = checkpoint.splitInto(size);
+ assertThat(Iterables.concat(split)).containsOnly(a, b, c);
+ assertThat(split).hasSize(Math.min(size, 3));
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderIT.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderIT.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderIT.java
index 5781033..8eb6546 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderIT.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderIT.java
@@ -23,7 +23,6 @@ import static java.util.concurrent.Executors.newSingleThreadExecutor;
import static org.assertj.core.api.Assertions.assertThat;
import com.amazonaws.regions.Regions;
-
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
@@ -32,7 +31,6 @@ import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
-
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.testing.PAssert;
@@ -52,75 +50,72 @@ import org.junit.Test;
* You need to provide all {@link KinesisTestOptions} in order to run this.
*/
public class KinesisReaderIT {
-
- private static final long PIPELINE_STARTUP_TIME = TimeUnit.SECONDS.toMillis(10);
- private ExecutorService singleThreadExecutor = newSingleThreadExecutor();
-
- @Rule
- public final transient TestPipeline p = TestPipeline.create();
-
- @Ignore
- @Test
- public void readsDataFromRealKinesisStream()
- throws IOException, InterruptedException, ExecutionException {
- KinesisTestOptions options = readKinesisOptions();
- List<String> testData = prepareTestData(1000);
-
- Future<?> future = startTestPipeline(testData, options);
- KinesisUploader.uploadAll(testData, options);
- future.get();
- }
-
- private List<String> prepareTestData(int count) {
- List<String> data = newArrayList();
- for (int i = 0; i < count; ++i) {
- data.add(RandomStringUtils.randomAlphabetic(32));
+ private static final long PIPELINE_STARTUP_TIME = TimeUnit.SECONDS.toMillis(10);
+ private ExecutorService singleThreadExecutor = newSingleThreadExecutor();
+
+ @Rule
+ public final transient TestPipeline p = TestPipeline.create();
+
+ @Ignore
+ @Test
+ public void readsDataFromRealKinesisStream()
+ throws IOException, InterruptedException, ExecutionException {
+ KinesisTestOptions options = readKinesisOptions();
+ List<String> testData = prepareTestData(1000);
+
+ Future<?> future = startTestPipeline(testData, options);
+ KinesisUploader.uploadAll(testData, options);
+ future.get();
}
- return data;
- }
- private Future<?> startTestPipeline(List<String> testData, KinesisTestOptions options)
- throws InterruptedException {
-
- PCollection<String> result = p.
- apply(KinesisIO.read()
- .from(options.getAwsKinesisStream(), Instant.now())
- .withClientProvider(options.getAwsAccessKey(), options.getAwsSecretKey(),
- Regions.fromName(options.getAwsKinesisRegion()))
- .withMaxReadTime(Duration.standardMinutes(3))
- ).
- apply(ParDo.of(new RecordDataToString()));
- PAssert.that(result).containsInAnyOrder(testData);
-
- Future<?> future = singleThreadExecutor.submit(new Callable<Void>() {
-
- @Override
- public Void call() throws Exception {
- PipelineResult result = p.run();
- PipelineResult.State state = result.getState();
- while (state != PipelineResult.State.DONE && state != PipelineResult.State.FAILED) {
- Thread.sleep(1000);
- state = result.getState();
+ private List<String> prepareTestData(int count) {
+ List<String> data = newArrayList();
+ for (int i = 0; i < count; ++i) {
+ data.add(RandomStringUtils.randomAlphabetic(32));
}
- assertThat(state).isEqualTo(PipelineResult.State.DONE);
- return null;
- }
- });
- Thread.sleep(PIPELINE_STARTUP_TIME);
- return future;
- }
+ return data;
+ }
- private KinesisTestOptions readKinesisOptions() {
- PipelineOptionsFactory.register(KinesisTestOptions.class);
- return TestPipeline.testingPipelineOptions().as(KinesisTestOptions.class);
- }
+ private Future<?> startTestPipeline(List<String> testData, KinesisTestOptions options)
+ throws InterruptedException {
+
+ PCollection<String> result = p.
+ apply(KinesisIO.read()
+ .from(options.getAwsKinesisStream(), Instant.now())
+ .withClientProvider(options.getAwsAccessKey(), options.getAwsSecretKey(),
+ Regions.fromName(options.getAwsKinesisRegion()))
+ .withMaxReadTime(Duration.standardMinutes(3))
+ ).
+ apply(ParDo.of(new RecordDataToString()));
+ PAssert.that(result).containsInAnyOrder(testData);
+
+ Future<?> future = singleThreadExecutor.submit(new Callable<Void>() {
+ @Override
+ public Void call() throws Exception {
+ PipelineResult result = p.run();
+ PipelineResult.State state = result.getState();
+ while (state != PipelineResult.State.DONE && state != PipelineResult.State.FAILED) {
+ Thread.sleep(1000);
+ state = result.getState();
+ }
+ assertThat(state).isEqualTo(PipelineResult.State.DONE);
+ return null;
+ }
+ });
+ Thread.sleep(PIPELINE_STARTUP_TIME);
+ return future;
+ }
- private static class RecordDataToString extends DoFn<KinesisRecord, String> {
+ private KinesisTestOptions readKinesisOptions() {
+ PipelineOptionsFactory.register(KinesisTestOptions.class);
+ return TestPipeline.testingPipelineOptions().as(KinesisTestOptions.class);
+ }
- @ProcessElement
- public void processElement(ProcessContext c) throws Exception {
- checkNotNull(c.element(), "Null record given");
- c.output(new String(c.element().getData().array(), StandardCharsets.UTF_8));
+ private static class RecordDataToString extends DoFn<KinesisRecord, String> {
+ @ProcessElement
+ public void processElement(ProcessContext c) throws Exception {
+ checkNotNull(c.element(), "Null record given");
+ c.output(new String(c.element().getData().array(), StandardCharsets.UTF_8));
+ }
}
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderTest.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderTest.java
index a26501a..3111029 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderTest.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisReaderTest.java
@@ -23,7 +23,6 @@ import static org.mockito.Mockito.when;
import java.io.IOException;
import java.util.NoSuchElementException;
-
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -35,88 +34,87 @@ import org.mockito.runners.MockitoJUnitRunner;
*/
@RunWith(MockitoJUnitRunner.class)
public class KinesisReaderTest {
-
- @Mock
- private SimplifiedKinesisClient kinesis;
- @Mock
- private CheckpointGenerator generator;
- @Mock
- private ShardCheckpoint firstCheckpoint, secondCheckpoint;
- @Mock
- private ShardRecordsIterator firstIterator, secondIterator;
- @Mock
- private KinesisRecord a, b, c, d;
-
- private KinesisReader reader;
-
- @Before
- public void setUp() throws IOException, TransientKinesisException {
- when(generator.generate(kinesis)).thenReturn(new KinesisReaderCheckpoint(
- asList(firstCheckpoint, secondCheckpoint)
- ));
- when(firstCheckpoint.getShardRecordsIterator(kinesis)).thenReturn(firstIterator);
- when(secondCheckpoint.getShardRecordsIterator(kinesis)).thenReturn(secondIterator);
- when(firstIterator.next()).thenReturn(CustomOptional.<KinesisRecord>absent());
- when(secondIterator.next()).thenReturn(CustomOptional.<KinesisRecord>absent());
-
- reader = new KinesisReader(kinesis, generator, null);
- }
-
- @Test
- public void startReturnsFalseIfNoDataAtTheBeginning() throws IOException {
- assertThat(reader.start()).isFalse();
- }
-
- @Test(expected = NoSuchElementException.class)
- public void throwsNoSuchElementExceptionIfNoData() throws IOException {
- reader.start();
- reader.getCurrent();
- }
-
- @Test
- public void startReturnsTrueIfSomeDataAvailable() throws IOException,
- TransientKinesisException {
- when(firstIterator.next()).
- thenReturn(CustomOptional.of(a)).
- thenReturn(CustomOptional.<KinesisRecord>absent());
-
- assertThat(reader.start()).isTrue();
- }
-
- @Test
- public void advanceReturnsFalseIfThereIsTransientExceptionInKinesis()
- throws IOException, TransientKinesisException {
- reader.start();
-
- when(firstIterator.next()).thenThrow(TransientKinesisException.class);
-
- assertThat(reader.advance()).isFalse();
- }
-
- @Test
- public void readsThroughAllDataAvailable() throws IOException, TransientKinesisException {
- when(firstIterator.next()).
- thenReturn(CustomOptional.<KinesisRecord>absent()).
- thenReturn(CustomOptional.of(a)).
- thenReturn(CustomOptional.<KinesisRecord>absent()).
- thenReturn(CustomOptional.of(b)).
- thenReturn(CustomOptional.<KinesisRecord>absent());
-
- when(secondIterator.next()).
- thenReturn(CustomOptional.of(c)).
- thenReturn(CustomOptional.<KinesisRecord>absent()).
- thenReturn(CustomOptional.of(d)).
- thenReturn(CustomOptional.<KinesisRecord>absent());
-
- assertThat(reader.start()).isTrue();
- assertThat(reader.getCurrent()).isEqualTo(c);
- assertThat(reader.advance()).isTrue();
- assertThat(reader.getCurrent()).isEqualTo(a);
- assertThat(reader.advance()).isTrue();
- assertThat(reader.getCurrent()).isEqualTo(d);
- assertThat(reader.advance()).isTrue();
- assertThat(reader.getCurrent()).isEqualTo(b);
- assertThat(reader.advance()).isFalse();
- }
+ @Mock
+ private SimplifiedKinesisClient kinesis;
+ @Mock
+ private CheckpointGenerator generator;
+ @Mock
+ private ShardCheckpoint firstCheckpoint, secondCheckpoint;
+ @Mock
+ private ShardRecordsIterator firstIterator, secondIterator;
+ @Mock
+ private KinesisRecord a, b, c, d;
+
+ private KinesisReader reader;
+
+ @Before
+ public void setUp() throws IOException, TransientKinesisException {
+ when(generator.generate(kinesis)).thenReturn(new KinesisReaderCheckpoint(
+ asList(firstCheckpoint, secondCheckpoint)
+ ));
+ when(firstCheckpoint.getShardRecordsIterator(kinesis)).thenReturn(firstIterator);
+ when(secondCheckpoint.getShardRecordsIterator(kinesis)).thenReturn(secondIterator);
+ when(firstIterator.next()).thenReturn(CustomOptional.<KinesisRecord>absent());
+ when(secondIterator.next()).thenReturn(CustomOptional.<KinesisRecord>absent());
+
+ reader = new KinesisReader(kinesis, generator, null);
+ }
+
+ @Test
+ public void startReturnsFalseIfNoDataAtTheBeginning() throws IOException {
+ assertThat(reader.start()).isFalse();
+ }
+
+ @Test(expected = NoSuchElementException.class)
+ public void throwsNoSuchElementExceptionIfNoData() throws IOException {
+ reader.start();
+ reader.getCurrent();
+ }
+
+ @Test
+ public void startReturnsTrueIfSomeDataAvailable() throws IOException,
+ TransientKinesisException {
+ when(firstIterator.next()).
+ thenReturn(CustomOptional.of(a)).
+ thenReturn(CustomOptional.<KinesisRecord>absent());
+
+ assertThat(reader.start()).isTrue();
+ }
+
+ @Test
+ public void advanceReturnsFalseIfThereIsTransientExceptionInKinesis()
+ throws IOException, TransientKinesisException {
+ reader.start();
+
+ when(firstIterator.next()).thenThrow(TransientKinesisException.class);
+
+ assertThat(reader.advance()).isFalse();
+ }
+
+ @Test
+ public void readsThroughAllDataAvailable() throws IOException, TransientKinesisException {
+ when(firstIterator.next()).
+ thenReturn(CustomOptional.<KinesisRecord>absent()).
+ thenReturn(CustomOptional.of(a)).
+ thenReturn(CustomOptional.<KinesisRecord>absent()).
+ thenReturn(CustomOptional.of(b)).
+ thenReturn(CustomOptional.<KinesisRecord>absent());
+
+ when(secondIterator.next()).
+ thenReturn(CustomOptional.of(c)).
+ thenReturn(CustomOptional.<KinesisRecord>absent()).
+ thenReturn(CustomOptional.of(d)).
+ thenReturn(CustomOptional.<KinesisRecord>absent());
+
+ assertThat(reader.start()).isTrue();
+ assertThat(reader.getCurrent()).isEqualTo(c);
+ assertThat(reader.advance()).isTrue();
+ assertThat(reader.getCurrent()).isEqualTo(a);
+ assertThat(reader.advance()).isTrue();
+ assertThat(reader.getCurrent()).isEqualTo(d);
+ assertThat(reader.advance()).isTrue();
+ assertThat(reader.getCurrent()).isEqualTo(b);
+ assertThat(reader.advance()).isFalse();
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisRecordCoderTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisRecordCoderTest.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisRecordCoderTest.java
index c9f01bb..8771c86 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisRecordCoderTest.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisRecordCoderTest.java
@@ -18,7 +18,6 @@
package org.apache.beam.sdk.io.kinesis;
import java.nio.ByteBuffer;
-
import org.apache.beam.sdk.testing.CoderProperties;
import org.joda.time.Instant;
import org.junit.Test;
@@ -27,21 +26,20 @@ import org.junit.Test;
* Tests {@link KinesisRecordCoder}.
*/
public class KinesisRecordCoderTest {
-
- @Test
- public void encodingAndDecodingWorks() throws Exception {
- KinesisRecord record = new KinesisRecord(
- ByteBuffer.wrap("data".getBytes()),
- "sequence",
- 128L,
- "partition",
- Instant.now(),
- Instant.now(),
- "stream",
- "shard"
- );
- CoderProperties.coderDecodeEncodeEqual(
- new KinesisRecordCoder(), record
- );
- }
+ @Test
+ public void encodingAndDecodingWorks() throws Exception {
+ KinesisRecord record = new KinesisRecord(
+ ByteBuffer.wrap("data".getBytes()),
+ "sequence",
+ 128L,
+ "partition",
+ Instant.now(),
+ Instant.now(),
+ "stream",
+ "shard"
+ );
+ CoderProperties.coderDecodeEncodeEqual(
+ new KinesisRecordCoder(), record
+ );
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisTestOptions.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisTestOptions.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisTestOptions.java
index 76bcb27..324de46 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisTestOptions.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisTestOptions.java
@@ -25,28 +25,23 @@ import org.apache.beam.sdk.testing.TestPipelineOptions;
* Options for Kinesis integration tests.
*/
public interface KinesisTestOptions extends TestPipelineOptions {
-
- @Description("AWS region where Kinesis stream resided")
- @Default.String("aws-kinesis-region")
- String getAwsKinesisRegion();
-
- void setAwsKinesisRegion(String value);
-
- @Description("Kinesis stream name")
- @Default.String("aws-kinesis-stream")
- String getAwsKinesisStream();
-
- void setAwsKinesisStream(String value);
-
- @Description("AWS secret key")
- @Default.String("aws-secret-key")
- String getAwsSecretKey();
-
- void setAwsSecretKey(String value);
-
- @Description("AWS access key")
- @Default.String("aws-access-key")
- String getAwsAccessKey();
-
- void setAwsAccessKey(String value);
+ @Description("AWS region where Kinesis stream resided")
+ @Default.String("aws-kinesis-region")
+ String getAwsKinesisRegion();
+ void setAwsKinesisRegion(String value);
+
+ @Description("Kinesis stream name")
+ @Default.String("aws-kinesis-stream")
+ String getAwsKinesisStream();
+ void setAwsKinesisStream(String value);
+
+ @Description("AWS secret key")
+ @Default.String("aws-secret-key")
+ String getAwsSecretKey();
+ void setAwsSecretKey(String value);
+
+ @Description("AWS access key")
+ @Default.String("aws-access-key")
+ String getAwsAccessKey();
+ void setAwsAccessKey(String value);
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisUploader.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisUploader.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisUploader.java
index 7a7cb02..7518ff7 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisUploader.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/KinesisUploader.java
@@ -29,7 +29,6 @@ import com.amazonaws.services.kinesis.model.PutRecordsResult;
import com.amazonaws.services.kinesis.model.PutRecordsResultEntry;
import com.google.common.base.Charsets;
import com.google.common.collect.Lists;
-
import java.nio.ByteBuffer;
import java.util.List;
@@ -38,46 +37,47 @@ import java.util.List;
*/
public class KinesisUploader {
- public static final int MAX_NUMBER_OF_RECORDS_IN_BATCH = 499;
+ public static final int MAX_NUMBER_OF_RECORDS_IN_BATCH = 499;
- public static void uploadAll(List<String> data, KinesisTestOptions options) {
- AmazonKinesisClient client = new AmazonKinesisClient(
- new StaticCredentialsProvider(
- new BasicAWSCredentials(
- options.getAwsAccessKey(), options.getAwsSecretKey()))
- ).withRegion(Regions.fromName(options.getAwsKinesisRegion()));
+ public static void uploadAll(List<String> data, KinesisTestOptions options) {
+ AmazonKinesisClient client = new AmazonKinesisClient(
+ new StaticCredentialsProvider(
+ new BasicAWSCredentials(
+ options.getAwsAccessKey(), options.getAwsSecretKey()))
+ ).withRegion(Regions.fromName(options.getAwsKinesisRegion()));
- List<List<String>> partitions = Lists.partition(data, MAX_NUMBER_OF_RECORDS_IN_BATCH);
+ List<List<String>> partitions = Lists.partition(data, MAX_NUMBER_OF_RECORDS_IN_BATCH);
- for (List<String> partition : partitions) {
- List<PutRecordsRequestEntry> allRecords = newArrayList();
- for (String row : partition) {
- allRecords.add(new PutRecordsRequestEntry().
- withData(ByteBuffer.wrap(row.getBytes(Charsets.UTF_8))).
- withPartitionKey(Integer.toString(row.hashCode()))
- );
- }
+ for (List<String> partition : partitions) {
+ List<PutRecordsRequestEntry> allRecords = newArrayList();
+ for (String row : partition) {
+ allRecords.add(new PutRecordsRequestEntry().
+ withData(ByteBuffer.wrap(row.getBytes(Charsets.UTF_8))).
+ withPartitionKey(Integer.toString(row.hashCode()))
- PutRecordsResult result;
- do {
- result = client.putRecords(
- new PutRecordsRequest().
- withStreamName(options.getAwsKinesisStream()).
- withRecords(allRecords));
- List<PutRecordsRequestEntry> failedRecords = newArrayList();
- int i = 0;
- for (PutRecordsResultEntry row : result.getRecords()) {
- if (row.getErrorCode() != null) {
- failedRecords.add(allRecords.get(i));
- }
- ++i;
- }
- allRecords = failedRecords;
- }
+ );
+ }
- while (result.getFailedRecordCount() > 0);
+ PutRecordsResult result;
+ do {
+ result = client.putRecords(
+ new PutRecordsRequest().
+ withStreamName(options.getAwsKinesisStream()).
+ withRecords(allRecords));
+ List<PutRecordsRequestEntry> failedRecords = newArrayList();
+ int i = 0;
+ for (PutRecordsResultEntry row : result.getRecords()) {
+ if (row.getErrorCode() != null) {
+ failedRecords.add(allRecords.get(i));
+ }
+ ++i;
+ }
+ allRecords = failedRecords;
+ }
+
+ while (result.getFailedRecordCount() > 0);
+ }
}
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/RecordFilterTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/RecordFilterTest.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/RecordFilterTest.java
index cb32562..f979c01 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/RecordFilterTest.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/RecordFilterTest.java
@@ -20,49 +20,47 @@ package org.apache.beam.sdk.io.kinesis;
import static org.mockito.BDDMockito.given;
import com.google.common.collect.Lists;
-
import java.util.Collections;
import java.util.List;
-
import org.assertj.core.api.Assertions;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.runners.MockitoJUnitRunner;
+
/***
*/
@RunWith(MockitoJUnitRunner.class)
public class RecordFilterTest {
+ @Mock
+ private ShardCheckpoint checkpoint;
+ @Mock
+ private KinesisRecord record1, record2, record3, record4, record5;
- @Mock
- private ShardCheckpoint checkpoint;
- @Mock
- private KinesisRecord record1, record2, record3, record4, record5;
-
- @Test
- public void shouldFilterOutRecordsBeforeOrAtCheckpoint() {
- given(checkpoint.isBeforeOrAt(record1)).willReturn(false);
- given(checkpoint.isBeforeOrAt(record2)).willReturn(true);
- given(checkpoint.isBeforeOrAt(record3)).willReturn(true);
- given(checkpoint.isBeforeOrAt(record4)).willReturn(false);
- given(checkpoint.isBeforeOrAt(record5)).willReturn(true);
- List<KinesisRecord> records = Lists.newArrayList(record1, record2,
- record3, record4, record5);
- RecordFilter underTest = new RecordFilter();
+ @Test
+ public void shouldFilterOutRecordsBeforeOrAtCheckpoint() {
+ given(checkpoint.isBeforeOrAt(record1)).willReturn(false);
+ given(checkpoint.isBeforeOrAt(record2)).willReturn(true);
+ given(checkpoint.isBeforeOrAt(record3)).willReturn(true);
+ given(checkpoint.isBeforeOrAt(record4)).willReturn(false);
+ given(checkpoint.isBeforeOrAt(record5)).willReturn(true);
+ List<KinesisRecord> records = Lists.newArrayList(record1, record2,
+ record3, record4, record5);
+ RecordFilter underTest = new RecordFilter();
- List<KinesisRecord> retainedRecords = underTest.apply(records, checkpoint);
+ List<KinesisRecord> retainedRecords = underTest.apply(records, checkpoint);
- Assertions.assertThat(retainedRecords).containsOnly(record2, record3, record5);
- }
+ Assertions.assertThat(retainedRecords).containsOnly(record2, record3, record5);
+ }
- @Test
- public void shouldNotFailOnEmptyList() {
- List<KinesisRecord> records = Collections.emptyList();
- RecordFilter underTest = new RecordFilter();
+ @Test
+ public void shouldNotFailOnEmptyList() {
+ List<KinesisRecord> records = Collections.emptyList();
+ RecordFilter underTest = new RecordFilter();
- List<KinesisRecord> retainedRecords = underTest.apply(records, checkpoint);
+ List<KinesisRecord> retainedRecords = underTest.apply(records, checkpoint);
- Assertions.assertThat(retainedRecords).isEmpty();
- }
+ Assertions.assertThat(retainedRecords).isEmpty();
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/RoundRobinTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/RoundRobinTest.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/RoundRobinTest.java
index e4abce4..f032eea 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/RoundRobinTest.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/RoundRobinTest.java
@@ -22,38 +22,36 @@ import static org.assertj.core.api.Assertions.assertThat;
import java.util.Collections;
import java.util.List;
-
import org.junit.Test;
/**
* Tests {@link RoundRobin}.
*/
public class RoundRobinTest {
+ @Test(expected = IllegalArgumentException.class)
+ public void doesNotAllowCreationWithEmptyCollection() {
+ new RoundRobin<>(Collections.emptyList());
+ }
- @Test(expected = IllegalArgumentException.class)
- public void doesNotAllowCreationWithEmptyCollection() {
- new RoundRobin<>(Collections.emptyList());
- }
-
- @Test
- public void goesThroughElementsInCycle() {
- List<String> input = newArrayList("a", "b", "c");
+ @Test
+ public void goesThroughElementsInCycle() {
+ List<String> input = newArrayList("a", "b", "c");
- RoundRobin<String> roundRobin = new RoundRobin<>(newArrayList(input));
+ RoundRobin<String> roundRobin = new RoundRobin<>(newArrayList(input));
- input.addAll(input); // duplicate the input
- for (String element : input) {
- assertThat(roundRobin.getCurrent()).isEqualTo(element);
- assertThat(roundRobin.getCurrent()).isEqualTo(element);
- roundRobin.moveForward();
+ input.addAll(input); // duplicate the input
+ for (String element : input) {
+ assertThat(roundRobin.getCurrent()).isEqualTo(element);
+ assertThat(roundRobin.getCurrent()).isEqualTo(element);
+ roundRobin.moveForward();
+ }
}
- }
- @Test
- public void usualIteratorGoesThroughElementsOnce() {
- List<String> input = newArrayList("a", "b", "c");
+ @Test
+ public void usualIteratorGoesThroughElementsOnce() {
+ List<String> input = newArrayList("a", "b", "c");
- RoundRobin<String> roundRobin = new RoundRobin<>(input);
- assertThat(roundRobin).hasSize(3).containsOnly(input.toArray(new String[0]));
- }
+ RoundRobin<String> roundRobin = new RoundRobin<>(input);
+ assertThat(roundRobin).hasSize(3).containsOnly(input.toArray(new String[0]));
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/ShardCheckpointTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/ShardCheckpointTest.java b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/ShardCheckpointTest.java
index d4784c4..39ab36f 100644
--- a/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/ShardCheckpointTest.java
+++ b/sdks/java/io/kinesis/src/test/java/org/apache/beam/sdk/io/kinesis/ShardCheckpointTest.java
@@ -32,9 +32,7 @@ import static org.mockito.Mockito.when;
import com.amazonaws.services.kinesis.clientlibrary.types.ExtendedSequenceNumber;
import com.amazonaws.services.kinesis.model.ShardIteratorType;
-
import java.io.IOException;
-
import org.joda.time.DateTime;
import org.joda.time.Instant;
import org.junit.Before;
@@ -48,105 +46,104 @@ import org.mockito.runners.MockitoJUnitRunner;
*/
@RunWith(MockitoJUnitRunner.class)
public class ShardCheckpointTest {
-
- private static final String AT_SEQUENCE_SHARD_IT = "AT_SEQUENCE_SHARD_IT";
- private static final String AFTER_SEQUENCE_SHARD_IT = "AFTER_SEQUENCE_SHARD_IT";
- private static final String STREAM_NAME = "STREAM";
- private static final String SHARD_ID = "SHARD_ID";
- @Mock
- private SimplifiedKinesisClient client;
-
- @Before
- public void setUp() throws IOException, TransientKinesisException {
- when(client.getShardIterator(
- eq(STREAM_NAME), eq(SHARD_ID), eq(AT_SEQUENCE_NUMBER),
- anyString(), isNull(Instant.class))).
- thenReturn(AT_SEQUENCE_SHARD_IT);
- when(client.getShardIterator(
- eq(STREAM_NAME), eq(SHARD_ID), eq(AFTER_SEQUENCE_NUMBER),
- anyString(), isNull(Instant.class))).
- thenReturn(AFTER_SEQUENCE_SHARD_IT);
- }
-
- @Test
- public void testProvidingShardIterator() throws IOException, TransientKinesisException {
- assertThat(checkpoint(AT_SEQUENCE_NUMBER, "100", null).getShardIterator(client))
- .isEqualTo(AT_SEQUENCE_SHARD_IT);
- assertThat(checkpoint(AFTER_SEQUENCE_NUMBER, "100", null).getShardIterator(client))
- .isEqualTo(AFTER_SEQUENCE_SHARD_IT);
- assertThat(checkpoint(AT_SEQUENCE_NUMBER, "100", 10L).getShardIterator(client)).isEqualTo
- (AT_SEQUENCE_SHARD_IT);
- assertThat(checkpoint(AFTER_SEQUENCE_NUMBER, "100", 10L).getShardIterator(client))
- .isEqualTo(AT_SEQUENCE_SHARD_IT);
- }
-
- @Test
- public void testComparisonWithExtendedSequenceNumber() {
- assertThat(new ShardCheckpoint("", "", new StartingPoint(LATEST)).isBeforeOrAt(
- recordWith(new ExtendedSequenceNumber("100", 0L))
- )).isTrue();
-
- assertThat(new ShardCheckpoint("", "", new StartingPoint(TRIM_HORIZON)).isBeforeOrAt(
- recordWith(new ExtendedSequenceNumber("100", 0L))
- )).isTrue();
-
- assertThat(checkpoint(AFTER_SEQUENCE_NUMBER, "10", 1L).isBeforeOrAt(
- recordWith(new ExtendedSequenceNumber("100", 0L))
- )).isTrue();
-
- assertThat(checkpoint(AT_SEQUENCE_NUMBER, "100", 0L).isBeforeOrAt(
- recordWith(new ExtendedSequenceNumber("100", 0L))
- )).isTrue();
-
- assertThat(checkpoint(AFTER_SEQUENCE_NUMBER, "100", 0L).isBeforeOrAt(
- recordWith(new ExtendedSequenceNumber("100", 0L))
- )).isFalse();
-
- assertThat(checkpoint(AT_SEQUENCE_NUMBER, "100", 1L).isBeforeOrAt(
- recordWith(new ExtendedSequenceNumber("100", 0L))
- )).isFalse();
-
- assertThat(checkpoint(AFTER_SEQUENCE_NUMBER, "100", 0L).isBeforeOrAt(
- recordWith(new ExtendedSequenceNumber("99", 1L))
- )).isFalse();
- }
-
- @Test
- public void testComparisonWithTimestamp() {
- DateTime referenceTimestamp = DateTime.now();
-
- assertThat(checkpoint(AT_TIMESTAMP, referenceTimestamp.toInstant())
- .isBeforeOrAt(recordWith(referenceTimestamp.minusMillis(10).toInstant()))
- ).isFalse();
-
- assertThat(checkpoint(AT_TIMESTAMP, referenceTimestamp.toInstant())
- .isBeforeOrAt(recordWith(referenceTimestamp.toInstant()))
- ).isTrue();
-
- assertThat(checkpoint(AT_TIMESTAMP, referenceTimestamp.toInstant())
- .isBeforeOrAt(recordWith(referenceTimestamp.plusMillis(10).toInstant()))
- ).isTrue();
- }
-
- private KinesisRecord recordWith(ExtendedSequenceNumber extendedSequenceNumber) {
- KinesisRecord record = mock(KinesisRecord.class);
- given(record.getExtendedSequenceNumber()).willReturn(extendedSequenceNumber);
- return record;
- }
-
- private ShardCheckpoint checkpoint(ShardIteratorType iteratorType, String sequenceNumber,
- Long subSequenceNumber) {
- return new ShardCheckpoint(STREAM_NAME, SHARD_ID, iteratorType, sequenceNumber,
- subSequenceNumber);
- }
-
- private KinesisRecord recordWith(Instant approximateArrivalTimestamp) {
- KinesisRecord record = mock(KinesisRecord.class);
- given(record.getApproximateArrivalTimestamp()).willReturn(approximateArrivalTimestamp);
- return record;
- }
-
- private ShardCheckpoint checkpoint(ShardIteratorType iteratorType, Instant timestamp) {
- return new ShardCheckpoint(STREAM_NAME, SHARD_ID, iteratorType, timestamp);
- }
+ private static final String AT_SEQUENCE_SHARD_IT = "AT_SEQUENCE_SHARD_IT";
+ private static final String AFTER_SEQUENCE_SHARD_IT = "AFTER_SEQUENCE_SHARD_IT";
+ private static final String STREAM_NAME = "STREAM";
+ private static final String SHARD_ID = "SHARD_ID";
+ @Mock
+ private SimplifiedKinesisClient client;
+
+ @Before
+ public void setUp() throws IOException, TransientKinesisException {
+ when(client.getShardIterator(
+ eq(STREAM_NAME), eq(SHARD_ID), eq(AT_SEQUENCE_NUMBER),
+ anyString(), isNull(Instant.class))).
+ thenReturn(AT_SEQUENCE_SHARD_IT);
+ when(client.getShardIterator(
+ eq(STREAM_NAME), eq(SHARD_ID), eq(AFTER_SEQUENCE_NUMBER),
+ anyString(), isNull(Instant.class))).
+ thenReturn(AFTER_SEQUENCE_SHARD_IT);
+ }
+
+ @Test
+ public void testProvidingShardIterator() throws IOException, TransientKinesisException {
+ assertThat(checkpoint(AT_SEQUENCE_NUMBER, "100", null).getShardIterator(client))
+ .isEqualTo(AT_SEQUENCE_SHARD_IT);
+ assertThat(checkpoint(AFTER_SEQUENCE_NUMBER, "100", null).getShardIterator(client))
+ .isEqualTo(AFTER_SEQUENCE_SHARD_IT);
+ assertThat(checkpoint(AT_SEQUENCE_NUMBER, "100", 10L).getShardIterator(client)).isEqualTo
+ (AT_SEQUENCE_SHARD_IT);
+ assertThat(checkpoint(AFTER_SEQUENCE_NUMBER, "100", 10L).getShardIterator(client))
+ .isEqualTo(AT_SEQUENCE_SHARD_IT);
+ }
+
+ @Test
+ public void testComparisonWithExtendedSequenceNumber() {
+ assertThat(new ShardCheckpoint("", "", new StartingPoint(LATEST)).isBeforeOrAt(
+ recordWith(new ExtendedSequenceNumber("100", 0L))
+ )).isTrue();
+
+ assertThat(new ShardCheckpoint("", "", new StartingPoint(TRIM_HORIZON)).isBeforeOrAt(
+ recordWith(new ExtendedSequenceNumber("100", 0L))
+ )).isTrue();
+
+ assertThat(checkpoint(AFTER_SEQUENCE_NUMBER, "10", 1L).isBeforeOrAt(
+ recordWith(new ExtendedSequenceNumber("100", 0L))
+ )).isTrue();
+
+ assertThat(checkpoint(AT_SEQUENCE_NUMBER, "100", 0L).isBeforeOrAt(
+ recordWith(new ExtendedSequenceNumber("100", 0L))
+ )).isTrue();
+
+ assertThat(checkpoint(AFTER_SEQUENCE_NUMBER, "100", 0L).isBeforeOrAt(
+ recordWith(new ExtendedSequenceNumber("100", 0L))
+ )).isFalse();
+
+ assertThat(checkpoint(AT_SEQUENCE_NUMBER, "100", 1L).isBeforeOrAt(
+ recordWith(new ExtendedSequenceNumber("100", 0L))
+ )).isFalse();
+
+ assertThat(checkpoint(AFTER_SEQUENCE_NUMBER, "100", 0L).isBeforeOrAt(
+ recordWith(new ExtendedSequenceNumber("99", 1L))
+ )).isFalse();
+ }
+
+ @Test
+ public void testComparisonWithTimestamp() {
+ DateTime referenceTimestamp = DateTime.now();
+
+ assertThat(checkpoint(AT_TIMESTAMP, referenceTimestamp.toInstant())
+ .isBeforeOrAt(recordWith(referenceTimestamp.minusMillis(10).toInstant()))
+ ).isFalse();
+
+ assertThat(checkpoint(AT_TIMESTAMP, referenceTimestamp.toInstant())
+ .isBeforeOrAt(recordWith(referenceTimestamp.toInstant()))
+ ).isTrue();
+
+ assertThat(checkpoint(AT_TIMESTAMP, referenceTimestamp.toInstant())
+ .isBeforeOrAt(recordWith(referenceTimestamp.plusMillis(10).toInstant()))
+ ).isTrue();
+ }
+
+ private KinesisRecord recordWith(ExtendedSequenceNumber extendedSequenceNumber) {
+ KinesisRecord record = mock(KinesisRecord.class);
+ given(record.getExtendedSequenceNumber()).willReturn(extendedSequenceNumber);
+ return record;
+ }
+
+ private ShardCheckpoint checkpoint(ShardIteratorType iteratorType, String sequenceNumber,
+ Long subSequenceNumber) {
+ return new ShardCheckpoint(STREAM_NAME, SHARD_ID, iteratorType, sequenceNumber,
+ subSequenceNumber);
+ }
+
+ private KinesisRecord recordWith(Instant approximateArrivalTimestamp) {
+ KinesisRecord record = mock(KinesisRecord.class);
+ given(record.getApproximateArrivalTimestamp()).willReturn(approximateArrivalTimestamp);
+ return record;
+ }
+
+ private ShardCheckpoint checkpoint(ShardIteratorType iteratorType, Instant timestamp) {
+ return new ShardCheckpoint(STREAM_NAME, SHARD_ID, iteratorType, timestamp);
+ }
}
[03/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/dataflow/internal/dependency.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/dataflow/internal/dependency.py b/sdks/python/apache_beam/runners/dataflow/internal/dependency.py
index 62c09ed..e656600 100644
--- a/sdks/python/apache_beam/runners/dataflow/internal/dependency.py
+++ b/sdks/python/apache_beam/runners/dataflow/internal/dependency.py
@@ -69,34 +69,20 @@ from apache_beam.utils import processes
from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import SetupOptions
-# All constants are for internal use only; no backwards-compatibility
-# guarantees.
-# In a released version BEAM_CONTAINER_VERSION and BEAM_FNAPI_CONTAINER_VERSION
-# should match each other, and should be in the same format as the SDK version
-# (i.e. MAJOR.MINOR.PATCH). For non-released (dev) versions, read below.
# Update this version to the next version whenever there is a change that will
-# require changes to legacy Dataflow worker execution environment.
+# require changes to the execution environment.
# This should be in the beam-[version]-[date] format, date is optional.
-BEAM_CONTAINER_VERSION = 'beam-2.1.0-20170626'
-# Update this version to the next version whenever there is a change that
-# requires changes to SDK harness container or SDK harness launcher.
-# This should be in the beam-[version]-[date] format, date is optional.
-BEAM_FNAPI_CONTAINER_VERSION = 'beam-2.1.0-20170621'
+BEAM_CONTAINER_VERSION = 'beam-2.1.0-20170601'
# Standard file names used for staging files.
WORKFLOW_TARBALL_FILE = 'workflow.tar.gz'
REQUIREMENTS_FILE = 'requirements.txt'
EXTRA_PACKAGES_FILE = 'extra_packages.txt'
-# Package names for different distributions
GOOGLE_PACKAGE_NAME = 'google-cloud-dataflow'
BEAM_PACKAGE_NAME = 'apache-beam'
-# SDK identifiers for different distributions
-GOOGLE_SDK_NAME = 'Google Cloud Dataflow SDK for Python'
-BEAM_SDK_NAME = 'Apache Beam SDK for Python'
-
def _dependency_file_copy(from_path, to_path):
"""Copies a local file to a GCS file or vice versa."""
@@ -488,33 +474,10 @@ def _stage_beam_sdk_tarball(sdk_remote_location, staged_path, temp_dir):
'type of location: %s' % sdk_remote_location)
-def get_default_container_image_for_current_sdk(job_type):
+def get_required_container_version():
"""For internal use only; no backwards-compatibility guarantees.
- Args:
- job_type (str): BEAM job type.
-
- Returns:
- str: Google Cloud Dataflow container image for remote execution.
- """
- # TODO(tvalentyn): Use enumerated type instead of strings for job types.
- if job_type == 'FNAPI_BATCH' or job_type == 'FNAPI_STREAMING':
- image_name = 'dataflow.gcr.io/v1beta3/python-fnapi'
- else:
- image_name = 'dataflow.gcr.io/v1beta3/python'
- image_tag = _get_required_container_version(job_type)
- return image_name + ':' + image_tag
-
-
-def _get_required_container_version(job_type=None):
- """For internal use only; no backwards-compatibility guarantees.
-
- Args:
- job_type (str, optional): BEAM job type. Defaults to None.
-
- Returns:
- str: The tag of worker container images in GCR that corresponds to
- current version of the SDK.
+ Returns the Google Cloud Dataflow container version for remote execution.
"""
# TODO(silviuc): Handle apache-beam versions when we have official releases.
import pkg_resources as pkg
@@ -530,34 +493,28 @@ def _get_required_container_version(job_type=None):
except pkg.DistributionNotFound:
# This case covers Apache Beam end-to-end testing scenarios. All these tests
# will run with a special container version.
- if job_type == 'FNAPI_BATCH' or job_type == 'FNAPI_STREAMING':
- return BEAM_FNAPI_CONTAINER_VERSION
- else:
- return BEAM_CONTAINER_VERSION
+ return BEAM_CONTAINER_VERSION
def get_sdk_name_and_version():
"""For internal use only; no backwards-compatibility guarantees.
Returns name and version of SDK reported to Google Cloud Dataflow."""
- import pkg_resources as pkg
- container_version = _get_required_container_version()
- try:
- pkg.get_distribution(GOOGLE_PACKAGE_NAME)
- return (GOOGLE_SDK_NAME, container_version)
- except pkg.DistributionNotFound:
- return (BEAM_SDK_NAME, beam_version.__version__)
+ # TODO(ccy): Make this check cleaner.
+ container_version = get_required_container_version()
+ if container_version == BEAM_CONTAINER_VERSION:
+ return ('Apache Beam SDK for Python', beam_version.__version__)
+ return ('Google Cloud Dataflow SDK for Python', container_version)
def get_sdk_package_name():
"""For internal use only; no backwards-compatibility guarantees.
Returns the PyPI package name to be staged to Google Cloud Dataflow."""
- sdk_name, _ = get_sdk_name_and_version()
- if sdk_name == GOOGLE_SDK_NAME:
- return GOOGLE_PACKAGE_NAME
- else:
+ container_version = get_required_container_version()
+ if container_version == BEAM_CONTAINER_VERSION:
return BEAM_PACKAGE_NAME
+ return GOOGLE_PACKAGE_NAME
def _download_pypi_sdk_package(temp_dir):
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/dataflow/native_io/iobase_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/dataflow/native_io/iobase_test.py b/sdks/python/apache_beam/runners/dataflow/native_io/iobase_test.py
index 3d8c24f..7610baf 100644
--- a/sdks/python/apache_beam/runners/dataflow/native_io/iobase_test.py
+++ b/sdks/python/apache_beam/runners/dataflow/native_io/iobase_test.py
@@ -20,9 +20,7 @@
import unittest
-from apache_beam import Create
-from apache_beam import error
-from apache_beam import pvalue
+from apache_beam import error, pvalue
from apache_beam.runners.dataflow.native_io.iobase import (
_dict_printable_fields,
_NativeWrite,
@@ -30,12 +28,10 @@ from apache_beam.runners.dataflow.native_io.iobase import (
DynamicSplitRequest,
DynamicSplitResultWithPosition,
NativeSink,
- NativeSinkWriter,
NativeSource,
ReaderPosition,
ReaderProgress
)
-from apache_beam.testing.test_pipeline import TestPipeline
class TestHelperFunctions(unittest.TestCase):
@@ -158,39 +154,6 @@ class TestNativeSink(unittest.TestCase):
fake_sink = FakeSink()
self.assertEqual(fake_sink.__repr__(), "<FakeSink ['validate=False']>")
- def test_on_direct_runner(self):
- class FakeSink(NativeSink):
- """A fake sink outputing a number of elements."""
-
- def __init__(self):
- self.written_values = []
- self.writer_instance = FakeSinkWriter(self.written_values)
-
- def writer(self):
- return self.writer_instance
-
- class FakeSinkWriter(NativeSinkWriter):
- """A fake sink writer for testing."""
-
- def __init__(self, written_values):
- self.written_values = written_values
-
- def __enter__(self):
- return self
-
- def __exit__(self, *unused_args):
- pass
-
- def Write(self, value):
- self.written_values.append(value)
-
- p = TestPipeline()
- sink = FakeSink()
- p | Create(['a', 'b', 'c']) | _NativeWrite(sink) # pylint: disable=expression-not-assigned
- p.run()
-
- self.assertEqual(['a', 'b', 'c'], sink.written_values)
-
class Test_NativeWrite(unittest.TestCase):
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/dataflow/native_io/streaming_create.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/dataflow/native_io/streaming_create.py b/sdks/python/apache_beam/runners/dataflow/native_io/streaming_create.py
deleted file mode 100644
index 8c6c8d6..0000000
--- a/sdks/python/apache_beam/runners/dataflow/native_io/streaming_create.py
+++ /dev/null
@@ -1,72 +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.
-#
-
-"""Create transform for streaming."""
-
-from apache_beam import pvalue
-from apache_beam import DoFn
-from apache_beam import ParDo
-from apache_beam import PTransform
-from apache_beam import Windowing
-from apache_beam.transforms.window import GlobalWindows
-
-
-class StreamingCreate(PTransform):
- """A specialized implementation for ``Create`` transform in streaming mode.
-
- Note: There is no unbounded source API in python to wrap the Create source,
- so we map this to composite of Impulse primitive and an SDF.
- """
-
- def __init__(self, values, coder):
- self.coder = coder
- self.encoded_values = map(coder.encode, values)
-
- class DecodeAndEmitDoFn(DoFn):
- """A DoFn which stores encoded versions of elements.
-
- It also stores a Coder to decode and emit those elements.
- TODO: BEAM-2422 - Make this a SplittableDoFn.
- """
-
- def __init__(self, encoded_values, coder):
- self.encoded_values = encoded_values
- self.coder = coder
-
- def process(self, unused_element):
- for encoded_value in self.encoded_values:
- yield self.coder.decode(encoded_value)
-
- class Impulse(PTransform):
- """The Dataflow specific override for the impulse primitive."""
-
- def expand(self, pbegin):
- assert isinstance(pbegin, pvalue.PBegin), (
- 'Input to Impulse transform must be a PBegin but found %s' % pbegin)
- return pvalue.PCollection(pbegin.pipeline)
-
- def get_windowing(self, inputs):
- return Windowing(GlobalWindows())
-
- def infer_output_type(self, unused_input_type):
- return bytes
-
- def expand(self, pbegin):
- return (pbegin
- | 'Impulse' >> self.Impulse()
- | 'Decode Values' >> ParDo(
- self.DecodeAndEmitDoFn(self.encoded_values, self.coder)))
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/dataflow/ptransform_overrides.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/dataflow/ptransform_overrides.py b/sdks/python/apache_beam/runners/dataflow/ptransform_overrides.py
deleted file mode 100644
index 680a4b7..0000000
--- a/sdks/python/apache_beam/runners/dataflow/ptransform_overrides.py
+++ /dev/null
@@ -1,52 +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.
-#
-
-"""Ptransform overrides for DataflowRunner."""
-
-from apache_beam.coders import typecoders
-from apache_beam.pipeline import PTransformOverride
-
-
-class CreatePTransformOverride(PTransformOverride):
- """A ``PTransformOverride`` for ``Create`` in streaming mode."""
-
- def get_matcher(self):
- return self.is_streaming_create
-
- @staticmethod
- def is_streaming_create(applied_ptransform):
- # Imported here to avoid circular dependencies.
- # pylint: disable=wrong-import-order, wrong-import-position
- from apache_beam import Create
- from apache_beam.options.pipeline_options import StandardOptions
-
- if isinstance(applied_ptransform.transform, Create):
- standard_options = (applied_ptransform
- .outputs[None]
- .pipeline._options
- .view_as(StandardOptions))
- return standard_options.streaming
- else:
- return False
-
- def get_replacement_transform(self, ptransform):
- # Imported here to avoid circular dependencies.
- # pylint: disable=wrong-import-order, wrong-import-position
- from apache_beam.runners.dataflow.native_io.streaming_create import \
- StreamingCreate
- coder = typecoders.registry.get_coder(ptransform.get_output_type())
- return StreamingCreate(ptransform.value, coder)
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/direct/bundle_factory.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/direct/bundle_factory.py b/sdks/python/apache_beam/runners/direct/bundle_factory.py
index 0182b4c..ed00b03 100644
--- a/sdks/python/apache_beam/runners/direct/bundle_factory.py
+++ b/sdks/python/apache_beam/runners/direct/bundle_factory.py
@@ -108,7 +108,7 @@ class _Bundle(object):
self._initial_windowed_value.windows)
def __init__(self, pcollection, stacked=True):
- assert isinstance(pcollection, (pvalue.PBegin, pvalue.PCollection))
+ assert isinstance(pcollection, pvalue.PCollection)
self._pcollection = pcollection
self._elements = []
self._stacked = stacked
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/direct/direct_runner.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/direct/direct_runner.py b/sdks/python/apache_beam/runners/direct/direct_runner.py
index 1a94b3d..ecf5114 100644
--- a/sdks/python/apache_beam/runners/direct/direct_runner.py
+++ b/sdks/python/apache_beam/runners/direct/direct_runner.py
@@ -26,55 +26,22 @@ from __future__ import absolute_import
import collections
import logging
-import apache_beam as beam
-from apache_beam import typehints
from apache_beam.metrics.execution import MetricsEnvironment
-from apache_beam.pvalue import PCollection
from apache_beam.runners.direct.bundle_factory import BundleFactory
from apache_beam.runners.runner import PipelineResult
from apache_beam.runners.runner import PipelineRunner
from apache_beam.runners.runner import PipelineState
from apache_beam.runners.runner import PValueCache
-from apache_beam.transforms.core import _GroupAlsoByWindow
-from apache_beam.transforms.core import _GroupByKeyOnly
from apache_beam.options.pipeline_options import DirectOptions
-from apache_beam.options.pipeline_options import StandardOptions
from apache_beam.options.value_provider import RuntimeValueProvider
__all__ = ['DirectRunner']
-# Type variables.
-K = typehints.TypeVariable('K')
-V = typehints.TypeVariable('V')
-
-
-@typehints.with_input_types(typehints.KV[K, V])
-@typehints.with_output_types(typehints.KV[K, typehints.Iterable[V]])
-class _StreamingGroupByKeyOnly(_GroupByKeyOnly):
- """Streaming GroupByKeyOnly placeholder for overriding in DirectRunner."""
- pass
-
-
-@typehints.with_input_types(typehints.KV[K, typehints.Iterable[V]])
-@typehints.with_output_types(typehints.KV[K, typehints.Iterable[V]])
-class _StreamingGroupAlsoByWindow(_GroupAlsoByWindow):
- """Streaming GroupAlsoByWindow placeholder for overriding in DirectRunner."""
- pass
-
-
class DirectRunner(PipelineRunner):
"""Executes a single pipeline on the local machine."""
- # A list of PTransformOverride objects to be applied before running a pipeline
- # using DirectRunner.
- # Currently this only works for overrides where the input and output types do
- # not change.
- # For internal SDK use only. This should not be updated by Beam pipeline
- # authors.
- _PTRANSFORM_OVERRIDES = []
-
def __init__(self):
self._cache = None
@@ -89,84 +56,9 @@ class DirectRunner(PipelineRunner):
except NotImplementedError:
return transform.expand(pcoll)
- def apply__GroupByKeyOnly(self, transform, pcoll):
- if (transform.__class__ == _GroupByKeyOnly and
- pcoll.pipeline._options.view_as(StandardOptions).streaming):
- # Use specialized streaming implementation, if requested.
- type_hints = transform.get_type_hints()
- return pcoll | (_StreamingGroupByKeyOnly()
- .with_input_types(*type_hints.input_types[0])
- .with_output_types(*type_hints.output_types[0]))
- return transform.expand(pcoll)
-
- def apply__GroupAlsoByWindow(self, transform, pcoll):
- if (transform.__class__ == _GroupAlsoByWindow and
- pcoll.pipeline._options.view_as(StandardOptions).streaming):
- # Use specialized streaming implementation, if requested.
- type_hints = transform.get_type_hints()
- return pcoll | (_StreamingGroupAlsoByWindow(transform.windowing)
- .with_input_types(*type_hints.input_types[0])
- .with_output_types(*type_hints.output_types[0]))
- return transform.expand(pcoll)
-
- def apply_ReadStringsFromPubSub(self, transform, pcoll):
- try:
- from google.cloud import pubsub as unused_pubsub
- except ImportError:
- raise ImportError('Google Cloud PubSub not available, please install '
- 'apache_beam[gcp]')
- # Execute this as a native transform.
- output = PCollection(pcoll.pipeline)
- output.element_type = unicode
- return output
-
- def apply_WriteStringsToPubSub(self, transform, pcoll):
- try:
- from google.cloud import pubsub
- except ImportError:
- raise ImportError('Google Cloud PubSub not available, please install '
- 'apache_beam[gcp]')
- project = transform._sink.project
- topic_name = transform._sink.topic_name
-
- class DirectWriteToPubSub(beam.DoFn):
- _topic = None
-
- def __init__(self, project, topic_name):
- self.project = project
- self.topic_name = topic_name
-
- def start_bundle(self):
- if self._topic is None:
- self._topic = pubsub.Client(project=self.project).topic(
- self.topic_name)
- self._buffer = []
-
- def process(self, elem):
- self._buffer.append(elem.encode('utf-8'))
- if len(self._buffer) >= 100:
- self._flush()
-
- def finish_bundle(self):
- self._flush()
-
- def _flush(self):
- if self._buffer:
- with self._topic.batch() as batch:
- for datum in self._buffer:
- batch.publish(datum)
- self._buffer = []
-
- output = pcoll | beam.ParDo(DirectWriteToPubSub(project, topic_name))
- output.element_type = unicode
- return output
-
def run(self, pipeline):
"""Execute the entire pipeline and returns an DirectPipelineResult."""
- # Performing configured PTransform overrides.
- pipeline.replace_all(DirectRunner._PTRANSFORM_OVERRIDES)
-
# TODO: Move imports to top. Pipeline <-> Runner dependency cause problems
# with resolving imports when they are at top.
# pylint: disable=wrong-import-position
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/direct/evaluation_context.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/direct/evaluation_context.py b/sdks/python/apache_beam/runners/direct/evaluation_context.py
index 54c407c..68d99d3 100644
--- a/sdks/python/apache_beam/runners/direct/evaluation_context.py
+++ b/sdks/python/apache_beam/runners/direct/evaluation_context.py
@@ -27,22 +27,22 @@ from apache_beam.runners.direct.clock import Clock
from apache_beam.runners.direct.watermark_manager import WatermarkManager
from apache_beam.runners.direct.executor import TransformExecutor
from apache_beam.runners.direct.direct_metrics import DirectMetrics
-from apache_beam.transforms.trigger import InMemoryUnmergedState
from apache_beam.utils import counters
class _ExecutionContext(object):
- def __init__(self, watermarks, keyed_states):
- self.watermarks = watermarks
- self.keyed_states = keyed_states
+ def __init__(self, watermarks, existing_state):
+ self._watermarks = watermarks
+ self._existing_state = existing_state
- self._step_context = None
+ @property
+ def watermarks(self):
+ return self._watermarks
- def get_step_context(self):
- if not self._step_context:
- self._step_context = DirectStepContext(self.keyed_states)
- return self._step_context
+ @property
+ def existing_state(self):
+ return self._existing_state
class _SideInputView(object):
@@ -145,11 +145,11 @@ class EvaluationContext(object):
self._pcollection_to_views = collections.defaultdict(list)
for view in views:
self._pcollection_to_views[view.pvalue].append(view)
- self._transform_keyed_states = self._initialize_keyed_states(
- root_transforms, value_to_consumers)
+
+ # AppliedPTransform -> Evaluator specific state objects
+ self._application_state_interals = {}
self._watermark_manager = WatermarkManager(
- Clock(), root_transforms, value_to_consumers,
- self._transform_keyed_states)
+ Clock(), root_transforms, value_to_consumers)
self._side_inputs_container = _SideInputsContainer(views)
self._pending_unblocked_tasks = []
self._counter_factory = counters.CounterFactory()
@@ -158,15 +158,6 @@ class EvaluationContext(object):
self._lock = threading.Lock()
- def _initialize_keyed_states(self, root_transforms, value_to_consumers):
- transform_keyed_states = {}
- for transform in root_transforms:
- transform_keyed_states[transform] = {}
- for consumers in value_to_consumers.values():
- for consumer in consumers:
- transform_keyed_states[consumer] = {}
- return transform_keyed_states
-
def use_pvalue_cache(self, cache):
assert not self._cache
self._cache = cache
@@ -208,12 +199,11 @@ class EvaluationContext(object):
the committed bundles contained within the handled result.
"""
with self._lock:
- committed_bundles, unprocessed_bundles = self._commit_bundles(
- result.uncommitted_output_bundles,
- result.unprocessed_bundles)
+ committed_bundles = self._commit_bundles(
+ result.uncommitted_output_bundles)
self._watermark_manager.update_watermarks(
completed_bundle, result.transform, completed_timers,
- committed_bundles, unprocessed_bundles, result.keyed_watermark_holds)
+ committed_bundles, result.watermark_hold)
self._metrics.commit_logical(completed_bundle,
result.logical_metric_updates)
@@ -241,6 +231,7 @@ class EvaluationContext(object):
counter.name, counter.combine_fn)
merged_counter.accumulator.merge([counter.accumulator])
+ self._application_state_interals[result.transform] = result.state
return committed_bundles
def get_aggregator_values(self, aggregator_or_name):
@@ -253,22 +244,19 @@ class EvaluationContext(object):
executor_service.submit(task)
self._pending_unblocked_tasks = []
- def _commit_bundles(self, uncommitted_bundles, unprocessed_bundles):
+ def _commit_bundles(self, uncommitted_bundles):
"""Commits bundles and returns a immutable set of committed bundles."""
for in_progress_bundle in uncommitted_bundles:
producing_applied_ptransform = in_progress_bundle.pcollection.producer
watermarks = self._watermark_manager.get_watermarks(
producing_applied_ptransform)
in_progress_bundle.commit(watermarks.synchronized_processing_output_time)
-
- for unprocessed_bundle in unprocessed_bundles:
- unprocessed_bundle.commit(None)
- return tuple(uncommitted_bundles), tuple(unprocessed_bundles)
+ return tuple(uncommitted_bundles)
def get_execution_context(self, applied_ptransform):
return _ExecutionContext(
self._watermark_manager.get_watermarks(applied_ptransform),
- self._transform_keyed_states[applied_ptransform])
+ self._application_state_interals.get(applied_ptransform))
def create_bundle(self, output_pcollection):
"""Create an uncommitted bundle for the specified PCollection."""
@@ -308,24 +296,3 @@ class EvaluationContext(object):
assert isinstance(task, TransformExecutor)
return self._side_inputs_container.get_value_or_schedule_after_output(
side_input, task)
-
-
-class DirectUnmergedState(InMemoryUnmergedState):
- """UnmergedState implementation for the DirectRunner."""
-
- def __init__(self):
- super(DirectUnmergedState, self).__init__(defensive_copy=False)
-
-
-class DirectStepContext(object):
- """Context for the currently-executing step."""
-
- def __init__(self, keyed_existing_state):
- self.keyed_existing_state = keyed_existing_state
-
- def get_keyed_state(self, key):
- # TODO(ccy): consider implementing transactional copy on write semantics
- # for state so that work items can be safely retried.
- if not self.keyed_existing_state.get(key):
- self.keyed_existing_state[key] = DirectUnmergedState()
- return self.keyed_existing_state[key]
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/direct/executor.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/direct/executor.py b/sdks/python/apache_beam/runners/direct/executor.py
index e70e326..86db291 100644
--- a/sdks/python/apache_beam/runners/direct/executor.py
+++ b/sdks/python/apache_beam/runners/direct/executor.py
@@ -20,7 +20,6 @@
from __future__ import absolute_import
import collections
-import itertools
import logging
import Queue
import sys
@@ -222,30 +221,22 @@ class _CompletionCallback(object):
or for a source transform.
"""
- def __init__(self, evaluation_context, all_updates, timer_firings=None):
+ def __init__(self, evaluation_context, all_updates, timers=None):
self._evaluation_context = evaluation_context
self._all_updates = all_updates
- self._timer_firings = timer_firings or []
+ self._timers = timers
- def handle_result(self, transform_executor, input_committed_bundle,
- transform_result):
+ def handle_result(self, input_committed_bundle, transform_result):
output_committed_bundles = self._evaluation_context.handle_result(
- input_committed_bundle, self._timer_firings, transform_result)
+ input_committed_bundle, self._timers, transform_result)
for output_committed_bundle in output_committed_bundles:
self._all_updates.offer(_ExecutorServiceParallelExecutor._ExecutorUpdate(
- transform_executor,
- committed_bundle=output_committed_bundle))
- for unprocessed_bundle in transform_result.unprocessed_bundles:
- self._all_updates.offer(
- _ExecutorServiceParallelExecutor._ExecutorUpdate(
- transform_executor,
- unprocessed_bundle=unprocessed_bundle))
+ output_committed_bundle, None))
return output_committed_bundles
- def handle_exception(self, transform_executor, exception):
+ def handle_exception(self, exception):
self._all_updates.offer(
- _ExecutorServiceParallelExecutor._ExecutorUpdate(
- transform_executor, exception=exception))
+ _ExecutorServiceParallelExecutor._ExecutorUpdate(None, exception))
class TransformExecutor(_ExecutorService.CallableTask):
@@ -259,13 +250,12 @@ class TransformExecutor(_ExecutorService.CallableTask):
"""
def __init__(self, transform_evaluator_registry, evaluation_context,
- input_bundle, fired_timers, applied_ptransform,
- completion_callback, transform_evaluation_state):
+ input_bundle, applied_transform, completion_callback,
+ transform_evaluation_state):
self._transform_evaluator_registry = transform_evaluator_registry
self._evaluation_context = evaluation_context
self._input_bundle = input_bundle
- self._fired_timers = fired_timers
- self._applied_ptransform = applied_ptransform
+ self._applied_transform = applied_transform
self._completion_callback = completion_callback
self._transform_evaluation_state = transform_evaluation_state
self._side_input_values = {}
@@ -274,11 +264,11 @@ class TransformExecutor(_ExecutorService.CallableTask):
def call(self):
self._call_count += 1
- assert self._call_count <= (1 + len(self._applied_ptransform.side_inputs))
- metrics_container = MetricsContainer(self._applied_ptransform.full_label)
+ assert self._call_count <= (1 + len(self._applied_transform.side_inputs))
+ metrics_container = MetricsContainer(self._applied_transform.full_label)
scoped_metrics_container = ScopedMetricsContainer(metrics_container)
- for side_input in self._applied_ptransform.side_inputs:
+ for side_input in self._applied_transform.side_inputs:
if side_input not in self._side_input_values:
has_result, value = (
self._evaluation_context.get_value_or_schedule_after_output(
@@ -290,17 +280,13 @@ class TransformExecutor(_ExecutorService.CallableTask):
self._side_input_values[side_input] = value
side_input_values = [self._side_input_values[side_input]
- for side_input in self._applied_ptransform.side_inputs]
+ for side_input in self._applied_transform.side_inputs]
try:
- evaluator = self._transform_evaluator_registry.get_evaluator(
- self._applied_ptransform, self._input_bundle,
+ evaluator = self._transform_evaluator_registry.for_application(
+ self._applied_transform, self._input_bundle,
side_input_values, scoped_metrics_container)
- if self._fired_timers:
- for timer_firing in self._fired_timers:
- evaluator.process_timer_wrapper(timer_firing)
-
if self._input_bundle:
for value in self._input_bundle.get_elements_iterable():
evaluator.process_element(value)
@@ -312,18 +298,18 @@ class TransformExecutor(_ExecutorService.CallableTask):
if self._evaluation_context.has_cache:
for uncommitted_bundle in result.uncommitted_output_bundles:
self._evaluation_context.append_to_cache(
- self._applied_ptransform, uncommitted_bundle.tag,
+ self._applied_transform, uncommitted_bundle.tag,
uncommitted_bundle.get_elements_iterable())
undeclared_tag_values = result.undeclared_tag_values
if undeclared_tag_values:
for tag, value in undeclared_tag_values.iteritems():
self._evaluation_context.append_to_cache(
- self._applied_ptransform, tag, value)
+ self._applied_transform, tag, value)
- self._completion_callback.handle_result(self, self._input_bundle, result)
+ self._completion_callback.handle_result(self._input_bundle, result)
return result
except Exception as e: # pylint: disable=broad-except
- self._completion_callback.handle_exception(self, e)
+ self._completion_callback.handle_exception(e)
finally:
self._evaluation_context.metrics().commit_physical(
self._input_bundle,
@@ -367,15 +353,6 @@ class _ExecutorServiceParallelExecutor(object):
def start(self, roots):
self.root_nodes = frozenset(roots)
- self.all_nodes = frozenset(
- itertools.chain(
- roots,
- *itertools.chain(self.value_to_consumers.values())))
- self.node_to_pending_bundles = {}
- for root_node in self.root_nodes:
- provider = (self.transform_evaluator_registry
- .get_root_bundle_provider(root_node))
- self.node_to_pending_bundles[root_node] = provider.get_root_bundles()
self.executor_service.submit(
_ExecutorServiceParallelExecutor._MonitorTask(self))
@@ -392,30 +369,26 @@ class _ExecutorServiceParallelExecutor(object):
if committed_bundle.pcollection in self.value_to_consumers:
consumers = self.value_to_consumers[committed_bundle.pcollection]
for applied_ptransform in consumers:
- self.schedule_consumption(applied_ptransform, committed_bundle, [],
+ self.schedule_consumption(applied_ptransform, committed_bundle,
self.default_completion_callback)
- def schedule_unprocessed_bundle(self, applied_ptransform,
- unprocessed_bundle):
- self.node_to_pending_bundles[applied_ptransform].append(unprocessed_bundle)
-
- def schedule_consumption(self, consumer_applied_ptransform, committed_bundle,
- fired_timers, on_complete):
+ def schedule_consumption(self, consumer_applied_transform, committed_bundle,
+ on_complete):
"""Schedules evaluation of the given bundle with the transform."""
- assert consumer_applied_ptransform
- assert committed_bundle
- assert on_complete
- if self.transform_evaluator_registry.should_execute_serially(
- consumer_applied_ptransform):
+ assert all([consumer_applied_transform, on_complete])
+ assert committed_bundle or consumer_applied_transform in self.root_nodes
+ if (committed_bundle
+ and self.transform_evaluator_registry.should_execute_serially(
+ consumer_applied_transform)):
transform_executor_service = self.transform_executor_services.serial(
- consumer_applied_ptransform)
+ consumer_applied_transform)
else:
transform_executor_service = self.transform_executor_services.parallel()
transform_executor = TransformExecutor(
self.transform_evaluator_registry, self.evaluation_context,
- committed_bundle, fired_timers, consumer_applied_ptransform,
- on_complete, transform_executor_service)
+ committed_bundle, consumer_applied_transform, on_complete,
+ transform_executor_service)
transform_executor_service.schedule(transform_executor)
class _TypedUpdateQueue(object):
@@ -445,16 +418,10 @@ class _ExecutorServiceParallelExecutor(object):
class _ExecutorUpdate(object):
"""An internal status update on the state of the executor."""
- def __init__(self, transform_executor, committed_bundle=None,
- unprocessed_bundle=None, exception=None):
- self.transform_executor = transform_executor
+ def __init__(self, produced_bundle=None, exception=None):
# Exactly one of them should be not-None
- assert sum([
- bool(committed_bundle),
- bool(unprocessed_bundle),
- bool(exception)]) == 1
- self.committed_bundle = committed_bundle
- self.unprocessed_bundle = unprocessed_bundle
+ assert bool(produced_bundle) != bool(exception)
+ self.committed_bundle = produced_bundle
self.exception = exception
self.exc_info = sys.exc_info()
if self.exc_info[1] is not exception:
@@ -489,10 +456,6 @@ class _ExecutorServiceParallelExecutor(object):
while update:
if update.committed_bundle:
self._executor.schedule_consumers(update.committed_bundle)
- elif update.unprocessed_bundle:
- self._executor.schedule_unprocessed_bundle(
- update.transform_executor._applied_ptransform,
- update.unprocessed_bundle)
else:
assert update.exception
logging.warning('A task failed with exception.\n %s',
@@ -554,21 +517,19 @@ class _ExecutorServiceParallelExecutor(object):
Returns:
True if timers fired.
"""
- transform_fired_timers = (
- self._executor.evaluation_context.extract_fired_timers())
- for applied_ptransform, fired_timers in transform_fired_timers:
+ fired_timers = self._executor.evaluation_context.extract_fired_timers()
+ for applied_ptransform in fired_timers:
# Use an empty committed bundle. just to trigger.
empty_bundle = (
self._executor.evaluation_context.create_empty_committed_bundle(
applied_ptransform.inputs[0]))
timer_completion_callback = _CompletionCallback(
self._executor.evaluation_context, self._executor.all_updates,
- timer_firings=fired_timers)
+ applied_ptransform)
self._executor.schedule_consumption(
- applied_ptransform, empty_bundle, fired_timers,
- timer_completion_callback)
- return bool(transform_fired_timers)
+ applied_ptransform, empty_bundle, timer_completion_callback)
+ return bool(fired_timers)
def _is_executing(self):
"""Returns True if there is at least one non-blocked TransformExecutor."""
@@ -603,14 +564,10 @@ class _ExecutorServiceParallelExecutor(object):
# additional work.
return
- # All current TransformExecutors are blocked; add more work from any
- # pending bundles.
- for applied_ptransform in self._executor.all_nodes:
- if not self._executor.evaluation_context.is_done(applied_ptransform):
- pending_bundles = self._executor.node_to_pending_bundles.get(
- applied_ptransform, [])
- for bundle in pending_bundles:
- self._executor.schedule_consumption(
- applied_ptransform, bundle, [],
- self._executor.default_completion_callback)
- self._executor.node_to_pending_bundles[applied_ptransform] = []
+ # All current TransformExecutors are blocked; add more work from the
+ # roots.
+ for applied_transform in self._executor.root_nodes:
+ if not self._executor.evaluation_context.is_done(applied_transform):
+ self._executor.schedule_consumption(
+ applied_transform, None,
+ self._executor.default_completion_callback)
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/direct/transform_evaluator.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/direct/transform_evaluator.py b/sdks/python/apache_beam/runners/direct/transform_evaluator.py
index cb2ace2..b1cb626 100644
--- a/sdks/python/apache_beam/runners/direct/transform_evaluator.py
+++ b/sdks/python/apache_beam/runners/direct/transform_evaluator.py
@@ -20,8 +20,6 @@
from __future__ import absolute_import
import collections
-import random
-import time
from apache_beam import coders
from apache_beam import pvalue
@@ -29,29 +27,16 @@ from apache_beam.internal import pickler
import apache_beam.io as io
from apache_beam.runners.common import DoFnRunner
from apache_beam.runners.common import DoFnState
-from apache_beam.runners.direct.direct_runner import _StreamingGroupByKeyOnly
-from apache_beam.runners.direct.direct_runner import _StreamingGroupAlsoByWindow
from apache_beam.runners.direct.watermark_manager import WatermarkManager
-from apache_beam.runners.direct.util import KeyedWorkItem
-from apache_beam.runners.direct.util import TransformResult
+from apache_beam.runners.direct.transform_result import TransformResult
from apache_beam.runners.dataflow.native_io.iobase import _NativeWrite # pylint: disable=protected-access
-from apache_beam.testing.test_stream import TestStream
-from apache_beam.testing.test_stream import ElementEvent
-from apache_beam.testing.test_stream import WatermarkEvent
-from apache_beam.testing.test_stream import ProcessingTimeEvent
from apache_beam.transforms import core
from apache_beam.transforms.window import GlobalWindows
from apache_beam.transforms.window import WindowedValue
-from apache_beam.transforms.trigger import create_trigger_driver
-from apache_beam.transforms.trigger import _CombiningValueStateTag
-from apache_beam.transforms.trigger import _ListStateTag
-from apache_beam.transforms.trigger import TimeDomain
from apache_beam.typehints.typecheck import OutputCheckWrapperDoFn
from apache_beam.typehints.typecheck import TypeCheckError
from apache_beam.typehints.typecheck import TypeCheckWrapperDoFn
from apache_beam.utils import counters
-from apache_beam.utils.timestamp import Timestamp
-from apache_beam.utils.timestamp import MIN_TIMESTAMP
from apache_beam.options.pipeline_options import TypeOptions
@@ -66,21 +51,13 @@ class TransformEvaluatorRegistry(object):
self._evaluation_context = evaluation_context
self._evaluators = {
io.Read: _BoundedReadEvaluator,
- io.ReadStringsFromPubSub: _PubSubReadEvaluator,
core.Flatten: _FlattenEvaluator,
core.ParDo: _ParDoEvaluator,
core._GroupByKeyOnly: _GroupByKeyOnlyEvaluator,
- _StreamingGroupByKeyOnly: _StreamingGroupByKeyOnlyEvaluator,
- _StreamingGroupAlsoByWindow: _StreamingGroupAlsoByWindowEvaluator,
_NativeWrite: _NativeWriteEvaluator,
- TestStream: _TestStreamEvaluator,
- }
- self._root_bundle_providers = {
- core.PTransform: DefaultRootBundleProvider,
- TestStream: _TestStreamRootBundleProvider,
}
- def get_evaluator(
+ def for_application(
self, applied_ptransform, input_committed_bundle,
side_inputs, scoped_metrics_container):
"""Returns a TransformEvaluator suitable for processing given inputs."""
@@ -102,18 +79,6 @@ class TransformEvaluatorRegistry(object):
input_committed_bundle, side_inputs,
scoped_metrics_container)
- def get_root_bundle_provider(self, applied_ptransform):
- provider_cls = None
- for cls in applied_ptransform.transform.__class__.mro():
- provider_cls = self._root_bundle_providers.get(cls)
- if provider_cls:
- break
- if not provider_cls:
- raise NotImplementedError(
- 'Root provider for [%s] not implemented in runner %s' % (
- type(applied_ptransform.transform), self))
- return provider_cls(self._evaluation_context, applied_ptransform)
-
def should_execute_serially(self, applied_ptransform):
"""Returns True if this applied_ptransform should run one bundle at a time.
@@ -134,48 +99,7 @@ class TransformEvaluatorRegistry(object):
True if executor should execute applied_ptransform serially.
"""
return isinstance(applied_ptransform.transform,
- (core._GroupByKeyOnly,
- _StreamingGroupByKeyOnly,
- _StreamingGroupAlsoByWindow,
- _NativeWrite,))
-
-
-class RootBundleProvider(object):
- """Provides bundles for the initial execution of a root transform."""
-
- def __init__(self, evaluation_context, applied_ptransform):
- self._evaluation_context = evaluation_context
- self._applied_ptransform = applied_ptransform
-
- def get_root_bundles(self):
- raise NotImplementedError
-
-
-class DefaultRootBundleProvider(RootBundleProvider):
- """Provides an empty bundle by default for root transforms."""
-
- def get_root_bundles(self):
- input_node = pvalue.PBegin(self._applied_ptransform.transform.pipeline)
- empty_bundle = (
- self._evaluation_context.create_empty_committed_bundle(input_node))
- return [empty_bundle]
-
-
-class _TestStreamRootBundleProvider(RootBundleProvider):
- """Provides an initial bundle for the TestStream evaluator."""
-
- def get_root_bundles(self):
- test_stream = self._applied_ptransform.transform
- bundles = []
- if len(test_stream.events) > 0:
- bundle = self._evaluation_context.create_bundle(
- pvalue.PBegin(self._applied_ptransform.transform.pipeline))
- # Explicitly set timestamp to MIN_TIMESTAMP to ensure that we hold the
- # watermark.
- bundle.add(GlobalWindows.windowed_value(0, timestamp=MIN_TIMESTAMP))
- bundle.commit(None)
- bundles.append(bundle)
- return bundles
+ (core._GroupByKeyOnly, _NativeWrite))
class _TransformEvaluator(object):
@@ -237,27 +161,6 @@ class _TransformEvaluator(object):
"""Starts a new bundle."""
pass
- def process_timer_wrapper(self, timer_firing):
- """Process timer by clearing and then calling process_timer().
-
- This method is called with any timer firing and clears the delivered
- timer from the keyed state and then calls process_timer(). The default
- process_timer() implementation emits a KeyedWorkItem for the particular
- timer and passes it to process_element(). Evaluator subclasses which
- desire different timer delivery semantics can override process_timer().
- """
- state = self.step_context.get_keyed_state(timer_firing.encoded_key)
- state.clear_timer(
- timer_firing.window, timer_firing.name, timer_firing.time_domain)
- self.process_timer(timer_firing)
-
- def process_timer(self, timer_firing):
- """Default process_timer() impl. generating KeyedWorkItem element."""
- self.process_element(
- GlobalWindows.windowed_value(
- KeyedWorkItem(timer_firing.encoded_key,
- timer_firings=[timer_firing])))
-
def process_element(self, element):
"""Processes a new element as part of the current bundle."""
raise NotImplementedError('%s do not process elements.', type(self))
@@ -275,6 +178,7 @@ class _BoundedReadEvaluator(_TransformEvaluator):
def __init__(self, evaluation_context, applied_ptransform,
input_committed_bundle, side_inputs, scoped_metrics_container):
+ assert not input_committed_bundle
assert not side_inputs
self._source = applied_ptransform.transform.source
self._source.pipeline_options = evaluation_context.pipeline_options
@@ -303,148 +207,7 @@ class _BoundedReadEvaluator(_TransformEvaluator):
bundles = _read_values_to_bundles(reader)
return TransformResult(
- self._applied_ptransform, bundles, [], None, None)
-
-
-class _TestStreamEvaluator(_TransformEvaluator):
- """TransformEvaluator for the TestStream transform."""
-
- def __init__(self, evaluation_context, applied_ptransform,
- input_committed_bundle, side_inputs, scoped_metrics_container):
- assert not side_inputs
- self.test_stream = applied_ptransform.transform
- super(_TestStreamEvaluator, self).__init__(
- evaluation_context, applied_ptransform, input_committed_bundle,
- side_inputs, scoped_metrics_container)
-
- def start_bundle(self):
- self.current_index = -1
- self.watermark = MIN_TIMESTAMP
- self.bundles = []
-
- def process_element(self, element):
- index = element.value
- self.watermark = element.timestamp
- assert isinstance(index, int)
- assert 0 <= index <= len(self.test_stream.events)
- self.current_index = index
- event = self.test_stream.events[self.current_index]
- if isinstance(event, ElementEvent):
- assert len(self._outputs) == 1
- output_pcollection = list(self._outputs)[0]
- bundle = self._evaluation_context.create_bundle(output_pcollection)
- for tv in event.timestamped_values:
- bundle.output(
- GlobalWindows.windowed_value(tv.value, timestamp=tv.timestamp))
- self.bundles.append(bundle)
- elif isinstance(event, WatermarkEvent):
- assert event.new_watermark >= self.watermark
- self.watermark = event.new_watermark
- elif isinstance(event, ProcessingTimeEvent):
- # TODO(ccy): advance processing time in the context's mock clock.
- pass
- else:
- raise ValueError('Invalid TestStream event: %s.' % event)
-
- def finish_bundle(self):
- unprocessed_bundles = []
- hold = None
- if self.current_index < len(self.test_stream.events) - 1:
- unprocessed_bundle = self._evaluation_context.create_bundle(
- pvalue.PBegin(self._applied_ptransform.transform.pipeline))
- unprocessed_bundle.add(GlobalWindows.windowed_value(
- self.current_index + 1, timestamp=self.watermark))
- unprocessed_bundles.append(unprocessed_bundle)
- hold = self.watermark
- return TransformResult(
- self._applied_ptransform, self.bundles, unprocessed_bundles, None,
- {None: hold})
-
-
-class _PubSubSubscriptionWrapper(object):
- """Wrapper for garbage-collecting temporary PubSub subscriptions."""
-
- def __init__(self, subscription, should_cleanup):
- self.subscription = subscription
- self.should_cleanup = should_cleanup
-
- def __del__(self):
- if self.should_cleanup:
- self.subscription.delete()
-
-
-class _PubSubReadEvaluator(_TransformEvaluator):
- """TransformEvaluator for PubSub read."""
-
- _subscription_cache = {}
-
- def __init__(self, evaluation_context, applied_ptransform,
- input_committed_bundle, side_inputs, scoped_metrics_container):
- assert not side_inputs
- super(_PubSubReadEvaluator, self).__init__(
- evaluation_context, applied_ptransform, input_committed_bundle,
- side_inputs, scoped_metrics_container)
-
- source = self._applied_ptransform.transform._source
- self._subscription = _PubSubReadEvaluator.get_subscription(
- self._applied_ptransform, source.project, source.topic_name,
- source.subscription_name)
-
- @classmethod
- def get_subscription(cls, transform, project, topic, subscription_name):
- if transform not in cls._subscription_cache:
- from google.cloud import pubsub
- should_create = not subscription_name
- if should_create:
- subscription_name = 'beam_%d_%x' % (
- int(time.time()), random.randrange(1 << 32))
- cls._subscription_cache[transform] = _PubSubSubscriptionWrapper(
- pubsub.Client(project=project).topic(topic).subscription(
- subscription_name),
- should_create)
- if should_create:
- cls._subscription_cache[transform].subscription.create()
- return cls._subscription_cache[transform].subscription
-
- def start_bundle(self):
- pass
-
- def process_element(self, element):
- pass
-
- def _read_from_pubsub(self):
- from google.cloud import pubsub
- # Because of the AutoAck, we are not able to reread messages if this
- # evaluator fails with an exception before emitting a bundle. However,
- # the DirectRunner currently doesn't retry work items anyway, so the
- # pipeline would enter an inconsistent state on any error.
- with pubsub.subscription.AutoAck(
- self._subscription, return_immediately=True,
- max_messages=10) as results:
- return [message.data for unused_ack_id, message in results.items()]
-
- def finish_bundle(self):
- data = self._read_from_pubsub()
- if data:
- output_pcollection = list(self._outputs)[0]
- bundle = self._evaluation_context.create_bundle(output_pcollection)
- # TODO(ccy): we currently do not use the PubSub message timestamp or
- # respect the PubSub source's id_label field.
- now = Timestamp.of(time.time())
- for message_data in data:
- bundle.output(GlobalWindows.windowed_value(message_data, timestamp=now))
- bundles = [bundle]
- else:
- bundles = []
- if self._applied_ptransform.inputs:
- input_pvalue = self._applied_ptransform.inputs[0]
- else:
- input_pvalue = pvalue.PBegin(self._applied_ptransform.transform.pipeline)
- unprocessed_bundle = self._evaluation_context.create_bundle(
- input_pvalue)
- return TransformResult(
- self._applied_ptransform, bundles,
- [unprocessed_bundle], None, {None: Timestamp.of(time.time())})
+ self._applied_ptransform, bundles, None, None, None, None)
class _FlattenEvaluator(_TransformEvaluator):
@@ -468,7 +231,7 @@ class _FlattenEvaluator(_TransformEvaluator):
def finish_bundle(self):
bundles = [self.bundle]
return TransformResult(
- self._applied_ptransform, bundles, [], None, None)
+ self._applied_ptransform, bundles, None, None, None, None)
class _TaggedReceivers(dict):
@@ -557,7 +320,7 @@ class _ParDoEvaluator(_TransformEvaluator):
bundles = self._tagged_receivers.values()
result_counters = self._counter_factory.get_counters()
return TransformResult(
- self._applied_ptransform, bundles, [], result_counters, None,
+ self._applied_ptransform, bundles, None, None, result_counters, None,
self._tagged_receivers.undeclared_in_memory_tag_values)
@@ -565,8 +328,13 @@ class _GroupByKeyOnlyEvaluator(_TransformEvaluator):
"""TransformEvaluator for _GroupByKeyOnly transform."""
MAX_ELEMENT_PER_BUNDLE = None
- ELEMENTS_TAG = _ListStateTag('elements')
- COMPLETION_TAG = _CombiningValueStateTag('completed', any)
+
+ class _GroupByKeyOnlyEvaluatorState(object):
+
+ def __init__(self):
+ # output: {} key -> [values]
+ self.output = collections.defaultdict(list)
+ self.completed = False
def __init__(self, evaluation_context, applied_ptransform,
input_committed_bundle, side_inputs, scoped_metrics_container):
@@ -575,13 +343,15 @@ class _GroupByKeyOnlyEvaluator(_TransformEvaluator):
evaluation_context, applied_ptransform, input_committed_bundle,
side_inputs, scoped_metrics_container)
+ @property
def _is_final_bundle(self):
return (self._execution_context.watermarks.input_watermark
== WatermarkManager.WATERMARK_POS_INF)
def start_bundle(self):
- self.step_context = self._execution_context.get_step_context()
- self.global_state = self.step_context.get_keyed_state(None)
+ self.state = (self._execution_context.existing_state
+ if self._execution_context.existing_state
+ else _GroupByKeyOnlyEvaluator._GroupByKeyOnlyEvaluatorState())
assert len(self._outputs) == 1
self.output_pcollection = list(self._outputs)[0]
@@ -591,44 +361,29 @@ class _GroupByKeyOnlyEvaluator(_TransformEvaluator):
self._applied_ptransform.transform.get_type_hints().input_types[0])
self.key_coder = coders.registry.get_coder(kv_type_hint[0].tuple_types[0])
- def process_timer(self, timer_firing):
- # We do not need to emit a KeyedWorkItem to process_element().
- pass
-
def process_element(self, element):
- assert not self.global_state.get_state(
- None, _GroupByKeyOnlyEvaluator.COMPLETION_TAG)
+ assert not self.state.completed
if (isinstance(element, WindowedValue)
and isinstance(element.value, collections.Iterable)
and len(element.value) == 2):
k, v = element.value
- encoded_k = self.key_coder.encode(k)
- state = self.step_context.get_keyed_state(encoded_k)
- state.add_state(None, _GroupByKeyOnlyEvaluator.ELEMENTS_TAG, v)
+ self.state.output[self.key_coder.encode(k)].append(v)
else:
raise TypeCheckError('Input to _GroupByKeyOnly must be a PCollection of '
'windowed key-value pairs. Instead received: %r.'
% element)
def finish_bundle(self):
- if self._is_final_bundle():
- if self.global_state.get_state(
- None, _GroupByKeyOnlyEvaluator.COMPLETION_TAG):
+ if self._is_final_bundle:
+ if self.state.completed:
# Ignore empty bundles after emitting output. (This may happen because
# empty bundles do not affect input watermarks.)
bundles = []
else:
- gbk_result = []
- # TODO(ccy): perhaps we can clean this up to not use this
- # internal attribute of the DirectStepContext.
- for encoded_k in self.step_context.keyed_existing_state:
- # Ignore global state.
- if encoded_k is None:
- continue
- k = self.key_coder.decode(encoded_k)
- state = self.step_context.get_keyed_state(encoded_k)
- vs = state.get_state(None, _GroupByKeyOnlyEvaluator.ELEMENTS_TAG)
- gbk_result.append(GlobalWindows.windowed_value((k, vs)))
+ gbk_result = (
+ map(GlobalWindows.windowed_value, (
+ (self.key_coder.decode(k), v)
+ for k, v in self.state.output.iteritems())))
def len_element_fn(element):
_, v = element.value
@@ -638,139 +393,21 @@ class _GroupByKeyOnlyEvaluator(_TransformEvaluator):
self.output_pcollection, gbk_result,
_GroupByKeyOnlyEvaluator.MAX_ELEMENT_PER_BUNDLE, len_element_fn)
- self.global_state.add_state(
- None, _GroupByKeyOnlyEvaluator.COMPLETION_TAG, True)
+ self.state.completed = True
+ state = self.state
hold = WatermarkManager.WATERMARK_POS_INF
else:
bundles = []
+ state = self.state
hold = WatermarkManager.WATERMARK_NEG_INF
- self.global_state.set_timer(
- None, '', TimeDomain.WATERMARK, WatermarkManager.WATERMARK_POS_INF)
return TransformResult(
- self._applied_ptransform, bundles, [], None, {None: hold})
-
-
-class _StreamingGroupByKeyOnlyEvaluator(_TransformEvaluator):
- """TransformEvaluator for _StreamingGroupByKeyOnly transform.
-
- The _GroupByKeyOnlyEvaluator buffers elements until its input watermark goes
- to infinity, which is suitable for batch mode execution. During streaming
- mode execution, we emit each bundle as it comes to the next transform.
- """
-
- MAX_ELEMENT_PER_BUNDLE = None
-
- def __init__(self, evaluation_context, applied_ptransform,
- input_committed_bundle, side_inputs, scoped_metrics_container):
- assert not side_inputs
- super(_StreamingGroupByKeyOnlyEvaluator, self).__init__(
- evaluation_context, applied_ptransform, input_committed_bundle,
- side_inputs, scoped_metrics_container)
-
- def start_bundle(self):
- self.gbk_items = collections.defaultdict(list)
-
- assert len(self._outputs) == 1
- self.output_pcollection = list(self._outputs)[0]
-
- # The input type of a GroupByKey will be KV[Any, Any] or more specific.
- kv_type_hint = (
- self._applied_ptransform.transform.get_type_hints().input_types[0])
- self.key_coder = coders.registry.get_coder(kv_type_hint[0].tuple_types[0])
-
- def process_element(self, element):
- if (isinstance(element, WindowedValue)
- and isinstance(element.value, collections.Iterable)
- and len(element.value) == 2):
- k, v = element.value
- self.gbk_items[self.key_coder.encode(k)].append(v)
- else:
- raise TypeCheckError('Input to _GroupByKeyOnly must be a PCollection of '
- 'windowed key-value pairs. Instead received: %r.'
- % element)
-
- def finish_bundle(self):
- bundles = []
- bundle = None
- for encoded_k, vs in self.gbk_items.iteritems():
- if not bundle:
- bundle = self._evaluation_context.create_bundle(
- self.output_pcollection)
- bundles.append(bundle)
- kwi = KeyedWorkItem(encoded_k, elements=vs)
- bundle.add(GlobalWindows.windowed_value(kwi))
-
- return TransformResult(
- self._applied_ptransform, bundles, [], None, None)
-
-
-class _StreamingGroupAlsoByWindowEvaluator(_TransformEvaluator):
- """TransformEvaluator for the _StreamingGroupAlsoByWindow transform.
-
- This evaluator is only used in streaming mode. In batch mode, the
- GroupAlsoByWindow operation is evaluated as a normal DoFn, as defined
- in transforms/core.py.
- """
-
- def __init__(self, evaluation_context, applied_ptransform,
- input_committed_bundle, side_inputs, scoped_metrics_container):
- assert not side_inputs
- super(_StreamingGroupAlsoByWindowEvaluator, self).__init__(
- evaluation_context, applied_ptransform, input_committed_bundle,
- side_inputs, scoped_metrics_container)
-
- def start_bundle(self):
- assert len(self._outputs) == 1
- self.output_pcollection = list(self._outputs)[0]
- self.step_context = self._execution_context.get_step_context()
- self.driver = create_trigger_driver(
- self._applied_ptransform.transform.windowing)
- self.gabw_items = []
- self.keyed_holds = {}
-
- # The input type of a GroupAlsoByWindow will be KV[Any, Iter[Any]] or more
- # specific.
- kv_type_hint = (
- self._applied_ptransform.transform.get_type_hints().input_types[0])
- self.key_coder = coders.registry.get_coder(kv_type_hint[0].tuple_types[0])
-
- def process_element(self, element):
- kwi = element.value
- assert isinstance(kwi, KeyedWorkItem), kwi
- encoded_k, timer_firings, vs = (
- kwi.encoded_key, kwi.timer_firings, kwi.elements)
- k = self.key_coder.decode(encoded_k)
- state = self.step_context.get_keyed_state(encoded_k)
-
- for timer_firing in timer_firings:
- for wvalue in self.driver.process_timer(
- timer_firing.window, timer_firing.name, timer_firing.time_domain,
- timer_firing.timestamp, state):
- self.gabw_items.append(wvalue.with_value((k, wvalue.value)))
- if vs:
- for wvalue in self.driver.process_elements(state, vs, MIN_TIMESTAMP):
- self.gabw_items.append(wvalue.with_value((k, wvalue.value)))
-
- self.keyed_holds[encoded_k] = state.get_earliest_hold()
-
- def finish_bundle(self):
- bundles = []
- if self.gabw_items:
- bundle = self._evaluation_context.create_bundle(self.output_pcollection)
- for item in self.gabw_items:
- bundle.add(item)
- bundles.append(bundle)
-
- return TransformResult(
- self._applied_ptransform, bundles, [], None, self.keyed_holds)
+ self._applied_ptransform, bundles, state, None, None, hold)
class _NativeWriteEvaluator(_TransformEvaluator):
"""TransformEvaluator for _NativeWrite transform."""
- ELEMENTS_TAG = _ListStateTag('elements')
-
def __init__(self, evaluation_context, applied_ptransform,
input_committed_bundle, side_inputs, scoped_metrics_container):
assert not side_inputs
@@ -792,16 +429,12 @@ class _NativeWriteEvaluator(_TransformEvaluator):
== WatermarkManager.WATERMARK_POS_INF)
def start_bundle(self):
- self.step_context = self._execution_context.get_step_context()
- self.global_state = self.step_context.get_keyed_state(None)
-
- def process_timer(self, timer_firing):
- # We do not need to emit a KeyedWorkItem to process_element().
- pass
+ # state: [values]
+ self.state = (self._execution_context.existing_state
+ if self._execution_context.existing_state else [])
def process_element(self, element):
- self.global_state.add_state(
- None, _NativeWriteEvaluator.ELEMENTS_TAG, element)
+ self.state.append(element)
def finish_bundle(self):
# finish_bundle will append incoming bundles in memory until all the bundles
@@ -811,21 +444,19 @@ class _NativeWriteEvaluator(_TransformEvaluator):
# ignored and would not generate additional output files.
# TODO(altay): Do not wait until the last bundle to write in a single shard.
if self._is_final_bundle:
- elements = self.global_state.get_state(
- None, _NativeWriteEvaluator.ELEMENTS_TAG)
if self._has_already_produced_output:
# Ignore empty bundles that arrive after the output is produced.
- assert elements == []
+ assert self.state == []
else:
self._sink.pipeline_options = self._evaluation_context.pipeline_options
with self._sink.writer() as writer:
- for v in elements:
+ for v in self.state:
writer.Write(v.value)
+ state = None
hold = WatermarkManager.WATERMARK_POS_INF
else:
+ state = self.state
hold = WatermarkManager.WATERMARK_NEG_INF
- self.global_state.set_timer(
- None, '', TimeDomain.WATERMARK, WatermarkManager.WATERMARK_POS_INF)
return TransformResult(
- self._applied_ptransform, [], [], None, {None: hold})
+ self._applied_ptransform, [], state, None, None, hold)
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/direct/transform_result.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/direct/transform_result.py b/sdks/python/apache_beam/runners/direct/transform_result.py
new file mode 100644
index 0000000..febdd20
--- /dev/null
+++ b/sdks/python/apache_beam/runners/direct/transform_result.py
@@ -0,0 +1,41 @@
+#
+# 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.
+#
+
+"""The result of evaluating an AppliedPTransform with a TransformEvaluator."""
+
+from __future__ import absolute_import
+
+
+class TransformResult(object):
+ """For internal use only; no backwards-compatibility guarantees.
+
+ The result of evaluating an AppliedPTransform with a TransformEvaluator."""
+
+ def __init__(self, applied_ptransform, uncommitted_output_bundles, state,
+ timer_update, counters, watermark_hold,
+ undeclared_tag_values=None):
+ self.transform = applied_ptransform
+ self.uncommitted_output_bundles = uncommitted_output_bundles
+ self.state = state
+ # TODO: timer update is currently unused.
+ self.timer_update = timer_update
+ self.counters = counters
+ self.watermark_hold = watermark_hold
+ # Only used when caching (materializing) all values is requested.
+ self.undeclared_tag_values = undeclared_tag_values
+ # Populated by the TransformExecutor.
+ self.logical_metric_updates = None
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/direct/util.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/direct/util.py b/sdks/python/apache_beam/runners/direct/util.py
deleted file mode 100644
index 10f7b29..0000000
--- a/sdks/python/apache_beam/runners/direct/util.py
+++ /dev/null
@@ -1,67 +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.
-#
-
-"""Utility classes used by the DirectRunner.
-
-For internal use only. No backwards compatibility guarantees.
-"""
-
-from __future__ import absolute_import
-
-
-class TransformResult(object):
- """Result of evaluating an AppliedPTransform with a TransformEvaluator."""
-
- def __init__(self, applied_ptransform, uncommitted_output_bundles,
- unprocessed_bundles, counters, keyed_watermark_holds,
- undeclared_tag_values=None):
- self.transform = applied_ptransform
- self.uncommitted_output_bundles = uncommitted_output_bundles
- self.unprocessed_bundles = unprocessed_bundles
- self.counters = counters
- # Mapping of key -> earliest hold timestamp or None. Keys should be
- # strings or None.
- #
- # For each key, we receive as its corresponding value the earliest
- # watermark hold for that key (the key can be None for global state), past
- # which the output watermark for the currently-executing step will not
- # advance. If the value is None or utils.timestamp.MAX_TIMESTAMP, the
- # watermark hold will be removed.
- self.keyed_watermark_holds = keyed_watermark_holds or {}
- # Only used when caching (materializing) all values is requested.
- self.undeclared_tag_values = undeclared_tag_values
- # Populated by the TransformExecutor.
- self.logical_metric_updates = None
-
-
-class TimerFiring(object):
- """A single instance of a fired timer."""
-
- def __init__(self, encoded_key, window, name, time_domain, timestamp):
- self.encoded_key = encoded_key
- self.window = window
- self.name = name
- self.time_domain = time_domain
- self.timestamp = timestamp
-
-
-class KeyedWorkItem(object):
- """A keyed item that can either be a timer firing or a list of elements."""
- def __init__(self, encoded_key, timer_firings=None, elements=None):
- self.encoded_key = encoded_key
- self.timer_firings = timer_firings or []
- self.elements = elements or []
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/direct/watermark_manager.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/direct/watermark_manager.py b/sdks/python/apache_beam/runners/direct/watermark_manager.py
index 935998d..3a13539 100644
--- a/sdks/python/apache_beam/runners/direct/watermark_manager.py
+++ b/sdks/python/apache_beam/runners/direct/watermark_manager.py
@@ -23,10 +23,8 @@ import threading
from apache_beam import pipeline
from apache_beam import pvalue
-from apache_beam.runners.direct.util import TimerFiring
from apache_beam.utils.timestamp import MAX_TIMESTAMP
from apache_beam.utils.timestamp import MIN_TIMESTAMP
-from apache_beam.utils.timestamp import TIME_GRANULARITY
class WatermarkManager(object):
@@ -37,23 +35,21 @@ class WatermarkManager(object):
WATERMARK_POS_INF = MAX_TIMESTAMP
WATERMARK_NEG_INF = MIN_TIMESTAMP
- def __init__(self, clock, root_transforms, value_to_consumers,
- transform_keyed_states):
+ def __init__(self, clock, root_transforms, value_to_consumers):
self._clock = clock # processing time clock
- self._root_transforms = root_transforms
self._value_to_consumers = value_to_consumers
- self._transform_keyed_states = transform_keyed_states
+ self._root_transforms = root_transforms
# AppliedPTransform -> TransformWatermarks
self._transform_to_watermarks = {}
for root_transform in root_transforms:
self._transform_to_watermarks[root_transform] = _TransformWatermarks(
- self._clock, transform_keyed_states[root_transform], root_transform)
+ self._clock)
for consumers in value_to_consumers.values():
for consumer in consumers:
self._transform_to_watermarks[consumer] = _TransformWatermarks(
- self._clock, transform_keyed_states[consumer], consumer)
+ self._clock)
for consumers in value_to_consumers.values():
for consumer in consumers:
@@ -93,19 +89,16 @@ class WatermarkManager(object):
return self._transform_to_watermarks[applied_ptransform]
def update_watermarks(self, completed_committed_bundle, applied_ptransform,
- completed_timers, outputs, unprocessed_bundles,
- keyed_earliest_holds):
+ timer_update, outputs, earliest_hold):
assert isinstance(applied_ptransform, pipeline.AppliedPTransform)
self._update_pending(
- completed_committed_bundle, applied_ptransform, completed_timers,
- outputs, unprocessed_bundles)
+ completed_committed_bundle, applied_ptransform, timer_update, outputs)
tw = self.get_watermarks(applied_ptransform)
- tw.hold(keyed_earliest_holds)
+ tw.hold(earliest_hold)
self._refresh_watermarks(applied_ptransform)
def _update_pending(self, input_committed_bundle, applied_ptransform,
- completed_timers, output_committed_bundles,
- unprocessed_bundles):
+ timer_update, output_committed_bundles):
"""Updated list of pending bundles for the given AppliedPTransform."""
# Update pending elements. Filter out empty bundles. They do not impact
@@ -119,10 +112,7 @@ class WatermarkManager(object):
consumer_tw.add_pending(output)
completed_tw = self._transform_to_watermarks[applied_ptransform]
- completed_tw.update_timers(completed_timers)
-
- for unprocessed_bundle in unprocessed_bundles:
- completed_tw.add_pending(unprocessed_bundle)
+ completed_tw.update_timers(timer_update)
assert input_committed_bundle or applied_ptransform in self._root_transforms
if input_committed_bundle and input_committed_bundle.has_elements():
@@ -146,36 +136,33 @@ class WatermarkManager(object):
def extract_fired_timers(self):
all_timers = []
for applied_ptransform, tw in self._transform_to_watermarks.iteritems():
- fired_timers = tw.extract_fired_timers()
- if fired_timers:
- all_timers.append((applied_ptransform, fired_timers))
+ if tw.extract_fired_timers():
+ all_timers.append(applied_ptransform)
return all_timers
class _TransformWatermarks(object):
- """Tracks input and output watermarks for an AppliedPTransform."""
+ """Tracks input and output watermarks for aan AppliedPTransform."""
- def __init__(self, clock, keyed_states, transform):
+ def __init__(self, clock):
self._clock = clock
- self._keyed_states = keyed_states
self._input_transform_watermarks = []
self._input_watermark = WatermarkManager.WATERMARK_NEG_INF
self._output_watermark = WatermarkManager.WATERMARK_NEG_INF
- self._keyed_earliest_holds = {}
+ self._earliest_hold = WatermarkManager.WATERMARK_POS_INF
self._pending = set() # Scheduled bundles targeted for this transform.
- self._fired_timers = set()
+ self._fired_timers = False
self._lock = threading.Lock()
- self._label = str(transform)
-
def update_input_transform_watermarks(self, input_transform_watermarks):
with self._lock:
self._input_transform_watermarks = input_transform_watermarks
- def update_timers(self, completed_timers):
+ def update_timers(self, timer_update):
with self._lock:
- for timer_firing in completed_timers:
- self._fired_timers.remove(timer_firing)
+ if timer_update:
+ assert self._fired_timers
+ self._fired_timers = False
@property
def input_watermark(self):
@@ -187,13 +174,11 @@ class _TransformWatermarks(object):
with self._lock:
return self._output_watermark
- def hold(self, keyed_earliest_holds):
+ def hold(self, value):
with self._lock:
- for key, hold_value in keyed_earliest_holds.iteritems():
- self._keyed_earliest_holds[key] = hold_value
- if (hold_value is None or
- hold_value == WatermarkManager.WATERMARK_POS_INF):
- del self._keyed_earliest_holds[key]
+ if value is None:
+ value = WatermarkManager.WATERMARK_POS_INF
+ self._earliest_hold = value
def add_pending(self, pending):
with self._lock:
@@ -208,22 +193,9 @@ class _TransformWatermarks(object):
def refresh(self):
with self._lock:
- min_pending_timestamp = WatermarkManager.WATERMARK_POS_INF
- has_pending_elements = False
- for input_bundle in self._pending:
- # TODO(ccy): we can have the Bundle class keep track of the minimum
- # timestamp so we don't have to do an iteration here.
- for wv in input_bundle.get_elements_iterable():
- has_pending_elements = True
- if wv.timestamp < min_pending_timestamp:
- min_pending_timestamp = wv.timestamp
-
- # If there is a pending element with a certain timestamp, we can at most
- # advance our watermark to the maximum timestamp less than that
- # timestamp.
- pending_holder = WatermarkManager.WATERMARK_POS_INF
- if has_pending_elements:
- pending_holder = min_pending_timestamp - TIME_GRANULARITY
+ pending_holder = (WatermarkManager.WATERMARK_NEG_INF
+ if self._pending else
+ WatermarkManager.WATERMARK_POS_INF)
input_watermarks = [
tw.output_watermark for tw in self._input_transform_watermarks]
@@ -232,11 +204,7 @@ class _TransformWatermarks(object):
self._input_watermark = max(self._input_watermark,
min(pending_holder, producer_watermark))
- earliest_hold = WatermarkManager.WATERMARK_POS_INF
- for hold in self._keyed_earliest_holds.values():
- if hold < earliest_hold:
- earliest_hold = hold
- new_output_watermark = min(self._input_watermark, earliest_hold)
+ new_output_watermark = min(self._input_watermark, self._earliest_hold)
advanced = new_output_watermark > self._output_watermark
self._output_watermark = new_output_watermark
@@ -251,12 +219,8 @@ class _TransformWatermarks(object):
if self._fired_timers:
return False
- fired_timers = []
- for encoded_key, state in self._keyed_states.iteritems():
- timers = state.get_timers(watermark=self._input_watermark)
- for expired in timers:
- window, (name, time_domain, timestamp) = expired
- fired_timers.append(
- TimerFiring(encoded_key, window, name, time_domain, timestamp))
- self._fired_timers.update(fired_timers)
- return fired_timers
+ should_fire = (
+ self._earliest_hold < WatermarkManager.WATERMARK_POS_INF and
+ self._input_watermark == WatermarkManager.WATERMARK_POS_INF)
+ self._fired_timers = should_fire
+ return should_fire
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/pipeline_context.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/pipeline_context.py b/sdks/python/apache_beam/runners/pipeline_context.py
index a40069b..1c89d06 100644
--- a/sdks/python/apache_beam/runners/pipeline_context.py
+++ b/sdks/python/apache_beam/runners/pipeline_context.py
@@ -24,8 +24,7 @@ For internal use only; no backwards-compatibility guarantees.
from apache_beam import pipeline
from apache_beam import pvalue
from apache_beam import coders
-from apache_beam.portability.api import beam_fn_api_pb2
-from apache_beam.portability.api import beam_runner_api_pb2
+from apache_beam.runners.api import beam_runner_api_pb2
from apache_beam.transforms import core
@@ -43,10 +42,9 @@ class _PipelineContextMap(object):
self._id_to_proto = proto_map if proto_map else {}
self._counter = 0
- def _unique_ref(self, obj=None):
+ def _unique_ref(self):
self._counter += 1
- return "ref_%s_%s_%s" % (
- self._obj_type.__name__, type(obj).__name__, self._counter)
+ return "ref_%s_%s" % (self._obj_type.__name__, self._counter)
def populate_map(self, proto_map):
for id, proto in self._id_to_proto.items():
@@ -54,7 +52,7 @@ class _PipelineContextMap(object):
def get_id(self, obj):
if obj not in self._obj_to_id:
- id = self._unique_ref(obj)
+ id = self._unique_ref()
self._id_to_obj[id] = obj
self._obj_to_id[obj] = id
self._id_to_proto[id] = obj.to_runner_api(self._pipeline_context)
@@ -81,16 +79,11 @@ class PipelineContext(object):
# TODO: environment
}
- def __init__(self, proto=None):
- if isinstance(proto, beam_fn_api_pb2.ProcessBundleDescriptor):
- proto = beam_runner_api_pb2.Components(
- coders=dict(proto.coders.items()),
- windowing_strategies=dict(proto.windowing_strategies.items()),
- environments=dict(proto.environments.items()))
+ def __init__(self, context_proto=None):
for name, cls in self._COMPONENT_TYPES.items():
setattr(
self, name, _PipelineContextMap(
- self, cls, getattr(proto, name, None)))
+ self, cls, getattr(context_proto, name, None)))
@staticmethod
def from_runner_api(proto):
[07/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
index 482f5a2..b69bc83 100644
--- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
+++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
@@ -83,7 +83,6 @@ import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.kafka.clients.consumer.Consumer;
-import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.MockConsumer;
import org.apache.kafka.clients.consumer.OffsetResetStrategy;
@@ -365,35 +364,6 @@ public class KafkaIOTest {
}
@Test
- public void testUnreachableKafkaBrokers() {
- // Expect an exception when the Kafka brokers are not reachable on the workers.
- // We specify partitions explicitly so that splitting does not involve server interaction.
- // Set request timeout to 10ms so that test does not take long.
-
- thrown.expect(Exception.class);
- thrown.expectMessage("Reader-0: Timeout while initializing partition 'test-0'");
-
- int numElements = 1000;
- PCollection<Long> input = p
- .apply(KafkaIO.<Integer, Long>read()
- .withBootstrapServers("8.8.8.8:9092") // Google public DNS ip.
- .withTopicPartitions(ImmutableList.of(new TopicPartition("test", 0)))
- .withKeyDeserializer(IntegerDeserializer.class)
- .withValueDeserializer(LongDeserializer.class)
- .updateConsumerProperties(ImmutableMap.<String, Object>of(
- ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG, 10,
- ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 5,
- ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, 8,
- ConsumerConfig.FETCH_MAX_WAIT_MS_CONFIG, 8))
- .withMaxNumRecords(10)
- .withoutMetadata())
- .apply(Values.<Long>create());
-
- addCountingAsserts(input, numElements);
- p.run();
- }
-
- @Test
public void testUnboundedSourceWithSingleTopic() {
// same as testUnboundedSource, but with single topic
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/pom.xml b/sdks/java/io/kinesis/pom.xml
index 46d5e26..cb7064b 100644
--- a/sdks/java/io/kinesis/pom.xml
+++ b/sdks/java/io/kinesis/pom.xml
@@ -21,7 +21,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/CheckpointGenerator.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/CheckpointGenerator.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/CheckpointGenerator.java
index 2629c57..919d85a 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/CheckpointGenerator.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/CheckpointGenerator.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.sdk.io.kinesis;
+
import java.io.Serializable;
/**
@@ -24,7 +25,6 @@ import java.io.Serializable;
* How exactly the checkpoint is generated is up to implementing class.
*/
interface CheckpointGenerator extends Serializable {
-
- KinesisReaderCheckpoint generate(SimplifiedKinesisClient client)
- throws TransientKinesisException;
+ KinesisReaderCheckpoint generate(SimplifiedKinesisClient client)
+ throws TransientKinesisException;
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/CustomOptional.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/CustomOptional.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/CustomOptional.java
index 5a28214..4bed0e3 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/CustomOptional.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/CustomOptional.java
@@ -24,79 +24,76 @@ import java.util.Objects;
* Similar to Guava {@code Optional}, but throws {@link NoSuchElementException} for missing element.
*/
abstract class CustomOptional<T> {
+ @SuppressWarnings("unchecked")
+ public static <T> CustomOptional<T> absent() {
+ return (Absent<T>) Absent.INSTANCE;
+ }
- @SuppressWarnings("unchecked")
- public static <T> CustomOptional<T> absent() {
- return (Absent<T>) Absent.INSTANCE;
- }
+ public static <T> CustomOptional<T> of(T v) {
+ return new Present<>(v);
+ }
- public static <T> CustomOptional<T> of(T v) {
- return new Present<>(v);
- }
+ public abstract boolean isPresent();
- public abstract boolean isPresent();
+ public abstract T get();
- public abstract T get();
+ private static class Present<T> extends CustomOptional<T> {
+ private final T value;
- private static class Present<T> extends CustomOptional<T> {
+ private Present(T value) {
+ this.value = value;
+ }
- private final T value;
+ @Override
+ public boolean isPresent() {
+ return true;
+ }
- private Present(T value) {
- this.value = value;
- }
+ @Override
+ public T get() {
+ return value;
+ }
- @Override
- public boolean isPresent() {
- return true;
- }
+ @Override
+ public boolean equals(Object o) {
+ if (!(o instanceof Present)) {
+ return false;
+ }
- @Override
- public T get() {
- return value;
- }
-
- @Override
- public boolean equals(Object o) {
- if (!(o instanceof Present)) {
- return false;
- }
+ Present<?> present = (Present<?>) o;
+ return Objects.equals(value, present.value);
+ }
- Present<?> present = (Present<?>) o;
- return Objects.equals(value, present.value);
+ @Override
+ public int hashCode() {
+ return Objects.hash(value);
+ }
}
- @Override
- public int hashCode() {
- return Objects.hash(value);
- }
- }
+ private static class Absent<T> extends CustomOptional<T> {
+ private static final Absent<Object> INSTANCE = new Absent<>();
- private static class Absent<T> extends CustomOptional<T> {
+ private Absent() {
+ }
- private static final Absent<Object> INSTANCE = new Absent<>();
+ @Override
+ public boolean isPresent() {
+ return false;
+ }
- private Absent() {
- }
-
- @Override
- public boolean isPresent() {
- return false;
- }
+ @Override
+ public T get() {
+ throw new NoSuchElementException();
+ }
- @Override
- public T get() {
- throw new NoSuchElementException();
- }
-
- @Override
- public boolean equals(Object o) {
- return o instanceof Absent;
- }
+ @Override
+ public boolean equals(Object o) {
+ return o instanceof Absent;
+ }
- @Override
- public int hashCode() {
- return 0;
+ @Override
+ public int hashCode() {
+ return 0;
+ }
}
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/DynamicCheckpointGenerator.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/DynamicCheckpointGenerator.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/DynamicCheckpointGenerator.java
index 9933019..2ec293c 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/DynamicCheckpointGenerator.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/DynamicCheckpointGenerator.java
@@ -28,31 +28,29 @@ import com.google.common.base.Function;
* List of shards is obtained dynamically on call to {@link #generate(SimplifiedKinesisClient)}.
*/
class DynamicCheckpointGenerator implements CheckpointGenerator {
-
- private final String streamName;
- private final StartingPoint startingPoint;
-
- public DynamicCheckpointGenerator(String streamName, StartingPoint startingPoint) {
- this.streamName = checkNotNull(streamName, "streamName");
- this.startingPoint = checkNotNull(startingPoint, "startingPoint");
- }
-
- @Override
- public KinesisReaderCheckpoint generate(SimplifiedKinesisClient kinesis)
- throws TransientKinesisException {
- return new KinesisReaderCheckpoint(
- transform(kinesis.listShards(streamName), new Function<Shard, ShardCheckpoint>() {
-
- @Override
- public ShardCheckpoint apply(Shard shard) {
- return new ShardCheckpoint(streamName, shard.getShardId(), startingPoint);
- }
- })
- );
- }
-
- @Override
- public String toString() {
- return String.format("Checkpoint generator for %s: %s", streamName, startingPoint);
- }
+ private final String streamName;
+ private final StartingPoint startingPoint;
+
+ public DynamicCheckpointGenerator(String streamName, StartingPoint startingPoint) {
+ this.streamName = checkNotNull(streamName, "streamName");
+ this.startingPoint = checkNotNull(startingPoint, "startingPoint");
+ }
+
+ @Override
+ public KinesisReaderCheckpoint generate(SimplifiedKinesisClient kinesis)
+ throws TransientKinesisException {
+ return new KinesisReaderCheckpoint(
+ transform(kinesis.listShards(streamName), new Function<Shard, ShardCheckpoint>() {
+ @Override
+ public ShardCheckpoint apply(Shard shard) {
+ return new ShardCheckpoint(streamName, shard.getShardId(), startingPoint);
+ }
+ })
+ );
+ }
+
+ @Override
+ public String toString() {
+ return String.format("Checkpoint generator for %s: %s", streamName, startingPoint);
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/GetKinesisRecordsResult.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/GetKinesisRecordsResult.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/GetKinesisRecordsResult.java
index f605f55..5a34d7d 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/GetKinesisRecordsResult.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/GetKinesisRecordsResult.java
@@ -21,7 +21,6 @@ import static com.google.common.collect.Lists.transform;
import com.amazonaws.services.kinesis.clientlibrary.types.UserRecord;
import com.google.common.base.Function;
-
import java.util.List;
import javax.annotation.Nullable;
@@ -29,29 +28,27 @@ import javax.annotation.Nullable;
* Represents the output of 'get' operation on Kinesis stream.
*/
class GetKinesisRecordsResult {
-
- private final List<KinesisRecord> records;
- private final String nextShardIterator;
-
- public GetKinesisRecordsResult(List<UserRecord> records, String nextShardIterator,
- final String streamName, final String shardId) {
- this.records = transform(records, new Function<UserRecord, KinesisRecord>() {
-
- @Nullable
- @Override
- public KinesisRecord apply(@Nullable UserRecord input) {
- assert input != null; // to make FindBugs happy
- return new KinesisRecord(input, streamName, shardId);
- }
- });
- this.nextShardIterator = nextShardIterator;
- }
-
- public List<KinesisRecord> getRecords() {
- return records;
- }
-
- public String getNextShardIterator() {
- return nextShardIterator;
- }
+ private final List<KinesisRecord> records;
+ private final String nextShardIterator;
+
+ public GetKinesisRecordsResult(List<UserRecord> records, String nextShardIterator,
+ final String streamName, final String shardId) {
+ this.records = transform(records, new Function<UserRecord, KinesisRecord>() {
+ @Nullable
+ @Override
+ public KinesisRecord apply(@Nullable UserRecord input) {
+ assert input != null; // to make FindBugs happy
+ return new KinesisRecord(input, streamName, shardId);
+ }
+ });
+ this.nextShardIterator = nextShardIterator;
+ }
+
+ public List<KinesisRecord> getRecords() {
+ return records;
+ }
+
+ public String getNextShardIterator() {
+ return nextShardIterator;
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisClientProvider.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisClientProvider.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisClientProvider.java
index b5b721e..c7fd7f6 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisClientProvider.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisClientProvider.java
@@ -18,7 +18,6 @@
package org.apache.beam.sdk.io.kinesis;
import com.amazonaws.services.kinesis.AmazonKinesis;
-
import java.io.Serializable;
/**
@@ -28,6 +27,5 @@ import java.io.Serializable;
* {@link Serializable} to ensure it can be sent to worker machines.
*/
interface KinesisClientProvider extends Serializable {
-
- AmazonKinesis get();
+ AmazonKinesis get();
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisIO.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisIO.java
index bc8ada1..c97316d 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisIO.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisIO.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.sdk.io.kinesis;
+
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
@@ -28,9 +29,7 @@ import com.amazonaws.services.kinesis.AmazonKinesis;
import com.amazonaws.services.kinesis.AmazonKinesisClient;
import com.amazonaws.services.kinesis.clientlibrary.lib.worker.InitialPositionInStream;
import com.google.auto.value.AutoValue;
-
import javax.annotation.Nullable;
-
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.io.BoundedReadFromUnboundedSource;
import org.apache.beam.sdk.transforms.PTransform;
@@ -101,150 +100,144 @@ import org.joda.time.Instant;
* }</pre>
*
*/
-@Experimental(Experimental.Kind.SOURCE_SINK)
+@Experimental
public final class KinesisIO {
-
- /** Returns a new {@link Read} transform for reading from Kinesis. */
- public static Read read() {
- return new AutoValue_KinesisIO_Read.Builder().setMaxNumRecords(-1).build();
- }
-
- /** Implementation of {@link #read}. */
- @AutoValue
- public abstract static class Read extends PTransform<PBegin, PCollection<KinesisRecord>> {
-
- @Nullable
- abstract String getStreamName();
-
- @Nullable
- abstract StartingPoint getInitialPosition();
-
- @Nullable
- abstract KinesisClientProvider getClientProvider();
-
- abstract int getMaxNumRecords();
-
- @Nullable
- abstract Duration getMaxReadTime();
-
- abstract Builder toBuilder();
-
- @AutoValue.Builder
- abstract static class Builder {
-
- abstract Builder setStreamName(String streamName);
-
- abstract Builder setInitialPosition(StartingPoint startingPoint);
-
- abstract Builder setClientProvider(KinesisClientProvider clientProvider);
-
- abstract Builder setMaxNumRecords(int maxNumRecords);
-
- abstract Builder setMaxReadTime(Duration maxReadTime);
-
- abstract Read build();
+ /** Returns a new {@link Read} transform for reading from Kinesis. */
+ public static Read read() {
+ return new AutoValue_KinesisIO_Read.Builder().setMaxNumRecords(-1).build();
}
- /**
- * Specify reading from streamName at some initial position.
- */
- public Read from(String streamName, InitialPositionInStream initialPosition) {
- return toBuilder()
- .setStreamName(streamName)
- .setInitialPosition(
- new StartingPoint(checkNotNull(initialPosition, "initialPosition")))
- .build();
- }
-
- /**
- * Specify reading from streamName beginning at given {@link Instant}.
- * This {@link Instant} must be in the past, i.e. before {@link Instant#now()}.
- */
- public Read from(String streamName, Instant initialTimestamp) {
- return toBuilder()
- .setStreamName(streamName)
- .setInitialPosition(
- new StartingPoint(checkNotNull(initialTimestamp, "initialTimestamp")))
- .build();
- }
-
- /**
- * Allows to specify custom {@link KinesisClientProvider}.
- * {@link KinesisClientProvider} provides {@link AmazonKinesis} instances which are later
- * used for communication with Kinesis.
- * You should use this method if {@link Read#withClientProvider(String, String, Regions)}
- * does not suit your needs.
- */
- public Read withClientProvider(KinesisClientProvider kinesisClientProvider) {
- return toBuilder().setClientProvider(kinesisClientProvider).build();
- }
-
- /**
- * Specify credential details and region to be used to read from Kinesis.
- * If you need more sophisticated credential protocol, then you should look at
- * {@link Read#withClientProvider(KinesisClientProvider)}.
- */
- public Read withClientProvider(String awsAccessKey, String awsSecretKey, Regions region) {
- return withClientProvider(new BasicKinesisProvider(awsAccessKey, awsSecretKey, region));
- }
-
- /** Specifies to read at most a given number of records. */
- public Read withMaxNumRecords(int maxNumRecords) {
- checkArgument(
- maxNumRecords > 0, "maxNumRecords must be positive, but was: %s", maxNumRecords);
- return toBuilder().setMaxNumRecords(maxNumRecords).build();
- }
-
- /** Specifies to read at most a given number of records. */
- public Read withMaxReadTime(Duration maxReadTime) {
- checkNotNull(maxReadTime, "maxReadTime");
- return toBuilder().setMaxReadTime(maxReadTime).build();
- }
-
- @Override
- public PCollection<KinesisRecord> expand(PBegin input) {
- org.apache.beam.sdk.io.Read.Unbounded<KinesisRecord> read =
- org.apache.beam.sdk.io.Read.from(
- new KinesisSource(getClientProvider(), getStreamName(), getInitialPosition()));
- if (getMaxNumRecords() > 0) {
- BoundedReadFromUnboundedSource<KinesisRecord> bounded =
- read.withMaxNumRecords(getMaxNumRecords());
- return getMaxReadTime() == null
- ? input.apply(bounded)
- : input.apply(bounded.withMaxReadTime(getMaxReadTime()));
- } else {
- return getMaxReadTime() == null
- ? input.apply(read)
- : input.apply(read.withMaxReadTime(getMaxReadTime()));
- }
- }
-
- private static final class BasicKinesisProvider implements KinesisClientProvider {
-
- private final String accessKey;
- private final String secretKey;
- private final Regions region;
-
- private BasicKinesisProvider(String accessKey, String secretKey, Regions region) {
- this.accessKey = checkNotNull(accessKey, "accessKey");
- this.secretKey = checkNotNull(secretKey, "secretKey");
- this.region = checkNotNull(region, "region");
- }
-
- private AWSCredentialsProvider getCredentialsProvider() {
- return new StaticCredentialsProvider(new BasicAWSCredentials(
- accessKey,
- secretKey
- ));
-
- }
-
- @Override
- public AmazonKinesis get() {
- AmazonKinesisClient client = new AmazonKinesisClient(getCredentialsProvider());
- client.withRegion(region);
- return client;
- }
+ /** Implementation of {@link #read}. */
+ @AutoValue
+ public abstract static class Read extends PTransform<PBegin, PCollection<KinesisRecord>> {
+ @Nullable
+ abstract String getStreamName();
+
+ @Nullable
+ abstract StartingPoint getInitialPosition();
+
+ @Nullable
+ abstract KinesisClientProvider getClientProvider();
+
+ abstract int getMaxNumRecords();
+
+ @Nullable
+ abstract Duration getMaxReadTime();
+
+ abstract Builder toBuilder();
+
+ @AutoValue.Builder
+ abstract static class Builder {
+ abstract Builder setStreamName(String streamName);
+ abstract Builder setInitialPosition(StartingPoint startingPoint);
+ abstract Builder setClientProvider(KinesisClientProvider clientProvider);
+ abstract Builder setMaxNumRecords(int maxNumRecords);
+ abstract Builder setMaxReadTime(Duration maxReadTime);
+
+ abstract Read build();
+ }
+
+ /**
+ * Specify reading from streamName at some initial position.
+ */
+ public Read from(String streamName, InitialPositionInStream initialPosition) {
+ return toBuilder()
+ .setStreamName(streamName)
+ .setInitialPosition(
+ new StartingPoint(checkNotNull(initialPosition, "initialPosition")))
+ .build();
+ }
+
+ /**
+ * Specify reading from streamName beginning at given {@link Instant}.
+ * This {@link Instant} must be in the past, i.e. before {@link Instant#now()}.
+ */
+ public Read from(String streamName, Instant initialTimestamp) {
+ return toBuilder()
+ .setStreamName(streamName)
+ .setInitialPosition(
+ new StartingPoint(checkNotNull(initialTimestamp, "initialTimestamp")))
+ .build();
+ }
+
+ /**
+ * Allows to specify custom {@link KinesisClientProvider}.
+ * {@link KinesisClientProvider} provides {@link AmazonKinesis} instances which are later
+ * used for communication with Kinesis.
+ * You should use this method if {@link Read#withClientProvider(String, String, Regions)}
+ * does not suit your needs.
+ */
+ public Read withClientProvider(KinesisClientProvider kinesisClientProvider) {
+ return toBuilder().setClientProvider(kinesisClientProvider).build();
+ }
+
+ /**
+ * Specify credential details and region to be used to read from Kinesis.
+ * If you need more sophisticated credential protocol, then you should look at
+ * {@link Read#withClientProvider(KinesisClientProvider)}.
+ */
+ public Read withClientProvider(String awsAccessKey, String awsSecretKey, Regions region) {
+ return withClientProvider(new BasicKinesisProvider(awsAccessKey, awsSecretKey, region));
+ }
+
+ /** Specifies to read at most a given number of records. */
+ public Read withMaxNumRecords(int maxNumRecords) {
+ checkArgument(
+ maxNumRecords > 0, "maxNumRecords must be positive, but was: %s", maxNumRecords);
+ return toBuilder().setMaxNumRecords(maxNumRecords).build();
+ }
+
+ /** Specifies to read at most a given number of records. */
+ public Read withMaxReadTime(Duration maxReadTime) {
+ checkNotNull(maxReadTime, "maxReadTime");
+ return toBuilder().setMaxReadTime(maxReadTime).build();
+ }
+
+ @Override
+ public PCollection<KinesisRecord> expand(PBegin input) {
+ org.apache.beam.sdk.io.Read.Unbounded<KinesisRecord> read =
+ org.apache.beam.sdk.io.Read.from(
+ new KinesisSource(getClientProvider(), getStreamName(), getInitialPosition()));
+ if (getMaxNumRecords() > 0) {
+ BoundedReadFromUnboundedSource<KinesisRecord> bounded =
+ read.withMaxNumRecords(getMaxNumRecords());
+ return getMaxReadTime() == null
+ ? input.apply(bounded)
+ : input.apply(bounded.withMaxReadTime(getMaxReadTime()));
+ } else {
+ return getMaxReadTime() == null
+ ? input.apply(read)
+ : input.apply(read.withMaxReadTime(getMaxReadTime()));
+ }
+ }
+
+ private static final class BasicKinesisProvider implements KinesisClientProvider {
+
+ private final String accessKey;
+ private final String secretKey;
+ private final Regions region;
+
+ private BasicKinesisProvider(String accessKey, String secretKey, Regions region) {
+ this.accessKey = checkNotNull(accessKey, "accessKey");
+ this.secretKey = checkNotNull(secretKey, "secretKey");
+ this.region = checkNotNull(region, "region");
+ }
+
+
+ private AWSCredentialsProvider getCredentialsProvider() {
+ return new StaticCredentialsProvider(new BasicAWSCredentials(
+ accessKey,
+ secretKey
+ ));
+
+ }
+
+ @Override
+ public AmazonKinesis get() {
+ AmazonKinesisClient client = new AmazonKinesisClient(getCredentialsProvider());
+ client.withRegion(region);
+ return client;
+ }
+ }
}
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisReader.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisReader.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisReader.java
index e5c32d2..2138094 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisReader.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisReader.java
@@ -17,129 +17,129 @@
*/
package org.apache.beam.sdk.io.kinesis;
+
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.collect.Lists.newArrayList;
import java.io.IOException;
import java.util.List;
import java.util.NoSuchElementException;
-
import org.apache.beam.sdk.io.UnboundedSource;
import org.joda.time.Instant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+
/**
* Reads data from multiple kinesis shards in a single thread.
* It uses simple round robin algorithm when fetching data from shards.
*/
class KinesisReader extends UnboundedSource.UnboundedReader<KinesisRecord> {
+ private static final Logger LOG = LoggerFactory.getLogger(KinesisReader.class);
+
+ private final SimplifiedKinesisClient kinesis;
+ private final UnboundedSource<KinesisRecord, ?> source;
+ private final CheckpointGenerator initialCheckpointGenerator;
+ private RoundRobin<ShardRecordsIterator> shardIterators;
+ private CustomOptional<KinesisRecord> currentRecord = CustomOptional.absent();
+
+ public KinesisReader(SimplifiedKinesisClient kinesis,
+ CheckpointGenerator initialCheckpointGenerator,
+ UnboundedSource<KinesisRecord, ?> source) {
+ this.kinesis = checkNotNull(kinesis, "kinesis");
+ this.initialCheckpointGenerator =
+ checkNotNull(initialCheckpointGenerator, "initialCheckpointGenerator");
+ this.source = source;
+ }
+
+ /**
+ * Generates initial checkpoint and instantiates iterators for shards.
+ */
+ @Override
+ public boolean start() throws IOException {
+ LOG.info("Starting reader using {}", initialCheckpointGenerator);
+
+ try {
+ KinesisReaderCheckpoint initialCheckpoint =
+ initialCheckpointGenerator.generate(kinesis);
+ List<ShardRecordsIterator> iterators = newArrayList();
+ for (ShardCheckpoint checkpoint : initialCheckpoint) {
+ iterators.add(checkpoint.getShardRecordsIterator(kinesis));
+ }
+ shardIterators = new RoundRobin<>(iterators);
+ } catch (TransientKinesisException e) {
+ throw new IOException(e);
+ }
- private static final Logger LOG = LoggerFactory.getLogger(KinesisReader.class);
-
- private final SimplifiedKinesisClient kinesis;
- private final UnboundedSource<KinesisRecord, ?> source;
- private final CheckpointGenerator initialCheckpointGenerator;
- private RoundRobin<ShardRecordsIterator> shardIterators;
- private CustomOptional<KinesisRecord> currentRecord = CustomOptional.absent();
-
- public KinesisReader(SimplifiedKinesisClient kinesis,
- CheckpointGenerator initialCheckpointGenerator,
- UnboundedSource<KinesisRecord, ?> source) {
- this.kinesis = checkNotNull(kinesis, "kinesis");
- this.initialCheckpointGenerator =
- checkNotNull(initialCheckpointGenerator, "initialCheckpointGenerator");
- this.source = source;
- }
-
- /**
- * Generates initial checkpoint and instantiates iterators for shards.
- */
- @Override
- public boolean start() throws IOException {
- LOG.info("Starting reader using {}", initialCheckpointGenerator);
-
- try {
- KinesisReaderCheckpoint initialCheckpoint =
- initialCheckpointGenerator.generate(kinesis);
- List<ShardRecordsIterator> iterators = newArrayList();
- for (ShardCheckpoint checkpoint : initialCheckpoint) {
- iterators.add(checkpoint.getShardRecordsIterator(kinesis));
- }
- shardIterators = new RoundRobin<>(iterators);
- } catch (TransientKinesisException e) {
- throw new IOException(e);
+ return advance();
}
- return advance();
- }
-
- /**
- * Moves to the next record in one of the shards.
- * If current shard iterator can be move forward (i.e. there's a record present) then we do it.
- * If not, we iterate over shards in a round-robin manner.
- */
- @Override
- public boolean advance() throws IOException {
- try {
- for (int i = 0; i < shardIterators.size(); ++i) {
- currentRecord = shardIterators.getCurrent().next();
- if (currentRecord.isPresent()) {
- return true;
- } else {
- shardIterators.moveForward();
+ /**
+ * Moves to the next record in one of the shards.
+ * If current shard iterator can be move forward (i.e. there's a record present) then we do it.
+ * If not, we iterate over shards in a round-robin manner.
+ */
+ @Override
+ public boolean advance() throws IOException {
+ try {
+ for (int i = 0; i < shardIterators.size(); ++i) {
+ currentRecord = shardIterators.getCurrent().next();
+ if (currentRecord.isPresent()) {
+ return true;
+ } else {
+ shardIterators.moveForward();
+ }
+ }
+ } catch (TransientKinesisException e) {
+ LOG.warn("Transient exception occurred", e);
}
- }
- } catch (TransientKinesisException e) {
- LOG.warn("Transient exception occurred", e);
+ return false;
+ }
+
+ @Override
+ public byte[] getCurrentRecordId() throws NoSuchElementException {
+ return currentRecord.get().getUniqueId();
+ }
+
+ @Override
+ public KinesisRecord getCurrent() throws NoSuchElementException {
+ return currentRecord.get();
+ }
+
+ /**
+ * When {@link KinesisReader} was advanced to the current record.
+ * We cannot use approximate arrival timestamp given for each record by Kinesis as it
+ * is not guaranteed to be accurate - this could lead to mark some records as "late"
+ * even if they were not.
+ */
+ @Override
+ public Instant getCurrentTimestamp() throws NoSuchElementException {
+ return currentRecord.get().getReadTime();
+ }
+
+ @Override
+ public void close() throws IOException {
+ }
+
+ /**
+ * Current time.
+ * We cannot give better approximation of the watermark with current semantics of
+ * {@link KinesisReader#getCurrentTimestamp()}, because we don't know when the next
+ * {@link KinesisReader#advance()} will be called.
+ */
+ @Override
+ public Instant getWatermark() {
+ return Instant.now();
+ }
+
+ @Override
+ public UnboundedSource.CheckpointMark getCheckpointMark() {
+ return KinesisReaderCheckpoint.asCurrentStateOf(shardIterators);
+ }
+
+ @Override
+ public UnboundedSource<KinesisRecord, ?> getCurrentSource() {
+ return source;
}
- return false;
- }
-
- @Override
- public byte[] getCurrentRecordId() throws NoSuchElementException {
- return currentRecord.get().getUniqueId();
- }
-
- @Override
- public KinesisRecord getCurrent() throws NoSuchElementException {
- return currentRecord.get();
- }
-
- /**
- * When {@link KinesisReader} was advanced to the current record.
- * We cannot use approximate arrival timestamp given for each record by Kinesis as it
- * is not guaranteed to be accurate - this could lead to mark some records as "late"
- * even if they were not.
- */
- @Override
- public Instant getCurrentTimestamp() throws NoSuchElementException {
- return currentRecord.get().getReadTime();
- }
-
- @Override
- public void close() throws IOException {
- }
-
- /**
- * Current time.
- * We cannot give better approximation of the watermark with current semantics of
- * {@link KinesisReader#getCurrentTimestamp()}, because we don't know when the next
- * {@link KinesisReader#advance()} will be called.
- */
- @Override
- public Instant getWatermark() {
- return Instant.now();
- }
-
- @Override
- public UnboundedSource.CheckpointMark getCheckpointMark() {
- return KinesisReaderCheckpoint.asCurrentStateOf(shardIterators);
- }
-
- @Override
- public UnboundedSource<KinesisRecord, ?> getCurrentSource() {
- return source;
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisReaderCheckpoint.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisReaderCheckpoint.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisReaderCheckpoint.java
index d995e75..f0fa45d 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisReaderCheckpoint.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisReaderCheckpoint.java
@@ -23,13 +23,11 @@ import static com.google.common.collect.Lists.partition;
import com.google.common.base.Function;
import com.google.common.collect.ImmutableList;
-
import java.io.IOException;
import java.io.Serializable;
import java.util.Iterator;
import java.util.List;
import javax.annotation.Nullable;
-
import org.apache.beam.sdk.io.UnboundedSource;
/**
@@ -39,61 +37,60 @@ import org.apache.beam.sdk.io.UnboundedSource;
* This class is immutable.
*/
class KinesisReaderCheckpoint implements Iterable<ShardCheckpoint>, UnboundedSource
- .CheckpointMark, Serializable {
-
- private final List<ShardCheckpoint> shardCheckpoints;
+ .CheckpointMark, Serializable {
+ private final List<ShardCheckpoint> shardCheckpoints;
- public KinesisReaderCheckpoint(Iterable<ShardCheckpoint> shardCheckpoints) {
- this.shardCheckpoints = ImmutableList.copyOf(shardCheckpoints);
- }
-
- public static KinesisReaderCheckpoint asCurrentStateOf(Iterable<ShardRecordsIterator>
- iterators) {
- return new KinesisReaderCheckpoint(transform(iterators,
- new Function<ShardRecordsIterator, ShardCheckpoint>() {
-
- @Nullable
- @Override
- public ShardCheckpoint apply(@Nullable
- ShardRecordsIterator shardRecordsIterator) {
- assert shardRecordsIterator != null;
- return shardRecordsIterator.getCheckpoint();
- }
- }));
- }
+ public KinesisReaderCheckpoint(Iterable<ShardCheckpoint> shardCheckpoints) {
+ this.shardCheckpoints = ImmutableList.copyOf(shardCheckpoints);
+ }
- /**
- * Splits given multi-shard checkpoint into partitions of approximately equal size.
- *
- * @param desiredNumSplits - upper limit for number of partitions to generate.
- * @return list of checkpoints covering consecutive partitions of current checkpoint.
- */
- public List<KinesisReaderCheckpoint> splitInto(int desiredNumSplits) {
- int partitionSize = divideAndRoundUp(shardCheckpoints.size(), desiredNumSplits);
+ public static KinesisReaderCheckpoint asCurrentStateOf(Iterable<ShardRecordsIterator>
+ iterators) {
+ return new KinesisReaderCheckpoint(transform(iterators,
+ new Function<ShardRecordsIterator, ShardCheckpoint>() {
+
+ @Nullable
+ @Override
+ public ShardCheckpoint apply(@Nullable
+ ShardRecordsIterator shardRecordsIterator) {
+ assert shardRecordsIterator != null;
+ return shardRecordsIterator.getCheckpoint();
+ }
+ }));
+ }
- List<KinesisReaderCheckpoint> checkpoints = newArrayList();
- for (List<ShardCheckpoint> shardPartition : partition(shardCheckpoints, partitionSize)) {
- checkpoints.add(new KinesisReaderCheckpoint(shardPartition));
+ /**
+ * Splits given multi-shard checkpoint into partitions of approximately equal size.
+ *
+ * @param desiredNumSplits - upper limit for number of partitions to generate.
+ * @return list of checkpoints covering consecutive partitions of current checkpoint.
+ */
+ public List<KinesisReaderCheckpoint> splitInto(int desiredNumSplits) {
+ int partitionSize = divideAndRoundUp(shardCheckpoints.size(), desiredNumSplits);
+
+ List<KinesisReaderCheckpoint> checkpoints = newArrayList();
+ for (List<ShardCheckpoint> shardPartition : partition(shardCheckpoints, partitionSize)) {
+ checkpoints.add(new KinesisReaderCheckpoint(shardPartition));
+ }
+ return checkpoints;
}
- return checkpoints;
- }
- private int divideAndRoundUp(int nominator, int denominator) {
- return (nominator + denominator - 1) / denominator;
- }
+ private int divideAndRoundUp(int nominator, int denominator) {
+ return (nominator + denominator - 1) / denominator;
+ }
- @Override
- public void finalizeCheckpoint() throws IOException {
+ @Override
+ public void finalizeCheckpoint() throws IOException {
- }
+ }
- @Override
- public String toString() {
- return shardCheckpoints.toString();
- }
+ @Override
+ public String toString() {
+ return shardCheckpoints.toString();
+ }
- @Override
- public Iterator<ShardCheckpoint> iterator() {
- return shardCheckpoints.iterator();
- }
+ @Override
+ public Iterator<ShardCheckpoint> iterator() {
+ return shardCheckpoints.iterator();
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisRecord.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisRecord.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisRecord.java
index 057b7bb..02b5370 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisRecord.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisRecord.java
@@ -22,9 +22,7 @@ import static org.apache.commons.lang.builder.HashCodeBuilder.reflectionHashCode
import com.amazonaws.services.kinesis.clientlibrary.types.ExtendedSequenceNumber;
import com.amazonaws.services.kinesis.clientlibrary.types.UserRecord;
import com.google.common.base.Charsets;
-
import java.nio.ByteBuffer;
-
import org.apache.commons.lang.builder.EqualsBuilder;
import org.joda.time.Instant;
@@ -32,92 +30,91 @@ import org.joda.time.Instant;
* {@link UserRecord} enhanced with utility methods.
*/
public class KinesisRecord {
-
- private Instant readTime;
- private String streamName;
- private String shardId;
- private long subSequenceNumber;
- private String sequenceNumber;
- private Instant approximateArrivalTimestamp;
- private ByteBuffer data;
- private String partitionKey;
-
- public KinesisRecord(UserRecord record, String streamName, String shardId) {
- this(record.getData(), record.getSequenceNumber(), record.getSubSequenceNumber(),
- record.getPartitionKey(),
- new Instant(record.getApproximateArrivalTimestamp()),
- Instant.now(),
- streamName, shardId);
- }
-
- public KinesisRecord(ByteBuffer data, String sequenceNumber, long subSequenceNumber,
- String partitionKey, Instant approximateArrivalTimestamp,
- Instant readTime,
- String streamName, String shardId) {
- this.data = data;
- this.sequenceNumber = sequenceNumber;
- this.subSequenceNumber = subSequenceNumber;
- this.partitionKey = partitionKey;
- this.approximateArrivalTimestamp = approximateArrivalTimestamp;
- this.readTime = readTime;
- this.streamName = streamName;
- this.shardId = shardId;
- }
-
- public ExtendedSequenceNumber getExtendedSequenceNumber() {
- return new ExtendedSequenceNumber(getSequenceNumber(), getSubSequenceNumber());
- }
-
- /***
- * @return unique id of the record based on its position in the stream
- */
- public byte[] getUniqueId() {
- return getExtendedSequenceNumber().toString().getBytes(Charsets.UTF_8);
- }
-
- public Instant getReadTime() {
- return readTime;
- }
-
- public String getStreamName() {
- return streamName;
- }
-
- public String getShardId() {
- return shardId;
- }
-
- public byte[] getDataAsBytes() {
- return getData().array();
- }
-
- @Override
- public boolean equals(Object obj) {
- return EqualsBuilder.reflectionEquals(this, obj);
- }
-
- @Override
- public int hashCode() {
- return reflectionHashCode(this);
- }
-
- public long getSubSequenceNumber() {
- return subSequenceNumber;
- }
-
- public String getSequenceNumber() {
- return sequenceNumber;
- }
-
- public Instant getApproximateArrivalTimestamp() {
- return approximateArrivalTimestamp;
- }
-
- public ByteBuffer getData() {
- return data;
- }
-
- public String getPartitionKey() {
- return partitionKey;
- }
+ private Instant readTime;
+ private String streamName;
+ private String shardId;
+ private long subSequenceNumber;
+ private String sequenceNumber;
+ private Instant approximateArrivalTimestamp;
+ private ByteBuffer data;
+ private String partitionKey;
+
+ public KinesisRecord(UserRecord record, String streamName, String shardId) {
+ this(record.getData(), record.getSequenceNumber(), record.getSubSequenceNumber(),
+ record.getPartitionKey(),
+ new Instant(record.getApproximateArrivalTimestamp()),
+ Instant.now(),
+ streamName, shardId);
+ }
+
+ public KinesisRecord(ByteBuffer data, String sequenceNumber, long subSequenceNumber,
+ String partitionKey, Instant approximateArrivalTimestamp,
+ Instant readTime,
+ String streamName, String shardId) {
+ this.data = data;
+ this.sequenceNumber = sequenceNumber;
+ this.subSequenceNumber = subSequenceNumber;
+ this.partitionKey = partitionKey;
+ this.approximateArrivalTimestamp = approximateArrivalTimestamp;
+ this.readTime = readTime;
+ this.streamName = streamName;
+ this.shardId = shardId;
+ }
+
+ public ExtendedSequenceNumber getExtendedSequenceNumber() {
+ return new ExtendedSequenceNumber(getSequenceNumber(), getSubSequenceNumber());
+ }
+
+ /***
+ * @return unique id of the record based on its position in the stream
+ */
+ public byte[] getUniqueId() {
+ return getExtendedSequenceNumber().toString().getBytes(Charsets.UTF_8);
+ }
+
+ public Instant getReadTime() {
+ return readTime;
+ }
+
+ public String getStreamName() {
+ return streamName;
+ }
+
+ public String getShardId() {
+ return shardId;
+ }
+
+ public byte[] getDataAsBytes() {
+ return getData().array();
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ return EqualsBuilder.reflectionEquals(this, obj);
+ }
+
+ @Override
+ public int hashCode() {
+ return reflectionHashCode(this);
+ }
+
+ public long getSubSequenceNumber() {
+ return subSequenceNumber;
+ }
+
+ public String getSequenceNumber() {
+ return sequenceNumber;
+ }
+
+ public Instant getApproximateArrivalTimestamp() {
+ return approximateArrivalTimestamp;
+ }
+
+ public ByteBuffer getData() {
+ return data;
+ }
+
+ public String getPartitionKey() {
+ return partitionKey;
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisRecordCoder.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisRecordCoder.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisRecordCoder.java
index dcf564d..f233e27 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisRecordCoder.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisRecordCoder.java
@@ -21,7 +21,6 @@ import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.ByteBuffer;
-
import org.apache.beam.sdk.coders.AtomicCoder;
import org.apache.beam.sdk.coders.ByteArrayCoder;
import org.apache.beam.sdk.coders.Coder;
@@ -34,41 +33,40 @@ import org.joda.time.Instant;
* A {@link Coder} for {@link KinesisRecord}.
*/
class KinesisRecordCoder extends AtomicCoder<KinesisRecord> {
+ private static final StringUtf8Coder STRING_CODER = StringUtf8Coder.of();
+ private static final ByteArrayCoder BYTE_ARRAY_CODER = ByteArrayCoder.of();
+ private static final InstantCoder INSTANT_CODER = InstantCoder.of();
+ private static final VarLongCoder VAR_LONG_CODER = VarLongCoder.of();
- private static final StringUtf8Coder STRING_CODER = StringUtf8Coder.of();
- private static final ByteArrayCoder BYTE_ARRAY_CODER = ByteArrayCoder.of();
- private static final InstantCoder INSTANT_CODER = InstantCoder.of();
- private static final VarLongCoder VAR_LONG_CODER = VarLongCoder.of();
-
- public static KinesisRecordCoder of() {
- return new KinesisRecordCoder();
- }
+ public static KinesisRecordCoder of() {
+ return new KinesisRecordCoder();
+ }
- @Override
- public void encode(KinesisRecord value, OutputStream outStream) throws
- IOException {
- BYTE_ARRAY_CODER.encode(value.getData().array(), outStream);
- STRING_CODER.encode(value.getSequenceNumber(), outStream);
- STRING_CODER.encode(value.getPartitionKey(), outStream);
- INSTANT_CODER.encode(value.getApproximateArrivalTimestamp(), outStream);
- VAR_LONG_CODER.encode(value.getSubSequenceNumber(), outStream);
- INSTANT_CODER.encode(value.getReadTime(), outStream);
- STRING_CODER.encode(value.getStreamName(), outStream);
- STRING_CODER.encode(value.getShardId(), outStream);
- }
+ @Override
+ public void encode(KinesisRecord value, OutputStream outStream) throws
+ IOException {
+ BYTE_ARRAY_CODER.encode(value.getData().array(), outStream);
+ STRING_CODER.encode(value.getSequenceNumber(), outStream);
+ STRING_CODER.encode(value.getPartitionKey(), outStream);
+ INSTANT_CODER.encode(value.getApproximateArrivalTimestamp(), outStream);
+ VAR_LONG_CODER.encode(value.getSubSequenceNumber(), outStream);
+ INSTANT_CODER.encode(value.getReadTime(), outStream);
+ STRING_CODER.encode(value.getStreamName(), outStream);
+ STRING_CODER.encode(value.getShardId(), outStream);
+ }
- @Override
- public KinesisRecord decode(InputStream inStream) throws IOException {
- ByteBuffer data = ByteBuffer.wrap(BYTE_ARRAY_CODER.decode(inStream));
- String sequenceNumber = STRING_CODER.decode(inStream);
- String partitionKey = STRING_CODER.decode(inStream);
- Instant approximateArrivalTimestamp = INSTANT_CODER.decode(inStream);
- long subSequenceNumber = VAR_LONG_CODER.decode(inStream);
- Instant readTimestamp = INSTANT_CODER.decode(inStream);
- String streamName = STRING_CODER.decode(inStream);
- String shardId = STRING_CODER.decode(inStream);
- return new KinesisRecord(data, sequenceNumber, subSequenceNumber, partitionKey,
- approximateArrivalTimestamp, readTimestamp, streamName, shardId
- );
- }
+ @Override
+ public KinesisRecord decode(InputStream inStream) throws IOException {
+ ByteBuffer data = ByteBuffer.wrap(BYTE_ARRAY_CODER.decode(inStream));
+ String sequenceNumber = STRING_CODER.decode(inStream);
+ String partitionKey = STRING_CODER.decode(inStream);
+ Instant approximateArrivalTimestamp = INSTANT_CODER.decode(inStream);
+ long subSequenceNumber = VAR_LONG_CODER.decode(inStream);
+ Instant readTimestamp = INSTANT_CODER.decode(inStream);
+ String streamName = STRING_CODER.decode(inStream);
+ String shardId = STRING_CODER.decode(inStream);
+ return new KinesisRecord(data, sequenceNumber, subSequenceNumber, partitionKey,
+ approximateArrivalTimestamp, readTimestamp, streamName, shardId
+ );
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisSource.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisSource.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisSource.java
index 362792b..7e67d07 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisSource.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/KinesisSource.java
@@ -21,7 +21,6 @@ import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.collect.Lists.newArrayList;
import java.util.List;
-
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.SerializableCoder;
import org.apache.beam.sdk.io.UnboundedSource;
@@ -29,85 +28,85 @@ import org.apache.beam.sdk.options.PipelineOptions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+
/**
* Represents source for single stream in Kinesis.
*/
class KinesisSource extends UnboundedSource<KinesisRecord, KinesisReaderCheckpoint> {
+ private static final Logger LOG = LoggerFactory.getLogger(KinesisSource.class);
+
+ private final KinesisClientProvider kinesis;
+ private CheckpointGenerator initialCheckpointGenerator;
- private static final Logger LOG = LoggerFactory.getLogger(KinesisSource.class);
-
- private final KinesisClientProvider kinesis;
- private CheckpointGenerator initialCheckpointGenerator;
-
- public KinesisSource(KinesisClientProvider kinesis, String streamName,
- StartingPoint startingPoint) {
- this(kinesis, new DynamicCheckpointGenerator(streamName, startingPoint));
- }
-
- private KinesisSource(KinesisClientProvider kinesisClientProvider,
- CheckpointGenerator initialCheckpoint) {
- this.kinesis = kinesisClientProvider;
- this.initialCheckpointGenerator = initialCheckpoint;
- validate();
- }
-
- /**
- * Generate splits for reading from the stream.
- * Basically, it'll try to evenly split set of shards in the stream into
- * {@code desiredNumSplits} partitions. Each partition is then a split.
- */
- @Override
- public List<KinesisSource> split(int desiredNumSplits,
- PipelineOptions options) throws Exception {
- KinesisReaderCheckpoint checkpoint =
- initialCheckpointGenerator.generate(SimplifiedKinesisClient.from(kinesis));
-
- List<KinesisSource> sources = newArrayList();
-
- for (KinesisReaderCheckpoint partition : checkpoint.splitInto(desiredNumSplits)) {
- sources.add(new KinesisSource(
- kinesis,
- new StaticCheckpointGenerator(partition)));
+ public KinesisSource(KinesisClientProvider kinesis, String streamName,
+ StartingPoint startingPoint) {
+ this(kinesis, new DynamicCheckpointGenerator(streamName, startingPoint));
}
- return sources;
- }
-
- /**
- * Creates reader based on given {@link KinesisReaderCheckpoint}.
- * If {@link KinesisReaderCheckpoint} is not given, then we use
- * {@code initialCheckpointGenerator} to generate new checkpoint.
- */
- @Override
- public UnboundedReader<KinesisRecord> createReader(PipelineOptions options,
- KinesisReaderCheckpoint checkpointMark) {
-
- CheckpointGenerator checkpointGenerator = initialCheckpointGenerator;
-
- if (checkpointMark != null) {
- checkpointGenerator = new StaticCheckpointGenerator(checkpointMark);
+
+ private KinesisSource(KinesisClientProvider kinesisClientProvider,
+ CheckpointGenerator initialCheckpoint) {
+ this.kinesis = kinesisClientProvider;
+ this.initialCheckpointGenerator = initialCheckpoint;
+ validate();
}
- LOG.info("Creating new reader using {}", checkpointGenerator);
-
- return new KinesisReader(
- SimplifiedKinesisClient.from(kinesis),
- checkpointGenerator,
- this);
- }
-
- @Override
- public Coder<KinesisReaderCheckpoint> getCheckpointMarkCoder() {
- return SerializableCoder.of(KinesisReaderCheckpoint.class);
- }
-
- @Override
- public void validate() {
- checkNotNull(kinesis);
- checkNotNull(initialCheckpointGenerator);
- }
-
- @Override
- public Coder<KinesisRecord> getDefaultOutputCoder() {
- return KinesisRecordCoder.of();
- }
+ /**
+ * Generate splits for reading from the stream.
+ * Basically, it'll try to evenly split set of shards in the stream into
+ * {@code desiredNumSplits} partitions. Each partition is then a split.
+ */
+ @Override
+ public List<KinesisSource> split(int desiredNumSplits,
+ PipelineOptions options) throws Exception {
+ KinesisReaderCheckpoint checkpoint =
+ initialCheckpointGenerator.generate(SimplifiedKinesisClient.from(kinesis));
+
+ List<KinesisSource> sources = newArrayList();
+
+ for (KinesisReaderCheckpoint partition : checkpoint.splitInto(desiredNumSplits)) {
+ sources.add(new KinesisSource(
+ kinesis,
+ new StaticCheckpointGenerator(partition)));
+ }
+ return sources;
+ }
+
+ /**
+ * Creates reader based on given {@link KinesisReaderCheckpoint}.
+ * If {@link KinesisReaderCheckpoint} is not given, then we use
+ * {@code initialCheckpointGenerator} to generate new checkpoint.
+ */
+ @Override
+ public UnboundedReader<KinesisRecord> createReader(PipelineOptions options,
+ KinesisReaderCheckpoint checkpointMark) {
+
+ CheckpointGenerator checkpointGenerator = initialCheckpointGenerator;
+
+ if (checkpointMark != null) {
+ checkpointGenerator = new StaticCheckpointGenerator(checkpointMark);
+ }
+
+ LOG.info("Creating new reader using {}", checkpointGenerator);
+
+ return new KinesisReader(
+ SimplifiedKinesisClient.from(kinesis),
+ checkpointGenerator,
+ this);
+ }
+
+ @Override
+ public Coder<KinesisReaderCheckpoint> getCheckpointMarkCoder() {
+ return SerializableCoder.of(KinesisReaderCheckpoint.class);
+ }
+
+ @Override
+ public void validate() {
+ checkNotNull(kinesis);
+ checkNotNull(initialCheckpointGenerator);
+ }
+
+ @Override
+ public Coder<KinesisRecord> getDefaultOutputCoder() {
+ return KinesisRecordCoder.of();
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/RecordFilter.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/RecordFilter.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/RecordFilter.java
index eca725c..40e65fc 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/RecordFilter.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/RecordFilter.java
@@ -21,6 +21,7 @@ import static com.google.common.collect.Lists.newArrayList;
import java.util.List;
+
/**
* Filters out records, which were already processed and checkpointed.
*
@@ -28,14 +29,13 @@ import java.util.List;
* accuracy, not with "subSequenceNumber" accuracy.
*/
class RecordFilter {
-
- public List<KinesisRecord> apply(List<KinesisRecord> records, ShardCheckpoint checkpoint) {
- List<KinesisRecord> filteredRecords = newArrayList();
- for (KinesisRecord record : records) {
- if (checkpoint.isBeforeOrAt(record)) {
- filteredRecords.add(record);
- }
+ public List<KinesisRecord> apply(List<KinesisRecord> records, ShardCheckpoint checkpoint) {
+ List<KinesisRecord> filteredRecords = newArrayList();
+ for (KinesisRecord record : records) {
+ if (checkpoint.isBeforeOrAt(record)) {
+ filteredRecords.add(record);
+ }
+ }
+ return filteredRecords;
}
- return filteredRecords;
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/RoundRobin.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/RoundRobin.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/RoundRobin.java
index 806d982..e4ff541 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/RoundRobin.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/RoundRobin.java
@@ -27,28 +27,27 @@ import java.util.Iterator;
* Very simple implementation of round robin algorithm.
*/
class RoundRobin<T> implements Iterable<T> {
+ private final Deque<T> deque;
- private final Deque<T> deque;
+ public RoundRobin(Iterable<T> collection) {
+ this.deque = newArrayDeque(collection);
+ checkArgument(!deque.isEmpty(), "Tried to initialize RoundRobin with empty collection");
+ }
- public RoundRobin(Iterable<T> collection) {
- this.deque = newArrayDeque(collection);
- checkArgument(!deque.isEmpty(), "Tried to initialize RoundRobin with empty collection");
- }
+ public T getCurrent() {
+ return deque.getFirst();
+ }
- public T getCurrent() {
- return deque.getFirst();
- }
+ public void moveForward() {
+ deque.addLast(deque.removeFirst());
+ }
- public void moveForward() {
- deque.addLast(deque.removeFirst());
- }
+ public int size() {
+ return deque.size();
+ }
- public int size() {
- return deque.size();
- }
-
- @Override
- public Iterator<T> iterator() {
- return deque.iterator();
- }
+ @Override
+ public Iterator<T> iterator() {
+ return deque.iterator();
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/ShardCheckpoint.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/ShardCheckpoint.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/ShardCheckpoint.java
index 95f97b8..6aa3504 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/ShardCheckpoint.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/ShardCheckpoint.java
@@ -17,6 +17,7 @@
*/
package org.apache.beam.sdk.io.kinesis;
+
import static com.amazonaws.services.kinesis.model.ShardIteratorType.AFTER_SEQUENCE_NUMBER;
import static com.amazonaws.services.kinesis.model.ShardIteratorType.AT_SEQUENCE_NUMBER;
import static com.amazonaws.services.kinesis.model.ShardIteratorType.AT_TIMESTAMP;
@@ -26,11 +27,10 @@ import static com.google.common.base.Preconditions.checkNotNull;
import com.amazonaws.services.kinesis.clientlibrary.types.ExtendedSequenceNumber;
import com.amazonaws.services.kinesis.model.Record;
import com.amazonaws.services.kinesis.model.ShardIteratorType;
-
import java.io.Serializable;
-
import org.joda.time.Instant;
+
/**
* Checkpoint mark for single shard in the stream.
* Current position in the shard is determined by either:
@@ -45,132 +45,131 @@ import org.joda.time.Instant;
* This class is immutable.
*/
class ShardCheckpoint implements Serializable {
+ private final String streamName;
+ private final String shardId;
+ private final String sequenceNumber;
+ private final ShardIteratorType shardIteratorType;
+ private final Long subSequenceNumber;
+ private final Instant timestamp;
+
+ public ShardCheckpoint(String streamName, String shardId, StartingPoint
+ startingPoint) {
+ this(streamName, shardId,
+ ShardIteratorType.fromValue(startingPoint.getPositionName()),
+ startingPoint.getTimestamp());
+ }
+
+ public ShardCheckpoint(String streamName, String shardId, ShardIteratorType
+ shardIteratorType, Instant timestamp) {
+ this(streamName, shardId, shardIteratorType, null, null, timestamp);
+ }
+
+ public ShardCheckpoint(String streamName, String shardId, ShardIteratorType
+ shardIteratorType, String sequenceNumber, Long subSequenceNumber) {
+ this(streamName, shardId, shardIteratorType, sequenceNumber, subSequenceNumber, null);
+ }
+
+ private ShardCheckpoint(String streamName, String shardId, ShardIteratorType shardIteratorType,
+ String sequenceNumber, Long subSequenceNumber, Instant timestamp) {
+ this.shardIteratorType = checkNotNull(shardIteratorType, "shardIteratorType");
+ this.streamName = checkNotNull(streamName, "streamName");
+ this.shardId = checkNotNull(shardId, "shardId");
+ if (shardIteratorType == AT_SEQUENCE_NUMBER || shardIteratorType == AFTER_SEQUENCE_NUMBER) {
+ checkNotNull(sequenceNumber,
+ "You must provide sequence number for AT_SEQUENCE_NUMBER"
+ + " or AFTER_SEQUENCE_NUMBER");
+ } else {
+ checkArgument(sequenceNumber == null,
+ "Sequence number must be null for LATEST, TRIM_HORIZON or AT_TIMESTAMP");
+ }
+ if (shardIteratorType == AT_TIMESTAMP) {
+ checkNotNull(timestamp,
+ "You must provide timestamp for AT_SEQUENCE_NUMBER"
+ + " or AFTER_SEQUENCE_NUMBER");
+ } else {
+ checkArgument(timestamp == null,
+ "Timestamp must be null for an iterator type other than AT_TIMESTAMP");
+ }
+
+ this.subSequenceNumber = subSequenceNumber;
+ this.sequenceNumber = sequenceNumber;
+ this.timestamp = timestamp;
+ }
+
+ /**
+ * Used to compare {@link ShardCheckpoint} object to {@link KinesisRecord}. Depending
+ * on the the underlying shardIteratorType, it will either compare the timestamp or the
+ * {@link ExtendedSequenceNumber}.
+ *
+ * @param other
+ * @return if current checkpoint mark points before or at given {@link ExtendedSequenceNumber}
+ */
+ public boolean isBeforeOrAt(KinesisRecord other) {
+ if (shardIteratorType == AT_TIMESTAMP) {
+ return timestamp.compareTo(other.getApproximateArrivalTimestamp()) <= 0;
+ }
+ int result = extendedSequenceNumber().compareTo(other.getExtendedSequenceNumber());
+ if (result == 0) {
+ return shardIteratorType == AT_SEQUENCE_NUMBER;
+ }
+ return result < 0;
+ }
+
+ private ExtendedSequenceNumber extendedSequenceNumber() {
+ String fullSequenceNumber = sequenceNumber;
+ if (fullSequenceNumber == null) {
+ fullSequenceNumber = shardIteratorType.toString();
+ }
+ return new ExtendedSequenceNumber(fullSequenceNumber, subSequenceNumber);
+ }
- private final String streamName;
- private final String shardId;
- private final String sequenceNumber;
- private final ShardIteratorType shardIteratorType;
- private final Long subSequenceNumber;
- private final Instant timestamp;
-
- public ShardCheckpoint(String streamName, String shardId, StartingPoint
- startingPoint) {
- this(streamName, shardId,
- ShardIteratorType.fromValue(startingPoint.getPositionName()),
- startingPoint.getTimestamp());
- }
-
- public ShardCheckpoint(String streamName, String shardId, ShardIteratorType
- shardIteratorType, Instant timestamp) {
- this(streamName, shardId, shardIteratorType, null, null, timestamp);
- }
-
- public ShardCheckpoint(String streamName, String shardId, ShardIteratorType
- shardIteratorType, String sequenceNumber, Long subSequenceNumber) {
- this(streamName, shardId, shardIteratorType, sequenceNumber, subSequenceNumber, null);
- }
-
- private ShardCheckpoint(String streamName, String shardId, ShardIteratorType shardIteratorType,
- String sequenceNumber, Long subSequenceNumber, Instant timestamp) {
- this.shardIteratorType = checkNotNull(shardIteratorType, "shardIteratorType");
- this.streamName = checkNotNull(streamName, "streamName");
- this.shardId = checkNotNull(shardId, "shardId");
- if (shardIteratorType == AT_SEQUENCE_NUMBER || shardIteratorType == AFTER_SEQUENCE_NUMBER) {
- checkNotNull(sequenceNumber,
- "You must provide sequence number for AT_SEQUENCE_NUMBER"
- + " or AFTER_SEQUENCE_NUMBER");
- } else {
- checkArgument(sequenceNumber == null,
- "Sequence number must be null for LATEST, TRIM_HORIZON or AT_TIMESTAMP");
+ @Override
+ public String toString() {
+ return String.format("Checkpoint %s for stream %s, shard %s: %s", shardIteratorType,
+ streamName, shardId,
+ sequenceNumber);
}
- if (shardIteratorType == AT_TIMESTAMP) {
- checkNotNull(timestamp,
- "You must provide timestamp for AT_SEQUENCE_NUMBER"
- + " or AFTER_SEQUENCE_NUMBER");
- } else {
- checkArgument(timestamp == null,
- "Timestamp must be null for an iterator type other than AT_TIMESTAMP");
+
+ public ShardRecordsIterator getShardRecordsIterator(SimplifiedKinesisClient kinesis)
+ throws TransientKinesisException {
+ return new ShardRecordsIterator(this, kinesis);
}
- this.subSequenceNumber = subSequenceNumber;
- this.sequenceNumber = sequenceNumber;
- this.timestamp = timestamp;
- }
-
- /**
- * Used to compare {@link ShardCheckpoint} object to {@link KinesisRecord}. Depending
- * on the the underlying shardIteratorType, it will either compare the timestamp or the
- * {@link ExtendedSequenceNumber}.
- *
- * @param other
- * @return if current checkpoint mark points before or at given {@link ExtendedSequenceNumber}
- */
- public boolean isBeforeOrAt(KinesisRecord other) {
- if (shardIteratorType == AT_TIMESTAMP) {
- return timestamp.compareTo(other.getApproximateArrivalTimestamp()) <= 0;
+ public String getShardIterator(SimplifiedKinesisClient kinesisClient)
+ throws TransientKinesisException {
+ if (checkpointIsInTheMiddleOfAUserRecord()) {
+ return kinesisClient.getShardIterator(streamName,
+ shardId, AT_SEQUENCE_NUMBER,
+ sequenceNumber, null);
+ }
+ return kinesisClient.getShardIterator(streamName,
+ shardId, shardIteratorType,
+ sequenceNumber, timestamp);
}
- int result = extendedSequenceNumber().compareTo(other.getExtendedSequenceNumber());
- if (result == 0) {
- return shardIteratorType == AT_SEQUENCE_NUMBER;
+
+ private boolean checkpointIsInTheMiddleOfAUserRecord() {
+ return shardIteratorType == AFTER_SEQUENCE_NUMBER && subSequenceNumber != null;
}
- return result < 0;
- }
- private ExtendedSequenceNumber extendedSequenceNumber() {
- String fullSequenceNumber = sequenceNumber;
- if (fullSequenceNumber == null) {
- fullSequenceNumber = shardIteratorType.toString();
+ /**
+ * Used to advance checkpoint mark to position after given {@link Record}.
+ *
+ * @param record
+ * @return new checkpoint object pointing directly after given {@link Record}
+ */
+ public ShardCheckpoint moveAfter(KinesisRecord record) {
+ return new ShardCheckpoint(
+ streamName, shardId,
+ AFTER_SEQUENCE_NUMBER,
+ record.getSequenceNumber(),
+ record.getSubSequenceNumber());
}
- return new ExtendedSequenceNumber(fullSequenceNumber, subSequenceNumber);
- }
-
- @Override
- public String toString() {
- return String.format("Checkpoint %s for stream %s, shard %s: %s", shardIteratorType,
- streamName, shardId,
- sequenceNumber);
- }
-
- public ShardRecordsIterator getShardRecordsIterator(SimplifiedKinesisClient kinesis)
- throws TransientKinesisException {
- return new ShardRecordsIterator(this, kinesis);
- }
-
- public String getShardIterator(SimplifiedKinesisClient kinesisClient)
- throws TransientKinesisException {
- if (checkpointIsInTheMiddleOfAUserRecord()) {
- return kinesisClient.getShardIterator(streamName,
- shardId, AT_SEQUENCE_NUMBER,
- sequenceNumber, null);
+
+ public String getStreamName() {
+ return streamName;
+ }
+
+ public String getShardId() {
+ return shardId;
}
- return kinesisClient.getShardIterator(streamName,
- shardId, shardIteratorType,
- sequenceNumber, timestamp);
- }
-
- private boolean checkpointIsInTheMiddleOfAUserRecord() {
- return shardIteratorType == AFTER_SEQUENCE_NUMBER && subSequenceNumber != null;
- }
-
- /**
- * Used to advance checkpoint mark to position after given {@link Record}.
- *
- * @param record
- * @return new checkpoint object pointing directly after given {@link Record}
- */
- public ShardCheckpoint moveAfter(KinesisRecord record) {
- return new ShardCheckpoint(
- streamName, shardId,
- AFTER_SEQUENCE_NUMBER,
- record.getSequenceNumber(),
- record.getSubSequenceNumber());
- }
-
- public String getStreamName() {
- return streamName;
- }
-
- public String getShardId() {
- return shardId;
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/ShardRecordsIterator.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/ShardRecordsIterator.java b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/ShardRecordsIterator.java
index a69c6c1..872f604 100644
--- a/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/ShardRecordsIterator.java
+++ b/sdks/java/io/kinesis/src/main/java/org/apache/beam/sdk/io/kinesis/ShardRecordsIterator.java
@@ -21,9 +21,7 @@ import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.collect.Queues.newArrayDeque;
import com.amazonaws.services.kinesis.model.ExpiredIteratorException;
-
import java.util.Deque;
-
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -33,68 +31,68 @@ import org.slf4j.LoggerFactory;
* Then the caller of {@link ShardRecordsIterator#next()} can read from queue one by one.
*/
class ShardRecordsIterator {
+ private static final Logger LOG = LoggerFactory.getLogger(ShardRecordsIterator.class);
- private static final Logger LOG = LoggerFactory.getLogger(ShardRecordsIterator.class);
+ private final SimplifiedKinesisClient kinesis;
+ private final RecordFilter filter;
+ private ShardCheckpoint checkpoint;
+ private String shardIterator;
+ private Deque<KinesisRecord> data = newArrayDeque();
- private final SimplifiedKinesisClient kinesis;
- private final RecordFilter filter;
- private ShardCheckpoint checkpoint;
- private String shardIterator;
- private Deque<KinesisRecord> data = newArrayDeque();
+ public ShardRecordsIterator(final ShardCheckpoint initialCheckpoint,
+ SimplifiedKinesisClient simplifiedKinesisClient) throws
+ TransientKinesisException {
+ this(initialCheckpoint, simplifiedKinesisClient, new RecordFilter());
+ }
- public ShardRecordsIterator(final ShardCheckpoint initialCheckpoint,
- SimplifiedKinesisClient simplifiedKinesisClient) throws
- TransientKinesisException {
- this(initialCheckpoint, simplifiedKinesisClient, new RecordFilter());
- }
+ public ShardRecordsIterator(final ShardCheckpoint initialCheckpoint,
+ SimplifiedKinesisClient simplifiedKinesisClient,
+ RecordFilter filter) throws
+ TransientKinesisException {
- public ShardRecordsIterator(final ShardCheckpoint initialCheckpoint,
- SimplifiedKinesisClient simplifiedKinesisClient,
- RecordFilter filter) throws
- TransientKinesisException {
+ this.checkpoint = checkNotNull(initialCheckpoint, "initialCheckpoint");
+ this.filter = checkNotNull(filter, "filter");
+ this.kinesis = checkNotNull(simplifiedKinesisClient, "simplifiedKinesisClient");
+ shardIterator = checkpoint.getShardIterator(kinesis);
+ }
- this.checkpoint = checkNotNull(initialCheckpoint, "initialCheckpoint");
- this.filter = checkNotNull(filter, "filter");
- this.kinesis = checkNotNull(simplifiedKinesisClient, "simplifiedKinesisClient");
- shardIterator = checkpoint.getShardIterator(kinesis);
- }
+ /**
+ * Returns record if there's any present.
+ * Returns absent() if there are no new records at this time in the shard.
+ */
+ public CustomOptional<KinesisRecord> next() throws TransientKinesisException {
+ readMoreIfNecessary();
- /**
- * Returns record if there's any present.
- * Returns absent() if there are no new records at this time in the shard.
- */
- public CustomOptional<KinesisRecord> next() throws TransientKinesisException {
- readMoreIfNecessary();
+ if (data.isEmpty()) {
+ return CustomOptional.absent();
+ } else {
+ KinesisRecord record = data.removeFirst();
+ checkpoint = checkpoint.moveAfter(record);
+ return CustomOptional.of(record);
+ }
+ }
- if (data.isEmpty()) {
- return CustomOptional.absent();
- } else {
- KinesisRecord record = data.removeFirst();
- checkpoint = checkpoint.moveAfter(record);
- return CustomOptional.of(record);
+ private void readMoreIfNecessary() throws TransientKinesisException {
+ if (data.isEmpty()) {
+ GetKinesisRecordsResult response;
+ try {
+ response = kinesis.getRecords(shardIterator, checkpoint.getStreamName(),
+ checkpoint.getShardId());
+ } catch (ExpiredIteratorException e) {
+ LOG.info("Refreshing expired iterator", e);
+ shardIterator = checkpoint.getShardIterator(kinesis);
+ response = kinesis.getRecords(shardIterator, checkpoint.getStreamName(),
+ checkpoint.getShardId());
+ }
+ LOG.debug("Fetched {} new records", response.getRecords().size());
+ shardIterator = response.getNextShardIterator();
+ data.addAll(filter.apply(response.getRecords(), checkpoint));
+ }
}
- }
- private void readMoreIfNecessary() throws TransientKinesisException {
- if (data.isEmpty()) {
- GetKinesisRecordsResult response;
- try {
- response = kinesis.getRecords(shardIterator, checkpoint.getStreamName(),
- checkpoint.getShardId());
- } catch (ExpiredIteratorException e) {
- LOG.info("Refreshing expired iterator", e);
- shardIterator = checkpoint.getShardIterator(kinesis);
- response = kinesis.getRecords(shardIterator, checkpoint.getStreamName(),
- checkpoint.getShardId());
- }
- LOG.debug("Fetched {} new records", response.getRecords().size());
- shardIterator = response.getNextShardIterator();
- data.addAll(filter.apply(response.getRecords(), checkpoint));
+ public ShardCheckpoint getCheckpoint() {
+ return checkpoint;
}
- }
- public ShardCheckpoint getCheckpoint() {
- return checkpoint;
- }
}
[21/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/state/FlinkStateInternals.java
----------------------------------------------------------------------
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/state/FlinkStateInternals.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/state/FlinkStateInternals.java
index a0b015b..f0d3278 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/state/FlinkStateInternals.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/state/FlinkStateInternals.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.runners.flink.translation.wrappers.streaming.state;
-import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import java.nio.ByteBuffer;
import java.util.Collections;
@@ -26,6 +25,7 @@ import java.util.Map;
import org.apache.beam.runners.core.StateInternals;
import org.apache.beam.runners.core.StateNamespace;
import org.apache.beam.runners.core.StateTag;
+import org.apache.beam.runners.flink.translation.types.CoderTypeInformation;
import org.apache.beam.runners.flink.translation.types.CoderTypeSerializer;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderException;
@@ -34,7 +34,6 @@ import org.apache.beam.sdk.state.BagState;
import org.apache.beam.sdk.state.CombiningState;
import org.apache.beam.sdk.state.MapState;
import org.apache.beam.sdk.state.ReadableState;
-import org.apache.beam.sdk.state.ReadableStates;
import org.apache.beam.sdk.state.SetState;
import org.apache.beam.sdk.state.State;
import org.apache.beam.sdk.state.StateContext;
@@ -50,7 +49,6 @@ import org.apache.beam.sdk.util.CombineContextFactory;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ValueStateDescriptor;
-import org.apache.flink.api.common.typeutils.base.BooleanSerializer;
import org.apache.flink.api.common.typeutils.base.StringSerializer;
import org.apache.flink.runtime.state.KeyedStateBackend;
import org.joda.time.Instant;
@@ -130,8 +128,8 @@ public class FlinkStateInternals<K> implements StateInternals {
@Override
public <T> SetState<T> bindSet(
StateTag<SetState<T>> address, Coder<T> elemCoder) {
- return new FlinkSetState<>(
- flinkStateBackend, address, namespace, elemCoder);
+ throw new UnsupportedOperationException(
+ String.format("%s is not supported", SetState.class.getSimpleName()));
}
@Override
@@ -198,8 +196,9 @@ public class FlinkStateInternals<K> implements StateInternals {
this.address = address;
this.flinkStateBackend = flinkStateBackend;
- flinkStateDescriptor = new ValueStateDescriptor<>(
- address.getId(), new CoderTypeSerializer<>(coder));
+ CoderTypeInformation<T> typeInfo = new CoderTypeInformation<>(coder);
+
+ flinkStateDescriptor = new ValueStateDescriptor<>(address.getId(), typeInfo, null);
}
@Override
@@ -283,8 +282,9 @@ public class FlinkStateInternals<K> implements StateInternals {
this.address = address;
this.flinkStateBackend = flinkStateBackend;
- flinkStateDescriptor = new ListStateDescriptor<>(
- address.getId(), new CoderTypeSerializer<>(coder));
+ CoderTypeInformation<T> typeInfo = new CoderTypeInformation<>(coder);
+
+ flinkStateDescriptor = new ListStateDescriptor<>(address.getId(), typeInfo);
}
@Override
@@ -398,8 +398,9 @@ public class FlinkStateInternals<K> implements StateInternals {
this.combineFn = combineFn;
this.flinkStateBackend = flinkStateBackend;
- flinkStateDescriptor = new ValueStateDescriptor<>(
- address.getId(), new CoderTypeSerializer<>(accumCoder));
+ CoderTypeInformation<AccumT> typeInfo = new CoderTypeInformation<>(accumCoder);
+
+ flinkStateDescriptor = new ValueStateDescriptor<>(address.getId(), typeInfo, null);
}
@Override
@@ -544,6 +545,179 @@ public class FlinkStateInternals<K> implements StateInternals {
}
}
+ private static class FlinkKeyedCombiningState<K, InputT, AccumT, OutputT>
+ implements CombiningState<InputT, AccumT, OutputT> {
+
+ private final StateNamespace namespace;
+ private final StateTag<CombiningState<InputT, AccumT, OutputT>> address;
+ private final Combine.CombineFn<InputT, AccumT, OutputT> combineFn;
+ private final ValueStateDescriptor<AccumT> flinkStateDescriptor;
+ private final KeyedStateBackend<ByteBuffer> flinkStateBackend;
+ private final FlinkStateInternals<K> flinkStateInternals;
+
+ FlinkKeyedCombiningState(
+ KeyedStateBackend<ByteBuffer> flinkStateBackend,
+ StateTag<CombiningState<InputT, AccumT, OutputT>> address,
+ Combine.CombineFn<InputT, AccumT, OutputT> combineFn,
+ StateNamespace namespace,
+ Coder<AccumT> accumCoder,
+ FlinkStateInternals<K> flinkStateInternals) {
+
+ this.namespace = namespace;
+ this.address = address;
+ this.combineFn = combineFn;
+ this.flinkStateBackend = flinkStateBackend;
+ this.flinkStateInternals = flinkStateInternals;
+
+ CoderTypeInformation<AccumT> typeInfo = new CoderTypeInformation<>(accumCoder);
+
+ flinkStateDescriptor = new ValueStateDescriptor<>(address.getId(), typeInfo, null);
+ }
+
+ @Override
+ public CombiningState<InputT, AccumT, OutputT> readLater() {
+ return this;
+ }
+
+ @Override
+ public void add(InputT value) {
+ try {
+ org.apache.flink.api.common.state.ValueState<AccumT> state =
+ flinkStateBackend.getPartitionedState(
+ namespace.stringKey(),
+ StringSerializer.INSTANCE,
+ flinkStateDescriptor);
+
+ AccumT current = state.value();
+ if (current == null) {
+ current = combineFn.createAccumulator();
+ }
+ current = combineFn.addInput(current, value);
+ state.update(current);
+ } catch (RuntimeException re) {
+ throw re;
+ } catch (Exception e) {
+ throw new RuntimeException("Error adding to state." , e);
+ }
+ }
+
+ @Override
+ public void addAccum(AccumT accum) {
+ try {
+ org.apache.flink.api.common.state.ValueState<AccumT> state =
+ flinkStateBackend.getPartitionedState(
+ namespace.stringKey(),
+ StringSerializer.INSTANCE,
+ flinkStateDescriptor);
+
+ AccumT current = state.value();
+ if (current == null) {
+ state.update(accum);
+ } else {
+ current = combineFn.mergeAccumulators(Lists.newArrayList(current, accum));
+ state.update(current);
+ }
+ } catch (Exception e) {
+ throw new RuntimeException("Error adding to state.", e);
+ }
+ }
+
+ @Override
+ public AccumT getAccum() {
+ try {
+ return flinkStateBackend.getPartitionedState(
+ namespace.stringKey(),
+ StringSerializer.INSTANCE,
+ flinkStateDescriptor).value();
+ } catch (Exception e) {
+ throw new RuntimeException("Error reading state.", e);
+ }
+ }
+
+ @Override
+ public AccumT mergeAccumulators(Iterable<AccumT> accumulators) {
+ return combineFn.mergeAccumulators(accumulators);
+ }
+
+ @Override
+ public OutputT read() {
+ try {
+ org.apache.flink.api.common.state.ValueState<AccumT> state =
+ flinkStateBackend.getPartitionedState(
+ namespace.stringKey(),
+ StringSerializer.INSTANCE,
+ flinkStateDescriptor);
+
+ AccumT accum = state.value();
+ if (accum != null) {
+ return combineFn.extractOutput(accum);
+ } else {
+ return combineFn.extractOutput(combineFn.createAccumulator());
+ }
+ } catch (Exception e) {
+ throw new RuntimeException("Error reading state.", e);
+ }
+ }
+
+ @Override
+ public ReadableState<Boolean> isEmpty() {
+ return new ReadableState<Boolean>() {
+ @Override
+ public Boolean read() {
+ try {
+ return flinkStateBackend.getPartitionedState(
+ namespace.stringKey(),
+ StringSerializer.INSTANCE,
+ flinkStateDescriptor).value() == null;
+ } catch (Exception e) {
+ throw new RuntimeException("Error reading state.", e);
+ }
+
+ }
+
+ @Override
+ public ReadableState<Boolean> readLater() {
+ return this;
+ }
+ };
+ }
+
+ @Override
+ public void clear() {
+ try {
+ flinkStateBackend.getPartitionedState(
+ namespace.stringKey(),
+ StringSerializer.INSTANCE,
+ flinkStateDescriptor).clear();
+ } catch (Exception e) {
+ throw new RuntimeException("Error clearing state.", e);
+ }
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+
+ FlinkKeyedCombiningState<?, ?, ?, ?> that =
+ (FlinkKeyedCombiningState<?, ?, ?, ?>) o;
+
+ return namespace.equals(that.namespace) && address.equals(that.address);
+
+ }
+
+ @Override
+ public int hashCode() {
+ int result = namespace.hashCode();
+ result = 31 * result + address.hashCode();
+ return result;
+ }
+ }
+
private static class FlinkCombiningStateWithContext<K, InputT, AccumT, OutputT>
implements CombiningState<InputT, AccumT, OutputT> {
@@ -571,8 +745,9 @@ public class FlinkStateInternals<K> implements StateInternals {
this.flinkStateInternals = flinkStateInternals;
this.context = context;
- flinkStateDescriptor = new ValueStateDescriptor<>(
- address.getId(), new CoderTypeSerializer<>(accumCoder));
+ CoderTypeInformation<AccumT> typeInfo = new CoderTypeInformation<>(accumCoder);
+
+ flinkStateDescriptor = new ValueStateDescriptor<>(address.getId(), typeInfo, null);
}
@Override
@@ -738,8 +913,8 @@ public class FlinkStateInternals<K> implements StateInternals {
this.flinkStateBackend = flinkStateBackend;
this.flinkStateInternals = flinkStateInternals;
- flinkStateDescriptor = new ValueStateDescriptor<>(
- address.getId(), new CoderTypeSerializer<>(InstantCoder.of()));
+ CoderTypeInformation<Instant> typeInfo = new CoderTypeInformation<>(InstantCoder.of());
+ flinkStateDescriptor = new ValueStateDescriptor<>(address.getId(), typeInfo, null);
}
@Override
@@ -878,15 +1053,24 @@ public class FlinkStateInternals<K> implements StateInternals {
@Override
public ReadableState<ValueT> get(final KeyT input) {
- try {
- return ReadableStates.immediate(
- flinkStateBackend.getPartitionedState(
+ return new ReadableState<ValueT>() {
+ @Override
+ public ValueT read() {
+ try {
+ return flinkStateBackend.getPartitionedState(
namespace.stringKey(),
StringSerializer.INSTANCE,
- flinkStateDescriptor).get(input));
- } catch (Exception e) {
- throw new RuntimeException("Error get from state.", e);
- }
+ flinkStateDescriptor).get(input);
+ } catch (Exception e) {
+ throw new RuntimeException("Error get from state.", e);
+ }
+ }
+
+ @Override
+ public ReadableState<ValueT> readLater() {
+ return this;
+ }
+ };
}
@Override
@@ -903,22 +1087,32 @@ public class FlinkStateInternals<K> implements StateInternals {
@Override
public ReadableState<ValueT> putIfAbsent(final KeyT key, final ValueT value) {
- try {
- ValueT current = flinkStateBackend.getPartitionedState(
- namespace.stringKey(),
- StringSerializer.INSTANCE,
- flinkStateDescriptor).get(key);
+ return new ReadableState<ValueT>() {
+ @Override
+ public ValueT read() {
+ try {
+ ValueT current = flinkStateBackend.getPartitionedState(
+ namespace.stringKey(),
+ StringSerializer.INSTANCE,
+ flinkStateDescriptor).get(key);
+
+ if (current == null) {
+ flinkStateBackend.getPartitionedState(
+ namespace.stringKey(),
+ StringSerializer.INSTANCE,
+ flinkStateDescriptor).put(key, value);
+ }
+ return current;
+ } catch (Exception e) {
+ throw new RuntimeException("Error put kv to state.", e);
+ }
+ }
- if (current == null) {
- flinkStateBackend.getPartitionedState(
- namespace.stringKey(),
- StringSerializer.INSTANCE,
- flinkStateDescriptor).put(key, value);
+ @Override
+ public ReadableState<ValueT> readLater() {
+ return this;
}
- return ReadableStates.immediate(current);
- } catch (Exception e) {
- throw new RuntimeException("Error put kv to state.", e);
- }
+ };
}
@Override
@@ -939,11 +1133,10 @@ public class FlinkStateInternals<K> implements StateInternals {
@Override
public Iterable<KeyT> read() {
try {
- Iterable<KeyT> result = flinkStateBackend.getPartitionedState(
+ return flinkStateBackend.getPartitionedState(
namespace.stringKey(),
StringSerializer.INSTANCE,
flinkStateDescriptor).keys();
- return result != null ? result : Collections.<KeyT>emptyList();
} catch (Exception e) {
throw new RuntimeException("Error get map state keys.", e);
}
@@ -962,11 +1155,10 @@ public class FlinkStateInternals<K> implements StateInternals {
@Override
public Iterable<ValueT> read() {
try {
- Iterable<ValueT> result = flinkStateBackend.getPartitionedState(
+ return flinkStateBackend.getPartitionedState(
namespace.stringKey(),
StringSerializer.INSTANCE,
flinkStateDescriptor).values();
- return result != null ? result : Collections.<ValueT>emptyList();
} catch (Exception e) {
throw new RuntimeException("Error get map state values.", e);
}
@@ -985,11 +1177,10 @@ public class FlinkStateInternals<K> implements StateInternals {
@Override
public Iterable<Map.Entry<KeyT, ValueT>> read() {
try {
- Iterable<Map.Entry<KeyT, ValueT>> result = flinkStateBackend.getPartitionedState(
+ return flinkStateBackend.getPartitionedState(
namespace.stringKey(),
StringSerializer.INSTANCE,
flinkStateDescriptor).entries();
- return result != null ? result : Collections.<Map.Entry<KeyT, ValueT>>emptyList();
} catch (Exception e) {
throw new RuntimeException("Error get map state entries.", e);
}
@@ -1037,154 +1228,4 @@ public class FlinkStateInternals<K> implements StateInternals {
}
}
- private static class FlinkSetState<T> implements SetState<T> {
-
- private final StateNamespace namespace;
- private final StateTag<SetState<T>> address;
- private final MapStateDescriptor<T, Boolean> flinkStateDescriptor;
- private final KeyedStateBackend<ByteBuffer> flinkStateBackend;
-
- FlinkSetState(
- KeyedStateBackend<ByteBuffer> flinkStateBackend,
- StateTag<SetState<T>> address,
- StateNamespace namespace,
- Coder<T> coder) {
- this.namespace = namespace;
- this.address = address;
- this.flinkStateBackend = flinkStateBackend;
- this.flinkStateDescriptor = new MapStateDescriptor<>(address.getId(),
- new CoderTypeSerializer<>(coder), new BooleanSerializer());
- }
-
- @Override
- public ReadableState<Boolean> contains(final T t) {
- try {
- Boolean result = flinkStateBackend.getPartitionedState(
- namespace.stringKey(),
- StringSerializer.INSTANCE,
- flinkStateDescriptor).get(t);
- return ReadableStates.immediate(result != null ? result : false);
- } catch (Exception e) {
- throw new RuntimeException("Error contains value from state.", e);
- }
- }
-
- @Override
- public ReadableState<Boolean> addIfAbsent(final T t) {
- try {
- org.apache.flink.api.common.state.MapState<T, Boolean> state =
- flinkStateBackend.getPartitionedState(
- namespace.stringKey(),
- StringSerializer.INSTANCE,
- flinkStateDescriptor);
- boolean alreadyContained = state.contains(t);
- if (!alreadyContained) {
- state.put(t, true);
- }
- return ReadableStates.immediate(!alreadyContained);
- } catch (Exception e) {
- throw new RuntimeException("Error addIfAbsent value to state.", e);
- }
- }
-
- @Override
- public void remove(T t) {
- try {
- flinkStateBackend.getPartitionedState(
- namespace.stringKey(),
- StringSerializer.INSTANCE,
- flinkStateDescriptor).remove(t);
- } catch (Exception e) {
- throw new RuntimeException("Error remove value to state.", e);
- }
- }
-
- @Override
- public SetState<T> readLater() {
- return this;
- }
-
- @Override
- public void add(T value) {
- try {
- flinkStateBackend.getPartitionedState(
- namespace.stringKey(),
- StringSerializer.INSTANCE,
- flinkStateDescriptor).put(value, true);
- } catch (Exception e) {
- throw new RuntimeException("Error add value to state.", e);
- }
- }
-
- @Override
- public ReadableState<Boolean> isEmpty() {
- return new ReadableState<Boolean>() {
- @Override
- public Boolean read() {
- try {
- Iterable<T> result = flinkStateBackend.getPartitionedState(
- namespace.stringKey(),
- StringSerializer.INSTANCE,
- flinkStateDescriptor).keys();
- return result == null || Iterables.isEmpty(result);
- } catch (Exception e) {
- throw new RuntimeException("Error isEmpty from state.", e);
- }
- }
-
- @Override
- public ReadableState<Boolean> readLater() {
- return this;
- }
- };
- }
-
- @Override
- public Iterable<T> read() {
- try {
- Iterable<T> result = flinkStateBackend.getPartitionedState(
- namespace.stringKey(),
- StringSerializer.INSTANCE,
- flinkStateDescriptor).keys();
- return result != null ? result : Collections.<T>emptyList();
- } catch (Exception e) {
- throw new RuntimeException("Error read from state.", e);
- }
- }
-
- @Override
- public void clear() {
- try {
- flinkStateBackend.getPartitionedState(
- namespace.stringKey(),
- StringSerializer.INSTANCE,
- flinkStateDescriptor).clear();
- } catch (Exception e) {
- throw new RuntimeException("Error clearing state.", e);
- }
- }
-
- @Override
- public boolean equals(Object o) {
- if (this == o) {
- return true;
- }
- if (o == null || getClass() != o.getClass()) {
- return false;
- }
-
- FlinkSetState<?> that = (FlinkSetState<?>) o;
-
- return namespace.equals(that.namespace) && address.equals(that.address);
-
- }
-
- @Override
- public int hashCode() {
- int result = namespace.hashCode();
- result = 31 * result + address.hashCode();
- return result;
- }
- }
-
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkBroadcastStateInternalsTest.java
----------------------------------------------------------------------
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkBroadcastStateInternalsTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkBroadcastStateInternalsTest.java
index 3409d27..2b96d91 100644
--- a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkBroadcastStateInternalsTest.java
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkBroadcastStateInternalsTest.java
@@ -17,87 +17,229 @@
*/
package org.apache.beam.runners.flink.streaming;
-import org.apache.beam.runners.core.StateInternals;
-import org.apache.beam.runners.core.StateInternalsTest;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertThat;
+
+import java.util.Arrays;
+import org.apache.beam.runners.core.StateMerging;
+import org.apache.beam.runners.core.StateNamespace;
+import org.apache.beam.runners.core.StateNamespaceForTest;
+import org.apache.beam.runners.core.StateTag;
+import org.apache.beam.runners.core.StateTags;
import org.apache.beam.runners.flink.translation.wrappers.streaming.state.FlinkBroadcastStateInternals;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.CombiningState;
+import org.apache.beam.sdk.state.GroupingState;
+import org.apache.beam.sdk.state.ReadableState;
+import org.apache.beam.sdk.state.ValueState;
+import org.apache.beam.sdk.transforms.Sum;
import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
import org.apache.flink.runtime.state.OperatorStateBackend;
import org.apache.flink.runtime.state.memory.MemoryStateBackend;
-import org.junit.Ignore;
+import org.hamcrest.Matchers;
+import org.junit.Before;
+import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
/**
* Tests for {@link FlinkBroadcastStateInternals}. This is based on the tests for
- * {@code StateInternalsTest}.
- *
- * <p>Just test value, bag and combining.
+ * {@code InMemoryStateInternals}.
*/
@RunWith(JUnit4.class)
-public class FlinkBroadcastStateInternalsTest extends StateInternalsTest {
-
- @Override
- protected StateInternals createStateInternals() {
+public class FlinkBroadcastStateInternalsTest {
+ private static final StateNamespace NAMESPACE_1 = new StateNamespaceForTest("ns1");
+ private static final StateNamespace NAMESPACE_2 = new StateNamespaceForTest("ns2");
+ private static final StateNamespace NAMESPACE_3 = new StateNamespaceForTest("ns3");
+
+ private static final StateTag<ValueState<String>> STRING_VALUE_ADDR =
+ StateTags.value("stringValue", StringUtf8Coder.of());
+ private static final StateTag<CombiningState<Integer, int[], Integer>>
+ SUM_INTEGER_ADDR = StateTags.combiningValueFromInputInternal(
+ "sumInteger", VarIntCoder.of(), Sum.ofIntegers());
+ private static final StateTag<BagState<String>> STRING_BAG_ADDR =
+ StateTags.bag("stringBag", StringUtf8Coder.of());
+
+ FlinkBroadcastStateInternals<String> underTest;
+
+ @Before
+ public void initStateInternals() {
MemoryStateBackend backend = new MemoryStateBackend();
try {
OperatorStateBackend operatorStateBackend =
backend.createOperatorStateBackend(new DummyEnvironment("test", 1, 0), "");
- return new FlinkBroadcastStateInternals<>(1, operatorStateBackend);
+ underTest = new FlinkBroadcastStateInternals<>(1, operatorStateBackend);
+
} catch (Exception e) {
throw new RuntimeException(e);
}
}
- @Override
- @Ignore
- public void testSet() {}
+ @Test
+ public void testValue() throws Exception {
+ ValueState<String> value = underTest.state(NAMESPACE_1, STRING_VALUE_ADDR);
+
+ assertEquals(underTest.state(NAMESPACE_1, STRING_VALUE_ADDR), value);
+ assertNotEquals(
+ underTest.state(NAMESPACE_2, STRING_VALUE_ADDR),
+ value);
+
+ assertThat(value.read(), Matchers.nullValue());
+ value.write("hello");
+ assertThat(value.read(), Matchers.equalTo("hello"));
+ value.write("world");
+ assertThat(value.read(), Matchers.equalTo("world"));
+
+ value.clear();
+ assertThat(value.read(), Matchers.nullValue());
+ assertEquals(underTest.state(NAMESPACE_1, STRING_VALUE_ADDR), value);
+
+ }
+
+ @Test
+ public void testBag() throws Exception {
+ BagState<String> value = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
- @Override
- @Ignore
- public void testSetIsEmpty() {}
+ assertEquals(value, underTest.state(NAMESPACE_1, STRING_BAG_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, STRING_BAG_ADDR)));
- @Override
- @Ignore
- public void testMergeSetIntoSource() {}
+ assertThat(value.read(), Matchers.emptyIterable());
+ value.add("hello");
+ assertThat(value.read(), Matchers.containsInAnyOrder("hello"));
- @Override
- @Ignore
- public void testMergeSetIntoNewNamespace() {}
+ value.add("world");
+ assertThat(value.read(), Matchers.containsInAnyOrder("hello", "world"));
- @Override
- @Ignore
- public void testMap() {}
+ value.clear();
+ assertThat(value.read(), Matchers.emptyIterable());
+ assertEquals(underTest.state(NAMESPACE_1, STRING_BAG_ADDR), value);
- @Override
- @Ignore
- public void testWatermarkEarliestState() {}
+ }
+
+ @Test
+ public void testBagIsEmpty() throws Exception {
+ BagState<String> value = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
- @Override
- @Ignore
- public void testWatermarkLatestState() {}
+ assertThat(value.isEmpty().read(), Matchers.is(true));
+ ReadableState<Boolean> readFuture = value.isEmpty();
+ value.add("hello");
+ assertThat(readFuture.read(), Matchers.is(false));
- @Override
- @Ignore
- public void testWatermarkEndOfWindowState() {}
+ value.clear();
+ assertThat(readFuture.read(), Matchers.is(true));
+ }
- @Override
- @Ignore
- public void testWatermarkStateIsEmpty() {}
+ @Test
+ public void testMergeBagIntoSource() throws Exception {
+ BagState<String> bag1 = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
+ BagState<String> bag2 = underTest.state(NAMESPACE_2, STRING_BAG_ADDR);
- @Override
- @Ignore
- public void testMergeEarliestWatermarkIntoSource() {}
+ bag1.add("Hello");
+ bag2.add("World");
+ bag1.add("!");
- @Override
- @Ignore
- public void testMergeLatestWatermarkIntoSource() {}
+ StateMerging.mergeBags(Arrays.asList(bag1, bag2), bag1);
+
+ // Reading the merged bag gets both the contents
+ assertThat(bag1.read(), Matchers.containsInAnyOrder("Hello", "World", "!"));
+ assertThat(bag2.read(), Matchers.emptyIterable());
+ }
- @Override
- @Ignore
- public void testSetReadable() {}
+ @Test
+ public void testMergeBagIntoNewNamespace() throws Exception {
+ BagState<String> bag1 = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
+ BagState<String> bag2 = underTest.state(NAMESPACE_2, STRING_BAG_ADDR);
+ BagState<String> bag3 = underTest.state(NAMESPACE_3, STRING_BAG_ADDR);
+
+ bag1.add("Hello");
+ bag2.add("World");
+ bag1.add("!");
+
+ StateMerging.mergeBags(Arrays.asList(bag1, bag2, bag3), bag3);
+
+ // Reading the merged bag gets both the contents
+ assertThat(bag3.read(), Matchers.containsInAnyOrder("Hello", "World", "!"));
+ assertThat(bag1.read(), Matchers.emptyIterable());
+ assertThat(bag2.read(), Matchers.emptyIterable());
+ }
- @Override
- @Ignore
- public void testMapReadable() {}
+ @Test
+ public void testCombiningValue() throws Exception {
+ GroupingState<Integer, Integer> value = underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
+
+ // State instances are cached, but depend on the namespace.
+ assertEquals(value, underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, SUM_INTEGER_ADDR)));
+
+ assertThat(value.read(), Matchers.equalTo(0));
+ value.add(2);
+ assertThat(value.read(), Matchers.equalTo(2));
+
+ value.add(3);
+ assertThat(value.read(), Matchers.equalTo(5));
+
+ value.clear();
+ assertThat(value.read(), Matchers.equalTo(0));
+ assertEquals(underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR), value);
+ }
+
+ @Test
+ public void testCombiningIsEmpty() throws Exception {
+ GroupingState<Integer, Integer> value = underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
+
+ assertThat(value.isEmpty().read(), Matchers.is(true));
+ ReadableState<Boolean> readFuture = value.isEmpty();
+ value.add(5);
+ assertThat(readFuture.read(), Matchers.is(false));
+
+ value.clear();
+ assertThat(readFuture.read(), Matchers.is(true));
+ }
+
+ @Test
+ public void testMergeCombiningValueIntoSource() throws Exception {
+ CombiningState<Integer, int[], Integer> value1 =
+ underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
+ CombiningState<Integer, int[], Integer> value2 =
+ underTest.state(NAMESPACE_2, SUM_INTEGER_ADDR);
+
+ value1.add(5);
+ value2.add(10);
+ value1.add(6);
+
+ assertThat(value1.read(), Matchers.equalTo(11));
+ assertThat(value2.read(), Matchers.equalTo(10));
+
+ // Merging clears the old values and updates the result value.
+ StateMerging.mergeCombiningValues(Arrays.asList(value1, value2), value1);
+
+ assertThat(value1.read(), Matchers.equalTo(21));
+ assertThat(value2.read(), Matchers.equalTo(0));
+ }
+
+ @Test
+ public void testMergeCombiningValueIntoNewNamespace() throws Exception {
+ CombiningState<Integer, int[], Integer> value1 =
+ underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
+ CombiningState<Integer, int[], Integer> value2 =
+ underTest.state(NAMESPACE_2, SUM_INTEGER_ADDR);
+ CombiningState<Integer, int[], Integer> value3 =
+ underTest.state(NAMESPACE_3, SUM_INTEGER_ADDR);
+
+ value1.add(5);
+ value2.add(10);
+ value1.add(6);
+
+ StateMerging.mergeCombiningValues(Arrays.asList(value1, value2), value3);
+
+ // Merging clears the old values and updates the result value.
+ assertThat(value1.read(), Matchers.equalTo(0));
+ assertThat(value2.read(), Matchers.equalTo(0));
+ assertThat(value3.read(), Matchers.equalTo(21));
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkKeyGroupStateInternalsTest.java
----------------------------------------------------------------------
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkKeyGroupStateInternalsTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkKeyGroupStateInternalsTest.java
index aed14f3..4012373 100644
--- a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkKeyGroupStateInternalsTest.java
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkKeyGroupStateInternalsTest.java
@@ -17,6 +17,8 @@
*/
package org.apache.beam.runners.flink.streaming;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertThat;
import java.io.ByteArrayInputStream;
@@ -24,8 +26,8 @@ import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.nio.ByteBuffer;
-import org.apache.beam.runners.core.StateInternals;
-import org.apache.beam.runners.core.StateInternalsTest;
+import java.util.Arrays;
+import org.apache.beam.runners.core.StateMerging;
import org.apache.beam.runners.core.StateNamespace;
import org.apache.beam.runners.core.StateNamespaceForTest;
import org.apache.beam.runners.core.StateTag;
@@ -33,6 +35,7 @@ import org.apache.beam.runners.core.StateTags;
import org.apache.beam.runners.flink.translation.wrappers.streaming.state.FlinkKeyGroupStateInternals;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.ReadableState;
import org.apache.beam.sdk.util.CoderUtils;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.JobID;
@@ -44,219 +47,215 @@ import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.memory.MemoryStateBackend;
+import org.apache.flink.streaming.api.operators.KeyContext;
import org.hamcrest.Matchers;
-import org.junit.Ignore;
+import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
-import org.junit.runners.Suite;
/**
* Tests for {@link FlinkKeyGroupStateInternals}. This is based on the tests for
- * {@code StateInternalsTest}.
+ * {@code InMemoryStateInternals}.
*/
-@RunWith(Suite.class)
-@Suite.SuiteClasses({
- FlinkKeyGroupStateInternalsTest.StandardStateInternalsTests.class,
- FlinkKeyGroupStateInternalsTest.OtherTests.class
-})
+@RunWith(JUnit4.class)
public class FlinkKeyGroupStateInternalsTest {
+ private static final StateNamespace NAMESPACE_1 = new StateNamespaceForTest("ns1");
+ private static final StateNamespace NAMESPACE_2 = new StateNamespaceForTest("ns2");
+ private static final StateNamespace NAMESPACE_3 = new StateNamespaceForTest("ns3");
- /**
- * A standard StateInternals test. Just test BagState.
- */
- @RunWith(JUnit4.class)
- public static class StandardStateInternalsTests extends StateInternalsTest {
- @Override
- protected StateInternals createStateInternals() {
- KeyedStateBackend keyedStateBackend =
- getKeyedStateBackend(2, new KeyGroupRange(0, 1));
- return new FlinkKeyGroupStateInternals<>(StringUtf8Coder.of(), keyedStateBackend);
+ private static final StateTag<BagState<String>> STRING_BAG_ADDR =
+ StateTags.bag("stringBag", StringUtf8Coder.of());
+
+ FlinkKeyGroupStateInternals<String> underTest;
+ private KeyedStateBackend keyedStateBackend;
+
+ @Before
+ public void initStateInternals() {
+ try {
+ keyedStateBackend = getKeyedStateBackend(2, new KeyGroupRange(0, 1));
+ underTest = new FlinkKeyGroupStateInternals<>(StringUtf8Coder.of(), keyedStateBackend);
+ } catch (Exception e) {
+ throw new RuntimeException(e);
}
+ }
- @Override
- @Ignore
- public void testValue() {}
+ private KeyedStateBackend getKeyedStateBackend(int numberOfKeyGroups,
+ KeyGroupRange keyGroupRange) {
+ MemoryStateBackend backend = new MemoryStateBackend();
+ try {
+ AbstractKeyedStateBackend<ByteBuffer> keyedStateBackend = backend.createKeyedStateBackend(
+ new DummyEnvironment("test", 1, 0),
+ new JobID(),
+ "test_op",
+ new GenericTypeInfo<>(ByteBuffer.class).createSerializer(new ExecutionConfig()),
+ numberOfKeyGroups,
+ keyGroupRange,
+ new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()));
+ keyedStateBackend.setCurrentKey(ByteBuffer.wrap(
+ CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "1")));
+ return keyedStateBackend;
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
- @Override
- @Ignore
- public void testSet() {}
+ @Test
+ public void testBag() throws Exception {
+ BagState<String> value = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
- @Override
- @Ignore
- public void testSetIsEmpty() {}
+ assertEquals(value, underTest.state(NAMESPACE_1, STRING_BAG_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, STRING_BAG_ADDR)));
- @Override
- @Ignore
- public void testMergeSetIntoSource() {}
+ assertThat(value.read(), Matchers.emptyIterable());
+ value.add("hello");
+ assertThat(value.read(), Matchers.containsInAnyOrder("hello"));
- @Override
- @Ignore
- public void testMergeSetIntoNewNamespace() {}
+ value.add("world");
+ assertThat(value.read(), Matchers.containsInAnyOrder("hello", "world"));
- @Override
- @Ignore
- public void testMap() {}
+ value.clear();
+ assertThat(value.read(), Matchers.emptyIterable());
+ assertEquals(underTest.state(NAMESPACE_1, STRING_BAG_ADDR), value);
- @Override
- @Ignore
- public void testCombiningValue() {}
+ }
- @Override
- @Ignore
- public void testCombiningIsEmpty() {}
+ @Test
+ public void testBagIsEmpty() throws Exception {
+ BagState<String> value = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
- @Override
- @Ignore
- public void testMergeCombiningValueIntoSource() {}
+ assertThat(value.isEmpty().read(), Matchers.is(true));
+ ReadableState<Boolean> readFuture = value.isEmpty();
+ value.add("hello");
+ assertThat(readFuture.read(), Matchers.is(false));
- @Override
- @Ignore
- public void testMergeCombiningValueIntoNewNamespace() {}
+ value.clear();
+ assertThat(readFuture.read(), Matchers.is(true));
+ }
- @Override
- @Ignore
- public void testWatermarkEarliestState() {}
+ @Test
+ public void testMergeBagIntoSource() throws Exception {
+ BagState<String> bag1 = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
+ BagState<String> bag2 = underTest.state(NAMESPACE_2, STRING_BAG_ADDR);
- @Override
- @Ignore
- public void testWatermarkLatestState() {}
+ bag1.add("Hello");
+ bag2.add("World");
+ bag1.add("!");
- @Override
- @Ignore
- public void testWatermarkEndOfWindowState() {}
+ StateMerging.mergeBags(Arrays.asList(bag1, bag2), bag1);
- @Override
- @Ignore
- public void testWatermarkStateIsEmpty() {}
+ // Reading the merged bag gets both the contents
+ assertThat(bag1.read(), Matchers.containsInAnyOrder("Hello", "World", "!"));
+ assertThat(bag2.read(), Matchers.emptyIterable());
+ }
- @Override
- @Ignore
- public void testMergeEarliestWatermarkIntoSource() {}
+ @Test
+ public void testMergeBagIntoNewNamespace() throws Exception {
+ BagState<String> bag1 = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
+ BagState<String> bag2 = underTest.state(NAMESPACE_2, STRING_BAG_ADDR);
+ BagState<String> bag3 = underTest.state(NAMESPACE_3, STRING_BAG_ADDR);
- @Override
- @Ignore
- public void testMergeLatestWatermarkIntoSource() {}
+ bag1.add("Hello");
+ bag2.add("World");
+ bag1.add("!");
- @Override
- @Ignore
- public void testSetReadable() {}
+ StateMerging.mergeBags(Arrays.asList(bag1, bag2, bag3), bag3);
- @Override
- @Ignore
- public void testMapReadable() {}
+ // Reading the merged bag gets both the contents
+ assertThat(bag3.read(), Matchers.containsInAnyOrder("Hello", "World", "!"));
+ assertThat(bag1.read(), Matchers.emptyIterable());
+ assertThat(bag2.read(), Matchers.emptyIterable());
}
- /**
- * A specific test of FlinkKeyGroupStateInternalsTest.
- */
- @RunWith(JUnit4.class)
- public static class OtherTests {
-
- private static final StateNamespace NAMESPACE_1 = new StateNamespaceForTest("ns1");
- private static final StateNamespace NAMESPACE_2 = new StateNamespaceForTest("ns2");
- private static final StateTag<BagState<String>> STRING_BAG_ADDR =
- StateTags.bag("stringBag", StringUtf8Coder.of());
-
- @Test
- public void testKeyGroupAndCheckpoint() throws Exception {
- // assign to keyGroup 0
- ByteBuffer key0 = ByteBuffer.wrap(
- CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "11111111"));
- // assign to keyGroup 1
- ByteBuffer key1 = ByteBuffer.wrap(
- CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "22222222"));
- FlinkKeyGroupStateInternals<String> allState;
- {
- KeyedStateBackend<ByteBuffer> keyedStateBackend =
- getKeyedStateBackend(2, new KeyGroupRange(0, 1));
- allState = new FlinkKeyGroupStateInternals<>(
- StringUtf8Coder.of(), keyedStateBackend);
- BagState<String> valueForNamespace0 = allState.state(NAMESPACE_1, STRING_BAG_ADDR);
- BagState<String> valueForNamespace1 = allState.state(NAMESPACE_2, STRING_BAG_ADDR);
- keyedStateBackend.setCurrentKey(key0);
- valueForNamespace0.add("0");
- valueForNamespace1.add("2");
- keyedStateBackend.setCurrentKey(key1);
- valueForNamespace0.add("1");
- valueForNamespace1.add("3");
- assertThat(valueForNamespace0.read(), Matchers.containsInAnyOrder("0", "1"));
- assertThat(valueForNamespace1.read(), Matchers.containsInAnyOrder("2", "3"));
- }
-
- ClassLoader classLoader = FlinkKeyGroupStateInternalsTest.class.getClassLoader();
-
- // 1. scale up
- ByteArrayOutputStream out0 = new ByteArrayOutputStream();
- allState.snapshotKeyGroupState(0, new DataOutputStream(out0));
- DataInputStream in0 = new DataInputStream(
- new ByteArrayInputStream(out0.toByteArray()));
- {
- KeyedStateBackend<ByteBuffer> keyedStateBackend =
- getKeyedStateBackend(2, new KeyGroupRange(0, 0));
- FlinkKeyGroupStateInternals<String> state0 =
- new FlinkKeyGroupStateInternals<>(
- StringUtf8Coder.of(), keyedStateBackend);
- state0.restoreKeyGroupState(0, in0, classLoader);
- BagState<String> valueForNamespace0 = state0.state(NAMESPACE_1, STRING_BAG_ADDR);
- BagState<String> valueForNamespace1 = state0.state(NAMESPACE_2, STRING_BAG_ADDR);
- assertThat(valueForNamespace0.read(), Matchers.containsInAnyOrder("0"));
- assertThat(valueForNamespace1.read(), Matchers.containsInAnyOrder("2"));
- }
-
- ByteArrayOutputStream out1 = new ByteArrayOutputStream();
- allState.snapshotKeyGroupState(1, new DataOutputStream(out1));
- DataInputStream in1 = new DataInputStream(
- new ByteArrayInputStream(out1.toByteArray()));
- {
- KeyedStateBackend<ByteBuffer> keyedStateBackend =
- getKeyedStateBackend(2, new KeyGroupRange(1, 1));
- FlinkKeyGroupStateInternals<String> state1 =
- new FlinkKeyGroupStateInternals<>(
- StringUtf8Coder.of(), keyedStateBackend);
- state1.restoreKeyGroupState(1, in1, classLoader);
- BagState<String> valueForNamespace0 = state1.state(NAMESPACE_1, STRING_BAG_ADDR);
- BagState<String> valueForNamespace1 = state1.state(NAMESPACE_2, STRING_BAG_ADDR);
- assertThat(valueForNamespace0.read(), Matchers.containsInAnyOrder("1"));
- assertThat(valueForNamespace1.read(), Matchers.containsInAnyOrder("3"));
- }
-
- // 2. scale down
- {
- KeyedStateBackend<ByteBuffer> keyedStateBackend =
- getKeyedStateBackend(2, new KeyGroupRange(0, 1));
- FlinkKeyGroupStateInternals<String> newAllState = new FlinkKeyGroupStateInternals<>(
- StringUtf8Coder.of(), keyedStateBackend);
- in0.reset();
- in1.reset();
- newAllState.restoreKeyGroupState(0, in0, classLoader);
- newAllState.restoreKeyGroupState(1, in1, classLoader);
- BagState<String> valueForNamespace0 = newAllState.state(NAMESPACE_1, STRING_BAG_ADDR);
- BagState<String> valueForNamespace1 = newAllState.state(NAMESPACE_2, STRING_BAG_ADDR);
- assertThat(valueForNamespace0.read(), Matchers.containsInAnyOrder("0", "1"));
- assertThat(valueForNamespace1.read(), Matchers.containsInAnyOrder("2", "3"));
- }
+ @Test
+ public void testKeyGroupAndCheckpoint() throws Exception {
+ // assign to keyGroup 0
+ ByteBuffer key0 = ByteBuffer.wrap(
+ CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "11111111"));
+ // assign to keyGroup 1
+ ByteBuffer key1 = ByteBuffer.wrap(
+ CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "22222222"));
+ FlinkKeyGroupStateInternals<String> allState;
+ {
+ KeyedStateBackend keyedStateBackend = getKeyedStateBackend(2, new KeyGroupRange(0, 1));
+ allState = new FlinkKeyGroupStateInternals<>(
+ StringUtf8Coder.of(), keyedStateBackend);
+ BagState<String> valueForNamespace0 = allState.state(NAMESPACE_1, STRING_BAG_ADDR);
+ BagState<String> valueForNamespace1 = allState.state(NAMESPACE_2, STRING_BAG_ADDR);
+ keyedStateBackend.setCurrentKey(key0);
+ valueForNamespace0.add("0");
+ valueForNamespace1.add("2");
+ keyedStateBackend.setCurrentKey(key1);
+ valueForNamespace0.add("1");
+ valueForNamespace1.add("3");
+ assertThat(valueForNamespace0.read(), Matchers.containsInAnyOrder("0", "1"));
+ assertThat(valueForNamespace1.read(), Matchers.containsInAnyOrder("2", "3"));
+ }
+
+ ClassLoader classLoader = FlinkKeyGroupStateInternalsTest.class.getClassLoader();
+
+ // 1. scale up
+ ByteArrayOutputStream out0 = new ByteArrayOutputStream();
+ allState.snapshotKeyGroupState(0, new DataOutputStream(out0));
+ DataInputStream in0 = new DataInputStream(
+ new ByteArrayInputStream(out0.toByteArray()));
+ {
+ KeyedStateBackend keyedStateBackend = getKeyedStateBackend(2, new KeyGroupRange(0, 0));
+ FlinkKeyGroupStateInternals<String> state0 =
+ new FlinkKeyGroupStateInternals<>(
+ StringUtf8Coder.of(), keyedStateBackend);
+ state0.restoreKeyGroupState(0, in0, classLoader);
+ BagState<String> valueForNamespace0 = state0.state(NAMESPACE_1, STRING_BAG_ADDR);
+ BagState<String> valueForNamespace1 = state0.state(NAMESPACE_2, STRING_BAG_ADDR);
+ assertThat(valueForNamespace0.read(), Matchers.containsInAnyOrder("0"));
+ assertThat(valueForNamespace1.read(), Matchers.containsInAnyOrder("2"));
+ }
+
+ ByteArrayOutputStream out1 = new ByteArrayOutputStream();
+ allState.snapshotKeyGroupState(1, new DataOutputStream(out1));
+ DataInputStream in1 = new DataInputStream(
+ new ByteArrayInputStream(out1.toByteArray()));
+ {
+ KeyedStateBackend keyedStateBackend = getKeyedStateBackend(2, new KeyGroupRange(1, 1));
+ FlinkKeyGroupStateInternals<String> state1 =
+ new FlinkKeyGroupStateInternals<>(
+ StringUtf8Coder.of(), keyedStateBackend);
+ state1.restoreKeyGroupState(1, in1, classLoader);
+ BagState<String> valueForNamespace0 = state1.state(NAMESPACE_1, STRING_BAG_ADDR);
+ BagState<String> valueForNamespace1 = state1.state(NAMESPACE_2, STRING_BAG_ADDR);
+ assertThat(valueForNamespace0.read(), Matchers.containsInAnyOrder("1"));
+ assertThat(valueForNamespace1.read(), Matchers.containsInAnyOrder("3"));
+ }
+
+ // 2. scale down
+ {
+ KeyedStateBackend keyedStateBackend = getKeyedStateBackend(2, new KeyGroupRange(0, 1));
+ FlinkKeyGroupStateInternals<String> newAllState = new FlinkKeyGroupStateInternals<>(
+ StringUtf8Coder.of(), keyedStateBackend);
+ in0.reset();
+ in1.reset();
+ newAllState.restoreKeyGroupState(0, in0, classLoader);
+ newAllState.restoreKeyGroupState(1, in1, classLoader);
+ BagState<String> valueForNamespace0 = newAllState.state(NAMESPACE_1, STRING_BAG_ADDR);
+ BagState<String> valueForNamespace1 = newAllState.state(NAMESPACE_2, STRING_BAG_ADDR);
+ assertThat(valueForNamespace0.read(), Matchers.containsInAnyOrder("0", "1"));
+ assertThat(valueForNamespace1.read(), Matchers.containsInAnyOrder("2", "3"));
}
}
- private static KeyedStateBackend<ByteBuffer> getKeyedStateBackend(int numberOfKeyGroups,
- KeyGroupRange keyGroupRange) {
- MemoryStateBackend backend = new MemoryStateBackend();
- try {
- AbstractKeyedStateBackend<ByteBuffer> keyedStateBackend = backend.createKeyedStateBackend(
- new DummyEnvironment("test", 1, 0),
- new JobID(),
- "test_op",
- new GenericTypeInfo<>(ByteBuffer.class).createSerializer(new ExecutionConfig()),
- numberOfKeyGroups,
- keyGroupRange,
- new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()));
- keyedStateBackend.setCurrentKey(ByteBuffer.wrap(
- CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "1")));
- return keyedStateBackend;
- } catch (Exception e) {
- throw new RuntimeException(e);
+ private static class TestKeyContext implements KeyContext {
+
+ private Object key;
+
+ @Override
+ public void setCurrentKey(Object key) {
+ this.key = key;
+ }
+
+ @Override
+ public Object getCurrentKey() {
+ return key;
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkSplitStateInternalsTest.java
----------------------------------------------------------------------
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkSplitStateInternalsTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkSplitStateInternalsTest.java
index 667b5ba..17cd3f5 100644
--- a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkSplitStateInternalsTest.java
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkSplitStateInternalsTest.java
@@ -17,115 +17,85 @@
*/
package org.apache.beam.runners.flink.streaming;
-import org.apache.beam.runners.core.StateInternals;
-import org.apache.beam.runners.core.StateInternalsTest;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThat;
+
+import org.apache.beam.runners.core.StateNamespace;
+import org.apache.beam.runners.core.StateNamespaceForTest;
+import org.apache.beam.runners.core.StateTag;
+import org.apache.beam.runners.core.StateTags;
import org.apache.beam.runners.flink.translation.wrappers.streaming.state.FlinkSplitStateInternals;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.ReadableState;
import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
import org.apache.flink.runtime.state.OperatorStateBackend;
import org.apache.flink.runtime.state.memory.MemoryStateBackend;
-import org.junit.Ignore;
+import org.hamcrest.Matchers;
+import org.junit.Before;
+import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
/**
* Tests for {@link FlinkSplitStateInternals}. This is based on the tests for
- * {@code StateInternalsTest}.
- *
- * <p>Just test testBag and testBagIsEmpty.
+ * {@code InMemoryStateInternals}.
*/
@RunWith(JUnit4.class)
-public class FlinkSplitStateInternalsTest extends StateInternalsTest {
+public class FlinkSplitStateInternalsTest {
+ private static final StateNamespace NAMESPACE_1 = new StateNamespaceForTest("ns1");
+ private static final StateNamespace NAMESPACE_2 = new StateNamespaceForTest("ns2");
+
+ private static final StateTag<BagState<String>> STRING_BAG_ADDR =
+ StateTags.bag("stringBag", StringUtf8Coder.of());
+
+ FlinkSplitStateInternals<String> underTest;
- @Override
- protected StateInternals createStateInternals() {
+ @Before
+ public void initStateInternals() {
MemoryStateBackend backend = new MemoryStateBackend();
try {
OperatorStateBackend operatorStateBackend =
backend.createOperatorStateBackend(new DummyEnvironment("test", 1, 0), "");
- return new FlinkSplitStateInternals<>(operatorStateBackend);
+ underTest = new FlinkSplitStateInternals<>(operatorStateBackend);
+
} catch (Exception e) {
throw new RuntimeException(e);
}
}
- @Override
- @Ignore
- public void testMergeBagIntoSource() {}
-
- @Override
- @Ignore
- public void testMergeBagIntoNewNamespace() {}
-
- @Override
- @Ignore
- public void testValue() {}
-
- @Override
- @Ignore
- public void testSet() {}
-
- @Override
- @Ignore
- public void testSetIsEmpty() {}
-
- @Override
- @Ignore
- public void testMergeSetIntoSource() {}
+ @Test
+ public void testBag() throws Exception {
+ BagState<String> value = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
- @Override
- @Ignore
- public void testMergeSetIntoNewNamespace() {}
+ assertEquals(value, underTest.state(NAMESPACE_1, STRING_BAG_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, STRING_BAG_ADDR)));
- @Override
- @Ignore
- public void testMap() {}
+ assertThat(value.read(), Matchers.emptyIterable());
+ value.add("hello");
+ assertThat(value.read(), Matchers.containsInAnyOrder("hello"));
- @Override
- @Ignore
- public void testCombiningValue() {}
+ value.add("world");
+ assertThat(value.read(), Matchers.containsInAnyOrder("hello", "world"));
- @Override
- @Ignore
- public void testCombiningIsEmpty() {}
+ value.clear();
+ assertThat(value.read(), Matchers.emptyIterable());
+ assertEquals(underTest.state(NAMESPACE_1, STRING_BAG_ADDR), value);
- @Override
- @Ignore
- public void testMergeCombiningValueIntoSource() {}
-
- @Override
- @Ignore
- public void testMergeCombiningValueIntoNewNamespace() {}
-
- @Override
- @Ignore
- public void testWatermarkEarliestState() {}
-
- @Override
- @Ignore
- public void testWatermarkLatestState() {}
-
- @Override
- @Ignore
- public void testWatermarkEndOfWindowState() {}
-
- @Override
- @Ignore
- public void testWatermarkStateIsEmpty() {}
-
- @Override
- @Ignore
- public void testMergeEarliestWatermarkIntoSource() {}
+ }
- @Override
- @Ignore
- public void testMergeLatestWatermarkIntoSource() {}
+ @Test
+ public void testBagIsEmpty() throws Exception {
+ BagState<String> value = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
- @Override
- @Ignore
- public void testSetReadable() {}
+ assertThat(value.isEmpty().read(), Matchers.is(true));
+ ReadableState<Boolean> readFuture = value.isEmpty();
+ value.add("hello");
+ assertThat(readFuture.read(), Matchers.is(false));
- @Override
- @Ignore
- public void testMapReadable() {}
+ value.clear();
+ assertThat(readFuture.read(), Matchers.is(true));
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkStateInternalsTest.java
----------------------------------------------------------------------
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkStateInternalsTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkStateInternalsTest.java
index b8d41de..35d2b78 100644
--- a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkStateInternalsTest.java
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/FlinkStateInternalsTest.java
@@ -17,11 +17,31 @@
*/
package org.apache.beam.runners.flink.streaming;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertThat;
+
import java.nio.ByteBuffer;
-import org.apache.beam.runners.core.StateInternals;
-import org.apache.beam.runners.core.StateInternalsTest;
+import java.util.Arrays;
+import org.apache.beam.runners.core.StateMerging;
+import org.apache.beam.runners.core.StateNamespace;
+import org.apache.beam.runners.core.StateNamespaceForTest;
+import org.apache.beam.runners.core.StateTag;
+import org.apache.beam.runners.core.StateTags;
import org.apache.beam.runners.flink.translation.wrappers.streaming.state.FlinkStateInternals;
import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.state.BagState;
+import org.apache.beam.sdk.state.CombiningState;
+import org.apache.beam.sdk.state.GroupingState;
+import org.apache.beam.sdk.state.ReadableState;
+import org.apache.beam.sdk.state.ValueState;
+import org.apache.beam.sdk.state.WatermarkHoldState;
+import org.apache.beam.sdk.transforms.Sum;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
+import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
import org.apache.beam.sdk.util.CoderUtils;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.JobID;
@@ -32,17 +52,42 @@ import org.apache.flink.runtime.query.KvStateRegistry;
import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.memory.MemoryStateBackend;
+import org.hamcrest.Matchers;
+import org.joda.time.Instant;
+import org.junit.Before;
+import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
/**
- * Tests for {@link FlinkStateInternals}. This is based on {@link StateInternalsTest}.
+ * Tests for {@link FlinkStateInternals}. This is based on the tests for
+ * {@code InMemoryStateInternals}.
*/
@RunWith(JUnit4.class)
-public class FlinkStateInternalsTest extends StateInternalsTest {
+public class FlinkStateInternalsTest {
+ private static final BoundedWindow WINDOW_1 = new IntervalWindow(new Instant(0), new Instant(10));
+ private static final StateNamespace NAMESPACE_1 = new StateNamespaceForTest("ns1");
+ private static final StateNamespace NAMESPACE_2 = new StateNamespaceForTest("ns2");
+ private static final StateNamespace NAMESPACE_3 = new StateNamespaceForTest("ns3");
+
+ private static final StateTag<ValueState<String>> STRING_VALUE_ADDR =
+ StateTags.value("stringValue", StringUtf8Coder.of());
+ private static final StateTag<CombiningState<Integer, int[], Integer>>
+ SUM_INTEGER_ADDR = StateTags.combiningValueFromInputInternal(
+ "sumInteger", VarIntCoder.of(), Sum.ofIntegers());
+ private static final StateTag<BagState<String>> STRING_BAG_ADDR =
+ StateTags.bag("stringBag", StringUtf8Coder.of());
+ private static final StateTag<WatermarkHoldState> WATERMARK_EARLIEST_ADDR =
+ StateTags.watermarkStateInternal("watermark", TimestampCombiner.EARLIEST);
+ private static final StateTag<WatermarkHoldState> WATERMARK_LATEST_ADDR =
+ StateTags.watermarkStateInternal("watermark", TimestampCombiner.LATEST);
+ private static final StateTag<WatermarkHoldState> WATERMARK_EOW_ADDR =
+ StateTags.watermarkStateInternal("watermark", TimestampCombiner.END_OF_WINDOW);
+
+ FlinkStateInternals<String> underTest;
- @Override
- protected StateInternals createStateInternals() {
+ @Before
+ public void initStateInternals() {
MemoryStateBackend backend = new MemoryStateBackend();
try {
AbstractKeyedStateBackend<ByteBuffer> keyedStateBackend = backend.createKeyedStateBackend(
@@ -53,14 +98,296 @@ public class FlinkStateInternalsTest extends StateInternalsTest {
1,
new KeyGroupRange(0, 0),
new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()));
+ underTest = new FlinkStateInternals<>(keyedStateBackend, StringUtf8Coder.of());
keyedStateBackend.setCurrentKey(
ByteBuffer.wrap(CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "Hello")));
-
- return new FlinkStateInternals<>(keyedStateBackend, StringUtf8Coder.of());
} catch (Exception e) {
throw new RuntimeException(e);
}
}
+ @Test
+ public void testValue() throws Exception {
+ ValueState<String> value = underTest.state(NAMESPACE_1, STRING_VALUE_ADDR);
+
+ assertEquals(underTest.state(NAMESPACE_1, STRING_VALUE_ADDR), value);
+ assertNotEquals(
+ underTest.state(NAMESPACE_2, STRING_VALUE_ADDR),
+ value);
+
+ assertThat(value.read(), Matchers.nullValue());
+ value.write("hello");
+ assertThat(value.read(), Matchers.equalTo("hello"));
+ value.write("world");
+ assertThat(value.read(), Matchers.equalTo("world"));
+
+ value.clear();
+ assertThat(value.read(), Matchers.nullValue());
+ assertEquals(underTest.state(NAMESPACE_1, STRING_VALUE_ADDR), value);
+
+ }
+
+ @Test
+ public void testBag() throws Exception {
+ BagState<String> value = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
+
+ assertEquals(value, underTest.state(NAMESPACE_1, STRING_BAG_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, STRING_BAG_ADDR)));
+
+ assertThat(value.read(), Matchers.emptyIterable());
+ value.add("hello");
+ assertThat(value.read(), Matchers.containsInAnyOrder("hello"));
+
+ value.add("world");
+ assertThat(value.read(), Matchers.containsInAnyOrder("hello", "world"));
+
+ value.clear();
+ assertThat(value.read(), Matchers.emptyIterable());
+ assertEquals(underTest.state(NAMESPACE_1, STRING_BAG_ADDR), value);
+
+ }
+
+ @Test
+ public void testBagIsEmpty() throws Exception {
+ BagState<String> value = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
+
+ assertThat(value.isEmpty().read(), Matchers.is(true));
+ ReadableState<Boolean> readFuture = value.isEmpty();
+ value.add("hello");
+ assertThat(readFuture.read(), Matchers.is(false));
+
+ value.clear();
+ assertThat(readFuture.read(), Matchers.is(true));
+ }
+
+ @Test
+ public void testMergeBagIntoSource() throws Exception {
+ BagState<String> bag1 = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
+ BagState<String> bag2 = underTest.state(NAMESPACE_2, STRING_BAG_ADDR);
+
+ bag1.add("Hello");
+ bag2.add("World");
+ bag1.add("!");
+
+ StateMerging.mergeBags(Arrays.asList(bag1, bag2), bag1);
+
+ // Reading the merged bag gets both the contents
+ assertThat(bag1.read(), Matchers.containsInAnyOrder("Hello", "World", "!"));
+ assertThat(bag2.read(), Matchers.emptyIterable());
+ }
+
+ @Test
+ public void testMergeBagIntoNewNamespace() throws Exception {
+ BagState<String> bag1 = underTest.state(NAMESPACE_1, STRING_BAG_ADDR);
+ BagState<String> bag2 = underTest.state(NAMESPACE_2, STRING_BAG_ADDR);
+ BagState<String> bag3 = underTest.state(NAMESPACE_3, STRING_BAG_ADDR);
+
+ bag1.add("Hello");
+ bag2.add("World");
+ bag1.add("!");
+
+ StateMerging.mergeBags(Arrays.asList(bag1, bag2, bag3), bag3);
+
+ // Reading the merged bag gets both the contents
+ assertThat(bag3.read(), Matchers.containsInAnyOrder("Hello", "World", "!"));
+ assertThat(bag1.read(), Matchers.emptyIterable());
+ assertThat(bag2.read(), Matchers.emptyIterable());
+ }
+
+ @Test
+ public void testCombiningValue() throws Exception {
+ GroupingState<Integer, Integer> value = underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
+
+ // State instances are cached, but depend on the namespace.
+ assertEquals(value, underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, SUM_INTEGER_ADDR)));
+
+ assertThat(value.read(), Matchers.equalTo(0));
+ value.add(2);
+ assertThat(value.read(), Matchers.equalTo(2));
+
+ value.add(3);
+ assertThat(value.read(), Matchers.equalTo(5));
+
+ value.clear();
+ assertThat(value.read(), Matchers.equalTo(0));
+ assertEquals(underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR), value);
+ }
+
+ @Test
+ public void testCombiningIsEmpty() throws Exception {
+ GroupingState<Integer, Integer> value = underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
+
+ assertThat(value.isEmpty().read(), Matchers.is(true));
+ ReadableState<Boolean> readFuture = value.isEmpty();
+ value.add(5);
+ assertThat(readFuture.read(), Matchers.is(false));
+
+ value.clear();
+ assertThat(readFuture.read(), Matchers.is(true));
+ }
+
+ @Test
+ public void testMergeCombiningValueIntoSource() throws Exception {
+ CombiningState<Integer, int[], Integer> value1 =
+ underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
+ CombiningState<Integer, int[], Integer> value2 =
+ underTest.state(NAMESPACE_2, SUM_INTEGER_ADDR);
+
+ value1.add(5);
+ value2.add(10);
+ value1.add(6);
+
+ assertThat(value1.read(), Matchers.equalTo(11));
+ assertThat(value2.read(), Matchers.equalTo(10));
+
+ // Merging clears the old values and updates the result value.
+ StateMerging.mergeCombiningValues(Arrays.asList(value1, value2), value1);
+
+ assertThat(value1.read(), Matchers.equalTo(21));
+ assertThat(value2.read(), Matchers.equalTo(0));
+ }
+
+ @Test
+ public void testMergeCombiningValueIntoNewNamespace() throws Exception {
+ CombiningState<Integer, int[], Integer> value1 =
+ underTest.state(NAMESPACE_1, SUM_INTEGER_ADDR);
+ CombiningState<Integer, int[], Integer> value2 =
+ underTest.state(NAMESPACE_2, SUM_INTEGER_ADDR);
+ CombiningState<Integer, int[], Integer> value3 =
+ underTest.state(NAMESPACE_3, SUM_INTEGER_ADDR);
+
+ value1.add(5);
+ value2.add(10);
+ value1.add(6);
+
+ StateMerging.mergeCombiningValues(Arrays.asList(value1, value2), value3);
+
+ // Merging clears the old values and updates the result value.
+ assertThat(value1.read(), Matchers.equalTo(0));
+ assertThat(value2.read(), Matchers.equalTo(0));
+ assertThat(value3.read(), Matchers.equalTo(21));
+ }
+
+ @Test
+ public void testWatermarkEarliestState() throws Exception {
+ WatermarkHoldState value =
+ underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR);
+
+ // State instances are cached, but depend on the namespace.
+ assertEquals(value, underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, WATERMARK_EARLIEST_ADDR)));
+
+ assertThat(value.read(), Matchers.nullValue());
+ value.add(new Instant(2000));
+ assertThat(value.read(), Matchers.equalTo(new Instant(2000)));
+
+ value.add(new Instant(3000));
+ assertThat(value.read(), Matchers.equalTo(new Instant(2000)));
+
+ value.add(new Instant(1000));
+ assertThat(value.read(), Matchers.equalTo(new Instant(1000)));
+
+ value.clear();
+ assertThat(value.read(), Matchers.equalTo(null));
+ assertEquals(underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR), value);
+ }
+
+ @Test
+ public void testWatermarkLatestState() throws Exception {
+ WatermarkHoldState value =
+ underTest.state(NAMESPACE_1, WATERMARK_LATEST_ADDR);
+
+ // State instances are cached, but depend on the namespace.
+ assertEquals(value, underTest.state(NAMESPACE_1, WATERMARK_LATEST_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, WATERMARK_LATEST_ADDR)));
+
+ assertThat(value.read(), Matchers.nullValue());
+ value.add(new Instant(2000));
+ assertThat(value.read(), Matchers.equalTo(new Instant(2000)));
+
+ value.add(new Instant(3000));
+ assertThat(value.read(), Matchers.equalTo(new Instant(3000)));
+
+ value.add(new Instant(1000));
+ assertThat(value.read(), Matchers.equalTo(new Instant(3000)));
+
+ value.clear();
+ assertThat(value.read(), Matchers.equalTo(null));
+ assertEquals(underTest.state(NAMESPACE_1, WATERMARK_LATEST_ADDR), value);
+ }
+
+ @Test
+ public void testWatermarkEndOfWindowState() throws Exception {
+ WatermarkHoldState value = underTest.state(NAMESPACE_1, WATERMARK_EOW_ADDR);
+
+ // State instances are cached, but depend on the namespace.
+ assertEquals(value, underTest.state(NAMESPACE_1, WATERMARK_EOW_ADDR));
+ assertFalse(value.equals(underTest.state(NAMESPACE_2, WATERMARK_EOW_ADDR)));
+
+ assertThat(value.read(), Matchers.nullValue());
+ value.add(new Instant(2000));
+ assertThat(value.read(), Matchers.equalTo(new Instant(2000)));
+
+ value.clear();
+ assertThat(value.read(), Matchers.equalTo(null));
+ assertEquals(underTest.state(NAMESPACE_1, WATERMARK_EOW_ADDR), value);
+ }
+
+ @Test
+ public void testWatermarkStateIsEmpty() throws Exception {
+ WatermarkHoldState value =
+ underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR);
+
+ assertThat(value.isEmpty().read(), Matchers.is(true));
+ ReadableState<Boolean> readFuture = value.isEmpty();
+ value.add(new Instant(1000));
+ assertThat(readFuture.read(), Matchers.is(false));
+
+ value.clear();
+ assertThat(readFuture.read(), Matchers.is(true));
+ }
+
+ @Test
+ public void testMergeEarliestWatermarkIntoSource() throws Exception {
+ WatermarkHoldState value1 =
+ underTest.state(NAMESPACE_1, WATERMARK_EARLIEST_ADDR);
+ WatermarkHoldState value2 =
+ underTest.state(NAMESPACE_2, WATERMARK_EARLIEST_ADDR);
+
+ value1.add(new Instant(3000));
+ value2.add(new Instant(5000));
+ value1.add(new Instant(4000));
+ value2.add(new Instant(2000));
+
+ // Merging clears the old values and updates the merged value.
+ StateMerging.mergeWatermarks(Arrays.asList(value1, value2), value1, WINDOW_1);
+
+ assertThat(value1.read(), Matchers.equalTo(new Instant(2000)));
+ assertThat(value2.read(), Matchers.equalTo(null));
+ }
+
+ @Test
+ public void testMergeLatestWatermarkIntoSource() throws Exception {
+ WatermarkHoldState value1 =
+ underTest.state(NAMESPACE_1, WATERMARK_LATEST_ADDR);
+ WatermarkHoldState value2 =
+ underTest.state(NAMESPACE_2, WATERMARK_LATEST_ADDR);
+ WatermarkHoldState value3 =
+ underTest.state(NAMESPACE_3, WATERMARK_LATEST_ADDR);
+
+ value1.add(new Instant(3000));
+ value2.add(new Instant(5000));
+ value1.add(new Instant(4000));
+ value2.add(new Instant(2000));
+
+ // Merging clears the old values and updates the result value.
+ StateMerging.mergeWatermarks(Arrays.asList(value1, value2), value3, WINDOW_1);
+
+ // Merging clears the old values and updates the result value.
+ assertThat(value3.read(), Matchers.equalTo(new Instant(5000)));
+ assertThat(value1.read(), Matchers.equalTo(null));
+ assertThat(value2.read(), Matchers.equalTo(null));
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/google-cloud-dataflow-java/pom.xml
----------------------------------------------------------------------
diff --git a/runners/google-cloud-dataflow-java/pom.xml b/runners/google-cloud-dataflow-java/pom.xml
index c8d63ac..92c94a8 100644
--- a/runners/google-cloud-dataflow-java/pom.xml
+++ b/runners/google-cloud-dataflow-java/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-runners-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
@@ -33,7 +33,7 @@
<packaging>jar</packaging>
<properties>
- <dataflow.container_version>beam-master-20170706</dataflow.container_version>
+ <dataflow.container_version>beam-master-20170530</dataflow.container_version>
<dataflow.fnapi_environment_major_version>1</dataflow.fnapi_environment_major_version>
<dataflow.legacy_environment_major_version>6</dataflow.legacy_environment_major_version>
</properties>
@@ -216,17 +216,13 @@
<execution>
<id>validates-runner-tests</id>
<configuration>
- <!--
- UsesSplittableParDoWithWindowedSideInputs because of
- https://issues.apache.org/jira/browse/BEAM-2476
- -->
<excludedGroups>
org.apache.beam.sdk.testing.LargeKeys$Above10MB,
org.apache.beam.sdk.testing.UsesDistributionMetrics,
org.apache.beam.sdk.testing.UsesGaugeMetrics,
org.apache.beam.sdk.testing.UsesSetState,
org.apache.beam.sdk.testing.UsesMapState,
- org.apache.beam.sdk.testing.UsesSplittableParDoWithWindowedSideInputs,
+ org.apache.beam.sdk.testing.UsesSplittableParDo,
org.apache.beam.sdk.testing.UsesUnboundedPCollections,
org.apache.beam.sdk.testing.UsesTestStream,
</excludedGroups>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchStatefulParDoOverrides.java
----------------------------------------------------------------------
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchStatefulParDoOverrides.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchStatefulParDoOverrides.java
index 7309f61..4d9a57f 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchStatefulParDoOverrides.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchStatefulParDoOverrides.java
@@ -145,8 +145,6 @@ public class BatchStatefulParDoOverrides {
public PCollection<OutputT> expand(PCollection<KV<K, InputT>> input) {
DoFn<KV<K, InputT>, OutputT> fn = originalParDo.getFn();
verifyFnIsStateful(fn);
- DataflowRunner.verifyStateSupported(fn);
- DataflowRunner.verifyStateSupportForWindowingStrategy(input.getWindowingStrategy());
PTransform<
PCollection<? extends KV<K, Iterable<KV<Instant, WindowedValue<KV<K, InputT>>>>>>,
@@ -171,8 +169,6 @@ public class BatchStatefulParDoOverrides {
public PCollectionTuple expand(PCollection<KV<K, InputT>> input) {
DoFn<KV<K, InputT>, OutputT> fn = originalParDo.getFn();
verifyFnIsStateful(fn);
- DataflowRunner.verifyStateSupported(fn);
- DataflowRunner.verifyStateSupportForWindowingStrategy(input.getWindowingStrategy());
PTransform<
PCollection<? extends KV<K, Iterable<KV<Instant, WindowedValue<KV<K, InputT>>>>>>,
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchViewOverrides.java
----------------------------------------------------------------------
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchViewOverrides.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchViewOverrides.java
index ad3faed..b4a6e64 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchViewOverrides.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchViewOverrides.java
@@ -39,6 +39,8 @@ import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
+import org.apache.beam.runners.core.construction.PTransformReplacements;
+import org.apache.beam.runners.core.construction.SingleInputOutputOverrideFactory;
import org.apache.beam.runners.dataflow.internal.IsmFormat;
import org.apache.beam.runners.dataflow.internal.IsmFormat.IsmRecord;
import org.apache.beam.runners.dataflow.internal.IsmFormat.IsmRecordCoder;
@@ -55,11 +57,17 @@ import org.apache.beam.sdk.coders.SerializableCoder;
import org.apache.beam.sdk.coders.StructuredCoder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.coders.VarLongCoder;
+import org.apache.beam.sdk.runners.AppliedPTransform;
+import org.apache.beam.sdk.transforms.Combine;
+import org.apache.beam.sdk.transforms.Combine.GloballyAsSingletonView;
+import org.apache.beam.sdk.transforms.CombineFnBase.GlobalCombineFn;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.Flatten;
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.View;
+import org.apache.beam.sdk.transforms.View.AsSingleton;
import org.apache.beam.sdk.transforms.View.CreatePCollectionView;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
@@ -75,6 +83,7 @@ import org.apache.beam.sdk.values.PCollection.IsBounded;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PCollectionViews;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
import org.apache.beam.sdk.values.WindowingStrategy;
@@ -183,13 +192,12 @@ class BatchViewOverrides {
}
private final DataflowRunner runner;
- private final PCollectionView<Map<K, V>> view;
- /** Builds an instance of this class from the overridden transform. */
+ /**
+ * Builds an instance of this class from the overridden transform.
+ */
@SuppressWarnings("unused") // used via reflection in DataflowRunner#apply()
- public BatchViewAsMap(
- DataflowRunner runner, CreatePCollectionView<KV<K, V>, Map<K, V>> transform) {
+ public BatchViewAsMap(DataflowRunner runner, View.AsMap<K, V> transform) {
this.runner = runner;
- this.view = transform.getView();
}
@Override
@@ -199,7 +207,12 @@ class BatchViewOverrides {
private <W extends BoundedWindow> PCollectionView<Map<K, V>>
applyInternal(PCollection<KV<K, V>> input) {
+
+ @SuppressWarnings({"rawtypes", "unchecked"})
+ KvCoder<K, V> inputCoder = (KvCoder) input.getCoder();
try {
+ PCollectionView<Map<K, V>> view = PCollectionViews.mapView(
+ input, input.getWindowingStrategy(), inputCoder);
return BatchViewAsMultimap.applyForMapLike(runner, input, view, true /* unique keys */);
} catch (NonDeterministicException e) {
runner.recordViewUsesNonDeterministicKeyCoder(this);
@@ -236,14 +249,19 @@ class BatchViewOverrides {
inputCoder.getKeyCoder(),
FullWindowedValueCoder.of(inputCoder.getValueCoder(), windowCoder)));
+ TransformedMap<K, WindowedValue<V>, V> defaultValue = new TransformedMap<>(
+ WindowedValueToValue.<V>of(),
+ ImmutableMap.<K, WindowedValue<V>>of());
+
return BatchViewAsSingleton.<KV<K, V>, TransformedMap<K, WindowedValue<V>, V>,
Map<K, V>,
W> applyForSingleton(
runner,
input,
new ToMapDoFn<K, V, W>(windowCoder),
- finalValueCoder,
- view);
+ true,
+ defaultValue,
+ finalValueCoder);
}
}
@@ -662,13 +680,12 @@ class BatchViewOverrides {
}
private final DataflowRunner runner;
- private final PCollectionView<Map<K, Iterable<V>>> view;
- /** Builds an instance of this class from the overridden transform. */
+ /**
+ * Builds an instance of this class from the overridden transform.
+ */
@SuppressWarnings("unused") // used via reflection in DataflowRunner#apply()
- public BatchViewAsMultimap(
- DataflowRunner runner, CreatePCollectionView<KV<K, V>, Map<K, Iterable<V>>> transform) {
+ public BatchViewAsMultimap(DataflowRunner runner, View.AsMultimap<K, V> transform) {
this.runner = runner;
- this.view = transform.getView();
}
@Override
@@ -678,7 +695,12 @@ class BatchViewOverrides {
private <W extends BoundedWindow> PCollectionView<Map<K, Iterable<V>>>
applyInternal(PCollection<KV<K, V>> input) {
+ @SuppressWarnings({"rawtypes", "unchecked"})
+ KvCoder<K, V> inputCoder = (KvCoder) input.getCoder();
try {
+ PCollectionView<Map<K, Iterable<V>>> view = PCollectionViews.multimapView(
+ input, input.getWindowingStrategy(), inputCoder);
+
return applyForMapLike(runner, input, view, false /* unique keys not expected */);
} catch (NonDeterministicException e) {
runner.recordViewUsesNonDeterministicKeyCoder(this);
@@ -716,15 +738,16 @@ class BatchViewOverrides {
IterableWithWindowedValuesToIterable.<V>of(),
ImmutableMap.<K, Iterable<WindowedValue<V>>>of());
- return BatchViewAsSingleton
- .<KV<K, V>, TransformedMap<K, Iterable<WindowedValue<V>>, Iterable<V>>,
- Map<K, Iterable<V>>, W>
- applyForSingleton(
- runner,
- input,
- new ToMultimapDoFn<K, V, W>(windowCoder),
- finalValueCoder,
- view);
+ return BatchViewAsSingleton.<KV<K, V>,
+ TransformedMap<K, Iterable<WindowedValue<V>>, Iterable<V>>,
+ Map<K, Iterable<V>>,
+ W> applyForSingleton(
+ runner,
+ input,
+ new ToMultimapDoFn<K, V, W>(windowCoder),
+ true,
+ defaultValue,
+ finalValueCoder);
}
private static <K, V, W extends BoundedWindow, ViewT> PCollectionView<ViewT> applyForMapLike(
@@ -804,9 +827,10 @@ class BatchViewOverrides {
PCollectionList.of(ImmutableList.of(
perHashWithReifiedWindows, windowMapSizeMetadata, windowMapKeysMetadata));
- Pipeline.applyTransform(outputs, Flatten.<IsmRecord<WindowedValue<V>>>pCollections())
- .apply(CreateDataflowView.<IsmRecord<WindowedValue<V>>, ViewT>of(view));
- return view;
+ return Pipeline.applyTransform(outputs,
+ Flatten.<IsmRecord<WindowedValue<V>>>pCollections())
+ .apply(CreateDataflowView.<IsmRecord<WindowedValue<V>>,
+ ViewT>of(view));
}
@Override
@@ -891,12 +915,14 @@ class BatchViewOverrides {
}
private final DataflowRunner runner;
- private final PCollectionView<T> view;
- /** Builds an instance of this class from the overridden transform. */
+ private final View.AsSingleton<T> transform;
+ /**
+ * Builds an instance of this class from the overridden transform.
+ */
@SuppressWarnings("unused") // used via reflection in DataflowRunner#apply()
- public BatchViewAsSingleton(DataflowRunner runner, CreatePCollectionView<T, T> transform) {
+ public BatchViewAsSingleton(DataflowRunner runner, View.AsSingleton<T> transform) {
this.runner = runner;
- this.view = transform.getView();
+ this.transform = transform;
}
@Override
@@ -909,8 +935,9 @@ class BatchViewOverrides {
runner,
input,
new IsmRecordForSingularValuePerWindowDoFn<T, BoundedWindow>(windowCoder),
- input.getCoder(),
- view);
+ transform.hasDefaultValue(),
+ transform.defaultValue(),
+ input.getCoder());
}
static <T, FinalT, ViewT, W extends BoundedWindow> PCollectionView<ViewT>
@@ -919,13 +946,23 @@ class BatchViewOverrides {
PCollection<T> input,
DoFn<KV<Integer, Iterable<KV<W, WindowedValue<T>>>>,
IsmRecord<WindowedValue<FinalT>>> doFn,
- Coder<FinalT> defaultValueCoder,
- PCollectionView<ViewT> view) {
+ boolean hasDefault,
+ FinalT defaultValue,
+ Coder<FinalT> defaultValueCoder) {
@SuppressWarnings("unchecked")
Coder<W> windowCoder = (Coder<W>)
input.getWindowingStrategy().getWindowFn().windowCoder();
+ @SuppressWarnings({"rawtypes", "unchecked"})
+ PCollectionView<ViewT> view =
+ (PCollectionView<ViewT>) PCollectionViews.<FinalT, W>singletonView(
+ (PCollection) input,
+ (WindowingStrategy) input.getWindowingStrategy(),
+ hasDefault,
+ defaultValue,
+ defaultValueCoder);
+
IsmRecordCoder<WindowedValue<FinalT>> ismCoder =
coderForSingleton(windowCoder, defaultValueCoder);
@@ -935,9 +972,8 @@ class BatchViewOverrides {
reifiedPerWindowAndSorted.setCoder(ismCoder);
runner.addPCollectionRequiringIndexedFormat(reifiedPerWindowAndSorted);
- reifiedPerWindowAndSorted.apply(
+ return reifiedPerWindowAndSorted.apply(
CreateDataflowView.<IsmRecord<WindowedValue<FinalT>>, ViewT>of(view));
- return view;
}
@Override
@@ -1043,18 +1079,18 @@ class BatchViewOverrides {
}
private final DataflowRunner runner;
- private final PCollectionView<List<T>> view;
/**
* Builds an instance of this class from the overridden transform.
*/
@SuppressWarnings("unused") // used via reflection in DataflowRunner#apply()
- public BatchViewAsList(DataflowRunner runner, CreatePCollectionView<T, List<T>> transform) {
+ public BatchViewAsList(DataflowRunner runner, View.AsList<T> transform) {
this.runner = runner;
- this.view = transform.getView();
}
@Override
public PCollectionView<List<T>> expand(PCollection<T> input) {
+ PCollectionView<List<T>> view = PCollectionViews.listView(
+ input, input.getWindowingStrategy(), input.getCoder());
return applyForIterableLike(runner, input, view);
}
@@ -1080,9 +1116,8 @@ class BatchViewOverrides {
reifiedPerWindowAndSorted.setCoder(ismCoder);
runner.addPCollectionRequiringIndexedFormat(reifiedPerWindowAndSorted);
- reifiedPerWindowAndSorted.apply(
+ return reifiedPerWindowAndSorted.apply(
CreateDataflowView.<IsmRecord<WindowedValue<T>>, ViewT>of(view));
- return view;
}
PCollection<IsmRecord<WindowedValue<T>>> reifiedPerWindowAndSorted = input
@@ -1091,9 +1126,8 @@ class BatchViewOverrides {
reifiedPerWindowAndSorted.setCoder(ismCoder);
runner.addPCollectionRequiringIndexedFormat(reifiedPerWindowAndSorted);
- reifiedPerWindowAndSorted.apply(
+ return reifiedPerWindowAndSorted.apply(
CreateDataflowView.<IsmRecord<WindowedValue<T>>, ViewT>of(view));
- return view;
}
@Override
@@ -1130,17 +1164,18 @@ class BatchViewOverrides {
extends PTransform<PCollection<T>, PCollectionView<Iterable<T>>> {
private final DataflowRunner runner;
- private final PCollectionView<Iterable<T>> view;
- /** Builds an instance of this class from the overridden transform. */
+ /**
+ * Builds an instance of this class from the overridden transform.
+ */
@SuppressWarnings("unused") // used via reflection in DataflowRunner#apply()
- public BatchViewAsIterable(
- DataflowRunner runner, CreatePCollectionView<T, Iterable<T>> transform) {
+ public BatchViewAsIterable(DataflowRunner runner, View.AsIterable<T> transform) {
this.runner = runner;
- this.view = transform.getView();
}
@Override
public PCollectionView<Iterable<T>> expand(PCollection<T> input) {
+ PCollectionView<Iterable<T>> view = PCollectionViews.iterableView(
+ input, input.getWindowingStrategy(), input.getCoder());
return BatchViewAsList.applyForIterableLike(runner, input, view);
}
}
@@ -1342,4 +1377,59 @@ class BatchViewOverrides {
verifyDeterministic(this, "Expected map coder to be deterministic.", originalMapCoder);
}
}
+
+ static class BatchCombineGloballyAsSingletonViewFactory<ElemT, ViewT>
+ extends SingleInputOutputOverrideFactory<
+ PCollection<ElemT>, PCollectionView<ViewT>,
+ Combine.GloballyAsSingletonView<ElemT, ViewT>> {
+ private final DataflowRunner runner;
+
+ BatchCombineGloballyAsSingletonViewFactory(DataflowRunner runner) {
+ this.runner = runner;
+ }
+
+ @Override
+ public PTransformReplacement<PCollection<ElemT>, PCollectionView<ViewT>>
+ getReplacementTransform(
+ AppliedPTransform<
+ PCollection<ElemT>, PCollectionView<ViewT>,
+ GloballyAsSingletonView<ElemT, ViewT>>
+ transform) {
+ GloballyAsSingletonView<ElemT, ViewT> combine = transform.getTransform();
+ return PTransformReplacement.of(
+ PTransformReplacements.getSingletonMainInput(transform),
+ new BatchCombineGloballyAsSingletonView<>(
+ runner, combine.getCombineFn(), combine.getFanout(), combine.getInsertDefault()));
+ }
+
+ private static class BatchCombineGloballyAsSingletonView<ElemT, ViewT>
+ extends PTransform<PCollection<ElemT>, PCollectionView<ViewT>> {
+ private final DataflowRunner runner;
+ private final GlobalCombineFn<? super ElemT, ?, ViewT> combineFn;
+ private final int fanout;
+ private final boolean insertDefault;
+
+ BatchCombineGloballyAsSingletonView(
+ DataflowRunner runner,
+ GlobalCombineFn<? super ElemT, ?, ViewT> combineFn,
+ int fanout,
+ boolean insertDefault) {
+ this.runner = runner;
+ this.combineFn = combineFn;
+ this.fanout = fanout;
+ this.insertDefault = insertDefault;
+ }
+
+ @Override
+ public PCollectionView<ViewT> expand(PCollection<ElemT> input) {
+ PCollection<ViewT> combined =
+ input.apply(Combine.globally(combineFn).withoutDefaults().withFanout(fanout));
+ AsSingleton<ViewT> viewAsSingleton = View.asSingleton();
+ if (insertDefault) {
+ viewAsSingleton.withDefaultValue(combineFn.defaultValue());
+ }
+ return combined.apply(new BatchViewAsSingleton<>(runner, viewAsSingleton));
+ }
+ }
+ }
}
[08/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseIO.java b/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseIO.java
index 90ede4c..849873c 100644
--- a/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseIO.java
+++ b/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseIO.java
@@ -31,7 +31,10 @@ import java.util.Set;
import java.util.TreeSet;
import javax.annotation.Nullable;
import org.apache.beam.sdk.annotations.Experimental;
+import org.apache.beam.sdk.coders.ByteArrayCoder;
import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.IterableCoder;
+import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.io.BoundedSource;
import org.apache.beam.sdk.io.hadoop.SerializableConfiguration;
import org.apache.beam.sdk.io.range.ByteKey;
@@ -41,6 +44,7 @@ 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.display.DisplayData;
+import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PDone;
@@ -118,15 +122,17 @@ import org.slf4j.LoggerFactory;
* <h3>Writing to HBase</h3>
*
* <p>The HBase sink executes a set of row mutations on a single table. It takes as input a
- * {@link PCollection PCollection<Mutation>}, where each {@link Mutation} represents an
- * idempotent transformation on a row.
+ * {@link PCollection PCollection<KV<byte[], Iterable<Mutation>>>}, where the
+ * {@code byte[]} is the key of the row being mutated, and each {@link Mutation} represents an
+ * idempotent transformation to that row.
*
* <p>To configure a HBase sink, you must supply a table id and a {@link Configuration}
* to identify the HBase instance, for example:
*
* <pre>{@code
* Configuration configuration = ...;
- * PCollection<Mutation> data = ...;
+ * PCollection<KV<byte[], Iterable<Mutation>>> data = ...;
+ * data.setCoder(HBaseIO.WRITE_CODER);
*
* data.apply("write",
* HBaseIO.write()
@@ -140,7 +146,7 @@ import org.slf4j.LoggerFactory;
* it can evolve or be different in some aspects, but the idea is that users can easily migrate
* from one to the other</p>.
*/
-@Experimental(Experimental.Kind.SOURCE_SINK)
+@Experimental
public class HBaseIO {
private static final Logger LOG = LoggerFactory.getLogger(HBaseIO.class);
@@ -539,7 +545,9 @@ public class HBaseIO {
*
* @see HBaseIO
*/
- public static class Write extends PTransform<PCollection<Mutation>, PDone> {
+ public static class Write
+ extends PTransform<PCollection<KV<byte[], Iterable<Mutation>>>, PDone> {
+
/**
* Returns a new {@link HBaseIO.Write} that will write to the HBase instance
* indicated by the given Configuration, and using any other specified customizations.
@@ -567,7 +575,7 @@ public class HBaseIO {
}
@Override
- public PDone expand(PCollection<Mutation> input) {
+ public PDone expand(PCollection<KV<byte[], Iterable<Mutation>>> input) {
input.apply(ParDo.of(new HBaseWriterFn(tableId, serializableConfiguration)));
return PDone.in(input.getPipeline());
}
@@ -605,7 +613,7 @@ public class HBaseIO {
private final String tableId;
private final SerializableConfiguration serializableConfiguration;
- private class HBaseWriterFn extends DoFn<Mutation, Void> {
+ private class HBaseWriterFn extends DoFn<KV<byte[], Iterable<Mutation>>, Void> {
public HBaseWriterFn(String tableId,
SerializableConfiguration serializableConfiguration) {
@@ -616,27 +624,31 @@ public class HBaseIO {
@Setup
public void setup() throws Exception {
- connection = ConnectionFactory.createConnection(serializableConfiguration.get());
- }
+ Configuration configuration = this.serializableConfiguration.get();
+ connection = ConnectionFactory.createConnection(configuration);
- @StartBundle
- public void startBundle(StartBundleContext c) throws IOException {
+ TableName tableName = TableName.valueOf(tableId);
BufferedMutatorParams params =
- new BufferedMutatorParams(TableName.valueOf(tableId));
+ new BufferedMutatorParams(tableName);
mutator = connection.getBufferedMutator(params);
+
recordsWritten = 0;
}
@ProcessElement
- public void processElement(ProcessContext c) throws Exception {
- mutator.mutate(c.element());
- ++recordsWritten;
+ public void processElement(ProcessContext ctx) throws Exception {
+ KV<byte[], Iterable<Mutation>> record = ctx.element();
+ List<Mutation> mutations = new ArrayList<>();
+ for (Mutation mutation : record.getValue()) {
+ mutations.add(mutation);
+ ++recordsWritten;
+ }
+ mutator.mutate(mutations);
}
@FinishBundle
public void finishBundle() throws Exception {
mutator.flush();
- LOG.debug("Wrote {} records", recordsWritten);
}
@Teardown
@@ -649,6 +661,7 @@ public class HBaseIO {
connection.close();
connection = null;
}
+ LOG.debug("Wrote {} records", recordsWritten);
}
@Override
@@ -665,4 +678,7 @@ public class HBaseIO {
private long recordsWritten;
}
}
+
+ public static final Coder<KV<byte[], Iterable<Mutation>>> WRITE_CODER =
+ KvCoder.of(ByteArrayCoder.of(), IterableCoder.of(HBaseMutationCoder.of()));
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseMutationCoder.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseMutationCoder.java b/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseMutationCoder.java
index ee83114..501fe09 100644
--- a/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseMutationCoder.java
+++ b/sdks/java/io/hbase/src/main/java/org/apache/beam/sdk/io/hbase/HBaseMutationCoder.java
@@ -21,12 +21,8 @@ import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.Serializable;
-import java.util.List;
import org.apache.beam.sdk.coders.AtomicCoder;
-import org.apache.beam.sdk.coders.CannotProvideCoderException;
import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.CoderProvider;
-import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.hadoop.hbase.client.Delete;
import org.apache.hadoop.hbase.client.Mutation;
import org.apache.hadoop.hbase.client.Put;
@@ -69,42 +65,4 @@ class HBaseMutationCoder extends AtomicCoder<Mutation> implements Serializable {
throw new IllegalArgumentException("Only Put and Delete are supported");
}
}
-
- /**
- * Returns a {@link CoderProvider} which uses the {@link HBaseMutationCoder} for
- * {@link Mutation mutations}.
- */
- static CoderProvider getCoderProvider() {
- return HBASE_MUTATION_CODER_PROVIDER;
- }
-
- private static final CoderProvider HBASE_MUTATION_CODER_PROVIDER =
- new HBaseMutationCoderProvider();
-
- /**
- * A {@link CoderProvider} for {@link Mutation mutations}.
- */
- private static class HBaseMutationCoderProvider extends CoderProvider {
- @Override
- public <T> Coder<T> coderFor(TypeDescriptor<T> typeDescriptor,
- List<? extends Coder<?>> componentCoders) throws CannotProvideCoderException {
- if (!typeDescriptor.isSubtypeOf(HBASE_MUTATION_TYPE_DESCRIPTOR)) {
- throw new CannotProvideCoderException(
- String.format(
- "Cannot provide %s because %s is not a subclass of %s",
- HBaseMutationCoder.class.getSimpleName(),
- typeDescriptor,
- Mutation.class.getName()));
- }
-
- try {
- return (Coder<T>) HBaseMutationCoder.of();
- } catch (IllegalArgumentException e) {
- throw new CannotProvideCoderException(e);
- }
- }
- }
-
- private static final TypeDescriptor<Mutation> HBASE_MUTATION_TYPE_DESCRIPTOR =
- new TypeDescriptor<Mutation>() {};
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/HBaseCoderProviderRegistrarTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/HBaseCoderProviderRegistrarTest.java b/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/HBaseCoderProviderRegistrarTest.java
deleted file mode 100644
index 5b2e138..0000000
--- a/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/HBaseCoderProviderRegistrarTest.java
+++ /dev/null
@@ -1,45 +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.sdk.io.hbase;
-
-import org.apache.beam.sdk.coders.CoderRegistry;
-import org.apache.hadoop.hbase.client.Delete;
-import org.apache.hadoop.hbase.client.Mutation;
-import org.apache.hadoop.hbase.client.Put;
-import org.apache.hadoop.hbase.client.Result;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-/**
- * Tests for {@link HBaseCoderProviderRegistrar}.
- */
-@RunWith(JUnit4.class)
-public class HBaseCoderProviderRegistrarTest {
- @Test
- public void testResultCoderIsRegistered() throws Exception {
- CoderRegistry.createDefault().getCoder(Result.class);
- }
-
- @Test
- public void testMutationCoderIsRegistered() throws Exception {
- CoderRegistry.createDefault().getCoder(Mutation.class);
- CoderRegistry.createDefault().getCoder(Put.class);
- CoderRegistry.createDefault().getCoder(Delete.class);
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/HBaseIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/HBaseIOTest.java b/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/HBaseIOTest.java
index 806a27f..4a06789 100644
--- a/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/HBaseIOTest.java
+++ b/sdks/java/io/hbase/src/test/java/org/apache/beam/sdk/io/hbase/HBaseIOTest.java
@@ -38,6 +38,7 @@ import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Count;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.commons.lang3.StringUtils;
import org.apache.hadoop.conf.Configuration;
@@ -46,7 +47,6 @@ import org.apache.hadoop.hbase.HBaseTestingUtility;
import org.apache.hadoop.hbase.HColumnDescriptor;
import org.apache.hadoop.hbase.HConstants;
import org.apache.hadoop.hbase.HTableDescriptor;
-import org.apache.hadoop.hbase.MiniHBaseCluster;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.BufferedMutator;
import org.apache.hadoop.hbase.client.Connection;
@@ -96,12 +96,7 @@ public class HBaseIOTest {
conf.setStrings("hbase.master.hostname", "localhost");
conf.setStrings("hbase.regionserver.hostname", "localhost");
htu = new HBaseTestingUtility(conf);
-
- // We don't use the full htu.startMiniCluster() to avoid starting unneeded HDFS/MR daemons
- htu.startMiniZKCluster();
- MiniHBaseCluster hbm = htu.startMiniHBaseCluster(1, 4);
- hbm.waitForActiveAndReadyMaster();
-
+ htu.startMiniCluster(1, 4);
admin = htu.getHBaseAdmin();
}
@@ -112,8 +107,7 @@ public class HBaseIOTest {
admin = null;
}
if (htu != null) {
- htu.shutdownMiniHBaseCluster();
- htu.shutdownMiniZKCluster();
+ htu.shutdownMiniCluster();
htu = null;
}
}
@@ -291,16 +285,15 @@ public class HBaseIOTest {
final String table = "table";
final String key = "key";
final String value = "value";
- final int numMutations = 100;
createTable(table);
- p.apply("multiple rows", Create.of(makeMutations(key, value, numMutations)))
- .apply("write", HBaseIO.write().withConfiguration(conf).withTableId(table));
+ p.apply("single row", Create.of(makeWrite(key, value)).withCoder(HBaseIO.WRITE_CODER))
+ .apply("write", HBaseIO.write().withConfiguration(conf).withTableId(table));
p.run().waitUntilFinish();
List<Result> results = readTable(table, new Scan());
- assertEquals(numMutations, results.size());
+ assertEquals(1, results.size());
}
/** Tests that when writing to a non-existent table, the write fails. */
@@ -308,8 +301,10 @@ public class HBaseIOTest {
public void testWritingFailsTableDoesNotExist() throws Exception {
final String table = "TEST-TABLE-DOES-NOT-EXIST";
- p.apply(Create.empty(HBaseMutationCoder.of()))
- .apply("write", HBaseIO.write().withConfiguration(conf).withTableId(table));
+ PCollection<KV<byte[], Iterable<Mutation>>> emptyInput =
+ p.apply(Create.empty(HBaseIO.WRITE_CODER));
+
+ emptyInput.apply("write", HBaseIO.write().withConfiguration(conf).withTableId(table));
// Exception will be thrown by write.validate() when write is applied.
thrown.expect(IllegalArgumentException.class);
@@ -324,8 +319,8 @@ public class HBaseIOTest {
final String key = "KEY";
createTable(table);
- p.apply(Create.of(makeBadMutation(key)))
- .apply(HBaseIO.write().withConfiguration(conf).withTableId(table));
+ p.apply(Create.of(makeBadWrite(key)).withCoder(HBaseIO.WRITE_CODER))
+ .apply(HBaseIO.write().withConfiguration(conf).withTableId(table));
thrown.expect(Pipeline.PipelineExecutionException.class);
thrown.expectCause(Matchers.<Throwable>instanceOf(IllegalArgumentException.class));
@@ -403,22 +398,26 @@ public class HBaseIOTest {
// Beam helper methods
/** Helper function to make a single row mutation to be written. */
- private static Iterable<Mutation> makeMutations(String key, String value, int numMutations) {
+ private static KV<byte[], Iterable<Mutation>> makeWrite(String key, String value) {
+ byte[] rowKey = key.getBytes(StandardCharsets.UTF_8);
List<Mutation> mutations = new ArrayList<>();
- for (int i = 0; i < numMutations; i++) {
- mutations.add(makeMutation(key + i, value));
- }
- return mutations;
+ mutations.add(makeMutation(key, value));
+ return KV.of(rowKey, (Iterable<Mutation>) mutations);
}
+
private static Mutation makeMutation(String key, String value) {
- return new Put(key.getBytes(StandardCharsets.UTF_8))
+ byte[] rowKey = key.getBytes(StandardCharsets.UTF_8);
+ return new Put(rowKey)
.addColumn(COLUMN_FAMILY, COLUMN_NAME, Bytes.toBytes(value))
.addColumn(COLUMN_FAMILY, COLUMN_EMAIL, Bytes.toBytes(value + "@email.com"));
}
- private static Mutation makeBadMutation(String key) {
- return new Put(key.getBytes());
+ private static KV<byte[], Iterable<Mutation>> makeBadWrite(String key) {
+ Put put = new Put(key.getBytes());
+ List<Mutation> mutations = new ArrayList<>();
+ mutations.add(put);
+ return KV.of(key.getBytes(StandardCharsets.UTF_8), (Iterable<Mutation>) mutations);
}
private void runReadTest(HBaseIO.Read read, List<Result> expected) {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hcatalog/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/hcatalog/pom.xml b/sdks/java/io/hcatalog/pom.xml
deleted file mode 100644
index 2aa661e..0000000
--- a/sdks/java/io/hcatalog/pom.xml
+++ /dev/null
@@ -1,175 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<!--
- 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.
--->
-<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
-
- <modelVersion>4.0.0</modelVersion>
-
- <parent>
- <groupId>org.apache.beam</groupId>
- <artifactId>beam-sdks-java-io-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
- <relativePath>../pom.xml</relativePath>
- </parent>
-
- <artifactId>beam-sdks-java-io-hcatalog</artifactId>
- <name>Apache Beam :: SDKs :: Java :: IO :: HCatalog</name>
- <description>IO to read and write for HCatalog source.</description>
-
- <properties>
- <hive.version>2.1.0</hive.version>
- <apache.commons.version>2.5</apache.commons.version>
- </properties>
-
- <build>
- <plugins>
- <plugin>
- <groupId>org.apache.maven.plugins</groupId>
- <artifactId>maven-shade-plugin</artifactId>
- <configuration>
- <createDependencyReducedPom>false</createDependencyReducedPom>
- </configuration>
- </plugin>
- </plugins>
- </build>
-
- <dependencies>
- <dependency>
- <groupId>org.apache.beam</groupId>
- <artifactId>beam-sdks-java-core</artifactId>
- </dependency>
-
- <dependency>
- <groupId>org.apache.beam</groupId>
- <artifactId>beam-sdks-java-io-hadoop-common</artifactId>
- </dependency>
-
- <dependency>
- <groupId>org.apache.hadoop</groupId>
- <artifactId>hadoop-common</artifactId>
- <exclusions>
- <!-- Fix build on JDK-9 -->
- <exclusion>
- <groupId>jdk.tools</groupId>
- <artifactId>jdk.tools</artifactId>
- </exclusion>
- </exclusions>
- </dependency>
-
- <dependency>
- <groupId>commons-io</groupId>
- <artifactId>commons-io</artifactId>
- <version>${apache.commons.version}</version>
- </dependency>
-
- <dependency>
- <groupId>org.slf4j</groupId>
- <artifactId>slf4j-api</artifactId>
- </dependency>
-
- <dependency>
- <groupId>com.google.guava</groupId>
- <artifactId>guava</artifactId>
- </dependency>
-
- <dependency>
- <groupId>com.google.code.findbugs</groupId>
- <artifactId>jsr305</artifactId>
- </dependency>
-
- <dependency>
- <groupId>org.apache.hive</groupId>
- <artifactId>hive-exec</artifactId>
- <version>${hive.version}</version>
- </dependency>
-
- <dependency>
- <groupId>com.google.auto.value</groupId>
- <artifactId>auto-value</artifactId>
- <scope>provided</scope>
- </dependency>
-
- <dependency>
- <groupId>org.apache.hive.hcatalog</groupId>
- <artifactId>hive-hcatalog-core</artifactId>
- <version>${hive.version}</version>
- <exclusions>
- <exclusion>
- <groupId>org.apache.hive</groupId>
- <artifactId>hive-exec</artifactId>
- </exclusion>
- <exclusion>
- <groupId>com.google.protobuf</groupId>
- <artifactId>protobuf-java</artifactId>
- </exclusion>
- <!-- Fix build on JDK-9 -->
- <exclusion>
- <groupId>jdk.tools</groupId>
- <artifactId>jdk.tools</artifactId>
- </exclusion>
- </exclusions>
- </dependency>
-
- <dependency>
- <groupId>org.apache.hive.hcatalog</groupId>
- <artifactId>hive-hcatalog-core</artifactId>
- <classifier>tests</classifier>
- <scope>test</scope>
- <version>${hive.version}</version>
- </dependency>
-
- <dependency>
- <groupId>junit</groupId>
- <artifactId>junit</artifactId>
- <scope>test</scope>
- </dependency>
-
- <dependency>
- <groupId>org.apache.hive</groupId>
- <artifactId>hive-exec</artifactId>
- <version>${hive.version}</version>
- <type>test-jar</type>
- <scope>test</scope>
- </dependency>
-
- <dependency>
- <groupId>org.apache.hive</groupId>
- <artifactId>hive-common</artifactId>
- <version>${hive.version}</version>
- <scope>test</scope>
- </dependency>
-
- <dependency>
- <groupId>org.apache.hive</groupId>
- <artifactId>hive-cli</artifactId>
- <version>${hive.version}</version>
- <scope>test</scope>
- </dependency>
-
- <dependency>
- <groupId>org.apache.beam</groupId>
- <artifactId>beam-runners-direct-java</artifactId>
- <scope>test</scope>
- </dependency>
-
- <dependency>
- <groupId>org.hamcrest</groupId>
- <artifactId>hamcrest-all</artifactId>
- <scope>test</scope>
- </dependency>
- </dependencies>
-</project>
\ No newline at end of file
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hcatalog/src/main/java/org/apache/beam/sdk/io/hcatalog/HCatalogIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hcatalog/src/main/java/org/apache/beam/sdk/io/hcatalog/HCatalogIO.java b/sdks/java/io/hcatalog/src/main/java/org/apache/beam/sdk/io/hcatalog/HCatalogIO.java
deleted file mode 100644
index 4199b80..0000000
--- a/sdks/java/io/hcatalog/src/main/java/org/apache/beam/sdk/io/hcatalog/HCatalogIO.java
+++ /dev/null
@@ -1,492 +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.sdk.io.hcatalog;
-
-import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Preconditions.checkNotNull;
-
-import com.google.auto.value.AutoValue;
-import com.google.common.annotations.VisibleForTesting;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.Map.Entry;
-import java.util.NoSuchElementException;
-import javax.annotation.Nullable;
-import org.apache.beam.sdk.annotations.Experimental;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.io.BoundedSource;
-import org.apache.beam.sdk.io.hadoop.WritableCoder;
-import org.apache.beam.sdk.options.PipelineOptions;
-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.display.DisplayData;
-import org.apache.beam.sdk.values.PBegin;
-import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PDone;
-import org.apache.hadoop.conf.Configuration;
-import org.apache.hadoop.hive.conf.HiveConf;
-import org.apache.hadoop.hive.metastore.IMetaStoreClient;
-import org.apache.hadoop.hive.ql.metadata.Table;
-import org.apache.hadoop.hive.ql.stats.StatsUtils;
-import org.apache.hive.hcatalog.common.HCatConstants;
-import org.apache.hive.hcatalog.common.HCatException;
-import org.apache.hive.hcatalog.common.HCatUtil;
-import org.apache.hive.hcatalog.data.DefaultHCatRecord;
-import org.apache.hive.hcatalog.data.HCatRecord;
-import org.apache.hive.hcatalog.data.transfer.DataTransferFactory;
-import org.apache.hive.hcatalog.data.transfer.HCatReader;
-import org.apache.hive.hcatalog.data.transfer.HCatWriter;
-import org.apache.hive.hcatalog.data.transfer.ReadEntity;
-import org.apache.hive.hcatalog.data.transfer.ReaderContext;
-import org.apache.hive.hcatalog.data.transfer.WriteEntity;
-import org.apache.hive.hcatalog.data.transfer.WriterContext;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/**
- * IO to read and write data using HCatalog.
- *
- * <h3>Reading using HCatalog</h3>
- *
- * <p>HCatalog source supports reading of HCatRecord from a HCatalog managed source, for eg. Hive.
- *
- * <p>To configure a HCatalog source, you must specify a metastore URI and a table name. Other
- * optional parameters are database & filter For instance:
- *
- * <pre>{@code
- * Map<String, String> configProperties = new HashMap<String, String>();
- * configProperties.put("hive.metastore.uris","thrift://metastore-host:port");
- *
- * pipeline
- * .apply(HCatalogIO.read()
- * .withConfigProperties(configProperties)
- * .withDatabase("default") //optional, assumes default if none specified
- * .withTable("employee")
- * .withFilter(filterString) //optional, may be specified if the table is partitioned
- * }</pre>
- *
- * <h3>Writing using HCatalog</h3>
- *
- * <p>HCatalog sink supports writing of HCatRecord to a HCatalog managed source, for eg. Hive.
- *
- * <p>To configure a HCatalog sink, you must specify a metastore URI and a table name. Other
- * optional parameters are database, partition & batchsize The destination table should exist
- * beforehand, the transform does not create a new table if it does not exist For instance:
- *
- * <pre>{@code
- * Map<String, String> configProperties = new HashMap<String, String>();
- * configProperties.put("hive.metastore.uris","thrift://metastore-host:port");
- *
- * pipeline
- * .apply(...)
- * .apply(HiveIO.write()
- * .withConfigProperties(configProperties)
- * .withDatabase("default") //optional, assumes default if none specified
- * .withTable("employee")
- * .withPartition(partitionValues) //optional, may be specified if the table is partitioned
- * .withBatchSize(1024L)) //optional, assumes a default batch size of 1024 if none specified
- * }</pre>
- */
-@Experimental(Experimental.Kind.SOURCE_SINK)
-public class HCatalogIO {
-
- private static final Logger LOG = LoggerFactory.getLogger(HCatalogIO.class);
-
- private static final long BATCH_SIZE = 1024L;
- private static final String DEFAULT_DATABASE = "default";
-
- /** Write data to Hive. */
- public static Write write() {
- return new AutoValue_HCatalogIO_Write.Builder().setBatchSize(BATCH_SIZE).build();
- }
-
- /** Read data from Hive. */
- public static Read read() {
- return new AutoValue_HCatalogIO_Read.Builder().setDatabase(DEFAULT_DATABASE).build();
- }
-
- private HCatalogIO() {}
-
- /** A {@link PTransform} to read data using HCatalog. */
- @VisibleForTesting
- @AutoValue
- public abstract static class Read extends PTransform<PBegin, PCollection<HCatRecord>> {
- @Nullable abstract Map<String, String> getConfigProperties();
- @Nullable abstract String getDatabase();
- @Nullable abstract String getTable();
- @Nullable abstract String getFilter();
- @Nullable abstract ReaderContext getContext();
- @Nullable abstract Integer getSplitId();
- abstract Builder toBuilder();
-
- @AutoValue.Builder
- abstract static class Builder {
- abstract Builder setConfigProperties(Map<String, String> configProperties);
- abstract Builder setDatabase(String database);
- abstract Builder setTable(String table);
- abstract Builder setFilter(String filter);
- abstract Builder setSplitId(Integer splitId);
- abstract Builder setContext(ReaderContext context);
- abstract Read build();
- }
-
- /** Sets the configuration properties like metastore URI. */
- public Read withConfigProperties(Map<String, String> configProperties) {
- return toBuilder().setConfigProperties(new HashMap<>(configProperties)).build();
- }
-
- /** Sets the database name. This is optional, assumes 'default' database if none specified */
- public Read withDatabase(String database) {
- return toBuilder().setDatabase(database).build();
- }
-
- /** Sets the table name to read from. */
- public Read withTable(String table) {
- return toBuilder().setTable(table).build();
- }
-
- /** Sets the filter details. This is optional, assumes none if not specified */
- public Read withFilter(String filter) {
- return toBuilder().setFilter(filter).build();
- }
-
- Read withSplitId(int splitId) {
- checkArgument(splitId >= 0, "Invalid split id-" + splitId);
- return toBuilder().setSplitId(splitId).build();
- }
-
- Read withContext(ReaderContext context) {
- return toBuilder().setContext(context).build();
- }
-
- @Override
- public PCollection<HCatRecord> expand(PBegin input) {
- return input.apply(org.apache.beam.sdk.io.Read.from(new BoundedHCatalogSource(this)));
- }
-
- @Override
- public void validate(PipelineOptions options) {
- checkNotNull(getTable(), "table");
- checkNotNull(getConfigProperties(), "configProperties");
- }
-
- @Override
- public void populateDisplayData(DisplayData.Builder builder) {
- super.populateDisplayData(builder);
- builder.add(DisplayData.item("configProperties", getConfigProperties().toString()));
- builder.add(DisplayData.item("table", getTable()));
- builder.addIfNotNull(DisplayData.item("database", getDatabase()));
- builder.addIfNotNull(DisplayData.item("filter", getFilter()));
- }
- }
-
- /** A HCatalog {@link BoundedSource} reading {@link HCatRecord} from a given instance. */
- @VisibleForTesting
- static class BoundedHCatalogSource extends BoundedSource<HCatRecord> {
- private final Read spec;
-
- BoundedHCatalogSource(Read spec) {
- this.spec = spec;
- }
-
- @Override
- @SuppressWarnings({"unchecked", "rawtypes"})
- public Coder<HCatRecord> getDefaultOutputCoder() {
- return (Coder) WritableCoder.of(DefaultHCatRecord.class);
- }
-
- @Override
- public void validate() {
- spec.validate(null);
- }
-
- @Override
- public void populateDisplayData(DisplayData.Builder builder) {
- spec.populateDisplayData(builder);
- }
-
- @Override
- public BoundedReader<HCatRecord> createReader(PipelineOptions options) {
- return new BoundedHCatalogReader(this);
- }
-
- /**
- * Returns the size of the table in bytes, does not take into consideration filter/partition
- * details passed, if any.
- */
- @Override
- public long getEstimatedSizeBytes(PipelineOptions pipelineOptions) throws Exception {
- Configuration conf = new Configuration();
- for (Entry<String, String> entry : spec.getConfigProperties().entrySet()) {
- conf.set(entry.getKey(), entry.getValue());
- }
- IMetaStoreClient client = null;
- try {
- HiveConf hiveConf = HCatUtil.getHiveConf(conf);
- client = HCatUtil.getHiveMetastoreClient(hiveConf);
- Table table = HCatUtil.getTable(client, spec.getDatabase(), spec.getTable());
- return StatsUtils.getFileSizeForTable(hiveConf, table);
- } finally {
- // IMetaStoreClient is not AutoCloseable, closing it manually
- if (client != null) {
- client.close();
- }
- }
- }
-
- /**
- * Calculates the 'desired' number of splits based on desiredBundleSizeBytes which is passed as
- * a hint to native API. Retrieves the actual splits generated by native API, which could be
- * different from the 'desired' split count calculated using desiredBundleSizeBytes
- */
- @Override
- public List<BoundedSource<HCatRecord>> split(
- long desiredBundleSizeBytes, PipelineOptions options) throws Exception {
- int desiredSplitCount = 1;
- long estimatedSizeBytes = getEstimatedSizeBytes(options);
- if (desiredBundleSizeBytes > 0 && estimatedSizeBytes > 0) {
- desiredSplitCount = (int) Math.ceil((double) estimatedSizeBytes / desiredBundleSizeBytes);
- }
- ReaderContext readerContext = getReaderContext(desiredSplitCount);
- //process the splits returned by native API
- //this could be different from 'desiredSplitCount' calculated above
- LOG.info(
- "Splitting into bundles of {} bytes: "
- + "estimated size {}, desired split count {}, actual split count {}",
- desiredBundleSizeBytes,
- estimatedSizeBytes,
- desiredSplitCount,
- readerContext.numSplits());
-
- List<BoundedSource<HCatRecord>> res = new ArrayList<>();
- for (int split = 0; split < readerContext.numSplits(); split++) {
- res.add(new BoundedHCatalogSource(spec.withContext(readerContext).withSplitId(split)));
- }
- return res;
- }
-
- private ReaderContext getReaderContext(long desiredSplitCount) throws HCatException {
- ReadEntity entity =
- new ReadEntity.Builder()
- .withDatabase(spec.getDatabase())
- .withTable(spec.getTable())
- .withFilter(spec.getFilter())
- .build();
- // pass the 'desired' split count as an hint to the API
- Map<String, String> configProps = new HashMap<>(spec.getConfigProperties());
- configProps.put(
- HCatConstants.HCAT_DESIRED_PARTITION_NUM_SPLITS, String.valueOf(desiredSplitCount));
- return DataTransferFactory.getHCatReader(entity, configProps).prepareRead();
- }
-
- static class BoundedHCatalogReader extends BoundedSource.BoundedReader<HCatRecord> {
- private final BoundedHCatalogSource source;
- private HCatRecord current;
- private Iterator<HCatRecord> hcatIterator;
-
- public BoundedHCatalogReader(BoundedHCatalogSource source) {
- this.source = source;
- }
-
- @Override
- public boolean start() throws HCatException {
- HCatReader reader =
- DataTransferFactory.getHCatReader(source.spec.getContext(), source.spec.getSplitId());
- hcatIterator = reader.read();
- return advance();
- }
-
- @Override
- public boolean advance() {
- if (hcatIterator.hasNext()) {
- current = hcatIterator.next();
- return true;
- } else {
- current = null;
- return false;
- }
- }
-
- @Override
- public BoundedHCatalogSource getCurrentSource() {
- return source;
- }
-
- @Override
- public HCatRecord getCurrent() {
- if (current == null) {
- throw new NoSuchElementException("Current element is null");
- }
- return current;
- }
-
- @Override
- public void close() {
- // nothing to close/release
- }
- }
- }
-
- /** A {@link PTransform} to write to a HCatalog managed source. */
- @AutoValue
- public abstract static class Write extends PTransform<PCollection<HCatRecord>, PDone> {
- @Nullable abstract Map<String, String> getConfigProperties();
- @Nullable abstract String getDatabase();
- @Nullable abstract String getTable();
- @Nullable abstract Map<String, String> getPartition();
- abstract long getBatchSize();
- abstract Builder toBuilder();
-
- @AutoValue.Builder
- abstract static class Builder {
- abstract Builder setConfigProperties(Map<String, String> configProperties);
- abstract Builder setDatabase(String database);
- abstract Builder setTable(String table);
- abstract Builder setPartition(Map<String, String> partition);
- abstract Builder setBatchSize(long batchSize);
- abstract Write build();
- }
-
- /** Sets the configuration properties like metastore URI. */
- public Write withConfigProperties(Map<String, String> configProperties) {
- return toBuilder().setConfigProperties(new HashMap<>(configProperties)).build();
- }
-
- /** Sets the database name. This is optional, assumes 'default' database if none specified */
- public Write withDatabase(String database) {
- return toBuilder().setDatabase(database).build();
- }
-
- /** Sets the table name to write to, the table should exist beforehand. */
- public Write withTable(String table) {
- return toBuilder().setTable(table).build();
- }
-
- /** Sets the partition details. */
- public Write withPartition(Map<String, String> partition) {
- return toBuilder().setPartition(partition).build();
- }
-
- /**
- * Sets batch size for the write operation. This is optional, assumes a default batch size of
- * 1024 if not set
- */
- public Write withBatchSize(long batchSize) {
- return toBuilder().setBatchSize(batchSize).build();
- }
-
- @Override
- public PDone expand(PCollection<HCatRecord> input) {
- input.apply(ParDo.of(new WriteFn(this)));
- return PDone.in(input.getPipeline());
- }
-
- @Override
- public void validate(PipelineOptions options) {
- checkNotNull(getConfigProperties(), "configProperties");
- checkNotNull(getTable(), "table");
- }
-
- private static class WriteFn extends DoFn<HCatRecord, Void> {
- private final Write spec;
- private WriterContext writerContext;
- private HCatWriter slaveWriter;
- private HCatWriter masterWriter;
- private List<HCatRecord> hCatRecordsBatch;
-
- public WriteFn(Write spec) {
- this.spec = spec;
- }
-
- @Override
- public void populateDisplayData(DisplayData.Builder builder) {
- super.populateDisplayData(builder);
- builder.addIfNotNull(DisplayData.item("database", spec.getDatabase()));
- builder.add(DisplayData.item("table", spec.getTable()));
- builder.addIfNotNull(DisplayData.item("partition", String.valueOf(spec.getPartition())));
- builder.add(DisplayData.item("configProperties", spec.getConfigProperties().toString()));
- builder.add(DisplayData.item("batchSize", spec.getBatchSize()));
- }
-
- @Setup
- public void initiateWrite() throws HCatException {
- WriteEntity entity =
- new WriteEntity.Builder()
- .withDatabase(spec.getDatabase())
- .withTable(spec.getTable())
- .withPartition(spec.getPartition())
- .build();
- masterWriter = DataTransferFactory.getHCatWriter(entity, spec.getConfigProperties());
- writerContext = masterWriter.prepareWrite();
- slaveWriter = DataTransferFactory.getHCatWriter(writerContext);
- }
-
- @StartBundle
- public void startBundle() {
- hCatRecordsBatch = new ArrayList<>();
- }
-
- @ProcessElement
- public void processElement(ProcessContext ctx) throws HCatException {
- hCatRecordsBatch.add(ctx.element());
- if (hCatRecordsBatch.size() >= spec.getBatchSize()) {
- flush();
- }
- }
-
- @FinishBundle
- public void finishBundle() throws HCatException {
- flush();
- }
-
- private void flush() throws HCatException {
- if (hCatRecordsBatch.isEmpty()) {
- return;
- }
- try {
- slaveWriter.write(hCatRecordsBatch.iterator());
- masterWriter.commit(writerContext);
- } catch (HCatException e) {
- LOG.error("Exception in flush - write/commit data to Hive", e);
- //abort on exception
- masterWriter.abort(writerContext);
- throw e;
- } finally {
- hCatRecordsBatch.clear();
- }
- }
-
- @Teardown
- public void tearDown() throws Exception {
- if (slaveWriter != null) {
- slaveWriter = null;
- }
- if (masterWriter != null) {
- masterWriter = null;
- }
- if (writerContext != null) {
- writerContext = null;
- }
- }
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hcatalog/src/main/java/org/apache/beam/sdk/io/hcatalog/package-info.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hcatalog/src/main/java/org/apache/beam/sdk/io/hcatalog/package-info.java b/sdks/java/io/hcatalog/src/main/java/org/apache/beam/sdk/io/hcatalog/package-info.java
deleted file mode 100644
index dff5bd1..0000000
--- a/sdks/java/io/hcatalog/src/main/java/org/apache/beam/sdk/io/hcatalog/package-info.java
+++ /dev/null
@@ -1,22 +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.
- */
-
-/**
- * Transforms for reading and writing using HCatalog.
- */
-package org.apache.beam.sdk.io.hcatalog;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hcatalog/src/test/java/org/apache/beam/sdk/io/hcatalog/EmbeddedMetastoreService.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hcatalog/src/test/java/org/apache/beam/sdk/io/hcatalog/EmbeddedMetastoreService.java b/sdks/java/io/hcatalog/src/test/java/org/apache/beam/sdk/io/hcatalog/EmbeddedMetastoreService.java
deleted file mode 100644
index 31e5b1c..0000000
--- a/sdks/java/io/hcatalog/src/test/java/org/apache/beam/sdk/io/hcatalog/EmbeddedMetastoreService.java
+++ /dev/null
@@ -1,87 +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.sdk.io.hcatalog;
-
-import static org.apache.hive.hcatalog.common.HCatUtil.makePathASafeFileName;
-
-import java.io.File;
-import java.io.IOException;
-import org.apache.commons.io.FileUtils;
-import org.apache.hadoop.hive.cli.CliSessionState;
-import org.apache.hadoop.hive.conf.HiveConf;
-import org.apache.hadoop.hive.ql.CommandNeedRetryException;
-import org.apache.hadoop.hive.ql.Driver;
-import org.apache.hadoop.hive.ql.session.SessionState;
-
-/**
- * Implementation of a light-weight embedded metastore. This class is a trimmed-down version of <a
- * href="https://github.com/apache/hive/blob/master
- * /hcatalog/core/src/test/java/org/apache/hive/hcatalog/mapreduce/HCatBaseTest.java">
- * https://github.com/apache/hive/blob/master/hcatalog/core/src/test/java/org/apache/hive/hcatalog/mapreduce
- * /HCatBaseTest.java </a>
- */
-final class EmbeddedMetastoreService implements AutoCloseable {
- private final Driver driver;
- private final HiveConf hiveConf;
- private final SessionState sessionState;
-
- EmbeddedMetastoreService(String baseDirPath) throws IOException {
- FileUtils.forceDeleteOnExit(new File(baseDirPath));
-
- String hiveDirPath = makePathASafeFileName(baseDirPath + "/hive");
- String testDataDirPath =
- makePathASafeFileName(
- hiveDirPath
- + "/data/"
- + EmbeddedMetastoreService.class.getCanonicalName()
- + System.currentTimeMillis());
- String testWarehouseDirPath = makePathASafeFileName(testDataDirPath + "/warehouse");
-
- hiveConf = new HiveConf(getClass());
- hiveConf.setVar(HiveConf.ConfVars.PREEXECHOOKS, "");
- hiveConf.setVar(HiveConf.ConfVars.POSTEXECHOOKS, "");
- hiveConf.setBoolVar(HiveConf.ConfVars.HIVE_SUPPORT_CONCURRENCY, false);
- hiveConf.setVar(HiveConf.ConfVars.METASTOREWAREHOUSE, testWarehouseDirPath);
- hiveConf.setBoolVar(HiveConf.ConfVars.HIVEOPTIMIZEMETADATAQUERIES, true);
- hiveConf.setVar(
- HiveConf.ConfVars.HIVE_AUTHORIZATION_MANAGER,
- "org.apache.hadoop.hive.ql.security.authorization.plugin.sqlstd."
- + "SQLStdHiveAuthorizerFactory");
- hiveConf.set("test.tmp.dir", hiveDirPath);
-
- System.setProperty("derby.stream.error.file", "/dev/null");
- driver = new Driver(hiveConf);
- sessionState = SessionState.start(new CliSessionState(hiveConf));
- }
-
- /** Executes the passed query on the embedded metastore service. */
- void executeQuery(String query) throws CommandNeedRetryException {
- driver.run(query);
- }
-
- /** Returns the HiveConf object for the embedded metastore. */
- HiveConf getHiveConf() {
- return hiveConf;
- }
-
- @Override
- public void close() throws Exception {
- driver.close();
- sessionState.close();
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hcatalog/src/test/java/org/apache/beam/sdk/io/hcatalog/HCatalogIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hcatalog/src/test/java/org/apache/beam/sdk/io/hcatalog/HCatalogIOTest.java b/sdks/java/io/hcatalog/src/test/java/org/apache/beam/sdk/io/hcatalog/HCatalogIOTest.java
deleted file mode 100644
index 91671a5..0000000
--- a/sdks/java/io/hcatalog/src/test/java/org/apache/beam/sdk/io/hcatalog/HCatalogIOTest.java
+++ /dev/null
@@ -1,277 +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.sdk.io.hcatalog;
-
-import static org.apache.beam.sdk.io.hcatalog.HCatalogIOTestUtils.TEST_DATABASE;
-import static org.apache.beam.sdk.io.hcatalog.HCatalogIOTestUtils.TEST_FILTER;
-import static org.apache.beam.sdk.io.hcatalog.HCatalogIOTestUtils.TEST_RECORDS_COUNT;
-import static org.apache.beam.sdk.io.hcatalog.HCatalogIOTestUtils.TEST_TABLE;
-import static org.apache.beam.sdk.io.hcatalog.HCatalogIOTestUtils.getConfigPropertiesAsMap;
-import static org.apache.beam.sdk.io.hcatalog.HCatalogIOTestUtils.getExpectedRecords;
-import static org.apache.beam.sdk.io.hcatalog.HCatalogIOTestUtils.getHCatRecords;
-import static org.apache.beam.sdk.io.hcatalog.HCatalogIOTestUtils.getReaderContext;
-import static org.apache.beam.sdk.io.hcatalog.HCatalogIOTestUtils.insertTestData;
-import static org.hamcrest.Matchers.containsInAnyOrder;
-import static org.hamcrest.Matchers.containsString;
-import static org.hamcrest.Matchers.isA;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertThat;
-import static org.junit.Assert.assertTrue;
-
-import java.io.IOException;
-import java.io.Serializable;
-import java.lang.annotation.ElementType;
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
-import java.lang.annotation.Target;
-import java.util.ArrayList;
-import java.util.List;
-import org.apache.beam.sdk.io.BoundedSource;
-import org.apache.beam.sdk.io.hcatalog.HCatalogIO.BoundedHCatalogSource;
-import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.sdk.testing.PAssert;
-import org.apache.beam.sdk.testing.SourceTestUtils;
-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.ParDo;
-import org.apache.beam.sdk.util.UserCodeException;
-import org.apache.beam.sdk.values.PCollection;
-import org.apache.hadoop.hive.metastore.api.NoSuchObjectException;
-import org.apache.hadoop.hive.ql.CommandNeedRetryException;
-import org.apache.hive.hcatalog.data.HCatRecord;
-import org.apache.hive.hcatalog.data.transfer.ReaderContext;
-import org.junit.AfterClass;
-import org.junit.BeforeClass;
-import org.junit.ClassRule;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.rules.ExpectedException;
-import org.junit.rules.TemporaryFolder;
-import org.junit.rules.TestRule;
-import org.junit.rules.TestWatcher;
-import org.junit.runner.Description;
-import org.junit.runners.model.Statement;
-
-/** Test for HCatalogIO. */
-public class HCatalogIOTest implements Serializable {
- private static final PipelineOptions OPTIONS = PipelineOptionsFactory.create();
-
- @ClassRule
- public static final TemporaryFolder TMP_FOLDER = new TemporaryFolder();
-
- @Rule public final transient TestPipeline defaultPipeline = TestPipeline.create();
-
- @Rule public final transient TestPipeline readAfterWritePipeline = TestPipeline.create();
-
- @Rule public transient ExpectedException thrown = ExpectedException.none();
-
- @Rule
- public final transient TestRule testDataSetupRule =
- new TestWatcher() {
- public Statement apply(final Statement base, final Description description) {
- return new Statement() {
- @Override
- public void evaluate() throws Throwable {
- if (description.getAnnotation(NeedsTestData.class) != null) {
- prepareTestData();
- } else if (description.getAnnotation(NeedsEmptyTestTables.class) != null) {
- reCreateTestTable();
- }
- base.evaluate();
- }
- };
- }
- };
-
- private static EmbeddedMetastoreService service;
-
- /** Use this annotation to setup complete test data(table populated with records). */
- @Retention(RetentionPolicy.RUNTIME)
- @Target({ElementType.METHOD})
- private @interface NeedsTestData {}
-
- /** Use this annotation to setup test tables alone(empty tables, no records are populated). */
- @Retention(RetentionPolicy.RUNTIME)
- @Target({ElementType.METHOD})
- private @interface NeedsEmptyTestTables {}
-
- @BeforeClass
- public static void setupEmbeddedMetastoreService () throws IOException {
- service = new EmbeddedMetastoreService(TMP_FOLDER.getRoot().getAbsolutePath());
- }
-
- @AfterClass
- public static void shutdownEmbeddedMetastoreService () throws Exception {
- service.executeQuery("drop table " + TEST_TABLE);
- service.close();
- }
-
- /** Perform end-to-end test of Write-then-Read operation. */
- @Test
- @NeedsEmptyTestTables
- public void testWriteThenReadSuccess() throws Exception {
- defaultPipeline
- .apply(Create.of(getHCatRecords(TEST_RECORDS_COUNT)))
- .apply(
- HCatalogIO.write()
- .withConfigProperties(getConfigPropertiesAsMap(service.getHiveConf()))
- .withDatabase(TEST_DATABASE)
- .withTable(TEST_TABLE)
- .withPartition(new java.util.HashMap<String, String>())
- .withBatchSize(512L));
- defaultPipeline.run();
-
- PCollection<String> output = readAfterWritePipeline
- .apply(
- HCatalogIO.read()
- .withConfigProperties(getConfigPropertiesAsMap(service.getHiveConf()))
- .withDatabase(TEST_DATABASE)
- .withTable(TEST_TABLE)
- .withFilter(TEST_FILTER))
- .apply(
- ParDo.of(
- new DoFn<HCatRecord, String>() {
- @ProcessElement
- public void processElement(ProcessContext c) {
- c.output(c.element().get(0).toString());
- }
- }));
- PAssert.that(output).containsInAnyOrder(getExpectedRecords(TEST_RECORDS_COUNT));
- readAfterWritePipeline.run();
- }
-
- /** Test of Write to a non-existent table. */
- @Test
- public void testWriteFailureTableDoesNotExist() throws Exception {
- thrown.expectCause(isA(UserCodeException.class));
- thrown.expectMessage(containsString("org.apache.hive.hcatalog.common.HCatException"));
- thrown.expectMessage(containsString("NoSuchObjectException"));
- defaultPipeline
- .apply(Create.of(getHCatRecords(TEST_RECORDS_COUNT)))
- .apply(
- HCatalogIO.write()
- .withConfigProperties(getConfigPropertiesAsMap(service.getHiveConf()))
- .withTable("myowntable"));
- defaultPipeline.run();
- }
-
- /** Test of Write without specifying a table. */
- @Test
- public void testWriteFailureValidationTable() throws Exception {
- thrown.expect(NullPointerException.class);
- thrown.expectMessage(containsString("table"));
- HCatalogIO.write()
- .withConfigProperties(getConfigPropertiesAsMap(service.getHiveConf()))
- .validate(null);
- }
-
- /** Test of Write without specifying configuration properties. */
- @Test
- public void testWriteFailureValidationConfigProp() throws Exception {
- thrown.expect(NullPointerException.class);
- thrown.expectMessage(containsString("configProperties"));
- HCatalogIO.write().withTable("myowntable").validate(null);
- }
-
- /** Test of Read from a non-existent table. */
- @Test
- public void testReadFailureTableDoesNotExist() throws Exception {
- defaultPipeline.apply(
- HCatalogIO.read()
- .withConfigProperties(getConfigPropertiesAsMap(service.getHiveConf()))
- .withTable("myowntable"));
- thrown.expectCause(isA(NoSuchObjectException.class));
- defaultPipeline.run();
- }
-
- /** Test of Read without specifying configuration properties. */
- @Test
- public void testReadFailureValidationConfig() throws Exception {
- thrown.expect(NullPointerException.class);
- thrown.expectMessage(containsString("configProperties"));
- HCatalogIO.read().withTable("myowntable").validate(null);
- }
-
- /** Test of Read without specifying a table. */
- @Test
- public void testReadFailureValidationTable() throws Exception {
- thrown.expect(NullPointerException.class);
- thrown.expectMessage(containsString("table"));
- HCatalogIO.read()
- .withConfigProperties(getConfigPropertiesAsMap(service.getHiveConf()))
- .validate(null);
- }
-
- /** Test of Read using SourceTestUtils.readFromSource(..). */
- @Test
- @NeedsTestData
- public void testReadFromSource() throws Exception {
- ReaderContext context = getReaderContext(getConfigPropertiesAsMap(service.getHiveConf()));
- HCatalogIO.Read spec =
- HCatalogIO.read()
- .withConfigProperties(getConfigPropertiesAsMap(service.getHiveConf()))
- .withContext(context)
- .withTable(TEST_TABLE);
-
- List<String> records = new ArrayList<>();
- for (int i = 0; i < context.numSplits(); i++) {
- BoundedHCatalogSource source = new BoundedHCatalogSource(spec.withSplitId(i));
- for (HCatRecord record : SourceTestUtils.readFromSource(source, OPTIONS)) {
- records.add(record.get(0).toString());
- }
- }
- assertThat(records, containsInAnyOrder(getExpectedRecords(TEST_RECORDS_COUNT).toArray()));
- }
-
- /** Test of Read using SourceTestUtils.assertSourcesEqualReferenceSource(..). */
- @Test
- @NeedsTestData
- public void testSourceEqualsSplits() throws Exception {
- final int numRows = 1500;
- final int numSamples = 10;
- final long bytesPerRow = 15;
- ReaderContext context = getReaderContext(getConfigPropertiesAsMap(service.getHiveConf()));
- HCatalogIO.Read spec =
- HCatalogIO.read()
- .withConfigProperties(getConfigPropertiesAsMap(service.getHiveConf()))
- .withContext(context)
- .withTable(TEST_TABLE);
-
- BoundedHCatalogSource source = new BoundedHCatalogSource(spec);
- List<BoundedSource<HCatRecord>> unSplitSource = source.split(-1, OPTIONS);
- assertEquals(1, unSplitSource.size());
-
- List<BoundedSource<HCatRecord>> splits =
- source.split(numRows * bytesPerRow / numSamples, OPTIONS);
- assertTrue(splits.size() >= 1);
-
- SourceTestUtils.assertSourcesEqualReferenceSource(unSplitSource.get(0), splits, OPTIONS);
- }
-
- private void reCreateTestTable() throws CommandNeedRetryException {
- service.executeQuery("drop table " + TEST_TABLE);
- service.executeQuery("create table " + TEST_TABLE + "(mycol1 string, mycol2 int)");
- }
-
- private void prepareTestData() throws Exception {
- reCreateTestTable();
- insertTestData(getConfigPropertiesAsMap(service.getHiveConf()));
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hcatalog/src/test/java/org/apache/beam/sdk/io/hcatalog/HCatalogIOTestUtils.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hcatalog/src/test/java/org/apache/beam/sdk/io/hcatalog/HCatalogIOTestUtils.java b/sdks/java/io/hcatalog/src/test/java/org/apache/beam/sdk/io/hcatalog/HCatalogIOTestUtils.java
deleted file mode 100644
index ae1eb50..0000000
--- a/sdks/java/io/hcatalog/src/test/java/org/apache/beam/sdk/io/hcatalog/HCatalogIOTestUtils.java
+++ /dev/null
@@ -1,108 +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.sdk.io.hcatalog;
-
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Map.Entry;
-import org.apache.hadoop.hive.conf.HiveConf;
-import org.apache.hive.hcatalog.common.HCatException;
-import org.apache.hive.hcatalog.data.DefaultHCatRecord;
-import org.apache.hive.hcatalog.data.HCatRecord;
-import org.apache.hive.hcatalog.data.transfer.DataTransferFactory;
-import org.apache.hive.hcatalog.data.transfer.ReadEntity;
-import org.apache.hive.hcatalog.data.transfer.ReaderContext;
-import org.apache.hive.hcatalog.data.transfer.WriteEntity;
-import org.apache.hive.hcatalog.data.transfer.WriterContext;
-
-/** Utility class for HCatalogIOTest. */
-class HCatalogIOTestUtils {
- static final String TEST_DATABASE = "default";
- static final String TEST_TABLE = "mytable";
- static final String TEST_FILTER = "myfilter";
- static final int TEST_RECORDS_COUNT = 1000;
-
- private static final ReadEntity READ_ENTITY =
- new ReadEntity.Builder().withTable(TEST_TABLE).build();
- private static final WriteEntity WRITE_ENTITY =
- new WriteEntity.Builder().withTable(TEST_TABLE).build();
-
- /** Returns a ReaderContext instance for the passed datastore config params. */
- static ReaderContext getReaderContext(Map<String, String> config) throws HCatException {
- return DataTransferFactory.getHCatReader(READ_ENTITY, config).prepareRead();
- }
-
- /** Returns a WriterContext instance for the passed datastore config params. */
- private static WriterContext getWriterContext(Map<String, String> config) throws HCatException {
- return DataTransferFactory.getHCatWriter(WRITE_ENTITY, config).prepareWrite();
- }
-
- /** Writes records to the table using the passed WriterContext. */
- private static void writeRecords(WriterContext context) throws HCatException {
- DataTransferFactory.getHCatWriter(context).write(getHCatRecords(TEST_RECORDS_COUNT).iterator());
- }
-
- /** Commits the pending writes to the database. */
- private static void commitRecords(Map<String, String> config, WriterContext context)
- throws IOException {
- DataTransferFactory.getHCatWriter(WRITE_ENTITY, config).commit(context);
- }
-
- /** Returns a list of strings containing 'expected' test data for verification. */
- static List<String> getExpectedRecords(int count) {
- List<String> expected = new ArrayList<>();
- for (int i = 0; i < count; i++) {
- expected.add("record " + i);
- }
- return expected;
- }
-
- /** Returns a list of HCatRecords of passed size. */
- static List<HCatRecord> getHCatRecords(int size) {
- List<HCatRecord> expected = new ArrayList<>();
- for (int i = 0; i < size; i++) {
- expected.add(toHCatRecord(i));
- }
- return expected;
- }
-
- /** Inserts data into test datastore. */
- static void insertTestData(Map<String, String> configMap) throws Exception {
- WriterContext cntxt = getWriterContext(configMap);
- writeRecords(cntxt);
- commitRecords(configMap, cntxt);
- }
-
- /** Returns config params for the test datastore as a Map. */
- static Map<String, String> getConfigPropertiesAsMap(HiveConf hiveConf) {
- Map<String, String> map = new HashMap<>();
- for (Entry<String, String> kv : hiveConf) {
- map.put(kv.getKey(), kv.getValue());
- }
- return map;
- }
-
- /** returns a DefaultHCatRecord instance for passed value. */
- private static DefaultHCatRecord toHCatRecord(int value) {
- return new DefaultHCatRecord(Arrays.<Object>asList("record " + value, value));
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/hcatalog/src/test/resources/hive-site.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/hcatalog/src/test/resources/hive-site.xml b/sdks/java/io/hcatalog/src/test/resources/hive-site.xml
deleted file mode 100644
index 5bb1496..0000000
--- a/sdks/java/io/hcatalog/src/test/resources/hive-site.xml
+++ /dev/null
@@ -1,301 +0,0 @@
-<?xml version="1.0"?>
-<?xml-stylesheet type="text/xsl" href="configuration.xsl"?>
-<!--
- 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.
--->
-<!-- This file is a copy of https://github.com/apache/hive/blob/master/data/conf/hive-site.xml used to support embedded Hive metastore-->
-<configuration>
-
-<property>
- <name>hive.in.test</name>
- <value>true</value>
- <description>Internal marker for test. Used for masking env-dependent values</description>
-</property>
-
-<!-- Hive Configuration can either be stored in this file or in the hadoop configuration files -->
-<!-- that are implied by Hadoop setup variables. -->
-<!-- Aside from Hadoop setup variables - this file is provided as a convenience so that Hive -->
-<!-- users do not have to edit hadoop configuration files (that may be managed as a centralized -->
-<!-- resource). -->
-
-<!-- Hive Execution Parameters -->
-<property>
- <name>hadoop.tmp.dir</name>
- <value>${test.tmp.dir}/hadoop-tmp</value>
- <description>A base for other temporary directories.</description>
-</property>
-
-<!--
-<property>
- <name>hive.exec.reducers.max</name>
- <value>1</value>
- <description>maximum number of reducers</description>
-</property>
--->
-
-<property>
- <name>hive.exec.scratchdir</name>
- <value>${test.tmp.dir}/scratchdir</value>
- <description>Scratch space for Hive jobs</description>
-</property>
-
-<property>
- <name>hive.exec.local.scratchdir</name>
- <value>${test.tmp.dir}/localscratchdir/</value>
- <description>Local scratch space for Hive jobs</description>
-</property>
-
-<property>
- <name>datanucleus.schema.autoCreateAll</name>
- <value>true</value>
-</property>
-
-<property>
- <name>javax.jdo.option.ConnectionURL</name>
- <value>jdbc:derby:;databaseName=${test.tmp.dir}/junit_metastore_db;create=true</value>
-</property>
-
-<property>
- <name>javax.jdo.option.ConnectionDriverName</name>
- <value>org.apache.derby.jdbc.EmbeddedDriver</value>
-</property>
-
-<property>
- <name>javax.jdo.option.ConnectionUserName</name>
- <value>APP</value>
-</property>
-
-<property>
- <name>javax.jdo.option.ConnectionPassword</name>
- <value>mine</value>
-</property>
-
-<property>
- <!-- this should eventually be deprecated since the metastore should supply this -->
- <name>hive.metastore.warehouse.dir</name>
- <value>${test.warehouse.dir}</value>
- <description></description>
-</property>
-
-<property>
- <name>hive.metastore.metadb.dir</name>
- <value>file://${test.tmp.dir}/metadb/</value>
- <description>
- Required by metastore server or if the uris argument below is not supplied
- </description>
-</property>
-
-<property>
- <name>test.log.dir</name>
- <value>${test.tmp.dir}/log/</value>
- <description></description>
-</property>
-
-<property>
- <name>test.data.files</name>
- <value>${hive.root}/data/files</value>
- <description></description>
-</property>
-
-<property>
- <name>test.data.scripts</name>
- <value>${hive.root}/data/scripts</value>
- <description></description>
-</property>
-
-<property>
- <name>hive.jar.path</name>
- <value>${maven.local.repository}/org/apache/hive/hive-exec/${hive.version}/hive-exec-${hive.version}.jar</value>
- <description></description>
-</property>
-
-<property>
- <name>hive.metastore.rawstore.impl</name>
- <value>org.apache.hadoop.hive.metastore.ObjectStore</value>
- <description>Name of the class that implements org.apache.hadoop.hive.metastore.rawstore interface. This class is used to store and retrieval of raw metadata objects such as table, database</description>
-</property>
-
-<property>
- <name>hive.querylog.location</name>
- <value>${test.tmp.dir}/tmp</value>
- <description>Location of the structured hive logs</description>
-</property>
-
-<property>
- <name>hive.exec.pre.hooks</name>
- <value>org.apache.hadoop.hive.ql.hooks.PreExecutePrinter, org.apache.hadoop.hive.ql.hooks.EnforceReadOnlyTables</value>
- <description>Pre Execute Hook for Tests</description>
-</property>
-
-<property>
- <name>hive.exec.post.hooks</name>
- <value>org.apache.hadoop.hive.ql.hooks.PostExecutePrinter</value>
- <description>Post Execute Hook for Tests</description>
-</property>
-
-<property>
- <name>hive.support.concurrency</name>
- <value>true</value>
- <description>Whether hive supports concurrency or not. A zookeeper instance must be up and running for the default hive lock manager to support read-write locks.</description>
-</property>
-
-<property>
- <key>hive.unlock.numretries</key>
- <value>2</value>
- <description>The number of times you want to retry to do one unlock</description>
-</property>
-
-<property>
- <key>hive.lock.sleep.between.retries</key>
- <value>2</value>
- <description>The sleep time (in seconds) between various retries</description>
-</property>
-
-
-<property>
- <name>fs.pfile.impl</name>
- <value>org.apache.hadoop.fs.ProxyLocalFileSystem</value>
- <description>A proxy for local file system used for cross file system testing</description>
-</property>
-
-<property>
- <name>hive.exec.mode.local.auto</name>
- <value>false</value>
- <description>
- Let hive determine whether to run in local mode automatically
- Disabling this for tests so that minimr is not affected
- </description>
-</property>
-
-<property>
- <name>hive.auto.convert.join</name>
- <value>false</value>
- <description>Whether Hive enable the optimization about converting common join into mapjoin based on the input file size</description>
-</property>
-
-<property>
- <name>hive.ignore.mapjoin.hint</name>
- <value>false</value>
- <description>Whether Hive ignores the mapjoin hint</description>
-</property>
-
-<property>
- <name>hive.input.format</name>
- <value>org.apache.hadoop.hive.ql.io.CombineHiveInputFormat</value>
- <description>The default input format, if it is not specified, the system assigns it. It is set to HiveInputFormat for hadoop versions 17, 18 and 19, whereas it is set to CombineHiveInputFormat for hadoop 20. The user can always overwrite it - if there is a bug in CombineHiveInputFormat, it can always be manually set to HiveInputFormat. </description>
-</property>
-
-<property>
- <name>hive.default.rcfile.serde</name>
- <value>org.apache.hadoop.hive.serde2.columnar.ColumnarSerDe</value>
- <description>The default SerDe hive will use for the rcfile format</description>
-</property>
-
-<property>
- <name>hive.stats.key.prefix.reserve.length</name>
- <value>0</value>
-</property>
-
-<property>
- <name>hive.conf.restricted.list</name>
- <value>dummy.config.value</value>
- <description>Using dummy config value above because you cannot override config with empty value</description>
-</property>
-
-<property>
- <name>hive.exec.submit.local.task.via.child</name>
- <value>false</value>
-</property>
-
-
-<property>
- <name>hive.dummyparam.test.server.specific.config.override</name>
- <value>from.hive-site.xml</value>
- <description>Using dummy param to test server specific configuration</description>
-</property>
-
-<property>
- <name>hive.dummyparam.test.server.specific.config.hivesite</name>
- <value>from.hive-site.xml</value>
- <description>Using dummy param to test server specific configuration</description>
-</property>
-
-<property>
- <name>test.var.hiveconf.property</name>
- <value>${hive.exec.default.partition.name}</value>
- <description>Test hiveconf property substitution</description>
-</property>
-
-<property>
- <name>test.property1</name>
- <value>value1</value>
- <description>Test property defined in hive-site.xml only</description>
-</property>
-
-<property>
- <name>hive.test.dummystats.aggregator</name>
- <value>value2</value>
-</property>
-
-<property>
- <name>hive.fetch.task.conversion</name>
- <value>minimal</value>
-</property>
-
-<property>
- <name>hive.users.in.admin.role</name>
- <value>hive_admin_user</value>
-</property>
-
-<property>
- <name>hive.llap.io.cache.orc.size</name>
- <value>8388608</value>
-</property>
-
-<property>
- <name>hive.llap.io.cache.orc.arena.size</name>
- <value>8388608</value>
-</property>
-
-<property>
- <name>hive.llap.io.cache.orc.alloc.max</name>
- <value>2097152</value>
-</property>
-
-
-<property>
- <name>hive.llap.io.cache.orc.alloc.min</name>
- <value>32768</value>
-</property>
-
-<property>
- <name>hive.llap.cache.allow.synthetic.fileid</name>
- <value>true</value>
-</property>
-
-<property>
- <name>hive.llap.io.use.lrfu</name>
- <value>true</value>
-</property>
-
-
-<property>
- <name>hive.llap.io.allocator.direct</name>
- <value>false</value>
-</property>
-
-
-</configuration>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/jdbc/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/jdbc/pom.xml b/sdks/java/io/jdbc/pom.xml
index 050fc6a..17c26a0 100644
--- a/sdks/java/io/jdbc/pom.xml
+++ b/sdks/java/io/jdbc/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
@@ -49,11 +49,13 @@
<dependency>
<groupId>org.apache.spark</groupId>
<artifactId>spark-streaming_2.10</artifactId>
+ <version>${spark.version}</version>
<scope>runtime</scope>
</dependency>
<dependency>
<groupId>org.apache.spark</groupId>
<artifactId>spark-core_2.10</artifactId>
+ <version>${spark.version}</version>
<scope>runtime</scope>
<exclusions>
<exclusion>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java b/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java
index bf73dbe..8092da6 100644
--- a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java
+++ b/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java
@@ -133,7 +133,7 @@ import org.apache.commons.dbcp2.BasicDataSource;
* Consider using <a href="https://en.wikipedia.org/wiki/Merge_(SQL)">MERGE ("upsert")
* statements</a> supported by your database instead.
*/
-@Experimental(Experimental.Kind.SOURCE_SINK)
+@Experimental
public class JdbcIO {
/**
* Read data from a JDBC datasource.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/jms/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/jms/pom.xml b/sdks/java/io/jms/pom.xml
index c2074af..58009a1 100644
--- a/sdks/java/io/jms/pom.xml
+++ b/sdks/java/io/jms/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java
index f8cba5e..c5e5150 100644
--- a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java
+++ b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java
@@ -98,7 +98,7 @@ import org.slf4j.LoggerFactory;
*
* }</pre>
*/
-@Experimental(Experimental.Kind.SOURCE_SINK)
+@Experimental
public class JmsIO {
private static final Logger LOG = LoggerFactory.getLogger(JmsIO.class);
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kafka/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/pom.xml b/sdks/java/io/kafka/pom.xml
index 1256c46..29350cc 100644
--- a/sdks/java/io/kafka/pom.xml
+++ b/sdks/java/io/kafka/pom.xml
@@ -21,7 +21,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-io-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index 026313a..4d2a358 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -49,11 +49,9 @@ import java.util.Random;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
-import java.util.concurrent.Future;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import javax.annotation.Nullable;
import org.apache.beam.sdk.annotations.Experimental;
@@ -237,7 +235,7 @@ import org.slf4j.LoggerFactory;
* Note that {@link KafkaRecord#getTimestamp()} reflects timestamp provided by Kafka if any,
* otherwise it is set to processing time.
*/
-@Experimental(Experimental.Kind.SOURCE_SINK)
+@Experimental
public class KafkaIO {
/**
@@ -906,22 +904,6 @@ public class KafkaIO {
return name;
}
- // Maintains approximate average over last 1000 elements
- private static class MovingAvg {
- private static final int MOVING_AVG_WINDOW = 1000;
- private double avg = 0;
- private long numUpdates = 0;
-
- void update(double quantity) {
- numUpdates++;
- avg += (quantity - avg) / Math.min(MOVING_AVG_WINDOW, numUpdates);
- }
-
- double get() {
- return avg;
- }
- }
-
// maintains state of each assigned partition (buffered records, consumed offset, etc)
private static class PartitionState {
private final TopicPartition topicPartition;
@@ -929,8 +911,9 @@ public class KafkaIO {
private long latestOffset;
private Iterator<ConsumerRecord<byte[], byte[]>> recordIter = Collections.emptyIterator();
- private MovingAvg avgRecordSize = new MovingAvg();
- private MovingAvg avgOffsetGap = new MovingAvg(); // > 0 only when log compaction is enabled.
+ // simple moving average for size of each record in bytes
+ private double avgRecordSize = 0;
+ private static final int movingAvgWindow = 1000; // very roughly avg of last 1000 elements
PartitionState(TopicPartition partition, long nextOffset) {
this.topicPartition = partition;
@@ -938,13 +921,17 @@ public class KafkaIO {
this.latestOffset = UNINITIALIZED_OFFSET;
}
- // Update consumedOffset, avgRecordSize, and avgOffsetGap
- void recordConsumed(long offset, int size, long offsetGap) {
+ // update consumedOffset and avgRecordSize
+ void recordConsumed(long offset, int size) {
nextOffset = offset + 1;
- // This is always updated from single thread. Probably not worth making atomic.
- avgRecordSize.update(size);
- avgOffsetGap.update(offsetGap);
+ // this is always updated from single thread. probably not worth making it an AtomicDouble
+ if (avgRecordSize <= 0) {
+ avgRecordSize = size;
+ } else {
+ // initially, first record heavily contributes to average.
+ avgRecordSize += ((size - avgRecordSize) / movingAvgWindow);
+ }
}
synchronized void setLatestOffset(long latestOffset) {
@@ -957,15 +944,14 @@ public class KafkaIO {
if (backlogMessageCount == UnboundedReader.BACKLOG_UNKNOWN) {
return UnboundedReader.BACKLOG_UNKNOWN;
}
- return (long) (backlogMessageCount * avgRecordSize.get());
+ return (long) (backlogMessageCount * avgRecordSize);
}
synchronized long backlogMessageCount() {
if (latestOffset < 0 || nextOffset < 0) {
return UnboundedReader.BACKLOG_UNKNOWN;
}
- double remaining = (latestOffset - nextOffset) / (1 + avgOffsetGap.get());
- return Math.max(0, (long) Math.ceil(remaining));
+ return Math.max(0, (latestOffset - nextOffset));
}
}
@@ -1063,32 +1049,8 @@ public class KafkaIO {
curBatch = Iterators.cycle(nonEmpty);
}
- private void setupInitialOffset(PartitionState pState) {
- Read<K, V> spec = source.spec;
-
- if (pState.nextOffset != UNINITIALIZED_OFFSET) {
- consumer.seek(pState.topicPartition, pState.nextOffset);
- } else {
- // nextOffset is unininitialized here, meaning start reading from latest record as of now
- // ('latest' is the default, and is configurable) or 'look up offset by startReadTime.
- // Remember the current position without waiting until the first record is read. This
- // ensures checkpoint is accurate even if the reader is closed before reading any records.
- Instant startReadTime = spec.getStartReadTime();
- if (startReadTime != null) {
- pState.nextOffset =
- consumerSpEL.offsetForTime(consumer, pState.topicPartition, spec.getStartReadTime());
- consumer.seek(pState.topicPartition, pState.nextOffset);
- } else {
- pState.nextOffset = consumer.position(pState.topicPartition);
- }
- }
- }
-
@Override
public boolean start() throws IOException {
- final int defaultPartitionInitTimeout = 60 * 1000;
- final int kafkaRequestTimeoutMultiple = 2;
-
Read<K, V> spec = source.spec;
consumer = spec.getConsumerFactoryFn().apply(spec.getConsumerConfig());
consumerSpEL.evaluateAssign(consumer, spec.getTopicPartitions());
@@ -1103,38 +1065,25 @@ public class KafkaIO {
keyDeserializerInstance.configure(spec.getConsumerConfig(), true);
valueDeserializerInstance.configure(spec.getConsumerConfig(), false);
- // Seek to start offset for each partition. This is the first interaction with the server.
- // Unfortunately it can block forever in case of network issues like incorrect ACLs.
- // Initialize partition in a separate thread and cancel it if takes longer than a minute.
- for (final PartitionState pState : partitionStates) {
- Future<?> future = consumerPollThread.submit(new Runnable() {
- public void run() {
- setupInitialOffset(pState);
+ for (PartitionState p : partitionStates) {
+ if (p.nextOffset != UNINITIALIZED_OFFSET) {
+ consumer.seek(p.topicPartition, p.nextOffset);
+ } else {
+ // nextOffset is unininitialized here, meaning start reading from latest record as of now
+ // ('latest' is the default, and is configurable) or 'look up offset by startReadTime.
+ // Remember the current position without waiting until the first record is read. This
+ // ensures checkpoint is accurate even if the reader is closed before reading any records.
+ Instant startReadTime = spec.getStartReadTime();
+ if (startReadTime != null) {
+ p.nextOffset =
+ consumerSpEL.offsetForTime(consumer, p.topicPartition, spec.getStartReadTime());
+ consumer.seek(p.topicPartition, p.nextOffset);
+ } else {
+ p.nextOffset = consumer.position(p.topicPartition);
}
- });
-
- try {
- // Timeout : 1 minute OR 2 * Kafka consumer request timeout if it is set.
- Integer reqTimeout = (Integer) source.spec.getConsumerConfig().get(
- ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG);
- future.get(reqTimeout != null ? kafkaRequestTimeoutMultiple * reqTimeout
- : defaultPartitionInitTimeout,
- TimeUnit.MILLISECONDS);
- } catch (TimeoutException e) {
- consumer.wakeup(); // This unblocks consumer stuck on network I/O.
- // Likely reason : Kafka servers are configured to advertise internal ips, but
- // those ips are not accessible from workers outside.
- String msg = String.format(
- "%s: Timeout while initializing partition '%s'. "
- + "Kafka client may not be able to connect to servers.",
- this, pState.topicPartition);
- LOG.error("{}", msg);
- throw new IOException(msg);
- } catch (Exception e) {
- throw new IOException(e);
}
- LOG.info("{}: reading from {} starting at offset {}",
- name, pState.topicPartition, pState.nextOffset);
+
+ LOG.info("{}: reading from {} starting at offset {}", name, p.topicPartition, p.nextOffset);
}
// Start consumer read loop.
@@ -1205,11 +1154,14 @@ public class KafkaIO {
continue;
}
- long offsetGap = offset - expected; // could be > 0 when Kafka log compaction is enabled.
+ // sanity check
+ if (offset != expected) {
+ LOG.warn("{}: gap in offsets for {} at {}. {} records missing.",
+ this, pState.topicPartition, expected, offset - expected);
+ }
if (curRecord == null) {
LOG.info("{}: first record offset {}", name, offset);
- offsetGap = 0;
}
curRecord = null; // user coders below might throw.
@@ -1230,7 +1182,7 @@ public class KafkaIO {
int recordSize = (rawRecord.key() == null ? 0 : rawRecord.key().length)
+ (rawRecord.value() == null ? 0 : rawRecord.value().length);
- pState.recordConsumed(offset, recordSize, offsetGap);
+ pState.recordConsumed(offset, recordSize);
bytesRead.inc(recordSize);
bytesReadBySplit.inc(recordSize);
return true;
@@ -1368,12 +1320,8 @@ public class KafkaIO {
// might block to enqueue right after availableRecordsQueue.poll() below.
while (!isShutdown) {
- if (consumer != null) {
- consumer.wakeup();
- }
- if (offsetConsumer != null) {
- offsetConsumer.wakeup();
- }
+ consumer.wakeup();
+ offsetConsumer.wakeup();
availableRecordsQueue.poll(); // drain unread batch, this unblocks consumer thread.
try {
isShutdown = consumerPollThread.awaitTermination(10, TimeUnit.SECONDS)
[20/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/CreateDataflowView.java
----------------------------------------------------------------------
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/CreateDataflowView.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/CreateDataflowView.java
index caad7f8..e7542cb 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/CreateDataflowView.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/CreateDataflowView.java
@@ -24,7 +24,7 @@ import org.apache.beam.sdk.values.PCollectionView;
/** A {@link DataflowRunner} marker class for creating a {@link PCollectionView}. */
public class CreateDataflowView<ElemT, ViewT>
- extends PTransform<PCollection<ElemT>, PCollection<ElemT>> {
+ extends PTransform<PCollection<ElemT>, PCollectionView<ViewT>> {
public static <ElemT, ViewT> CreateDataflowView<ElemT, ViewT> of(PCollectionView<ViewT> view) {
return new CreateDataflowView<>(view);
}
@@ -36,10 +36,8 @@ public class CreateDataflowView<ElemT, ViewT>
}
@Override
- public PCollection<ElemT> expand(PCollection<ElemT> input) {
- return PCollection.<ElemT>createPrimitiveOutputInternal(
- input.getPipeline(), input.getWindowingStrategy(), input.isBounded())
- .setCoder(input.getCoder());
+ public PCollectionView<ViewT> expand(PCollection<ElemT> input) {
+ return view;
}
public PCollectionView<ViewT> getView() {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslator.java
----------------------------------------------------------------------
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslator.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslator.java
index f1783de..8eaf61b 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslator.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslator.java
@@ -56,8 +56,6 @@ import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;
import javax.annotation.Nullable;
-import org.apache.beam.runners.core.construction.SplittableParDo;
-import org.apache.beam.runners.core.construction.TransformInputs;
import org.apache.beam.runners.core.construction.WindowingStrategyTranslation;
import org.apache.beam.runners.dataflow.BatchViewOverrides.GroupByKeyAndSortValuesOnly;
import org.apache.beam.runners.dataflow.DataflowRunner.CombineGroupedValues;
@@ -397,9 +395,7 @@ public class DataflowPipelineTranslator {
@Override
public <InputT extends PValue> InputT getInput(PTransform<InputT, ?> transform) {
- return (InputT)
- Iterables.getOnlyElement(
- TransformInputs.nonAdditionalInputs(getCurrentTransform(transform)));
+ return (InputT) Iterables.getOnlyElement(getInputs(transform).values());
}
@Override
@@ -444,14 +440,6 @@ public class DataflowPipelineTranslator {
public void visitValue(PValue value, TransformHierarchy.Node producer) {
LOG.debug("Checking translation of {}", value);
// Primitive transforms are the only ones assigned step names.
- if (producer.getTransform() instanceof CreateDataflowView) {
- // CreateDataflowView produces a dummy output (as it must be a primitive transform) but
- // in the Dataflow Job graph produces only the view and not the output PCollection.
- asOutputReference(
- ((CreateDataflowView) producer.getTransform()).getView(),
- producer.toAppliedPTransform(getPipeline()));
- return;
- }
asOutputReference(value, producer.toAppliedPTransform(getPipeline()));
}
@@ -477,7 +465,6 @@ public class DataflowPipelineTranslator {
StepTranslator stepContext = new StepTranslator(this, step);
stepContext.addInput(PropertyNames.USER_NAME, getFullName(transform));
stepContext.addDisplayData(step, stepName, transform);
- LOG.info("Adding {} as step {}", getCurrentTransform(transform).getFullName(), stepName);
return stepContext;
}
@@ -690,7 +677,7 @@ public class DataflowPipelineTranslator {
context.addStep(transform, "CollectionToSingleton");
PCollection<ElemT> input = context.getInput(transform);
stepContext.addInput(PropertyNames.PARALLEL_INPUT, input);
- stepContext.addCollectionToSingletonOutput(input, transform.getView());
+ stepContext.addCollectionToSingletonOutput(input, context.getOutput(transform));
}
});
@@ -793,7 +780,6 @@ public class DataflowPipelineTranslator {
context.getPipelineOptions().as(StreamingOptions.class).isStreaming();
boolean disallowCombinerLifting =
!windowingStrategy.getWindowFn().isNonMerging()
- || !windowingStrategy.getWindowFn().assignsToOneWindow()
|| (isStreaming && !transform.fewKeys())
// TODO: Allow combiner lifting on the non-default trigger, as appropriate.
|| !(windowingStrategy.getTrigger() instanceof DefaultTrigger);
@@ -888,45 +874,6 @@ public class DataflowPipelineTranslator {
// IO Translation.
registerTransformTranslator(Read.Bounded.class, new ReadTranslator());
-
- ///////////////////////////////////////////////////////////////////////////
- // Splittable DoFn translation.
-
- registerTransformTranslator(
- SplittableParDo.ProcessKeyedElements.class,
- new TransformTranslator<SplittableParDo.ProcessKeyedElements>() {
- @Override
- public void translate(
- SplittableParDo.ProcessKeyedElements transform, TranslationContext context) {
- translateTyped(transform, context);
- }
-
- private <InputT, OutputT, RestrictionT> void translateTyped(
- SplittableParDo.ProcessKeyedElements<InputT, OutputT, RestrictionT> transform,
- TranslationContext context) {
- StepTranslationContext stepContext =
- context.addStep(transform, "SplittableProcessKeyed");
-
- translateInputs(
- stepContext, context.getInput(transform), transform.getSideInputs(), context);
- BiMap<Long, TupleTag<?>> outputMap =
- translateOutputs(context.getOutputs(transform), stepContext);
- stepContext.addInput(
- PropertyNames.SERIALIZED_FN,
- byteArrayToJsonString(
- serializeToByteArray(
- DoFnInfo.forFn(
- transform.getFn(),
- transform.getInputWindowingStrategy(),
- transform.getSideInputs(),
- transform.getElementCoder(),
- outputMap.inverse().get(transform.getMainOutputTag()),
- outputMap))));
- stepContext.addInput(
- PropertyNames.RESTRICTION_CODER,
- CloudObjects.asCloudObject(transform.getRestrictionCoder()));
- }
- });
}
private static void translateInputs(
@@ -973,11 +920,6 @@ public class DataflowPipelineTranslator {
fn));
}
- if (signature.usesState() || signature.usesTimers()) {
- DataflowRunner.verifyStateSupported(fn);
- DataflowRunner.verifyStateSupportForWindowingStrategy(windowingStrategy);
- }
-
stepContext.addInput(PropertyNames.USER_FN, fn.getClass().getName());
stepContext.addInput(
PropertyNames.SERIALIZED_FN,
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java
----------------------------------------------------------------------
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java
index 8935759..3e7c8ce 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java
@@ -67,12 +67,11 @@ import org.apache.beam.runners.core.construction.ReplacementOutputs;
import org.apache.beam.runners.core.construction.SingleInputOutputOverrideFactory;
import org.apache.beam.runners.core.construction.UnboundedReadFromBoundedSource;
import org.apache.beam.runners.core.construction.UnconsumedReads;
-import org.apache.beam.runners.core.construction.WriteFilesTranslation;
+import org.apache.beam.runners.dataflow.BatchViewOverrides.BatchCombineGloballyAsSingletonViewFactory;
import org.apache.beam.runners.dataflow.DataflowPipelineTranslator.JobSpecification;
import org.apache.beam.runners.dataflow.StreamingViewOverrides.StreamingCreatePCollectionViewFactory;
import org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions;
import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
-import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions;
import org.apache.beam.runners.dataflow.util.DataflowTemplateJob;
import org.apache.beam.runners.dataflow.util.DataflowTransport;
import org.apache.beam.runners.dataflow.util.MonitoringUtil;
@@ -93,7 +92,6 @@ import org.apache.beam.sdk.io.BoundedSource;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.Read;
import org.apache.beam.sdk.io.UnboundedSource;
-import org.apache.beam.sdk.io.WriteFiles;
import org.apache.beam.sdk.io.fs.ResourceId;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubMessageWithAttributesCoder;
@@ -107,8 +105,6 @@ import org.apache.beam.sdk.runners.PTransformOverride;
import org.apache.beam.sdk.runners.PTransformOverrideFactory;
import org.apache.beam.sdk.runners.TransformHierarchy;
import org.apache.beam.sdk.runners.TransformHierarchy.Node;
-import org.apache.beam.sdk.state.MapState;
-import org.apache.beam.sdk.state.SetState;
import org.apache.beam.sdk.transforms.Combine;
import org.apache.beam.sdk.transforms.Combine.GroupedValues;
import org.apache.beam.sdk.transforms.Create;
@@ -121,8 +117,6 @@ import org.apache.beam.sdk.transforms.SimpleFunction;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.WithKeys;
import org.apache.beam.sdk.transforms.display.DisplayData;
-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.util.CoderUtils;
import org.apache.beam.sdk.util.InstanceBuilder;
@@ -135,12 +129,10 @@ import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollection.IsBounded;
import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.PCollectionViews;
import org.apache.beam.sdk.values.PDone;
import org.apache.beam.sdk.values.PInput;
import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.sdk.values.ValueWithRecordId;
import org.apache.beam.sdk.values.WindowingStrategy;
import org.joda.time.DateTimeUtils;
@@ -333,24 +325,6 @@ public class DataflowRunner extends PipelineRunner<DataflowPipelineJob> {
new StreamingFnApiCreateOverrideFactory()));
}
overridesBuilder
- // Support Splittable DoFn for now only in streaming mode.
- // The order of the following overrides is important because they are applied in order.
-
- // By default Dataflow runner replaces single-output ParDo with a ParDoSingle override.
- // However, we want a different expansion for single-output splittable ParDo.
- .add(
- PTransformOverride.of(
- PTransformMatchers.splittableParDoSingle(),
- new ReflectiveOneToOneOverrideFactory(
- SplittableParDoOverrides.ParDoSingleViaMulti.class, this)))
- .add(
- PTransformOverride.of(
- PTransformMatchers.splittableParDoMulti(),
- new SplittableParDoOverrides.SplittableParDoOverrideFactory()))
- .add(
- PTransformOverride.of(
- PTransformMatchers.writeWithRunnerDeterminedSharding(),
- new StreamingShardedWriteFactory(options)))
.add(
// Streaming Bounded Read is implemented in terms of Streaming Unbounded Read, and
// must precede it
@@ -376,29 +350,34 @@ public class DataflowRunner extends PipelineRunner<DataflowPipelineJob> {
PTransformOverride.of(
PTransformMatchers.stateOrTimerParDoSingle(),
BatchStatefulParDoOverrides.singleOutputOverrideFactory()))
+
+ .add(
+ PTransformOverride.of(
+ PTransformMatchers.classEqualTo(Combine.GloballyAsSingletonView.class),
+ new BatchCombineGloballyAsSingletonViewFactory(this)))
.add(
PTransformOverride.of(
- PTransformMatchers.createViewWithViewFn(PCollectionViews.MapViewFn.class),
+ PTransformMatchers.classEqualTo(View.AsMap.class),
new ReflectiveOneToOneOverrideFactory(
BatchViewOverrides.BatchViewAsMap.class, this)))
.add(
PTransformOverride.of(
- PTransformMatchers.createViewWithViewFn(PCollectionViews.MultimapViewFn.class),
+ PTransformMatchers.classEqualTo(View.AsMultimap.class),
new ReflectiveOneToOneOverrideFactory(
BatchViewOverrides.BatchViewAsMultimap.class, this)))
.add(
PTransformOverride.of(
- PTransformMatchers.createViewWithViewFn(PCollectionViews.SingletonViewFn.class),
+ PTransformMatchers.classEqualTo(View.AsSingleton.class),
new ReflectiveOneToOneOverrideFactory(
BatchViewOverrides.BatchViewAsSingleton.class, this)))
.add(
PTransformOverride.of(
- PTransformMatchers.createViewWithViewFn(PCollectionViews.ListViewFn.class),
+ PTransformMatchers.classEqualTo(View.AsList.class),
new ReflectiveOneToOneOverrideFactory(
BatchViewOverrides.BatchViewAsList.class, this)))
.add(
PTransformOverride.of(
- PTransformMatchers.createViewWithViewFn(PCollectionViews.IterableViewFn.class),
+ PTransformMatchers.classEqualTo(View.AsIterable.class),
new ReflectiveOneToOneOverrideFactory(
BatchViewOverrides.BatchViewAsIterable.class, this)));
}
@@ -1455,61 +1434,6 @@ public class DataflowRunner extends PipelineRunner<DataflowPipelineJob> {
}
@VisibleForTesting
- static class StreamingShardedWriteFactory<UserT, DestinationT, OutputT>
- implements PTransformOverrideFactory<
- PCollection<UserT>, PDone, WriteFiles<UserT, DestinationT, OutputT>> {
- // We pick 10 as a a default, as it works well with the default number of workers started
- // by Dataflow.
- static final int DEFAULT_NUM_SHARDS = 10;
- DataflowPipelineWorkerPoolOptions options;
-
- StreamingShardedWriteFactory(PipelineOptions options) {
- this.options = options.as(DataflowPipelineWorkerPoolOptions.class);
- }
-
- @Override
- public PTransformReplacement<PCollection<UserT>, PDone> getReplacementTransform(
- AppliedPTransform<PCollection<UserT>, PDone, WriteFiles<UserT, DestinationT, OutputT>>
- transform) {
- // By default, if numShards is not set WriteFiles will produce one file per bundle. In
- // streaming, there are large numbers of small bundles, resulting in many tiny files.
- // Instead we pick max workers * 2 to ensure full parallelism, but prevent too-many files.
- // (current_num_workers * 2 might be a better choice, but that value is not easily available
- // today).
- // If the user does not set either numWorkers or maxNumWorkers, default to 10 shards.
- int numShards;
- if (options.getMaxNumWorkers() > 0) {
- numShards = options.getMaxNumWorkers() * 2;
- } else if (options.getNumWorkers() > 0) {
- numShards = options.getNumWorkers() * 2;
- } else {
- numShards = DEFAULT_NUM_SHARDS;
- }
-
- try {
- WriteFiles<UserT, DestinationT, OutputT> replacement =
- WriteFiles.<UserT, DestinationT, OutputT>to(
- WriteFilesTranslation.<UserT, DestinationT, OutputT>getSink(transform),
- WriteFilesTranslation.<UserT, OutputT>getFormatFunction(transform));
- if (WriteFilesTranslation.isWindowedWrites(transform)) {
- replacement = replacement.withWindowedWrites();
- }
- return PTransformReplacement.of(
- PTransformReplacements.getSingletonMainInput(transform),
- replacement.withNumShards(numShards));
- } catch (Exception e) {
- throw new RuntimeException(e);
- }
- }
-
- @Override
- public Map<PValue, ReplacementOutput> mapOutputs(Map<TupleTag<?>, PValue> outputs,
- PDone newOutput) {
- return Collections.emptyMap();
- }
- }
-
- @VisibleForTesting
static String getContainerImageForJob(DataflowPipelineOptions options) {
String workerHarnessContainerImage = options.getWorkerHarnessContainerImage();
if (!workerHarnessContainerImage.contains("IMAGE")) {
@@ -1522,39 +1446,4 @@ public class DataflowRunner extends PipelineRunner<DataflowPipelineJob> {
return workerHarnessContainerImage.replace("IMAGE", "beam-java-batch");
}
}
-
- static void verifyStateSupported(DoFn<?, ?> fn) {
- DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());
-
- for (DoFnSignature.StateDeclaration stateDecl : signature.stateDeclarations().values()) {
-
- // https://issues.apache.org/jira/browse/BEAM-1474
- if (stateDecl.stateType().isSubtypeOf(TypeDescriptor.of(MapState.class))) {
- throw new UnsupportedOperationException(String.format(
- "%s does not currently support %s",
- DataflowRunner.class.getSimpleName(),
- MapState.class.getSimpleName()
- ));
- }
-
- // https://issues.apache.org/jira/browse/BEAM-1479
- if (stateDecl.stateType().isSubtypeOf(TypeDescriptor.of(SetState.class))) {
- throw new UnsupportedOperationException(String.format(
- "%s does not currently support %s",
- DataflowRunner.class.getSimpleName(),
- SetState.class.getSimpleName()
- ));
- }
- }
- }
-
- static void verifyStateSupportForWindowingStrategy(WindowingStrategy strategy) {
- // https://issues.apache.org/jira/browse/BEAM-2507
- if (!strategy.getWindowFn().isNonMerging()) {
- throw new UnsupportedOperationException(
- String.format(
- "%s does not currently support state or timers with merging windows",
- DataflowRunner.class.getSimpleName()));
- }
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/SplittableParDoOverrides.java
----------------------------------------------------------------------
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/SplittableParDoOverrides.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/SplittableParDoOverrides.java
deleted file mode 100644
index fc010f8..0000000
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/SplittableParDoOverrides.java
+++ /dev/null
@@ -1,76 +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.dataflow;
-
-import java.util.Map;
-import org.apache.beam.runners.core.construction.ForwardingPTransform;
-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.runners.AppliedPTransform;
-import org.apache.beam.sdk.runners.PTransformOverrideFactory;
-import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PCollectionTuple;
-import org.apache.beam.sdk.values.PValue;
-import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.sdk.values.TupleTagList;
-
-/** Transform overrides for supporting {@link SplittableParDo} in the Dataflow runner. */
-class SplittableParDoOverrides {
- static class ParDoSingleViaMulti<InputT, OutputT>
- extends ForwardingPTransform<PCollection<? extends InputT>, PCollection<OutputT>> {
- private final ParDo.SingleOutput<InputT, OutputT> original;
-
- public ParDoSingleViaMulti(
- DataflowRunner ignored, ParDo.SingleOutput<InputT, OutputT> original) {
- this.original = original;
- }
-
- @Override
- protected PTransform<PCollection<? extends InputT>, PCollection<OutputT>> delegate() {
- return original;
- }
-
- @Override
- public PCollection<OutputT> expand(PCollection<? extends InputT> input) {
- TupleTag<OutputT> mainOutput = new TupleTag<>();
- return input.apply(original.withOutputTags(mainOutput, TupleTagList.empty())).get(mainOutput);
- }
- }
-
- static class SplittableParDoOverrideFactory<InputT, OutputT, RestrictionT>
- implements PTransformOverrideFactory<
- PCollection<InputT>, PCollectionTuple, ParDo.MultiOutput<InputT, OutputT>> {
- @Override
- public PTransformReplacement<PCollection<InputT>, PCollectionTuple> getReplacementTransform(
- AppliedPTransform<PCollection<InputT>, PCollectionTuple, ParDo.MultiOutput<InputT, OutputT>>
- appliedTransform) {
- return PTransformReplacement.of(
- PTransformReplacements.getSingletonMainInput(appliedTransform),
- SplittableParDo.forJavaParDo(appliedTransform.getTransform()));
- }
-
- @Override
- public Map<PValue, ReplacementOutput> mapOutputs(
- Map<TupleTag<?>, PValue> outputs, PCollectionTuple newOutput) {
- return ReplacementOutputs.tagged(outputs, newOutput);
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/StreamingViewOverrides.java
----------------------------------------------------------------------
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/StreamingViewOverrides.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/StreamingViewOverrides.java
index 1853248..6c385d7 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/StreamingViewOverrides.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/StreamingViewOverrides.java
@@ -42,12 +42,12 @@ import org.apache.beam.sdk.values.PCollectionView;
class StreamingViewOverrides {
static class StreamingCreatePCollectionViewFactory<ElemT, ViewT>
extends SingleInputOutputOverrideFactory<
- PCollection<ElemT>, PCollection<ElemT>, CreatePCollectionView<ElemT, ViewT>> {
+ PCollection<ElemT>, PCollectionView<ViewT>, CreatePCollectionView<ElemT, ViewT>> {
@Override
- public PTransformReplacement<PCollection<ElemT>, PCollection<ElemT>>
+ public PTransformReplacement<PCollection<ElemT>, PCollectionView<ViewT>>
getReplacementTransform(
AppliedPTransform<
- PCollection<ElemT>, PCollection<ElemT>, CreatePCollectionView<ElemT, ViewT>>
+ PCollection<ElemT>, PCollectionView<ViewT>, CreatePCollectionView<ElemT, ViewT>>
transform) {
StreamingCreatePCollectionView<ElemT, ViewT> streamingView =
new StreamingCreatePCollectionView<>(transform.getTransform().getView());
@@ -56,7 +56,7 @@ class StreamingViewOverrides {
}
private static class StreamingCreatePCollectionView<ElemT, ViewT>
- extends PTransform<PCollection<ElemT>, PCollection<ElemT>> {
+ extends PTransform<PCollection<ElemT>, PCollectionView<ViewT>> {
private final PCollectionView<ViewT> view;
private StreamingCreatePCollectionView(PCollectionView<ViewT> view) {
@@ -64,7 +64,7 @@ class StreamingViewOverrides {
}
@Override
- public PCollection<ElemT> expand(PCollection<ElemT> input) {
+ public PCollectionView<ViewT> expand(PCollection<ElemT> input) {
return input
.apply(Combine.globally(new Concatenate<ElemT>()).withoutDefaults())
.apply(ParDo.of(StreamingPCollectionViewWriterFn.create(view, input.getCoder())))
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/TransformTranslator.java
----------------------------------------------------------------------
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/TransformTranslator.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/TransformTranslator.java
index 06ed1e0..a7452b2 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/TransformTranslator.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/TransformTranslator.java
@@ -22,7 +22,6 @@ import java.util.Map;
import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
import org.apache.beam.runners.dataflow.util.OutputReference;
import org.apache.beam.sdk.Pipeline;
-import org.apache.beam.sdk.annotations.Internal;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.transforms.PTransform;
@@ -37,8 +36,7 @@ import org.apache.beam.sdk.values.TupleTag;
* A {@link TransformTranslator} knows how to translate a particular subclass of {@link PTransform}
* for the Cloud Dataflow service. It does so by mutating the {@link TranslationContext}.
*/
-@Internal
-public interface TransformTranslator<TransformT extends PTransform> {
+interface TransformTranslator<TransformT extends PTransform> {
void translate(TransformT transform, TranslationContext context);
/**
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/PropertyNames.java
----------------------------------------------------------------------
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/PropertyNames.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/PropertyNames.java
index 55e0c4e..f82f1f1 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/PropertyNames.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/PropertyNames.java
@@ -63,5 +63,4 @@ public class PropertyNames {
public static final String USES_KEYED_STATE = "uses_keyed_state";
public static final String VALUE = "value";
public static final String DISPLAY_DATA = "display_data";
- public static final String RESTRICTION_CODER = "restriction_coder";
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/TimeUtil.java
----------------------------------------------------------------------
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/TimeUtil.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/TimeUtil.java
index 172dc6e..bff379f 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/TimeUtil.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/TimeUtil.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.runners.dataflow.util;
-import com.google.common.base.Strings;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import javax.annotation.Nullable;
@@ -99,19 +98,26 @@ public final class TimeUtil {
int hour = Integer.valueOf(matcher.group(4));
int minute = Integer.valueOf(matcher.group(5));
int second = Integer.valueOf(matcher.group(6));
- int millis = computeMillis(matcher.group(7));
+ int millis = 0;
+
+ String frac = matcher.group(7);
+ if (frac != null) {
+ int fracs = Integer.valueOf(frac);
+ if (frac.length() == 3) { // millisecond resolution
+ millis = fracs;
+ } else if (frac.length() == 6) { // microsecond resolution
+ millis = fracs / 1000;
+ } else if (frac.length() == 9) { // nanosecond resolution
+ millis = fracs / 1000000;
+ } else {
+ return null;
+ }
+ }
return new DateTime(year, month, day, hour, minute, second, millis,
ISOChronology.getInstanceUTC()).toInstant();
}
- private static int computeMillis(String frac) {
- if (frac == null) {
- return 0;
- }
- return Integer.valueOf(frac.length() > 3 ? frac.substring(0, 3) : Strings.padEnd(frac, 3, '0'));
- }
-
/**
* Converts a {@link ReadableDuration} into a Dataflow API duration string.
*/
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslatorTest.java
----------------------------------------------------------------------
diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslatorTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslatorTest.java
index 43b2788..89dc2d5 100644
--- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslatorTest.java
+++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslatorTest.java
@@ -18,14 +18,11 @@
package org.apache.beam.runners.dataflow;
import static org.apache.beam.runners.dataflow.util.Structs.getString;
-import static org.apache.beam.sdk.util.StringUtils.jsonStringToByteArray;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.hasEntry;
import static org.hamcrest.Matchers.hasKey;
-import static org.hamcrest.Matchers.instanceOf;
import static org.hamcrest.Matchers.not;
import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.assertTrue;
@@ -69,22 +66,17 @@ import java.util.Set;
import org.apache.beam.runners.dataflow.DataflowPipelineTranslator.JobSpecification;
import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions;
-import org.apache.beam.runners.dataflow.util.CloudObject;
-import org.apache.beam.runners.dataflow.util.CloudObjects;
-import org.apache.beam.runners.dataflow.util.DoFnInfo;
import org.apache.beam.runners.dataflow.util.OutputReference;
import org.apache.beam.runners.dataflow.util.PropertyNames;
import org.apache.beam.runners.dataflow.util.Structs;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.SerializableCoder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.extensions.gcp.auth.TestCredential;
import org.apache.beam.sdk.extensions.gcp.storage.GcsPathValidator;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.TextIO;
-import org.apache.beam.sdk.io.range.OffsetRange;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.ValueProvider;
@@ -99,12 +91,7 @@ import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.Sum;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.display.DisplayData;
-import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
-import org.apache.beam.sdk.transforms.windowing.FixedWindows;
-import org.apache.beam.sdk.transforms.windowing.Window;
-import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.util.GcsUtil;
-import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.util.gcsfs.GcsPath;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
@@ -113,8 +100,6 @@ import org.apache.beam.sdk.values.PDone;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
import org.apache.beam.sdk.values.WindowingStrategy;
-import org.hamcrest.Matchers;
-import org.joda.time.Duration;
import org.junit.Assert;
import org.junit.Rule;
import org.junit.Test;
@@ -911,68 +896,6 @@ public class DataflowPipelineTranslatorTest implements Serializable {
not(equalTo("true")));
}
- /**
- * Smoke test to fail fast if translation of a splittable ParDo
- * in streaming breaks.
- */
- @Test
- public void testStreamingSplittableParDoTranslation() throws Exception {
- DataflowPipelineOptions options = buildPipelineOptions();
- DataflowRunner runner = DataflowRunner.fromOptions(options);
- options.setStreaming(true);
- DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options);
-
- Pipeline pipeline = Pipeline.create(options);
-
- PCollection<String> windowedInput = pipeline
- .apply(Create.of("a"))
- .apply(Window.<String>into(FixedWindows.of(Duration.standardMinutes(1))));
- windowedInput.apply(ParDo.of(new TestSplittableFn()));
-
- runner.replaceTransforms(pipeline);
-
- Job job =
- translator
- .translate(
- pipeline,
- runner,
- Collections.<DataflowPackage>emptyList())
- .getJob();
-
- // The job should contain a SplittableParDo.ProcessKeyedElements step, translated as
- // "SplittableProcessKeyed".
-
- List<Step> steps = job.getSteps();
- Step processKeyedStep = null;
- for (Step step : steps) {
- if (step.getKind().equals("SplittableProcessKeyed")) {
- assertNull(processKeyedStep);
- processKeyedStep = step;
- }
- }
- assertNotNull(processKeyedStep);
-
- @SuppressWarnings({"unchecked", "rawtypes"})
- DoFnInfo<String, Integer> fnInfo =
- (DoFnInfo<String, Integer>)
- SerializableUtils.deserializeFromByteArray(
- jsonStringToByteArray(
- Structs.getString(
- processKeyedStep.getProperties(), PropertyNames.SERIALIZED_FN)),
- "DoFnInfo");
- assertThat(fnInfo.getDoFn(), instanceOf(TestSplittableFn.class));
- assertThat(
- fnInfo.getWindowingStrategy().getWindowFn(),
- Matchers.<WindowFn>equalTo(FixedWindows.of(Duration.standardMinutes(1))));
- Coder<?> restrictionCoder =
- CloudObjects.coderFromCloudObject(
- (CloudObject)
- Structs.getObject(
- processKeyedStep.getProperties(), PropertyNames.RESTRICTION_CODER));
-
- assertEquals(SerializableCoder.of(OffsetRange.class), restrictionCoder);
- }
-
@Test
public void testToSingletonTranslationWithIsmSideInput() throws Exception {
// A "change detector" test that makes sure the translation
@@ -997,15 +920,15 @@ public class DataflowPipelineTranslatorTest implements Serializable {
assertAllStepOutputsHaveUniqueIds(job);
List<Step> steps = job.getSteps();
- assertEquals(9, steps.size());
+ assertEquals(5, steps.size());
@SuppressWarnings("unchecked")
List<Map<String, Object>> toIsmRecordOutputs =
- (List<Map<String, Object>>) steps.get(7).getProperties().get(PropertyNames.OUTPUT_INFO);
+ (List<Map<String, Object>>) steps.get(3).getProperties().get(PropertyNames.OUTPUT_INFO);
assertTrue(
Structs.getBoolean(Iterables.getOnlyElement(toIsmRecordOutputs), "use_indexed_format"));
- Step collectionToSingletonStep = steps.get(8);
+ Step collectionToSingletonStep = steps.get(4);
assertEquals("CollectionToSingleton", collectionToSingletonStep.getKind());
}
@@ -1167,16 +1090,4 @@ public class DataflowPipelineTranslatorTest implements Serializable {
assertTrue(String.format("Found duplicate output ids %s", outputIds),
outputIds.size() == 0);
}
-
- private static class TestSplittableFn extends DoFn<String, Integer> {
- @ProcessElement
- public void process(ProcessContext c, OffsetRangeTracker tracker) {
- // noop
- }
-
- @GetInitialRestriction
- public OffsetRange getInitialRange(String element) {
- return null;
- }
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java
----------------------------------------------------------------------
diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java
index 94985f8..8f10b18 100644
--- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java
+++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java
@@ -23,7 +23,6 @@ import static org.hamcrest.Matchers.containsString;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.hasItem;
import static org.hamcrest.Matchers.is;
-import static org.hamcrest.Matchers.not;
import static org.hamcrest.Matchers.startsWith;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
@@ -50,7 +49,6 @@ import com.google.common.collect.ImmutableList;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
-import java.io.Serializable;
import java.net.URL;
import java.net.URLClassLoader;
import java.nio.channels.FileChannel;
@@ -64,60 +62,36 @@ import java.util.LinkedList;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.regex.Pattern;
-import org.apache.beam.runners.dataflow.DataflowRunner.StreamingShardedWriteFactory;
import org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions;
import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
-import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.Pipeline.PipelineVisitor;
import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.extensions.gcp.auth.NoopCredentialFactory;
import org.apache.beam.sdk.extensions.gcp.auth.TestCredential;
import org.apache.beam.sdk.extensions.gcp.storage.NoopPathValidator;
-import org.apache.beam.sdk.io.DynamicFileDestinations;
-import org.apache.beam.sdk.io.FileBasedSink;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.TextIO;
-import org.apache.beam.sdk.io.WriteFiles;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptions.CheckEnabled;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.sdk.options.StreamingOptions;
import org.apache.beam.sdk.options.ValueProvider;
-import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
-import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.runners.TransformHierarchy;
import org.apache.beam.sdk.runners.TransformHierarchy.Node;
-import org.apache.beam.sdk.state.MapState;
-import org.apache.beam.sdk.state.SetState;
-import org.apache.beam.sdk.state.StateSpec;
-import org.apache.beam.sdk.state.StateSpecs;
-import org.apache.beam.sdk.state.ValueState;
import org.apache.beam.sdk.testing.ExpectedLogs;
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.PTransform;
-import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.transforms.SerializableFunctions;
-import org.apache.beam.sdk.transforms.windowing.Sessions;
-import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.util.GcsUtil;
import org.apache.beam.sdk.util.ReleaseInfo;
import org.apache.beam.sdk.util.gcsfs.GcsPath;
-import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PDone;
-import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TimestampedValue;
-import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
import org.hamcrest.Description;
import org.hamcrest.Matchers;
import org.hamcrest.TypeSafeMatcher;
-import org.joda.time.Duration;
import org.joda.time.Instant;
import org.junit.Before;
import org.junit.Rule;
@@ -133,11 +107,9 @@ import org.mockito.stubbing.Answer;
/**
* Tests for the {@link DataflowRunner}.
- *
- * <p>Implements {@link Serializable} because it is caught in closures.
*/
@RunWith(JUnit4.class)
-public class DataflowRunnerTest implements Serializable {
+public class DataflowRunnerTest {
private static final String VALID_STAGING_BUCKET = "gs://valid-bucket/staging";
private static final String VALID_TEMP_BUCKET = "gs://valid-bucket/temp";
@@ -147,12 +119,15 @@ public class DataflowRunnerTest implements Serializable {
private static final String PROJECT_ID = "some-project";
private static final String REGION_ID = "some-region-1";
- @Rule public transient TemporaryFolder tmpFolder = new TemporaryFolder();
- @Rule public transient ExpectedException thrown = ExpectedException.none();
- @Rule public transient ExpectedLogs expectedLogs = ExpectedLogs.none(DataflowRunner.class);
+ @Rule
+ public TemporaryFolder tmpFolder = new TemporaryFolder();
+ @Rule
+ public ExpectedException thrown = ExpectedException.none();
+ @Rule
+ public ExpectedLogs expectedLogs = ExpectedLogs.none(DataflowRunner.class);
- private transient Dataflow.Projects.Locations.Jobs mockJobs;
- private transient GcsUtil mockGcsUtil;
+ private Dataflow.Projects.Locations.Jobs mockJobs;
+ private GcsUtil mockGcsUtil;
// Asserts that the given Job has all expected fields set.
private static void assertValidJob(Job job) {
@@ -848,6 +823,7 @@ public class DataflowRunnerTest implements Serializable {
DataflowRunner.fromOptions(options);
}
+
@Test
public void testValidProfileLocation() throws IOException {
DataflowPipelineOptions options = buildPipelineOptions();
@@ -1015,71 +991,6 @@ public class DataflowRunnerTest implements Serializable {
assertTrue(transform.translated);
}
- private void verifyMapStateUnsupported(PipelineOptions options) throws Exception {
- Pipeline p = Pipeline.create(options);
- p.apply(Create.of(KV.of(13, 42)))
- .apply(
- ParDo.of(
- new DoFn<KV<Integer, Integer>, Void>() {
- @StateId("fizzle")
- private final StateSpec<MapState<Void, Void>> voidState = StateSpecs.map();
-
- @ProcessElement
- public void process() {}
- }));
-
- thrown.expectMessage("MapState");
- thrown.expect(UnsupportedOperationException.class);
- p.run();
- }
-
- @Test
- public void testMapStateUnsupportedInBatch() throws Exception {
- PipelineOptions options = buildPipelineOptions();
- options.as(StreamingOptions.class).setStreaming(false);
- verifyMapStateUnsupported(options);
- }
-
- @Test
- public void testMapStateUnsupportedInStreaming() throws Exception {
- PipelineOptions options = buildPipelineOptions();
- options.as(StreamingOptions.class).setStreaming(true);
- verifyMapStateUnsupported(options);
- }
-
- private void verifySetStateUnsupported(PipelineOptions options) throws Exception {
- Pipeline p = Pipeline.create(options);
- p.apply(Create.of(KV.of(13, 42)))
- .apply(
- ParDo.of(
- new DoFn<KV<Integer, Integer>, Void>() {
- @StateId("fizzle")
- private final StateSpec<SetState<Void>> voidState = StateSpecs.set();
-
- @ProcessElement
- public void process() {}
- }));
-
- thrown.expectMessage("SetState");
- thrown.expect(UnsupportedOperationException.class);
- p.run();
- }
-
- @Test
- public void testSetStateUnsupportedInBatch() throws Exception {
- PipelineOptions options = buildPipelineOptions();
- options.as(StreamingOptions.class).setStreaming(false);
- Pipeline p = Pipeline.create(options);
- verifySetStateUnsupported(options);
- }
-
- @Test
- public void testSetStateUnsupportedInStreaming() throws Exception {
- PipelineOptions options = buildPipelineOptions();
- options.as(StreamingOptions.class).setStreaming(true);
- verifySetStateUnsupported(options);
- }
-
/** Records all the composite transforms visited within the Pipeline. */
private static class CompositeTransformRecorder extends PipelineVisitor.Defaults {
private List<PTransform<?, ?>> transforms = new ArrayList<>();
@@ -1136,8 +1047,8 @@ public class DataflowRunnerTest implements Serializable {
}
/**
- * Tests that the {@link DataflowRunner} with {@code --templateLocation} returns normally when the
- * runner is successfully run.
+ * Tests that the {@link DataflowRunner} with {@code --templateLocation} returns normally
+ * when the runner issuccessfully run.
*/
@Test
public void testTemplateRunnerFullCompletion() throws Exception {
@@ -1216,89 +1127,4 @@ public class DataflowRunnerTest implements Serializable {
assertThat(
getContainerImageForJob(options), equalTo("gcr.io/java/foo"));
}
-
- @Test
- public void testStreamingWriteWithNoShardingReturnsNewTransform() {
- PipelineOptions options = TestPipeline.testingPipelineOptions();
- options.as(DataflowPipelineWorkerPoolOptions.class).setMaxNumWorkers(10);
- testStreamingWriteOverride(options, 20);
- }
-
- @Test
- public void testStreamingWriteWithNoShardingReturnsNewTransformMaxWorkersUnset() {
- PipelineOptions options = TestPipeline.testingPipelineOptions();
- testStreamingWriteOverride(options, StreamingShardedWriteFactory.DEFAULT_NUM_SHARDS);
- }
-
- private void verifyMergingStatefulParDoRejected(PipelineOptions options) throws Exception {
- Pipeline p = Pipeline.create(options);
-
- p.apply(Create.of(KV.of(13, 42)))
- .apply(Window.<KV<Integer, Integer>>into(Sessions.withGapDuration(Duration.millis(1))))
- .apply(ParDo.of(new DoFn<KV<Integer, Integer>, Void>() {
- @StateId("fizzle")
- private final StateSpec<ValueState<Void>> voidState = StateSpecs.value();
-
- @ProcessElement
- public void process() {}
- }));
-
- thrown.expectMessage("merging");
- thrown.expect(UnsupportedOperationException.class);
- p.run();
- }
-
- @Test
- public void testMergingStatefulRejectedInStreaming() throws Exception {
- PipelineOptions options = buildPipelineOptions();
- options.as(StreamingOptions.class).setStreaming(true);
- verifyMergingStatefulParDoRejected(options);
- }
-
- @Test
- public void testMergingStatefulRejectedInBatch() throws Exception {
- PipelineOptions options = buildPipelineOptions();
- options.as(StreamingOptions.class).setStreaming(false);
- verifyMergingStatefulParDoRejected(options);
- }
-
- private void testStreamingWriteOverride(PipelineOptions options, int expectedNumShards) {
- TestPipeline p = TestPipeline.fromOptions(options);
-
- StreamingShardedWriteFactory<Object, Void, Object> factory =
- new StreamingShardedWriteFactory<>(p.getOptions());
- WriteFiles<Object, Void, Object> original =
- WriteFiles.to(new TestSink(tmpFolder.toString()), SerializableFunctions.identity());
- PCollection<Object> objs = (PCollection) p.apply(Create.empty(VoidCoder.of()));
- AppliedPTransform<PCollection<Object>, PDone, WriteFiles<Object, Void, Object>>
- originalApplication =
- AppliedPTransform.of(
- "writefiles",
- objs.expand(),
- Collections.<TupleTag<?>, PValue>emptyMap(),
- original,
- p);
-
- WriteFiles<Object, Void, Object> replacement =
- (WriteFiles<Object, Void, Object>)
- factory.getReplacementTransform(originalApplication).getTransform();
- assertThat(replacement, not(equalTo((Object) original)));
- assertThat(replacement.getNumShards().get(), equalTo(expectedNumShards));
- }
-
- private static class TestSink extends FileBasedSink<Object, Void> {
- @Override
- public void validate(PipelineOptions options) {}
-
- TestSink(String tmpFolder) {
- super(
- StaticValueProvider.of(FileSystems.matchNewResource(tmpFolder, true)),
- DynamicFileDestinations.constant(null));
- }
-
- @Override
- public WriteOperation<Object, Void> createWriteOperation() {
- throw new IllegalArgumentException("Should not be used");
- }
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/TimeUtilTest.java
----------------------------------------------------------------------
diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/TimeUtilTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/TimeUtilTest.java
index 1ac9fab..e0785d4 100644
--- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/TimeUtilTest.java
+++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/TimeUtilTest.java
@@ -47,14 +47,8 @@ public final class TimeUtilTest {
assertEquals(new Instant(1), fromCloudTime("1970-01-01T00:00:00.001001Z"));
assertEquals(new Instant(1), fromCloudTime("1970-01-01T00:00:00.001000000Z"));
assertEquals(new Instant(1), fromCloudTime("1970-01-01T00:00:00.001000001Z"));
- assertEquals(new Instant(0), fromCloudTime("1970-01-01T00:00:00.0Z"));
- assertEquals(new Instant(0), fromCloudTime("1970-01-01T00:00:00.00Z"));
- assertEquals(new Instant(420), fromCloudTime("1970-01-01T00:00:00.42Z"));
- assertEquals(new Instant(300), fromCloudTime("1970-01-01T00:00:00.3Z"));
- assertEquals(new Instant(20), fromCloudTime("1970-01-01T00:00:00.02Z"));
assertNull(fromCloudTime(""));
assertNull(fromCloudTime("1970-01-01T00:00:00"));
- assertNull(fromCloudTime("1970-01-01T00:00:00.1e3Z"));
}
@Test
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/pom.xml
----------------------------------------------------------------------
diff --git a/runners/pom.xml b/runners/pom.xml
index b00ba9c..38aada8 100644
--- a/runners/pom.xml
+++ b/runners/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/spark/pom.xml
----------------------------------------------------------------------
diff --git a/runners/spark/pom.xml b/runners/spark/pom.xml
index 7f70204..ddb4aca 100644
--- a/runners/spark/pom.xml
+++ b/runners/spark/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-runners-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
@@ -34,6 +34,8 @@
<properties>
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
<project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding>
+ <spark.version>1.6.3</spark.version>
+ <hadoop.version>2.2.0</hadoop.version>
<kafka.version>0.9.0.1</kafka.version>
<jackson.version>2.4.4</jackson.version>
<dropwizard.metrics.version>3.1.2</dropwizard.metrics.version>
@@ -101,27 +103,6 @@
<threadCount>4</threadCount>
</configuration>
</execution>
- <execution>
- <id>streaming-tests</id>
- <phase>test</phase>
- <goals>
- <goal>test</goal>
- </goals>
- <configuration>
- <groups>
- org.apache.beam.runners.spark.StreamingTest
- </groups>
- <systemPropertyVariables>
- <beamTestPipelineOptions>
- [
- "--runner=TestSparkRunner",
- "--forceStreaming=true",
- "--enableSparkMetricSinks=true"
- ]
- </beamTestPipelineOptions>
- </systemPropertyVariables>
- </configuration>
- </execution>
</executions>
</plugin>
</plugins>
@@ -133,33 +114,31 @@
<dependency>
<groupId>org.apache.spark</groupId>
<artifactId>spark-core_2.10</artifactId>
+ <version>${spark.version}</version>
<scope>provided</scope>
</dependency>
<dependency>
<groupId>org.apache.spark</groupId>
<artifactId>spark-streaming_2.10</artifactId>
+ <version>${spark.version}</version>
<scope>provided</scope>
</dependency>
<dependency>
<groupId>org.apache.spark</groupId>
<artifactId>spark-network-common_2.10</artifactId>
+ <version>${spark.version}</version>
<scope>provided</scope>
</dependency>
<dependency>
<groupId>org.apache.hadoop</groupId>
<artifactId>hadoop-common</artifactId>
+ <version>${hadoop.version}</version>
<scope>provided</scope>
- <exclusions>
- <!-- Fix build on JDK-9 -->
- <exclusion>
- <groupId>jdk.tools</groupId>
- <artifactId>jdk.tools</artifactId>
- </exclusion>
- </exclusions>
</dependency>
<dependency>
<groupId>org.apache.hadoop</groupId>
<artifactId>hadoop-mapreduce-client-core</artifactId>
+ <version>${hadoop.version}</version>
<scope>provided</scope>
</dependency>
<dependency>
@@ -216,11 +195,6 @@
<dependency>
<groupId>org.apache.commons</groupId>
<artifactId>commons-lang3</artifactId>
- <scope>provided</scope>
- </dependency>
- <dependency>
- <groupId>org.apache.commons</groupId>
- <artifactId>commons-text</artifactId>
</dependency>
<dependency>
<groupId>commons-io</groupId>
@@ -347,13 +321,6 @@
<type>test-jar</type>
<scope>test</scope>
</dependency>
-
- <dependency>
- <groupId>org.apache.beam</groupId>
- <artifactId>beam-runners-core-java</artifactId>
- <type>test-jar</type>
- <scope>test</scope>
- </dependency>
</dependencies>
<build>
@@ -393,6 +360,27 @@
</systemPropertyVariables>
</configuration>
</execution>
+ <execution>
+ <id>streaming-tests</id>
+ <phase>test</phase>
+ <goals>
+ <goal>test</goal>
+ </goals>
+ <configuration>
+ <groups>
+ org.apache.beam.runners.spark.StreamingTest
+ </groups>
+ <systemPropertyVariables>
+ <beamTestPipelineOptions>
+ [
+ "--runner=TestSparkRunner",
+ "--forceStreaming=true",
+ "--enableSparkMetricSinks=true"
+ ]
+ </beamTestPipelineOptions>
+ </systemPropertyVariables>
+ </configuration>
+ </execution>
</executions>
</plugin>
<plugin>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkNativePipelineVisitor.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkNativePipelineVisitor.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkNativePipelineVisitor.java
index 6972acb..d75c955 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkNativePipelineVisitor.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkNativePipelineVisitor.java
@@ -35,7 +35,8 @@ import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.values.PInput;
import org.apache.beam.sdk.values.POutput;
-import org.apache.commons.text.WordUtils;
+import org.apache.commons.lang3.text.WordUtils;
+
/**
* Pipeline visitor for translating a Beam pipeline into equivalent Spark operations.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkRunner.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkRunner.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkRunner.java
index 595521f..9e2426e 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkRunner.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/SparkRunner.java
@@ -26,7 +26,6 @@ import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
-import org.apache.beam.runners.core.construction.TransformInputs;
import org.apache.beam.runners.spark.aggregators.AggregatorsAccumulator;
import org.apache.beam.runners.spark.io.CreateStream;
import org.apache.beam.runners.spark.metrics.AggregatorMetricSource;
@@ -171,7 +170,7 @@ public final class SparkRunner extends PipelineRunner<SparkPipelineResult> {
}
// register Watermarks listener to broadcast the advanced WMs.
- jssc.addStreamingListener(new JavaStreamingListenerWrapper(new WatermarksListener()));
+ jssc.addStreamingListener(new JavaStreamingListenerWrapper(new WatermarksListener(jssc)));
// The reason we call initAccumulators here even though it is called in
// SparkRunnerStreamingContextFactory is because the factory is not called when resuming
@@ -360,12 +359,10 @@ public final class SparkRunner extends PipelineRunner<SparkPipelineResult> {
protected boolean shouldDefer(TransformHierarchy.Node node) {
// if the input is not a PCollection, or it is but with non merging windows, don't defer.
- Collection<PValue> nonAdditionalInputs =
- TransformInputs.nonAdditionalInputs(node.toAppliedPTransform(getPipeline()));
- if (nonAdditionalInputs.size() != 1) {
+ if (node.getInputs().size() != 1) {
return false;
}
- PValue input = Iterables.getOnlyElement(nonAdditionalInputs);
+ PValue input = Iterables.getOnlyElement(node.getInputs().values());
if (!(input instanceof PCollection)
|| ((PCollection) input).getWindowingStrategy().getWindowFn().isNonMerging()) {
return false;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkRunner.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkRunner.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkRunner.java
index a13a3b1..eccee57 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkRunner.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/TestSparkRunner.java
@@ -169,7 +169,7 @@ public final class TestSparkRunner extends PipelineRunner<SparkPipelineResult> {
result.waitUntilFinish(Duration.millis(batchDurationMillis));
do {
SparkTimerInternals sparkTimerInternals =
- SparkTimerInternals.global(GlobalWatermarkHolder.get(batchDurationMillis));
+ SparkTimerInternals.global(GlobalWatermarkHolder.get());
sparkTimerInternals.advanceWatermark();
globalWatermark = sparkTimerInternals.currentInputWatermarkTime();
// let another batch-interval period of execution, just to reason about WM propagation.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkGroupAlsoByWindowViaWindowSet.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkGroupAlsoByWindowViaWindowSet.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkGroupAlsoByWindowViaWindowSet.java
index 1385e07..be4f3f6 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkGroupAlsoByWindowViaWindowSet.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkGroupAlsoByWindowViaWindowSet.java
@@ -104,15 +104,13 @@ public class SparkGroupAlsoByWindowViaWindowSet {
public static <K, InputT, W extends BoundedWindow>
JavaDStream<WindowedValue<KV<K, Iterable<InputT>>>> groupAlsoByWindow(
- final JavaDStream<WindowedValue<KV<K, Iterable<WindowedValue<InputT>>>>> inputDStream,
+ JavaDStream<WindowedValue<KV<K, Iterable<WindowedValue<InputT>>>>> inputDStream,
final Coder<K> keyCoder,
final Coder<WindowedValue<InputT>> wvCoder,
final WindowingStrategy<?, W> windowingStrategy,
final SparkRuntimeContext runtimeContext,
final List<Integer> sourceIds) {
- final long batchDurationMillis =
- runtimeContext.getPipelineOptions().as(SparkPipelineOptions.class).getBatchIntervalMillis();
final IterableCoder<WindowedValue<InputT>> itrWvCoder = IterableCoder.of(wvCoder);
final Coder<InputT> iCoder = ((FullWindowedValueCoder<InputT>) wvCoder).getValueCoder();
final Coder<? extends BoundedWindow> wCoder =
@@ -241,7 +239,7 @@ public class SparkGroupAlsoByWindowViaWindowSet {
SparkStateInternals<K> stateInternals;
SparkTimerInternals timerInternals = SparkTimerInternals.forStreamFromSources(
- sourceIds, GlobalWatermarkHolder.get(batchDurationMillis));
+ sourceIds, GlobalWatermarkHolder.get());
// get state(internals) per key.
if (prevStateAndTimersOpt.isEmpty()) {
// no previous state.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java
index a68da55..107915f 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java
@@ -34,6 +34,7 @@ import org.apache.beam.runners.spark.coders.CoderHelpers;
import org.apache.beam.runners.spark.util.GlobalWatermarkHolder.SparkWatermarks;
import org.apache.beam.sdk.state.TimeDomain;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.spark.broadcast.Broadcast;
import org.joda.time.Instant;
@@ -57,10 +58,10 @@ public class SparkTimerInternals implements TimerInternals {
/** Build the {@link TimerInternals} according to the feeding streams. */
public static SparkTimerInternals forStreamFromSources(
List<Integer> sourceIds,
- Map<Integer, SparkWatermarks> watermarks) {
- // if watermarks are invalid for the specific ids, use defaults.
- if (watermarks == null || watermarks.isEmpty()
- || Collections.disjoint(sourceIds, watermarks.keySet())) {
+ @Nullable Broadcast<Map<Integer, SparkWatermarks>> broadcast) {
+ // if broadcast is invalid for the specific ids, use defaults.
+ if (broadcast == null || broadcast.getValue().isEmpty()
+ || Collections.disjoint(sourceIds, broadcast.getValue().keySet())) {
return new SparkTimerInternals(
BoundedWindow.TIMESTAMP_MIN_VALUE, BoundedWindow.TIMESTAMP_MIN_VALUE, new Instant(0));
}
@@ -70,7 +71,7 @@ public class SparkTimerInternals implements TimerInternals {
// synchronized processing time should clearly be synchronized.
Instant synchronizedProcessingTime = null;
for (Integer sourceId: sourceIds) {
- SparkWatermarks sparkWatermarks = watermarks.get(sourceId);
+ SparkWatermarks sparkWatermarks = broadcast.getValue().get(sourceId);
if (sparkWatermarks != null) {
// keep slowest WMs.
slowestLowWatermark = slowestLowWatermark.isBefore(sparkWatermarks.getLowWatermark())
@@ -93,9 +94,10 @@ public class SparkTimerInternals implements TimerInternals {
}
/** Build a global {@link TimerInternals} for all feeding streams.*/
- public static SparkTimerInternals global(Map<Integer, SparkWatermarks> watermarks) {
- return watermarks == null ? forStreamFromSources(Collections.<Integer>emptyList(), null)
- : forStreamFromSources(Lists.newArrayList(watermarks.keySet()), watermarks);
+ public static SparkTimerInternals global(
+ @Nullable Broadcast<Map<Integer, SparkWatermarks>> broadcast) {
+ return broadcast == null ? forStreamFromSources(Collections.<Integer>emptyList(), null)
+ : forStreamFromSources(Lists.newArrayList(broadcast.getValue().keySet()), broadcast);
}
Collection<TimerData> getTimers() {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/EvaluationContext.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/EvaluationContext.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/EvaluationContext.java
index 0c6c4d1..8102926 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/EvaluationContext.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/EvaluationContext.java
@@ -26,7 +26,6 @@ import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;
-import org.apache.beam.runners.core.construction.TransformInputs;
import org.apache.beam.runners.spark.SparkPipelineOptions;
import org.apache.beam.runners.spark.coders.CoderHelpers;
import org.apache.beam.sdk.Pipeline;
@@ -104,8 +103,7 @@ public class EvaluationContext {
public <T extends PValue> T getInput(PTransform<T, ?> transform) {
@SuppressWarnings("unchecked")
- T input =
- (T) Iterables.getOnlyElement(TransformInputs.nonAdditionalInputs(getCurrentTransform()));
+ T input = (T) Iterables.getOnlyElement(getInputs(transform).values());
return input;
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TransformTranslator.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TransformTranslator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TransformTranslator.java
index ac5e0cd..64aa35a 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TransformTranslator.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TransformTranslator.java
@@ -508,6 +508,50 @@ public final class TransformTranslator {
};
}
+ private static <T> TransformEvaluator<View.AsSingleton<T>> viewAsSingleton() {
+ return new TransformEvaluator<View.AsSingleton<T>>() {
+ @Override
+ public void evaluate(View.AsSingleton<T> transform, EvaluationContext context) {
+ Iterable<? extends WindowedValue<?>> iter =
+ context.getWindowedValues(context.getInput(transform));
+ PCollectionView<T> output = context.getOutput(transform);
+ Coder<Iterable<WindowedValue<?>>> coderInternal = output.getCoderInternal();
+
+ @SuppressWarnings("unchecked")
+ Iterable<WindowedValue<?>> iterCast = (Iterable<WindowedValue<?>>) iter;
+
+ context.putPView(output, iterCast, coderInternal);
+ }
+
+ @Override
+ public String toNativeString() {
+ return "collect()";
+ }
+ };
+ }
+
+ private static <T> TransformEvaluator<View.AsIterable<T>> viewAsIter() {
+ return new TransformEvaluator<View.AsIterable<T>>() {
+ @Override
+ public void evaluate(View.AsIterable<T> transform, EvaluationContext context) {
+ Iterable<? extends WindowedValue<?>> iter =
+ context.getWindowedValues(context.getInput(transform));
+ PCollectionView<Iterable<T>> output = context.getOutput(transform);
+ Coder<Iterable<WindowedValue<?>>> coderInternal = output.getCoderInternal();
+
+ @SuppressWarnings("unchecked")
+ Iterable<WindowedValue<?>> iterCast = (Iterable<WindowedValue<?>>) iter;
+
+ context.putPView(output, iterCast, coderInternal);
+ }
+
+ @Override
+ public String toNativeString() {
+ return "collect()";
+ }
+ };
+ }
+
private static <ReadT, WriteT> TransformEvaluator<View.CreatePCollectionView<ReadT, WriteT>>
createPCollView() {
return new TransformEvaluator<View.CreatePCollectionView<ReadT, WriteT>>() {
@@ -516,7 +560,7 @@ public final class TransformTranslator {
EvaluationContext context) {
Iterable<? extends WindowedValue<?>> iter =
context.getWindowedValues(context.getInput(transform));
- PCollectionView<WriteT> output = transform.getView();
+ PCollectionView<WriteT> output = context.getOutput(transform);
Coder<Iterable<WindowedValue<?>>> coderInternal = output.getCoderInternal();
@SuppressWarnings("unchecked")
@@ -601,8 +645,8 @@ public final class TransformTranslator {
EVALUATORS.put(Combine.PerKey.class, combinePerKey());
EVALUATORS.put(Flatten.PCollections.class, flattenPColl());
EVALUATORS.put(Create.Values.class, create());
-// EVALUATORS.put(View.AsSingleton.class, viewAsSingleton());
-// EVALUATORS.put(View.AsIterable.class, viewAsIter());
+ EVALUATORS.put(View.AsSingleton.class, viewAsSingleton());
+ EVALUATORS.put(View.AsIterable.class, viewAsIter());
EVALUATORS.put(View.CreatePCollectionView.class, createPCollView());
EVALUATORS.put(Window.Assign.class, window());
EVALUATORS.put(Reshuffle.class, reshuffle());
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/GlobalWatermarkHolder.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/GlobalWatermarkHolder.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/util/GlobalWatermarkHolder.java
index 2cb6f26..8b384d8 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/GlobalWatermarkHolder.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/util/GlobalWatermarkHolder.java
@@ -21,43 +21,31 @@ package org.apache.beam.runners.spark.util;
import static com.google.common.base.Preconditions.checkState;
import com.google.common.annotations.VisibleForTesting;
-import com.google.common.cache.CacheBuilder;
-import com.google.common.cache.CacheLoader;
-import com.google.common.cache.LoadingCache;
-import com.google.common.collect.Maps;
import java.io.Serializable;
import java.util.HashMap;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeUnit;
-import javax.annotation.Nonnull;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.spark.SparkEnv;
+import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.broadcast.Broadcast;
-import org.apache.spark.storage.BlockId;
-import org.apache.spark.storage.BlockManager;
-import org.apache.spark.storage.BlockResult;
-import org.apache.spark.storage.BlockStore;
-import org.apache.spark.storage.StorageLevel;
+import org.apache.spark.streaming.api.java.JavaStreamingContext;
import org.apache.spark.streaming.api.java.JavaStreamingListener;
import org.apache.spark.streaming.api.java.JavaStreamingListenerBatchCompleted;
import org.joda.time.Instant;
-import scala.Option;
+
/**
- * A {@link BlockStore} variable to hold the global watermarks for a micro-batch.
+ * A {@link Broadcast} variable to hold the global watermarks for a micro-batch.
*
* <p>For each source, holds a queue for the watermarks of each micro-batch that was read,
* and advances the watermarks according to the queue (first-in-first-out).
*/
public class GlobalWatermarkHolder {
- private static final Map<Integer, Queue<SparkWatermarks>> sourceTimes = new HashMap<>();
- private static final BlockId WATERMARKS_BLOCK_ID = BlockId.apply("broadcast_0WATERMARKS");
-
- private static volatile Map<Integer, SparkWatermarks> driverWatermarks = null;
- private static volatile LoadingCache<String, Map<Integer, SparkWatermarks>> watermarkCache = null;
+ // the broadcast is broadcasted to the workers.
+ private static volatile Broadcast<Map<Integer, SparkWatermarks>> broadcast = null;
+ // this should only live in the driver so transient.
+ private static final transient Map<Integer, Queue<SparkWatermarks>> sourceTimes = new HashMap<>();
public static void add(int sourceId, SparkWatermarks sparkWatermarks) {
Queue<SparkWatermarks> timesQueue = sourceTimes.get(sourceId);
@@ -83,48 +71,22 @@ public class GlobalWatermarkHolder {
* Returns the {@link Broadcast} containing the {@link SparkWatermarks} mapped
* to their sources.
*/
- @SuppressWarnings("unchecked")
- public static Map<Integer, SparkWatermarks> get(Long cacheInterval) {
- if (driverWatermarks != null) {
- // if we are executing in local mode simply return the local values.
- return driverWatermarks;
- } else {
- if (watermarkCache == null) {
- initWatermarkCache(cacheInterval);
- }
- try {
- return watermarkCache.get("SINGLETON");
- } catch (ExecutionException e) {
- throw new RuntimeException(e);
- }
- }
- }
-
- private static synchronized void initWatermarkCache(Long batchDuration) {
- if (watermarkCache == null) {
- watermarkCache =
- CacheBuilder.newBuilder()
- // expire watermarks every half batch duration to ensure they update in every batch.
- .expireAfterWrite(batchDuration / 2, TimeUnit.MILLISECONDS)
- .build(new WatermarksLoader());
- }
+ public static Broadcast<Map<Integer, SparkWatermarks>> get() {
+ return broadcast;
}
/**
* Advances the watermarks to the next-in-line watermarks.
* SparkWatermarks are monotonically increasing.
*/
- @SuppressWarnings("unchecked")
- public static void advance() {
- synchronized (GlobalWatermarkHolder.class) {
- BlockManager blockManager = SparkEnv.get().blockManager();
-
+ public static void advance(JavaSparkContext jsc) {
+ synchronized (GlobalWatermarkHolder.class){
if (sourceTimes.isEmpty()) {
return;
}
// update all sources' watermarks into the new broadcast.
- Map<Integer, SparkWatermarks> newValues = new HashMap<>();
+ Map<Integer, SparkWatermarks> newBroadcast = new HashMap<>();
for (Map.Entry<Integer, Queue<SparkWatermarks>> en: sourceTimes.entrySet()) {
if (en.getValue().isEmpty()) {
@@ -137,22 +99,8 @@ public class GlobalWatermarkHolder {
Instant currentLowWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE;
Instant currentHighWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE;
Instant currentSynchronizedProcessingTime = BoundedWindow.TIMESTAMP_MIN_VALUE;
-
- Option<BlockResult> currentOption = blockManager.getRemote(WATERMARKS_BLOCK_ID);
- Map<Integer, SparkWatermarks> current;
- if (currentOption.isDefined()) {
- current = (Map<Integer, SparkWatermarks>) currentOption.get().data().next();
- } else {
- current = Maps.newHashMap();
- blockManager.putSingle(
- WATERMARKS_BLOCK_ID,
- current,
- StorageLevel.MEMORY_ONLY(),
- true);
- }
-
- if (current.containsKey(sourceId)) {
- SparkWatermarks currentTimes = current.get(sourceId);
+ if (broadcast != null && broadcast.getValue().containsKey(sourceId)) {
+ SparkWatermarks currentTimes = broadcast.getValue().get(sourceId);
currentLowWatermark = currentTimes.getLowWatermark();
currentHighWatermark = currentTimes.getHighWatermark();
currentSynchronizedProcessingTime = currentTimes.getSynchronizedProcessingTime();
@@ -171,21 +119,20 @@ public class GlobalWatermarkHolder {
nextLowWatermark, nextHighWatermark));
checkState(nextSynchronizedProcessingTime.isAfter(currentSynchronizedProcessingTime),
"Synchronized processing time must advance.");
- newValues.put(
+ newBroadcast.put(
sourceId,
new SparkWatermarks(
nextLowWatermark, nextHighWatermark, nextSynchronizedProcessingTime));
}
// update the watermarks broadcast only if something has changed.
- if (!newValues.isEmpty()) {
- driverWatermarks = newValues;
- blockManager.removeBlock(WATERMARKS_BLOCK_ID, true);
- blockManager.putSingle(
- WATERMARKS_BLOCK_ID,
- newValues,
- StorageLevel.MEMORY_ONLY(),
- true);
+ if (!newBroadcast.isEmpty()) {
+ if (broadcast != null) {
+ // for now this is blocking, we could make this asynchronous
+ // but it could slow down WM propagation.
+ broadcast.destroy();
+ }
+ broadcast = jsc.broadcast(newBroadcast);
}
}
}
@@ -193,12 +140,7 @@ public class GlobalWatermarkHolder {
@VisibleForTesting
public static synchronized void clear() {
sourceTimes.clear();
- driverWatermarks = null;
- SparkEnv sparkEnv = SparkEnv.get();
- if (sparkEnv != null) {
- BlockManager blockManager = sparkEnv.blockManager();
- blockManager.removeBlock(WATERMARKS_BLOCK_ID, true);
- }
+ broadcast = null;
}
/**
@@ -243,24 +185,15 @@ public class GlobalWatermarkHolder {
/** Advance the WMs onBatchCompleted event. */
public static class WatermarksListener extends JavaStreamingListener {
- @Override
- public void onBatchCompleted(JavaStreamingListenerBatchCompleted batchCompleted) {
- GlobalWatermarkHolder.advance();
- }
- }
+ private final JavaStreamingContext jssc;
- private static class WatermarksLoader extends CacheLoader<String, Map<Integer, SparkWatermarks>> {
+ public WatermarksListener(JavaStreamingContext jssc) {
+ this.jssc = jssc;
+ }
- @SuppressWarnings("unchecked")
@Override
- public Map<Integer, SparkWatermarks> load(@Nonnull String key) throws Exception {
- Option<BlockResult> blockResultOption =
- SparkEnv.get().blockManager().getRemote(WATERMARKS_BLOCK_ID);
- if (blockResultOption.isDefined()) {
- return (Map<Integer, SparkWatermarks>) blockResultOption.get().data().next();
- } else {
- return Maps.newHashMap();
- }
+ public void onBatchCompleted(JavaStreamingListenerBatchCompleted batchCompleted) {
+ GlobalWatermarkHolder.advance(jssc.sparkContext());
}
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/spark/src/test/java/org/apache/beam/runners/spark/GlobalWatermarkHolderTest.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/GlobalWatermarkHolderTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/GlobalWatermarkHolderTest.java
index 1708123..47a6e3f 100644
--- a/runners/spark/src/test/java/org/apache/beam/runners/spark/GlobalWatermarkHolderTest.java
+++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/GlobalWatermarkHolderTest.java
@@ -65,17 +65,17 @@ public class GlobalWatermarkHolderTest {
instant.plus(Duration.millis(5)),
instant.plus(Duration.millis(5)),
instant));
- GlobalWatermarkHolder.advance();
+ GlobalWatermarkHolder.advance(jsc);
// low < high.
GlobalWatermarkHolder.add(1,
new SparkWatermarks(
instant.plus(Duration.millis(10)),
instant.plus(Duration.millis(15)),
instant.plus(Duration.millis(100))));
- GlobalWatermarkHolder.advance();
+ GlobalWatermarkHolder.advance(jsc);
// assert watermarks in Broadcast.
- SparkWatermarks currentWatermarks = GlobalWatermarkHolder.get(0L).get(1);
+ SparkWatermarks currentWatermarks = GlobalWatermarkHolder.get().getValue().get(1);
assertThat(currentWatermarks.getLowWatermark(), equalTo(instant.plus(Duration.millis(10))));
assertThat(currentWatermarks.getHighWatermark(), equalTo(instant.plus(Duration.millis(15))));
assertThat(currentWatermarks.getSynchronizedProcessingTime(),
@@ -93,7 +93,7 @@ public class GlobalWatermarkHolderTest {
instant.plus(Duration.millis(25)),
instant.plus(Duration.millis(20)),
instant.plus(Duration.millis(200))));
- GlobalWatermarkHolder.advance();
+ GlobalWatermarkHolder.advance(jsc);
}
@Test
@@ -106,7 +106,7 @@ public class GlobalWatermarkHolderTest {
instant.plus(Duration.millis(5)),
instant.plus(Duration.millis(10)),
instant));
- GlobalWatermarkHolder.advance();
+ GlobalWatermarkHolder.advance(jsc);
thrown.expect(IllegalStateException.class);
thrown.expectMessage("Synchronized processing time must advance.");
@@ -117,7 +117,7 @@ public class GlobalWatermarkHolderTest {
instant.plus(Duration.millis(5)),
instant.plus(Duration.millis(10)),
instant));
- GlobalWatermarkHolder.advance();
+ GlobalWatermarkHolder.advance(jsc);
}
@Test
@@ -136,15 +136,15 @@ public class GlobalWatermarkHolderTest {
instant.plus(Duration.millis(6)),
instant));
- GlobalWatermarkHolder.advance();
+ GlobalWatermarkHolder.advance(jsc);
// assert watermarks for source 1.
- SparkWatermarks watermarksForSource1 = GlobalWatermarkHolder.get(0L).get(1);
+ SparkWatermarks watermarksForSource1 = GlobalWatermarkHolder.get().getValue().get(1);
assertThat(watermarksForSource1.getLowWatermark(), equalTo(instant.plus(Duration.millis(5))));
assertThat(watermarksForSource1.getHighWatermark(), equalTo(instant.plus(Duration.millis(10))));
// assert watermarks for source 2.
- SparkWatermarks watermarksForSource2 = GlobalWatermarkHolder.get(0L).get(2);
+ SparkWatermarks watermarksForSource2 = GlobalWatermarkHolder.get().getValue().get(2);
assertThat(watermarksForSource2.getLowWatermark(), equalTo(instant.plus(Duration.millis(3))));
assertThat(watermarksForSource2.getHighWatermark(), equalTo(instant.plus(Duration.millis(6))));
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/spark/src/test/java/org/apache/beam/runners/spark/SparkRunnerDebuggerTest.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/SparkRunnerDebuggerTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/SparkRunnerDebuggerTest.java
index 246eb81..64ff98c 100644
--- a/runners/spark/src/test/java/org/apache/beam/runners/spark/SparkRunnerDebuggerTest.java
+++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/SparkRunnerDebuggerTest.java
@@ -52,6 +52,7 @@ import org.hamcrest.Matchers;
import org.joda.time.Duration;
import org.junit.Test;
+
/**
* Test {@link SparkRunnerDebugger} with different pipelines.
*/
@@ -84,20 +85,17 @@ public class SparkRunnerDebuggerTest {
.apply(MapElements.via(new WordCount.FormatAsTextFn()))
.apply(TextIO.write().to("!!PLACEHOLDER-OUTPUT-DIR!!").withNumShards(3).withSuffix(".txt"));
- final String expectedPipeline =
- "sparkContext.parallelize(Arrays.asList(...))\n"
- + "_.mapPartitions("
- + "new org.apache.beam.runners.spark.examples.WordCount$ExtractWordsFn())\n"
- + "_.mapPartitions(new org.apache.beam.sdk.transforms.Count$PerElement$1())\n"
- + "_.combineByKey(..., new org.apache.beam.sdk.transforms.Count$CountFn(), ...)\n"
- + "_.groupByKey()\n"
- + "_.map(new org.apache.beam.sdk.transforms.Sum$SumLongFn())\n"
- + "_.mapPartitions(new org.apache.beam.runners.spark"
- + ".SparkRunnerDebuggerTest$PlusOne())\n"
- + "sparkContext.union(...)\n"
- + "_.mapPartitions("
- + "new org.apache.beam.runners.spark.examples.WordCount$FormatAsTextFn())\n"
- + "_.<org.apache.beam.sdk.io.TextIO$Write>";
+ final String expectedPipeline = "sparkContext.parallelize(Arrays.asList(...))\n"
+ + "_.mapPartitions(new org.apache.beam.runners.spark.examples.WordCount$ExtractWordsFn())\n"
+ + "_.mapPartitions(new org.apache.beam.sdk.transforms.Count$PerElement$1())\n"
+ + "_.combineByKey(..., new org.apache.beam.sdk.transforms.Count$CountFn(), ...)\n"
+ + "_.groupByKey()\n"
+ + "_.map(new org.apache.beam.sdk.transforms.Sum$SumLongFn())\n"
+ + "_.mapPartitions(new org.apache.beam.runners.spark"
+ + ".SparkRunnerDebuggerTest$PlusOne())\n"
+ + "sparkContext.union(...)\n"
+ + "_.mapPartitions(new org.apache.beam.runners.spark.examples.WordCount$FormatAsTextFn())\n"
+ + "_.<org.apache.beam.sdk.io.AutoValue_TextIO_Write>";
SparkRunnerDebugger.DebugSparkPipelineResult result =
(SparkRunnerDebugger.DebugSparkPipelineResult) pipeline.run();
[02/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/portability/fn_api_runner.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/portability/fn_api_runner.py b/sdks/python/apache_beam/runners/portability/fn_api_runner.py
index f88fe53..db34ef9 100644
--- a/sdks/python/apache_beam/runners/portability/fn_api_runner.py
+++ b/sdks/python/apache_beam/runners/portability/fn_api_runner.py
@@ -17,32 +17,27 @@
"""A PipelineRunner using the SDK harness.
"""
-import base64
import collections
+import json
import logging
import Queue as queue
import threading
from concurrent import futures
-from google.protobuf import wrappers_pb2
import grpc
-import apache_beam as beam # pylint: disable=ungrouped-imports
+import apache_beam as beam
from apache_beam.coders import WindowedValueCoder
from apache_beam.coders.coder_impl import create_InputStream
from apache_beam.coders.coder_impl import create_OutputStream
from apache_beam.internal import pickler
from apache_beam.io import iobase
from apache_beam.transforms.window import GlobalWindows
-from apache_beam.portability.api import beam_fn_api_pb2
-from apache_beam.portability.api import beam_runner_api_pb2
-from apache_beam.runners import pipeline_context
+from apache_beam.runners.api import beam_fn_api_pb2
from apache_beam.runners.portability import maptask_executor_runner
-from apache_beam.runners.worker import bundle_processor
from apache_beam.runners.worker import data_plane
from apache_beam.runners.worker import operation_specs
from apache_beam.runners.worker import sdk_worker
-from apache_beam.utils import proto_utils
# This module is experimental. No backwards-compatibility guarantees.
@@ -128,145 +123,191 @@ class FnApiRunner(maptask_executor_runner.MapTaskExecutorRunner):
def _map_task_registration(self, map_task, state_handler,
data_operation_spec):
- input_data, side_input_data, runner_sinks, process_bundle_descriptor = (
- self._map_task_to_protos(map_task, data_operation_spec))
- # Side inputs will be accessed over the state API.
- for key, elements_data in side_input_data.items():
- state_key = beam_fn_api_pb2.StateKey.MultimapSideInput(key=key)
- state_handler.Clear(state_key)
- state_handler.Append(state_key, [elements_data])
- return beam_fn_api_pb2.InstructionRequest(
- instruction_id=self._next_uid(),
- register=beam_fn_api_pb2.RegisterRequest(
- process_bundle_descriptor=[process_bundle_descriptor])
- ), runner_sinks, input_data
-
- def _map_task_to_protos(self, map_task, data_operation_spec):
input_data = {}
- side_input_data = {}
runner_sinks = {}
-
- context = pipeline_context.PipelineContext()
- transform_protos = {}
- used_pcollections = {}
-
- def uniquify(*names):
- # An injective mapping from string* to string.
- return ':'.join("%s:%d" % (name, len(name)) for name in names)
-
- def pcollection_id(op_ix, out_ix):
- if (op_ix, out_ix) not in used_pcollections:
- used_pcollections[op_ix, out_ix] = uniquify(
- map_task[op_ix][0], 'out', str(out_ix))
- return used_pcollections[op_ix, out_ix]
-
- def get_inputs(op):
- if hasattr(op, 'inputs'):
- inputs = op.inputs
- elif hasattr(op, 'input'):
- inputs = [op.input]
- else:
- inputs = []
- return {'in%s' % ix: pcollection_id(*input)
- for ix, input in enumerate(inputs)}
-
- def get_outputs(op_ix):
- op = map_task[op_ix][1]
- return {tag: pcollection_id(op_ix, out_ix)
- for out_ix, tag in enumerate(getattr(op, 'output_tags', ['out']))}
-
- def only_element(iterable):
- element, = iterable
- return element
+ transforms = []
+ transform_index_to_id = {}
+
+ # Maps coders to new coder objects and references.
+ coders = {}
+
+ def coder_id(coder):
+ if coder not in coders:
+ coders[coder] = beam_fn_api_pb2.Coder(
+ function_spec=sdk_worker.pack_function_spec_data(
+ json.dumps(coder.as_cloud_object()),
+ sdk_worker.PYTHON_CODER_URN, id=self._next_uid()))
+
+ return coders[coder].function_spec.id
+
+ def output_tags(op):
+ return getattr(op, 'output_tags', ['out'])
+
+ def as_target(op_input):
+ input_op_index, input_output_index = op_input
+ input_op = map_task[input_op_index][1]
+ return {
+ 'ignored_input_tag':
+ beam_fn_api_pb2.Target.List(target=[
+ beam_fn_api_pb2.Target(
+ primitive_transform_reference=transform_index_to_id[
+ input_op_index],
+ name=output_tags(input_op)[input_output_index])
+ ])
+ }
+
+ def outputs(op):
+ return {
+ tag: beam_fn_api_pb2.PCollection(coder_reference=coder_id(coder))
+ for tag, coder in zip(output_tags(op), op.output_coders)
+ }
for op_ix, (stage_name, operation) in enumerate(map_task):
- transform_id = uniquify(stage_name)
-
+ transform_id = transform_index_to_id[op_ix] = self._next_uid()
if isinstance(operation, operation_specs.WorkerInMemoryWrite):
# Write this data back to the runner.
- target_name = only_element(get_inputs(operation).keys())
- runner_sinks[(transform_id, target_name)] = operation
- transform_spec = beam_runner_api_pb2.FunctionSpec(
- urn=bundle_processor.DATA_OUTPUT_URN,
- parameter=proto_utils.pack_Any(data_operation_spec))
+ fn = beam_fn_api_pb2.FunctionSpec(urn=sdk_worker.DATA_OUTPUT_URN,
+ id=self._next_uid())
+ if data_operation_spec:
+ fn.data.Pack(data_operation_spec)
+ inputs = as_target(operation.input)
+ side_inputs = {}
+ runner_sinks[(transform_id, 'out')] = operation
elif isinstance(operation, operation_specs.WorkerRead):
- # A Read from an in-memory source is done over the data plane.
+ # A Read is either translated to a direct injection of windowed values
+ # into the sdk worker, or an injection of the source object into the
+ # sdk worker as data followed by an SDF that reads that source.
if (isinstance(operation.source.source,
- maptask_executor_runner.InMemorySource)
+ worker_runner_base.InMemorySource)
and isinstance(operation.source.source.default_output_coder(),
WindowedValueCoder)):
- target_name = only_element(get_outputs(op_ix).keys())
- input_data[(transform_id, target_name)] = self._reencode_elements(
- operation.source.source.read(None),
- operation.source.source.default_output_coder())
- transform_spec = beam_runner_api_pb2.FunctionSpec(
- urn=bundle_processor.DATA_INPUT_URN,
- parameter=proto_utils.pack_Any(data_operation_spec))
-
+ output_stream = create_OutputStream()
+ element_coder = (
+ operation.source.source.default_output_coder().get_impl())
+ # Re-encode the elements in the nested context and
+ # concatenate them together
+ for element in operation.source.source.read(None):
+ element_coder.encode_to_stream(element, output_stream, True)
+ target_name = self._next_uid()
+ input_data[(transform_id, target_name)] = output_stream.get()
+ fn = beam_fn_api_pb2.FunctionSpec(urn=sdk_worker.DATA_INPUT_URN,
+ id=self._next_uid())
+ if data_operation_spec:
+ fn.data.Pack(data_operation_spec)
+ inputs = {target_name: beam_fn_api_pb2.Target.List()}
+ side_inputs = {}
else:
- # Otherwise serialize the source and execute it there.
- # TODO: Use SDFs with an initial impulse.
- # The Dataflow runner harness strips the base64 encoding. do the same
- # here until we get the same thing back that we sent in.
- transform_spec = beam_runner_api_pb2.FunctionSpec(
- urn=bundle_processor.PYTHON_SOURCE_URN,
- parameter=proto_utils.pack_Any(
- wrappers_pb2.BytesValue(
- value=base64.b64decode(
- pickler.dumps(operation.source.source)))))
+ # Read the source object from the runner.
+ source_coder = beam.coders.DillCoder()
+ input_transform_id = self._next_uid()
+ output_stream = create_OutputStream()
+ source_coder.get_impl().encode_to_stream(
+ GlobalWindows.windowed_value(operation.source),
+ output_stream,
+ True)
+ target_name = self._next_uid()
+ input_data[(input_transform_id, target_name)] = output_stream.get()
+ input_ptransform = beam_fn_api_pb2.PrimitiveTransform(
+ id=input_transform_id,
+ function_spec=beam_fn_api_pb2.FunctionSpec(
+ urn=sdk_worker.DATA_INPUT_URN,
+ id=self._next_uid()),
+ # TODO(robertwb): Possible name collision.
+ step_name=stage_name + '/inject_source',
+ inputs={target_name: beam_fn_api_pb2.Target.List()},
+ outputs={
+ 'out':
+ beam_fn_api_pb2.PCollection(
+ coder_reference=coder_id(source_coder))
+ })
+ if data_operation_spec:
+ input_ptransform.function_spec.data.Pack(data_operation_spec)
+ transforms.append(input_ptransform)
+
+ # Read the elements out of the source.
+ fn = sdk_worker.pack_function_spec_data(
+ OLDE_SOURCE_SPLITTABLE_DOFN_DATA,
+ sdk_worker.PYTHON_DOFN_URN,
+ id=self._next_uid())
+ inputs = {
+ 'ignored_input_tag':
+ beam_fn_api_pb2.Target.List(target=[
+ beam_fn_api_pb2.Target(
+ primitive_transform_reference=input_transform_id,
+ name='out')
+ ])
+ }
+ side_inputs = {}
elif isinstance(operation, operation_specs.WorkerDoFn):
- # Record the contents of each side input for access via the state api.
- side_input_extras = []
+ fn = sdk_worker.pack_function_spec_data(
+ operation.serialized_fn,
+ sdk_worker.PYTHON_DOFN_URN,
+ id=self._next_uid())
+ inputs = as_target(operation.input)
+ # Store the contents of each side input for state access.
for si in operation.side_inputs:
assert isinstance(si.source, iobase.BoundedSource)
element_coder = si.source.default_output_coder()
+ view_id = self._next_uid()
# TODO(robertwb): Actually flesh out the ViewFn API.
- side_input_extras.append((si.tag, element_coder))
- side_input_data[
- bundle_processor.side_input_tag(transform_id, si.tag)] = (
- self._reencode_elements(
- si.source.read(si.source.get_range_tracker(None, None)),
- element_coder))
- augmented_serialized_fn = pickler.dumps(
- (operation.serialized_fn, side_input_extras))
- transform_spec = beam_runner_api_pb2.FunctionSpec(
- urn=bundle_processor.PYTHON_DOFN_URN,
- parameter=proto_utils.pack_Any(
- wrappers_pb2.BytesValue(value=augmented_serialized_fn)))
+ side_inputs[si.tag] = beam_fn_api_pb2.SideInput(
+ view_fn=sdk_worker.serialize_and_pack_py_fn(
+ element_coder, urn=sdk_worker.PYTHON_ITERABLE_VIEWFN_URN,
+ id=view_id))
+ # Re-encode the elements in the nested context and
+ # concatenate them together
+ output_stream = create_OutputStream()
+ for element in si.source.read(
+ si.source.get_range_tracker(None, None)):
+ element_coder.get_impl().encode_to_stream(
+ element, output_stream, True)
+ elements_data = output_stream.get()
+ state_key = beam_fn_api_pb2.StateKey(function_spec_reference=view_id)
+ state_handler.Clear(state_key)
+ state_handler.Append(
+ beam_fn_api_pb2.SimpleStateAppendRequest(
+ state_key=state_key, data=[elements_data]))
elif isinstance(operation, operation_specs.WorkerFlatten):
- # Flatten is nice and simple.
- transform_spec = beam_runner_api_pb2.FunctionSpec(
- urn=bundle_processor.IDENTITY_DOFN_URN)
+ fn = sdk_worker.pack_function_spec_data(
+ operation.serialized_fn,
+ sdk_worker.IDENTITY_DOFN_URN,
+ id=self._next_uid())
+ inputs = {
+ 'ignored_input_tag':
+ beam_fn_api_pb2.Target.List(target=[
+ beam_fn_api_pb2.Target(
+ primitive_transform_reference=transform_index_to_id[
+ input_op_index],
+ name=output_tags(map_task[input_op_index][1])[
+ input_output_index])
+ for input_op_index, input_output_index in operation.inputs
+ ])
+ }
+ side_inputs = {}
else:
- raise NotImplementedError(operation)
-
- transform_protos[transform_id] = beam_runner_api_pb2.PTransform(
- unique_name=stage_name,
- spec=transform_spec,
- inputs=get_inputs(operation),
- outputs=get_outputs(op_ix))
-
- pcollection_protos = {
- name: beam_runner_api_pb2.PCollection(
- unique_name=name,
- coder_id=context.coders.get_id(
- map_task[op_id][1].output_coders[out_id]))
- for (op_id, out_id), name in used_pcollections.items()
- }
- # Must follow creation of pcollection_protos to capture used coders.
- context_proto = context.to_runner_api()
+ raise TypeError(operation)
+
+ ptransform = beam_fn_api_pb2.PrimitiveTransform(
+ id=transform_id,
+ function_spec=fn,
+ step_name=stage_name,
+ inputs=inputs,
+ side_inputs=side_inputs,
+ outputs=outputs(operation))
+ transforms.append(ptransform)
+
process_bundle_descriptor = beam_fn_api_pb2.ProcessBundleDescriptor(
- id=self._next_uid(),
- transforms=transform_protos,
- pcollections=pcollection_protos,
- coders=dict(context_proto.coders.items()),
- windowing_strategies=dict(context_proto.windowing_strategies.items()),
- environments=dict(context_proto.environments.items()))
- return input_data, side_input_data, runner_sinks, process_bundle_descriptor
+ id=self._next_uid(), coders=coders.values(),
+ primitive_transform=transforms)
+ return beam_fn_api_pb2.InstructionRequest(
+ instruction_id=self._next_uid(),
+ register=beam_fn_api_pb2.RegisterRequest(
+ process_bundle_descriptor=[process_bundle_descriptor
+ ])), runner_sinks, input_data
def _run_map_task(
self, map_task, control_handler, state_handler, data_plane_handler,
@@ -317,7 +358,7 @@ class FnApiRunner(maptask_executor_runner.MapTaskExecutorRunner):
sink_op.output_buffer.append(e)
return
- def execute_map_tasks(self, ordered_map_tasks, direct=False):
+ def execute_map_tasks(self, ordered_map_tasks, direct=True):
if direct:
controller = FnApiRunner.DirectController()
else:
@@ -341,8 +382,9 @@ class FnApiRunner(maptask_executor_runner.MapTaskExecutorRunner):
return beam_fn_api_pb2.Elements.Data(
data=''.join(self._all[self._to_key(state_key)]))
- def Append(self, state_key, data):
- self._all[self._to_key(state_key)].extend(data)
+ def Append(self, append_request):
+ self._all[self._to_key(append_request.state_key)].extend(
+ append_request.data)
def Clear(self, state_key):
try:
@@ -352,7 +394,8 @@ class FnApiRunner(maptask_executor_runner.MapTaskExecutorRunner):
@staticmethod
def _to_key(state_key):
- return state_key.window, state_key.key
+ return (state_key.function_spec_reference, state_key.window,
+ state_key.key)
class DirectController(object):
"""An in-memory controller for fn API control, state and data planes."""
@@ -428,10 +471,3 @@ class FnApiRunner(maptask_executor_runner.MapTaskExecutorRunner):
self.data_plane_handler.close()
self.control_server.stop(5).wait()
self.data_server.stop(5).wait()
-
- @staticmethod
- def _reencode_elements(elements, element_coder):
- output_stream = create_OutputStream()
- for element in elements:
- element_coder.get_impl().encode_to_stream(element, output_stream, True)
- return output_stream.get()
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/portability/fn_api_runner_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/portability/fn_api_runner_test.py b/sdks/python/apache_beam/runners/portability/fn_api_runner_test.py
index 163e980..633602f 100644
--- a/sdks/python/apache_beam/runners/portability/fn_api_runner_test.py
+++ b/sdks/python/apache_beam/runners/portability/fn_api_runner_test.py
@@ -20,41 +20,18 @@ import unittest
import apache_beam as beam
from apache_beam.runners.portability import fn_api_runner
-from apache_beam.runners.portability import maptask_executor_runner_test
-from apache_beam.testing.util import assert_that
-from apache_beam.testing.util import equal_to
+from apache_beam.runners.portability import maptask_executor_runner
-class FnApiRunnerTest(
- maptask_executor_runner_test.MapTaskExecutorRunnerTest):
+class FnApiRunnerTest(maptask_executor_runner.MapTaskExecutorRunner):
def create_pipeline(self):
- return beam.Pipeline(
- runner=fn_api_runner.FnApiRunner())
+ return beam.Pipeline(runner=fn_api_runner.FnApiRunner())
def test_combine_per_key(self):
- # TODO(BEAM-1348): Enable once Partial GBK is supported in fn API.
+ # TODO(robertwb): Implement PGBKCV operation.
pass
- def test_combine_per_key(self):
- # TODO(BEAM-1348): Enable once Partial GBK is supported in fn API.
- pass
-
- def test_pardo_side_inputs(self):
- # TODO(BEAM-1348): Enable once side inputs are supported in fn API.
- pass
-
- def test_pardo_unfusable_side_inputs(self):
- # TODO(BEAM-1348): Enable once side inputs are supported in fn API.
- pass
-
- def test_assert_that(self):
- # TODO: figure out a way for fn_api_runner to parse and raise the
- # underlying exception.
- with self.assertRaisesRegexp(RuntimeError, 'BeamAssertException'):
- with self.create_pipeline() as p:
- assert_that(p | beam.Create(['a', 'b']), equal_to(['a']))
-
# Inherits all tests from maptask_executor_runner.MapTaskExecutorRunner
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/worker/bundle_processor.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor.py b/sdks/python/apache_beam/runners/worker/bundle_processor.py
deleted file mode 100644
index 2669bfc..0000000
--- a/sdks/python/apache_beam/runners/worker/bundle_processor.py
+++ /dev/null
@@ -1,426 +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.
-#
-
-"""SDK harness for executing Python Fns via the Fn API."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import base64
-import collections
-import json
-import logging
-
-from google.protobuf import wrappers_pb2
-
-from apache_beam.coders import coder_impl
-from apache_beam.coders import WindowedValueCoder
-from apache_beam.internal import pickler
-from apache_beam.io import iobase
-from apache_beam.portability.api import beam_fn_api_pb2
-from apache_beam.runners.dataflow.native_io import iobase as native_iobase
-from apache_beam.runners import pipeline_context
-from apache_beam.runners.worker import operation_specs
-from apache_beam.runners.worker import operations
-from apache_beam.utils import counters
-from apache_beam.utils import proto_utils
-
-# This module is experimental. No backwards-compatibility guarantees.
-
-
-try:
- from apache_beam.runners.worker import statesampler
-except ImportError:
- from apache_beam.runners.worker import statesampler_fake as statesampler
-
-
-DATA_INPUT_URN = 'urn:org.apache.beam:source:runner:0.1'
-DATA_OUTPUT_URN = 'urn:org.apache.beam:sink:runner:0.1'
-IDENTITY_DOFN_URN = 'urn:org.apache.beam:dofn:identity:0.1'
-PYTHON_ITERABLE_VIEWFN_URN = 'urn:org.apache.beam:viewfn:iterable:python:0.1'
-PYTHON_CODER_URN = 'urn:org.apache.beam:coder:python:0.1'
-# TODO(vikasrk): Fix this once runner sends appropriate python urns.
-PYTHON_DOFN_URN = 'urn:org.apache.beam:dofn:java:0.1'
-PYTHON_SOURCE_URN = 'urn:org.apache.beam:source:java:0.1'
-
-
-def side_input_tag(transform_id, tag):
- return str("%d[%s][%s]" % (len(transform_id), transform_id, tag))
-
-
-class RunnerIOOperation(operations.Operation):
- """Common baseclass for runner harness IO operations."""
-
- def __init__(self, operation_name, step_name, consumers, counter_factory,
- state_sampler, windowed_coder, target, data_channel):
- super(RunnerIOOperation, self).__init__(
- operation_name, None, counter_factory, state_sampler)
- self.windowed_coder = windowed_coder
- self.step_name = step_name
- # target represents the consumer for the bytes in the data plane for a
- # DataInputOperation or a producer of these bytes for a DataOutputOperation.
- self.target = target
- self.data_channel = data_channel
- for _, consumer_ops in consumers.items():
- for consumer in consumer_ops:
- self.add_receiver(consumer, 0)
-
-
-class DataOutputOperation(RunnerIOOperation):
- """A sink-like operation that gathers outputs to be sent back to the runner.
- """
-
- def set_output_stream(self, output_stream):
- self.output_stream = output_stream
-
- def process(self, windowed_value):
- self.windowed_coder.get_impl().encode_to_stream(
- windowed_value, self.output_stream, True)
-
- def finish(self):
- self.output_stream.close()
- super(DataOutputOperation, self).finish()
-
-
-class DataInputOperation(RunnerIOOperation):
- """A source-like operation that gathers input from the runner.
- """
-
- def __init__(self, operation_name, step_name, consumers, counter_factory,
- state_sampler, windowed_coder, input_target, data_channel):
- super(DataInputOperation, self).__init__(
- operation_name, step_name, consumers, counter_factory, state_sampler,
- windowed_coder, target=input_target, data_channel=data_channel)
- # We must do this manually as we don't have a spec or spec.output_coders.
- self.receivers = [
- operations.ConsumerSet(self.counter_factory, self.step_name, 0,
- consumers.itervalues().next(),
- self.windowed_coder)]
-
- def process(self, windowed_value):
- self.output(windowed_value)
-
- def process_encoded(self, encoded_windowed_values):
- input_stream = coder_impl.create_InputStream(encoded_windowed_values)
- while input_stream.size() > 0:
- decoded_value = self.windowed_coder.get_impl().decode_from_stream(
- input_stream, True)
- self.output(decoded_value)
-
-
-# TODO(robertwb): Revise side input API to not be in terms of native sources.
-# This will enable lookups, but there's an open question as to how to handle
-# custom sources without forcing intermediate materialization. This seems very
-# related to the desire to inject key and window preserving [Splittable]DoFns
-# into the view computation.
-class SideInputSource(native_iobase.NativeSource,
- native_iobase.NativeSourceReader):
- """A 'source' for reading side inputs via state API calls.
- """
-
- def __init__(self, state_handler, state_key, coder):
- self._state_handler = state_handler
- self._state_key = state_key
- self._coder = coder
-
- def reader(self):
- return self
-
- @property
- def returns_windowed_values(self):
- return True
-
- def __enter__(self):
- return self
-
- def __exit__(self, *exn_info):
- pass
-
- def __iter__(self):
- # TODO(robertwb): Support pagination.
- input_stream = coder_impl.create_InputStream(
- self._state_handler.Get(self._state_key).data)
- while input_stream.size() > 0:
- yield self._coder.get_impl().decode_from_stream(input_stream, True)
-
-
-def memoize(func):
- cache = {}
- missing = object()
-
- def wrapper(*args):
- result = cache.get(args, missing)
- if result is missing:
- result = cache[args] = func(*args)
- return result
- return wrapper
-
-
-def only_element(iterable):
- element, = iterable
- return element
-
-
-class BundleProcessor(object):
- """A class for processing bundles of elements.
- """
- def __init__(
- self, process_bundle_descriptor, state_handler, data_channel_factory):
- self.process_bundle_descriptor = process_bundle_descriptor
- self.state_handler = state_handler
- self.data_channel_factory = data_channel_factory
-
- def create_execution_tree(self, descriptor):
- # TODO(robertwb): Figure out the correct prefix to use for output counters
- # from StateSampler.
- counter_factory = counters.CounterFactory()
- state_sampler = statesampler.StateSampler(
- 'fnapi-step%s-' % descriptor.id, counter_factory)
-
- transform_factory = BeamTransformFactory(
- descriptor, self.data_channel_factory, counter_factory, state_sampler,
- self.state_handler)
-
- pcoll_consumers = collections.defaultdict(list)
- for transform_id, transform_proto in descriptor.transforms.items():
- for pcoll_id in transform_proto.inputs.values():
- pcoll_consumers[pcoll_id].append(transform_id)
-
- @memoize
- def get_operation(transform_id):
- transform_consumers = {
- tag: [get_operation(op) for op in pcoll_consumers[pcoll_id]]
- for tag, pcoll_id
- in descriptor.transforms[transform_id].outputs.items()
- }
- return transform_factory.create_operation(
- transform_id, transform_consumers)
-
- # Operations must be started (hence returned) in order.
- @memoize
- def topological_height(transform_id):
- return 1 + max(
- [0] +
- [topological_height(consumer)
- for pcoll in descriptor.transforms[transform_id].outputs.values()
- for consumer in pcoll_consumers[pcoll]])
-
- return [get_operation(transform_id)
- for transform_id in sorted(
- descriptor.transforms, key=topological_height, reverse=True)]
-
- def process_bundle(self, instruction_id):
- ops = self.create_execution_tree(self.process_bundle_descriptor)
-
- expected_inputs = []
- for op in ops:
- if isinstance(op, DataOutputOperation):
- # TODO(robertwb): Is there a better way to pass the instruction id to
- # the operation?
- op.set_output_stream(op.data_channel.output_stream(
- instruction_id, op.target))
- elif isinstance(op, DataInputOperation):
- # We must wait until we receive "end of stream" for each of these ops.
- expected_inputs.append(op)
-
- # Start all operations.
- for op in reversed(ops):
- logging.info('start %s', op)
- op.start()
-
- # Inject inputs from data plane.
- for input_op in expected_inputs:
- for data in input_op.data_channel.input_elements(
- instruction_id, [input_op.target]):
- # ignores input name
- input_op.process_encoded(data.data)
-
- # Finish all operations.
- for op in ops:
- logging.info('finish %s', op)
- op.finish()
-
-
-class BeamTransformFactory(object):
- """Factory for turning transform_protos into executable operations."""
- def __init__(self, descriptor, data_channel_factory, counter_factory,
- state_sampler, state_handler):
- self.descriptor = descriptor
- self.data_channel_factory = data_channel_factory
- self.counter_factory = counter_factory
- self.state_sampler = state_sampler
- self.state_handler = state_handler
- self.context = pipeline_context.PipelineContext(descriptor)
-
- _known_urns = {}
-
- @classmethod
- def register_urn(cls, urn, parameter_type):
- def wrapper(func):
- cls._known_urns[urn] = func, parameter_type
- return func
- return wrapper
-
- def create_operation(self, transform_id, consumers):
- transform_proto = self.descriptor.transforms[transform_id]
- creator, parameter_type = self._known_urns[transform_proto.spec.urn]
- parameter = proto_utils.unpack_Any(
- transform_proto.spec.parameter, parameter_type)
- return creator(self, transform_id, transform_proto, parameter, consumers)
-
- def get_coder(self, coder_id):
- coder_proto = self.descriptor.coders[coder_id]
- if coder_proto.spec.spec.urn:
- return self.context.coders.get_by_id(coder_id)
- else:
- # No URN, assume cloud object encoding json bytes.
- return operation_specs.get_coder_from_spec(
- json.loads(
- proto_utils.unpack_Any(coder_proto.spec.spec.parameter,
- wrappers_pb2.BytesValue).value))
-
- def get_output_coders(self, transform_proto):
- return {
- tag: self.get_coder(self.descriptor.pcollections[pcoll_id].coder_id)
- for tag, pcoll_id in transform_proto.outputs.items()
- }
-
- def get_only_output_coder(self, transform_proto):
- return only_element(self.get_output_coders(transform_proto).values())
-
- def get_input_coders(self, transform_proto):
- return {
- tag: self.get_coder(self.descriptor.pcollections[pcoll_id].coder_id)
- for tag, pcoll_id in transform_proto.inputs.items()
- }
-
- def get_only_input_coder(self, transform_proto):
- return only_element(self.get_input_coders(transform_proto).values())
-
- # TODO(robertwb): Update all operations to take these in the constructor.
- @staticmethod
- def augment_oldstyle_op(op, step_name, consumers, tag_list=None):
- op.step_name = step_name
- for tag, op_consumers in consumers.items():
- for consumer in op_consumers:
- op.add_receiver(consumer, tag_list.index(tag) if tag_list else 0)
- return op
-
-
-@BeamTransformFactory.register_urn(
- DATA_INPUT_URN, beam_fn_api_pb2.RemoteGrpcPort)
-def create(factory, transform_id, transform_proto, grpc_port, consumers):
- target = beam_fn_api_pb2.Target(
- primitive_transform_reference=transform_id,
- name=only_element(transform_proto.outputs.keys()))
- return DataInputOperation(
- transform_proto.unique_name,
- transform_proto.unique_name,
- consumers,
- factory.counter_factory,
- factory.state_sampler,
- factory.get_only_output_coder(transform_proto),
- input_target=target,
- data_channel=factory.data_channel_factory.create_data_channel(grpc_port))
-
-
-@BeamTransformFactory.register_urn(
- DATA_OUTPUT_URN, beam_fn_api_pb2.RemoteGrpcPort)
-def create(factory, transform_id, transform_proto, grpc_port, consumers):
- target = beam_fn_api_pb2.Target(
- primitive_transform_reference=transform_id,
- name=only_element(transform_proto.inputs.keys()))
- return DataOutputOperation(
- transform_proto.unique_name,
- transform_proto.unique_name,
- consumers,
- factory.counter_factory,
- factory.state_sampler,
- # TODO(robertwb): Perhaps this could be distinct from the input coder?
- factory.get_only_input_coder(transform_proto),
- target=target,
- data_channel=factory.data_channel_factory.create_data_channel(grpc_port))
-
-
-@BeamTransformFactory.register_urn(PYTHON_SOURCE_URN, wrappers_pb2.BytesValue)
-def create(factory, transform_id, transform_proto, parameter, consumers):
- # The Dataflow runner harness strips the base64 encoding.
- source = pickler.loads(base64.b64encode(parameter.value))
- spec = operation_specs.WorkerRead(
- iobase.SourceBundle(1.0, source, None, None),
- [WindowedValueCoder(source.default_output_coder())])
- return factory.augment_oldstyle_op(
- operations.ReadOperation(
- transform_proto.unique_name,
- spec,
- factory.counter_factory,
- factory.state_sampler),
- transform_proto.unique_name,
- consumers)
-
-
-@BeamTransformFactory.register_urn(PYTHON_DOFN_URN, wrappers_pb2.BytesValue)
-def create(factory, transform_id, transform_proto, parameter, consumers):
- dofn_data = pickler.loads(parameter.value)
- if len(dofn_data) == 2:
- # Has side input data.
- serialized_fn, side_input_data = dofn_data
- else:
- # No side input data.
- serialized_fn, side_input_data = parameter.value, []
-
- def create_side_input(tag, coder):
- # TODO(robertwb): Extract windows (and keys) out of element data.
- # TODO(robertwb): Extract state key from ParDoPayload.
- return operation_specs.WorkerSideInputSource(
- tag=tag,
- source=SideInputSource(
- factory.state_handler,
- beam_fn_api_pb2.StateKey.MultimapSideInput(
- key=side_input_tag(transform_id, tag)),
- coder=coder))
- output_tags = list(transform_proto.outputs.keys())
- output_coders = factory.get_output_coders(transform_proto)
- spec = operation_specs.WorkerDoFn(
- serialized_fn=serialized_fn,
- output_tags=output_tags,
- input=None,
- side_inputs=[
- create_side_input(tag, coder) for tag, coder in side_input_data],
- output_coders=[output_coders[tag] for tag in output_tags])
- return factory.augment_oldstyle_op(
- operations.DoOperation(
- transform_proto.unique_name,
- spec,
- factory.counter_factory,
- factory.state_sampler),
- transform_proto.unique_name,
- consumers,
- output_tags)
-
-
-@BeamTransformFactory.register_urn(IDENTITY_DOFN_URN, None)
-def create(factory, transform_id, transform_proto, unused_parameter, consumers):
- return factory.augment_oldstyle_op(
- operations.FlattenOperation(
- transform_proto.unique_name,
- None,
- factory.counter_factory,
- factory.state_sampler),
- transform_proto.unique_name,
- consumers)
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/worker/data_plane.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/worker/data_plane.py b/sdks/python/apache_beam/runners/worker/data_plane.py
index 26f65ee..5edd0b4 100644
--- a/sdks/python/apache_beam/runners/worker/data_plane.py
+++ b/sdks/python/apache_beam/runners/worker/data_plane.py
@@ -28,7 +28,7 @@ import Queue as queue
import threading
from apache_beam.coders import coder_impl
-from apache_beam.portability.api import beam_fn_api_pb2
+from apache_beam.runners.api import beam_fn_api_pb2
import grpc
# This module is experimental. No backwards-compatibility guarantees.
@@ -167,18 +167,12 @@ class _GrpcDataChannel(DataChannel):
yield data
def output_stream(self, instruction_id, target):
- # TODO: Return an output stream that sends data
- # to the Runner once a fixed size buffer is full.
- # Currently we buffer all the data before sending
- # any messages.
def add_to_send_queue(data):
- if data:
- self._to_send.put(
- beam_fn_api_pb2.Elements.Data(
- instruction_reference=instruction_id,
- target=target,
- data=data))
- # End of stream marker.
+ self._to_send.put(
+ beam_fn_api_pb2.Elements.Data(
+ instruction_reference=instruction_id,
+ target=target,
+ data=data))
self._to_send.put(
beam_fn_api_pb2.Elements.Data(
instruction_reference=instruction_id,
@@ -246,8 +240,8 @@ class DataChannelFactory(object):
__metaclass__ = abc.ABCMeta
@abc.abstractmethod
- def create_data_channel(self, remote_grpc_port):
- """Returns a ``DataChannel`` from the given RemoteGrpcPort."""
+ def create_data_channel(self, function_spec):
+ """Returns a ``DataChannel`` from the given function_spec."""
raise NotImplementedError(type(self))
@abc.abstractmethod
@@ -265,7 +259,9 @@ class GrpcClientDataChannelFactory(DataChannelFactory):
def __init__(self):
self._data_channel_cache = {}
- def create_data_channel(self, remote_grpc_port):
+ def create_data_channel(self, function_spec):
+ remote_grpc_port = beam_fn_api_pb2.RemoteGrpcPort()
+ function_spec.data.Unpack(remote_grpc_port)
url = remote_grpc_port.api_service_descriptor.url
if url not in self._data_channel_cache:
logging.info('Creating channel for %s', url)
@@ -287,7 +283,7 @@ class InMemoryDataChannelFactory(DataChannelFactory):
def __init__(self, in_memory_data_channel):
self._in_memory_data_channel = in_memory_data_channel
- def create_data_channel(self, unused_remote_grpc_port):
+ def create_data_channel(self, unused_function_spec):
return self._in_memory_data_channel
def close(self):
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/worker/data_plane_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/worker/data_plane_test.py b/sdks/python/apache_beam/runners/worker/data_plane_test.py
index 360468a..e3e01ac 100644
--- a/sdks/python/apache_beam/runners/worker/data_plane_test.py
+++ b/sdks/python/apache_beam/runners/worker/data_plane_test.py
@@ -29,7 +29,7 @@ import unittest
from concurrent import futures
import grpc
-from apache_beam.portability.api import beam_fn_api_pb2
+from apache_beam.runners.api import beam_fn_api_pb2
from apache_beam.runners.worker import data_plane
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/worker/log_handler.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/worker/log_handler.py b/sdks/python/apache_beam/runners/worker/log_handler.py
index b8f6352..59ffbf4 100644
--- a/sdks/python/apache_beam/runners/worker/log_handler.py
+++ b/sdks/python/apache_beam/runners/worker/log_handler.py
@@ -21,7 +21,7 @@ import math
import Queue as queue
import threading
-from apache_beam.portability.api import beam_fn_api_pb2
+from apache_beam.runners.api import beam_fn_api_pb2
import grpc
# This module is experimental. No backwards-compatibility guarantees.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/worker/log_handler_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/worker/log_handler_test.py b/sdks/python/apache_beam/runners/worker/log_handler_test.py
index 2256bb5..8720ca8 100644
--- a/sdks/python/apache_beam/runners/worker/log_handler_test.py
+++ b/sdks/python/apache_beam/runners/worker/log_handler_test.py
@@ -22,7 +22,7 @@ import unittest
from concurrent import futures
import grpc
-from apache_beam.portability.api import beam_fn_api_pb2
+from apache_beam.runners.api import beam_fn_api_pb2
from apache_beam.runners.worker import log_handler
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/worker/operation_specs.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/worker/operation_specs.py b/sdks/python/apache_beam/runners/worker/operation_specs.py
index bdafbea..db5eb76 100644
--- a/sdks/python/apache_beam/runners/worker/operation_specs.py
+++ b/sdks/python/apache_beam/runners/worker/operation_specs.py
@@ -326,12 +326,11 @@ def get_coder_from_spec(coder_spec):
assert len(coder_spec['component_encodings']) == 2
value_coder, window_coder = [
get_coder_from_spec(c) for c in coder_spec['component_encodings']]
- return coders.coders.WindowedValueCoder(
- value_coder, window_coder=window_coder)
+ return coders.WindowedValueCoder(value_coder, window_coder=window_coder)
elif coder_spec['@type'] == 'kind:interval_window':
assert ('component_encodings' not in coder_spec
or not coder_spec['component_encodings'])
- return coders.coders.IntervalWindowCoder()
+ return coders.IntervalWindowCoder()
elif coder_spec['@type'] == 'kind:global_window':
assert ('component_encodings' not in coder_spec
or not coder_spec['component_encodings'])
@@ -340,10 +339,6 @@ def get_coder_from_spec(coder_spec):
assert len(coder_spec['component_encodings']) == 1
return coders.coders.LengthPrefixCoder(
get_coder_from_spec(coder_spec['component_encodings'][0]))
- elif coder_spec['@type'] == 'kind:bytes':
- assert ('component_encodings' not in coder_spec
- or len(coder_spec['component_encodings'] == 0))
- return coders.BytesCoder()
# We pass coders in the form "<coder_name>$<pickled_data>" to make the job
# description JSON more readable.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/worker/operations.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/worker/operations.py b/sdks/python/apache_beam/runners/worker/operations.py
index c4f945b..a44561d 100644
--- a/sdks/python/apache_beam/runners/worker/operations.py
+++ b/sdks/python/apache_beam/runners/worker/operations.py
@@ -129,7 +129,6 @@ class Operation(object):
self.operation_name + '-finish')
# TODO(ccy): the '-abort' state can be added when the abort is supported in
# Operations.
- self.scoped_metrics_container = None
def start(self):
"""Start operation."""
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/worker/sdk_worker.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker.py b/sdks/python/apache_beam/runners/worker/sdk_worker.py
index 6a23680..33c50ad 100644
--- a/sdks/python/apache_beam/runners/worker/sdk_worker.py
+++ b/sdks/python/apache_beam/runners/worker/sdk_worker.py
@@ -21,21 +21,198 @@ from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
+import collections
+import json
import logging
import Queue as queue
import threading
import traceback
+import zlib
-from apache_beam.portability.api import beam_fn_api_pb2
-from apache_beam.runners.worker import bundle_processor
-from apache_beam.runners.worker import data_plane
+import dill
+from google.protobuf import wrappers_pb2
+
+from apache_beam.coders import coder_impl
+from apache_beam.coders import WindowedValueCoder
+from apache_beam.internal import pickler
+from apache_beam.io import iobase
+from apache_beam.runners.dataflow.native_io import iobase as native_iobase
+from apache_beam.utils import counters
+from apache_beam.runners.api import beam_fn_api_pb2
+from apache_beam.runners.worker import operation_specs
+from apache_beam.runners.worker import operations
+
+# This module is experimental. No backwards-compatibility guarantees.
+
+
+try:
+ from apache_beam.runners.worker import statesampler
+except ImportError:
+ from apache_beam.runners.worker import statesampler_fake as statesampler
+from apache_beam.runners.worker.data_plane import GrpcClientDataChannelFactory
+
+
+DATA_INPUT_URN = 'urn:org.apache.beam:source:runner:0.1'
+DATA_OUTPUT_URN = 'urn:org.apache.beam:sink:runner:0.1'
+IDENTITY_DOFN_URN = 'urn:org.apache.beam:dofn:identity:0.1'
+PYTHON_ITERABLE_VIEWFN_URN = 'urn:org.apache.beam:viewfn:iterable:python:0.1'
+PYTHON_CODER_URN = 'urn:org.apache.beam:coder:python:0.1'
+# TODO(vikasrk): Fix this once runner sends appropriate python urns.
+PYTHON_DOFN_URN = 'urn:org.apache.beam:dofn:java:0.1'
+PYTHON_SOURCE_URN = 'urn:org.apache.beam:source:java:0.1'
+
+
+class RunnerIOOperation(operations.Operation):
+ """Common baseclass for runner harness IO operations."""
+
+ def __init__(self, operation_name, step_name, consumers, counter_factory,
+ state_sampler, windowed_coder, target, data_channel):
+ super(RunnerIOOperation, self).__init__(
+ operation_name, None, counter_factory, state_sampler)
+ self.windowed_coder = windowed_coder
+ self.step_name = step_name
+ # target represents the consumer for the bytes in the data plane for a
+ # DataInputOperation or a producer of these bytes for a DataOutputOperation.
+ self.target = target
+ self.data_channel = data_channel
+ for _, consumer_ops in consumers.items():
+ for consumer in consumer_ops:
+ self.add_receiver(consumer, 0)
+
+
+class DataOutputOperation(RunnerIOOperation):
+ """A sink-like operation that gathers outputs to be sent back to the runner.
+ """
+
+ def set_output_stream(self, output_stream):
+ self.output_stream = output_stream
+
+ def process(self, windowed_value):
+ self.windowed_coder.get_impl().encode_to_stream(
+ windowed_value, self.output_stream, True)
+
+ def finish(self):
+ self.output_stream.close()
+ super(DataOutputOperation, self).finish()
+
+
+class DataInputOperation(RunnerIOOperation):
+ """A source-like operation that gathers input from the runner.
+ """
+
+ def __init__(self, operation_name, step_name, consumers, counter_factory,
+ state_sampler, windowed_coder, input_target, data_channel):
+ super(DataInputOperation, self).__init__(
+ operation_name, step_name, consumers, counter_factory, state_sampler,
+ windowed_coder, target=input_target, data_channel=data_channel)
+ # We must do this manually as we don't have a spec or spec.output_coders.
+ self.receivers = [
+ operations.ConsumerSet(self.counter_factory, self.step_name, 0,
+ consumers.itervalues().next(),
+ self.windowed_coder)]
+
+ def process(self, windowed_value):
+ self.output(windowed_value)
+
+ def process_encoded(self, encoded_windowed_values):
+ input_stream = coder_impl.create_InputStream(encoded_windowed_values)
+ while input_stream.size() > 0:
+ decoded_value = self.windowed_coder.get_impl().decode_from_stream(
+ input_stream, True)
+ self.output(decoded_value)
+
+
+# TODO(robertwb): Revise side input API to not be in terms of native sources.
+# This will enable lookups, but there's an open question as to how to handle
+# custom sources without forcing intermediate materialization. This seems very
+# related to the desire to inject key and window preserving [Splittable]DoFns
+# into the view computation.
+class SideInputSource(native_iobase.NativeSource,
+ native_iobase.NativeSourceReader):
+ """A 'source' for reading side inputs via state API calls.
+ """
+
+ def __init__(self, state_handler, state_key, coder):
+ self._state_handler = state_handler
+ self._state_key = state_key
+ self._coder = coder
+
+ def reader(self):
+ return self
+
+ @property
+ def returns_windowed_values(self):
+ return True
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *exn_info):
+ pass
+
+ def __iter__(self):
+ # TODO(robertwb): Support pagination.
+ input_stream = coder_impl.create_InputStream(
+ self._state_handler.Get(self._state_key).data)
+ while input_stream.size() > 0:
+ yield self._coder.get_impl().decode_from_stream(input_stream, True)
+
+
+def unpack_and_deserialize_py_fn(function_spec):
+ """Returns unpacked and deserialized object from function spec proto."""
+ return pickler.loads(unpack_function_spec_data(function_spec))
+
+
+def unpack_function_spec_data(function_spec):
+ """Returns unpacked data from function spec proto."""
+ data = wrappers_pb2.BytesValue()
+ function_spec.data.Unpack(data)
+ return data.value
+
+
+# pylint: disable=redefined-builtin
+def serialize_and_pack_py_fn(fn, urn, id=None):
+ """Returns serialized and packed function in a function spec proto."""
+ return pack_function_spec_data(pickler.dumps(fn), urn, id)
+# pylint: enable=redefined-builtin
+
+
+# pylint: disable=redefined-builtin
+def pack_function_spec_data(value, urn, id=None):
+ """Returns packed data in a function spec proto."""
+ data = wrappers_pb2.BytesValue(value=value)
+ fn_proto = beam_fn_api_pb2.FunctionSpec(urn=urn)
+ fn_proto.data.Pack(data)
+ if id:
+ fn_proto.id = id
+ return fn_proto
+# pylint: enable=redefined-builtin
+
+
+# TODO(vikasrk): move this method to ``coders.py`` in the SDK.
+def load_compressed(compressed_data):
+ """Returns a decompressed and deserialized python object."""
+ # Note: SDK uses ``pickler.dumps`` to serialize certain python objects
+ # (like sources), which involves serialization, compression and base64
+ # encoding. We cannot directly use ``pickler.loads`` for
+ # deserialization, as the runner would have already base64 decoded the
+ # data. So we only need to decompress and deserialize.
+
+ data = zlib.decompress(compressed_data)
+ try:
+ return dill.loads(data)
+ except Exception: # pylint: disable=broad-except
+ dill.dill._trace(True) # pylint: disable=protected-access
+ return dill.loads(data)
+ finally:
+ dill.dill._trace(False) # pylint: disable=protected-access
class SdkHarness(object):
def __init__(self, control_channel):
self._control_channel = control_channel
- self._data_channel_factory = data_plane.GrpcClientDataChannelFactory()
+ self._data_channel_factory = GrpcClientDataChannelFactory()
def run(self):
contol_stub = beam_fn_api_pb2.BeamFnControlStub(self._control_channel)
@@ -59,10 +236,6 @@ class SdkHarness(object):
try:
response = self.worker.do_instruction(work_request)
except Exception: # pylint: disable=broad-except
- logging.error(
- 'Error processing instruction %s',
- work_request.instruction_id,
- exc_info=True)
response = beam_fn_api_pb2.InstructionResponse(
instruction_id=work_request.instruction_id,
error=traceback.format_exc())
@@ -100,12 +273,185 @@ class SdkWorker(object):
def register(self, request, unused_instruction_id=None):
for process_bundle_descriptor in request.process_bundle_descriptor:
self.fns[process_bundle_descriptor.id] = process_bundle_descriptor
+ for p_transform in list(process_bundle_descriptor.primitive_transform):
+ self.fns[p_transform.function_spec.id] = p_transform.function_spec
return beam_fn_api_pb2.RegisterResponse()
+ def initial_source_split(self, request, unused_instruction_id=None):
+ source_spec = self.fns[request.source_reference]
+ assert source_spec.urn == PYTHON_SOURCE_URN
+ source_bundle = unpack_and_deserialize_py_fn(
+ self.fns[request.source_reference])
+ splits = source_bundle.source.split(request.desired_bundle_size_bytes,
+ source_bundle.start_position,
+ source_bundle.stop_position)
+ response = beam_fn_api_pb2.InitialSourceSplitResponse()
+ response.splits.extend([
+ beam_fn_api_pb2.SourceSplit(
+ source=serialize_and_pack_py_fn(split, PYTHON_SOURCE_URN),
+ relative_size=split.weight,
+ )
+ for split in splits
+ ])
+ return response
+
+ def create_execution_tree(self, descriptor):
+ # TODO(vikasrk): Add an id field to Coder proto and use that instead.
+ coders = {coder.function_spec.id: operation_specs.get_coder_from_spec(
+ json.loads(unpack_function_spec_data(coder.function_spec)))
+ for coder in descriptor.coders}
+
+ counter_factory = counters.CounterFactory()
+ # TODO(robertwb): Figure out the correct prefix to use for output counters
+ # from StateSampler.
+ state_sampler = statesampler.StateSampler(
+ 'fnapi-step%s-' % descriptor.id, counter_factory)
+ consumers = collections.defaultdict(lambda: collections.defaultdict(list))
+ ops_by_id = {}
+ reversed_ops = []
+
+ for transform in reversed(descriptor.primitive_transform):
+ # TODO(robertwb): Figure out how to plumb through the operation name (e.g.
+ # "s3") from the service through the FnAPI so that msec counters can be
+ # reported and correctly plumbed through the service and the UI.
+ operation_name = 'fnapis%s' % transform.id
+
+ def only_element(iterable):
+ element, = iterable
+ return element
+
+ if transform.function_spec.urn == DATA_OUTPUT_URN:
+ target = beam_fn_api_pb2.Target(
+ primitive_transform_reference=transform.id,
+ name=only_element(transform.outputs.keys()))
+
+ op = DataOutputOperation(
+ operation_name,
+ transform.step_name,
+ consumers[transform.id],
+ counter_factory,
+ state_sampler,
+ coders[only_element(transform.outputs.values()).coder_reference],
+ target,
+ self.data_channel_factory.create_data_channel(
+ transform.function_spec))
+
+ elif transform.function_spec.urn == DATA_INPUT_URN:
+ target = beam_fn_api_pb2.Target(
+ primitive_transform_reference=transform.id,
+ name=only_element(transform.inputs.keys()))
+ op = DataInputOperation(
+ operation_name,
+ transform.step_name,
+ consumers[transform.id],
+ counter_factory,
+ state_sampler,
+ coders[only_element(transform.outputs.values()).coder_reference],
+ target,
+ self.data_channel_factory.create_data_channel(
+ transform.function_spec))
+
+ elif transform.function_spec.urn == PYTHON_DOFN_URN:
+ def create_side_input(tag, si):
+ # TODO(robertwb): Extract windows (and keys) out of element data.
+ return operation_specs.WorkerSideInputSource(
+ tag=tag,
+ source=SideInputSource(
+ self.state_handler,
+ beam_fn_api_pb2.StateKey(
+ function_spec_reference=si.view_fn.id),
+ coder=unpack_and_deserialize_py_fn(si.view_fn)))
+ output_tags = list(transform.outputs.keys())
+ spec = operation_specs.WorkerDoFn(
+ serialized_fn=unpack_function_spec_data(transform.function_spec),
+ output_tags=output_tags,
+ input=None,
+ side_inputs=[create_side_input(tag, si)
+ for tag, si in transform.side_inputs.items()],
+ output_coders=[coders[transform.outputs[out].coder_reference]
+ for out in output_tags])
+
+ op = operations.DoOperation(operation_name, spec, counter_factory,
+ state_sampler)
+ # TODO(robertwb): Move these to the constructor.
+ op.step_name = transform.step_name
+ for tag, op_consumers in consumers[transform.id].items():
+ for consumer in op_consumers:
+ op.add_receiver(
+ consumer, output_tags.index(tag))
+
+ elif transform.function_spec.urn == IDENTITY_DOFN_URN:
+ op = operations.FlattenOperation(operation_name, None, counter_factory,
+ state_sampler)
+ # TODO(robertwb): Move these to the constructor.
+ op.step_name = transform.step_name
+ for tag, op_consumers in consumers[transform.id].items():
+ for consumer in op_consumers:
+ op.add_receiver(consumer, 0)
+
+ elif transform.function_spec.urn == PYTHON_SOURCE_URN:
+ source = load_compressed(unpack_function_spec_data(
+ transform.function_spec))
+ # TODO(vikasrk): Remove this once custom source is implemented with
+ # splittable dofn via the data plane.
+ spec = operation_specs.WorkerRead(
+ iobase.SourceBundle(1.0, source, None, None),
+ [WindowedValueCoder(source.default_output_coder())])
+ op = operations.ReadOperation(operation_name, spec, counter_factory,
+ state_sampler)
+ op.step_name = transform.step_name
+ output_tags = list(transform.outputs.keys())
+ for tag, op_consumers in consumers[transform.id].items():
+ for consumer in op_consumers:
+ op.add_receiver(
+ consumer, output_tags.index(tag))
+
+ else:
+ raise NotImplementedError
+
+ # Record consumers.
+ for _, inputs in transform.inputs.items():
+ for target in inputs.target:
+ consumers[target.primitive_transform_reference][target.name].append(
+ op)
+
+ reversed_ops.append(op)
+ ops_by_id[transform.id] = op
+
+ return list(reversed(reversed_ops)), ops_by_id
+
def process_bundle(self, request, instruction_id):
- bundle_processor.BundleProcessor(
- self.fns[request.process_bundle_descriptor_reference],
- self.state_handler,
- self.data_channel_factory).process_bundle(instruction_id)
+ ops, ops_by_id = self.create_execution_tree(
+ self.fns[request.process_bundle_descriptor_reference])
+
+ expected_inputs = []
+ for _, op in ops_by_id.items():
+ if isinstance(op, DataOutputOperation):
+ # TODO(robertwb): Is there a better way to pass the instruction id to
+ # the operation?
+ op.set_output_stream(op.data_channel.output_stream(
+ instruction_id, op.target))
+ elif isinstance(op, DataInputOperation):
+ # We must wait until we receive "end of stream" for each of these ops.
+ expected_inputs.append(op)
+
+ # Start all operations.
+ for op in reversed(ops):
+ logging.info('start %s', op)
+ op.start()
+
+ # Inject inputs from data plane.
+ for input_op in expected_inputs:
+ for data in input_op.data_channel.input_elements(
+ instruction_id, [input_op.target]):
+ # ignores input name
+ target_op = ops_by_id[data.target.primitive_transform_reference]
+ # lacks coder for non-input ops
+ target_op.process_encoded(data.data)
+
+ # Finish all operations.
+ for op in ops:
+ logging.info('finish %s', op)
+ op.finish()
return beam_fn_api_pb2.ProcessBundleResponse()
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/worker/sdk_worker_main.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker_main.py b/sdks/python/apache_beam/runners/worker/sdk_worker_main.py
index f3f1e02..b891779 100644
--- a/sdks/python/apache_beam/runners/worker/sdk_worker_main.py
+++ b/sdks/python/apache_beam/runners/worker/sdk_worker_main.py
@@ -24,7 +24,7 @@ import sys
import grpc
from google.protobuf import text_format
-from apache_beam.portability.api import beam_fn_api_pb2
+from apache_beam.runners.api import beam_fn_api_pb2
from apache_beam.runners.worker.log_handler import FnApiLogRecordHandler
from apache_beam.runners.worker.sdk_worker import SdkHarness
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/worker/sdk_worker_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker_test.py b/sdks/python/apache_beam/runners/worker/sdk_worker_test.py
index dc72a5f..0d0811b 100644
--- a/sdks/python/apache_beam/runners/worker/sdk_worker_test.py
+++ b/sdks/python/apache_beam/runners/worker/sdk_worker_test.py
@@ -27,8 +27,10 @@ import unittest
from concurrent import futures
import grpc
-from apache_beam.portability.api import beam_fn_api_pb2
-from apache_beam.portability.api import beam_runner_api_pb2
+from apache_beam.io.concat_source_test import RangeSource
+from apache_beam.io.iobase import SourceBundle
+from apache_beam.runners.api import beam_fn_api_pb2
+from apache_beam.runners.worker import data_plane
from apache_beam.runners.worker import sdk_worker
@@ -62,12 +64,13 @@ class BeamFnControlServicer(beam_fn_api_pb2.BeamFnControlServicer):
class SdkWorkerTest(unittest.TestCase):
def test_fn_registration(self):
- process_bundle_descriptors = [
- beam_fn_api_pb2.ProcessBundleDescriptor(
- id=str(100+ix),
- transforms={
- str(ix): beam_runner_api_pb2.PTransform(unique_name=str(ix))})
- for ix in range(4)]
+ fns = [beam_fn_api_pb2.FunctionSpec(id=str(ix)) for ix in range(4)]
+
+ process_bundle_descriptors = [beam_fn_api_pb2.ProcessBundleDescriptor(
+ id=str(100+ix),
+ primitive_transform=[
+ beam_fn_api_pb2.PrimitiveTransform(function_spec=fn)])
+ for ix, fn in enumerate(fns)]
test_controller = BeamFnControlServicer([beam_fn_api_pb2.InstructionRequest(
register=beam_fn_api_pb2.RegisterRequest(
@@ -83,7 +86,81 @@ class SdkWorkerTest(unittest.TestCase):
harness.run()
self.assertEqual(
harness.worker.fns,
- {item.id: item for item in process_bundle_descriptors})
+ {item.id: item for item in fns + process_bundle_descriptors})
+
+ @unittest.skip("initial splitting not in proto")
+ def test_source_split(self):
+ source = RangeSource(0, 100)
+ expected_splits = list(source.split(30))
+
+ worker = sdk_harness.SdkWorker(
+ None, data_plane.GrpcClientDataChannelFactory())
+ worker.register(
+ beam_fn_api_pb2.RegisterRequest(
+ process_bundle_descriptor=[beam_fn_api_pb2.ProcessBundleDescriptor(
+ primitive_transform=[beam_fn_api_pb2.PrimitiveTransform(
+ function_spec=sdk_harness.serialize_and_pack_py_fn(
+ SourceBundle(1.0, source, None, None),
+ sdk_harness.PYTHON_SOURCE_URN,
+ id="src"))])]))
+ split_response = worker.initial_source_split(
+ beam_fn_api_pb2.InitialSourceSplitRequest(
+ desired_bundle_size_bytes=30,
+ source_reference="src"))
+
+ self.assertEqual(
+ expected_splits,
+ [sdk_harness.unpack_and_deserialize_py_fn(s.source)
+ for s in split_response.splits])
+
+ self.assertEqual(
+ [s.weight for s in expected_splits],
+ [s.relative_size for s in split_response.splits])
+
+ @unittest.skip("initial splitting not in proto")
+ def test_source_split_via_instruction(self):
+
+ source = RangeSource(0, 100)
+ expected_splits = list(source.split(30))
+
+ test_controller = BeamFnControlServicer([
+ beam_fn_api_pb2.InstructionRequest(
+ instruction_id="register_request",
+ register=beam_fn_api_pb2.RegisterRequest(
+ process_bundle_descriptor=[
+ beam_fn_api_pb2.ProcessBundleDescriptor(
+ primitive_transform=[beam_fn_api_pb2.PrimitiveTransform(
+ function_spec=sdk_harness.serialize_and_pack_py_fn(
+ SourceBundle(1.0, source, None, None),
+ sdk_harness.PYTHON_SOURCE_URN,
+ id="src"))])])),
+ beam_fn_api_pb2.InstructionRequest(
+ instruction_id="split_request",
+ initial_source_split=beam_fn_api_pb2.InitialSourceSplitRequest(
+ desired_bundle_size_bytes=30,
+ source_reference="src"))
+ ])
+
+ server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
+ beam_fn_api_pb2.add_BeamFnControlServicer_to_server(test_controller, server)
+ test_port = server.add_insecure_port("[::]:0")
+ server.start()
+
+ channel = grpc.insecure_channel("localhost:%s" % test_port)
+ harness = sdk_harness.SdkHarness(channel)
+ harness.run()
+
+ split_response = test_controller.responses[
+ "split_request"].initial_source_split
+
+ self.assertEqual(
+ expected_splits,
+ [sdk_harness.unpack_and_deserialize_py_fn(s.source)
+ for s in split_response.splits])
+
+ self.assertEqual(
+ [s.weight for s in expected_splits],
+ [s.relative_size for s in split_response.splits])
if __name__ == "__main__":
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/testing/test_stream.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/testing/test_stream.py b/sdks/python/apache_beam/testing/test_stream.py
index 7989fb2..a06bcd0 100644
--- a/sdks/python/apache_beam/testing/test_stream.py
+++ b/sdks/python/apache_beam/testing/test_stream.py
@@ -24,10 +24,8 @@ from abc import ABCMeta
from abc import abstractmethod
from apache_beam import coders
-from apache_beam import core
from apache_beam import pvalue
from apache_beam.transforms import PTransform
-from apache_beam.transforms import window
from apache_beam.transforms.window import TimestampedValue
from apache_beam.utils import timestamp
from apache_beam.utils.windowed_value import WindowedValue
@@ -101,9 +99,6 @@ class TestStream(PTransform):
self.current_watermark = timestamp.MIN_TIMESTAMP
self.events = []
- def get_windowing(self, unused_inputs):
- return core.Windowing(window.GlobalWindows())
-
def expand(self, pbegin):
assert isinstance(pbegin, pvalue.PBegin)
self.pipeline = pbegin.pipeline
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/testing/test_stream_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/testing/test_stream_test.py b/sdks/python/apache_beam/testing/test_stream_test.py
index b7ca141..e32dda2 100644
--- a/sdks/python/apache_beam/testing/test_stream_test.py
+++ b/sdks/python/apache_beam/testing/test_stream_test.py
@@ -19,16 +19,10 @@
import unittest
-import apache_beam as beam
-from apache_beam.options.pipeline_options import PipelineOptions
-from apache_beam.options.pipeline_options import StandardOptions
-from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.test_stream import ElementEvent
from apache_beam.testing.test_stream import ProcessingTimeEvent
from apache_beam.testing.test_stream import TestStream
from apache_beam.testing.test_stream import WatermarkEvent
-from apache_beam.testing.util import assert_that, equal_to
-from apache_beam.transforms.window import FixedWindows
from apache_beam.transforms.window import TimestampedValue
from apache_beam.utils import timestamp
from apache_beam.utils.windowed_value import WindowedValue
@@ -84,68 +78,6 @@ class TestStreamTest(unittest.TestCase):
TimestampedValue('a', timestamp.MAX_TIMESTAMP)
]))
- def test_basic_execution(self):
- test_stream = (TestStream()
- .advance_watermark_to(10)
- .add_elements(['a', 'b', 'c'])
- .advance_watermark_to(20)
- .add_elements(['d'])
- .add_elements(['e'])
- .advance_processing_time(10)
- .advance_watermark_to(300)
- .add_elements([TimestampedValue('late', 12)])
- .add_elements([TimestampedValue('last', 310)]))
-
- class RecordFn(beam.DoFn):
- def process(self, element=beam.DoFn.ElementParam,
- timestamp=beam.DoFn.TimestampParam):
- yield (element, timestamp)
-
- options = PipelineOptions()
- options.view_as(StandardOptions).streaming = True
- p = TestPipeline(options=options)
- my_record_fn = RecordFn()
- records = p | test_stream | beam.ParDo(my_record_fn)
- assert_that(records, equal_to([
- ('a', timestamp.Timestamp(10)),
- ('b', timestamp.Timestamp(10)),
- ('c', timestamp.Timestamp(10)),
- ('d', timestamp.Timestamp(20)),
- ('e', timestamp.Timestamp(20)),
- ('late', timestamp.Timestamp(12)),
- ('last', timestamp.Timestamp(310)),]))
- p.run()
-
- def test_gbk_execution(self):
- test_stream = (TestStream()
- .advance_watermark_to(10)
- .add_elements(['a', 'b', 'c'])
- .advance_watermark_to(20)
- .add_elements(['d'])
- .add_elements(['e'])
- .advance_processing_time(10)
- .advance_watermark_to(300)
- .add_elements([TimestampedValue('late', 12)])
- .add_elements([TimestampedValue('last', 310)]))
-
- options = PipelineOptions()
- options.view_as(StandardOptions).streaming = True
- p = TestPipeline(options=options)
- records = (p
- | test_stream
- | beam.WindowInto(FixedWindows(15))
- | beam.Map(lambda x: ('k', x))
- | beam.GroupByKey())
- # TODO(BEAM-2519): timestamp assignment for elements from a GBK should
- # respect the TimestampCombiner. The test below should also verify the
- # timestamps of the outputted elements once this is implemented.
- assert_that(records, equal_to([
- ('k', ['a', 'b', 'c']),
- ('k', ['d', 'e']),
- ('k', ['late']),
- ('k', ['last'])]))
- p.run()
-
if __name__ == '__main__':
unittest.main()
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/transforms/combiners.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/transforms/combiners.py b/sdks/python/apache_beam/transforms/combiners.py
index 875306f..fa0742d 100644
--- a/sdks/python/apache_beam/transforms/combiners.py
+++ b/sdks/python/apache_beam/transforms/combiners.py
@@ -149,7 +149,6 @@ class Top(object):
"""Combiners for obtaining extremal elements."""
# pylint: disable=no-self-argument
- @staticmethod
@ptransform.ptransform_fn
def Of(pcoll, n, compare=None, *args, **kwargs):
"""Obtain a list of the compare-most N elements in a PCollection.
@@ -178,7 +177,6 @@ class Top(object):
return pcoll | core.CombineGlobally(
TopCombineFn(n, compare, key, reverse), *args, **kwargs)
- @staticmethod
@ptransform.ptransform_fn
def PerKey(pcoll, n, compare=None, *args, **kwargs):
"""Identifies the compare-most N elements associated with each key.
@@ -212,25 +210,21 @@ class Top(object):
return pcoll | core.CombinePerKey(
TopCombineFn(n, compare, key, reverse), *args, **kwargs)
- @staticmethod
@ptransform.ptransform_fn
def Largest(pcoll, n):
"""Obtain a list of the greatest N elements in a PCollection."""
return pcoll | Top.Of(n)
- @staticmethod
@ptransform.ptransform_fn
def Smallest(pcoll, n):
"""Obtain a list of the least N elements in a PCollection."""
return pcoll | Top.Of(n, reverse=True)
- @staticmethod
@ptransform.ptransform_fn
def LargestPerKey(pcoll, n):
"""Identifies the N greatest elements associated with each key."""
return pcoll | Top.PerKey(n)
- @staticmethod
@ptransform.ptransform_fn
def SmallestPerKey(pcoll, n, reverse=True):
"""Identifies the N least elements associated with each key."""
@@ -375,12 +369,10 @@ class Sample(object):
"""Combiners for sampling n elements without replacement."""
# pylint: disable=no-self-argument
- @staticmethod
@ptransform.ptransform_fn
def FixedSizeGlobally(pcoll, n):
return pcoll | core.CombineGlobally(SampleCombineFn(n))
- @staticmethod
@ptransform.ptransform_fn
def FixedSizePerKey(pcoll, n):
return pcoll | core.CombinePerKey(SampleCombineFn(n))
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/transforms/combiners_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/transforms/combiners_test.py b/sdks/python/apache_beam/transforms/combiners_test.py
index cd2b595..c79fec8 100644
--- a/sdks/python/apache_beam/transforms/combiners_test.py
+++ b/sdks/python/apache_beam/transforms/combiners_test.py
@@ -156,11 +156,14 @@ class CombineTest(unittest.TestCase):
def test_combine_sample_display_data(self):
def individual_test_per_key_dd(sampleFn, args, kwargs):
- trs = [sampleFn(*args, **kwargs)]
+ trs = [beam.CombinePerKey(sampleFn(*args, **kwargs)),
+ beam.CombineGlobally(sampleFn(*args, **kwargs))]
for transform in trs:
dd = DisplayData.create_from(transform)
expected_items = [
- DisplayDataItemMatcher('fn', transform._fn.__name__)]
+ DisplayDataItemMatcher('fn', sampleFn.fn.__name__),
+ DisplayDataItemMatcher('combine_fn',
+ transform.fn.__class__)]
if args:
expected_items.append(
DisplayDataItemMatcher('args', str(args)))
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/transforms/core.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/transforms/core.py b/sdks/python/apache_beam/transforms/core.py
index 8018219..0e497f9 100644
--- a/sdks/python/apache_beam/transforms/core.py
+++ b/sdks/python/apache_beam/transforms/core.py
@@ -27,7 +27,7 @@ from apache_beam import pvalue
from apache_beam import typehints
from apache_beam.coders import typecoders
from apache_beam.internal import util
-from apache_beam.portability.api import beam_runner_api_pb2
+from apache_beam.runners.api import beam_runner_api_pb2
from apache_beam.transforms import ptransform
from apache_beam.transforms.display import DisplayDataItem
from apache_beam.transforms.display import HasDisplayData
@@ -1078,6 +1078,40 @@ class GroupByKey(PTransform):
key_type, value_type = trivial_inference.key_value_types(input_type)
return Iterable[KV[key_type, typehints.WindowedValue[value_type]]]
+ class GroupAlsoByWindow(DoFn):
+ # TODO(robertwb): Support combiner lifting.
+
+ def __init__(self, windowing):
+ super(GroupByKey.GroupAlsoByWindow, self).__init__()
+ self.windowing = windowing
+
+ def infer_output_type(self, input_type):
+ key_type, windowed_value_iter_type = trivial_inference.key_value_types(
+ input_type)
+ value_type = windowed_value_iter_type.inner_type.inner_type
+ return Iterable[KV[key_type, Iterable[value_type]]]
+
+ def start_bundle(self):
+ # pylint: disable=wrong-import-order, wrong-import-position
+ from apache_beam.transforms.trigger import InMemoryUnmergedState
+ from apache_beam.transforms.trigger import create_trigger_driver
+ # pylint: enable=wrong-import-order, wrong-import-position
+ self.driver = create_trigger_driver(self.windowing, True)
+ self.state_type = InMemoryUnmergedState
+
+ def process(self, element):
+ k, vs = element
+ state = self.state_type()
+ # TODO(robertwb): Conditionally process in smaller chunks.
+ for wvalue in self.driver.process_elements(state, vs, MIN_TIMESTAMP):
+ yield wvalue.with_value((k, wvalue.value))
+ while state.timers:
+ fired = state.get_and_clear_timers()
+ for timer_window, (name, time_domain, fire_time) in fired:
+ for wvalue in self.driver.process_timer(
+ timer_window, name, time_domain, fire_time, state):
+ yield wvalue.with_value((k, wvalue.value))
+
def expand(self, pcoll):
# This code path is only used in the local direct runner. For Dataflow
# runner execution, the GroupByKey transform is expanded on the service.
@@ -1102,7 +1136,8 @@ class GroupByKey(PTransform):
| 'GroupByKey' >> (_GroupByKeyOnly()
.with_input_types(reify_output_type)
.with_output_types(gbk_input_type))
- | ('GroupByWindow' >> _GroupAlsoByWindow(pcoll.windowing)
+ | ('GroupByWindow' >> ParDo(
+ self.GroupAlsoByWindow(pcoll.windowing))
.with_input_types(gbk_input_type)
.with_output_types(gbk_output_type)))
else:
@@ -1110,7 +1145,8 @@ class GroupByKey(PTransform):
return (pcoll
| 'ReifyWindows' >> ParDo(self.ReifyWindows())
| 'GroupByKey' >> _GroupByKeyOnly()
- | 'GroupByWindow' >> _GroupAlsoByWindow(pcoll.windowing))
+ | 'GroupByWindow' >> ParDo(
+ self.GroupAlsoByWindow(pcoll.windowing)))
@typehints.with_input_types(typehints.KV[K, V])
@@ -1126,55 +1162,6 @@ class _GroupByKeyOnly(PTransform):
return pvalue.PCollection(pcoll.pipeline)
-@typehints.with_input_types(typehints.KV[K, typehints.Iterable[V]])
-@typehints.with_output_types(typehints.KV[K, typehints.Iterable[V]])
-class _GroupAlsoByWindow(ParDo):
- """The GroupAlsoByWindow transform."""
- def __init__(self, windowing):
- super(_GroupAlsoByWindow, self).__init__(
- _GroupAlsoByWindowDoFn(windowing))
- self.windowing = windowing
-
- def expand(self, pcoll):
- self._check_pcollection(pcoll)
- return pvalue.PCollection(pcoll.pipeline)
-
-
-class _GroupAlsoByWindowDoFn(DoFn):
- # TODO(robertwb): Support combiner lifting.
-
- def __init__(self, windowing):
- super(_GroupAlsoByWindowDoFn, self).__init__()
- self.windowing = windowing
-
- def infer_output_type(self, input_type):
- key_type, windowed_value_iter_type = trivial_inference.key_value_types(
- input_type)
- value_type = windowed_value_iter_type.inner_type.inner_type
- return Iterable[KV[key_type, Iterable[value_type]]]
-
- def start_bundle(self):
- # pylint: disable=wrong-import-order, wrong-import-position
- from apache_beam.transforms.trigger import InMemoryUnmergedState
- from apache_beam.transforms.trigger import create_trigger_driver
- # pylint: enable=wrong-import-order, wrong-import-position
- self.driver = create_trigger_driver(self.windowing, True)
- self.state_type = InMemoryUnmergedState
-
- def process(self, element):
- k, vs = element
- state = self.state_type()
- # TODO(robertwb): Conditionally process in smaller chunks.
- for wvalue in self.driver.process_elements(state, vs, MIN_TIMESTAMP):
- yield wvalue.with_value((k, wvalue.value))
- while state.timers:
- fired = state.get_and_clear_timers()
- for timer_window, (name, time_domain, fire_time) in fired:
- for wvalue in self.driver.process_timer(
- timer_window, name, time_domain, fire_time, state):
- yield wvalue.with_value((k, wvalue.value))
-
-
class Partition(PTransformWithSideInputs):
"""Split a PCollection into several partitions.
@@ -1444,18 +1431,15 @@ class Create(PTransform):
return Any
return Union[[trivial_inference.instance_to_type(v) for v in self.value]]
- def get_output_type(self):
- return (self.get_type_hints().simple_output_type(self.label) or
- self.infer_output_type(None))
-
def expand(self, pbegin):
from apache_beam.io import iobase
assert isinstance(pbegin, pvalue.PBegin)
self.pipeline = pbegin.pipeline
- coder = typecoders.registry.get_coder(self.get_output_type())
+ ouput_type = (self.get_type_hints().simple_output_type(self.label) or
+ self.infer_output_type(None))
+ coder = typecoders.registry.get_coder(ouput_type)
source = self._create_source_from_iterable(self.value, coder)
- return (pbegin.pipeline
- | iobase.Read(source).with_output_types(self.get_output_type()))
+ return pbegin.pipeline | iobase.Read(source).with_output_types(ouput_type)
def get_windowing(self, unused_inputs):
return Windowing(GlobalWindows())
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/transforms/ptransform.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/transforms/ptransform.py b/sdks/python/apache_beam/transforms/ptransform.py
index cd84122..bd2a120 100644
--- a/sdks/python/apache_beam/transforms/ptransform.py
+++ b/sdks/python/apache_beam/transforms/ptransform.py
@@ -430,7 +430,7 @@ class PTransform(WithTypeHints, HasDisplayData):
cls._known_urns[urn] = parameter_type, constructor
def to_runner_api(self, context):
- from apache_beam.portability.api import beam_runner_api_pb2
+ from apache_beam.runners.api import beam_runner_api_pb2
urn, typed_param = self.to_runner_api_parameter(context)
return beam_runner_api_pb2.FunctionSpec(
urn=urn,
@@ -595,23 +595,32 @@ class PTransformWithSideInputs(PTransform):
return '%s(%s)' % (self.__class__.__name__, self.fn.default_label())
-class _PTransformFnPTransform(PTransform):
+class CallablePTransform(PTransform):
"""A class wrapper for a function-based transform."""
- def __init__(self, fn, *args, **kwargs):
- super(_PTransformFnPTransform, self).__init__()
- self._fn = fn
- self._args = args
- self._kwargs = kwargs
+ def __init__(self, fn):
+ # pylint: disable=super-init-not-called
+ # This is a helper class for a function decorator. Only when the class
+ # is called (and __call__ invoked) we will have all the information
+ # needed to initialize the super class.
+ self.fn = fn
+ self._args = ()
+ self._kwargs = {}
def display_data(self):
- res = {'fn': (self._fn.__name__
- if hasattr(self._fn, '__name__')
- else self._fn.__class__),
+ res = {'fn': (self.fn.__name__
+ if hasattr(self.fn, '__name__')
+ else self.fn.__class__),
'args': DisplayDataItem(str(self._args)).drop_if_default('()'),
'kwargs': DisplayDataItem(str(self._kwargs)).drop_if_default('{}')}
return res
+ def __call__(self, *args, **kwargs):
+ super(CallablePTransform, self).__init__()
+ self._args = args
+ self._kwargs = kwargs
+ return self
+
def expand(self, pcoll):
# Since the PTransform will be implemented entirely as a function
# (once called), we need to pass through any type-hinting information that
@@ -620,18 +629,18 @@ class _PTransformFnPTransform(PTransform):
kwargs = dict(self._kwargs)
args = tuple(self._args)
try:
- if 'type_hints' in inspect.getargspec(self._fn).args:
+ if 'type_hints' in inspect.getargspec(self.fn).args:
args = (self.get_type_hints(),) + args
except TypeError:
# Might not be a function.
pass
- return self._fn(pcoll, *args, **kwargs)
+ return self.fn(pcoll, *args, **kwargs)
def default_label(self):
if self._args:
return '%s(%s)' % (
- label_from_callable(self._fn), label_from_callable(self._args[0]))
- return label_from_callable(self._fn)
+ label_from_callable(self.fn), label_from_callable(self._args[0]))
+ return label_from_callable(self.fn)
def ptransform_fn(fn):
@@ -675,11 +684,7 @@ def ptransform_fn(fn):
operator (i.e., `|`) will inject the pcoll argument in its proper place
(first argument if no label was specified and second argument otherwise).
"""
- # TODO(robertwb): Consider removing staticmethod to allow for self parameter.
-
- def callable_ptransform_factory(*args, **kwargs):
- return _PTransformFnPTransform(fn, *args, **kwargs)
- return callable_ptransform_factory
+ return CallablePTransform(fn)
def label_from_callable(fn):
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/transforms/trigger.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/transforms/trigger.py b/sdks/python/apache_beam/transforms/trigger.py
index f77fa1a..4200995 100644
--- a/sdks/python/apache_beam/transforms/trigger.py
+++ b/sdks/python/apache_beam/transforms/trigger.py
@@ -33,10 +33,9 @@ from apache_beam.transforms.window import GlobalWindow
from apache_beam.transforms.window import TimestampCombiner
from apache_beam.transforms.window import WindowedValue
from apache_beam.transforms.window import WindowFn
-from apache_beam.portability.api import beam_runner_api_pb2
+from apache_beam.runners.api import beam_runner_api_pb2
from apache_beam.utils.timestamp import MAX_TIMESTAMP
from apache_beam.utils.timestamp import MIN_TIMESTAMP
-from apache_beam.utils.timestamp import TIME_GRANULARITY
# AfterCount is experimental. No backwards compatibility guarantees.
@@ -1067,8 +1066,6 @@ class InMemoryUnmergedState(UnmergedState):
def clear_timer(self, window, name, time_domain):
self.timers[window].pop((name, time_domain), None)
- if not self.timers[window]:
- del self.timers[window]
def get_window(self, window_id):
return window_id
@@ -1105,34 +1102,17 @@ class InMemoryUnmergedState(UnmergedState):
if not self.state[window]:
self.state.pop(window, None)
- def get_timers(self, clear=False, watermark=MAX_TIMESTAMP):
+ def get_and_clear_timers(self, watermark=MAX_TIMESTAMP):
expired = []
for window, timers in list(self.timers.items()):
for (name, time_domain), timestamp in list(timers.items()):
if timestamp <= watermark:
expired.append((window, (name, time_domain, timestamp)))
- if clear:
- del timers[(name, time_domain)]
- if not timers and clear:
+ del timers[(name, time_domain)]
+ if not timers:
del self.timers[window]
return expired
- def get_and_clear_timers(self, watermark=MAX_TIMESTAMP):
- return self.get_timers(clear=True, watermark=watermark)
-
- def get_earliest_hold(self):
- earliest_hold = MAX_TIMESTAMP
- for unused_window, tagged_states in self.state.iteritems():
- # TODO(BEAM-2519): currently, this assumes that the watermark hold tag is
- # named "watermark". This is currently only true because the only place
- # watermark holds are set is in the GeneralTriggerDriver, where we use
- # this name. We should fix this by allowing enumeration of the tag types
- # used in adding state.
- if 'watermark' in tagged_states and tagged_states['watermark']:
- hold = min(tagged_states['watermark']) - TIME_GRANULARITY
- earliest_hold = min(earliest_hold, hold)
- return earliest_hold
-
def __repr__(self):
state_str = '\n'.join('%s: %s' % (key, dict(state))
for key, state in self.state.items())
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/transforms/window.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/transforms/window.py b/sdks/python/apache_beam/transforms/window.py
index 458fb74..e87a007 100644
--- a/sdks/python/apache_beam/transforms/window.py
+++ b/sdks/python/apache_beam/transforms/window.py
@@ -55,8 +55,8 @@ from google.protobuf import duration_pb2
from google.protobuf import timestamp_pb2
from apache_beam.coders import coders
-from apache_beam.portability.api import beam_runner_api_pb2
-from apache_beam.portability.api import standard_window_fns_pb2
+from apache_beam.runners.api import beam_runner_api_pb2
+from apache_beam.runners.api import standard_window_fns_pb2
from apache_beam.transforms import timeutil
from apache_beam.utils import proto_utils
from apache_beam.utils import urns
[14/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOTest.java
index a6be4fb..9468893 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOTest.java
@@ -42,9 +42,7 @@ import static org.junit.Assert.assertThat;
import static org.junit.Assert.assertTrue;
import com.google.common.base.Function;
-import com.google.common.base.Functions;
import com.google.common.base.Predicate;
-import com.google.common.base.Predicates;
import com.google.common.collect.FluentIterable;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
@@ -71,31 +69,22 @@ import java.util.zip.GZIPOutputStream;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;
import javax.annotation.Nullable;
-import org.apache.beam.sdk.coders.AvroCoder;
import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.DefaultCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.io.BoundedSource.BoundedReader;
-import org.apache.beam.sdk.io.DefaultFilenamePolicy.Params;
-import org.apache.beam.sdk.io.FileBasedSink.DynamicDestinations;
-import org.apache.beam.sdk.io.FileBasedSink.FilenamePolicy;
import org.apache.beam.sdk.io.FileBasedSink.WritableByteChannelFactory;
import org.apache.beam.sdk.io.TextIO.CompressionType;
import org.apache.beam.sdk.io.fs.MatchResult;
import org.apache.beam.sdk.io.fs.MatchResult.Metadata;
-import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions;
-import org.apache.beam.sdk.io.fs.ResourceId;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.ValueProvider;
-import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
import org.apache.beam.sdk.testing.NeedsRunner;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.SourceTestUtils;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.testing.ValidatesRunner;
import org.apache.beam.sdk.transforms.Create;
-import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.display.DisplayDataEvaluator;
import org.apache.beam.sdk.util.CoderUtils;
@@ -120,10 +109,10 @@ import org.junit.runners.JUnit4;
public class TextIOTest {
private static final String MY_HEADER = "myHeader";
private static final String MY_FOOTER = "myFooter";
- private static final List<String> EMPTY = Collections.emptyList();
- private static final List<String> TINY =
- Arrays.asList("Irritable eagle", "Optimistic jay", "Fanciful hawk");
- private static final List<String> LARGE = makeLines(1000);
+ private static final String[] EMPTY = new String[] {};
+ private static final String[] TINY =
+ new String[] {"Irritable eagle", "Optimistic jay", "Fanciful hawk"};
+ private static final String[] LARGE = makeLines(1000);
private static Path tempFolder;
private static File emptyTxt;
@@ -148,7 +137,7 @@ public class TextIOTest {
@Rule
public ExpectedException expectedException = ExpectedException.none();
- private static File writeToFile(List<String> lines, String filename, CompressionType compression)
+ private static File writeToFile(String[] lines, String filename, CompressionType compression)
throws IOException {
File file = tempFolder.resolve(filename).toFile();
OutputStream output = new FileOutputStream(file);
@@ -216,7 +205,7 @@ public class TextIOTest {
});
}
- private void runTestRead(String[] expected) throws Exception {
+ private <T> void runTestRead(String[] expected) throws Exception {
File tmpFile = Files.createTempFile(tempFolder, "file", "txt").toFile();
String filename = tmpFile.getPath();
@@ -285,213 +274,6 @@ public class TextIOTest {
displayData, hasItem(hasDisplayItem(hasValue(startsWith("foobar")))));
}
- static class TestDynamicDestinations extends DynamicDestinations<String, String> {
- ResourceId baseDir;
-
- TestDynamicDestinations(ResourceId baseDir) {
- this.baseDir = baseDir;
- }
-
- @Override
- public String getDestination(String element) {
- // Destination is based on first character of string.
- return element.substring(0, 1);
- }
-
- @Override
- public String getDefaultDestination() {
- return "";
- }
-
- @Nullable
- @Override
- public Coder<String> getDestinationCoder() {
- return StringUtf8Coder.of();
- }
-
- @Override
- public FilenamePolicy getFilenamePolicy(String destination) {
- return DefaultFilenamePolicy.fromStandardParameters(
- StaticValueProvider.of(
- baseDir.resolve("file_" + destination + ".txt", StandardResolveOptions.RESOLVE_FILE)),
- null,
- null,
- false);
- }
- }
-
- class StartsWith implements Predicate<String> {
- String prefix;
-
- StartsWith(String prefix) {
- this.prefix = prefix;
- }
-
- @Override
- public boolean apply(@Nullable String input) {
- return input.startsWith(prefix);
- }
- }
-
- @Test
- @Category(NeedsRunner.class)
- public void testDynamicDestinations() throws Exception {
- ResourceId baseDir =
- FileSystems.matchNewResource(
- Files.createTempDirectory(tempFolder, "testDynamicDestinations").toString(), true);
-
- List<String> elements = Lists.newArrayList("aaaa", "aaab", "baaa", "baab", "caaa", "caab");
- PCollection<String> input = p.apply(Create.of(elements).withCoder(StringUtf8Coder.of()));
- input.apply(
- TextIO.write()
- .to(new TestDynamicDestinations(baseDir))
- .withTempDirectory(FileSystems.matchNewResource(baseDir.toString(), true)));
- p.run();
-
- assertOutputFiles(
- Iterables.toArray(Iterables.filter(elements, new StartsWith("a")), String.class),
- null,
- null,
- 0,
- baseDir.resolve("file_a.txt", StandardResolveOptions.RESOLVE_FILE),
- DefaultFilenamePolicy.DEFAULT_UNWINDOWED_SHARD_TEMPLATE);
- assertOutputFiles(
- Iterables.toArray(Iterables.filter(elements, new StartsWith("b")), String.class),
- null,
- null,
- 0,
- baseDir.resolve("file_b.txt", StandardResolveOptions.RESOLVE_FILE),
- DefaultFilenamePolicy.DEFAULT_UNWINDOWED_SHARD_TEMPLATE);
- assertOutputFiles(
- Iterables.toArray(Iterables.filter(elements, new StartsWith("c")), String.class),
- null,
- null,
- 0,
- baseDir.resolve("file_c.txt", StandardResolveOptions.RESOLVE_FILE),
- DefaultFilenamePolicy.DEFAULT_UNWINDOWED_SHARD_TEMPLATE);
- }
-
- @DefaultCoder(AvroCoder.class)
- private static class UserWriteType {
- String destination;
- String metadata;
-
- UserWriteType() {
- this.destination = "";
- this.metadata = "";
- }
-
- UserWriteType(String destination, String metadata) {
- this.destination = destination;
- this.metadata = metadata;
- }
-
- @Override
- public String toString() {
- return String.format("destination: %s metadata : %s", destination, metadata);
- }
- }
-
- private static class SerializeUserWrite implements SerializableFunction<UserWriteType, String> {
- @Override
- public String apply(UserWriteType input) {
- return input.toString();
- }
- }
-
- private static class UserWriteDestination implements SerializableFunction<UserWriteType, Params> {
- private ResourceId baseDir;
-
- UserWriteDestination(ResourceId baseDir) {
- this.baseDir = baseDir;
- }
-
- @Override
- public Params apply(UserWriteType input) {
- return new Params()
- .withBaseFilename(
- baseDir.resolve(
- "file_" + input.destination.substring(0, 1) + ".txt",
- StandardResolveOptions.RESOLVE_FILE));
- }
- }
-
- private static class ExtractWriteDestination implements Function<UserWriteType, String> {
- @Override
- public String apply(@Nullable UserWriteType input) {
- return input.destination;
- }
- }
-
- @Test
- @Category(NeedsRunner.class)
- public void testDynamicDefaultFilenamePolicy() throws Exception {
- ResourceId baseDir =
- FileSystems.matchNewResource(
- Files.createTempDirectory(tempFolder, "testDynamicDestinations").toString(), true);
-
- List<UserWriteType> elements =
- Lists.newArrayList(
- new UserWriteType("aaaa", "first"),
- new UserWriteType("aaab", "second"),
- new UserWriteType("baaa", "third"),
- new UserWriteType("baab", "fourth"),
- new UserWriteType("caaa", "fifth"),
- new UserWriteType("caab", "sixth"));
- PCollection<UserWriteType> input = p.apply(Create.of(elements));
- input.apply(
- TextIO.writeCustomType(new SerializeUserWrite())
- .to(new UserWriteDestination(baseDir), new Params())
- .withTempDirectory(FileSystems.matchNewResource(baseDir.toString(), true)));
- p.run();
-
- String[] aElements =
- Iterables.toArray(
- Iterables.transform(
- Iterables.filter(
- elements,
- Predicates.compose(new StartsWith("a"), new ExtractWriteDestination())),
- Functions.toStringFunction()),
- String.class);
- String[] bElements =
- Iterables.toArray(
- Iterables.transform(
- Iterables.filter(
- elements,
- Predicates.compose(new StartsWith("b"), new ExtractWriteDestination())),
- Functions.toStringFunction()),
- String.class);
- String[] cElements =
- Iterables.toArray(
- Iterables.transform(
- Iterables.filter(
- elements,
- Predicates.compose(new StartsWith("c"), new ExtractWriteDestination())),
- Functions.toStringFunction()),
- String.class);
- assertOutputFiles(
- aElements,
- null,
- null,
- 0,
- baseDir.resolve("file_a.txt", StandardResolveOptions.RESOLVE_FILE),
- DefaultFilenamePolicy.DEFAULT_UNWINDOWED_SHARD_TEMPLATE);
- assertOutputFiles(
- bElements,
- null,
- null,
- 0,
- baseDir.resolve("file_b.txt", StandardResolveOptions.RESOLVE_FILE),
- DefaultFilenamePolicy.DEFAULT_UNWINDOWED_SHARD_TEMPLATE);
- assertOutputFiles(
- cElements,
- null,
- null,
- 0,
- baseDir.resolve("file_c.txt", StandardResolveOptions.RESOLVE_FILE),
- DefaultFilenamePolicy.DEFAULT_UNWINDOWED_SHARD_TEMPLATE);
- }
-
private void runTestWrite(String[] elems) throws Exception {
runTestWrite(elems, null, null, 1);
}
@@ -509,8 +291,7 @@ public class TextIOTest {
String[] elems, String header, String footer, int numShards) throws Exception {
String outputName = "file.txt";
Path baseDir = Files.createTempDirectory(tempFolder, "testwrite");
- ResourceId baseFilename =
- FileBasedSink.convertToFileResourceIfPossible(baseDir.resolve(outputName).toString());
+ String baseFilename = baseDir.resolve(outputName).toString();
PCollection<String> input =
p.apply(Create.of(Arrays.asList(elems)).withCoder(StringUtf8Coder.of()));
@@ -530,14 +311,8 @@ public class TextIOTest {
p.run();
- assertOutputFiles(
- elems,
- header,
- footer,
- numShards,
- baseFilename,
- firstNonNull(
- write.inner.getShardTemplate(),
+ assertOutputFiles(elems, header, footer, numShards, baseDir, outputName,
+ firstNonNull(write.getShardTemplate(),
DefaultFilenamePolicy.DEFAULT_UNWINDOWED_SHARD_TEMPLATE));
}
@@ -546,12 +321,13 @@ public class TextIOTest {
final String header,
final String footer,
int numShards,
- ResourceId outputPrefix,
+ Path rootLocation,
+ String outputName,
String shardNameTemplate)
throws Exception {
List<File> expectedFiles = new ArrayList<>();
if (numShards == 0) {
- String pattern = outputPrefix.toString() + "*";
+ String pattern = rootLocation.toAbsolutePath().resolve(outputName + "*").toString();
List<MatchResult> matches = FileSystems.match(Collections.singletonList(pattern));
for (Metadata expectedFile : Iterables.getOnlyElement(matches).metadata()) {
expectedFiles.add(new File(expectedFile.resourceId().toString()));
@@ -560,9 +336,9 @@ public class TextIOTest {
for (int i = 0; i < numShards; i++) {
expectedFiles.add(
new File(
+ rootLocation.toString(),
DefaultFilenamePolicy.constructName(
- outputPrefix, shardNameTemplate, "", i, numShards, null, null)
- .toString()));
+ outputName, shardNameTemplate, "", i, numShards, null, null)));
}
}
@@ -680,19 +456,14 @@ public class TextIOTest {
public void testWriteWithWritableByteChannelFactory() throws Exception {
Coder<String> coder = StringUtf8Coder.of();
String outputName = "file.txt";
- ResourceId baseDir =
- FileSystems.matchNewResource(
- Files.createTempDirectory(tempFolder, "testwrite").toString(), true);
+ Path baseDir = Files.createTempDirectory(tempFolder, "testwrite");
PCollection<String> input = p.apply(Create.of(Arrays.asList(LINES2_ARRAY)).withCoder(coder));
final WritableByteChannelFactory writableByteChannelFactory =
new DrunkWritableByteChannelFactory();
- TextIO.Write write =
- TextIO.write()
- .to(baseDir.resolve(outputName, StandardResolveOptions.RESOLVE_FILE).toString())
- .withoutSharding()
- .withWritableByteChannelFactory(writableByteChannelFactory);
+ TextIO.Write write = TextIO.write().to(baseDir.resolve(outputName).toString())
+ .withoutSharding().withWritableByteChannelFactory(writableByteChannelFactory);
DisplayData displayData = DisplayData.from(write);
assertThat(displayData, hasDisplayItem("writableByteChannelFactory", "DRUNK"));
@@ -705,15 +476,8 @@ public class TextIOTest {
drunkElems.add(elem);
drunkElems.add(elem);
}
- assertOutputFiles(
- drunkElems.toArray(new String[0]),
- null,
- null,
- 1,
- baseDir.resolve(
- outputName + writableByteChannelFactory.getSuggestedFilenameSuffix(),
- StandardResolveOptions.RESOLVE_FILE),
- write.inner.getShardTemplate());
+ assertOutputFiles(drunkElems.toArray(new String[0]), null, null, 1, baseDir,
+ outputName + writableByteChannelFactory.getFilenameSuffix(), write.getShardTemplate());
}
@Test
@@ -791,7 +555,7 @@ public class TextIOTest {
* Helper that writes the given lines (adding a newline in between) to a stream, then closes the
* stream.
*/
- private static void writeToStreamAndClose(List<String> lines, OutputStream outputStream) {
+ private static void writeToStreamAndClose(String[] lines, OutputStream outputStream) {
try (PrintStream writer = new PrintStream(outputStream)) {
for (String line : lines) {
writer.println(line);
@@ -800,33 +564,27 @@ public class TextIOTest {
}
/**
- * Helper method that runs TextIO.read().from(filename).withCompressionType(compressionType) and
- * TextIO.readAll().withCompressionType(compressionType) applied to the single filename,
+ * Helper method that runs TextIO.read().from(filename).withCompressionType(compressionType)
* and asserts that the results match the given expected output.
*/
private void assertReadingCompressedFileMatchesExpected(
- File file, CompressionType compressionType, List<String> expected) {
-
- TextIO.Read read = TextIO.read().from(file.getPath()).withCompressionType(compressionType);
- PAssert.that(p.apply("Read_" + file + "_" + compressionType.toString(), read))
- .containsInAnyOrder(expected);
-
- TextIO.ReadAll readAll =
- TextIO.readAll().withCompressionType(compressionType).withDesiredBundleSizeBytes(10);
- PAssert.that(
- p.apply("Create_" + file, Create.of(file.getPath()))
- .apply("Read_" + compressionType.toString(), readAll))
- .containsInAnyOrder(expected);
+ File file, CompressionType compressionType, String[] expected) {
+
+ TextIO.Read read =
+ TextIO.read().from(file.getPath()).withCompressionType(compressionType);
+ PCollection<String> output = p.apply("Read_" + file + "_" + compressionType.toString(), read);
+
+ PAssert.that(output).containsInAnyOrder(expected);
p.run();
}
/**
* Helper to make an array of compressible strings. Returns ["word"i] for i in range(0,n).
*/
- private static List<String> makeLines(int n) {
- List<String> ret = new ArrayList<>();
+ private static String[] makeLines(int n) {
+ String[] ret = new String[n];
for (int i = 0; i < n; ++i) {
- ret.add("word" + i);
+ ret[i] = "word" + i;
}
return ret;
}
@@ -1010,7 +768,7 @@ public class TextIOTest {
String filename = createZipFile(expected, "multiple entries", entry0, entry1, entry2);
assertReadingCompressedFileMatchesExpected(
- new File(filename), CompressionType.ZIP, expected);
+ new File(filename), CompressionType.ZIP, expected.toArray(new String[]{}));
}
/**
@@ -1029,7 +787,7 @@ public class TextIOTest {
new String[]{"dog"});
assertReadingCompressedFileMatchesExpected(
- new File(filename), CompressionType.ZIP, Arrays.asList("cat", "dog"));
+ new File(filename), CompressionType.ZIP, new String[] {"cat", "dog"});
}
@Test
@@ -1346,21 +1104,5 @@ public class TextIOTest {
SourceTestUtils.assertSourcesEqualReferenceSource(source, splits, options);
}
-
- @Test
- @Category(NeedsRunner.class)
- public void testReadAll() throws IOException {
- writeToFile(TINY, "readAllTiny1.zip", ZIP);
- writeToFile(TINY, "readAllTiny2.zip", ZIP);
- writeToFile(LARGE, "readAllLarge1.zip", ZIP);
- writeToFile(LARGE, "readAllLarge2.zip", ZIP);
- PCollection<String> lines =
- p.apply(
- Create.of(
- tempFolder.resolve("readAllTiny*").toString(),
- tempFolder.resolve("readAllLarge*").toString()))
- .apply(TextIO.readAll().withCompressionType(AUTO));
- PAssert.that(lines).containsInAnyOrder(Iterables.concat(TINY, TINY, LARGE, LARGE));
- p.run();
- }
}
+
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java
index 55f2a87..a5dacd1 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.sdk.io;
-import static com.google.common.base.MoreObjects.firstNonNull;
import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem;
import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.includesDisplayDataFor;
import static org.hamcrest.Matchers.containsInAnyOrder;
@@ -42,11 +41,6 @@ import java.util.List;
import java.util.concurrent.ThreadLocalRandom;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.sdk.io.DefaultFilenamePolicy.Params;
-import org.apache.beam.sdk.io.FileBasedSink.CompressionType;
-import org.apache.beam.sdk.io.FileBasedSink.DynamicDestinations;
-import org.apache.beam.sdk.io.FileBasedSink.FilenamePolicy;
-import org.apache.beam.sdk.io.FileBasedSink.OutputFileHints;
import org.apache.beam.sdk.io.SimpleSink.SimpleWriter;
import org.apache.beam.sdk.io.fs.MatchResult.Metadata;
import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions;
@@ -63,24 +57,17 @@ import org.apache.beam.sdk.transforms.GroupByKey;
import org.apache.beam.sdk.transforms.MapElements;
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.SerializableFunctions;
import org.apache.beam.sdk.transforms.SimpleFunction;
import org.apache.beam.sdk.transforms.Top;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.display.DisplayData;
-import org.apache.beam.sdk.transforms.display.DisplayData.Builder;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
-import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
import org.apache.beam.sdk.transforms.windowing.Sessions;
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PCollection.IsBounded;
import org.apache.beam.sdk.values.PCollectionView;
import org.joda.time.Duration;
-import org.joda.time.format.DateTimeFormatter;
-import org.joda.time.format.ISODateTimeFormat;
import org.junit.Rule;
import org.junit.Test;
import org.junit.experimental.categories.Category;
@@ -173,11 +160,7 @@ public class WriteFilesTest {
public void testWrite() throws IOException {
List<String> inputs = Arrays.asList("Critical canary", "Apprehensive eagle",
"Intimidating pigeon", "Pedantic gull", "Frisky finch");
- runWrite(
- inputs,
- IDENTITY_MAP,
- getBaseOutputFilename(),
- WriteFiles.to(makeSimpleSink(), SerializableFunctions.<String>identity()));
+ runWrite(inputs, IDENTITY_MAP, getBaseOutputFilename());
}
/**
@@ -186,11 +169,7 @@ public class WriteFilesTest {
@Test
@Category(NeedsRunner.class)
public void testEmptyWrite() throws IOException {
- runWrite(
- Collections.<String>emptyList(),
- IDENTITY_MAP,
- getBaseOutputFilename(),
- WriteFiles.to(makeSimpleSink(), SerializableFunctions.<String>identity()));
+ runWrite(Collections.<String>emptyList(), IDENTITY_MAP, getBaseOutputFilename());
checkFileContents(getBaseOutputFilename(), Collections.<String>emptyList(),
Optional.of(1));
}
@@ -206,7 +185,7 @@ public class WriteFilesTest {
Arrays.asList("one", "two", "three", "four", "five", "six"),
IDENTITY_MAP,
getBaseOutputFilename(),
- WriteFiles.to(makeSimpleSink(), SerializableFunctions.<String>identity()).withNumShards(1));
+ Optional.of(1));
}
private ResourceId getBaseOutputDirectory() {
@@ -214,13 +193,8 @@ public class WriteFilesTest {
.resolve("output", StandardResolveOptions.RESOLVE_DIRECTORY);
}
-
- private SimpleSink<Void> makeSimpleSink() {
- FilenamePolicy filenamePolicy =
- new PerWindowFiles(
- getBaseOutputDirectory().resolve("file", StandardResolveOptions.RESOLVE_FILE),
- "simple");
- return SimpleSink.makeSimpleSink(getBaseOutputDirectory(), filenamePolicy);
+ private SimpleSink makeSimpleSink() {
+ return new SimpleSink(getBaseOutputDirectory(), "file", "-SS-of-NN", "simple");
}
@Test
@@ -239,10 +213,8 @@ public class WriteFilesTest {
timestamps.add(i + 1);
}
- SimpleSink<Void> sink = makeSimpleSink();
- WriteFiles<String, ?, String> write =
- WriteFiles.to(sink, SerializableFunctions.<String>identity())
- .withSharding(new LargestInt());
+ SimpleSink sink = makeSimpleSink();
+ WriteFiles<String> write = WriteFiles.to(sink).withSharding(new LargestInt());
p.apply(Create.timestamped(inputs, timestamps).withCoder(StringUtf8Coder.of()))
.apply(IDENTITY_MAP)
.apply(write);
@@ -263,8 +235,7 @@ public class WriteFilesTest {
Arrays.asList("one", "two", "three", "four", "five", "six"),
IDENTITY_MAP,
getBaseOutputFilename(),
- WriteFiles.to(makeSimpleSink(), SerializableFunctions.<String>identity())
- .withNumShards(20));
+ Optional.of(20));
}
/**
@@ -274,11 +245,7 @@ public class WriteFilesTest {
@Category(NeedsRunner.class)
public void testWriteWithEmptyPCollection() throws IOException {
List<String> inputs = new ArrayList<>();
- runWrite(
- inputs,
- IDENTITY_MAP,
- getBaseOutputFilename(),
- WriteFiles.to(makeSimpleSink(), SerializableFunctions.<String>identity()));
+ runWrite(inputs, IDENTITY_MAP, getBaseOutputFilename());
}
/**
@@ -290,10 +257,8 @@ public class WriteFilesTest {
List<String> inputs = Arrays.asList("Critical canary", "Apprehensive eagle",
"Intimidating pigeon", "Pedantic gull", "Frisky finch");
runWrite(
- inputs,
- new WindowAndReshuffle<>(Window.<String>into(FixedWindows.of(Duration.millis(2)))),
- getBaseOutputFilename(),
- WriteFiles.to(makeSimpleSink(), SerializableFunctions.<String>identity()));
+ inputs, new WindowAndReshuffle<>(Window.<String>into(FixedWindows.of(Duration.millis(2)))),
+ getBaseOutputFilename());
}
/**
@@ -307,32 +272,16 @@ public class WriteFilesTest {
runWrite(
inputs,
- new WindowAndReshuffle<>(Window.<String>into(Sessions.withGapDuration(Duration.millis(1)))),
- getBaseOutputFilename(),
- WriteFiles.to(makeSimpleSink(), SerializableFunctions.<String>identity()));
+ new WindowAndReshuffle<>(
+ Window.<String>into(Sessions.withGapDuration(Duration.millis(1)))),
+ getBaseOutputFilename());
}
@Test
- @Category(NeedsRunner.class)
- public void testWriteSpilling() throws IOException {
- List<String> inputs = Lists.newArrayList();
- for (int i = 0; i < 100; ++i) {
- inputs.add("mambo_number_" + i);
- }
- runWrite(
- inputs,
- Window.<String>into(FixedWindows.of(Duration.millis(2))),
- getBaseOutputFilename(),
- WriteFiles.to(makeSimpleSink(), SerializableFunctions.<String>identity())
- .withMaxNumWritersPerBundle(2)
- .withWindowedWrites());
- }
-
public void testBuildWrite() {
- SimpleSink<Void> sink = makeSimpleSink();
- WriteFiles<String, ?, String> write =
- WriteFiles.to(sink, SerializableFunctions.<String>identity()).withNumShards(3);
- assertThat((SimpleSink<Void>) write.getSink(), is(sink));
+ SimpleSink sink = makeSimpleSink();
+ WriteFiles<String> write = WriteFiles.to(sink).withNumShards(3);
+ assertThat((SimpleSink) write.getSink(), is(sink));
PTransform<PCollection<String>, PCollectionView<Integer>> originalSharding =
write.getSharding();
@@ -341,37 +290,25 @@ public class WriteFilesTest {
assertThat(write.getNumShards().get(), equalTo(3));
assertThat(write.getSharding(), equalTo(originalSharding));
- WriteFiles<String, ?, ?> write2 = write.withSharding(SHARDING_TRANSFORM);
- assertThat((SimpleSink<Void>) write2.getSink(), is(sink));
+ WriteFiles<String> write2 = write.withSharding(SHARDING_TRANSFORM);
+ assertThat((SimpleSink) write2.getSink(), is(sink));
assertThat(write2.getSharding(), equalTo(SHARDING_TRANSFORM));
// original unchanged
- WriteFiles<String, ?, ?> writeUnsharded = write2.withRunnerDeterminedSharding();
+ WriteFiles<String> writeUnsharded = write2.withRunnerDeterminedSharding();
assertThat(writeUnsharded.getSharding(), nullValue());
assertThat(write.getSharding(), equalTo(originalSharding));
}
@Test
public void testDisplayData() {
- DynamicDestinations<String, Void> dynamicDestinations =
- DynamicFileDestinations.constant(
- DefaultFilenamePolicy.fromParams(
- new Params()
- .withBaseFilename(
- getBaseOutputDirectory()
- .resolve("file", StandardResolveOptions.RESOLVE_FILE))
- .withShardTemplate("-SS-of-NN")));
- SimpleSink<Void> sink =
- new SimpleSink<Void>(
- getBaseOutputDirectory(), dynamicDestinations, CompressionType.UNCOMPRESSED) {
- @Override
- public void populateDisplayData(DisplayData.Builder builder) {
- builder.add(DisplayData.item("foo", "bar"));
- }
- };
- WriteFiles<String, ?, String> write =
- WriteFiles.to(sink, SerializableFunctions.<String>identity());
-
+ SimpleSink sink = new SimpleSink(getBaseOutputDirectory(), "file", "-SS-of-NN", "") {
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ builder.add(DisplayData.item("foo", "bar"));
+ }
+ };
+ WriteFiles<String> write = WriteFiles.to(sink);
DisplayData displayData = DisplayData.from(write);
assertThat(displayData, hasDisplayItem("sink", sink.getClass()));
@@ -379,145 +316,14 @@ public class WriteFilesTest {
}
@Test
- @Category(NeedsRunner.class)
- public void testUnboundedNeedsWindowed() {
- thrown.expect(IllegalArgumentException.class);
- thrown.expectMessage(
- "Must use windowed writes when applying WriteFiles to an unbounded PCollection");
-
- SimpleSink<Void> sink = makeSimpleSink();
- p.apply(Create.of("foo"))
- .setIsBoundedInternal(IsBounded.UNBOUNDED)
- .apply(WriteFiles.to(sink, SerializableFunctions.<String>identity()));
- p.run();
- }
-
- @Test
- @Category(NeedsRunner.class)
- public void testUnboundedNeedsSharding() {
- thrown.expect(IllegalArgumentException.class);
- thrown.expectMessage(
- "When applying WriteFiles to an unbounded PCollection, "
- + "must specify number of output shards explicitly");
-
- SimpleSink<Void> sink = makeSimpleSink();
- p.apply(Create.of("foo"))
- .setIsBoundedInternal(IsBounded.UNBOUNDED)
- .apply(WriteFiles.to(sink, SerializableFunctions.<String>identity()).withWindowedWrites());
- p.run();
- }
-
- // Test DynamicDestinations class. Expects user values to be string-encoded integers.
- // Stores the integer mod 5 as the destination, and uses that in the file prefix.
- static class TestDestinations extends DynamicDestinations<String, Integer> {
- private ResourceId baseOutputDirectory;
-
- TestDestinations(ResourceId baseOutputDirectory) {
- this.baseOutputDirectory = baseOutputDirectory;
- }
-
- @Override
- public Integer getDestination(String element) {
- return Integer.valueOf(element) % 5;
- }
-
- @Override
- public Integer getDefaultDestination() {
- return 0;
- }
-
- @Override
- public FilenamePolicy getFilenamePolicy(Integer destination) {
- return new PerWindowFiles(
- baseOutputDirectory.resolve("file_" + destination, StandardResolveOptions.RESOLVE_FILE),
- "simple");
- }
-
- @Override
- public void populateDisplayData(Builder builder) {
- super.populateDisplayData(builder);
- }
- }
-
- // Test format function. Prepend a string to each record before writing.
- static class TestDynamicFormatFunction implements SerializableFunction<String, String> {
- @Override
- public String apply(String input) {
- return "record_" + input;
- }
- }
-
- @Test
- @Category(NeedsRunner.class)
- public void testDynamicDestinationsBounded() throws Exception {
- testDynamicDestinationsHelper(true);
- }
-
- @Test
- @Category(NeedsRunner.class)
- public void testDynamicDestinationsUnbounded() throws Exception {
- testDynamicDestinationsHelper(false);
- }
-
- private void testDynamicDestinationsHelper(boolean bounded) throws IOException {
- TestDestinations dynamicDestinations = new TestDestinations(getBaseOutputDirectory());
- SimpleSink<Integer> sink =
- new SimpleSink<>(
- getBaseOutputDirectory(), dynamicDestinations, CompressionType.UNCOMPRESSED);
-
- // Flag to validate that the pipeline options are passed to the Sink.
- WriteOptions options = TestPipeline.testingPipelineOptions().as(WriteOptions.class);
- options.setTestFlag("test_value");
- Pipeline p = TestPipeline.create(options);
-
- List<String> inputs = Lists.newArrayList("0", "1", "2", "3", "4", "5", "6", "7", "8", "9");
- // Prepare timestamps for the elements.
- List<Long> timestamps = new ArrayList<>();
- for (long i = 0; i < inputs.size(); i++) {
- timestamps.add(i + 1);
- }
-
- WriteFiles<String, Integer, String> writeFiles =
- WriteFiles.to(sink, new TestDynamicFormatFunction()).withNumShards(1);
-
- PCollection<String> input = p.apply(Create.timestamped(inputs, timestamps));
- if (!bounded) {
- input.setIsBoundedInternal(IsBounded.UNBOUNDED);
- input = input.apply(Window.<String>into(FixedWindows.of(Duration.standardDays(1))));
- input.apply(writeFiles.withWindowedWrites());
- } else {
- input.apply(writeFiles);
- }
- p.run();
-
- for (int i = 0; i < 5; ++i) {
- ResourceId base =
- getBaseOutputDirectory().resolve("file_" + i, StandardResolveOptions.RESOLVE_FILE);
- List<String> expected = Lists.newArrayList("record_" + i, "record_" + (i + 5));
- checkFileContents(base.toString(), expected, Optional.of(1));
- }
- }
-
- @Test
public void testShardedDisplayData() {
- DynamicDestinations<String, Void> dynamicDestinations =
- DynamicFileDestinations.constant(
- DefaultFilenamePolicy.fromParams(
- new Params()
- .withBaseFilename(
- getBaseOutputDirectory()
- .resolve("file", StandardResolveOptions.RESOLVE_FILE))
- .withShardTemplate("-SS-of-NN")));
- SimpleSink<Void> sink =
- new SimpleSink<Void>(
- getBaseOutputDirectory(), dynamicDestinations, CompressionType.UNCOMPRESSED) {
- @Override
- public void populateDisplayData(DisplayData.Builder builder) {
- builder.add(DisplayData.item("foo", "bar"));
- }
- };
- WriteFiles<String, ?, String> write =
- WriteFiles.to(sink, SerializableFunctions.<String>identity()).withNumShards(1);
+ SimpleSink sink = new SimpleSink(getBaseOutputDirectory(), "file", "-SS-of-NN", "") {
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ builder.add(DisplayData.item("foo", "bar"));
+ }
+ };
+ WriteFiles<String> write = WriteFiles.to(sink).withNumShards(1);
DisplayData displayData = DisplayData.from(write);
assertThat(displayData, hasDisplayItem("sink", sink.getClass()));
assertThat(displayData, includesDisplayDataFor("sink", sink));
@@ -526,24 +332,14 @@ public class WriteFilesTest {
@Test
public void testCustomShardStrategyDisplayData() {
- DynamicDestinations<String, Void> dynamicDestinations =
- DynamicFileDestinations.constant(
- DefaultFilenamePolicy.fromParams(
- new Params()
- .withBaseFilename(
- getBaseOutputDirectory()
- .resolve("file", StandardResolveOptions.RESOLVE_FILE))
- .withShardTemplate("-SS-of-NN")));
- SimpleSink<Void> sink =
- new SimpleSink<Void>(
- getBaseOutputDirectory(), dynamicDestinations, CompressionType.UNCOMPRESSED) {
- @Override
- public void populateDisplayData(DisplayData.Builder builder) {
- builder.add(DisplayData.item("foo", "bar"));
- }
- };
- WriteFiles<String, ?, String> write =
- WriteFiles.to(sink, SerializableFunctions.<String>identity())
+ SimpleSink sink = new SimpleSink(getBaseOutputDirectory(), "file", "-SS-of-NN", "") {
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ builder.add(DisplayData.item("foo", "bar"));
+ }
+ };
+ WriteFiles<String> write =
+ WriteFiles.to(sink)
.withSharding(
new PTransform<PCollection<String>, PCollectionView<Integer>>() {
@Override
@@ -568,77 +364,22 @@ public class WriteFilesTest {
* PCollection are written to the sink.
*/
private void runWrite(
- List<String> inputs,
- PTransform<PCollection<String>, PCollection<String>> transform,
- String baseName,
- WriteFiles<String, ?, String> write)
- throws IOException {
- runShardedWrite(inputs, transform, baseName, write);
- }
-
- private static class PerWindowFiles extends FilenamePolicy {
- private static final DateTimeFormatter FORMATTER = ISODateTimeFormat.hourMinuteSecondMillis();
- private final ResourceId baseFilename;
- private final String suffix;
-
- public PerWindowFiles(ResourceId baseFilename, String suffix) {
- this.baseFilename = baseFilename;
- this.suffix = suffix;
- }
-
- public String filenamePrefixForWindow(IntervalWindow window) {
- String prefix =
- baseFilename.isDirectory() ? "" : firstNonNull(baseFilename.getFilename(), "");
- return String.format("%s%s-%s",
- prefix, FORMATTER.print(window.start()), FORMATTER.print(window.end()));
- }
-
- @Override
- public ResourceId windowedFilename(WindowedContext context, OutputFileHints outputFileHints) {
- IntervalWindow window = (IntervalWindow) context.getWindow();
- String filename =
- String.format(
- "%s-%s-of-%s%s%s",
- filenamePrefixForWindow(window),
- context.getShardNumber(),
- context.getNumShards(),
- outputFileHints.getSuggestedFilenameSuffix(),
- suffix);
- return baseFilename
- .getCurrentDirectory()
- .resolve(filename, StandardResolveOptions.RESOLVE_FILE);
- }
-
- @Override
- public ResourceId unwindowedFilename(Context context, OutputFileHints outputFileHints) {
- String prefix =
- baseFilename.isDirectory() ? "" : firstNonNull(baseFilename.getFilename(), "");
- String filename =
- String.format(
- "%s-%s-of-%s%s%s",
- prefix,
- context.getShardNumber(),
- context.getNumShards(),
- outputFileHints.getSuggestedFilenameSuffix(),
- suffix);
- return baseFilename
- .getCurrentDirectory()
- .resolve(filename, StandardResolveOptions.RESOLVE_FILE);
- }
+ List<String> inputs, PTransform<PCollection<String>, PCollection<String>> transform,
+ String baseName) throws IOException {
+ runShardedWrite(inputs, transform, baseName, Optional.<Integer>absent());
}
/**
* Performs a WriteFiles transform with the desired number of shards. Verifies the WriteFiles
* transform calls the appropriate methods on a test sink in the correct order, as well as
- * verifies that the elements of a PCollection are written to the sink. If numConfiguredShards is
- * not null, also verifies that the output number of shards is correct.
+ * verifies that the elements of a PCollection are written to the sink. If numConfiguredShards
+ * is not null, also verifies that the output number of shards is correct.
*/
private void runShardedWrite(
List<String> inputs,
PTransform<PCollection<String>, PCollection<String>> transform,
String baseName,
- WriteFiles<String, ?, String> write)
- throws IOException {
+ Optional<Integer> numConfiguredShards) throws IOException {
// Flag to validate that the pipeline options are passed to the Sink
WriteOptions options = TestPipeline.testingPipelineOptions().as(WriteOptions.class);
options.setTestFlag("test_value");
@@ -649,15 +390,18 @@ public class WriteFilesTest {
for (long i = 0; i < inputs.size(); i++) {
timestamps.add(i + 1);
}
+
+ SimpleSink sink = makeSimpleSink();
+ WriteFiles<String> write = WriteFiles.to(sink);
+ if (numConfiguredShards.isPresent()) {
+ write = write.withNumShards(numConfiguredShards.get());
+ }
p.apply(Create.timestamped(inputs, timestamps).withCoder(StringUtf8Coder.of()))
.apply(transform)
.apply(write);
p.run();
- Optional<Integer> numShards =
- (write.getNumShards() != null)
- ? Optional.of(write.getNumShards().get()) : Optional.<Integer>absent();
- checkFileContents(baseName, inputs, numShards);
+ checkFileContents(baseName, inputs, numConfiguredShards);
}
static void checkFileContents(String baseName, List<String> inputs,
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/options/PipelineOptionsValidatorTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/options/PipelineOptionsValidatorTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/options/PipelineOptionsValidatorTest.java
index f8cd00f..120d5ed 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/options/PipelineOptionsValidatorTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/options/PipelineOptionsValidatorTest.java
@@ -60,18 +60,6 @@ public class PipelineOptionsValidatorTest {
}
@Test
- public void testWhenRequiredOptionIsSetAndClearedCli() {
- expectedException.expect(IllegalArgumentException.class);
- expectedException.expectMessage("Missing required value for "
- + "[--object, \"Fake Description\"].");
-
- Required required = PipelineOptionsFactory.fromArgs(new String[]{"--object=blah"})
- .as(Required.class);
- required.setObject(null);
- PipelineOptionsValidator.validateCli(Required.class, required);
- }
-
- @Test
public void testWhenRequiredOptionIsNeverSet() {
expectedException.expect(IllegalArgumentException.class);
expectedException.expectMessage("Missing required value for "
@@ -82,17 +70,6 @@ public class PipelineOptionsValidatorTest {
PipelineOptionsValidator.validate(Required.class, required);
}
-
- @Test
- public void testWhenRequiredOptionIsNeverSetCli() {
- expectedException.expect(IllegalArgumentException.class);
- expectedException.expectMessage("Missing required value for "
- + "[--object, \"Fake Description\"].");
-
- Required required = PipelineOptionsFactory.fromArgs(new String[]{}).as(Required.class);
- PipelineOptionsValidator.validateCli(Required.class, required);
- }
-
@Test
public void testWhenRequiredOptionIsNeverSetOnSuperInterface() {
expectedException.expect(IllegalArgumentException.class);
@@ -104,16 +81,6 @@ public class PipelineOptionsValidatorTest {
PipelineOptionsValidator.validate(Required.class, options);
}
- @Test
- public void testWhenRequiredOptionIsNeverSetOnSuperInterfaceCli() {
- expectedException.expect(IllegalArgumentException.class);
- expectedException.expectMessage("Missing required value for "
- + "[--object, \"Fake Description\"].");
-
- PipelineOptions options = PipelineOptionsFactory.fromArgs(new String[]{}).create();
- PipelineOptionsValidator.validateCli(Required.class, options);
- }
-
/** A test interface that overrides the parent's method. */
public interface SubClassValidation extends Required {
@Override
@@ -133,17 +100,6 @@ public class PipelineOptionsValidatorTest {
PipelineOptionsValidator.validate(Required.class, required);
}
- @Test
- public void testValidationOnOverriddenMethodsCli() throws Exception {
- expectedException.expect(IllegalArgumentException.class);
- expectedException.expectMessage("Missing required value for "
- + "[--object, \"Fake Description\"].");
-
- SubClassValidation required = PipelineOptionsFactory.fromArgs(new String[]{})
- .as(SubClassValidation.class);
- PipelineOptionsValidator.validateCli(Required.class, required);
- }
-
/** A test interface with a required group. */
public interface GroupRequired extends PipelineOptions {
@Validation.Required(groups = {"ham"})
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/options/ProxyInvocationHandlerTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/options/ProxyInvocationHandlerTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/options/ProxyInvocationHandlerTest.java
index fb0a0d7..2c43f57 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/options/ProxyInvocationHandlerTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/options/ProxyInvocationHandlerTest.java
@@ -44,7 +44,6 @@ import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Maps;
import com.google.common.testing.EqualsTester;
import java.io.IOException;
-import java.io.NotSerializableException;
import java.io.Serializable;
import java.util.HashSet;
import java.util.List;
@@ -55,7 +54,6 @@ import org.apache.beam.sdk.testing.NeedsRunner;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.display.DisplayData;
-import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.util.common.ReflectHelpers;
import org.hamcrest.Matchers;
import org.joda.time.Instant;
@@ -1021,21 +1019,4 @@ public class ProxyInvocationHandlerTest {
DisplayData data = DisplayData.from(options);
assertThat(data, not(hasDisplayItem("value")));
}
-
- private static class CapturesOptions implements Serializable {
- PipelineOptions options = PipelineOptionsFactory.create();
- }
-
- @Test
- public void testOptionsAreNotSerializable() {
- expectedException.expectCause(Matchers.<Throwable>instanceOf(NotSerializableException.class));
- SerializableUtils.clone(new CapturesOptions());
- }
-
- @Test
- public void testGetOptionNameFromMethod() throws NoSuchMethodException {
- ProxyInvocationHandler handler = new ProxyInvocationHandler(Maps.<String, Object>newHashMap());
- handler.as(BaseOptions.class);
- assertEquals("foo", handler.getOptionName(BaseOptions.class.getMethod("getFoo")));
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/runners/TransformHierarchyTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/runners/TransformHierarchyTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/runners/TransformHierarchyTest.java
index 93650dd..1197d1b 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/runners/TransformHierarchyTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/runners/TransformHierarchyTest.java
@@ -19,7 +19,6 @@ package org.apache.beam.sdk.runners;
import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.hasItem;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.not;
import static org.junit.Assert.assertThat;
@@ -33,8 +32,6 @@ import java.util.Map.Entry;
import java.util.Set;
import org.apache.beam.sdk.Pipeline.PipelineVisitor;
import org.apache.beam.sdk.Pipeline.PipelineVisitor.Defaults;
-import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.io.CountingSource;
import org.apache.beam.sdk.io.GenerateSequence;
import org.apache.beam.sdk.io.Read;
@@ -495,198 +492,4 @@ public class TransformHierarchyTest implements Serializable {
assertThat(visitedPrimitiveNodes, containsInAnyOrder(upstreamNode, replacementParNode));
assertThat(visitedValues, Matchers.<PValue>containsInAnyOrder(upstream, output));
}
-
- @Test
- public void visitIsTopologicallyOrdered() {
- PCollection<String> one =
- PCollection.<String>createPrimitiveOutputInternal(
- pipeline, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)
- .setCoder(StringUtf8Coder.of());
- final PCollection<Integer> two =
- PCollection.<Integer>createPrimitiveOutputInternal(
- pipeline, WindowingStrategy.globalDefault(), IsBounded.UNBOUNDED)
- .setCoder(VarIntCoder.of());
- final PDone done = PDone.in(pipeline);
- final TupleTag<String> oneTag = new TupleTag<String>() {};
- final TupleTag<Integer> twoTag = new TupleTag<Integer>() {};
- final PCollectionTuple oneAndTwo = PCollectionTuple.of(oneTag, one).and(twoTag, two);
-
- PTransform<PCollection<String>, PDone> multiConsumer =
- new PTransform<PCollection<String>, PDone>() {
- @Override
- public PDone expand(PCollection<String> input) {
- return done;
- }
-
- @Override
- public Map<TupleTag<?>, PValue> getAdditionalInputs() {
- return Collections.<TupleTag<?>, PValue>singletonMap(twoTag, two);
- }
- };
- hierarchy.pushNode("consumes_both", one, multiConsumer);
- hierarchy.setOutput(done);
- hierarchy.popNode();
-
- final PTransform<PBegin, PCollectionTuple> producer =
- new PTransform<PBegin, PCollectionTuple>() {
- @Override
- public PCollectionTuple expand(PBegin input) {
- return oneAndTwo;
- }
- };
- hierarchy.pushNode(
- "encloses_producer",
- PBegin.in(pipeline),
- new PTransform<PBegin, PCollectionTuple>() {
- @Override
- public PCollectionTuple expand(PBegin input) {
- return input.apply(producer);
- }
- });
- hierarchy.pushNode(
- "creates_one_and_two",
- PBegin.in(pipeline), producer);
- hierarchy.setOutput(oneAndTwo);
- hierarchy.popNode();
- hierarchy.setOutput(oneAndTwo);
- hierarchy.popNode();
-
- hierarchy.pushNode("second_copy_of_consumes_both", one, multiConsumer);
- hierarchy.setOutput(done);
- hierarchy.popNode();
-
- final Set<Node> visitedNodes = new HashSet<>();
- final Set<Node> exitedNodes = new HashSet<>();
- final Set<PValue> visitedValues = new HashSet<>();
- hierarchy.visit(
- new PipelineVisitor.Defaults() {
-
- @Override
- public CompositeBehavior enterCompositeTransform(Node node) {
- for (PValue input : node.getInputs().values()) {
- assertThat(visitedValues, hasItem(input));
- }
- assertThat(
- "Nodes should not be visited more than once", visitedNodes, not(hasItem(node)));
- if (!node.isRootNode()) {
- assertThat(
- "Nodes should always be visited after their enclosing nodes",
- visitedNodes,
- hasItem(node.getEnclosingNode()));
- }
- visitedNodes.add(node);
- return CompositeBehavior.ENTER_TRANSFORM;
- }
-
- @Override
- public void leaveCompositeTransform(Node node) {
- assertThat(visitedNodes, hasItem(node));
- if (!node.isRootNode()) {
- assertThat(
- "Nodes should always be left before their enclosing nodes are left",
- exitedNodes,
- not(hasItem(node.getEnclosingNode())));
- }
- assertThat(exitedNodes, not(hasItem(node)));
- exitedNodes.add(node);
- }
-
- @Override
- public void visitPrimitiveTransform(Node node) {
- assertThat(visitedNodes, hasItem(node.getEnclosingNode()));
- assertThat(exitedNodes, not(hasItem(node.getEnclosingNode())));
- assertThat(
- "Nodes should not be visited more than once", visitedNodes, not(hasItem(node)));
- for (PValue input : node.getInputs().values()) {
- assertThat(visitedValues, hasItem(input));
- }
- visitedNodes.add(node);
- }
-
- @Override
- public void visitValue(PValue value, Node producer) {
- assertThat(visitedNodes, hasItem(producer));
- assertThat(visitedValues, not(hasItem(value)));
- visitedValues.add(value);
- }
- });
- assertThat("Should have visited all the nodes", visitedNodes.size(), equalTo(5));
- assertThat("Should have left all of the visited composites", exitedNodes.size(), equalTo(2));
- }
-
- @Test
- public void visitDoesNotVisitSkippedNodes() {
- PCollection<String> one =
- PCollection.<String>createPrimitiveOutputInternal(
- pipeline, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)
- .setCoder(StringUtf8Coder.of());
- final PCollection<Integer> two =
- PCollection.<Integer>createPrimitiveOutputInternal(
- pipeline, WindowingStrategy.globalDefault(), IsBounded.UNBOUNDED)
- .setCoder(VarIntCoder.of());
- final PDone done = PDone.in(pipeline);
- final TupleTag<String> oneTag = new TupleTag<String>() {};
- final TupleTag<Integer> twoTag = new TupleTag<Integer>() {};
- final PCollectionTuple oneAndTwo = PCollectionTuple.of(oneTag, one).and(twoTag, two);
-
- hierarchy.pushNode(
- "consumes_both",
- one,
- new PTransform<PCollection<String>, PDone>() {
- @Override
- public PDone expand(PCollection<String> input) {
- return done;
- }
-
- @Override
- public Map<TupleTag<?>, PValue> getAdditionalInputs() {
- return Collections.<TupleTag<?>, PValue>singletonMap(twoTag, two);
- }
- });
- hierarchy.setOutput(done);
- hierarchy.popNode();
-
- final PTransform<PBegin, PCollectionTuple> producer =
- new PTransform<PBegin, PCollectionTuple>() {
- @Override
- public PCollectionTuple expand(PBegin input) {
- return oneAndTwo;
- }
- };
- final Node enclosing =
- hierarchy.pushNode(
- "encloses_producer",
- PBegin.in(pipeline),
- new PTransform<PBegin, PCollectionTuple>() {
- @Override
- public PCollectionTuple expand(PBegin input) {
- return input.apply(producer);
- }
- });
- Node enclosed = hierarchy.pushNode("creates_one_and_two", PBegin.in(pipeline), producer);
- hierarchy.setOutput(oneAndTwo);
- hierarchy.popNode();
- hierarchy.setOutput(oneAndTwo);
- hierarchy.popNode();
-
- final Set<Node> visitedNodes = new HashSet<>();
- hierarchy.visit(
- new PipelineVisitor.Defaults() {
- @Override
- public CompositeBehavior enterCompositeTransform(Node node) {
- visitedNodes.add(node);
- return node.equals(enclosing)
- ? CompositeBehavior.DO_NOT_ENTER_TRANSFORM
- : CompositeBehavior.ENTER_TRANSFORM;
- }
-
- @Override
- public void visitPrimitiveTransform(Node node) {
- visitedNodes.add(node);
- }
- });
-
- assertThat(visitedNodes, hasItem(enclosing));
- assertThat(visitedNodes, not(hasItem(enclosed)));
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PCollectionViewTesting.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PCollectionViewTesting.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PCollectionViewTesting.java
index aaf8b91..adf27f8 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PCollectionViewTesting.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PCollectionViewTesting.java
@@ -22,9 +22,7 @@ import com.google.common.base.Function;
import com.google.common.base.MoreObjects;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
-import java.util.Collections;
import java.util.List;
-import java.util.Map;
import java.util.Objects;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.IterableCoder;
@@ -39,7 +37,6 @@ import org.apache.beam.sdk.transforms.windowing.WindowMappingFn;
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.PValue;
import org.apache.beam.sdk.values.PValueBase;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
@@ -352,10 +349,5 @@ public final class PCollectionViewTesting {
.add("viewFn", viewFn)
.toString();
}
-
- @Override
- public Map<TupleTag<?>, PValue> expand() {
- return Collections.<TupleTag<?>, PValue>singletonMap(tag, pCollection);
- }
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/CombineTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/CombineTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/CombineTest.java
index b24d82d..dc9788f 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/CombineTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/CombineTest.java
@@ -17,7 +17,7 @@
*/
package org.apache.beam.sdk.transforms;
-import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import static org.apache.beam.sdk.TestUtils.checkCombineFn;
import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem;
@@ -29,13 +29,11 @@ import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThat;
import com.google.common.base.MoreObjects;
-import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.Serializable;
-import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
@@ -47,6 +45,7 @@ import org.apache.beam.sdk.coders.AtomicCoder;
import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
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.CoderRegistry;
import org.apache.beam.sdk.coders.DoubleCoder;
import org.apache.beam.sdk.coders.KvCoder;
@@ -86,6 +85,7 @@ import org.junit.Test;
import org.junit.experimental.categories.Category;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
+import org.mockito.Mock;
/**
* Tests for Combine transforms.
@@ -95,8 +95,18 @@ public class CombineTest implements Serializable {
// This test is Serializable, just so that it's easy to have
// anonymous inner classes inside the non-static test methods.
+ static final List<KV<String, Integer>> TABLE = Arrays.asList(
+ KV.of("a", 1),
+ KV.of("a", 1),
+ KV.of("a", 4),
+ KV.of("b", 1),
+ KV.of("b", 13)
+ );
+
static final List<KV<String, Integer>> EMPTY_TABLE = Collections.emptyList();
+ @Mock private DoFn<?, ?>.ProcessContext processContext;
+
@Rule
public final transient TestPipeline pipeline = TestPipeline.create();
@@ -140,12 +150,12 @@ public class CombineTest implements Serializable {
PCollection<KV<String, String>> combinePerKey =
perKeyInput.apply(
Combine.<String, Integer, String>perKey(new TestCombineFnWithContext(globallySumView))
- .withSideInputs(globallySumView));
+ .withSideInputs(Arrays.asList(globallySumView)));
PCollection<String> combineGlobally = globallyInput
.apply(Combine.globally(new TestCombineFnWithContext(globallySumView))
.withoutDefaults()
- .withSideInputs(globallySumView));
+ .withSideInputs(Arrays.asList(globallySumView)));
PAssert.that(sum).containsInAnyOrder(globalSum);
PAssert.that(combinePerKey).containsInAnyOrder(perKeyCombines);
@@ -158,28 +168,16 @@ public class CombineTest implements Serializable {
@Category(ValidatesRunner.class)
@SuppressWarnings({"rawtypes", "unchecked"})
public void testSimpleCombine() {
- runTestSimpleCombine(Arrays.asList(
- KV.of("a", 1),
- KV.of("a", 1),
- KV.of("a", 4),
- KV.of("b", 1),
- KV.of("b", 13)
- ), 20, Arrays.asList(KV.of("a", "114"), KV.of("b", "113")));
+ runTestSimpleCombine(TABLE, 20, Arrays.asList(KV.of("a", "114"), KV.of("b", "113")));
}
@Test
@Category(ValidatesRunner.class)
@SuppressWarnings({"rawtypes", "unchecked"})
public void testSimpleCombineWithContext() {
- runTestSimpleCombineWithContext(Arrays.asList(
- KV.of("a", 1),
- KV.of("a", 1),
- KV.of("a", 4),
- KV.of("b", 1),
- KV.of("b", 13)
- ), 20,
- Arrays.asList(KV.of("a", "20:114"), KV.of("b", "20:113")),
- new String[] {"20:111134"});
+ runTestSimpleCombineWithContext(TABLE, 20,
+ Arrays.asList(KV.of("a", "01124"), KV.of("b", "01123")),
+ new String[] {"01111234"});
}
@Test
@@ -218,13 +216,7 @@ public class CombineTest implements Serializable {
@Test
@Category(ValidatesRunner.class)
public void testBasicCombine() {
- runTestBasicCombine(Arrays.asList(
- KV.of("a", 1),
- KV.of("a", 1),
- KV.of("a", 4),
- KV.of("b", 1),
- KV.of("b", 13)
- ), ImmutableSet.of(1, 13, 4), Arrays.asList(
+ runTestBasicCombine(TABLE, ImmutableSet.of(1, 13, 4), Arrays.asList(
KV.of("a", (Set<Integer>) ImmutableSet.of(1, 4)),
KV.of("b", (Set<Integer>) ImmutableSet.of(1, 13))));
}
@@ -259,16 +251,9 @@ public class CombineTest implements Serializable {
@Category(ValidatesRunner.class)
public void testFixedWindowsCombine() {
PCollection<KV<String, Integer>> input =
- pipeline
- .apply(
- Create.timestamped(
- TimestampedValue.of(KV.of("a", 1), new Instant(0L)),
- TimestampedValue.of(KV.of("a", 1), new Instant(1L)),
- TimestampedValue.of(KV.of("a", 4), new Instant(6L)),
- TimestampedValue.of(KV.of("b", 1), new Instant(7L)),
- TimestampedValue.of(KV.of("b", 13), new Instant(8L)))
- .withCoder(KvCoder.of(StringUtf8Coder.of(), BigEndianIntegerCoder.of())))
- .apply(Window.<KV<String, Integer>>into(FixedWindows.of(Duration.millis(2))));
+ pipeline.apply(Create.timestamped(TABLE, Arrays.asList(0L, 1L, 6L, 7L, 8L))
+ .withCoder(KvCoder.of(StringUtf8Coder.of(), BigEndianIntegerCoder.of())))
+ .apply(Window.<KV<String, Integer>>into(FixedWindows.of(Duration.millis(2))));
PCollection<Integer> sum = input
.apply(Values.<Integer>create())
@@ -278,9 +263,11 @@ public class CombineTest implements Serializable {
.apply(Combine.<String, Integer, String>perKey(new TestCombineFn()));
PAssert.that(sum).containsInAnyOrder(2, 5, 13);
- PAssert.that(sumPerKey)
- .containsInAnyOrder(
- Arrays.asList(KV.of("a", "11"), KV.of("a", "4"), KV.of("b", "1"), KV.of("b", "13")));
+ PAssert.that(sumPerKey).containsInAnyOrder(
+ KV.of("a", "11"),
+ KV.of("a", "4"),
+ KV.of("b", "1"),
+ KV.of("b", "13"));
pipeline.run();
}
@@ -288,16 +275,9 @@ public class CombineTest implements Serializable {
@Category(ValidatesRunner.class)
public void testFixedWindowsCombineWithContext() {
PCollection<KV<String, Integer>> perKeyInput =
- pipeline
- .apply(
- Create.timestamped(
- TimestampedValue.of(KV.of("a", 1), new Instant(0L)),
- TimestampedValue.of(KV.of("a", 1), new Instant(1L)),
- TimestampedValue.of(KV.of("a", 4), new Instant(6L)),
- TimestampedValue.of(KV.of("b", 1), new Instant(7L)),
- TimestampedValue.of(KV.of("b", 13), new Instant(8L)))
- .withCoder(KvCoder.of(StringUtf8Coder.of(), BigEndianIntegerCoder.of())))
- .apply(Window.<KV<String, Integer>>into(FixedWindows.of(Duration.millis(2))));
+ pipeline.apply(Create.timestamped(TABLE, Arrays.asList(0L, 1L, 6L, 7L, 8L))
+ .withCoder(KvCoder.of(StringUtf8Coder.of(), BigEndianIntegerCoder.of())))
+ .apply(Window.<KV<String, Integer>>into(FixedWindows.of(Duration.millis(2))));
PCollection<Integer> globallyInput = perKeyInput.apply(Values.<Integer>create());
@@ -309,129 +289,60 @@ public class CombineTest implements Serializable {
PCollection<KV<String, String>> combinePerKeyWithContext =
perKeyInput.apply(
Combine.<String, Integer, String>perKey(new TestCombineFnWithContext(globallySumView))
- .withSideInputs(globallySumView));
+ .withSideInputs(Arrays.asList(globallySumView)));
PCollection<String> combineGloballyWithContext = globallyInput
.apply(Combine.globally(new TestCombineFnWithContext(globallySumView))
.withoutDefaults()
- .withSideInputs(globallySumView));
+ .withSideInputs(Arrays.asList(globallySumView)));
PAssert.that(sum).containsInAnyOrder(2, 5, 13);
- PAssert.that(combinePerKeyWithContext)
- .containsInAnyOrder(
- Arrays.asList(
- KV.of("a", "2:11"), KV.of("a", "5:4"), KV.of("b", "5:1"), KV.of("b", "13:13")));
- PAssert.that(combineGloballyWithContext).containsInAnyOrder("2:11", "5:14", "13:13");
- pipeline.run();
- }
-
- @Test
- @Category(ValidatesRunner.class)
- public void testSlidingWindowsCombine() {
- PCollection<String> input =
- pipeline
- .apply(
- Create.timestamped(
- TimestampedValue.of("a", new Instant(1L)),
- TimestampedValue.of("b", new Instant(2L)),
- TimestampedValue.of("c", new Instant(3L))))
- .apply(
- Window.<String>into(
- SlidingWindows.of(Duration.millis(3)).every(Duration.millis(1L))));
- PCollection<List<String>> combined =
- input.apply(
- Combine.globally(
- new CombineFn<String, List<String>, List<String>>() {
- @Override
- public List<String> createAccumulator() {
- return new ArrayList<>();
- }
-
- @Override
- public List<String> addInput(List<String> accumulator, String input) {
- accumulator.add(input);
- return accumulator;
- }
-
- @Override
- public List<String> mergeAccumulators(Iterable<List<String>> accumulators) {
- // Mutate all of the accumulators. Instances should be used in only one
- // place, and not
- // reused after merging.
- List<String> cur = createAccumulator();
- for (List<String> accumulator : accumulators) {
- accumulator.addAll(cur);
- cur = accumulator;
- }
- return cur;
- }
-
- @Override
- public List<String> extractOutput(List<String> accumulator) {
- List<String> result = new ArrayList<>(accumulator);
- Collections.sort(result);
- return result;
- }
- })
- .withoutDefaults());
-
- PAssert.that(combined)
- .containsInAnyOrder(
- ImmutableList.of("a"),
- ImmutableList.of("a", "b"),
- ImmutableList.of("a", "b", "c"),
- ImmutableList.of("b", "c"),
- ImmutableList.of("c"));
-
+ PAssert.that(combinePerKeyWithContext).containsInAnyOrder(
+ KV.of("a", "112"),
+ KV.of("a", "45"),
+ KV.of("b", "15"),
+ KV.of("b", "1133"));
+ PAssert.that(combineGloballyWithContext).containsInAnyOrder("112", "145", "1133");
pipeline.run();
}
@Test
@Category(ValidatesRunner.class)
public void testSlidingWindowsCombineWithContext() {
- // [a: 1, 1], [a: 4; b: 1], [b: 13]
PCollection<KV<String, Integer>> perKeyInput =
- pipeline
- .apply(
- Create.timestamped(
- TimestampedValue.of(KV.of("a", 1), new Instant(2L)),
- TimestampedValue.of(KV.of("a", 1), new Instant(3L)),
- TimestampedValue.of(KV.of("a", 4), new Instant(8L)),
- TimestampedValue.of(KV.of("b", 1), new Instant(9L)),
- TimestampedValue.of(KV.of("b", 13), new Instant(10L)))
- .withCoder(KvCoder.of(StringUtf8Coder.of(), BigEndianIntegerCoder.of())))
- .apply(Window.<KV<String, Integer>>into(SlidingWindows.of(Duration.millis(2))));
+ pipeline.apply(Create.timestamped(TABLE, Arrays.asList(2L, 3L, 8L, 9L, 10L))
+ .withCoder(KvCoder.of(StringUtf8Coder.of(), BigEndianIntegerCoder.of())))
+ .apply(Window.<KV<String, Integer>>into(SlidingWindows.of(Duration.millis(2))));
PCollection<Integer> globallyInput = perKeyInput.apply(Values.<Integer>create());
- PCollection<Integer> sum = globallyInput.apply("Sum", Sum.integersGlobally().withoutDefaults());
+ PCollection<Integer> sum = globallyInput
+ .apply("Sum", Combine.globally(new SumInts()).withoutDefaults());
PCollectionView<Integer> globallySumView = sum.apply(View.<Integer>asSingleton());
PCollection<KV<String, String>> combinePerKeyWithContext =
perKeyInput.apply(
Combine.<String, Integer, String>perKey(new TestCombineFnWithContext(globallySumView))
- .withSideInputs(globallySumView));
+ .withSideInputs(Arrays.asList(globallySumView)));
PCollection<String> combineGloballyWithContext = globallyInput
.apply(Combine.globally(new TestCombineFnWithContext(globallySumView))
.withoutDefaults()
- .withSideInputs(globallySumView));
+ .withSideInputs(Arrays.asList(globallySumView)));
PAssert.that(sum).containsInAnyOrder(1, 2, 1, 4, 5, 14, 13);
- PAssert.that(combinePerKeyWithContext)
- .containsInAnyOrder(
- Arrays.asList(
- KV.of("a", "1:1"),
- KV.of("a", "2:11"),
- KV.of("a", "1:1"),
- KV.of("a", "4:4"),
- KV.of("a", "5:4"),
- KV.of("b", "5:1"),
- KV.of("b", "14:113"),
- KV.of("b", "13:13")));
+ PAssert.that(combinePerKeyWithContext).containsInAnyOrder(
+ KV.of("a", "11"),
+ KV.of("a", "112"),
+ KV.of("a", "11"),
+ KV.of("a", "44"),
+ KV.of("a", "45"),
+ KV.of("b", "15"),
+ KV.of("b", "11134"),
+ KV.of("b", "1133"));
PAssert.that(combineGloballyWithContext).containsInAnyOrder(
- "1:1", "2:11", "1:1", "4:4", "5:14", "14:113", "13:13");
+ "11", "112", "11", "44", "145", "11134", "1133");
pipeline.run();
}
@@ -472,16 +383,9 @@ public class CombineTest implements Serializable {
@Category(ValidatesRunner.class)
public void testSessionsCombine() {
PCollection<KV<String, Integer>> input =
- pipeline
- .apply(
- Create.timestamped(
- TimestampedValue.of(KV.of("a", 1), new Instant(0L)),
- TimestampedValue.of(KV.of("a", 1), new Instant(4L)),
- TimestampedValue.of(KV.of("a", 4), new Instant(7L)),
- TimestampedValue.of(KV.of("b", 1), new Instant(10L)),
- TimestampedValue.of(KV.of("b", 13), new Instant(16L)))
- .withCoder(KvCoder.of(StringUtf8Coder.of(), BigEndianIntegerCoder.of())))
- .apply(Window.<KV<String, Integer>>into(Sessions.withGapDuration(Duration.millis(5))));
+ pipeline.apply(Create.timestamped(TABLE, Arrays.asList(0L, 4L, 7L, 10L, 16L))
+ .withCoder(KvCoder.of(StringUtf8Coder.of(), BigEndianIntegerCoder.of())))
+ .apply(Window.<KV<String, Integer>>into(Sessions.withGapDuration(Duration.millis(5))));
PCollection<Integer> sum = input
.apply(Values.<Integer>create())
@@ -491,8 +395,10 @@ public class CombineTest implements Serializable {
.apply(Combine.<String, Integer, String>perKey(new TestCombineFn()));
PAssert.that(sum).containsInAnyOrder(7, 13);
- PAssert.that(sumPerKey)
- .containsInAnyOrder(Arrays.asList(KV.of("a", "114"), KV.of("b", "1"), KV.of("b", "13")));
+ PAssert.that(sumPerKey).containsInAnyOrder(
+ KV.of("a", "114"),
+ KV.of("b", "1"),
+ KV.of("b", "13"));
pipeline.run();
}
@@ -500,13 +406,7 @@ public class CombineTest implements Serializable {
@Category(ValidatesRunner.class)
public void testSessionsCombineWithContext() {
PCollection<KV<String, Integer>> perKeyInput =
- pipeline.apply(
- Create.timestamped(
- TimestampedValue.of(KV.of("a", 1), new Instant(0L)),
- TimestampedValue.of(KV.of("a", 1), new Instant(4L)),
- TimestampedValue.of(KV.of("a", 4), new Instant(7L)),
- TimestampedValue.of(KV.of("b", 1), new Instant(10L)),
- TimestampedValue.of(KV.of("b", 13), new Instant(16L)))
+ pipeline.apply(Create.timestamped(TABLE, Arrays.asList(0L, 4L, 7L, 10L, 16L))
.withCoder(KvCoder.of(StringUtf8Coder.of(), BigEndianIntegerCoder.of())));
PCollection<Integer> globallyInput = perKeyInput.apply(Values.<Integer>create());
@@ -527,23 +427,21 @@ public class CombineTest implements Serializable {
.apply(
Combine.<String, Integer, String>perKey(
new TestCombineFnWithContext(globallyFixedWindowsView))
- .withSideInputs(globallyFixedWindowsView));
+ .withSideInputs(Arrays.asList(globallyFixedWindowsView)));
- PCollection<String> sessionsCombineGlobally =
- globallyInput
- .apply(
- "Globally Input Sessions",
- Window.<Integer>into(Sessions.withGapDuration(Duration.millis(5))))
- .apply(
- Combine.globally(new TestCombineFnWithContext(globallyFixedWindowsView))
- .withoutDefaults()
- .withSideInputs(globallyFixedWindowsView));
+ PCollection<String> sessionsCombineGlobally = globallyInput
+ .apply("Globally Input Sessions",
+ Window.<Integer>into(Sessions.withGapDuration(Duration.millis(5))))
+ .apply(Combine.globally(new TestCombineFnWithContext(globallyFixedWindowsView))
+ .withoutDefaults()
+ .withSideInputs(Arrays.asList(globallyFixedWindowsView)));
PAssert.that(fixedWindowsSum).containsInAnyOrder(2, 4, 1, 13);
- PAssert.that(sessionsCombinePerKey)
- .containsInAnyOrder(
- Arrays.asList(KV.of("a", "1:114"), KV.of("b", "1:1"), KV.of("b", "0:13")));
- PAssert.that(sessionsCombineGlobally).containsInAnyOrder("1:1114", "0:13");
+ PAssert.that(sessionsCombinePerKey).containsInAnyOrder(
+ KV.of("a", "1114"),
+ KV.of("b", "11"),
+ KV.of("b", "013"));
+ PAssert.that(sessionsCombineGlobally).containsInAnyOrder("11114", "013");
pipeline.run();
}
@@ -563,13 +461,7 @@ public class CombineTest implements Serializable {
@Test
@Category(ValidatesRunner.class)
public void testAccumulatingCombine() {
- runTestAccumulatingCombine(Arrays.asList(
- KV.of("a", 1),
- KV.of("a", 1),
- KV.of("a", 4),
- KV.of("b", 1),
- KV.of("b", 13)
- ), 4.0, Arrays.asList(KV.of("a", 2.0), KV.of("b", 7.0)));
+ runTestAccumulatingCombine(TABLE, 4.0, Arrays.asList(KV.of("a", 2.0), KV.of("b", 7.0)));
}
@Test
@@ -611,13 +503,7 @@ public class CombineTest implements Serializable {
@Test
@Category(ValidatesRunner.class)
public void testHotKeyCombining() {
- PCollection<KV<String, Integer>> input = copy(createInput(pipeline, Arrays.asList(
- KV.of("a", 1),
- KV.of("a", 1),
- KV.of("a", 4),
- KV.of("b", 1),
- KV.of("b", 13)
- )), 10);
+ PCollection<KV<String, Integer>> input = copy(createInput(pipeline, TABLE), 10);
CombineFn<Integer, ?, Double> mean = new MeanInts();
PCollection<KV<String, Double>> coldMean = input.apply("ColdMean",
@@ -674,13 +560,7 @@ public class CombineTest implements Serializable {
@Test
@Category(NeedsRunner.class)
public void testBinaryCombineFn() {
- PCollection<KV<String, Integer>> input = copy(createInput(pipeline, Arrays.asList(
- KV.of("a", 1),
- KV.of("a", 1),
- KV.of("a", 4),
- KV.of("b", 1),
- KV.of("b", 13)
- )), 2);
+ PCollection<KV<String, Integer>> input = copy(createInput(pipeline, TABLE), 2);
PCollection<KV<String, Integer>> intProduct = input
.apply("IntProduct", Combine.<String, Integer, Integer>perKey(new TestProdInt()));
PCollection<KV<String, Integer>> objProduct = input
@@ -771,7 +651,7 @@ public class CombineTest implements Serializable {
pipeline
.apply(
"CreateMainInput",
- Create.timestamped(nonEmptyElement, emptyElement).withCoder(VoidCoder.of()))
+ Create.<Void>timestamped(nonEmptyElement, emptyElement).withCoder(VoidCoder.of()))
.apply("WindowMainInput", Window.<Void>into(windowFn))
.apply(
"OutputSideInput",
@@ -996,13 +876,15 @@ public class CombineTest implements Serializable {
*/
private class CountSumCoder extends AtomicCoder<CountSum> {
@Override
- public void encode(CountSum value, OutputStream outStream) throws IOException {
+ public void encode(CountSum value, OutputStream outStream)
+ throws CoderException, IOException {
LONG_CODER.encode(value.count, outStream);
DOUBLE_CODER.encode(value.sum, outStream);
}
@Override
- public CountSum decode(InputStream inStream) throws IOException {
+ public CountSum decode(InputStream inStream)
+ throws CoderException, IOException {
long count = LONG_CODER.decode(inStream);
double sum = DOUBLE_CODER.decode(inStream);
return new CountSum(count, sum);
@@ -1035,26 +917,34 @@ public class CombineTest implements Serializable {
// Not serializable.
static class Accumulator {
- final String seed;
String value;
- public Accumulator(String seed, String value) {
- this.seed = seed;
+ public Accumulator(String value) {
this.value = value;
}
public static Coder<Accumulator> getCoder() {
return new AtomicCoder<Accumulator>() {
@Override
- public void encode(Accumulator accumulator, OutputStream outStream) throws IOException {
- StringUtf8Coder.of().encode(accumulator.seed, outStream);
- StringUtf8Coder.of().encode(accumulator.value, outStream);
+ public void encode(Accumulator accumulator, OutputStream outStream)
+ throws CoderException, IOException {
+ encode(accumulator, outStream, Coder.Context.NESTED);
}
@Override
- public Accumulator decode(InputStream inStream) throws IOException {
- String seed = StringUtf8Coder.of().decode(inStream);
- String value = StringUtf8Coder.of().decode(inStream);
- return new Accumulator(seed, value);
+ public void encode(Accumulator accumulator, OutputStream outStream, Coder.Context context)
+ throws CoderException, IOException {
+ StringUtf8Coder.of().encode(accumulator.value, outStream, context);
+ }
+
+ @Override
+ public Accumulator decode(InputStream inStream) throws CoderException, IOException {
+ return decode(inStream, Coder.Context.NESTED);
+ }
+
+ @Override
+ public Accumulator decode(InputStream inStream, Coder.Context context)
+ throws CoderException, IOException {
+ return new Accumulator(StringUtf8Coder.of().decode(inStream, context));
}
};
}
@@ -1068,13 +958,13 @@ public class CombineTest implements Serializable {
@Override
public Accumulator createAccumulator() {
- return new Accumulator("", "");
+ return new Accumulator("");
}
@Override
public Accumulator addInput(Accumulator accumulator, Integer value) {
try {
- return new Accumulator(accumulator.seed, accumulator.value + String.valueOf(value));
+ return new Accumulator(accumulator.value + String.valueOf(value));
} finally {
accumulator.value = "cleared in addInput";
}
@@ -1082,22 +972,12 @@ public class CombineTest implements Serializable {
@Override
public Accumulator mergeAccumulators(Iterable<Accumulator> accumulators) {
- Accumulator seedAccumulator = null;
- StringBuilder all = new StringBuilder();
+ String all = "";
for (Accumulator accumulator : accumulators) {
- if (seedAccumulator == null) {
- seedAccumulator = accumulator;
- } else {
- assertEquals(
- String.format(
- "Different seed values in accumulator: %s vs. %s", seedAccumulator, accumulator),
- seedAccumulator.seed,
- accumulator.seed);
- }
- all.append(accumulator.value);
+ all += accumulator.value;
accumulator.value = "cleared in mergeAccumulators";
}
- return new Accumulator(checkNotNull(seedAccumulator).seed, all.toString());
+ return new Accumulator(all);
}
@Override
@@ -1127,47 +1007,40 @@ public class CombineTest implements Serializable {
@Override
public TestCombineFn.Accumulator createAccumulator(Context c) {
- Integer sideInputValue = c.sideInput(view);
- return new TestCombineFn.Accumulator(sideInputValue.toString(), "");
+ return new TestCombineFn.Accumulator(c.sideInput(view).toString());
}
@Override
public TestCombineFn.Accumulator addInput(
TestCombineFn.Accumulator accumulator, Integer value, Context c) {
try {
- assertThat(
- "Not expecting view contents to change",
- accumulator.seed,
- Matchers.equalTo(Integer.toString(c.sideInput(view))));
- return new TestCombineFn.Accumulator(
- accumulator.seed, accumulator.value + String.valueOf(value));
+ assertThat(accumulator.value, Matchers.startsWith(c.sideInput(view).toString()));
+ return new TestCombineFn.Accumulator(accumulator.value + String.valueOf(value));
} finally {
accumulator.value = "cleared in addInput";
}
+
}
@Override
public TestCombineFn.Accumulator mergeAccumulators(
Iterable<TestCombineFn.Accumulator> accumulators, Context c) {
- String sideInputValue = c.sideInput(view).toString();
- StringBuilder all = new StringBuilder();
+ String prefix = c.sideInput(view).toString();
+ String all = prefix;
for (TestCombineFn.Accumulator accumulator : accumulators) {
- assertThat(
- "Accumulators should all have the same Side Input Value",
- accumulator.seed,
- Matchers.equalTo(sideInputValue));
- all.append(accumulator.value);
+ assertThat(accumulator.value, Matchers.startsWith(prefix));
+ all += accumulator.value.substring(prefix.length());
accumulator.value = "cleared in mergeAccumulators";
}
- return new TestCombineFn.Accumulator(sideInputValue, all.toString());
+ return new TestCombineFn.Accumulator(all);
}
@Override
public String extractOutput(TestCombineFn.Accumulator accumulator, Context c) {
- assertThat(accumulator.seed, Matchers.startsWith(c.sideInput(view).toString()));
+ assertThat(accumulator.value, Matchers.startsWith(c.sideInput(view).toString()));
char[] chars = accumulator.value.toCharArray();
Arrays.sort(chars);
- return accumulator.seed + ":" + new String(chars);
+ return new String(chars);
}
}
@@ -1205,7 +1078,7 @@ public class CombineTest implements Serializable {
@Override
public void mergeAccumulator(Counter accumulator) {
checkState(outputs == 0);
- assertEquals(0, accumulator.outputs);
+ checkArgument(accumulator.outputs == 0);
merges += accumulator.merges + 1;
inputs += accumulator.inputs;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/DoFnTesterTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/DoFnTesterTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/DoFnTesterTest.java
index 5cb9e18..1bb71bb 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/DoFnTesterTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/DoFnTesterTest.java
@@ -360,38 +360,6 @@ public class DoFnTesterTest {
}
}
- @Test
- public void testSupportsFinishBundleOutput() throws Exception {
- for (DoFnTester.CloningBehavior cloning : DoFnTester.CloningBehavior.values()) {
- try (DoFnTester<Integer, Integer> tester = DoFnTester.of(new BundleCounterDoFn())) {
- tester.setCloningBehavior(cloning);
-
- assertThat(tester.processBundle(1, 2, 3, 4), contains(4));
- assertThat(tester.processBundle(5, 6, 7), contains(3));
- assertThat(tester.processBundle(8, 9), contains(2));
- }
- }
- }
-
- private static class BundleCounterDoFn extends DoFn<Integer, Integer> {
- private int elements;
-
- @StartBundle
- public void startBundle() {
- elements = 0;
- }
-
- @ProcessElement
- public void processElement(ProcessContext c) {
- elements++;
- }
-
- @FinishBundle
- public void finishBundle(FinishBundleContext c) {
- c.output(elements, Instant.now(), GlobalWindow.INSTANCE);
- }
- }
-
private static class SideInputDoFn extends DoFn<Integer, Integer> {
private final PCollectionView<Integer> value;
[04/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/io/gcp/pubsub.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/io/gcp/pubsub.py b/sdks/python/apache_beam/io/gcp/pubsub.py
index 32d388a..1ba8ac0 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub.py
@@ -24,35 +24,29 @@ This API is currently under development and is subject to change.
from __future__ import absolute_import
-import re
-
from apache_beam import coders
from apache_beam.io.iobase import Read
from apache_beam.io.iobase import Write
from apache_beam.runners.dataflow.native_io import iobase as dataflow_io
-from apache_beam.transforms import core
from apache_beam.transforms import PTransform
-from apache_beam.transforms import Map
-from apache_beam.transforms import window
+from apache_beam.transforms import ParDo
from apache_beam.transforms.display import DisplayDataItem
-__all__ = ['ReadStringsFromPubSub', 'WriteStringsToPubSub']
+__all__ = ['ReadStringsFromPubSub', 'WriteStringsToPubSub',
+ 'PubSubSource', 'PubSubSink']
class ReadStringsFromPubSub(PTransform):
"""A ``PTransform`` for reading utf-8 string payloads from Cloud Pub/Sub."""
- def __init__(self, topic=None, subscription=None, id_label=None):
+ def __init__(self, topic, subscription=None, id_label=None):
"""Initializes ``ReadStringsFromPubSub``.
Attributes:
- topic: Cloud Pub/Sub topic in the form "projects/<project>/topics/
- <topic>". If provided, subscription must be None.
- subscription: Existing Cloud Pub/Sub subscription to use in the
- form "projects/<project>/subscriptions/<subscription>". If not
- specified, a temporary subscription will be created from the specified
- topic. If provided, topic must be None.
+ topic: Cloud Pub/Sub topic in the form "/topics/<project>/<topic>".
+ subscription: Optional existing Cloud Pub/Sub subscription to use in the
+ form "projects/<project>/subscriptions/<subscription>".
id_label: The attribute on incoming Pub/Sub messages to use as a unique
record identifier. When specified, the value of this attribute (which
can be any string that uniquely identifies the record) will be used for
@@ -66,13 +60,10 @@ class ReadStringsFromPubSub(PTransform):
subscription=subscription,
id_label=id_label)
- def get_windowing(self, unused_inputs):
- return core.Windowing(window.GlobalWindows())
-
def expand(self, pvalue):
pcoll = pvalue.pipeline | Read(self._source)
pcoll.element_type = bytes
- pcoll = pcoll | 'DecodeString' >> Map(lambda b: b.decode('utf-8'))
+ pcoll = pcoll | 'decode string' >> ParDo(_decodeUtf8String)
pcoll.element_type = unicode
return pcoll
@@ -90,50 +81,18 @@ class WriteStringsToPubSub(PTransform):
self._sink = _PubSubPayloadSink(topic)
def expand(self, pcoll):
- pcoll = pcoll | 'EncodeString' >> Map(lambda s: s.encode('utf-8'))
+ pcoll = pcoll | 'encode string' >> ParDo(_encodeUtf8String)
pcoll.element_type = bytes
return pcoll | Write(self._sink)
-PROJECT_ID_REGEXP = '[a-z][-a-z0-9:.]{4,61}[a-z0-9]'
-SUBSCRIPTION_REGEXP = 'projects/([^/]+)/subscriptions/(.+)'
-TOPIC_REGEXP = 'projects/([^/]+)/topics/(.+)'
-
-
-def parse_topic(full_topic):
- match = re.match(TOPIC_REGEXP, full_topic)
- if not match:
- raise ValueError(
- 'PubSub topic must be in the form "projects/<project>/topics'
- '/<topic>" (got %r).' % full_topic)
- project, topic_name = match.group(1), match.group(2)
- if not re.match(PROJECT_ID_REGEXP, project):
- raise ValueError('Invalid PubSub project name: %r.' % project)
- return project, topic_name
-
-
-def parse_subscription(full_subscription):
- match = re.match(SUBSCRIPTION_REGEXP, full_subscription)
- if not match:
- raise ValueError(
- 'PubSub subscription must be in the form "projects/<project>'
- '/subscriptions/<subscription>" (got %r).' % full_subscription)
- project, subscription_name = match.group(1), match.group(2)
- if not re.match(PROJECT_ID_REGEXP, project):
- raise ValueError('Invalid PubSub project name: %r.' % project)
- return project, subscription_name
-
-
class _PubSubPayloadSource(dataflow_io.NativeSource):
"""Source for the payload of a message as bytes from a Cloud Pub/Sub topic.
Attributes:
- topic: Cloud Pub/Sub topic in the form "projects/<project>/topics/<topic>".
- If provided, subscription must be None.
- subscription: Existing Cloud Pub/Sub subscription to use in the
- form "projects/<project>/subscriptions/<subscription>". If not specified,
- a temporary subscription will be created from the specified topic. If
- provided, topic must be None.
+ topic: Cloud Pub/Sub topic in the form "/topics/<project>/<topic>".
+ subscription: Optional existing Cloud Pub/Sub subscription to use in the
+ form "projects/<project>/subscriptions/<subscription>".
id_label: The attribute on incoming Pub/Sub messages to use as a unique
record identifier. When specified, the value of this attribute (which can
be any string that uniquely identifies the record) will be used for
@@ -142,27 +101,11 @@ class _PubSubPayloadSource(dataflow_io.NativeSource):
case, deduplication of the stream will be strictly best effort.
"""
- def __init__(self, topic=None, subscription=None, id_label=None):
- # We are using this coder explicitly for portability reasons of PubsubIO
- # across implementations in languages.
- self.coder = coders.BytesCoder()
- self.full_topic = topic
- self.full_subscription = subscription
- self.topic_name = None
- self.subscription_name = None
+ def __init__(self, topic, subscription=None, id_label=None):
+ self.topic = topic
+ self.subscription = subscription
self.id_label = id_label
- # Perform some validation on the topic and subscription.
- if not (topic or subscription):
- raise ValueError('Either a topic or subscription must be provided.')
- if topic and subscription:
- raise ValueError('Only one of topic or subscription should be provided.')
-
- if topic:
- self.project, self.topic_name = parse_topic(topic)
- if subscription:
- self.project, self.subscription_name = parse_subscription(subscription)
-
@property
def format(self):
"""Source format name required for remote execution."""
@@ -173,10 +116,10 @@ class _PubSubPayloadSource(dataflow_io.NativeSource):
DisplayDataItem(self.id_label,
label='ID Label Attribute').drop_if_none(),
'topic':
- DisplayDataItem(self.full_topic,
- label='Pubsub Topic').drop_if_none(),
+ DisplayDataItem(self.topic,
+ label='Pubsub Topic'),
'subscription':
- DisplayDataItem(self.full_subscription,
+ DisplayDataItem(self.subscription,
label='Pubsub Subscription').drop_if_none()}
def reader(self):
@@ -188,12 +131,7 @@ class _PubSubPayloadSink(dataflow_io.NativeSink):
"""Sink for the payload of a message as bytes to a Cloud Pub/Sub topic."""
def __init__(self, topic):
- # we are using this coder explicitly for portability reasons of PubsubIO
- # across implementations in languages.
- self.coder = coders.BytesCoder()
- self.full_topic = topic
-
- self.project, self.topic_name = parse_topic(topic)
+ self.topic = topic
@property
def format(self):
@@ -201,8 +139,86 @@ class _PubSubPayloadSink(dataflow_io.NativeSink):
return 'pubsub'
def display_data(self):
- return {'topic': DisplayDataItem(self.full_topic, label='Pubsub Topic')}
+ return {'topic': DisplayDataItem(self.topic, label='Pubsub Topic')}
def writer(self):
raise NotImplementedError(
'PubSubPayloadSink is not supported in local execution.')
+
+
+def _decodeUtf8String(encoded_value):
+ """Decodes a string in utf-8 format from bytes"""
+ return encoded_value.decode('utf-8')
+
+
+def _encodeUtf8String(value):
+ """Encodes a string in utf-8 format to bytes"""
+ return value.encode('utf-8')
+
+
+class PubSubSource(dataflow_io.NativeSource):
+ """Deprecated: do not use.
+
+ Source for reading from a given Cloud Pub/Sub topic.
+
+ Attributes:
+ topic: Cloud Pub/Sub topic in the form "/topics/<project>/<topic>".
+ subscription: Optional existing Cloud Pub/Sub subscription to use in the
+ form "projects/<project>/subscriptions/<subscription>".
+ id_label: The attribute on incoming Pub/Sub messages to use as a unique
+ record identifier. When specified, the value of this attribute (which can
+ be any string that uniquely identifies the record) will be used for
+ deduplication of messages. If not provided, Dataflow cannot guarantee
+ that no duplicate data will be delivered on the Pub/Sub stream. In this
+ case, deduplication of the stream will be strictly best effort.
+ coder: The Coder to use for decoding incoming Pub/Sub messages.
+ """
+
+ def __init__(self, topic, subscription=None, id_label=None,
+ coder=coders.StrUtf8Coder()):
+ self.topic = topic
+ self.subscription = subscription
+ self.id_label = id_label
+ self.coder = coder
+
+ @property
+ def format(self):
+ """Source format name required for remote execution."""
+ return 'pubsub'
+
+ def display_data(self):
+ return {'id_label':
+ DisplayDataItem(self.id_label,
+ label='ID Label Attribute').drop_if_none(),
+ 'topic':
+ DisplayDataItem(self.topic,
+ label='Pubsub Topic'),
+ 'subscription':
+ DisplayDataItem(self.subscription,
+ label='Pubsub Subscription').drop_if_none()}
+
+ def reader(self):
+ raise NotImplementedError(
+ 'PubSubSource is not supported in local execution.')
+
+
+class PubSubSink(dataflow_io.NativeSink):
+ """Deprecated: do not use.
+
+ Sink for writing to a given Cloud Pub/Sub topic."""
+
+ def __init__(self, topic, coder=coders.StrUtf8Coder()):
+ self.topic = topic
+ self.coder = coder
+
+ @property
+ def format(self):
+ """Sink format name required for remote execution."""
+ return 'pubsub'
+
+ def display_data(self):
+ return {'topic': DisplayDataItem(self.topic, label='Pubsub Topic')}
+
+ def writer(self):
+ raise NotImplementedError(
+ 'PubSubSink is not supported in local execution.')
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/io/gcp/pubsub_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/io/gcp/pubsub_test.py b/sdks/python/apache_beam/io/gcp/pubsub_test.py
index 0dcc3c3..322d08a 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub_test.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub_test.py
@@ -22,6 +22,8 @@ import unittest
import hamcrest as hc
+from apache_beam.io.gcp.pubsub import _decodeUtf8String
+from apache_beam.io.gcp.pubsub import _encodeUtf8String
from apache_beam.io.gcp.pubsub import _PubSubPayloadSink
from apache_beam.io.gcp.pubsub import _PubSubPayloadSource
from apache_beam.io.gcp.pubsub import ReadStringsFromPubSub
@@ -31,112 +33,77 @@ from apache_beam.transforms.display import DisplayData
from apache_beam.transforms.display_test import DisplayDataItemMatcher
-# Protect against environments where the PubSub library is not available.
-# pylint: disable=wrong-import-order, wrong-import-position
-try:
- from google.cloud import pubsub
-except ImportError:
- pubsub = None
-# pylint: enable=wrong-import-order, wrong-import-position
-
-
-@unittest.skipIf(pubsub is None, 'GCP dependencies are not installed')
class TestReadStringsFromPubSub(unittest.TestCase):
- def test_expand_with_topic(self):
+ def test_expand(self):
p = TestPipeline()
- pcoll = p | ReadStringsFromPubSub('projects/fakeprj/topics/a_topic',
- None, 'a_label')
+ pcoll = p | ReadStringsFromPubSub('a_topic', 'a_subscription', 'a_label')
# Ensure that the output type is str
self.assertEqual(unicode, pcoll.element_type)
- # Ensure that the properties passed through correctly
- source = pcoll.producer.transform._source
- self.assertEqual('a_topic', source.topic_name)
- self.assertEqual('a_label', source.id_label)
-
- def test_expand_with_subscription(self):
- p = TestPipeline()
- pcoll = p | ReadStringsFromPubSub(
- None, 'projects/fakeprj/subscriptions/a_subscription', 'a_label')
- # Ensure that the output type is str
- self.assertEqual(unicode, pcoll.element_type)
+ # Ensure that the type on the intermediate read output PCollection is bytes
+ read_pcoll = pcoll.producer.inputs[0]
+ self.assertEqual(bytes, read_pcoll.element_type)
# Ensure that the properties passed through correctly
- source = pcoll.producer.transform._source
- self.assertEqual('a_subscription', source.subscription_name)
+ source = read_pcoll.producer.transform.source
+ self.assertEqual('a_topic', source.topic)
+ self.assertEqual('a_subscription', source.subscription)
self.assertEqual('a_label', source.id_label)
- def test_expand_with_no_topic_or_subscription(self):
- with self.assertRaisesRegexp(
- ValueError, "Either a topic or subscription must be provided."):
- ReadStringsFromPubSub(None, None, 'a_label')
-
- def test_expand_with_both_topic_and_subscription(self):
- with self.assertRaisesRegexp(
- ValueError, "Only one of topic or subscription should be provided."):
- ReadStringsFromPubSub('a_topic', 'a_subscription', 'a_label')
-
-@unittest.skipIf(pubsub is None, 'GCP dependencies are not installed')
class TestWriteStringsToPubSub(unittest.TestCase):
def test_expand(self):
p = TestPipeline()
- pdone = (p
- | ReadStringsFromPubSub('projects/fakeprj/topics/baz')
- | WriteStringsToPubSub('projects/fakeprj/topics/a_topic'))
+ pdone = p | ReadStringsFromPubSub('baz') | WriteStringsToPubSub('a_topic')
# Ensure that the properties passed through correctly
- self.assertEqual('a_topic', pdone.producer.transform.dofn.topic_name)
+ sink = pdone.producer.transform.sink
+ self.assertEqual('a_topic', sink.topic)
+ # Ensure that the type on the intermediate payload transformer output
+ # PCollection is bytes
+ write_pcoll = pdone.producer.inputs[0]
+ self.assertEqual(bytes, write_pcoll.element_type)
-@unittest.skipIf(pubsub is None, 'GCP dependencies are not installed')
-class TestPubSubSource(unittest.TestCase):
- def test_display_data_topic(self):
- source = _PubSubPayloadSource(
- 'projects/fakeprj/topics/a_topic',
- None,
- 'a_label')
- dd = DisplayData.create_from(source)
- expected_items = [
- DisplayDataItemMatcher(
- 'topic', 'projects/fakeprj/topics/a_topic'),
- DisplayDataItemMatcher('id_label', 'a_label')]
-
- hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
- def test_display_data_subscription(self):
- source = _PubSubPayloadSource(
- None,
- 'projects/fakeprj/subscriptions/a_subscription',
- 'a_label')
+class TestPubSubSource(unittest.TestCase):
+ def test_display_data(self):
+ source = _PubSubPayloadSource('a_topic', 'a_subscription', 'a_label')
dd = DisplayData.create_from(source)
expected_items = [
- DisplayDataItemMatcher(
- 'subscription', 'projects/fakeprj/subscriptions/a_subscription'),
+ DisplayDataItemMatcher('topic', 'a_topic'),
+ DisplayDataItemMatcher('subscription', 'a_subscription'),
DisplayDataItemMatcher('id_label', 'a_label')]
hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
def test_display_data_no_subscription(self):
- source = _PubSubPayloadSource('projects/fakeprj/topics/a_topic')
+ source = _PubSubPayloadSource('a_topic')
dd = DisplayData.create_from(source)
expected_items = [
- DisplayDataItemMatcher('topic', 'projects/fakeprj/topics/a_topic')]
+ DisplayDataItemMatcher('topic', 'a_topic')]
hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
-@unittest.skipIf(pubsub is None, 'GCP dependencies are not installed')
class TestPubSubSink(unittest.TestCase):
def test_display_data(self):
- sink = _PubSubPayloadSink('projects/fakeprj/topics/a_topic')
+ sink = _PubSubPayloadSink('a_topic')
dd = DisplayData.create_from(sink)
expected_items = [
- DisplayDataItemMatcher('topic', 'projects/fakeprj/topics/a_topic')]
+ DisplayDataItemMatcher('topic', 'a_topic')]
hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
+class TestEncodeDecodeUtf8String(unittest.TestCase):
+ def test_encode(self):
+ self.assertEqual(b'test_data', _encodeUtf8String('test_data'))
+
+ def test_decode(self):
+ self.assertEqual('test_data', _decodeUtf8String(b'test_data'))
+
+
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
unittest.main()
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py b/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py
index d6f0e97..844cbc5 100644
--- a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py
+++ b/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py
@@ -92,9 +92,9 @@ class BigqueryMatcher(BaseMatcher):
page_token = None
results = []
while True:
- for row in query.fetch_data(page_token=page_token):
- results.append(row)
- if results:
+ rows, _, page_token = query.fetch_data(page_token=page_token)
+ results.extend(rows)
+ if not page_token:
break
return results
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher_test.py b/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher_test.py
index 5b72285..f12293e 100644
--- a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher_test.py
+++ b/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher_test.py
@@ -53,7 +53,7 @@ class BigqueryMatcherTest(unittest.TestCase):
matcher = bq_verifier.BigqueryMatcher(
'mock_project',
'mock_query',
- '59f9d6bdee30d67ea73b8aded121c3a0280f9cd8')
+ 'da39a3ee5e6b4b0d3255bfef95601890afd80709')
hc_assert_that(self._mock_result, matcher)
@patch.object(bigquery, 'Client')
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/io/range_trackers.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/io/range_trackers.py b/sdks/python/apache_beam/io/range_trackers.py
index bef77d4..9cb36e7 100644
--- a/sdks/python/apache_beam/io/range_trackers.py
+++ b/sdks/python/apache_beam/io/range_trackers.py
@@ -193,6 +193,136 @@ class OffsetRangeTracker(iobase.RangeTracker):
self._split_points_unclaimed_callback = callback
+class GroupedShuffleRangeTracker(iobase.RangeTracker):
+ """For internal use only; no backwards-compatibility guarantees.
+
+ A 'RangeTracker' for positions used by'GroupedShuffleReader'.
+
+ These positions roughly correspond to hashes of keys. In case of hash
+ collisions, multiple groups can have the same position. In that case, the
+ first group at a particular position is considered a split point (because
+ it is the first to be returned when reading a position range starting at this
+ position), others are not.
+ """
+
+ def __init__(self, decoded_start_pos, decoded_stop_pos):
+ super(GroupedShuffleRangeTracker, self).__init__()
+ self._decoded_start_pos = decoded_start_pos
+ self._decoded_stop_pos = decoded_stop_pos
+ self._decoded_last_group_start = None
+ self._last_group_was_at_a_split_point = False
+ self._split_points_seen = 0
+ self._lock = threading.Lock()
+
+ def start_position(self):
+ return self._decoded_start_pos
+
+ def stop_position(self):
+ return self._decoded_stop_pos
+
+ def last_group_start(self):
+ return self._decoded_last_group_start
+
+ def _validate_decoded_group_start(self, decoded_group_start, split_point):
+ if self.start_position() and decoded_group_start < self.start_position():
+ raise ValueError('Trying to return record at %r which is before the'
+ ' starting position at %r' %
+ (decoded_group_start, self.start_position()))
+
+ if (self.last_group_start() and
+ decoded_group_start < self.last_group_start()):
+ raise ValueError('Trying to return group at %r which is before the'
+ ' last-returned group at %r' %
+ (decoded_group_start, self.last_group_start()))
+ if (split_point and self.last_group_start() and
+ self.last_group_start() == decoded_group_start):
+ raise ValueError('Trying to return a group at a split point with '
+ 'same position as the previous group: both at %r, '
+ 'last group was %sat a split point.' %
+ (decoded_group_start,
+ ('' if self._last_group_was_at_a_split_point
+ else 'not ')))
+ if not split_point:
+ if self.last_group_start() is None:
+ raise ValueError('The first group [at %r] must be at a split point' %
+ decoded_group_start)
+ if self.last_group_start() != decoded_group_start:
+ # This case is not a violation of general RangeTracker semantics, but it
+ # is contrary to how GroupingShuffleReader in particular works. Hitting
+ # it would mean it's behaving unexpectedly.
+ raise ValueError('Trying to return a group not at a split point, but '
+ 'with a different position than the previous group: '
+ 'last group was %r at %r, current at a %s split'
+ ' point.' %
+ (self.last_group_start()
+ , decoded_group_start
+ , ('' if self._last_group_was_at_a_split_point
+ else 'non-')))
+
+ def try_claim(self, decoded_group_start):
+ with self._lock:
+ self._validate_decoded_group_start(decoded_group_start, True)
+ if (self.stop_position()
+ and decoded_group_start >= self.stop_position()):
+ return False
+
+ self._decoded_last_group_start = decoded_group_start
+ self._last_group_was_at_a_split_point = True
+ self._split_points_seen += 1
+ return True
+
+ def set_current_position(self, decoded_group_start):
+ with self._lock:
+ self._validate_decoded_group_start(decoded_group_start, False)
+ self._decoded_last_group_start = decoded_group_start
+ self._last_group_was_at_a_split_point = False
+
+ def try_split(self, decoded_split_position):
+ with self._lock:
+ if self.last_group_start() is None:
+ logging.info('Refusing to split %r at %r: unstarted'
+ , self, decoded_split_position)
+ return
+
+ if decoded_split_position <= self.last_group_start():
+ logging.info('Refusing to split %r at %r: already past proposed split '
+ 'position'
+ , self, decoded_split_position)
+ return
+
+ if ((self.stop_position()
+ and decoded_split_position >= self.stop_position())
+ or (self.start_position()
+ and decoded_split_position <= self.start_position())):
+ logging.error('Refusing to split %r at %r: proposed split position out '
+ 'of range', self, decoded_split_position)
+ return
+
+ logging.debug('Agreeing to split %r at %r'
+ , self, decoded_split_position)
+ self._decoded_stop_pos = decoded_split_position
+
+ # Since GroupedShuffleRangeTracker cannot determine relative sizes of the
+ # two splits, returning 0.5 as the fraction below so that the framework
+ # assumes the splits to be of the same size.
+ return self._decoded_stop_pos, 0.5
+
+ def fraction_consumed(self):
+ # GroupingShuffle sources have special support on the service and the
+ # service will estimate progress from positions for us.
+ raise RuntimeError('GroupedShuffleRangeTracker does not measure fraction'
+ ' consumed due to positions being opaque strings'
+ ' that are interpreted by the service')
+
+ def split_points(self):
+ with self._lock:
+ splits_points_consumed = (
+ 0 if self._split_points_seen <= 1 else (self._split_points_seen - 1))
+
+ return (splits_points_consumed,
+ iobase.RangeTracker.SPLIT_POINTS_UNKNOWN)
+
+
class OrderedPositionRangeTracker(iobase.RangeTracker):
"""
An abstract base class for range trackers whose positions are comparable.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/io/range_trackers_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/io/range_trackers_test.py b/sdks/python/apache_beam/io/range_trackers_test.py
index 3e92663..edb6386 100644
--- a/sdks/python/apache_beam/io/range_trackers_test.py
+++ b/sdks/python/apache_beam/io/range_trackers_test.py
@@ -17,11 +17,14 @@
"""Unit tests for the range_trackers module."""
+import array
import copy
import logging
import math
import unittest
+
+from apache_beam.io import iobase
from apache_beam.io import range_trackers
@@ -186,6 +189,189 @@ class OffsetRangeTrackerTest(unittest.TestCase):
(3, 41))
+class GroupedShuffleRangeTrackerTest(unittest.TestCase):
+
+ def bytes_to_position(self, bytes_array):
+ return array.array('B', bytes_array).tostring()
+
+ def test_try_return_record_in_infinite_range(self):
+ tracker = range_trackers.GroupedShuffleRangeTracker('', '')
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([1, 2, 3])))
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([1, 2, 5])))
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([3, 6, 8, 10])))
+
+ def test_try_return_record_finite_range(self):
+ tracker = range_trackers.GroupedShuffleRangeTracker(
+ self.bytes_to_position([1, 0, 0]), self.bytes_to_position([5, 0, 0]))
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([1, 2, 3])))
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([1, 2, 5])))
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([3, 6, 8, 10])))
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([4, 255, 255, 255])))
+ # Should fail for positions that are lexicographically equal to or larger
+ # than the defined stop position.
+ self.assertFalse(copy.copy(tracker).try_claim(
+ self.bytes_to_position([5, 0, 0])))
+ self.assertFalse(copy.copy(tracker).try_claim(
+ self.bytes_to_position([5, 0, 1])))
+ self.assertFalse(copy.copy(tracker).try_claim(
+ self.bytes_to_position([6, 0, 0])))
+
+ def test_try_return_record_with_non_split_point(self):
+ tracker = range_trackers.GroupedShuffleRangeTracker(
+ self.bytes_to_position([1, 0, 0]), self.bytes_to_position([5, 0, 0]))
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([1, 2, 3])))
+ tracker.set_current_position(self.bytes_to_position([1, 2, 3]))
+ tracker.set_current_position(self.bytes_to_position([1, 2, 3]))
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([1, 2, 5])))
+ tracker.set_current_position(self.bytes_to_position([1, 2, 5]))
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([3, 6, 8, 10])))
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([4, 255, 255, 255])))
+
+ def test_first_record_non_split_point(self):
+ tracker = range_trackers.GroupedShuffleRangeTracker(
+ self.bytes_to_position([3, 0, 0]), self.bytes_to_position([5, 0, 0]))
+ with self.assertRaises(ValueError):
+ tracker.set_current_position(self.bytes_to_position([3, 4, 5]))
+
+ def test_non_split_point_record_with_different_position(self):
+ tracker = range_trackers.GroupedShuffleRangeTracker(
+ self.bytes_to_position([3, 0, 0]), self.bytes_to_position([5, 0, 0]))
+ self.assertTrue(tracker.try_claim(self.bytes_to_position([3, 4, 5])))
+ with self.assertRaises(ValueError):
+ tracker.set_current_position(self.bytes_to_position([3, 4, 6]))
+
+ def test_try_return_record_before_start(self):
+ tracker = range_trackers.GroupedShuffleRangeTracker(
+ self.bytes_to_position([3, 0, 0]), self.bytes_to_position([5, 0, 0]))
+ with self.assertRaises(ValueError):
+ tracker.try_claim(self.bytes_to_position([1, 2, 3]))
+
+ def test_try_return_non_monotonic(self):
+ tracker = range_trackers.GroupedShuffleRangeTracker(
+ self.bytes_to_position([3, 0, 0]), self.bytes_to_position([5, 0, 0]))
+ self.assertTrue(tracker.try_claim(self.bytes_to_position([3, 4, 5])))
+ self.assertTrue(tracker.try_claim(self.bytes_to_position([3, 4, 6])))
+ with self.assertRaises(ValueError):
+ tracker.try_claim(self.bytes_to_position([3, 2, 1]))
+
+ def test_try_return_identical_positions(self):
+ tracker = range_trackers.GroupedShuffleRangeTracker(
+ self.bytes_to_position([3, 0, 0]), self.bytes_to_position([5, 0, 0]))
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([3, 4, 5])))
+ with self.assertRaises(ValueError):
+ tracker.try_claim(self.bytes_to_position([3, 4, 5]))
+
+ def test_try_split_at_position_infinite_range(self):
+ tracker = range_trackers.GroupedShuffleRangeTracker('', '')
+ # Should fail before first record is returned.
+ self.assertFalse(tracker.try_split(
+ self.bytes_to_position([3, 4, 5, 6])))
+
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([1, 2, 3])))
+
+ # Should now succeed.
+ self.assertIsNotNone(tracker.try_split(
+ self.bytes_to_position([3, 4, 5, 6])))
+ # Should not split at same or larger position.
+ self.assertIsNone(tracker.try_split(
+ self.bytes_to_position([3, 4, 5, 6])))
+ self.assertIsNone(tracker.try_split(
+ self.bytes_to_position([3, 4, 5, 6, 7])))
+ self.assertIsNone(tracker.try_split(
+ self.bytes_to_position([4, 5, 6, 7])))
+
+ # Should split at smaller position.
+ self.assertIsNotNone(tracker.try_split(
+ self.bytes_to_position([3, 2, 1])))
+
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([2, 3, 4])))
+
+ # Should not split at a position we're already past.
+ self.assertIsNone(tracker.try_split(
+ self.bytes_to_position([2, 3, 4])))
+ self.assertIsNone(tracker.try_split(
+ self.bytes_to_position([2, 3, 3])))
+
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([3, 2, 0])))
+ self.assertFalse(tracker.try_claim(
+ self.bytes_to_position([3, 2, 1])))
+
+ def test_try_test_split_at_position_finite_range(self):
+ tracker = range_trackers.GroupedShuffleRangeTracker(
+ self.bytes_to_position([0, 0, 0]),
+ self.bytes_to_position([10, 20, 30]))
+ # Should fail before first record is returned.
+ self.assertFalse(tracker.try_split(
+ self.bytes_to_position([0, 0, 0])))
+ self.assertFalse(tracker.try_split(
+ self.bytes_to_position([3, 4, 5, 6])))
+
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([1, 2, 3])))
+
+ # Should now succeed.
+ self.assertTrue(tracker.try_split(
+ self.bytes_to_position([3, 4, 5, 6])))
+ # Should not split at same or larger position.
+ self.assertFalse(tracker.try_split(
+ self.bytes_to_position([3, 4, 5, 6])))
+ self.assertFalse(tracker.try_split(
+ self.bytes_to_position([3, 4, 5, 6, 7])))
+ self.assertFalse(tracker.try_split(
+ self.bytes_to_position([4, 5, 6, 7])))
+
+ # Should split at smaller position.
+ self.assertTrue(tracker.try_split(
+ self.bytes_to_position([3, 2, 1])))
+ # But not at a position at or before last returned record.
+ self.assertFalse(tracker.try_split(
+ self.bytes_to_position([1, 2, 3])))
+
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([2, 3, 4])))
+ self.assertTrue(tracker.try_claim(
+ self.bytes_to_position([3, 2, 0])))
+ self.assertFalse(tracker.try_claim(
+ self.bytes_to_position([3, 2, 1])))
+
+ def test_split_points(self):
+ tracker = range_trackers.GroupedShuffleRangeTracker(
+ self.bytes_to_position([1, 0, 0]),
+ self.bytes_to_position([5, 0, 0]))
+ self.assertEqual(tracker.split_points(),
+ (0, iobase.RangeTracker.SPLIT_POINTS_UNKNOWN))
+ self.assertTrue(tracker.try_claim(self.bytes_to_position([1, 2, 3])))
+ self.assertEqual(tracker.split_points(),
+ (0, iobase.RangeTracker.SPLIT_POINTS_UNKNOWN))
+ self.assertTrue(tracker.try_claim(self.bytes_to_position([1, 2, 5])))
+ self.assertEqual(tracker.split_points(),
+ (1, iobase.RangeTracker.SPLIT_POINTS_UNKNOWN))
+ self.assertTrue(tracker.try_claim(self.bytes_to_position([3, 6, 8])))
+ self.assertEqual(tracker.split_points(),
+ (2, iobase.RangeTracker.SPLIT_POINTS_UNKNOWN))
+ self.assertTrue(tracker.try_claim(self.bytes_to_position([4, 255, 255])))
+ self.assertEqual(tracker.split_points(),
+ (3, iobase.RangeTracker.SPLIT_POINTS_UNKNOWN))
+ self.assertFalse(tracker.try_claim(self.bytes_to_position([5, 1, 0])))
+ self.assertEqual(tracker.split_points(),
+ (3, iobase.RangeTracker.SPLIT_POINTS_UNKNOWN))
+
+
class OrderedPositionRangeTrackerTest(unittest.TestCase):
class DoubleRangeTracker(range_trackers.OrderedPositionRangeTracker):
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/options/pipeline_options.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py
index ea996a3..daef3a7 100644
--- a/sdks/python/apache_beam/options/pipeline_options.py
+++ b/sdks/python/apache_beam/options/pipeline_options.py
@@ -18,6 +18,7 @@
"""Pipeline options obtained from command line parsing."""
import argparse
+import warnings
from apache_beam.transforms.display import HasDisplayData
from apache_beam.options.value_provider import StaticValueProvider
@@ -278,6 +279,14 @@ class StandardOptions(PipelineOptions):
action='store_true',
help='Whether to enable streaming mode.')
+ # TODO(BEAM-1265): Remove this warning, once at least one runner supports
+ # streaming pipelines.
+ def validate(self, validator):
+ errors = []
+ if self.view_as(StandardOptions).streaming:
+ warnings.warn('Streaming pipelines are not supported.')
+ return errors
+
class TypeOptions(PipelineOptions):
@@ -465,14 +474,7 @@ class WorkerOptions(PipelineOptions):
parser.add_argument(
'--use_public_ips',
default=None,
- action='store_true',
- help='Whether to assign public IP addresses to the worker VMs.')
- parser.add_argument(
- '--no_use_public_ips',
- dest='use_public_ips',
- default=None,
- action='store_false',
- help='Whether to assign only private IP addresses to the worker VMs.')
+ help='Whether to assign public IP addresses to the worker machines.')
def validate(self, validator):
errors = []
@@ -552,18 +554,6 @@ class SetupOptions(PipelineOptions):
'worker will install the resulting package before running any custom '
'code.'))
parser.add_argument(
- '--beam_plugin', '--beam_plugin',
- dest='beam_plugins',
- action='append',
- default=None,
- help=
- ('Bootstrap the python process before executing any code by importing '
- 'all the plugins used in the pipeline. Please pass a comma separated'
- 'list of import paths to be included. This is currently an '
- 'experimental flag and provides no stability. Multiple '
- '--beam_plugin options can be specified if more than one plugin '
- 'is needed.'))
- parser.add_argument(
'--save_main_session',
default=False,
action='store_true',
@@ -609,11 +599,6 @@ class TestOptions(PipelineOptions):
help=('Verify state/output of e2e test pipeline. This is pickled '
'version of the matcher which should extends '
'hamcrest.core.base_matcher.BaseMatcher.'))
- parser.add_argument(
- '--dry_run',
- default=False,
- help=('Used in unit testing runners without submitting the '
- 'actual job.'))
def validate(self, validator):
errors = []
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/options/pipeline_options_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/options/pipeline_options_test.py b/sdks/python/apache_beam/options/pipeline_options_test.py
index f4dd4d9..1a644b4 100644
--- a/sdks/python/apache_beam/options/pipeline_options_test.py
+++ b/sdks/python/apache_beam/options/pipeline_options_test.py
@@ -192,52 +192,47 @@ class PipelineOptionsTest(unittest.TestCase):
options = PipelineOptions(['--redefined_flag'])
self.assertTrue(options.get_all_options()['redefined_flag'])
- # TODO(BEAM-1319): Require unique names only within a test.
- # For now, <file name acronym>_vp_arg<number> will be the convention
- # to name value-provider arguments in tests, as opposed to
- # <file name acronym>_non_vp_arg<number> for non-value-provider arguments.
- # The number will grow per file as tests are added.
def test_value_provider_options(self):
class UserOptions(PipelineOptions):
@classmethod
def _add_argparse_args(cls, parser):
parser.add_value_provider_argument(
- '--pot_vp_arg1',
+ '--vp_arg',
help='This flag is a value provider')
parser.add_value_provider_argument(
- '--pot_vp_arg2',
+ '--vp_arg2',
default=1,
type=int)
parser.add_argument(
- '--pot_non_vp_arg1',
+ '--non_vp_arg',
default=1,
type=int
)
# Provide values: if not provided, the option becomes of the type runtime vp
- options = UserOptions(['--pot_vp_arg1', 'hello'])
- self.assertIsInstance(options.pot_vp_arg1, StaticValueProvider)
- self.assertIsInstance(options.pot_vp_arg2, RuntimeValueProvider)
- self.assertIsInstance(options.pot_non_vp_arg1, int)
+ options = UserOptions(['--vp_arg', 'hello'])
+ self.assertIsInstance(options.vp_arg, StaticValueProvider)
+ self.assertIsInstance(options.vp_arg2, RuntimeValueProvider)
+ self.assertIsInstance(options.non_vp_arg, int)
# Values can be overwritten
- options = UserOptions(pot_vp_arg1=5,
- pot_vp_arg2=StaticValueProvider(value_type=str,
- value='bye'),
- pot_non_vp_arg1=RuntimeValueProvider(
+ options = UserOptions(vp_arg=5,
+ vp_arg2=StaticValueProvider(value_type=str,
+ value='bye'),
+ non_vp_arg=RuntimeValueProvider(
option_name='foo',
value_type=int,
default_value=10))
- self.assertEqual(options.pot_vp_arg1, 5)
- self.assertTrue(options.pot_vp_arg2.is_accessible(),
- '%s is not accessible' % options.pot_vp_arg2)
- self.assertEqual(options.pot_vp_arg2.get(), 'bye')
- self.assertFalse(options.pot_non_vp_arg1.is_accessible())
+ self.assertEqual(options.vp_arg, 5)
+ self.assertTrue(options.vp_arg2.is_accessible(),
+ '%s is not accessible' % options.vp_arg2)
+ self.assertEqual(options.vp_arg2.get(), 'bye')
+ self.assertFalse(options.non_vp_arg.is_accessible())
with self.assertRaises(RuntimeError):
- options.pot_non_vp_arg1.get()
+ options.non_vp_arg.get()
if __name__ == '__main__':
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/options/value_provider_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/options/value_provider_test.py b/sdks/python/apache_beam/options/value_provider_test.py
index 17e9590..3a45e8b 100644
--- a/sdks/python/apache_beam/options/value_provider_test.py
+++ b/sdks/python/apache_beam/options/value_provider_test.py
@@ -24,77 +24,72 @@ from apache_beam.options.value_provider import RuntimeValueProvider
from apache_beam.options.value_provider import StaticValueProvider
-# TODO(BEAM-1319): Require unique names only within a test.
-# For now, <file name acronym>_vp_arg<number> will be the convention
-# to name value-provider arguments in tests, as opposed to
-# <file name acronym>_non_vp_arg<number> for non-value-provider arguments.
-# The number will grow per file as tests are added.
class ValueProviderTests(unittest.TestCase):
def test_static_value_provider_keyword_argument(self):
class UserDefinedOptions(PipelineOptions):
@classmethod
def _add_argparse_args(cls, parser):
parser.add_value_provider_argument(
- '--vpt_vp_arg1',
+ '--vp_arg',
help='This keyword argument is a value provider',
default='some value')
- options = UserDefinedOptions(['--vpt_vp_arg1', 'abc'])
- self.assertTrue(isinstance(options.vpt_vp_arg1, StaticValueProvider))
- self.assertTrue(options.vpt_vp_arg1.is_accessible())
- self.assertEqual(options.vpt_vp_arg1.get(), 'abc')
+ options = UserDefinedOptions(['--vp_arg', 'abc'])
+ self.assertTrue(isinstance(options.vp_arg, StaticValueProvider))
+ self.assertTrue(options.vp_arg.is_accessible())
+ self.assertEqual(options.vp_arg.get(), 'abc')
def test_runtime_value_provider_keyword_argument(self):
class UserDefinedOptions(PipelineOptions):
@classmethod
def _add_argparse_args(cls, parser):
parser.add_value_provider_argument(
- '--vpt_vp_arg2',
+ '--vp_arg',
help='This keyword argument is a value provider')
options = UserDefinedOptions()
- self.assertTrue(isinstance(options.vpt_vp_arg2, RuntimeValueProvider))
- self.assertFalse(options.vpt_vp_arg2.is_accessible())
+ self.assertTrue(isinstance(options.vp_arg, RuntimeValueProvider))
+ self.assertFalse(options.vp_arg.is_accessible())
with self.assertRaises(RuntimeError):
- options.vpt_vp_arg2.get()
+ options.vp_arg.get()
def test_static_value_provider_positional_argument(self):
class UserDefinedOptions(PipelineOptions):
@classmethod
def _add_argparse_args(cls, parser):
parser.add_value_provider_argument(
- 'vpt_vp_arg3',
+ 'vp_pos_arg',
help='This positional argument is a value provider',
default='some value')
options = UserDefinedOptions(['abc'])
- self.assertTrue(isinstance(options.vpt_vp_arg3, StaticValueProvider))
- self.assertTrue(options.vpt_vp_arg3.is_accessible())
- self.assertEqual(options.vpt_vp_arg3.get(), 'abc')
+ self.assertTrue(isinstance(options.vp_pos_arg, StaticValueProvider))
+ self.assertTrue(options.vp_pos_arg.is_accessible())
+ self.assertEqual(options.vp_pos_arg.get(), 'abc')
def test_runtime_value_provider_positional_argument(self):
class UserDefinedOptions(PipelineOptions):
@classmethod
def _add_argparse_args(cls, parser):
parser.add_value_provider_argument(
- 'vpt_vp_arg4',
+ 'vp_pos_arg',
help='This positional argument is a value provider')
options = UserDefinedOptions([])
- self.assertTrue(isinstance(options.vpt_vp_arg4, RuntimeValueProvider))
- self.assertFalse(options.vpt_vp_arg4.is_accessible())
+ self.assertTrue(isinstance(options.vp_pos_arg, RuntimeValueProvider))
+ self.assertFalse(options.vp_pos_arg.is_accessible())
with self.assertRaises(RuntimeError):
- options.vpt_vp_arg4.get()
+ options.vp_pos_arg.get()
def test_static_value_provider_type_cast(self):
class UserDefinedOptions(PipelineOptions):
@classmethod
def _add_argparse_args(cls, parser):
parser.add_value_provider_argument(
- '--vpt_vp_arg5',
+ '--vp_arg',
type=int,
help='This flag is a value provider')
- options = UserDefinedOptions(['--vpt_vp_arg5', '123'])
- self.assertTrue(isinstance(options.vpt_vp_arg5, StaticValueProvider))
- self.assertTrue(options.vpt_vp_arg5.is_accessible())
- self.assertEqual(options.vpt_vp_arg5.get(), 123)
+ options = UserDefinedOptions(['--vp_arg', '123'])
+ self.assertTrue(isinstance(options.vp_arg, StaticValueProvider))
+ self.assertTrue(options.vp_arg.is_accessible())
+ self.assertEqual(options.vp_arg.get(), 123)
def test_set_runtime_option(self):
# define ValueProvider ptions, with and without default values
@@ -102,25 +97,25 @@ class ValueProviderTests(unittest.TestCase):
@classmethod
def _add_argparse_args(cls, parser):
parser.add_value_provider_argument(
- '--vpt_vp_arg6',
+ '--vp_arg',
help='This keyword argument is a value provider') # set at runtime
parser.add_value_provider_argument( # not set, had default int
- '-v', '--vpt_vp_arg7', # with short form
+ '-v', '--vp_arg2', # with short form
default=123,
type=int)
parser.add_value_provider_argument( # not set, had default str
- '--vpt_vp-arg8', # with dash in name
+ '--vp-arg3', # with dash in name
default='123',
type=str)
parser.add_value_provider_argument( # not set and no default
- '--vpt_vp_arg9',
+ '--vp_arg4',
type=float)
parser.add_value_provider_argument( # positional argument set
- 'vpt_vp_arg10', # default & runtime ignored
+ 'vp_pos_arg', # default & runtime ignored
help='This positional argument is a value provider',
type=float,
default=5.4)
@@ -128,23 +123,23 @@ class ValueProviderTests(unittest.TestCase):
# provide values at graph-construction time
# (options not provided here become of the type RuntimeValueProvider)
options = UserDefinedOptions1(['1.2'])
- self.assertFalse(options.vpt_vp_arg6.is_accessible())
- self.assertFalse(options.vpt_vp_arg7.is_accessible())
- self.assertFalse(options.vpt_vp_arg8.is_accessible())
- self.assertFalse(options.vpt_vp_arg9.is_accessible())
- self.assertTrue(options.vpt_vp_arg10.is_accessible())
+ self.assertFalse(options.vp_arg.is_accessible())
+ self.assertFalse(options.vp_arg2.is_accessible())
+ self.assertFalse(options.vp_arg3.is_accessible())
+ self.assertFalse(options.vp_arg4.is_accessible())
+ self.assertTrue(options.vp_pos_arg.is_accessible())
# provide values at job-execution time
# (options not provided here will use their default, if they have one)
- RuntimeValueProvider.set_runtime_options({'vpt_vp_arg6': 'abc',
- 'vpt_vp_arg10':'3.2'})
- self.assertTrue(options.vpt_vp_arg6.is_accessible())
- self.assertEqual(options.vpt_vp_arg6.get(), 'abc')
- self.assertTrue(options.vpt_vp_arg7.is_accessible())
- self.assertEqual(options.vpt_vp_arg7.get(), 123)
- self.assertTrue(options.vpt_vp_arg8.is_accessible())
- self.assertEqual(options.vpt_vp_arg8.get(), '123')
- self.assertTrue(options.vpt_vp_arg9.is_accessible())
- self.assertIsNone(options.vpt_vp_arg9.get())
- self.assertTrue(options.vpt_vp_arg10.is_accessible())
- self.assertEqual(options.vpt_vp_arg10.get(), 1.2)
+ RuntimeValueProvider.set_runtime_options({'vp_arg': 'abc',
+ 'vp_pos_arg':'3.2'})
+ self.assertTrue(options.vp_arg.is_accessible())
+ self.assertEqual(options.vp_arg.get(), 'abc')
+ self.assertTrue(options.vp_arg2.is_accessible())
+ self.assertEqual(options.vp_arg2.get(), 123)
+ self.assertTrue(options.vp_arg3.is_accessible())
+ self.assertEqual(options.vp_arg3.get(), '123')
+ self.assertTrue(options.vp_arg4.is_accessible())
+ self.assertIsNone(options.vp_arg4.get())
+ self.assertTrue(options.vp_pos_arg.is_accessible())
+ self.assertEqual(options.vp_pos_arg.get(), 1.2)
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/pipeline.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/pipeline.py b/sdks/python/apache_beam/pipeline.py
index fe36d85..9093abf 100644
--- a/sdks/python/apache_beam/pipeline.py
+++ b/sdks/python/apache_beam/pipeline.py
@@ -45,7 +45,6 @@ Typical usage:
from __future__ import absolute_import
-import abc
import collections
import logging
import os
@@ -54,7 +53,6 @@ import tempfile
from apache_beam import pvalue
from apache_beam.internal import pickler
-from apache_beam.pvalue import PCollection
from apache_beam.runners import create_runner
from apache_beam.runners import PipelineRunner
from apache_beam.transforms import ptransform
@@ -159,157 +157,6 @@ class Pipeline(object):
"""Returns the root transform of the transform stack."""
return self.transforms_stack[0]
- def _remove_labels_recursively(self, applied_transform):
- for part in applied_transform.parts:
- if part.full_label in self.applied_labels:
- self.applied_labels.remove(part.full_label)
- if part.parts:
- for part2 in part.parts:
- self._remove_labels_recursively(part2)
-
- def _replace(self, override):
-
- assert isinstance(override, PTransformOverride)
- matcher = override.get_matcher()
-
- output_map = {}
- output_replacements = {}
- input_replacements = {}
-
- class TransformUpdater(PipelineVisitor): # pylint: disable=used-before-assignment
- """"A visitor that replaces the matching PTransforms."""
-
- def __init__(self, pipeline):
- self.pipeline = pipeline
-
- def _replace_if_needed(self, transform_node):
- if matcher(transform_node):
- replacement_transform = override.get_replacement_transform(
- transform_node.transform)
- inputs = transform_node.inputs
- # TODO: Support replacing PTransforms with multiple inputs.
- if len(inputs) > 1:
- raise NotImplementedError(
- 'PTransform overriding is only supported for PTransforms that '
- 'have a single input. Tried to replace input of '
- 'AppliedPTransform %r that has %d inputs',
- transform_node, len(inputs))
- transform_node.transform = replacement_transform
- self.pipeline.transforms_stack.append(transform_node)
-
- # Keeping the same label for the replaced node but recursively
- # removing labels of child transforms since they will be replaced
- # during the expand below.
- self.pipeline._remove_labels_recursively(transform_node)
-
- new_output = replacement_transform.expand(inputs[0])
- if new_output.producer is None:
- # When current transform is a primitive, we set the producer here.
- new_output.producer = transform_node
-
- # We only support replacing transforms with a single output with
- # another transform that produces a single output.
- # TODO: Support replacing PTransforms with multiple outputs.
- if (len(transform_node.outputs) > 1 or
- not isinstance(transform_node.outputs[None], PCollection) or
- not isinstance(new_output, PCollection)):
- raise NotImplementedError(
- 'PTransform overriding is only supported for PTransforms that '
- 'have a single output. Tried to replace output of '
- 'AppliedPTransform %r with %r.'
- , transform_node, new_output)
-
- # Recording updated outputs. This cannot be done in the same visitor
- # since if we dynamically update output type here, we'll run into
- # errors when visiting child nodes.
- output_map[transform_node.outputs[None]] = new_output
-
- self.pipeline.transforms_stack.pop()
-
- def enter_composite_transform(self, transform_node):
- self._replace_if_needed(transform_node)
-
- def visit_transform(self, transform_node):
- self._replace_if_needed(transform_node)
-
- self.visit(TransformUpdater(self))
-
- # Adjusting inputs and outputs
- class InputOutputUpdater(PipelineVisitor): # pylint: disable=used-before-assignment
- """"A visitor that records input and output values to be replaced.
-
- Input and output values that should be updated are recorded in maps
- input_replacements and output_replacements respectively.
-
- We cannot update input and output values while visiting since that results
- in validation errors.
- """
-
- def __init__(self, pipeline):
- self.pipeline = pipeline
-
- def enter_composite_transform(self, transform_node):
- self.visit_transform(transform_node)
-
- def visit_transform(self, transform_node):
- if (None in transform_node.outputs and
- transform_node.outputs[None] in output_map):
- output_replacements[transform_node] = (
- output_map[transform_node.outputs[None]])
-
- replace_input = False
- for input in transform_node.inputs:
- if input in output_map:
- replace_input = True
- break
-
- if replace_input:
- new_input = [
- input if not input in output_map else output_map[input]
- for input in transform_node.inputs]
- input_replacements[transform_node] = new_input
-
- self.visit(InputOutputUpdater(self))
-
- for transform in output_replacements:
- transform.replace_output(output_replacements[transform])
-
- for transform in input_replacements:
- transform.inputs = input_replacements[transform]
-
- def _check_replacement(self, override):
- matcher = override.get_matcher()
-
- class ReplacementValidator(PipelineVisitor):
- def visit_transform(self, transform_node):
- if matcher(transform_node):
- raise RuntimeError('Transform node %r was not replaced as expected.',
- transform_node)
-
- self.visit(ReplacementValidator())
-
- def replace_all(self, replacements):
- """ Dynamically replaces PTransforms in the currently populated hierarchy.
-
- Currently this only works for replacements where input and output types
- are exactly the same.
- TODO: Update this to also work for transform overrides where input and
- output types are different.
-
- Args:
- replacements a list of PTransformOverride objects.
- """
- for override in replacements:
- assert isinstance(override, PTransformOverride)
- self._replace(override)
-
- # Checking if the PTransforms have been successfully replaced. This will
- # result in a failure if a PTransform that was replaced in a given override
- # gets re-added in a subsequent override. This is not allowed and ordering
- # of PTransformOverride objects in 'replacements' is important.
- for override in replacements:
- self._check_replacement(override)
-
def run(self, test_runner_api=True):
"""Runs the pipeline. Returns whatever our runner returns after running."""
@@ -466,20 +313,10 @@ class Pipeline(object):
self.transforms_stack.pop()
return pvalueish_result
- def __reduce__(self):
- # Some transforms contain a reference to their enclosing pipeline,
- # which in turn reference all other transforms (resulting in quadratic
- # time/space to pickle each transform individually). As we don't
- # require pickled pipelines to be executable, break the chain here.
- return str, ('Pickled pipeline stub.',)
-
def _verify_runner_api_compatible(self):
class Visitor(PipelineVisitor): # pylint: disable=used-before-assignment
ok = True # Really a nonlocal.
- def enter_composite_transform(self, transform_node):
- self.visit_transform(transform_node)
-
def visit_transform(self, transform_node):
if transform_node.side_inputs:
# No side inputs (yet).
@@ -502,7 +339,7 @@ class Pipeline(object):
def to_runner_api(self):
"""For internal use only; no backwards-compatibility guarantees."""
from apache_beam.runners import pipeline_context
- from apache_beam.portability.api import beam_runner_api_pb2
+ from apache_beam.runners.api import beam_runner_api_pb2
context = pipeline_context.PipelineContext()
# Mutates context; placing inline would force dependence on
# argument evaluation order.
@@ -525,18 +362,7 @@ class Pipeline(object):
p.applied_labels = set([
t.unique_name for t in proto.components.transforms.values()])
for id in proto.components.pcollections:
- pcollection = context.pcollections.get_by_id(id)
- pcollection.pipeline = p
-
- # Inject PBegin input where necessary.
- from apache_beam.io.iobase import Read
- from apache_beam.transforms.core import Create
- has_pbegin = [Read, Create]
- for id in proto.components.transforms:
- transform = context.transforms.get_by_id(id)
- if not transform.inputs and transform.transform.__class__ in has_pbegin:
- transform.inputs = (pvalue.PBegin(p),)
-
+ context.pcollections.get_by_id(id).pipeline = p
return p
@@ -558,7 +384,7 @@ class PipelineVisitor(object):
pass
def visit_transform(self, transform_node):
- """Callback for visiting a transform leaf node in the pipeline DAG."""
+ """Callback for visiting a transform node in the pipeline DAG."""
pass
def enter_composite_transform(self, transform_node):
@@ -615,20 +441,6 @@ class AppliedPTransform(object):
for side_input in self.side_inputs:
real_producer(side_input.pvalue).refcounts[side_input.pvalue.tag] += 1
- def replace_output(self, output, tag=None):
- """Replaces the output defined by the given tag with the given output.
-
- Args:
- output: replacement output
- tag: tag of the output to be replaced.
- """
- if isinstance(output, pvalue.DoOutputsTuple):
- self.replace_output(output[output._main_tag])
- elif isinstance(output, pvalue.PValue):
- self.outputs[tag] = output
- else:
- raise TypeError("Unexpected output type: %s" % output)
-
def add_output(self, output, tag=None):
if isinstance(output, pvalue.DoOutputsTuple):
self.add_output(output[output._main_tag])
@@ -713,7 +525,7 @@ class AppliedPTransform(object):
if isinstance(output, pvalue.PCollection)}
def to_runner_api(self, context):
- from apache_beam.portability.api import beam_runner_api_pb2
+ from apache_beam.runners.api import beam_runner_api_pb2
def transform_to_runner_api(transform, context):
if transform is None:
@@ -752,37 +564,3 @@ class AppliedPTransform(object):
pc.tag = tag
result.update_input_refcounts()
return result
-
-
-class PTransformOverride(object):
- """For internal use only; no backwards-compatibility guarantees.
-
- Gives a matcher and replacements for matching PTransforms.
-
- TODO: Update this to support cases where input and/our output types are
- different.
- """
- __metaclass__ = abc.ABCMeta
-
- @abc.abstractmethod
- def get_matcher(self):
- """Gives a matcher that will be used to to perform this override.
-
- Returns:
- a callable that takes an AppliedPTransform as a parameter and returns a
- boolean as a result.
- """
- raise NotImplementedError
-
- @abc.abstractmethod
- def get_replacement_transform(self, ptransform):
- """Provides a runner specific override for a given PTransform.
-
- Args:
- ptransform: PTransform to be replaced.
- Returns:
- A PTransform that will be the replacement for the PTransform given as an
- argument.
- """
- # Returns a PTransformReplacement
- raise NotImplementedError
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/pipeline_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/pipeline_test.py b/sdks/python/apache_beam/pipeline_test.py
index aad0143..e0775d1 100644
--- a/sdks/python/apache_beam/pipeline_test.py
+++ b/sdks/python/apache_beam/pipeline_test.py
@@ -28,11 +28,9 @@ import apache_beam as beam
from apache_beam.io import Read
from apache_beam.metrics import Metrics
from apache_beam.pipeline import Pipeline
-from apache_beam.pipeline import PTransformOverride
from apache_beam.pipeline import PipelineOptions
from apache_beam.pipeline import PipelineVisitor
from apache_beam.pvalue import AsSingleton
-from apache_beam.runners import DirectRunner
from apache_beam.runners.dataflow.native_io.iobase import NativeSource
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
@@ -77,18 +75,6 @@ class FakeSource(NativeSource):
return FakeSource._Reader(self._vals)
-class DoubleParDo(beam.PTransform):
- def expand(self, input):
- return input | 'Inner' >> beam.Map(lambda a: a * 2)
-
-
-class TripleParDo(beam.PTransform):
- def expand(self, input):
- # Keeping labels the same intentionally to make sure that there is no label
- # conflict due to replacement.
- return input | 'Inner' >> beam.Map(lambda a: a * 3)
-
-
class PipelineTest(unittest.TestCase):
@staticmethod
@@ -299,27 +285,6 @@ class PipelineTest(unittest.TestCase):
# p = Pipeline('EagerRunner')
# self.assertEqual([1, 4, 9], p | Create([1, 2, 3]) | Map(lambda x: x*x))
- def test_ptransform_overrides(self):
-
- def my_par_do_matcher(applied_ptransform):
- return isinstance(applied_ptransform.transform, DoubleParDo)
-
- class MyParDoOverride(PTransformOverride):
-
- def get_matcher(self):
- return my_par_do_matcher
-
- def get_replacement_transform(self, ptransform):
- if isinstance(ptransform, DoubleParDo):
- return TripleParDo()
- raise ValueError('Unsupported type of transform: %r', ptransform)
-
- # Using following private variable for testing.
- DirectRunner._PTRANSFORM_OVERRIDES.append(MyParDoOverride())
- with Pipeline() as p:
- pcoll = p | beam.Create([1, 2, 3]) | 'Multiply' >> DoubleParDo()
- assert_that(pcoll, equal_to([3, 6, 9]))
-
class DoFnTest(unittest.TestCase):
@@ -480,24 +445,6 @@ class RunnerApiTest(unittest.TestCase):
p2 = Pipeline.from_runner_api(proto, p.runner, p._options)
p2.run()
- def test_pickling(self):
- class MyPTransform(beam.PTransform):
- pickle_count = [0]
-
- def expand(self, p):
- self.p = p
- return p | beam.Create([None])
-
- def __reduce__(self):
- self.pickle_count[0] += 1
- return str, ()
-
- p = beam.Pipeline()
- for k in range(20):
- p | 'Iter%s' % k >> MyPTransform() # pylint: disable=expression-not-assigned
- p.to_runner_api()
- self.assertEqual(MyPTransform.pickle_count[0], 20)
-
if __name__ == '__main__':
logging.getLogger().setLevel(logging.DEBUG)
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/portability/__init__.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/portability/__init__.py b/sdks/python/apache_beam/portability/__init__.py
deleted file mode 100644
index 0bce5d6..0000000
--- a/sdks/python/apache_beam/portability/__init__.py
+++ /dev/null
@@ -1,18 +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.
-#
-
-"""For internal use only; no backwards-compatibility guarantees."""
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/portability/api/__init__.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/portability/api/__init__.py b/sdks/python/apache_beam/portability/api/__init__.py
deleted file mode 100644
index 2750859..0000000
--- a/sdks/python/apache_beam/portability/api/__init__.py
+++ /dev/null
@@ -1,21 +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.
-#
-
-"""For internal use only; no backwards-compatibility guarantees.
-
-Automatically generated when running setup.py sdist or build[_py].
-"""
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/pvalue.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/pvalue.py b/sdks/python/apache_beam/pvalue.py
index 34a483e..7385e82 100644
--- a/sdks/python/apache_beam/pvalue.py
+++ b/sdks/python/apache_beam/pvalue.py
@@ -128,7 +128,7 @@ class PCollection(PValue):
return _InvalidUnpickledPCollection, ()
def to_runner_api(self, context):
- from apache_beam.portability.api import beam_runner_api_pb2
+ from apache_beam.runners.api import beam_runner_api_pb2
from apache_beam.internal import pickler
return beam_runner_api_pb2.PCollection(
unique_name='%d%s.%s' % (
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/api/__init__.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/api/__init__.py b/sdks/python/apache_beam/runners/api/__init__.py
new file mode 100644
index 0000000..2750859
--- /dev/null
+++ b/sdks/python/apache_beam/runners/api/__init__.py
@@ -0,0 +1,21 @@
+#
+# 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.
+#
+
+"""For internal use only; no backwards-compatibility guarantees.
+
+Automatically generated when running setup.py sdist or build[_py].
+"""
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py b/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
index 059e139..3fc8983 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
@@ -39,17 +39,13 @@ from apache_beam.runners.dataflow.internal import names
from apache_beam.runners.dataflow.internal.clients import dataflow as dataflow_api
from apache_beam.runners.dataflow.internal.names import PropertyNames
from apache_beam.runners.dataflow.internal.names import TransformNames
-from apache_beam.runners.dataflow.ptransform_overrides import CreatePTransformOverride
from apache_beam.runners.runner import PValueCache
from apache_beam.runners.runner import PipelineResult
from apache_beam.runners.runner import PipelineRunner
from apache_beam.runners.runner import PipelineState
from apache_beam.transforms.display import DisplayData
from apache_beam.typehints import typehints
-from apache_beam.options.pipeline_options import SetupOptions
from apache_beam.options.pipeline_options import StandardOptions
-from apache_beam.options.pipeline_options import TestOptions
-from apache_beam.utils.plugin import BeamPlugin
__all__ = ['DataflowRunner']
@@ -65,15 +61,11 @@ class DataflowRunner(PipelineRunner):
if blocking is set to False.
"""
- # A list of PTransformOverride objects to be applied before running a pipeline
- # using DataflowRunner.
- # Currently this only works for overrides where the input and output types do
- # not change.
- # For internal SDK use only. This should not be updated by Beam pipeline
- # authors.
- _PTRANSFORM_OVERRIDES = [
- CreatePTransformOverride(),
- ]
+ # Environment version information. It is passed to the service during a
+ # a job submission and is used by the service to establish what features
+ # are expected by the workers.
+ BATCH_ENVIRONMENT_MAJOR_VERSION = '6'
+ STREAMING_ENVIRONMENT_MAJOR_VERSION = '1'
def __init__(self, cache=None):
# Cache of CloudWorkflowStep protos generated while the runner
@@ -223,6 +215,7 @@ class DataflowRunner(PipelineRunner):
return FlattenInputVisitor()
+ # TODO(mariagh): Make this method take pipepline_options
def run(self, pipeline):
"""Remotely executes entire pipeline or parts reachable from node."""
# Import here to avoid adding the dependency for local running scenarios.
@@ -233,17 +226,6 @@ class DataflowRunner(PipelineRunner):
raise ImportError(
'Google Cloud Dataflow runner not available, '
'please install apache_beam[gcp]')
-
- # Performing configured PTransform overrides.
- pipeline.replace_all(DataflowRunner._PTRANSFORM_OVERRIDES)
-
- # Add setup_options for all the BeamPlugin imports
- setup_options = pipeline._options.view_as(SetupOptions)
- plugins = BeamPlugin.get_all_plugin_paths()
- if setup_options.beam_plugins is not None:
- plugins = list(set(plugins + setup_options.beam_plugins))
- setup_options.beam_plugins = plugins
-
self.job = apiclient.Job(pipeline._options)
# Dataflow runner requires a KV type for GBK inputs, hence we enforce that
@@ -257,14 +239,15 @@ class DataflowRunner(PipelineRunner):
# The superclass's run will trigger a traversal of all reachable nodes.
super(DataflowRunner, self).run(pipeline)
- test_options = pipeline._options.view_as(TestOptions)
- # If it is a dry run, return without submitting the job.
- if test_options.dry_run:
- return None
+ standard_options = pipeline._options.view_as(StandardOptions)
+ if standard_options.streaming:
+ job_version = DataflowRunner.STREAMING_ENVIRONMENT_MAJOR_VERSION
+ else:
+ job_version = DataflowRunner.BATCH_ENVIRONMENT_MAJOR_VERSION
# Get a Dataflow API client and set its options
self.dataflow_client = apiclient.DataflowApplicationClient(
- pipeline._options)
+ pipeline._options, job_version)
# Create the job
result = DataflowPipelineResult(
@@ -377,26 +360,6 @@ class DataflowRunner(PipelineRunner):
PropertyNames.OUTPUT_NAME: PropertyNames.OUT}])
return step
- def run_Impulse(self, transform_node):
- standard_options = (
- transform_node.outputs[None].pipeline._options.view_as(StandardOptions))
- if standard_options.streaming:
- step = self._add_step(
- TransformNames.READ, transform_node.full_label, transform_node)
- step.add_property(PropertyNames.FORMAT, 'pubsub')
- step.add_property(PropertyNames.PUBSUB_SUBSCRIPTION, '_starting_signal/')
-
- step.encoding = self._get_encoded_output_coder(transform_node)
- step.add_property(
- PropertyNames.OUTPUT_INFO,
- [{PropertyNames.USER_NAME: (
- '%s.%s' % (
- transform_node.full_label, PropertyNames.OUT)),
- PropertyNames.ENCODING: step.encoding,
- PropertyNames.OUTPUT_NAME: PropertyNames.OUT}])
- else:
- ValueError('Impulse source for batch pipelines has not been defined.')
-
def run_Flatten(self, transform_node):
step = self._add_step(TransformNames.FLATTEN,
transform_node.full_label, transform_node)
@@ -655,13 +618,10 @@ class DataflowRunner(PipelineRunner):
if not standard_options.streaming:
raise ValueError('PubSubPayloadSource is currently available for use '
'only in streaming pipelines.')
- # Only one of topic or subscription should be set.
- if transform.source.full_subscription:
+ step.add_property(PropertyNames.PUBSUB_TOPIC, transform.source.topic)
+ if transform.source.subscription:
step.add_property(PropertyNames.PUBSUB_SUBSCRIPTION,
- transform.source.full_subscription)
- elif transform.source.full_topic:
- step.add_property(PropertyNames.PUBSUB_TOPIC,
- transform.source.full_topic)
+ transform.source.topic)
if transform.source.id_label:
step.add_property(PropertyNames.PUBSUB_ID_LABEL,
transform.source.id_label)
@@ -679,12 +639,7 @@ class DataflowRunner(PipelineRunner):
# step should be the type of value outputted by each step. Read steps
# automatically wrap output values in a WindowedValue wrapper, if necessary.
# This is also necessary for proper encoding for size estimation.
- # Using a GlobalWindowCoder as a place holder instead of the default
- # PickleCoder because GlobalWindowCoder is known coder.
- # TODO(robertwb): Query the collection for the windowfn to extract the
- # correct coder.
- coder = coders.WindowedValueCoder(transform._infer_output_coder(),
- coders.coders.GlobalWindowCoder()) # pylint: disable=protected-access
+ coder = coders.WindowedValueCoder(transform._infer_output_coder()) # pylint: disable=protected-access
step.encoding = self._get_cloud_encoding(coder)
step.add_property(
@@ -745,7 +700,7 @@ class DataflowRunner(PipelineRunner):
if not standard_options.streaming:
raise ValueError('PubSubPayloadSink is currently available for use '
'only in streaming pipelines.')
- step.add_property(PropertyNames.PUBSUB_TOPIC, transform.sink.full_topic)
+ step.add_property(PropertyNames.PUBSUB_TOPIC, transform.sink.topic)
else:
raise ValueError(
'Sink %r has unexpected format %s.' % (
@@ -753,12 +708,8 @@ class DataflowRunner(PipelineRunner):
step.add_property(PropertyNames.FORMAT, transform.sink.format)
# Wrap coder in WindowedValueCoder: this is necessary for proper encoding
- # for size estimation. Using a GlobalWindowCoder as a place holder instead
- # of the default PickleCoder because GlobalWindowCoder is known coder.
- # TODO(robertwb): Query the collection for the windowfn to extract the
- # correct coder.
- coder = coders.WindowedValueCoder(transform.sink.coder,
- coders.coders.GlobalWindowCoder())
+ # for size estimation.
+ coder = coders.WindowedValueCoder(transform.sink.coder)
step.encoding = self._get_cloud_encoding(coder)
step.add_property(PropertyNames.ENCODING, step.encoding)
step.add_property(
@@ -770,7 +721,7 @@ class DataflowRunner(PipelineRunner):
@classmethod
def serialize_windowing_strategy(cls, windowing):
from apache_beam.runners import pipeline_context
- from apache_beam.portability.api import beam_runner_api_pb2
+ from apache_beam.runners.api import beam_runner_api_pb2
context = pipeline_context.PipelineContext()
windowing_proto = windowing.to_runner_api(context)
return cls.byte_array_to_json_string(
@@ -783,7 +734,7 @@ class DataflowRunner(PipelineRunner):
# Imported here to avoid circular dependencies.
# pylint: disable=wrong-import-order, wrong-import-position
from apache_beam.runners import pipeline_context
- from apache_beam.portability.api import beam_runner_api_pb2
+ from apache_beam.runners.api import beam_runner_api_pb2
from apache_beam.transforms.core import Windowing
proto = beam_runner_api_pb2.MessageWithComponents()
proto.ParseFromString(cls.json_string_to_byte_array(serialized_data))
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py b/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py
index a9b8fdb..74fd01d 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py
@@ -59,8 +59,7 @@ class DataflowRunnerTest(unittest.TestCase):
'--project=test-project',
'--staging_location=ignored',
'--temp_location=/dev/null',
- '--no_auth=True',
- '--dry_run=True']
+ '--no_auth=True']
@mock.patch('time.sleep', return_value=None)
def test_wait_until_finish(self, patched_time_sleep):
@@ -109,22 +108,8 @@ class DataflowRunnerTest(unittest.TestCase):
(p | ptransform.Create([1, 2, 3]) # pylint: disable=expression-not-assigned
| 'Do' >> ptransform.FlatMap(lambda x: [(x, x)])
| ptransform.GroupByKey())
- p.run()
-
- def test_streaming_create_translation(self):
- remote_runner = DataflowRunner()
- self.default_properties.append("--streaming")
- p = Pipeline(remote_runner, PipelineOptions(self.default_properties))
- p | ptransform.Create([1]) # pylint: disable=expression-not-assigned
- p.run()
- job_dict = json.loads(str(remote_runner.job))
- self.assertEqual(len(job_dict[u'steps']), 2)
-
- self.assertEqual(job_dict[u'steps'][0][u'kind'], u'ParallelRead')
- self.assertEqual(
- job_dict[u'steps'][0][u'properties'][u'pubsub_subscription'],
- '_starting_signal/')
- self.assertEqual(job_dict[u'steps'][1][u'kind'], u'ParallelDo')
+ remote_runner.job = apiclient.Job(p._options)
+ super(DataflowRunner, remote_runner).run(p)
def test_remote_runner_display_data(self):
remote_runner = DataflowRunner()
@@ -157,7 +142,8 @@ class DataflowRunnerTest(unittest.TestCase):
(p | ptransform.Create([1, 2, 3, 4, 5])
| 'Do' >> SpecialParDo(SpecialDoFn(), now))
- p.run()
+ remote_runner.job = apiclient.Job(p._options)
+ super(DataflowRunner, remote_runner).run(p)
job_dict = json.loads(str(remote_runner.job))
steps = [step
for step in job_dict['steps']
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py b/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py
index 33dfe19..df1a3f2 100644
--- a/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py
+++ b/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py
@@ -38,6 +38,7 @@ from apache_beam.io.filesystems import FileSystems
from apache_beam.io.gcp.internal.clients import storage
from apache_beam.runners.dataflow.internal import dependency
from apache_beam.runners.dataflow.internal.clients import dataflow
+from apache_beam.runners.dataflow.internal.dependency import get_required_container_version
from apache_beam.runners.dataflow.internal.dependency import get_sdk_name_and_version
from apache_beam.runners.dataflow.internal.names import PropertyNames
from apache_beam.transforms import cy_combiners
@@ -49,13 +50,6 @@ from apache_beam.options.pipeline_options import StandardOptions
from apache_beam.options.pipeline_options import WorkerOptions
-# Environment version information. It is passed to the service during a
-# a job submission and is used by the service to establish what features
-# are expected by the workers.
-_LEGACY_ENVIRONMENT_MAJOR_VERSION = '6'
-_FNAPI_ENVIRONMENT_MAJOR_VERSION = '1'
-
-
class Step(object):
"""Wrapper for a dataflow Step protobuf."""
@@ -153,10 +147,7 @@ class Environment(object):
if self.standard_options.streaming:
job_type = 'FNAPI_STREAMING'
else:
- if _use_fnapi(options):
- job_type = 'FNAPI_BATCH'
- else:
- job_type = 'PYTHON_BATCH'
+ job_type = 'PYTHON_BATCH'
self.proto.version.additionalProperties.extend([
dataflow.Environment.VersionValue.AdditionalProperty(
key='job_type',
@@ -214,8 +205,11 @@ class Environment(object):
pool.workerHarnessContainerImage = (
self.worker_options.worker_harness_container_image)
else:
+ # Default to using the worker harness container image for the current SDK
+ # version.
pool.workerHarnessContainerImage = (
- dependency.get_default_container_image_for_current_sdk(job_type))
+ 'dataflow.gcr.io/v1beta3/python:%s' %
+ get_required_container_version())
if self.worker_options.use_public_ips is not None:
if self.worker_options.use_public_ips:
pool.ipConfiguration = (
@@ -370,16 +364,11 @@ class Job(object):
class DataflowApplicationClient(object):
"""A Dataflow API client used by application code to create and query jobs."""
- def __init__(self, options):
+ def __init__(self, options, environment_version):
"""Initializes a Dataflow API client object."""
self.standard_options = options.view_as(StandardOptions)
self.google_cloud_options = options.view_as(GoogleCloudOptions)
-
- if _use_fnapi(options):
- self.environment_version = _FNAPI_ENVIRONMENT_MAJOR_VERSION
- else:
- self.environment_version = _LEGACY_ENVIRONMENT_MAJOR_VERSION
-
+ self.environment_version = environment_version
if self.google_cloud_options.no_auth:
credentials = None
else:
@@ -721,14 +710,6 @@ def translate_mean(accumulator, metric_update):
metric_update.kind = None
-def _use_fnapi(pipeline_options):
- standard_options = pipeline_options.view_as(StandardOptions)
- debug_options = pipeline_options.view_as(DebugOptions)
-
- return standard_options.streaming or (
- debug_options.experiments and 'beam_fn_api' in debug_options.experiments)
-
-
# To enable a counter on the service, add it to this dictionary.
metric_translations = {
cy_combiners.CountCombineFn: ('sum', translate_scalar),
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py b/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py
index 407ffcf..67cf77f 100644
--- a/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py
+++ b/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py
@@ -22,6 +22,7 @@ from mock import Mock
from apache_beam.metrics.cells import DistributionData
from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.runners.dataflow.dataflow_runner import DataflowRunner
from apache_beam.runners.dataflow.internal.clients import dataflow
# Protect against environments where apitools library is not available.
@@ -39,7 +40,9 @@ class UtilTest(unittest.TestCase):
@unittest.skip("Enable once BEAM-1080 is fixed.")
def test_create_application_client(self):
pipeline_options = PipelineOptions()
- apiclient.DataflowApplicationClient(pipeline_options)
+ apiclient.DataflowApplicationClient(
+ pipeline_options,
+ DataflowRunner.BATCH_ENVIRONMENT_MAJOR_VERSION)
def test_set_network(self):
pipeline_options = PipelineOptions(
@@ -119,30 +122,6 @@ class UtilTest(unittest.TestCase):
self.assertEqual(
metric_update.floatingPointMean.count.lowBits, accumulator.count)
- def test_default_ip_configuration(self):
- pipeline_options = PipelineOptions(
- ['--temp_location', 'gs://any-location/temp'])
- env = apiclient.Environment([], pipeline_options, '2.0.0')
- self.assertEqual(env.proto.workerPools[0].ipConfiguration, None)
-
- def test_public_ip_configuration(self):
- pipeline_options = PipelineOptions(
- ['--temp_location', 'gs://any-location/temp',
- '--use_public_ips'])
- env = apiclient.Environment([], pipeline_options, '2.0.0')
- self.assertEqual(
- env.proto.workerPools[0].ipConfiguration,
- dataflow.WorkerPool.IpConfigurationValueValuesEnum.WORKER_IP_PUBLIC)
-
- def test_private_ip_configuration(self):
- pipeline_options = PipelineOptions(
- ['--temp_location', 'gs://any-location/temp',
- '--no_use_public_ips'])
- env = apiclient.Environment([], pipeline_options, '2.0.0')
- self.assertEqual(
- env.proto.workerPools[0].ipConfiguration,
- dataflow.WorkerPool.IpConfigurationValueValuesEnum.WORKER_IP_PRIVATE)
-
if __name__ == '__main__':
unittest.main()
[18/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileBasedSink.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileBasedSink.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileBasedSink.java
index 583af60..8102316 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileBasedSink.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileBasedSink.java
@@ -33,7 +33,6 @@ import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.Serializable;
-import java.lang.reflect.TypeVariable;
import java.nio.channels.Channels;
import java.nio.channels.WritableByteChannel;
import java.util.ArrayList;
@@ -50,10 +49,8 @@ import java.util.zip.GZIPOutputStream;
import javax.annotation.Nullable;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.annotations.Experimental.Kind;
-import org.apache.beam.sdk.coders.CannotProvideCoderException;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderException;
-import org.apache.beam.sdk.coders.CoderRegistry;
import org.apache.beam.sdk.coders.NullableCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.StructuredCoder;
@@ -76,7 +73,6 @@ import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.transforms.windowing.PaneInfo.PaneInfoCoder;
import org.apache.beam.sdk.util.MimeTypes;
-import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.commons.compress.compressors.bzip2.BZip2CompressorOutputStream;
import org.apache.commons.compress.compressors.deflate.DeflateCompressorOutputStream;
import org.joda.time.Instant;
@@ -86,43 +82,43 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
- * Abstract class for file-based output. An implementation of FileBasedSink writes file-based output
- * and defines the format of output files (how values are written, headers/footers, MIME type,
- * etc.).
+ * Abstract class for file-based output. An implementation of FileBasedSink writes file-based
+ * output and defines the format of output files (how values are written, headers/footers, MIME
+ * type, etc.).
*
* <p>At pipeline construction time, the methods of FileBasedSink are called to validate the sink
* and to create a {@link WriteOperation} that manages the process of writing to the sink.
*
* <p>The process of writing to file-based sink is as follows:
- *
* <ol>
- * <li>An optional subclass-defined initialization,
- * <li>a parallel write of bundles to temporary files, and finally,
- * <li>these temporary files are renamed with final output filenames.
+ * <li>An optional subclass-defined initialization,
+ * <li>a parallel write of bundles to temporary files, and finally,
+ * <li>these temporary files are renamed with final output filenames.
* </ol>
*
* <p>In order to ensure fault-tolerance, a bundle may be executed multiple times (e.g., in the
* event of failure/retry or for redundancy). However, exactly one of these executions will have its
- * result passed to the finalize method. Each call to {@link Writer#openWindowed} or {@link
- * Writer#openUnwindowed} is passed a unique <i>bundle id</i> when it is called by the WriteFiles
- * transform, so even redundant or retried bundles will have a unique way of identifying their
- * output.
+ * result passed to the finalize method. Each call to {@link Writer#openWindowed}
+ * or {@link Writer#openUnwindowed} is passed a unique <i>bundle id</i> when it is called
+ * by the WriteFiles transform, so even redundant or retried bundles will have a unique way of
+ * identifying
+ * their output.
*
* <p>The bundle id should be used to guarantee that a bundle's output is unique. This uniqueness
* guarantee is important; if a bundle is to be output to a file, for example, the name of the file
* will encode the unique bundle id to avoid conflicts with other writers.
*
- * <p>{@link FileBasedSink} can take a custom {@link FilenamePolicy} object to determine output
- * filenames, and this policy object can be used to write windowed or triggered PCollections into
- * separate files per window pane. This allows file output from unbounded PCollections, and also
- * works for bounded PCollecctions.
+ * {@link FileBasedSink} can take a custom {@link FilenamePolicy} object to determine output
+ * filenames, and this policy object can be used to write windowed or triggered
+ * PCollections into separate files per window pane. This allows file output from unbounded
+ * PCollections, and also works for bounded PCollecctions.
*
* <p>Supported file systems are those registered with {@link FileSystems}.
*
- * @param <OutputT> the type of values written to the sink.
+ * @param <T> the type of values written to the sink.
*/
@Experimental(Kind.FILESYSTEM)
-public abstract class FileBasedSink<OutputT, DestinationT> implements Serializable, HasDisplayData {
+public abstract class FileBasedSink<T> implements Serializable, HasDisplayData {
private static final Logger LOG = LoggerFactory.getLogger(FileBasedSink.class);
/**
@@ -177,7 +173,7 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
}
@Override
- public String getSuggestedFilenameSuffix() {
+ public String getFilenameSuffix() {
return filenameSuffix;
}
@@ -209,8 +205,6 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
}
}
- private final DynamicDestinations<?, DestinationT> dynamicDestinations;
-
/**
* The {@link WritableByteChannelFactory} that is used to wrap the raw data output to the
* underlying channel. The default is to not compress the output using
@@ -219,70 +213,8 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
private final WritableByteChannelFactory writableByteChannelFactory;
/**
- * A class that allows value-dependent writes in {@link FileBasedSink}.
- *
- * <p>Users can define a custom type to represent destinations, and provide a mapping to turn this
- * destination type into an instance of {@link FilenamePolicy}.
+ * A naming policy for output files.
*/
- @Experimental(Kind.FILESYSTEM)
- public abstract static class DynamicDestinations<UserT, DestinationT>
- implements HasDisplayData, Serializable {
- /**
- * Returns an object that represents at a high level the destination being written to. May not
- * return null.
- */
- public abstract DestinationT getDestination(UserT element);
-
- /**
- * Returns the default destination. This is used for collections that have no elements as the
- * destination to write empty files to.
- */
- public abstract DestinationT getDefaultDestination();
-
- /**
- * Returns the coder for {@link DestinationT}. If this is not overridden, then the coder
- * registry will be use to find a suitable coder. This must be a deterministic coder, as {@link
- * DestinationT} will be used as a key type in a {@link
- * org.apache.beam.sdk.transforms.GroupByKey}.
- */
- @Nullable
- public Coder<DestinationT> getDestinationCoder() {
- return null;
- }
-
- /** Converts a destination into a {@link FilenamePolicy}. May not return null. */
- public abstract FilenamePolicy getFilenamePolicy(DestinationT destination);
-
- /** Populates the display data. */
- @Override
- public void populateDisplayData(DisplayData.Builder builder) {}
-
- // Gets the destination coder. If the user does not provide one, try to find one in the coder
- // registry. If no coder can be found, throws CannotProvideCoderException.
- final Coder<DestinationT> getDestinationCoderWithDefault(CoderRegistry registry)
- throws CannotProvideCoderException {
- Coder<DestinationT> destinationCoder = getDestinationCoder();
- if (destinationCoder != null) {
- return destinationCoder;
- }
- // If dynamicDestinations doesn't provide a coder, try to find it in the coder registry.
- // We must first use reflection to figure out what the type parameter is.
- TypeDescriptor<?> superDescriptor =
- TypeDescriptor.of(getClass()).getSupertype(DynamicDestinations.class);
- if (!superDescriptor.getRawType().equals(DynamicDestinations.class)) {
- throw new AssertionError(
- "Couldn't find the DynamicDestinations superclass of " + this.getClass());
- }
- TypeVariable typeVariable = superDescriptor.getTypeParameter("DestinationT");
- @SuppressWarnings("unchecked")
- TypeDescriptor<DestinationT> descriptor =
- (TypeDescriptor<DestinationT>) superDescriptor.resolveType(typeVariable);
- return registry.getCoder(descriptor);
- }
- }
-
- /** A naming policy for output files. */
- @Experimental(Kind.FILESYSTEM)
public abstract static class FilenamePolicy implements Serializable {
/**
* Context used for generating a name based on shard number, and num shards.
@@ -355,28 +287,29 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
/**
* When a sink has requested windowed or triggered output, this method will be invoked to return
* the file {@link ResourceId resource} to be created given the base output directory and a
- * {@link OutputFileHints} containing information about the file, including a suggested
- * extension (e.g. coming from {@link CompressionType}).
+ * (possibly empty) extension from {@link FileBasedSink} configuration
+ * (e.g., {@link CompressionType}).
*
- * <p>The {@link WindowedContext} object gives access to the window and pane, as well as
- * sharding information. The policy must return unique and consistent filenames for different
- * windows and panes.
+ * <p>The {@link WindowedContext} object gives access to the window and pane,
+ * as well as sharding information. The policy must return unique and consistent filenames
+ * for different windows and panes.
*/
@Experimental(Kind.FILESYSTEM)
- public abstract ResourceId windowedFilename(WindowedContext c, OutputFileHints outputFileHints);
+ public abstract ResourceId windowedFilename(
+ ResourceId outputDirectory, WindowedContext c, String extension);
/**
* When a sink has not requested windowed or triggered output, this method will be invoked to
* return the file {@link ResourceId resource} to be created given the base output directory and
- * a {@link OutputFileHints} containing information about the file, including a suggested (e.g.
- * coming from {@link CompressionType}).
+ * a (possibly empty) extension applied by additional {@link FileBasedSink} configuration
+ * (e.g., {@link CompressionType}).
*
* <p>The {@link Context} object only provides sharding information, which is used by the policy
* to generate unique and consistent filenames.
*/
@Experimental(Kind.FILESYSTEM)
- @Nullable
- public abstract ResourceId unwindowedFilename(Context c, OutputFileHints outputFileHints);
+ @Nullable public abstract ResourceId unwindowedFilename(
+ ResourceId outputDirectory, Context c, String extension);
/**
* Populates the display data.
@@ -385,8 +318,19 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
}
}
+ /** The policy used to generate names of files to be produced. */
+ private final FilenamePolicy filenamePolicy;
/** The directory to which files will be written. */
- private final ValueProvider<ResourceId> tempDirectoryProvider;
+ private final ValueProvider<ResourceId> baseOutputDirectoryProvider;
+
+ /**
+ * Construct a {@link FileBasedSink} with the given filename policy, producing uncompressed files.
+ */
+ @Experimental(Kind.FILESYSTEM)
+ public FileBasedSink(
+ ValueProvider<ResourceId> baseOutputDirectoryProvider, FilenamePolicy filenamePolicy) {
+ this(baseOutputDirectoryProvider, filenamePolicy, CompressionType.UNCOMPRESSED);
+ }
private static class ExtractDirectory implements SerializableFunction<ResourceId, ResourceId> {
@Override
@@ -396,91 +340,95 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
}
/**
- * Construct a {@link FileBasedSink} with the given temp directory, producing uncompressed files.
+ * Construct a {@link FileBasedSink} with the given filename policy and output channel type.
*/
@Experimental(Kind.FILESYSTEM)
public FileBasedSink(
- ValueProvider<ResourceId> tempDirectoryProvider,
- DynamicDestinations<?, DestinationT> dynamicDestinations) {
- this(tempDirectoryProvider, dynamicDestinations, CompressionType.UNCOMPRESSED);
- }
-
- /** Construct a {@link FileBasedSink} with the given temp directory and output channel type. */
- @Experimental(Kind.FILESYSTEM)
- public FileBasedSink(
- ValueProvider<ResourceId> tempDirectoryProvider,
- DynamicDestinations<?, DestinationT> dynamicDestinations,
+ ValueProvider<ResourceId> baseOutputDirectoryProvider,
+ FilenamePolicy filenamePolicy,
WritableByteChannelFactory writableByteChannelFactory) {
- this.tempDirectoryProvider =
- NestedValueProvider.of(tempDirectoryProvider, new ExtractDirectory());
- this.dynamicDestinations = checkNotNull(dynamicDestinations);
+ this.baseOutputDirectoryProvider =
+ NestedValueProvider.of(baseOutputDirectoryProvider, new ExtractDirectory());
+ this.filenamePolicy = filenamePolicy;
this.writableByteChannelFactory = writableByteChannelFactory;
}
- /** Return the {@link DynamicDestinations} used. */
- @SuppressWarnings("unchecked")
- public <UserT> DynamicDestinations<UserT, DestinationT> getDynamicDestinations() {
- return (DynamicDestinations<UserT, DestinationT>) dynamicDestinations;
+ /**
+ * Returns the base directory inside which files will be written according to the configured
+ * {@link FilenamePolicy}.
+ */
+ @Experimental(Kind.FILESYSTEM)
+ public ValueProvider<ResourceId> getBaseOutputDirectoryProvider() {
+ return baseOutputDirectoryProvider;
}
/**
- * Returns the directory inside which temprary files will be written according to the configured
- * {@link FilenamePolicy}.
+ * Returns the policy by which files will be named inside of the base output directory. Note that
+ * the {@link FilenamePolicy} may itself specify one or more inner directories before each output
+ * file, say when writing windowed outputs in a {@code output/YYYY/MM/DD/file.txt} format.
*/
@Experimental(Kind.FILESYSTEM)
- public ValueProvider<ResourceId> getTempDirectoryProvider() {
- return tempDirectoryProvider;
+ public final FilenamePolicy getFilenamePolicy() {
+ return filenamePolicy;
}
public void validate(PipelineOptions options) {}
- /** Return a subclass of {@link WriteOperation} that will manage the write to the sink. */
- public abstract WriteOperation<OutputT, DestinationT> createWriteOperation();
+ /**
+ * Return a subclass of {@link WriteOperation} that will manage the write
+ * to the sink.
+ */
+ public abstract WriteOperation<T> createWriteOperation();
public void populateDisplayData(DisplayData.Builder builder) {
- getDynamicDestinations().populateDisplayData(builder);
+ getFilenamePolicy().populateDisplayData(builder);
}
/**
* Abstract operation that manages the process of writing to {@link FileBasedSink}.
*
- * <p>The primary responsibilities of the WriteOperation is the management of output files. During
- * a write, {@link Writer}s write bundles to temporary file locations. After the bundles have been
- * written,
- *
+ * <p>The primary responsibilities of the WriteOperation is the management of output
+ * files. During a write, {@link Writer}s write bundles to temporary file
+ * locations. After the bundles have been written,
* <ol>
- * <li>{@link WriteOperation#finalize} is given a list of the temporary files containing the
- * output bundles.
- * <li>During finalize, these temporary files are copied to final output locations and named
- * according to a file naming template.
- * <li>Finally, any temporary files that were created during the write are removed.
+ * <li>{@link WriteOperation#finalize} is given a list of the temporary
+ * files containing the output bundles.
+ * <li>During finalize, these temporary files are copied to final output locations and named
+ * according to a file naming template.
+ * <li>Finally, any temporary files that were created during the write are removed.
* </ol>
*
- * <p>Subclass implementations of WriteOperation must implement {@link
- * WriteOperation#createWriter} to return a concrete FileBasedSinkWriter.
+ * <p>Subclass implementations of WriteOperation must implement
+ * {@link WriteOperation#createWriter} to return a concrete
+ * FileBasedSinkWriter.
*
- * <h2>Temporary and Output File Naming:</h2>
+ * <h2>Temporary and Output File Naming:</h2> During the write, bundles are written to temporary
+ * files using the tempDirectory that can be provided via the constructor of
+ * WriteOperation. These temporary files will be named
+ * {@code {tempDirectory}/{bundleId}}, where bundleId is the unique id of the bundle.
+ * For example, if tempDirectory is "gs://my-bucket/my_temp_output", the output for a
+ * bundle with bundle id 15723 will be "gs://my-bucket/my_temp_output/15723".
*
- * <p>During the write, bundles are written to temporary files using the tempDirectory that can be
- * provided via the constructor of WriteOperation. These temporary files will be named {@code
- * {tempDirectory}/{bundleId}}, where bundleId is the unique id of the bundle. For example, if
- * tempDirectory is "gs://my-bucket/my_temp_output", the output for a bundle with bundle id 15723
- * will be "gs://my-bucket/my_temp_output/15723".
+ * <p>Final output files are written to baseOutputFilename with the format
+ * {@code {baseOutputFilename}-0000i-of-0000n.{extension}} where n is the total number of bundles
+ * written and extension is the file extension. Both baseOutputFilename and extension are required
+ * constructor arguments.
*
- * <p>Final output files are written to the location specified by the {@link FilenamePolicy}. If
- * no filename policy is specified, then the {@link DefaultFilenamePolicy} will be used. The
- * directory that the files are written to is determined by the {@link FilenamePolicy} instance.
+ * <p>Subclass implementations can change the file naming template by supplying a value for
+ * fileNamingTemplate.
*
* <p>Note that in the case of permanent failure of a bundle's write, no clean up of temporary
* files will occur.
*
* <p>If there are no elements in the PCollection being written, no output will be generated.
*
- * @param <OutputT> the type of values written to the sink.
+ * @param <T> the type of values written to the sink.
*/
- public abstract static class WriteOperation<OutputT, DestinationT> implements Serializable {
- /** The Sink that this WriteOperation will write to. */
- protected final FileBasedSink<OutputT, DestinationT> sink;
+ public abstract static class WriteOperation<T> implements Serializable {
+ /**
+ * The Sink that this WriteOperation will write to.
+ */
+ protected final FileBasedSink<T> sink;
/** Directory for temporary output files. */
protected final ValueProvider<ResourceId> tempDirectory;
@@ -497,19 +445,17 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
}
/**
- * Constructs a WriteOperation using the default strategy for generating a temporary directory
- * from the base output filename.
+ * Constructs a WriteOperation using the default strategy for generating a temporary
+ * directory from the base output filename.
*
- * <p>Default is a uniquely named subdirectory of the provided tempDirectory, e.g. if
- * tempDirectory is /path/to/foo/, the temporary directory will be
- * /path/to/foo/temp-beam-foo-$date.
+ * <p>Default is a uniquely named sibling of baseOutputFilename, e.g. if baseOutputFilename is
+ * /path/to/foo, the temporary directory will be /path/to/temp-beam-foo-$date.
*
* @param sink the FileBasedSink that will be used to configure this write operation.
*/
- public WriteOperation(FileBasedSink<OutputT, DestinationT> sink) {
- this(
- sink,
- NestedValueProvider.of(sink.getTempDirectoryProvider(), new TemporaryDirectoryBuilder()));
+ public WriteOperation(FileBasedSink<T> sink) {
+ this(sink, NestedValueProvider.of(
+ sink.getBaseOutputDirectoryProvider(), new TemporaryDirectoryBuilder()));
}
private static class TemporaryDirectoryBuilder
@@ -525,12 +471,10 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
private final Long tempId = TEMP_COUNT.getAndIncrement();
@Override
- public ResourceId apply(ResourceId tempDirectory) {
+ public ResourceId apply(ResourceId baseOutputDirectory) {
// Temp directory has a timestamp and a unique ID
String tempDirName = String.format(".temp-beam-%s-%s", timestamp, tempId);
- return tempDirectory
- .getCurrentDirectory()
- .resolve(tempDirName, StandardResolveOptions.RESOLVE_DIRECTORY);
+ return baseOutputDirectory.resolve(tempDirName, StandardResolveOptions.RESOLVE_DIRECTORY);
}
}
@@ -541,22 +485,22 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
* @param tempDirectory the base directory to be used for temporary output files.
*/
@Experimental(Kind.FILESYSTEM)
- public WriteOperation(FileBasedSink<OutputT, DestinationT> sink, ResourceId tempDirectory) {
+ public WriteOperation(FileBasedSink<T> sink, ResourceId tempDirectory) {
this(sink, StaticValueProvider.of(tempDirectory));
}
private WriteOperation(
- FileBasedSink<OutputT, DestinationT> sink, ValueProvider<ResourceId> tempDirectory) {
+ FileBasedSink<T> sink, ValueProvider<ResourceId> tempDirectory) {
this.sink = sink;
this.tempDirectory = tempDirectory;
this.windowedWrites = false;
}
/**
- * Clients must implement to return a subclass of {@link Writer}. This method must not mutate
- * the state of the object.
+ * Clients must implement to return a subclass of {@link Writer}. This
+ * method must not mutate the state of the object.
*/
- public abstract Writer<OutputT, DestinationT> createWriter() throws Exception;
+ public abstract Writer<T> createWriter() throws Exception;
/**
* Indicates that the operation will be performing windowed writes.
@@ -570,8 +514,8 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
* removing temporary files.
*
* <p>Finalization may be overridden by subclass implementations to perform customized
- * finalization (e.g., initiating some operation on output bundles, merging them, etc.). {@code
- * writerResults} contains the filenames of written bundles.
+ * finalization (e.g., initiating some operation on output bundles, merging them, etc.).
+ * {@code writerResults} contains the filenames of written bundles.
*
* <p>If subclasses override this method, they must guarantee that its implementation is
* idempotent, as it may be executed multiple times in the case of failure or for redundancy. It
@@ -579,7 +523,7 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
*
* @param writerResults the results of writes (FileResult).
*/
- public void finalize(Iterable<FileResult<DestinationT>> writerResults) throws Exception {
+ public void finalize(Iterable<FileResult> writerResults) throws Exception {
// Collect names of temporary files and rename them.
Map<ResourceId, ResourceId> outputFilenames = buildOutputFilenames(writerResults);
copyToOutputFiles(outputFilenames);
@@ -598,14 +542,17 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
@Experimental(Kind.FILESYSTEM)
protected final Map<ResourceId, ResourceId> buildOutputFilenames(
- Iterable<FileResult<DestinationT>> writerResults) {
+ Iterable<FileResult> writerResults) {
int numShards = Iterables.size(writerResults);
Map<ResourceId, ResourceId> outputFilenames = new HashMap<>();
+ FilenamePolicy policy = getSink().getFilenamePolicy();
+ ResourceId baseOutputDir = getSink().getBaseOutputDirectoryProvider().get();
+
// Either all results have a shard number set (if the sink is configured with a fixed
// number of shards), or they all don't (otherwise).
Boolean isShardNumberSetEverywhere = null;
- for (FileResult<DestinationT> result : writerResults) {
+ for (FileResult result : writerResults) {
boolean isShardNumberSetHere = (result.getShard() != UNKNOWN_SHARDNUM);
if (isShardNumberSetEverywhere == null) {
isShardNumberSetEverywhere = isShardNumberSetHere;
@@ -621,7 +568,7 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
isShardNumberSetEverywhere = true;
}
- List<FileResult<DestinationT>> resultsWithShardNumbers = Lists.newArrayList();
+ List<FileResult> resultsWithShardNumbers = Lists.newArrayList();
if (isShardNumberSetEverywhere) {
resultsWithShardNumbers = Lists.newArrayList(writerResults);
} else {
@@ -630,32 +577,29 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
// case of triggers, the list of FileResult objects in the Finalize iterable is not
// deterministic, and might change over retries. This breaks the assumption below that
// sorting the FileResult objects provides idempotency.
- List<FileResult<DestinationT>> sortedByTempFilename =
+ List<FileResult> sortedByTempFilename =
Ordering.from(
- new Comparator<FileResult<DestinationT>>() {
- @Override
- public int compare(
- FileResult<DestinationT> first, FileResult<DestinationT> second) {
- String firstFilename = first.getTempFilename().toString();
- String secondFilename = second.getTempFilename().toString();
- return firstFilename.compareTo(secondFilename);
- }
- })
+ new Comparator<FileResult>() {
+ @Override
+ public int compare(FileResult first, FileResult second) {
+ String firstFilename = first.getTempFilename().toString();
+ String secondFilename = second.getTempFilename().toString();
+ return firstFilename.compareTo(secondFilename);
+ }
+ })
.sortedCopy(writerResults);
for (int i = 0; i < sortedByTempFilename.size(); i++) {
resultsWithShardNumbers.add(sortedByTempFilename.get(i).withShard(i));
}
}
- for (FileResult<DestinationT> result : resultsWithShardNumbers) {
+ for (FileResult result : resultsWithShardNumbers) {
checkArgument(
result.getShard() != UNKNOWN_SHARDNUM, "Should have set shard number on %s", result);
outputFilenames.put(
result.getTempFilename(),
result.getDestinationFile(
- getSink().getDynamicDestinations(),
- numShards,
- getSink().getWritableByteChannelFactory()));
+ policy, baseOutputDir, numShards, getSink().getExtension()));
}
int numDistinctShards = new HashSet<>(outputFilenames.values()).size();
@@ -671,18 +615,18 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
*
* <p>Can be called from subclasses that override {@link WriteOperation#finalize}.
*
- * <p>Files will be named according to the {@link FilenamePolicy}. The order of the output files
- * will be the same as the sorted order of the input filenames. In other words (when using
- * {@link DefaultFilenamePolicy}), if the input filenames are ["C", "A", "B"], baseFilename (int
- * the policy) is "dir/file", the extension is ".txt", and the fileNamingTemplate is
- * "-SSS-of-NNN", the contents of A will be copied to dir/file-000-of-003.txt, the contents of B
- * will be copied to dir/file-001-of-003.txt, etc.
+ * <p>Files will be named according to the file naming template. The order of the output files
+ * will be the same as the sorted order of the input filenames. In other words, if the input
+ * filenames are ["C", "A", "B"], baseOutputFilename is "file", the extension is ".txt", and
+ * the fileNamingTemplate is "-SSS-of-NNN", the contents of A will be copied to
+ * file-000-of-003.txt, the contents of B will be copied to file-001-of-003.txt, etc.
*
* @param filenames the filenames of temporary files.
*/
@VisibleForTesting
@Experimental(Kind.FILESYSTEM)
- final void copyToOutputFiles(Map<ResourceId, ResourceId> filenames) throws IOException {
+ final void copyToOutputFiles(Map<ResourceId, ResourceId> filenames)
+ throws IOException {
int numFiles = filenames.size();
if (numFiles > 0) {
LOG.debug("Copying {} files.", numFiles);
@@ -754,8 +698,10 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
}
}
- /** Returns the FileBasedSink for this write operation. */
- public FileBasedSink<OutputT, DestinationT> getSink() {
+ /**
+ * Returns the FileBasedSink for this write operation.
+ */
+ public FileBasedSink<T> getSink() {
return sink;
}
@@ -773,28 +719,33 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
}
}
- /** Returns the {@link WritableByteChannelFactory} used. */
- protected final WritableByteChannelFactory getWritableByteChannelFactory() {
- return writableByteChannelFactory;
+ /** Returns the extension that will be written to the produced files. */
+ protected final String getExtension() {
+ String extension = MoreObjects.firstNonNull(writableByteChannelFactory.getFilenameSuffix(), "");
+ if (!extension.isEmpty() && !extension.startsWith(".")) {
+ extension = "." + extension;
+ }
+ return extension;
}
/**
- * Abstract writer that writes a bundle to a {@link FileBasedSink}. Subclass implementations
- * provide a method that can write a single value to a {@link WritableByteChannel}.
+ * Abstract writer that writes a bundle to a {@link FileBasedSink}. Subclass
+ * implementations provide a method that can write a single value to a
+ * {@link WritableByteChannel}.
*
* <p>Subclass implementations may also override methods that write headers and footers before and
* after the values in a bundle, respectively, as well as provide a MIME type for the output
* channel.
*
- * <p>Multiple {@link Writer} instances may be created on the same worker, and therefore any
- * access to static members or methods should be thread safe.
+ * <p>Multiple {@link Writer} instances may be created on the same worker, and therefore
+ * any access to static members or methods should be thread safe.
*
- * @param <OutputT> the type of values to write.
+ * @param <T> the type of values to write.
*/
- public abstract static class Writer<OutputT, DestinationT> {
+ public abstract static class Writer<T> {
private static final Logger LOG = LoggerFactory.getLogger(Writer.class);
- private final WriteOperation<OutputT, DestinationT> writeOperation;
+ private final WriteOperation<T> writeOperation;
/** Unique id for this output bundle. */
private String id;
@@ -802,7 +753,6 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
private BoundedWindow window;
private PaneInfo paneInfo;
private int shard = -1;
- private DestinationT destination;
/** The output file for this bundle. May be null if opening failed. */
private @Nullable ResourceId outputFile;
@@ -822,8 +772,10 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
*/
private final String mimeType;
- /** Construct a new {@link Writer} that will produce files of the given MIME type. */
- public Writer(WriteOperation<OutputT, DestinationT> writeOperation, String mimeType) {
+ /**
+ * Construct a new {@link Writer} that will produce files of the given MIME type.
+ */
+ public Writer(WriteOperation<T> writeOperation, String mimeType) {
checkNotNull(writeOperation);
this.writeOperation = writeOperation;
this.mimeType = mimeType;
@@ -866,29 +818,28 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
* id populated for the case of static sharding. In cases where the runner is dynamically
* picking sharding, shard might be set to -1.
*/
- public final void openWindowed(
- String uId, BoundedWindow window, PaneInfo paneInfo, int shard, DestinationT destination)
+ public final void openWindowed(String uId, BoundedWindow window, PaneInfo paneInfo, int shard)
throws Exception {
if (!getWriteOperation().windowedWrites) {
throw new IllegalStateException("openWindowed called a non-windowed sink.");
}
- open(uId, window, paneInfo, shard, destination);
+ open(uId, window, paneInfo, shard);
}
/**
* Called for each value in the bundle.
*/
- public abstract void write(OutputT value) throws Exception;
+ public abstract void write(T value) throws Exception;
/**
- * Similar to {@link #openWindowed} however for the case where unwindowed writes were requested.
+ * Similar to {@link #openWindowed} however for the case where unwindowed writes were
+ * requested.
*/
- public final void openUnwindowed(String uId, int shard, DestinationT destination)
- throws Exception {
+ public final void openUnwindowed(String uId, int shard) throws Exception {
if (getWriteOperation().windowedWrites) {
throw new IllegalStateException("openUnwindowed called a windowed sink.");
}
- open(uId, null, null, shard, destination);
+ open(uId, null, null, shard);
}
// Helper function to close a channel, on exception cases.
@@ -904,18 +855,14 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
}
}
- private void open(
- String uId,
- @Nullable BoundedWindow window,
- @Nullable PaneInfo paneInfo,
- int shard,
- DestinationT destination)
- throws Exception {
+ private void open(String uId,
+ @Nullable BoundedWindow window,
+ @Nullable PaneInfo paneInfo,
+ int shard) throws Exception {
this.id = uId;
this.window = window;
this.paneInfo = paneInfo;
this.shard = shard;
- this.destination = destination;
ResourceId tempDirectory = getWriteOperation().tempDirectory.get();
outputFile = tempDirectory.resolve(id, StandardResolveOptions.RESOLVE_FILE);
verifyNotNull(
@@ -961,7 +908,7 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
}
/** Closes the channel and returns the bundle result. */
- public final FileResult<DestinationT> close() throws Exception {
+ public final FileResult close() throws Exception {
checkState(outputFile != null, "FileResult.close cannot be called with a null outputFile");
LOG.debug("Writing footer to {}.", outputFile);
@@ -991,41 +938,35 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
throw new IOException(String.format("Failed closing channel to %s", outputFile), e);
}
- FileResult<DestinationT> result =
- new FileResult<>(outputFile, shard, window, paneInfo, destination);
+ FileResult result = new FileResult(outputFile, shard, window, paneInfo);
LOG.debug("Result for bundle {}: {}", this.id, outputFile);
return result;
}
- /** Return the WriteOperation that this Writer belongs to. */
- public WriteOperation<OutputT, DestinationT> getWriteOperation() {
+ /**
+ * Return the WriteOperation that this Writer belongs to.
+ */
+ public WriteOperation<T> getWriteOperation() {
return writeOperation;
}
}
/**
- * Result of a single bundle write. Contains the filename produced by the bundle, and if known the
- * final output filename.
+ * Result of a single bundle write. Contains the filename produced by the bundle, and if known
+ * the final output filename.
*/
- public static final class FileResult<DestinationT> {
+ public static final class FileResult {
private final ResourceId tempFilename;
private final int shard;
private final BoundedWindow window;
private final PaneInfo paneInfo;
- private final DestinationT destination;
@Experimental(Kind.FILESYSTEM)
- public FileResult(
- ResourceId tempFilename,
- int shard,
- BoundedWindow window,
- PaneInfo paneInfo,
- DestinationT destination) {
+ public FileResult(ResourceId tempFilename, int shard, BoundedWindow window, PaneInfo paneInfo) {
this.tempFilename = tempFilename;
this.shard = shard;
this.window = window;
this.paneInfo = paneInfo;
- this.destination = destination;
}
@Experimental(Kind.FILESYSTEM)
@@ -1037,8 +978,8 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
return shard;
}
- public FileResult<DestinationT> withShard(int shard) {
- return new FileResult<>(tempFilename, shard, window, paneInfo, destination);
+ public FileResult withShard(int shard) {
+ return new FileResult(tempFilename, shard, window, paneInfo);
}
public BoundedWindow getWindow() {
@@ -1049,24 +990,17 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
return paneInfo;
}
- public DestinationT getDestination() {
- return destination;
- }
-
@Experimental(Kind.FILESYSTEM)
- public ResourceId getDestinationFile(
- DynamicDestinations<?, DestinationT> dynamicDestinations,
- int numShards,
- OutputFileHints outputFileHints) {
+ public ResourceId getDestinationFile(FilenamePolicy policy, ResourceId outputDirectory,
+ int numShards, String extension) {
checkArgument(getShard() != UNKNOWN_SHARDNUM);
checkArgument(numShards > 0);
- FilenamePolicy policy = dynamicDestinations.getFilenamePolicy(destination);
if (getWindow() != null) {
- return policy.windowedFilename(
- new WindowedContext(getWindow(), getPaneInfo(), getShard(), numShards),
- outputFileHints);
+ return policy.windowedFilename(outputDirectory, new WindowedContext(
+ getWindow(), getPaneInfo(), getShard(), numShards), extension);
} else {
- return policy.unwindowedFilename(new Context(getShard(), numShards), outputFileHints);
+ return policy.unwindowedFilename(outputDirectory, new Context(getShard(), numShards),
+ extension);
}
}
@@ -1080,24 +1014,22 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
}
}
- /** A coder for {@link FileResult} objects. */
- public static final class FileResultCoder<DestinationT>
- extends StructuredCoder<FileResult<DestinationT>> {
+ /**
+ * A coder for {@link FileResult} objects.
+ */
+ public static final class FileResultCoder extends StructuredCoder<FileResult> {
private static final Coder<String> FILENAME_CODER = StringUtf8Coder.of();
private static final Coder<Integer> SHARD_CODER = VarIntCoder.of();
private static final Coder<PaneInfo> PANE_INFO_CODER = NullableCoder.of(PaneInfoCoder.INSTANCE);
+
private final Coder<BoundedWindow> windowCoder;
- private final Coder<DestinationT> destinationCoder;
- protected FileResultCoder(
- Coder<BoundedWindow> windowCoder, Coder<DestinationT> destinationCoder) {
+ protected FileResultCoder(Coder<BoundedWindow> windowCoder) {
this.windowCoder = NullableCoder.of(windowCoder);
- this.destinationCoder = destinationCoder;
}
- public static <DestinationT> FileResultCoder<DestinationT> of(
- Coder<BoundedWindow> windowCoder, Coder<DestinationT> destinationCoder) {
- return new FileResultCoder<>(windowCoder, destinationCoder);
+ public static FileResultCoder of(Coder<BoundedWindow> windowCoder) {
+ return new FileResultCoder(windowCoder);
}
@Override
@@ -1106,7 +1038,8 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
}
@Override
- public void encode(FileResult<DestinationT> value, OutputStream outStream) throws IOException {
+ public void encode(FileResult value, OutputStream outStream)
+ throws IOException {
if (value == null) {
throw new CoderException("cannot encode a null value");
}
@@ -1114,22 +1047,17 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
windowCoder.encode(value.getWindow(), outStream);
PANE_INFO_CODER.encode(value.getPaneInfo(), outStream);
SHARD_CODER.encode(value.getShard(), outStream);
- destinationCoder.encode(value.getDestination(), outStream);
}
@Override
- public FileResult<DestinationT> decode(InputStream inStream) throws IOException {
+ public FileResult decode(InputStream inStream)
+ throws IOException {
String tempFilename = FILENAME_CODER.decode(inStream);
BoundedWindow window = windowCoder.decode(inStream);
PaneInfo paneInfo = PANE_INFO_CODER.decode(inStream);
int shard = SHARD_CODER.decode(inStream);
- DestinationT destination = destinationCoder.decode(inStream);
- return new FileResult<>(
- FileSystems.matchNewResource(tempFilename, false /* isDirectory */),
- shard,
- window,
- paneInfo,
- destination);
+ return new FileResult(FileSystems.matchNewResource(tempFilename, false /* isDirectory */),
+ shard, window, paneInfo);
}
@Override
@@ -1138,15 +1066,25 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
windowCoder.verifyDeterministic();
PANE_INFO_CODER.verifyDeterministic();
SHARD_CODER.verifyDeterministic();
- destinationCoder.verifyDeterministic();
}
}
/**
- * Provides hints about how to generate output files, such as a suggested filename suffix (e.g.
- * based on the compression type), and the file MIME type.
+ * Implementations create instances of {@link WritableByteChannel} used by {@link FileBasedSink}
+ * and related classes to allow <em>decorating</em>, or otherwise transforming, the raw data that
+ * would normally be written directly to the {@link WritableByteChannel} passed into
+ * {@link WritableByteChannelFactory#create(WritableByteChannel)}.
+ *
+ * <p>Subclasses should override {@link #toString()} with something meaningful, as it is used when
+ * building {@link DisplayData}.
*/
- public interface OutputFileHints extends Serializable {
+ public interface WritableByteChannelFactory extends Serializable {
+ /**
+ * @param channel the {@link WritableByteChannel} to wrap
+ * @return the {@link WritableByteChannel} to be used during output
+ */
+ WritableByteChannel create(WritableByteChannel channel) throws IOException;
+
/**
* Returns the MIME type that should be used for the files that will hold the output data. May
* return {@code null} if this {@code WritableByteChannelFactory} does not meaningfully change
@@ -1163,23 +1101,6 @@ public abstract class FileBasedSink<OutputT, DestinationT> implements Serializab
* @return an optional filename suffix, eg, ".gz" is returned by {@link CompressionType#GZIP}
*/
@Nullable
- String getSuggestedFilenameSuffix();
- }
-
- /**
- * Implementations create instances of {@link WritableByteChannel} used by {@link FileBasedSink}
- * and related classes to allow <em>decorating</em>, or otherwise transforming, the raw data that
- * would normally be written directly to the {@link WritableByteChannel} passed into {@link
- * WritableByteChannelFactory#create(WritableByteChannel)}.
- *
- * <p>Subclasses should override {@link #toString()} with something meaningful, as it is used when
- * building {@link DisplayData}.
- */
- public interface WritableByteChannelFactory extends OutputFileHints {
- /**
- * @param channel the {@link WritableByteChannel} to wrap
- * @return the {@link WritableByteChannel} to be used during output
- */
- WritableByteChannel create(WritableByteChannel channel) throws IOException;
+ String getFilenameSuffix();
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/OffsetBasedSource.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/OffsetBasedSource.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/OffsetBasedSource.java
index c3687a9..05f0d97 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/OffsetBasedSource.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/OffsetBasedSource.java
@@ -23,7 +23,6 @@ import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.NoSuchElementException;
-import org.apache.beam.sdk.io.range.OffsetRange;
import org.apache.beam.sdk.io.range.OffsetRangeTracker;
import org.apache.beam.sdk.io.range.RangeTracker;
import org.apache.beam.sdk.options.PipelineOptions;
@@ -111,7 +110,8 @@ public abstract class OffsetBasedSource<T> extends BoundedSource<T> {
@Override
public List<? extends OffsetBasedSource<T>> split(
long desiredBundleSizeBytes, PipelineOptions options) throws Exception {
- // Split the range into bundles based on the desiredBundleSizeBytes. If the desired bundle
+ // Split the range into bundles based on the desiredBundleSizeBytes. Final bundle is adjusted to
+ // make sure that we do not end up with a too small bundle at the end. If the desired bundle
// size is smaller than the minBundleSize of the source then minBundleSize will be used instead.
long desiredBundleSizeOffsetUnits = Math.max(
@@ -119,10 +119,20 @@ public abstract class OffsetBasedSource<T> extends BoundedSource<T> {
minBundleSize);
List<OffsetBasedSource<T>> subSources = new ArrayList<>();
- for (OffsetRange range :
- new OffsetRange(startOffset, Math.min(endOffset, getMaxEndOffset(options)))
- .split(desiredBundleSizeOffsetUnits, minBundleSize)) {
- subSources.add(createSourceForSubrange(range.getFrom(), range.getTo()));
+ long start = startOffset;
+ long maxEnd = Math.min(endOffset, getMaxEndOffset(options));
+
+ while (start < maxEnd) {
+ long end = start + desiredBundleSizeOffsetUnits;
+ end = Math.min(end, maxEnd);
+ // Avoid having a too small bundle at the end and ensure that we respect minBundleSize.
+ long remaining = maxEnd - end;
+ if ((remaining < desiredBundleSizeOffsetUnits / 4) || (remaining < minBundleSize)) {
+ end = maxEnd;
+ }
+ subSources.add(createSourceForSubrange(start, end));
+
+ start = end;
}
return subSources;
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TFRecordIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TFRecordIO.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TFRecordIO.java
index 6e7b243..e288075 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TFRecordIO.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TFRecordIO.java
@@ -45,7 +45,6 @@ import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.SerializableFunction;
-import org.apache.beam.sdk.transforms.SerializableFunctions;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.util.MimeTypes;
import org.apache.beam.sdk.values.PBegin;
@@ -356,11 +355,12 @@ public class TFRecordIO {
public PDone expand(PCollection<byte[]> input) {
checkState(getOutputPrefix() != null,
"need to set the output prefix of a TFRecordIO.Write transform");
- WriteFiles<byte[], Void, byte[]> write =
- WriteFiles.<byte[], Void, byte[]>to(
+ WriteFiles<byte[]> write = WriteFiles.to(
new TFRecordSink(
- getOutputPrefix(), getShardTemplate(), getFilenameSuffix(), getCompressionType()),
- SerializableFunctions.<byte[]>identity());
+ getOutputPrefix(),
+ getShardTemplate(),
+ getFilenameSuffix(),
+ getCompressionType()));
if (getNumShards() > 0) {
write = write.withNumShards(getNumShards());
}
@@ -546,20 +546,20 @@ public class TFRecordIO {
}
}
- /** A {@link FileBasedSink} for TFRecord files. Produces TFRecord files. */
+ /**
+ * A {@link FileBasedSink} for TFRecord files. Produces TFRecord files.
+ */
@VisibleForTesting
- static class TFRecordSink extends FileBasedSink<byte[], Void> {
+ static class TFRecordSink extends FileBasedSink<byte[]> {
@VisibleForTesting
- TFRecordSink(
- ValueProvider<ResourceId> outputPrefix,
+ TFRecordSink(ValueProvider<ResourceId> outputPrefix,
@Nullable String shardTemplate,
@Nullable String suffix,
TFRecordIO.CompressionType compressionType) {
super(
outputPrefix,
- DynamicFileDestinations.constant(
- DefaultFilenamePolicy.fromStandardParameters(
- outputPrefix, shardTemplate, suffix, false)),
+ DefaultFilenamePolicy.constructUsingStandardParameters(
+ outputPrefix, shardTemplate, suffix, false),
writableByteChannelFactory(compressionType));
}
@@ -571,7 +571,7 @@ public class TFRecordIO {
}
@Override
- public WriteOperation<byte[], Void> createWriteOperation() {
+ public WriteOperation<byte[]> createWriteOperation() {
return new TFRecordWriteOperation(this);
}
@@ -590,24 +590,30 @@ public class TFRecordIO {
return CompressionType.UNCOMPRESSED;
}
- /** A {@link WriteOperation WriteOperation} for TFRecord files. */
- private static class TFRecordWriteOperation extends WriteOperation<byte[], Void> {
+ /**
+ * A {@link WriteOperation
+ * WriteOperation} for TFRecord files.
+ */
+ private static class TFRecordWriteOperation extends WriteOperation<byte[]> {
private TFRecordWriteOperation(TFRecordSink sink) {
super(sink);
}
@Override
- public Writer<byte[], Void> createWriter() throws Exception {
+ public Writer<byte[]> createWriter() throws Exception {
return new TFRecordWriter(this);
}
}
- /** A {@link Writer Writer} for TFRecord files. */
- private static class TFRecordWriter extends Writer<byte[], Void> {
+ /**
+ * A {@link Writer Writer}
+ * for TFRecord files.
+ */
+ private static class TFRecordWriter extends Writer<byte[]> {
private WritableByteChannel outChannel;
private TFRecordCodec codec;
- private TFRecordWriter(WriteOperation<byte[], Void> writeOperation) {
+ private TFRecordWriter(WriteOperation<byte[]> writeOperation) {
super(writeOperation, MimeTypes.BINARY);
}
[19/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/spark/src/test/java/org/apache/beam/runners/spark/stateful/SparkStateInternalsTest.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/stateful/SparkStateInternalsTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/stateful/SparkStateInternalsTest.java
deleted file mode 100644
index b4597f9..0000000
--- a/runners/spark/src/test/java/org/apache/beam/runners/spark/stateful/SparkStateInternalsTest.java
+++ /dev/null
@@ -1,66 +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.spark.stateful;
-
-import org.apache.beam.runners.core.StateInternals;
-import org.apache.beam.runners.core.StateInternalsTest;
-import org.junit.Ignore;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-/**
- * Tests for {@link SparkStateInternals}. This is based on {@link StateInternalsTest}.
- * Ignore set and map tests.
- */
-@RunWith(JUnit4.class)
-public class SparkStateInternalsTest extends StateInternalsTest {
-
- @Override
- protected StateInternals createStateInternals() {
- return SparkStateInternals.forKey("dummyKey");
- }
-
- @Override
- @Ignore
- public void testSet() {}
-
- @Override
- @Ignore
- public void testSetIsEmpty() {}
-
- @Override
- @Ignore
- public void testMergeSetIntoSource() {}
-
- @Override
- @Ignore
- public void testMergeSetIntoNewNamespace() {}
-
- @Override
- @Ignore
- public void testMap() {}
-
- @Override
- @Ignore
- public void testSetReadable() {}
-
- @Override
- @Ignore
- public void testMapReadable() {}
-
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/StorageLevelTest.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/StorageLevelTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/StorageLevelTest.java
index 8bd6dae..8f2e681 100644
--- a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/StorageLevelTest.java
+++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/StorageLevelTest.java
@@ -58,12 +58,12 @@ public class StorageLevelTest {
@Test
public void test() throws Exception {
- PCollection<String> pCollection = pipeline.apply("CreateFoo", Create.of("foo"));
+ PCollection<String> pCollection = pipeline.apply(Create.of("foo"));
// by default, the Spark runner doesn't cache the RDD if it accessed only one time.
// So, to "force" the caching of the RDD, we have to call the RDD at least two time.
// That's why we are using Count fn on the PCollection.
- pCollection.apply("CountAll", Count.<String>globally());
+ pCollection.apply(Count.<String>globally());
PCollection<String> output = pCollection.apply(new StorageLevelPTransform());
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/common/fn-api/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/common/fn-api/pom.xml b/sdks/common/fn-api/pom.xml
index 6810667..77a9ba5 100644
--- a/sdks/common/fn-api/pom.xml
+++ b/sdks/common/fn-api/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-common-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/common/fn-api/src/main/proto/beam_fn_api.proto
----------------------------------------------------------------------
diff --git a/sdks/common/fn-api/src/main/proto/beam_fn_api.proto b/sdks/common/fn-api/src/main/proto/beam_fn_api.proto
index 9da5afe..9fe2b2f 100644
--- a/sdks/common/fn-api/src/main/proto/beam_fn_api.proto
+++ b/sdks/common/fn-api/src/main/proto/beam_fn_api.proto
@@ -38,6 +38,7 @@ option java_package = "org.apache.beam.fn.v1";
option java_outer_classname = "BeamFnApi";
import "beam_runner_api.proto";
+import "google/protobuf/any.proto";
import "google/protobuf/timestamp.proto";
/*
@@ -66,6 +67,122 @@ message Target {
string name = 2;
}
+// Information defining a PCollection
+message PCollection {
+ // (Required) A reference to a coder.
+ string coder_reference = 1;
+
+ // TODO: Windowing strategy, ...
+}
+
+// A primitive transform within Apache Beam.
+message PrimitiveTransform {
+ // (Required) A pipeline level unique id which can be used as a reference to
+ // refer to this.
+ string id = 1;
+
+ // (Required) A function spec that is used by this primitive
+ // transform to process data.
+ FunctionSpec function_spec = 2;
+
+ // A map of distinct input names to target definitions.
+ // For example, in CoGbk this represents the tag name associated with each
+ // distinct input name and a list of primitive transforms that are associated
+ // with the specified input.
+ map<string, Target.List> inputs = 3;
+
+ // A map from local output name to PCollection definitions. For example, in
+ // DoFn this represents the tag name associated with each distinct output.
+ map<string, PCollection> outputs = 4;
+
+ // TODO: Should we model side inputs as a special type of input for a
+ // primitive transform or should it be modeled as the relationship that
+ // the predecessor input will be a view primitive transform.
+ // A map of from side input names to side inputs.
+ map<string, SideInput> side_inputs = 5;
+
+ // The user name of this step.
+ // TODO: This should really be in display data and not at this level
+ string step_name = 6;
+}
+
+/*
+ * User Definable Functions
+ *
+ * This is still unstable mainly due to how we model the side input.
+ */
+
+// Defines the common elements of user-definable functions, to allow the SDK to
+// express the information the runner needs to execute work.
+// Stable
+message FunctionSpec {
+ // (Required) A pipeline level unique id which can be used as a reference to
+ // refer to this.
+ string id = 1;
+
+ // (Required) A globally unique name representing this user definable
+ // function.
+ //
+ // User definable functions use the urn encodings registered such that another
+ // may implement the user definable function within another language.
+ //
+ // For example:
+ // urn:org.apache.beam:coder:kv:1.0
+ string urn = 2;
+
+ // (Required) Reference to specification of execution environment required to
+ // invoke this function.
+ string environment_reference = 3;
+
+ // Data used to parameterize this function. Depending on the urn, this may be
+ // optional or required.
+ google.protobuf.Any data = 4;
+}
+
+message SideInput {
+ // TODO: Coder?
+
+ // For RunnerAPI.
+ Target input = 1;
+
+ // For FnAPI.
+ FunctionSpec view_fn = 2;
+}
+
+// Defines how to encode values into byte streams and decode values from byte
+// streams. A coder can be parameterized by additional properties which may or
+// may not be language agnostic.
+//
+// Coders using the urn:org.apache.beam:coder namespace must have their
+// encodings registered such that another may implement the encoding within
+// another language.
+//
+// For example:
+// urn:org.apache.beam:coder:kv:1.0
+// urn:org.apache.beam:coder:iterable:1.0
+// Stable
+message Coder {
+ // TODO: This looks weird when compared to the other function specs
+ // which use URN to differentiate themselves. Should "Coder" be embedded
+ // inside the FunctionSpec data block.
+
+ // The data associated with this coder used to reconstruct it.
+ FunctionSpec function_spec = 1;
+
+ // A list of component coder references.
+ //
+ // For a key-value coder, there must be exactly two component coder references
+ // where the first reference represents the key coder and the second reference
+ // is the value coder.
+ //
+ // For an iterable coder, there must be exactly one component coder reference
+ // representing the value coder.
+ //
+ // TODO: Perhaps this is redundant with the data of the FunctionSpec
+ // for known coders?
+ repeated string component_coder_reference = 2;
+}
+
// A descriptor for connecting to a remote port using the Beam Fn Data API.
// Allows for communication between two environments (for example between the
// runner and the SDK).
@@ -154,20 +271,29 @@ message ProcessBundleDescriptor {
// refer to this.
string id = 1;
+ // (Deprecated) A list of primitive transforms that should
+ // be used to construct the bundle processing graph.
+ repeated PrimitiveTransform primitive_transform = 2;
+
+ // (Deprecated) The set of all coders referenced in this bundle.
+ repeated Coder coders = 4;
+
// (Required) A map from pipeline-scoped id to PTransform.
- map<string, org.apache.beam.runner_api.v1.PTransform> transforms = 2;
+ map<string, org.apache.beam.runner_api.v1.PTransform> transforms = 5;
// (Required) A map from pipeline-scoped id to PCollection.
- map<string, org.apache.beam.runner_api.v1.PCollection> pcollections = 3;
+ map<string, org.apache.beam.runner_api.v1.PCollection> pcollections = 6;
// (Required) A map from pipeline-scoped id to WindowingStrategy.
- map<string, org.apache.beam.runner_api.v1.WindowingStrategy> windowing_strategies = 4;
+ map<string, org.apache.beam.runner_api.v1.WindowingStrategy> windowing_strategies = 7;
// (Required) A map from pipeline-scoped id to Coder.
- map<string, org.apache.beam.runner_api.v1.Coder> coders = 5;
+ // TODO: Rename to "coders" once deprecated coders field is removed. Unique
+ // name is choosen to make it an easy search/replace
+ map<string, org.apache.beam.runner_api.v1.Coder> codersyyy = 8;
// (Required) A map from pipeline-scoped id to Environment.
- map<string, org.apache.beam.runner_api.v1.Environment> environments = 6;
+ map<string, org.apache.beam.runner_api.v1.Environment> environments = 9;
}
// A request to process a given bundle.
@@ -177,9 +303,9 @@ message ProcessBundleRequest {
// instantiated and executed by the SDK harness.
string process_bundle_descriptor_reference = 1;
- // (Optional) A list of cache tokens that can be used by an SDK to reuse
- // cached data returned by the State API across multiple bundles.
- repeated bytes cache_tokens = 2;
+ // (Optional) A list of cache tokens that can be used by an SDK to cache
+ // data looked up using the State API across multiple bundles.
+ repeated CacheToken cache_tokens = 2;
}
// Stable
@@ -248,14 +374,14 @@ message PrimitiveTransformSplit {
//
// For example, a remote GRPC source will have a specific urn and data
// block containing an ElementCountRestriction.
- org.apache.beam.runner_api.v1.FunctionSpec completed_restriction = 2;
+ FunctionSpec completed_restriction = 2;
// (Required) A function specification describing the restriction
// representing the remainder of work for the primitive transform.
//
// FOr example, a remote GRPC source will have a specific urn and data
// block contain an ElemntCountSkipRestriction.
- org.apache.beam.runner_api.v1.FunctionSpec remaining_restriction = 3;
+ FunctionSpec remaining_restriction = 3;
}
message ProcessBundleSplitResponse {
@@ -402,10 +528,6 @@ message StateResponse {
// failed.
string error = 2;
- // (Optional) If this is specified, then the result of this state request
- // can be cached using the supplied token.
- bytes cache_token = 3;
-
// A corresponding response matching the request will be populated.
oneof response {
// A response to getting state.
@@ -431,44 +553,49 @@ service BeamFnState {
) {}
}
+message CacheToken {
+ // (Required) Represents the function spec and tag associated with this state
+ // key.
+ //
+ // By combining the function_spec_reference with the tag representing:
+ // * the input, we refer to the iterable portion of a large GBK
+ // * the side input, we refer to the side input
+ // * the user state, we refer to user state
+ Target target = 1;
+
+ // (Required) An opaque identifier.
+ bytes token = 2;
+}
+
message StateKey {
- message Runner {
- // (Required) Opaque information supplied by the runner. Used to support
- // remote references.
- bytes key = 1;
- }
+ // (Required) Represents the function spec and tag associated with this state
+ // key.
+ //
+ // By combining the function_spec_reference with the tag representing:
+ // * the input, we refer to fetching the iterable portion of a large GBK
+ // * the side input, we refer to fetching the side input
+ // * the user state, we refer to fetching user state
+ Target target = 1;
- message MultimapSideInput {
- // (Required) The id of the PTransform containing a side input.
- string ptransform_id = 1;
- // (Required) The id of the side input.
- string side_input_id = 2;
- // (Required) The window (after mapping the currently executing elements
- // window into the side input windows domain) encoded in a nested context.
- bytes window = 3;
- // (Required) The key encoded in a nested context.
- bytes key = 4;
- }
+ // (Required) The bytes of the window which this state request is for encoded
+ // in the nested context.
+ bytes window = 2;
- message BagUserState {
- // (Required) The id of the PTransform containing user state.
- string ptransform_id = 1;
- // (Required) The id of the user state.
- string user_state_id = 2;
- // (Required) The window encoded in a nested context.
- bytes window = 3;
- // (Required) The key of the currently executing element encoded in a
- // nested context.
- bytes key = 4;
- }
+ // (Required) The user key encoded in the nested context.
+ bytes key = 3;
+}
- // (Required) One of the following state keys must be set.
- oneof type {
- Runner runner = 1;
- MultimapSideInput multimap_side_input = 2;
- BagUserState bag_user_state = 3;
- // TODO: represent a state key for user map state
- }
+// A logical byte stream which can be continued using the state API.
+message ContinuableStream {
+ // (Optional) If specified, represents a token which can be used with the
+ // state API to get the next chunk of this logical byte stream. The end of
+ // the logical byte stream is signalled by this field being unset.
+ bytes continuation_token = 1;
+
+ // Represents a part of a logical byte stream. Elements within
+ // the logical byte stream are encoded in the nested context and
+ // concatenated together.
+ bytes data = 2;
}
// A request to get state.
@@ -481,18 +608,10 @@ message StateGetRequest {
bytes continuation_token = 1;
}
-// A response to get state representing a logical byte stream which can be
-// continued using the state API.
+// A response to get state.
message StateGetResponse {
- // (Optional) If specified, represents a token which can be used with the
- // state API to get the next chunk of this logical byte stream. The end of
- // the logical byte stream is signalled by this field being unset.
- bytes continuation_token = 1;
-
- // Represents a part of a logical byte stream. Elements within
- // the logical byte stream are encoded in the nested context and
- // concatenated together.
- bytes data = 2;
+ // (Required) The response containing a continuable logical byte stream.
+ ContinuableStream stream = 1;
}
// A request to append state.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/common/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/common/pom.xml b/sdks/common/pom.xml
index 40eefa7..c621ed5 100644
--- a/sdks/common/pom.xml
+++ b/sdks/common/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/common/runner-api/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/common/runner-api/pom.xml b/sdks/common/runner-api/pom.xml
index 8bc4123..f5536a7 100644
--- a/sdks/common/runner-api/pom.xml
+++ b/sdks/common/runner-api/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-common-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/common/runner-api/src/main/proto/beam_runner_api.proto
----------------------------------------------------------------------
diff --git a/sdks/common/runner-api/src/main/proto/beam_runner_api.proto b/sdks/common/runner-api/src/main/proto/beam_runner_api.proto
index 711da2a..039ecb0 100644
--- a/sdks/common/runner-api/src/main/proto/beam_runner_api.proto
+++ b/sdks/common/runner-api/src/main/proto/beam_runner_api.proto
@@ -367,12 +367,9 @@ message WriteFilesPayload {
// (Required) The SdkFunctionSpec of the FileBasedSink.
SdkFunctionSpec sink = 1;
- // (Required) The format function.
- SdkFunctionSpec format_function = 2;
+ bool windowed_writes = 2;
- bool windowed_writes = 3;
-
- bool runner_determined_sharding = 4;
+ bool runner_determined_sharding = 3;
}
// A coder, the binary format for serialization and deserialization of data in
@@ -436,14 +433,6 @@ message WindowingStrategy {
// (Required) The duration, in milliseconds, beyond the end of a window at
// which the window becomes droppable.
int64 allowed_lateness = 8;
-
- // (Required) Indicate whether empty on-time panes should be omitted.
- OnTimeBehavior OnTimeBehavior = 9;
-
- // (Required) Whether or not the window fn assigns inputs to exactly one window
- //
- // This knowledge is required for some optimizations
- bool assigns_to_one_window = 10;
}
// Whether or not a PCollection's WindowFn is non-merging, merging, or
@@ -489,17 +478,6 @@ enum ClosingBehavior {
EMIT_IF_NONEMPTY = 1;
}
-// Controls whether or not an aggregating transform should output data
-// when an on-time pane is empty.
-enum OnTimeBehavior {
- // Always fire the on-time pane. Even if there is no new data since
- // the previous firing, an element will be produced.
- FIRE_ALWAYS = 0;
-
- // Only fire the on-time pane if there is new data since the previous firing.
- FIRE_IF_NONEMPTY = 1;
-}
-
// When a number of windowed, timestamped inputs are aggregated, the timestamp
// for the resulting output.
enum OutputTime {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/build-tools/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/build-tools/pom.xml b/sdks/java/build-tools/pom.xml
index d7d25f6..5a2c498 100644
--- a/sdks/java/build-tools/pom.xml
+++ b/sdks/java/build-tools/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../../../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/build-tools/src/main/resources/beam/findbugs-filter.xml
----------------------------------------------------------------------
diff --git a/sdks/java/build-tools/src/main/resources/beam/findbugs-filter.xml b/sdks/java/build-tools/src/main/resources/beam/findbugs-filter.xml
index 0c9080d..3430750 100644
--- a/sdks/java/build-tools/src/main/resources/beam/findbugs-filter.xml
+++ b/sdks/java/build-tools/src/main/resources/beam/findbugs-filter.xml
@@ -412,13 +412,4 @@
<!-- PTransforms do not actually support serialization. -->
</Match>
- <Match>
- <Class name="org.apache.beam.sdk.options.ProxyInvocationHandler"/>
- <Field name="~.*"/>
- <Bug pattern="SE_BAD_FIELD"/>
- <!--
- ProxyInvocationHandler implements Serializable only for the sake of throwing an informative
- exception in writeObject()
- -->
- </Match>
</FindBugsFilter>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/core/pom.xml b/sdks/java/core/pom.xml
index 3f12dc4..11b68e6 100644
--- a/sdks/java/core/pom.xml
+++ b/sdks/java/core/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/ShardedKeyCoder.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/ShardedKeyCoder.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/ShardedKeyCoder.java
deleted file mode 100644
index a86b198..0000000
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/ShardedKeyCoder.java
+++ /dev/null
@@ -1,66 +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.sdk.coders;
-
-import com.google.common.annotations.VisibleForTesting;
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
-import java.util.Arrays;
-import java.util.List;
-import org.apache.beam.sdk.values.ShardedKey;
-
-/** A {@link Coder} for {@link ShardedKey}, using a wrapped key {@link Coder}. */
-@VisibleForTesting
-public class ShardedKeyCoder<KeyT> extends StructuredCoder<ShardedKey<KeyT>> {
- public static <KeyT> ShardedKeyCoder<KeyT> of(Coder<KeyT> keyCoder) {
- return new ShardedKeyCoder<>(keyCoder);
- }
-
- private final Coder<KeyT> keyCoder;
- private final VarIntCoder shardNumberCoder;
-
- protected ShardedKeyCoder(Coder<KeyT> keyCoder) {
- this.keyCoder = keyCoder;
- this.shardNumberCoder = VarIntCoder.of();
- }
-
- @Override
- public List<? extends Coder<?>> getCoderArguments() {
- return Arrays.asList(keyCoder);
- }
-
- @Override
- public void encode(ShardedKey<KeyT> key, OutputStream outStream)
- throws IOException {
- keyCoder.encode(key.getKey(), outStream);
- shardNumberCoder.encode(key.getShardNumber(), outStream);
- }
-
- @Override
- public ShardedKey<KeyT> decode(InputStream inStream)
- throws IOException {
- return ShardedKey.of(keyCoder.decode(inStream), shardNumberCoder.decode(inStream));
- }
-
- @Override
- public void verifyDeterministic() throws NonDeterministicException {
- keyCoder.verifyDeterministic();
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/AvroIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/AvroIO.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/AvroIO.java
index 89cadbd..4143db2 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/AvroIO.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/AvroIO.java
@@ -18,6 +18,7 @@
package org.apache.beam.sdk.io;
import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkState;
import com.google.auto.value.AutoValue;
import com.google.common.collect.ImmutableMap;
@@ -34,7 +35,6 @@ import org.apache.beam.sdk.annotations.Experimental.Kind;
import org.apache.beam.sdk.coders.AvroCoder;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.VoidCoder;
-import org.apache.beam.sdk.io.FileBasedSink.DynamicDestinations;
import org.apache.beam.sdk.io.FileBasedSink.FilenamePolicy;
import org.apache.beam.sdk.io.Read.Bounded;
import org.apache.beam.sdk.io.fs.ResourceId;
@@ -43,7 +43,6 @@ import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider;
import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.SerializableFunction;
-import org.apache.beam.sdk.transforms.SerializableFunctions;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.display.HasDisplayData;
import org.apache.beam.sdk.values.PBegin;
@@ -53,19 +52,18 @@ import org.apache.beam.sdk.values.PDone;
/**
* {@link PTransform}s for reading and writing Avro files.
*
- * <p>To read a {@link PCollection} from one or more Avro files, use {@code AvroIO.read()}, using
- * {@link AvroIO.Read#from} to specify the filename or filepattern to read from. See {@link
- * FileSystems} for information on supported file systems and filepatterns.
+ * <p>To read a {@link PCollection} from one or more Avro files, use {@code AvroIO.read()},
+ * using {@link AvroIO.Read#from} to specify the filename or filepattern to read from.
+ * See {@link FileSystems} for information on supported file systems and filepatterns.
*
- * <p>To read specific records, such as Avro-generated classes, use {@link #read(Class)}. To read
- * {@link GenericRecord GenericRecords}, use {@link #readGenericRecords(Schema)} which takes a
- * {@link Schema} object, or {@link #readGenericRecords(String)} which takes an Avro schema in a
+ * <p>To read specific records, such as Avro-generated classes, use {@link #read(Class)}.
+ * To read {@link GenericRecord GenericRecords}, use {@link #readGenericRecords(Schema)} which takes
+ * a {@link Schema} object, or {@link #readGenericRecords(String)} which takes an Avro schema in a
* JSON-encoded string form. An exception will be thrown if a record doesn't match the specified
* schema.
*
* <p>For example:
- *
- * <pre>{@code
+ * <pre> {@code
* Pipeline p = ...;
*
* // A simple Read of a local file (only runs locally):
@@ -77,33 +75,34 @@ import org.apache.beam.sdk.values.PDone;
* PCollection<GenericRecord> records =
* p.apply(AvroIO.readGenericRecords(schema)
* .from("gs://my_bucket/path/to/records-*.avro"));
- * }</pre>
+ * } </pre>
*
* <p>To write a {@link PCollection} to one or more Avro files, use {@link AvroIO.Write}, using
- * {@code AvroIO.write().to(String)} to specify the output filename prefix. The default {@link
- * DefaultFilenamePolicy} will use this prefix, in conjunction with a {@link ShardNameTemplate} (set
- * via {@link Write#withShardNameTemplate(String)}) and optional filename suffix (set via {@link
- * Write#withSuffix(String)}, to generate output filenames in a sharded way. You can override this
- * default write filename policy using {@link Write#to(FileBasedSink.FilenamePolicy)} to specify a
- * custom file naming policy.
+ * {@code AvroIO.write().to(String)} to specify the output filename prefix. The default
+ * {@link DefaultFilenamePolicy} will use this prefix, in conjunction with a
+ * {@link ShardNameTemplate} (set via {@link Write#withShardNameTemplate(String)}) and optional
+ * filename suffix (set via {@link Write#withSuffix(String)}, to generate output filenames in a
+ * sharded way. You can override this default write filename policy using
+ * {@link Write#withFilenamePolicy(FileBasedSink.FilenamePolicy)} to specify a custom file naming
+ * policy.
*
* <p>By default, all input is put into the global window before writing. If per-window writes are
- * desired - for example, when using a streaming runner - {@link AvroIO.Write#withWindowedWrites()}
- * will cause windowing and triggering to be preserved. When producing windowed writes with a
- * streaming runner that supports triggers, the number of output shards must be set explicitly using
- * {@link AvroIO.Write#withNumShards(int)}; some runners may set this for you to a runner-chosen
- * value, so you may need not set it yourself. A {@link FileBasedSink.FilenamePolicy} must be set,
- * and unique windows and triggers must produce unique filenames.
+ * desired - for example, when using a streaming runner -
+ * {@link AvroIO.Write#withWindowedWrites()} will cause windowing and triggering to be
+ * preserved. When producing windowed writes, the number of output shards must be set explicitly
+ * using {@link AvroIO.Write#withNumShards(int)}; some runners may set this for you to a
+ * runner-chosen value, so you may need not set it yourself. A
+ * {@link FileBasedSink.FilenamePolicy} must be set, and unique windows and triggers must produce
+ * unique filenames.
*
- * <p>To write specific records, such as Avro-generated classes, use {@link #write(Class)}. To write
- * {@link GenericRecord GenericRecords}, use either {@link #writeGenericRecords(Schema)} which takes
- * a {@link Schema} object, or {@link #writeGenericRecords(String)} which takes a schema in a
- * JSON-encoded string form. An exception will be thrown if a record doesn't match the specified
- * schema.
+ * <p>To write specific records, such as Avro-generated classes, use {@link #write(Class)}.
+ * To write {@link GenericRecord GenericRecords}, use either {@link #writeGenericRecords(Schema)}
+ * which takes a {@link Schema} object, or {@link #writeGenericRecords(String)} which takes a schema
+ * in a JSON-encoded string form. An exception will be thrown if a record doesn't match the
+ * specified schema.
*
* <p>For example:
- *
- * <pre>{@code
+ * <pre> {@code
* // A simple Write to a local file (only runs locally):
* PCollection<AvroAutoGenClass> records = ...;
* records.apply(AvroIO.write(AvroAutoGenClass.class).to("/path/to/file.avro"));
@@ -114,11 +113,11 @@ import org.apache.beam.sdk.values.PDone;
* records.apply("WriteToAvro", AvroIO.writeGenericRecords(schema)
* .to("gs://my_bucket/path/to/numbers")
* .withSuffix(".avro"));
- * }</pre>
+ * } </pre>
*
- * <p>By default, {@link AvroIO.Write} produces output files that are compressed using the {@link
- * org.apache.avro.file.Codec CodecFactory.deflateCodec(6)}. This default can be changed or
- * overridden using {@link AvroIO.Write#withCodec}.
+ * <p>By default, {@link AvroIO.Write} produces output files that are compressed using the
+ * {@link org.apache.avro.file.Codec CodecFactory.deflateCodec(6)}. This default can
+ * be changed or overridden using {@link AvroIO.Write#withCodec}.
*/
public class AvroIO {
/**
@@ -259,16 +258,11 @@ public class AvroIO {
@Nullable abstract ValueProvider<ResourceId> getFilenamePrefix();
@Nullable abstract String getShardTemplate();
@Nullable abstract String getFilenameSuffix();
-
- @Nullable
- abstract ValueProvider<ResourceId> getTempDirectory();
-
abstract int getNumShards();
@Nullable abstract Class<T> getRecordClass();
@Nullable abstract Schema getSchema();
abstract boolean getWindowedWrites();
@Nullable abstract FilenamePolicy getFilenamePolicy();
-
/**
* The codec used to encode the blocks in the Avro file. String value drawn from those in
* https://avro.apache.org/docs/1.7.7/api/java/org/apache/avro/file/CodecFactory.html
@@ -283,9 +277,6 @@ public class AvroIO {
abstract static class Builder<T> {
abstract Builder<T> setFilenamePrefix(ValueProvider<ResourceId> filenamePrefix);
abstract Builder<T> setFilenameSuffix(String filenameSuffix);
-
- abstract Builder<T> setTempDirectory(ValueProvider<ResourceId> tempDirectory);
-
abstract Builder<T> setNumShards(int numShards);
abstract Builder<T> setShardTemplate(String shardTemplate);
abstract Builder<T> setRecordClass(Class<T> recordClass);
@@ -305,9 +296,9 @@ public class AvroIO {
* <p>The name of the output files will be determined by the {@link FilenamePolicy} used.
*
* <p>By default, a {@link DefaultFilenamePolicy} will build output filenames using the
- * specified prefix, a shard name template (see {@link #withShardNameTemplate(String)}, and a
- * common suffix (if supplied using {@link #withSuffix(String)}). This default can be overridden
- * using {@link #to(FilenamePolicy)}.
+ * specified prefix, a shard name template (see {@link #withShardNameTemplate(String)}, and
+ * a common suffix (if supplied using {@link #withSuffix(String)}). This default can be
+ * overridden using {@link #withFilenamePolicy(FilenamePolicy)}.
*/
public Write<T> to(String outputPrefix) {
return to(FileBasedSink.convertToFileResourceIfPossible(outputPrefix));
@@ -315,21 +306,14 @@ public class AvroIO {
/**
* Writes to file(s) with the given output prefix. See {@link FileSystems} for information on
- * supported file systems. This prefix is used by the {@link DefaultFilenamePolicy} to generate
- * filenames.
- *
- * <p>By default, a {@link DefaultFilenamePolicy} will build output filenames using the
- * specified prefix, a shard name template (see {@link #withShardNameTemplate(String)}, and a
- * common suffix (if supplied using {@link #withSuffix(String)}). This default can be overridden
- * using {@link #to(FilenamePolicy)}.
+ * supported file systems.
*
- * <p>This default policy can be overridden using {@link #to(FilenamePolicy)}, in which case
- * {@link #withShardNameTemplate(String)} and {@link #withSuffix(String)} should not be set.
- * Custom filename policies do not automatically see this prefix - you should explicitly pass
- * the prefix into your {@link FilenamePolicy} object if you need this.
+ * <p>The name of the output files will be determined by the {@link FilenamePolicy} used.
*
- * <p>If {@link #withTempDirectory} has not been called, this filename prefix will be used to
- * infer a directory for temporary files.
+ * <p>By default, a {@link DefaultFilenamePolicy} will build output filenames using the
+ * specified prefix, a shard name template (see {@link #withShardNameTemplate(String)}, and
+ * a common suffix (if supplied using {@link #withSuffix(String)}). This default can be
+ * overridden using {@link #withFilenamePolicy(FilenamePolicy)}.
*/
@Experimental(Kind.FILESYSTEM)
public Write<T> to(ResourceId outputPrefix) {
@@ -358,22 +342,15 @@ public class AvroIO {
}
/**
- * Writes to files named according to the given {@link FileBasedSink.FilenamePolicy}. A
- * directory for temporary files must be specified using {@link #withTempDirectory}.
+ * Configures the {@link FileBasedSink.FilenamePolicy} that will be used to name written files.
*/
- public Write<T> to(FilenamePolicy filenamePolicy) {
+ public Write<T> withFilenamePolicy(FilenamePolicy filenamePolicy) {
return toBuilder().setFilenamePolicy(filenamePolicy).build();
}
- /** Set the base directory used to generate temporary files. */
- @Experimental(Kind.FILESYSTEM)
- public Write<T> withTempDirectory(ValueProvider<ResourceId> tempDirectory) {
- return toBuilder().setTempDirectory(tempDirectory).build();
- }
-
/**
* Uses the given {@link ShardNameTemplate} for naming output files. This option may only be
- * used when using one of the default filename-prefix to() overrides.
+ * used when {@link #withFilenamePolicy(FilenamePolicy)} has not been configured.
*
* <p>See {@link DefaultFilenamePolicy} for how the prefix, shard name template, and suffix are
* used.
@@ -383,8 +360,8 @@ public class AvroIO {
}
/**
- * Configures the filename suffix for written files. This option may only be used when using one
- * of the default filename-prefix to() overrides.
+ * Configures the filename suffix for written files. This option may only be used when
+ * {@link #withFilenamePolicy(FilenamePolicy)} has not been configured.
*
* <p>See {@link DefaultFilenamePolicy} for how the prefix, shard name template, and suffix are
* used.
@@ -425,8 +402,9 @@ public class AvroIO {
/**
* Preserves windowing of input elements and writes them to files based on the element's window.
*
- * <p>If using {@link #to(FileBasedSink.FilenamePolicy)}. Filenames will be generated using
- * {@link FilenamePolicy#windowedFilename}. See also {@link WriteFiles#withWindowedWrites()}.
+ * <p>Requires use of {@link #withFilenamePolicy(FileBasedSink.FilenamePolicy)}. Filenames will
+ * be generated using {@link FilenamePolicy#windowedFilename}. See also
+ * {@link WriteFiles#withWindowedWrites()}.
*/
public Write<T> withWindowedWrites() {
return toBuilder().setWindowedWrites(true).build();
@@ -457,46 +435,32 @@ public class AvroIO {
return toBuilder().setMetadata(ImmutableMap.copyOf(metadata)).build();
}
- DynamicDestinations<T, Void> resolveDynamicDestinations() {
- FilenamePolicy usedFilenamePolicy = getFilenamePolicy();
- if (usedFilenamePolicy == null) {
- usedFilenamePolicy =
- DefaultFilenamePolicy.fromStandardParameters(
- getFilenamePrefix(), getShardTemplate(), getFilenameSuffix(), getWindowedWrites());
- }
- return DynamicFileDestinations.constant(usedFilenamePolicy);
- }
-
@Override
public PDone expand(PCollection<T> input) {
- checkArgument(
- getFilenamePrefix() != null || getTempDirectory() != null,
- "Need to set either the filename prefix or the tempDirectory of a AvroIO.Write "
- + "transform.");
- if (getFilenamePolicy() != null) {
- checkArgument(
- getShardTemplate() == null && getFilenameSuffix() == null,
- "shardTemplate and filenameSuffix should only be used with the default "
- + "filename policy");
- }
- return expandTyped(input, resolveDynamicDestinations());
- }
+ checkState(getFilenamePrefix() != null,
+ "Need to set the filename prefix of an AvroIO.Write transform.");
+ checkState(
+ (getFilenamePolicy() == null)
+ || (getShardTemplate() == null && getFilenameSuffix() == null),
+ "Cannot set a filename policy and also a filename template or suffix.");
+ checkState(getSchema() != null,
+ "Need to set the schema of an AvroIO.Write transform.");
+ checkState(!getWindowedWrites() || (getFilenamePolicy() != null),
+ "When using windowed writes, a filename policy must be set via withFilenamePolicy().");
- public <DestinationT> PDone expandTyped(
- PCollection<T> input, DynamicDestinations<T, DestinationT> dynamicDestinations) {
- ValueProvider<ResourceId> tempDirectory = getTempDirectory();
- if (tempDirectory == null) {
- tempDirectory = getFilenamePrefix();
+ FilenamePolicy usedFilenamePolicy = getFilenamePolicy();
+ if (usedFilenamePolicy == null) {
+ usedFilenamePolicy = DefaultFilenamePolicy.constructUsingStandardParameters(
+ getFilenamePrefix(), getShardTemplate(), getFilenameSuffix(), getWindowedWrites());
}
- WriteFiles<T, DestinationT, T> write =
- WriteFiles.to(
- new AvroSink<>(
- tempDirectory,
- dynamicDestinations,
- AvroCoder.of(getRecordClass(), getSchema()),
- getCodec(),
- getMetadata()),
- SerializableFunctions.<T>identity());
+
+ WriteFiles<T> write = WriteFiles.to(
+ new AvroSink<>(
+ getFilenamePrefix(),
+ usedFilenamePolicy,
+ AvroCoder.of(getRecordClass(), getSchema()),
+ getCodec(),
+ getMetadata()));
if (getNumShards() > 0) {
write = write.withNumShards(getNumShards());
}
@@ -509,25 +473,31 @@ public class AvroIO {
@Override
public void populateDisplayData(DisplayData.Builder builder) {
super.populateDisplayData(builder);
- resolveDynamicDestinations().populateDisplayData(builder);
-
- String tempDirectory = null;
- if (getTempDirectory() != null) {
- tempDirectory =
- getTempDirectory().isAccessible()
- ? getTempDirectory().get().toString()
- : getTempDirectory().toString();
+ checkState(
+ getFilenamePrefix() != null,
+ "Unable to populate DisplayData for invalid AvroIO.Write (unset output prefix).");
+ String outputPrefixString = null;
+ if (getFilenamePrefix().isAccessible()) {
+ ResourceId dir = getFilenamePrefix().get();
+ outputPrefixString = dir.toString();
+ } else {
+ outputPrefixString = getFilenamePrefix().toString();
}
builder
- .add(DisplayData.item("schema", getRecordClass()).withLabel("Record Schema"))
- .addIfNotDefault(
- DisplayData.item("numShards", getNumShards()).withLabel("Maximum Output Shards"), 0)
- .addIfNotDefault(
- DisplayData.item("codec", getCodec().toString()).withLabel("Avro Compression Codec"),
- DEFAULT_CODEC.toString())
- .addIfNotNull(
- DisplayData.item("tempDirectory", tempDirectory)
- .withLabel("Directory for temporary files"));
+ .add(DisplayData.item("schema", getRecordClass())
+ .withLabel("Record Schema"))
+ .addIfNotNull(DisplayData.item("filePrefix", outputPrefixString)
+ .withLabel("Output File Prefix"))
+ .addIfNotNull(DisplayData.item("shardNameTemplate", getShardTemplate())
+ .withLabel("Output Shard Name Template"))
+ .addIfNotNull(DisplayData.item("fileSuffix", getFilenameSuffix())
+ .withLabel("Output File Suffix"))
+ .addIfNotDefault(DisplayData.item("numShards", getNumShards())
+ .withLabel("Maximum Output Shards"),
+ 0)
+ .addIfNotDefault(DisplayData.item("codec", getCodec().toString())
+ .withLabel("Avro Compression Codec"),
+ DEFAULT_CODEC.toString());
builder.include("Metadata", new Metadata());
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/AvroSink.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/AvroSink.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/AvroSink.java
index c78870b..6c36266 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/AvroSink.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/AvroSink.java
@@ -32,40 +32,39 @@ import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.util.MimeTypes;
/** A {@link FileBasedSink} for Avro files. */
-class AvroSink<T, DestinationT> extends FileBasedSink<T, DestinationT> {
+class AvroSink<T> extends FileBasedSink<T> {
private final AvroCoder<T> coder;
private final SerializableAvroCodecFactory codec;
private final ImmutableMap<String, Object> metadata;
AvroSink(
ValueProvider<ResourceId> outputPrefix,
- DynamicDestinations<T, DestinationT> dynamicDestinations,
+ FilenamePolicy filenamePolicy,
AvroCoder<T> coder,
SerializableAvroCodecFactory codec,
ImmutableMap<String, Object> metadata) {
// Avro handle compression internally using the codec.
- super(outputPrefix, dynamicDestinations, CompressionType.UNCOMPRESSED);
+ super(outputPrefix, filenamePolicy, CompressionType.UNCOMPRESSED);
this.coder = coder;
this.codec = codec;
this.metadata = metadata;
}
@Override
- public WriteOperation<T, DestinationT> createWriteOperation() {
+ public WriteOperation<T> createWriteOperation() {
return new AvroWriteOperation<>(this, coder, codec, metadata);
}
/** A {@link WriteOperation WriteOperation} for Avro files. */
- private static class AvroWriteOperation<T, DestinationT> extends WriteOperation<T, DestinationT> {
+ private static class AvroWriteOperation<T> extends WriteOperation<T> {
private final AvroCoder<T> coder;
private final SerializableAvroCodecFactory codec;
private final ImmutableMap<String, Object> metadata;
- private AvroWriteOperation(
- AvroSink<T, DestinationT> sink,
- AvroCoder<T> coder,
- SerializableAvroCodecFactory codec,
- ImmutableMap<String, Object> metadata) {
+ private AvroWriteOperation(AvroSink<T> sink,
+ AvroCoder<T> coder,
+ SerializableAvroCodecFactory codec,
+ ImmutableMap<String, Object> metadata) {
super(sink);
this.coder = coder;
this.codec = codec;
@@ -73,23 +72,22 @@ class AvroSink<T, DestinationT> extends FileBasedSink<T, DestinationT> {
}
@Override
- public Writer<T, DestinationT> createWriter() throws Exception {
+ public Writer<T> createWriter() throws Exception {
return new AvroWriter<>(this, coder, codec, metadata);
}
}
/** A {@link Writer Writer} for Avro files. */
- private static class AvroWriter<T, DestinationT> extends Writer<T, DestinationT> {
+ private static class AvroWriter<T> extends Writer<T> {
private final AvroCoder<T> coder;
private DataFileWriter<T> dataFileWriter;
private SerializableAvroCodecFactory codec;
private final ImmutableMap<String, Object> metadata;
- public AvroWriter(
- WriteOperation<T, DestinationT> writeOperation,
- AvroCoder<T> coder,
- SerializableAvroCodecFactory codec,
- ImmutableMap<String, Object> metadata) {
+ public AvroWriter(WriteOperation<T> writeOperation,
+ AvroCoder<T> coder,
+ SerializableAvroCodecFactory codec,
+ ImmutableMap<String, Object> metadata) {
super(writeOperation, MimeTypes.BINARY);
this.coder = coder;
this.codec = codec;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/CompressedSource.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/CompressedSource.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/CompressedSource.java
index 4baac36..6ab8dec 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/CompressedSource.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/CompressedSource.java
@@ -96,6 +96,12 @@ public class CompressedSource<T> extends FileBasedSource<T> {
*/
ReadableByteChannel createDecompressingChannel(String fileName, ReadableByteChannel channel)
throws IOException;
+
+ /**
+ * Given a file name, returns true if the file name matches any supported compression
+ * scheme.
+ */
+ boolean isCompressed(String fileName);
}
/**
@@ -236,16 +242,6 @@ public class CompressedSource<T> extends FileBasedSource<T> {
@Override
public abstract ReadableByteChannel createDecompressingChannel(ReadableByteChannel channel)
throws IOException;
-
- /** Returns whether the file's extension matches of one of the known compression formats. */
- public static boolean isCompressed(String filename) {
- for (CompressionMode type : CompressionMode.values()) {
- if (type.matches(filename)) {
- return true;
- }
- }
- return false;
- }
}
/**
@@ -277,6 +273,16 @@ public class CompressedSource<T> extends FileBasedSource<T> {
ReadableByteChannel.class.getSimpleName(),
ReadableByteChannel.class.getSimpleName()));
}
+
+ @Override
+ public boolean isCompressed(String fileName) {
+ for (CompressionMode type : CompressionMode.values()) {
+ if (type.matches(fileName)) {
+ return true;
+ }
+ }
+ return false;
+ }
}
private final FileBasedSource<T> sourceDelegate;
@@ -360,9 +366,13 @@ public class CompressedSource<T> extends FileBasedSource<T> {
*/
@Override
protected final boolean isSplittable() throws Exception {
- return channelFactory instanceof FileNameBasedDecompressingChannelFactory
- && !CompressionMode.isCompressed(getFileOrPatternSpec())
- && sourceDelegate.isSplittable();
+ if (channelFactory instanceof FileNameBasedDecompressingChannelFactory) {
+ FileNameBasedDecompressingChannelFactory fileNameBasedChannelFactory =
+ (FileNameBasedDecompressingChannelFactory) channelFactory;
+ return !fileNameBasedChannelFactory.isCompressed(getFileOrPatternSpec())
+ && sourceDelegate.isSplittable();
+ }
+ return false;
}
/**
@@ -376,7 +386,9 @@ public class CompressedSource<T> extends FileBasedSource<T> {
@Override
protected final FileBasedReader<T> createSingleFileReader(PipelineOptions options) {
if (channelFactory instanceof FileNameBasedDecompressingChannelFactory) {
- if (!CompressionMode.isCompressed(getFileOrPatternSpec())) {
+ FileNameBasedDecompressingChannelFactory fileNameBasedChannelFactory =
+ (FileNameBasedDecompressingChannelFactory) channelFactory;
+ if (!fileNameBasedChannelFactory.isCompressed(getFileOrPatternSpec())) {
return sourceDelegate.createSingleFileReader(options);
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/DefaultFilenamePolicy.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/DefaultFilenamePolicy.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/DefaultFilenamePolicy.java
index 7a60e49..f9e4ac4 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/DefaultFilenamePolicy.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/DefaultFilenamePolicy.java
@@ -20,31 +20,25 @@ package org.apache.beam.sdk.io;
import static com.google.common.base.MoreObjects.firstNonNull;
import com.google.common.annotations.VisibleForTesting;
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
-import java.io.Serializable;
import java.text.DecimalFormat;
import java.util.Arrays;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import javax.annotation.Nullable;
-import org.apache.beam.sdk.coders.AtomicCoder;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.CoderException;
-import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.io.FileBasedSink.FilenamePolicy;
-import org.apache.beam.sdk.io.FileBasedSink.OutputFileHints;
import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions;
import org.apache.beam.sdk.io.fs.ResourceId;
import org.apache.beam.sdk.options.ValueProvider;
-import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
+import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider;
+import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
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.transforms.windowing.PaneInfo.Timing;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
* A default {@link FilenamePolicy} for windowed and unwindowed files. This policy is constructed
@@ -57,7 +51,10 @@ import org.apache.beam.sdk.transforms.windowing.PaneInfo.Timing;
* {@code WriteOneFilePerWindow} example pipeline.
*/
public final class DefaultFilenamePolicy extends FilenamePolicy {
- /** The default sharding name template. */
+
+ private static final Logger LOG = LoggerFactory.getLogger(DefaultFilenamePolicy.class);
+
+ /** The default sharding name template used in {@link #constructUsingStandardParameters}. */
public static final String DEFAULT_UNWINDOWED_SHARD_TEMPLATE = ShardNameTemplate.INDEX_OF_MAX;
/** The default windowed sharding name template used when writing windowed files.
@@ -70,184 +67,75 @@ public final class DefaultFilenamePolicy extends FilenamePolicy {
"W-P" + DEFAULT_UNWINDOWED_SHARD_TEMPLATE;
/*
- * pattern for both windowed and non-windowed file names.
+ * pattern for both windowed and non-windowed file names
*/
private static final Pattern SHARD_FORMAT_RE = Pattern.compile("(S+|N+|W|P)");
/**
- * Encapsulates constructor parameters to {@link DefaultFilenamePolicy}.
- *
- * <p>This is used as the {@code DestinationT} argument to allow {@link DefaultFilenamePolicy}
- * objects to be dynamically generated.
- */
- public static class Params implements Serializable {
- private final ValueProvider<ResourceId> baseFilename;
- private final String shardTemplate;
- private final boolean explicitTemplate;
- private final String suffix;
-
- /**
- * Construct a default Params object. The shard template will be set to the default {@link
- * #DEFAULT_UNWINDOWED_SHARD_TEMPLATE} value.
- */
- public Params() {
- this.baseFilename = null;
- this.shardTemplate = DEFAULT_UNWINDOWED_SHARD_TEMPLATE;
- this.suffix = "";
- this.explicitTemplate = false;
- }
-
- private Params(
- ValueProvider<ResourceId> baseFilename,
- String shardTemplate,
- String suffix,
- boolean explicitTemplate) {
- this.baseFilename = baseFilename;
- this.shardTemplate = shardTemplate;
- this.suffix = suffix;
- this.explicitTemplate = explicitTemplate;
- }
-
- /**
- * Specify that writes are windowed. This affects the default shard template, changing it to
- * {@link #DEFAULT_WINDOWED_SHARD_TEMPLATE}.
- */
- public Params withWindowedWrites() {
- String template = this.shardTemplate;
- if (!explicitTemplate) {
- template = DEFAULT_WINDOWED_SHARD_TEMPLATE;
- }
- return new Params(baseFilename, template, suffix, explicitTemplate);
- }
-
- /** Sets the base filename. */
- public Params withBaseFilename(ResourceId baseFilename) {
- return withBaseFilename(StaticValueProvider.of(baseFilename));
- }
-
- /** Like {@link #withBaseFilename(ResourceId)}, but takes in a {@link ValueProvider}. */
- public Params withBaseFilename(ValueProvider<ResourceId> baseFilename) {
- return new Params(baseFilename, shardTemplate, suffix, explicitTemplate);
- }
-
- /** Sets the shard template. */
- public Params withShardTemplate(String shardTemplate) {
- return new Params(baseFilename, shardTemplate, suffix, true);
- }
-
- /** Sets the suffix. */
- public Params withSuffix(String suffix) {
- return new Params(baseFilename, shardTemplate, suffix, explicitTemplate);
- }
- }
-
- /** A Coder for {@link Params}. */
- public static class ParamsCoder extends AtomicCoder<Params> {
- private static final ParamsCoder INSTANCE = new ParamsCoder();
- private Coder<String> stringCoder = StringUtf8Coder.of();
-
- public static ParamsCoder of() {
- return INSTANCE;
- }
-
- @Override
- public void encode(Params value, OutputStream outStream) throws IOException {
- if (value == null) {
- throw new CoderException("cannot encode a null value");
- }
- stringCoder.encode(value.baseFilename.get().toString(), outStream);
- stringCoder.encode(value.shardTemplate, outStream);
- stringCoder.encode(value.suffix, outStream);
- }
-
- @Override
- public Params decode(InputStream inStream) throws IOException {
- ResourceId prefix =
- FileBasedSink.convertToFileResourceIfPossible(stringCoder.decode(inStream));
- String shardTemplate = stringCoder.decode(inStream);
- String suffix = stringCoder.decode(inStream);
- return new Params()
- .withBaseFilename(prefix)
- .withShardTemplate(shardTemplate)
- .withSuffix(suffix);
- }
- }
-
- private final Params params;
- /**
* Constructs a new {@link DefaultFilenamePolicy}.
*
* @see DefaultFilenamePolicy for more information on the arguments to this function.
*/
@VisibleForTesting
- DefaultFilenamePolicy(Params params) {
- this.params = params;
+ DefaultFilenamePolicy(ValueProvider<String> prefix, String shardTemplate, String suffix) {
+ this.prefix = prefix;
+ this.shardTemplate = shardTemplate;
+ this.suffix = suffix;
}
/**
- * Construct a {@link DefaultFilenamePolicy}.
+ * A helper function to construct a {@link DefaultFilenamePolicy} using the standard filename
+ * parameters, namely a provided {@link ResourceId} for the output prefix, and possibly-null
+ * shard name template and suffix.
*
- * <p>This is a shortcut for:
+ * <p>Any filename component of the provided resource will be used as the filename prefix.
*
- * <pre>{@code
- * DefaultFilenamePolicy.fromParams(new Params()
- * .withBaseFilename(baseFilename)
- * .withShardTemplate(shardTemplate)
- * .withSuffix(filenameSuffix)
- * .withWindowedWrites())
- * }</pre>
+ * <p>If provided, the shard name template will be used; otherwise
+ * {@link #DEFAULT_UNWINDOWED_SHARD_TEMPLATE} will be used for non-windowed file names and
+ * {@link #DEFAULT_WINDOWED_SHARD_TEMPLATE} will be used for windowed file names.
*
- * <p>Where the respective {@code with} methods are invoked only if the value is non-null or true.
+ * <p>If provided, the suffix will be used; otherwise the files will have an empty suffix.
*/
- public static DefaultFilenamePolicy fromStandardParameters(
- ValueProvider<ResourceId> baseFilename,
+ public static DefaultFilenamePolicy constructUsingStandardParameters(
+ ValueProvider<ResourceId> outputPrefix,
@Nullable String shardTemplate,
@Nullable String filenameSuffix,
boolean windowedWrites) {
- Params params = new Params().withBaseFilename(baseFilename);
- if (shardTemplate != null) {
- params = params.withShardTemplate(shardTemplate);
- }
- if (filenameSuffix != null) {
- params = params.withSuffix(filenameSuffix);
- }
- if (windowedWrites) {
- params = params.withWindowedWrites();
- }
- return fromParams(params);
+ // Pick the appropriate default policy based on whether windowed writes are being performed.
+ String defaultTemplate =
+ windowedWrites ? DEFAULT_WINDOWED_SHARD_TEMPLATE : DEFAULT_UNWINDOWED_SHARD_TEMPLATE;
+ return new DefaultFilenamePolicy(
+ NestedValueProvider.of(outputPrefix, new ExtractFilename()),
+ firstNonNull(shardTemplate, defaultTemplate),
+ firstNonNull(filenameSuffix, ""));
}
- /** Construct a {@link DefaultFilenamePolicy} from a {@link Params} object. */
- public static DefaultFilenamePolicy fromParams(Params params) {
- return new DefaultFilenamePolicy(params);
- }
+ private final ValueProvider<String> prefix;
+ private final String shardTemplate;
+ private final String suffix;
/**
* Constructs a fully qualified name from components.
*
- * <p>The name is built from a base filename, shard template (with shard numbers applied), and a
- * suffix. All components are required, but may be empty strings.
+ * <p>The name is built from a prefix, shard template (with shard numbers
+ * applied), and a suffix. All components are required, but may be empty
+ * strings.
*
- * <p>Within a shard template, repeating sequences of the letters "S" or "N" are replaced with the
- * shard number, or number of shards respectively. "P" is replaced with by stringification of
- * current pane. "W" is replaced by stringification of current window.
+ * <p>Within a shard template, repeating sequences of the letters "S" or "N"
+ * are replaced with the shard number, or number of shards respectively.
+ * "P" is replaced with by stringification of current pane.
+ * "W" is replaced by stringification of current window.
*
- * <p>The numbers are formatted with leading zeros to match the length of the repeated sequence of
- * letters.
+ * <p>The numbers are formatted with leading zeros to match the length of the
+ * repeated sequence of letters.
*
- * <p>For example, if baseFilename = "path/to/output", shardTemplate = "-SSS-of-NNN", and suffix =
- * ".txt", with shardNum = 1 and numShards = 100, the following is produced:
- * "path/to/output-001-of-100.txt".
+ * <p>For example, if prefix = "output", shardTemplate = "-SSS-of-NNN", and
+ * suffix = ".txt", with shardNum = 1 and numShards = 100, the following is
+ * produced: "output-001-of-100.txt".
*/
- static ResourceId constructName(
- ResourceId baseFilename,
- String shardTemplate,
- String suffix,
- int shardNum,
- int numShards,
- String paneStr,
- String windowStr) {
- String prefix = extractFilename(baseFilename);
+ static String constructName(
+ String prefix, String shardTemplate, String suffix, int shardNum, int numShards,
+ String paneStr, String windowStr) {
// Matcher API works with StringBuffer, rather than StringBuilder.
StringBuffer sb = new StringBuffer();
sb.append(prefix);
@@ -277,37 +165,27 @@ public final class DefaultFilenamePolicy extends FilenamePolicy {
m.appendTail(sb);
sb.append(suffix);
- return baseFilename
- .getCurrentDirectory()
- .resolve(sb.toString(), StandardResolveOptions.RESOLVE_FILE);
+ return sb.toString();
}
@Override
@Nullable
- public ResourceId unwindowedFilename(Context context, OutputFileHints outputFileHints) {
- return constructName(
- params.baseFilename.get(),
- params.shardTemplate,
- params.suffix + outputFileHints.getSuggestedFilenameSuffix(),
- context.getShardNumber(),
- context.getNumShards(),
- null,
- null);
+ public ResourceId unwindowedFilename(ResourceId outputDirectory, Context context,
+ String extension) {
+ String filename = constructName(prefix.get(), shardTemplate, suffix, context.getShardNumber(),
+ context.getNumShards(), null, null) + extension;
+ return outputDirectory.resolve(filename, StandardResolveOptions.RESOLVE_FILE);
}
@Override
- public ResourceId windowedFilename(WindowedContext context, OutputFileHints outputFileHints) {
+ public ResourceId windowedFilename(ResourceId outputDirectory,
+ WindowedContext context, String extension) {
final PaneInfo paneInfo = context.getPaneInfo();
String paneStr = paneInfoToString(paneInfo);
String windowStr = windowToString(context.getWindow());
- return constructName(
- params.baseFilename.get(),
- params.shardTemplate,
- params.suffix + outputFileHints.getSuggestedFilenameSuffix(),
- context.getShardNumber(),
- context.getNumShards(),
- paneStr,
- windowStr);
+ String filename = constructName(prefix.get(), shardTemplate, suffix, context.getShardNumber(),
+ context.getNumShards(), paneStr, windowStr) + extension;
+ return outputDirectory.resolve(filename, StandardResolveOptions.RESOLVE_FILE);
}
/*
@@ -338,32 +216,24 @@ public final class DefaultFilenamePolicy extends FilenamePolicy {
@Override
public void populateDisplayData(DisplayData.Builder builder) {
String filenamePattern;
- if (params.baseFilename.isAccessible()) {
- filenamePattern =
- String.format("%s%s%s", params.baseFilename.get(), params.shardTemplate, params.suffix);
+ if (prefix.isAccessible()) {
+ filenamePattern = String.format("%s%s%s", prefix.get(), shardTemplate, suffix);
} else {
- filenamePattern =
- String.format("%s%s%s", params.baseFilename, params.shardTemplate, params.suffix);
+ filenamePattern = String.format("%s%s%s", prefix, shardTemplate, suffix);
}
-
- String outputPrefixString = null;
- outputPrefixString =
- params.baseFilename.isAccessible()
- ? params.baseFilename.get().toString()
- : params.baseFilename.toString();
- builder.add(DisplayData.item("filenamePattern", filenamePattern).withLabel("Filename Pattern"));
- builder.add(DisplayData.item("filePrefix", outputPrefixString).withLabel("Output File Prefix"));
- builder.add(DisplayData.item("fileSuffix", params.suffix).withLabel("Output file Suffix"));
builder.add(
- DisplayData.item("shardNameTemplate", params.shardTemplate)
- .withLabel("Output Shard Name Template"));
+ DisplayData.item("filenamePattern", filenamePattern)
+ .withLabel("Filename Pattern"));
}
- private static String extractFilename(ResourceId input) {
- if (input.isDirectory()) {
- return "";
- } else {
- return firstNonNull(input.getFilename(), "");
+ private static class ExtractFilename implements SerializableFunction<ResourceId, String> {
+ @Override
+ public String apply(ResourceId input) {
+ if (input.isDirectory()) {
+ return "";
+ } else {
+ return firstNonNull(input.getFilename(), "");
+ }
}
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/DynamicFileDestinations.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/DynamicFileDestinations.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/DynamicFileDestinations.java
deleted file mode 100644
index e7ef0f6..0000000
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/DynamicFileDestinations.java
+++ /dev/null
@@ -1,115 +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.sdk.io;
-
-import javax.annotation.Nullable;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.io.DefaultFilenamePolicy.Params;
-import org.apache.beam.sdk.io.DefaultFilenamePolicy.ParamsCoder;
-import org.apache.beam.sdk.io.FileBasedSink.DynamicDestinations;
-import org.apache.beam.sdk.io.FileBasedSink.FilenamePolicy;
-import org.apache.beam.sdk.transforms.SerializableFunction;
-import org.apache.beam.sdk.transforms.display.DisplayData;
-
-/** Some helper classes that derive from {@link FileBasedSink.DynamicDestinations}. */
-public class DynamicFileDestinations {
- /** Always returns a constant {@link FilenamePolicy}. */
- private static class ConstantFilenamePolicy<T> extends DynamicDestinations<T, Void> {
- private final FilenamePolicy filenamePolicy;
-
- public ConstantFilenamePolicy(FilenamePolicy filenamePolicy) {
- this.filenamePolicy = filenamePolicy;
- }
-
- @Override
- public Void getDestination(T element) {
- return (Void) null;
- }
-
- @Override
- public Coder<Void> getDestinationCoder() {
- return null;
- }
-
- @Override
- public Void getDefaultDestination() {
- return (Void) null;
- }
-
- @Override
- public FilenamePolicy getFilenamePolicy(Void destination) {
- return filenamePolicy;
- }
-
- @Override
- public void populateDisplayData(DisplayData.Builder builder) {
- filenamePolicy.populateDisplayData(builder);
- }
- }
-
- /**
- * A base class for a {@link DynamicDestinations} object that returns differently-configured
- * instances of {@link DefaultFilenamePolicy}.
- */
- private static class DefaultPolicyDestinations<UserT> extends DynamicDestinations<UserT, Params> {
- SerializableFunction<UserT, Params> destinationFunction;
- Params emptyDestination;
-
- public DefaultPolicyDestinations(
- SerializableFunction<UserT, Params> destinationFunction, Params emptyDestination) {
- this.destinationFunction = destinationFunction;
- this.emptyDestination = emptyDestination;
- }
-
- @Override
- public Params getDestination(UserT element) {
- return destinationFunction.apply(element);
- }
-
- @Override
- public Params getDefaultDestination() {
- return emptyDestination;
- }
-
- @Nullable
- @Override
- public Coder<Params> getDestinationCoder() {
- return ParamsCoder.of();
- }
-
- @Override
- public FilenamePolicy getFilenamePolicy(DefaultFilenamePolicy.Params params) {
- return DefaultFilenamePolicy.fromParams(params);
- }
- }
-
- /** Returns a {@link DynamicDestinations} that always returns the same {@link FilenamePolicy}. */
- public static <T> DynamicDestinations<T, Void> constant(FilenamePolicy filenamePolicy) {
- return new ConstantFilenamePolicy<>(filenamePolicy);
- }
-
- /**
- * Returns a {@link DynamicDestinations} that returns instances of {@link DefaultFilenamePolicy}
- * configured with the given {@link Params}.
- */
- public static <UserT> DynamicDestinations<UserT, Params> toDefaultPolicies(
- SerializableFunction<UserT, Params> destinationFunction, Params emptyDestination) {
- return new DefaultPolicyDestinations<>(destinationFunction, emptyDestination);
- }
-}
[16/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java
index 7013044..2fd10ac 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java
@@ -20,12 +20,9 @@ package org.apache.beam.sdk.io;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
-import com.google.common.base.Objects;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
-import com.google.common.hash.Hashing;
-import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.UUID;
@@ -33,12 +30,8 @@ import java.util.concurrent.ThreadLocalRandom;
import javax.annotation.Nullable;
import org.apache.beam.sdk.Pipeline;
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.Coder.NonDeterministicException;
import org.apache.beam.sdk.coders.KvCoder;
-import org.apache.beam.sdk.coders.ShardedKeyCoder;
-import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.io.FileBasedSink.FileResult;
import org.apache.beam.sdk.io.FileBasedSink.FileResultCoder;
@@ -49,11 +42,9 @@ import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.Flatten;
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.View;
import org.apache.beam.sdk.transforms.WithKeys;
import org.apache.beam.sdk.transforms.display.DisplayData;
@@ -62,17 +53,11 @@ import org.apache.beam.sdk.transforms.windowing.DefaultTrigger;
import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.transforms.windowing.Window;
-import org.apache.beam.sdk.util.CoderUtils;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollection.IsBounded;
-import org.apache.beam.sdk.values.PCollectionList;
-import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PDone;
-import org.apache.beam.sdk.values.ShardedKey;
-import org.apache.beam.sdk.values.TupleTag;
-import org.apache.beam.sdk.values.TupleTagList;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -81,12 +66,13 @@ import org.slf4j.LoggerFactory;
* global initialization of a sink, followed by a parallel write, and ends with a sequential
* finalization of the write. The output of a write is {@link PDone}.
*
- * <p>By default, every bundle in the input {@link PCollection} will be processed by a {@link
- * WriteOperation}, so the number of output will vary based on runner behavior, though at least 1
- * output will always be produced. The exact parallelism of the write stage can be controlled using
- * {@link WriteFiles#withNumShards}, typically used to control how many files are produced or to
- * globally limit the number of workers connecting to an external service. However, this option can
- * often hurt performance: it adds an additional {@link GroupByKey} to the pipeline.
+ * <p>By default, every bundle in the input {@link PCollection} will be processed by a
+ * {@link WriteOperation}, so the number of output
+ * will vary based on runner behavior, though at least 1 output will always be produced. The
+ * exact parallelism of the write stage can be controlled using {@link WriteFiles#withNumShards},
+ * typically used to control how many files are produced or to globally limit the number of
+ * workers connecting to an external service. However, this option can often hurt performance: it
+ * adds an additional {@link GroupByKey} to the pipeline.
*
* <p>Example usage with runner-determined sharding:
*
@@ -97,70 +83,44 @@ import org.slf4j.LoggerFactory;
* <pre>{@code p.apply(WriteFiles.to(new MySink(...)).withNumShards(3));}</pre>
*/
@Experimental(Experimental.Kind.SOURCE_SINK)
-public class WriteFiles<UserT, DestinationT, OutputT>
- extends PTransform<PCollection<UserT>, PDone> {
+public class WriteFiles<T> extends PTransform<PCollection<T>, PDone> {
private static final Logger LOG = LoggerFactory.getLogger(WriteFiles.class);
- // The maximum number of file writers to keep open in a single bundle at a time, since file
- // writers default to 64mb buffers. This comes into play when writing per-window files.
- // The first 20 files from a single WriteFiles transform will write files inline in the
- // transform. Anything beyond that might be shuffled.
- // Keep in mind that specific runners may decide to run multiple bundles in parallel, based on
- // their own policy.
- private static final int DEFAULT_MAX_NUM_WRITERS_PER_BUNDLE = 20;
-
- // When we spill records, shard the output keys to prevent hotspots.
- // We could consider making this a parameter.
- private static final int SPILLED_RECORD_SHARDING_FACTOR = 10;
-
static final int UNKNOWN_SHARDNUM = -1;
- private FileBasedSink<OutputT, DestinationT> sink;
- private SerializableFunction<UserT, OutputT> formatFunction;
- private WriteOperation<OutputT, DestinationT> writeOperation;
+ private FileBasedSink<T> sink;
+ private WriteOperation<T> writeOperation;
// This allows the number of shards to be dynamically computed based on the input
// PCollection.
- @Nullable private final PTransform<PCollection<UserT>, PCollectionView<Integer>> computeNumShards;
+ @Nullable
+ private final PTransform<PCollection<T>, PCollectionView<Integer>> computeNumShards;
// We don't use a side input for static sharding, as we want this value to be updatable
// when a pipeline is updated.
@Nullable
private final ValueProvider<Integer> numShardsProvider;
private final boolean windowedWrites;
- private int maxNumWritersPerBundle;
/**
* Creates a {@link WriteFiles} transform that writes to the given {@link FileBasedSink}, letting
* the runner control how many different shards are produced.
*/
- public static <UserT, DestinationT, OutputT> WriteFiles<UserT, DestinationT, OutputT> to(
- FileBasedSink<OutputT, DestinationT> sink,
- SerializableFunction<UserT, OutputT> formatFunction) {
+ public static <T> WriteFiles<T> to(FileBasedSink<T> sink) {
checkNotNull(sink, "sink");
- return new WriteFiles<>(
- sink,
- formatFunction,
- null /* runner-determined sharding */,
- null,
- false,
- DEFAULT_MAX_NUM_WRITERS_PER_BUNDLE);
+ return new WriteFiles<>(sink, null /* runner-determined sharding */, null, false);
}
private WriteFiles(
- FileBasedSink<OutputT, DestinationT> sink,
- SerializableFunction<UserT, OutputT> formatFunction,
- @Nullable PTransform<PCollection<UserT>, PCollectionView<Integer>> computeNumShards,
+ FileBasedSink<T> sink,
+ @Nullable PTransform<PCollection<T>, PCollectionView<Integer>> computeNumShards,
@Nullable ValueProvider<Integer> numShardsProvider,
- boolean windowedWrites,
- int maxNumWritersPerBundle) {
+ boolean windowedWrites) {
this.sink = sink;
- this.formatFunction = checkNotNull(formatFunction);
this.computeNumShards = computeNumShards;
this.numShardsProvider = numShardsProvider;
this.windowedWrites = windowedWrites;
- this.maxNumWritersPerBundle = maxNumWritersPerBundle;
}
@Override
- public PDone expand(PCollection<UserT> input) {
+ public PDone expand(PCollection<T> input) {
if (input.isBounded() == IsBounded.UNBOUNDED) {
checkArgument(windowedWrites,
"Must use windowed writes when applying %s to an unbounded PCollection",
@@ -199,16 +159,13 @@ public class WriteFiles<UserT, DestinationT, OutputT>
}
}
- /** Returns the {@link FileBasedSink} associated with this PTransform. */
- public FileBasedSink<OutputT, DestinationT> getSink() {
+ /**
+ * Returns the {@link FileBasedSink} associated with this PTransform.
+ */
+ public FileBasedSink<T> getSink() {
return sink;
}
- /** Returns the the format function that maps the user type to the record written to files. */
- public SerializableFunction<UserT, OutputT> getFormatFunction() {
- return formatFunction;
- }
-
/**
* Returns whether or not to perform windowed writes.
*/
@@ -223,7 +180,7 @@ public class WriteFiles<UserT, DestinationT, OutputT>
* #withRunnerDeterminedSharding()}.
*/
@Nullable
- public PTransform<PCollection<UserT>, PCollectionView<Integer>> getSharding() {
+ public PTransform<PCollection<T>, PCollectionView<Integer>> getSharding() {
return computeNumShards;
}
@@ -241,7 +198,7 @@ public class WriteFiles<UserT, DestinationT, OutputT>
* <p>A value less than or equal to 0 will be equivalent to the default behavior of
* runner-determined sharding.
*/
- public WriteFiles<UserT, DestinationT, OutputT> withNumShards(int numShards) {
+ public WriteFiles<T> withNumShards(int numShards) {
if (numShards > 0) {
return withNumShards(StaticValueProvider.of(numShards));
}
@@ -255,27 +212,8 @@ public class WriteFiles<UserT, DestinationT, OutputT>
* <p>This option should be used sparingly as it can hurt performance. See {@link WriteFiles} for
* more information.
*/
- public WriteFiles<UserT, DestinationT, OutputT> withNumShards(
- ValueProvider<Integer> numShardsProvider) {
- return new WriteFiles<>(
- sink,
- formatFunction,
- computeNumShards,
- numShardsProvider,
- windowedWrites,
- maxNumWritersPerBundle);
- }
-
- /** Set the maximum number of writers created in a bundle before spilling to shuffle. */
- public WriteFiles<UserT, DestinationT, OutputT> withMaxNumWritersPerBundle(
- int maxNumWritersPerBundle) {
- return new WriteFiles<>(
- sink,
- formatFunction,
- computeNumShards,
- numShardsProvider,
- windowedWrites,
- maxNumWritersPerBundle);
+ public WriteFiles<T> withNumShards(ValueProvider<Integer> numShardsProvider) {
+ return new WriteFiles<>(sink, null, numShardsProvider, windowedWrites);
}
/**
@@ -285,169 +223,127 @@ public class WriteFiles<UserT, DestinationT, OutputT>
* <p>This option should be used sparingly as it can hurt performance. See {@link WriteFiles} for
* more information.
*/
- public WriteFiles<UserT, DestinationT, OutputT> withSharding(
- PTransform<PCollection<UserT>, PCollectionView<Integer>> sharding) {
+ public WriteFiles<T> withSharding(PTransform<PCollection<T>, PCollectionView<Integer>> sharding) {
checkNotNull(
sharding, "Cannot provide null sharding. Use withRunnerDeterminedSharding() instead");
- return new WriteFiles<>(
- sink, formatFunction, sharding, null, windowedWrites, maxNumWritersPerBundle);
+ return new WriteFiles<>(sink, sharding, null, windowedWrites);
}
/**
* Returns a new {@link WriteFiles} that will write to the current {@link FileBasedSink} with
* runner-determined sharding.
*/
- public WriteFiles<UserT, DestinationT, OutputT> withRunnerDeterminedSharding() {
- return new WriteFiles<>(
- sink, formatFunction, null, null, windowedWrites, maxNumWritersPerBundle);
+ public WriteFiles<T> withRunnerDeterminedSharding() {
+ return new WriteFiles<>(sink, null, null, windowedWrites);
}
/**
* Returns a new {@link WriteFiles} that writes preserves windowing on it's input.
*
- * <p>If this option is not specified, windowing and triggering are replaced by {@link
- * GlobalWindows} and {@link DefaultTrigger}.
+ * <p>If this option is not specified, windowing and triggering are replaced by
+ * {@link GlobalWindows} and {@link DefaultTrigger}.
*
- * <p>If there is no data for a window, no output shards will be generated for that window. If a
- * window triggers multiple times, then more than a single output shard might be generated
- * multiple times; it's up to the sink implementation to keep these output shards unique.
+ * <p>If there is no data for a window, no output shards will be generated for that window.
+ * If a window triggers multiple times, then more than a single output shard might be
+ * generated multiple times; it's up to the sink implementation to keep these output shards
+ * unique.
*
- * <p>This option can only be used if {@link #withNumShards(int)} is also set to a positive value.
+ * <p>This option can only be used if {@link #withNumShards(int)} is also set to a
+ * positive value.
*/
- public WriteFiles<UserT, DestinationT, OutputT> withWindowedWrites() {
- return new WriteFiles<>(
- sink, formatFunction, computeNumShards, numShardsProvider, true, maxNumWritersPerBundle);
+ public WriteFiles<T> withWindowedWrites() {
+ return new WriteFiles<>(sink, computeNumShards, numShardsProvider, true);
}
- private static class WriterKey<DestinationT> {
- private final BoundedWindow window;
- private final PaneInfo paneInfo;
- private final DestinationT destination;
+ /**
+ * Writes all the elements in a bundle using a {@link Writer} produced by the
+ * {@link WriteOperation} associated with the {@link FileBasedSink} with windowed writes enabled.
+ */
+ private class WriteWindowedBundles extends DoFn<T, FileResult> {
+ private Map<KV<BoundedWindow, PaneInfo>, Writer<T>> windowedWriters;
- WriterKey(BoundedWindow window, PaneInfo paneInfo, DestinationT destination) {
- this.window = window;
- this.paneInfo = paneInfo;
- this.destination = destination;
+ @StartBundle
+ public void startBundle(StartBundleContext c) {
+ // Reset state in case of reuse. We need to make sure that each bundle gets unique writers.
+ windowedWriters = Maps.newHashMap();
}
- @Override
- public boolean equals(Object o) {
- if (!(o instanceof WriterKey)) {
- return false;
+ @ProcessElement
+ public void processElement(ProcessContext c, BoundedWindow window) throws Exception {
+ PaneInfo paneInfo = c.pane();
+ Writer<T> writer;
+ // If we are doing windowed writes, we need to ensure that we have separate files for
+ // data in different windows/panes.
+ KV<BoundedWindow, PaneInfo> key = KV.of(window, paneInfo);
+ writer = windowedWriters.get(key);
+ if (writer == null) {
+ String uuid = UUID.randomUUID().toString();
+ LOG.info(
+ "Opening writer {} for write operation {}, window {} pane {}",
+ uuid,
+ writeOperation,
+ window,
+ paneInfo);
+ writer = writeOperation.createWriter();
+ writer.openWindowed(uuid, window, paneInfo, UNKNOWN_SHARDNUM);
+ windowedWriters.put(key, writer);
+ LOG.debug("Done opening writer");
}
- WriterKey other = (WriterKey) o;
- return Objects.equal(window, other.window)
- && Objects.equal(paneInfo, other.paneInfo)
- && Objects.equal(destination, other.destination);
+
+ writeOrClose(writer, c.element());
}
- @Override
- public int hashCode() {
- return Objects.hashCode(window, paneInfo, destination);
+ @FinishBundle
+ public void finishBundle(FinishBundleContext c) throws Exception {
+ for (Map.Entry<KV<BoundedWindow, PaneInfo>, Writer<T>> entry : windowedWriters.entrySet()) {
+ FileResult result = entry.getValue().close();
+ BoundedWindow window = entry.getKey().getKey();
+ c.output(result, window.maxTimestamp(), window);
+ }
}
- }
- // Hash the destination in a manner that we can then use as a key in a GBK. Since Java's
- // hashCode isn't guaranteed to be stable across machines, we instead serialize the destination
- // and use murmur3_32 to hash it. We enforce that destinationCoder must be deterministic, so
- // this can be used as a key.
- private static <DestinationT> int hashDestination(
- DestinationT destination, Coder<DestinationT> destinationCoder) throws IOException {
- return Hashing.murmur3_32()
- .hashBytes(CoderUtils.encodeToByteArray(destinationCoder, destination))
- .asInt();
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ builder.delegate(WriteFiles.this);
+ }
}
/**
- * Writes all the elements in a bundle using a {@link Writer} produced by the {@link
- * WriteOperation} associated with the {@link FileBasedSink}.
+ * Writes all the elements in a bundle using a {@link Writer} produced by the
+ * {@link WriteOperation} associated with the {@link FileBasedSink} with windowed writes disabled.
*/
- private class WriteBundles extends DoFn<UserT, FileResult<DestinationT>> {
- private final TupleTag<KV<ShardedKey<Integer>, UserT>> unwrittenRecordsTag;
- private final Coder<DestinationT> destinationCoder;
- private final boolean windowedWrites;
-
- private Map<WriterKey<DestinationT>, Writer<OutputT, DestinationT>> writers;
- private int spilledShardNum = UNKNOWN_SHARDNUM;
-
- WriteBundles(
- boolean windowedWrites,
- TupleTag<KV<ShardedKey<Integer>, UserT>> unwrittenRecordsTag,
- Coder<DestinationT> destinationCoder) {
- this.windowedWrites = windowedWrites;
- this.unwrittenRecordsTag = unwrittenRecordsTag;
- this.destinationCoder = destinationCoder;
- }
+ private class WriteUnwindowedBundles extends DoFn<T, FileResult> {
+ // Writer that will write the records in this bundle. Lazily
+ // initialized in processElement.
+ private Writer<T> writer = null;
+ private BoundedWindow window = null;
@StartBundle
public void startBundle(StartBundleContext c) {
// Reset state in case of reuse. We need to make sure that each bundle gets unique writers.
- writers = Maps.newHashMap();
+ writer = null;
}
@ProcessElement
public void processElement(ProcessContext c, BoundedWindow window) throws Exception {
- PaneInfo paneInfo = c.pane();
- // If we are doing windowed writes, we need to ensure that we have separate files for
- // data in different windows/panes. Similar for dynamic writes, make sure that different
- // destinations go to different writers.
- // In the case of unwindowed writes, the window and the pane will always be the same, and
- // the map will only have a single element.
- DestinationT destination = sink.getDynamicDestinations().getDestination(c.element());
- WriterKey<DestinationT> key = new WriterKey<>(window, c.pane(), destination);
- Writer<OutputT, DestinationT> writer = writers.get(key);
+ // Cache a single writer for the bundle.
if (writer == null) {
- if (writers.size() <= maxNumWritersPerBundle) {
- String uuid = UUID.randomUUID().toString();
- LOG.info(
- "Opening writer {} for write operation {}, window {} pane {} destination {}",
- uuid,
- writeOperation,
- window,
- paneInfo,
- destination);
- writer = writeOperation.createWriter();
- if (windowedWrites) {
- writer.openWindowed(uuid, window, paneInfo, UNKNOWN_SHARDNUM, destination);
- } else {
- writer.openUnwindowed(uuid, UNKNOWN_SHARDNUM, destination);
- }
- writers.put(key, writer);
- LOG.debug("Done opening writer");
- } else {
- if (spilledShardNum == UNKNOWN_SHARDNUM) {
- // Cache the random value so we only call ThreadLocalRandom once per DoFn instance.
- spilledShardNum = ThreadLocalRandom.current().nextInt(SPILLED_RECORD_SHARDING_FACTOR);
- } else {
- spilledShardNum = (spilledShardNum + 1) % SPILLED_RECORD_SHARDING_FACTOR;
- }
- c.output(
- unwrittenRecordsTag,
- KV.of(
- ShardedKey.of(hashDestination(destination, destinationCoder), spilledShardNum),
- c.element()));
- return;
- }
+ LOG.info("Opening writer for write operation {}", writeOperation);
+ writer = writeOperation.createWriter();
+ writer.openUnwindowed(UUID.randomUUID().toString(), UNKNOWN_SHARDNUM);
+ LOG.debug("Done opening writer");
}
- writeOrClose(writer, formatFunction.apply(c.element()));
+ this.window = window;
+ writeOrClose(this.writer, c.element());
}
@FinishBundle
public void finishBundle(FinishBundleContext c) throws Exception {
- for (Map.Entry<WriterKey<DestinationT>, Writer<OutputT, DestinationT>> entry :
- writers.entrySet()) {
- Writer<OutputT, DestinationT> writer = entry.getValue();
- FileResult<DestinationT> result;
- try {
- result = writer.close();
- } catch (Exception e) {
- // If anything goes wrong, make sure to delete the temporary file.
- writer.cleanup();
- throw e;
- }
- BoundedWindow window = entry.getKey().window;
- c.output(result, window.maxTimestamp(), window);
+ if (writer == null) {
+ return;
}
+ FileResult result = writer.close();
+ c.output(result, window.maxTimestamp(), window);
}
@Override
@@ -456,62 +352,38 @@ public class WriteFiles<UserT, DestinationT, OutputT>
}
}
- enum ShardAssignment { ASSIGN_IN_FINALIZE, ASSIGN_WHEN_WRITING }
-
- /*
- * Like {@link WriteBundles}, but where the elements for each shard have been collected into a
- * single iterable.
+ /**
+ * Like {@link WriteWindowedBundles} and {@link WriteUnwindowedBundles}, but where the elements
+ * for each shard have been collected into a single iterable.
*/
- private class WriteShardedBundles
- extends DoFn<KV<ShardedKey<Integer>, Iterable<UserT>>, FileResult<DestinationT>> {
- ShardAssignment shardNumberAssignment;
- WriteShardedBundles(ShardAssignment shardNumberAssignment) {
- this.shardNumberAssignment = shardNumberAssignment;
- }
-
+ private class WriteShardedBundles extends DoFn<KV<Integer, Iterable<T>>, FileResult> {
@ProcessElement
public void processElement(ProcessContext c, BoundedWindow window) throws Exception {
- // Since we key by a 32-bit hash of the destination, there might be multiple destinations
- // in this iterable. The number of destinations is generally very small (1000s or less), so
- // there will rarely be hash collisions.
- Map<DestinationT, Writer<OutputT, DestinationT>> writers = Maps.newHashMap();
- for (UserT input : c.element().getValue()) {
- DestinationT destination = sink.getDynamicDestinations().getDestination(input);
- Writer<OutputT, DestinationT> writer = writers.get(destination);
- if (writer == null) {
- LOG.debug("Opening writer for write operation {}", writeOperation);
- writer = writeOperation.createWriter();
- if (windowedWrites) {
- int shardNumber =
- shardNumberAssignment == ShardAssignment.ASSIGN_WHEN_WRITING
- ? c.element().getKey().getShardNumber()
- : UNKNOWN_SHARDNUM;
- writer.openWindowed(
- UUID.randomUUID().toString(), window, c.pane(), shardNumber, destination);
- } else {
- writer.openUnwindowed(UUID.randomUUID().toString(), UNKNOWN_SHARDNUM, destination);
- }
- LOG.debug("Done opening writer");
- writers.put(destination, writer);
- }
- writeOrClose(writer, formatFunction.apply(input));
- }
+ // In a sharded write, single input element represents one shard. We can open and close
+ // the writer in each call to processElement.
+ LOG.info("Opening writer for write operation {}", writeOperation);
+ Writer<T> writer = writeOperation.createWriter();
+ if (windowedWrites) {
+ writer.openWindowed(UUID.randomUUID().toString(), window, c.pane(), c.element().getKey());
+ } else {
+ writer.openUnwindowed(UUID.randomUUID().toString(), UNKNOWN_SHARDNUM);
+ }
+ LOG.debug("Done opening writer");
- // Close all writers.
- for (Map.Entry<DestinationT, Writer<OutputT, DestinationT>> entry : writers.entrySet()) {
- Writer<OutputT, DestinationT> writer = entry.getValue();
- FileResult<DestinationT> result;
- try {
- // Close the writer; if this throws let the error propagate.
- result = writer.close();
- c.output(result);
- } catch (Exception e) {
- // If anything goes wrong, make sure to delete the temporary file.
- writer.cleanup();
- throw e;
+ try {
+ for (T t : c.element().getValue()) {
+ writeOrClose(writer, t);
}
+
+ // Close the writer; if this throws let the error propagate.
+ FileResult result = writer.close();
+ c.output(result);
+ } catch (Exception e) {
+ // If anything goes wrong, make sure to delete the temporary file.
+ writer.cleanup();
+ throw e;
}
- }
+ }
@Override
public void populateDisplayData(DisplayData.Builder builder) {
@@ -519,15 +391,12 @@ public class WriteFiles<UserT, DestinationT, OutputT>
}
}
- private static <OutputT, DestinationT> void writeOrClose(
- Writer<OutputT, DestinationT> writer, OutputT t) throws Exception {
+ private static <T> void writeOrClose(Writer<T> writer, T t) throws Exception {
try {
writer.write(t);
} catch (Exception e) {
try {
writer.close();
- // If anything goes wrong, make sure to delete the temporary file.
- writer.cleanup();
} catch (Exception closeException) {
if (closeException instanceof InterruptedException) {
// Do not silently ignore interrupted state.
@@ -540,25 +409,20 @@ public class WriteFiles<UserT, DestinationT, OutputT>
}
}
- private class ApplyShardingKey extends DoFn<UserT, KV<ShardedKey<Integer>, UserT>> {
+ private static class ApplyShardingKey<T> extends DoFn<T, KV<Integer, T>> {
private final PCollectionView<Integer> numShardsView;
private final ValueProvider<Integer> numShardsProvider;
- private final Coder<DestinationT> destinationCoder;
-
private int shardNumber;
- ApplyShardingKey(
- PCollectionView<Integer> numShardsView,
- ValueProvider<Integer> numShardsProvider,
- Coder<DestinationT> destinationCoder) {
- this.destinationCoder = destinationCoder;
+ ApplyShardingKey(PCollectionView<Integer> numShardsView,
+ ValueProvider<Integer> numShardsProvider) {
this.numShardsView = numShardsView;
this.numShardsProvider = numShardsProvider;
shardNumber = UNKNOWN_SHARDNUM;
}
@ProcessElement
- public void processElement(ProcessContext context) throws IOException {
+ public void processElement(ProcessContext context) {
final int shardCount;
if (numShardsView != null) {
shardCount = context.sideInput(numShardsView);
@@ -578,110 +442,65 @@ public class WriteFiles<UserT, DestinationT, OutputT>
} else {
shardNumber = (shardNumber + 1) % shardCount;
}
- // We avoid using destination itself as a sharding key, because destination is often large.
- // e.g. when using {@link DefaultFilenamePolicy}, the destination contains the entire path
- // to the file. Often most of the path is constant across all destinations, just the path
- // suffix is appended by the destination function. Instead we key by a 32-bit hash (carefully
- // chosen to be guaranteed stable), and call getDestination again in the next ParDo to resolve
- // the destinations. This does mean that multiple destinations might end up on the same shard,
- // however the number of collisions should be small, so there's no need to worry about memory
- // issues.
- DestinationT destination = sink.getDynamicDestinations().getDestination(context.element());
- context.output(
- KV.of(
- ShardedKey.of(hashDestination(destination, destinationCoder), shardNumber),
- context.element()));
+ context.output(KV.of(shardNumber, context.element()));
}
}
/**
* A write is performed as sequence of three {@link ParDo}'s.
*
- * <p>This singleton collection containing the WriteOperation is then used as a side input to a
- * ParDo over the PCollection of elements to write. In this bundle-writing phase, {@link
- * WriteOperation#createWriter} is called to obtain a {@link Writer}. {@link Writer#open} and
- * {@link Writer#close} are called in {@link DoFn.StartBundle} and {@link DoFn.FinishBundle},
- * respectively, and {@link Writer#write} method is called for every element in the bundle. The
- * output of this ParDo is a PCollection of <i>writer result</i> objects (see {@link
- * FileBasedSink} for a description of writer results)-one for each bundle.
+ * <p>This singleton collection containing the WriteOperation is then used as a side
+ * input to a ParDo over the PCollection of elements to write. In this bundle-writing phase,
+ * {@link WriteOperation#createWriter} is called to obtain a {@link Writer}.
+ * {@link Writer#open} and {@link Writer#close} are called in
+ * {@link DoFn.StartBundle} and {@link DoFn.FinishBundle}, respectively, and
+ * {@link Writer#write} method is called for every element in the bundle. The output
+ * of this ParDo is a PCollection of <i>writer result</i> objects (see {@link FileBasedSink}
+ * for a description of writer results)-one for each bundle.
*
* <p>The final do-once ParDo uses a singleton collection asinput and the collection of writer
- * results as a side-input. In this ParDo, {@link WriteOperation#finalize} is called to finalize
- * the write.
+ * results as a side-input. In this ParDo, {@link WriteOperation#finalize} is called
+ * to finalize the write.
*
- * <p>If the write of any element in the PCollection fails, {@link Writer#close} will be called
- * before the exception that caused the write to fail is propagated and the write result will be
- * discarded.
+ * <p>If the write of any element in the PCollection fails, {@link Writer#close} will be
+ * called before the exception that caused the write to fail is propagated and the write result
+ * will be discarded.
*
* <p>Since the {@link WriteOperation} is serialized after the initialization ParDo and
* deserialized in the bundle-writing and finalization phases, any state change to the
- * WriteOperation object that occurs during initialization is visible in the latter phases.
- * However, the WriteOperation is not serialized after the bundle-writing phase. This is why
- * implementations should guarantee that {@link WriteOperation#createWriter} does not mutate
- * WriteOperation).
+ * WriteOperation object that occurs during initialization is visible in the latter
+ * phases. However, the WriteOperation is not serialized after the bundle-writing
+ * phase. This is why implementations should guarantee that
+ * {@link WriteOperation#createWriter} does not mutate WriteOperation).
*/
- private PDone createWrite(PCollection<UserT> input) {
+ private PDone createWrite(PCollection<T> input) {
Pipeline p = input.getPipeline();
if (!windowedWrites) {
// Re-window the data into the global window and remove any existing triggers.
input =
input.apply(
- Window.<UserT>into(new GlobalWindows())
+ Window.<T>into(new GlobalWindows())
.triggering(DefaultTrigger.of())
.discardingFiredPanes());
}
+
// Perform the per-bundle writes as a ParDo on the input PCollection (with the
// WriteOperation as a side input) and collect the results of the writes in a
// PCollection. There is a dependency between this ParDo and the first (the
// WriteOperation PCollection as a side input), so this will happen after the
// initial ParDo.
- PCollection<FileResult<DestinationT>> results;
+ PCollection<FileResult> results;
final PCollectionView<Integer> numShardsView;
- @SuppressWarnings("unchecked")
Coder<BoundedWindow> shardedWindowCoder =
(Coder<BoundedWindow>) input.getWindowingStrategy().getWindowFn().windowCoder();
- final Coder<DestinationT> destinationCoder;
- try {
- destinationCoder =
- sink.getDynamicDestinations()
- .getDestinationCoderWithDefault(input.getPipeline().getCoderRegistry());
- destinationCoder.verifyDeterministic();
- } catch (CannotProvideCoderException | NonDeterministicException e) {
- throw new RuntimeException(e);
- }
-
if (computeNumShards == null && numShardsProvider == null) {
numShardsView = null;
- TupleTag<FileResult<DestinationT>> writtenRecordsTag = new TupleTag<>("writtenRecordsTag");
- TupleTag<KV<ShardedKey<Integer>, UserT>> unwrittedRecordsTag =
- new TupleTag<>("unwrittenRecordsTag");
- String writeName = windowedWrites ? "WriteWindowedBundles" : "WriteBundles";
- PCollectionTuple writeTuple =
- input.apply(
- writeName,
- ParDo.of(new WriteBundles(windowedWrites, unwrittedRecordsTag, destinationCoder))
- .withOutputTags(writtenRecordsTag, TupleTagList.of(unwrittedRecordsTag)));
- PCollection<FileResult<DestinationT>> writtenBundleFiles =
- writeTuple
- .get(writtenRecordsTag)
- .setCoder(FileResultCoder.of(shardedWindowCoder, destinationCoder));
- // Any "spilled" elements are written using WriteShardedBundles. Assign shard numbers in
- // finalize to stay consistent with what WriteWindowedBundles does.
- PCollection<FileResult<DestinationT>> writtenGroupedFiles =
- writeTuple
- .get(unwrittedRecordsTag)
- .setCoder(KvCoder.of(ShardedKeyCoder.of(VarIntCoder.of()), input.getCoder()))
- .apply("GroupUnwritten", GroupByKey.<ShardedKey<Integer>, UserT>create())
- .apply(
- "WriteUnwritten",
- ParDo.of(new WriteShardedBundles(ShardAssignment.ASSIGN_IN_FINALIZE)))
- .setCoder(FileResultCoder.of(shardedWindowCoder, destinationCoder));
results =
- PCollectionList.of(writtenBundleFiles)
- .and(writtenGroupedFiles)
- .apply(Flatten.<FileResult<DestinationT>>pCollections());
+ input.apply(
+ "WriteBundles",
+ ParDo.of(windowedWrites ? new WriteWindowedBundles() : new WriteUnwindowedBundles()));
} else {
List<PCollectionView<?>> sideInputs = Lists.newArrayList();
if (computeNumShards != null) {
@@ -690,31 +509,20 @@ public class WriteFiles<UserT, DestinationT, OutputT>
} else {
numShardsView = null;
}
- PCollection<KV<ShardedKey<Integer>, Iterable<UserT>>> sharded =
+
+ PCollection<KV<Integer, Iterable<T>>> sharded =
input
- .apply(
- "ApplyShardLabel",
- ParDo.of(
- new ApplyShardingKey(
- numShardsView,
- (numShardsView != null) ? null : numShardsProvider,
- destinationCoder))
- .withSideInputs(sideInputs))
- .setCoder(KvCoder.of(ShardedKeyCoder.of(VarIntCoder.of()), input.getCoder()))
- .apply("GroupIntoShards", GroupByKey.<ShardedKey<Integer>, UserT>create());
+ .apply("ApplyShardLabel", ParDo.of(
+ new ApplyShardingKey<T>(numShardsView,
+ (numShardsView != null) ? null : numShardsProvider))
+ .withSideInputs(sideInputs))
+ .apply("GroupIntoShards", GroupByKey.<Integer, T>create());
shardedWindowCoder =
(Coder<BoundedWindow>) sharded.getWindowingStrategy().getWindowFn().windowCoder();
- // Since this path might be used by streaming runners processing triggers, it's important
- // to assign shard numbers here so that they are deterministic. The ASSIGN_IN_FINALIZE
- // strategy works by sorting all FileResult objects and assigning them numbers, which is not
- // guaranteed to work well when processing triggers - if the finalize step retries it might
- // see a different Iterable of FileResult objects, and it will assign different shard numbers.
- results =
- sharded.apply(
- "WriteShardedBundles",
- ParDo.of(new WriteShardedBundles(ShardAssignment.ASSIGN_WHEN_WRITING)));
+
+ results = sharded.apply("WriteShardedBundles", ParDo.of(new WriteShardedBundles()));
}
- results.setCoder(FileResultCoder.of(shardedWindowCoder, destinationCoder));
+ results.setCoder(FileResultCoder.of(shardedWindowCoder));
if (windowedWrites) {
// When processing streaming windowed writes, results will arrive multiple times. This
@@ -722,31 +530,26 @@ public class WriteFiles<UserT, DestinationT, OutputT>
// as new data arriving into a side input does not trigger the listening DoFn. Instead
// we aggregate the result set using a singleton GroupByKey, so the DoFn will be triggered
// whenever new data arrives.
- PCollection<KV<Void, FileResult<DestinationT>>> keyedResults =
- results.apply(
- "AttachSingletonKey", WithKeys.<Void, FileResult<DestinationT>>of((Void) null));
- keyedResults.setCoder(
- KvCoder.of(VoidCoder.of(), FileResultCoder.of(shardedWindowCoder, destinationCoder)));
+ PCollection<KV<Void, FileResult>> keyedResults =
+ results.apply("AttachSingletonKey", WithKeys.<Void, FileResult>of((Void) null));
+ keyedResults.setCoder(KvCoder.of(VoidCoder.of(),
+ FileResultCoder.of(shardedWindowCoder)));
// Is the continuation trigger sufficient?
keyedResults
- .apply("FinalizeGroupByKey", GroupByKey.<Void, FileResult<DestinationT>>create())
- .apply(
- "Finalize",
- ParDo.of(
- new DoFn<KV<Void, Iterable<FileResult<DestinationT>>>, Integer>() {
- @ProcessElement
- public void processElement(ProcessContext c) throws Exception {
- LOG.info("Finalizing write operation {}.", writeOperation);
- List<FileResult<DestinationT>> results =
- Lists.newArrayList(c.element().getValue());
- writeOperation.finalize(results);
- LOG.debug("Done finalizing write operation");
- }
- }));
+ .apply("FinalizeGroupByKey", GroupByKey.<Void, FileResult>create())
+ .apply("Finalize", ParDo.of(new DoFn<KV<Void, Iterable<FileResult>>, Integer>() {
+ @ProcessElement
+ public void processElement(ProcessContext c) throws Exception {
+ LOG.info("Finalizing write operation {}.", writeOperation);
+ List<FileResult> results = Lists.newArrayList(c.element().getValue());
+ writeOperation.finalize(results);
+ LOG.debug("Done finalizing write operation");
+ }
+ }));
} else {
- final PCollectionView<Iterable<FileResult<DestinationT>>> resultsView =
- results.apply(View.<FileResult<DestinationT>>asIterable());
+ final PCollectionView<Iterable<FileResult>> resultsView =
+ results.apply(View.<FileResult>asIterable());
ImmutableList.Builder<PCollectionView<?>> sideInputs =
ImmutableList.<PCollectionView<?>>builder().add(resultsView);
if (numShardsView != null) {
@@ -762,53 +565,41 @@ public class WriteFiles<UserT, DestinationT, OutputT>
// set numShards, then all shards will be written out as empty files. For this reason we
// use a side input here.
PCollection<Void> singletonCollection = p.apply(Create.of((Void) null));
- singletonCollection.apply(
- "Finalize",
- ParDo.of(
- new DoFn<Void, Integer>() {
- @ProcessElement
- public void processElement(ProcessContext c) throws Exception {
- LOG.info("Finalizing write operation {}.", writeOperation);
- List<FileResult<DestinationT>> results =
- Lists.newArrayList(c.sideInput(resultsView));
- LOG.debug(
- "Side input initialized to finalize write operation {}.", writeOperation);
-
- // We must always output at least 1 shard, and honor user-specified numShards
- // if
- // set.
- int minShardsNeeded;
- if (numShardsView != null) {
- minShardsNeeded = c.sideInput(numShardsView);
- } else if (numShardsProvider != null) {
- minShardsNeeded = numShardsProvider.get();
- } else {
- minShardsNeeded = 1;
- }
- int extraShardsNeeded = minShardsNeeded - results.size();
- if (extraShardsNeeded > 0) {
- LOG.info(
- "Creating {} empty output shards in addition to {} written "
- + "for a total of {}.",
- extraShardsNeeded,
- results.size(),
- minShardsNeeded);
- for (int i = 0; i < extraShardsNeeded; ++i) {
- Writer<OutputT, DestinationT> writer = writeOperation.createWriter();
- writer.openUnwindowed(
- UUID.randomUUID().toString(),
- UNKNOWN_SHARDNUM,
- sink.getDynamicDestinations().getDefaultDestination());
- FileResult<DestinationT> emptyWrite = writer.close();
- results.add(emptyWrite);
- }
- LOG.debug("Done creating extra shards.");
- }
- writeOperation.finalize(results);
- LOG.debug("Done finalizing write operation {}", writeOperation);
- }
- })
- .withSideInputs(sideInputs.build()));
+ singletonCollection
+ .apply("Finalize", ParDo.of(new DoFn<Void, Integer>() {
+ @ProcessElement
+ public void processElement(ProcessContext c) throws Exception {
+ LOG.info("Finalizing write operation {}.", writeOperation);
+ List<FileResult> results = Lists.newArrayList(c.sideInput(resultsView));
+ LOG.debug("Side input initialized to finalize write operation {}.", writeOperation);
+
+ // We must always output at least 1 shard, and honor user-specified numShards if
+ // set.
+ int minShardsNeeded;
+ if (numShardsView != null) {
+ minShardsNeeded = c.sideInput(numShardsView);
+ } else if (numShardsProvider != null) {
+ minShardsNeeded = numShardsProvider.get();
+ } else {
+ minShardsNeeded = 1;
+ }
+ int extraShardsNeeded = minShardsNeeded - results.size();
+ if (extraShardsNeeded > 0) {
+ LOG.info(
+ "Creating {} empty output shards in addition to {} written for a total of {}.",
+ extraShardsNeeded, results.size(), minShardsNeeded);
+ for (int i = 0; i < extraShardsNeeded; ++i) {
+ Writer<T> writer = writeOperation.createWriter();
+ writer.openUnwindowed(UUID.randomUUID().toString(), UNKNOWN_SHARDNUM);
+ FileResult emptyWrite = writer.close();
+ results.add(emptyWrite);
+ }
+ LOG.debug("Done creating extra shards.");
+ }
+ writeOperation.finalize(results);
+ LOG.debug("Done finalizing write operation {}", writeOperation);
+ }
+ }).withSideInputs(sideInputs.build()));
}
return PDone.in(input.getPipeline());
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/range/ByteKeyRangeTracker.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/range/ByteKeyRangeTracker.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/range/ByteKeyRangeTracker.java
index b889ec7..99717a4 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/range/ByteKeyRangeTracker.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/range/ByteKeyRangeTracker.java
@@ -71,10 +71,6 @@ public final class ByteKeyRangeTracker implements RangeTracker<ByteKey> {
"Trying to return record which is before the last-returned record");
if (position == null) {
- LOG.info(
- "Adjusting range start from {} to {} as position of first returned record",
- range.getStartKey(),
- recordStart);
range = range.withStartKey(recordStart);
}
position = recordStart;
@@ -91,15 +87,6 @@ public final class ByteKeyRangeTracker implements RangeTracker<ByteKey> {
@Override
public synchronized boolean trySplitAtPosition(ByteKey splitPosition) {
- // Sanity check.
- if (!range.containsKey(splitPosition)) {
- LOG.warn(
- "{}: Rejecting split request at {} because it is not within the range.",
- this,
- splitPosition);
- return false;
- }
-
// Unstarted.
if (position == null) {
LOG.warn(
@@ -119,6 +106,15 @@ public final class ByteKeyRangeTracker implements RangeTracker<ByteKey> {
return false;
}
+ // Sanity check.
+ if (!range.containsKey(splitPosition)) {
+ LOG.warn(
+ "{}: Rejecting split request at {} because it is not within the range.",
+ this,
+ splitPosition);
+ return false;
+ }
+
range = range.withEndKey(splitPosition);
return true;
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/range/OffsetRange.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/range/OffsetRange.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/range/OffsetRange.java
deleted file mode 100644
index d3bff37..0000000
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/range/OffsetRange.java
+++ /dev/null
@@ -1,101 +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.sdk.io.range;
-
-import static com.google.common.base.Preconditions.checkArgument;
-
-import java.io.Serializable;
-import java.util.ArrayList;
-import java.util.List;
-import org.apache.beam.sdk.transforms.splittabledofn.HasDefaultTracker;
-
-/** A restriction represented by a range of integers [from, to). */
-public class OffsetRange
- implements Serializable,
- HasDefaultTracker<
- OffsetRange, org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker> {
- private final long from;
- private final long to;
-
- public OffsetRange(long from, long to) {
- checkArgument(from <= to, "Malformed range [%s, %s)", from, to);
- this.from = from;
- this.to = to;
- }
-
- public long getFrom() {
- return from;
- }
-
- public long getTo() {
- return to;
- }
-
- @Override
- public org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker newTracker() {
- return new org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker(this);
- }
-
- @Override
- public String toString() {
- return "[" + from + ", " + to + ')';
- }
-
- @Override
- public boolean equals(Object o) {
- if (this == o) {
- return true;
- }
- if (o == null || getClass() != o.getClass()) {
- return false;
- }
-
- OffsetRange that = (OffsetRange) o;
-
- if (from != that.from) {
- return false;
- }
- return to == that.to;
- }
-
- @Override
- public int hashCode() {
- int result = (int) (from ^ (from >>> 32));
- result = 31 * result + (int) (to ^ (to >>> 32));
- return result;
- }
-
- public List<OffsetRange> split(long desiredNumOffsetsPerSplit, long minNumOffsetPerSplit) {
- List<OffsetRange> res = new ArrayList<>();
- long start = getFrom();
- long maxEnd = getTo();
-
- while (start < maxEnd) {
- long end = start + desiredNumOffsetsPerSplit;
- end = Math.min(end, maxEnd);
- // Avoid having a too small range at the end and ensure that we respect minNumOffsetPerSplit.
- long remaining = maxEnd - end;
- if ((remaining < desiredNumOffsetsPerSplit / 4) || (remaining < minNumOffsetPerSplit)) {
- end = maxEnd;
- }
- res.add(new OffsetRange(start, end));
- start = end;
- }
- return res;
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/range/OffsetRangeTracker.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/range/OffsetRangeTracker.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/range/OffsetRangeTracker.java
index 8f0083e..51e2b1a 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/range/OffsetRangeTracker.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/range/OffsetRangeTracker.java
@@ -26,9 +26,6 @@ import org.slf4j.LoggerFactory;
/**
* A {@link RangeTracker} for non-negative positions of type {@code long}.
- *
- * <p>Not to be confused with {@link
- * org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker}.
*/
public class OffsetRangeTracker implements RangeTracker<Long> {
private static final Logger LOG = LoggerFactory.getLogger(OffsetRangeTracker.class);
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptionsFactory.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptionsFactory.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptionsFactory.java
index d7e6cc8..c0990cb 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptionsFactory.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptionsFactory.java
@@ -184,20 +184,18 @@ public class PipelineOptionsFactory {
private final String[] args;
private final boolean validation;
private final boolean strictParsing;
- private final boolean isCli;
// Do not allow direct instantiation
private Builder() {
- this(null, false, true, false);
+ this(null, false, true);
}
private Builder(String[] args, boolean validation,
- boolean strictParsing, boolean isCli) {
+ boolean strictParsing) {
this.defaultAppName = findCallersClassName();
this.args = args;
this.validation = validation;
this.strictParsing = strictParsing;
- this.isCli = isCli;
}
/**
@@ -239,7 +237,7 @@ public class PipelineOptionsFactory {
*/
public Builder fromArgs(String... args) {
checkNotNull(args, "Arguments should not be null.");
- return new Builder(args, validation, strictParsing, true);
+ return new Builder(args, validation, strictParsing);
}
/**
@@ -249,7 +247,7 @@ public class PipelineOptionsFactory {
* validation.
*/
public Builder withValidation() {
- return new Builder(args, true, strictParsing, isCli);
+ return new Builder(args, true, strictParsing);
}
/**
@@ -257,7 +255,7 @@ public class PipelineOptionsFactory {
* arguments.
*/
public Builder withoutStrictParsing() {
- return new Builder(args, validation, false, isCli);
+ return new Builder(args, validation, false);
}
/**
@@ -302,11 +300,7 @@ public class PipelineOptionsFactory {
}
if (validation) {
- if (isCli) {
- PipelineOptionsValidator.validateCli(klass, t);
- } else {
- PipelineOptionsValidator.validate(klass, t);
- }
+ PipelineOptionsValidator.validate(klass, t);
}
return t;
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptionsValidator.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptionsValidator.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptionsValidator.java
index fcffd74..bd54ec3 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptionsValidator.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptionsValidator.java
@@ -43,29 +43,9 @@ public class PipelineOptionsValidator {
*
* @param klass The interface to fetch validation criteria from.
* @param options The {@link PipelineOptions} to validate.
- * @return Validated options.
+ * @return The type
*/
public static <T extends PipelineOptions> T validate(Class<T> klass, PipelineOptions options) {
- return validate(klass, options, false);
- }
-
- /**
- * Validates that the passed {@link PipelineOptions} from command line interface (CLI)
- * conforms to all the validation criteria from the passed in interface.
- *
- * <p>Note that the interface requested must conform to the validation criteria specified on
- * {@link PipelineOptions#as(Class)}.
- *
- * @param klass The interface to fetch validation criteria from.
- * @param options The {@link PipelineOptions} to validate.
- * @return Validated options.
- */
- public static <T extends PipelineOptions> T validateCli(Class<T> klass, PipelineOptions options) {
- return validate(klass, options, true);
- }
-
- private static <T extends PipelineOptions> T validate(Class<T> klass, PipelineOptions options,
- boolean isCli) {
checkNotNull(klass);
checkNotNull(options);
checkArgument(Proxy.isProxyClass(options.getClass()));
@@ -87,15 +67,9 @@ public class PipelineOptionsValidator {
requiredGroups.put(requiredGroup, method);
}
} else {
- if (isCli) {
- checkArgument(handler.invoke(asClassOptions, method, null) != null,
- "Missing required value for [--%s, \"%s\"]. ",
- handler.getOptionName(method), getDescription(method));
- } else {
- checkArgument(handler.invoke(asClassOptions, method, null) != null,
- "Missing required value for [%s, \"%s\"]. ",
- method, getDescription(method));
- }
+ checkArgument(handler.invoke(asClassOptions, method, null) != null,
+ "Missing required value for [%s, \"%s\"]. ",
+ method, getDescription(method));
}
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/ProxyInvocationHandler.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/ProxyInvocationHandler.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/ProxyInvocationHandler.java
index 926a7b9..eda21a8 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/ProxyInvocationHandler.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/ProxyInvocationHandler.java
@@ -45,8 +45,6 @@ import com.google.common.collect.Multimap;
import com.google.common.collect.MutableClassToInstanceMap;
import java.beans.PropertyDescriptor;
import java.io.IOException;
-import java.io.NotSerializableException;
-import java.io.Serializable;
import java.lang.annotation.Annotation;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
@@ -89,7 +87,7 @@ import org.apache.beam.sdk.util.common.ReflectHelpers;
* {@link PipelineOptions#as(Class)}.
*/
@ThreadSafe
-class ProxyInvocationHandler implements InvocationHandler, Serializable {
+class ProxyInvocationHandler implements InvocationHandler {
/**
* No two instances of this class are considered equivalent hence we generate a random hash code.
*/
@@ -166,21 +164,6 @@ class ProxyInvocationHandler implements InvocationHandler, Serializable {
+ Arrays.toString(args) + "].");
}
- public String getOptionName(Method method) {
- return gettersToPropertyNames.get(method.getName());
- }
-
- private void writeObject(java.io.ObjectOutputStream stream)
- throws IOException {
- throw new NotSerializableException(
- "PipelineOptions objects are not serializable and should not be embedded into transforms "
- + "(did you capture a PipelineOptions object in a field or in an anonymous class?). "
- + "Instead, if you're using a DoFn, access PipelineOptions at runtime "
- + "via ProcessContext/StartBundleContext/FinishBundleContext.getPipelineOptions(), "
- + "or pre-extract necessary fields from PipelineOptions "
- + "at pipeline construction time.");
- }
-
/**
* Track whether options values are explicitly set, or retrieved from defaults.
*/
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/TransformHierarchy.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/TransformHierarchy.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/TransformHierarchy.java
index d8ff59e..2f0e8ef 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/TransformHierarchy.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/TransformHierarchy.java
@@ -24,23 +24,21 @@ import static com.google.common.base.Preconditions.checkState;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.MoreObjects;
import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Iterables;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
-import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
+import javax.annotation.Nullable;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.Pipeline.PipelineVisitor;
import org.apache.beam.sdk.Pipeline.PipelineVisitor.CompositeBehavior;
import org.apache.beam.sdk.annotations.Internal;
import org.apache.beam.sdk.runners.PTransformOverrideFactory.ReplacementOutput;
import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PInput;
import org.apache.beam.sdk.values.POutput;
import org.apache.beam.sdk.values.PValue;
@@ -70,7 +68,7 @@ public class TransformHierarchy {
producers = new HashMap<>();
producerInput = new HashMap<>();
unexpandedInputs = new HashMap<>();
- root = new Node();
+ root = new Node(null, null, "", null);
current = root;
}
@@ -145,6 +143,14 @@ public class TransformHierarchy {
Node producerNode = getProducer(inputValue);
PInput input = producerInput.remove(inputValue);
inputValue.finishSpecifying(input, producerNode.getTransform());
+ checkState(
+ producers.get(inputValue) != null,
+ "Producer unknown for input %s",
+ inputValue);
+ checkState(
+ producers.get(inputValue) != null,
+ "Producer unknown for input %s",
+ inputValue);
}
}
@@ -159,7 +165,7 @@ public class TransformHierarchy {
* nodes.
*/
public void setOutput(POutput output) {
- for (PCollection<?> value : fullyExpand(output).values()) {
+ for (PValue value : output.expand().values()) {
if (!producers.containsKey(value)) {
producers.put(value, current);
value.finishSpecifyingOutput(
@@ -193,13 +199,13 @@ public class TransformHierarchy {
}
Node getProducer(PValue produced) {
- return checkNotNull(producers.get(produced), "No producer found for %s", produced);
+ return producers.get(produced);
}
public Set<PValue> visit(PipelineVisitor visitor) {
finishSpecifying();
Set<PValue> visitedValues = new HashSet<>();
- root.visit(visitor, visitedValues, new HashSet<Node>(), new HashSet<Node>());
+ root.visit(visitor, visitedValues);
return visitedValues;
}
@@ -220,47 +226,6 @@ public class TransformHierarchy {
return current;
}
- private Map<TupleTag<?>, PCollection<?>> fullyExpand(POutput output) {
- Map<TupleTag<?>, PCollection<?>> result = new LinkedHashMap<>();
- for (Map.Entry<TupleTag<?>, PValue> value : output.expand().entrySet()) {
- if (value.getValue() instanceof PCollection) {
- PCollection<?> previous = result.put(value.getKey(), (PCollection<?>) value.getValue());
- checkArgument(
- previous == null,
- "Found conflicting %ss in flattened expansion of %s: %s maps to %s and %s",
- output,
- TupleTag.class.getSimpleName(),
- value.getKey(),
- previous,
- value.getValue());
- } else {
- if (value.getValue().expand().size() == 1
- && Iterables.getOnlyElement(value.getValue().expand().values())
- .equals(value.getValue())) {
- throw new IllegalStateException(
- String.format(
- "Non %s %s that expands into itself %s",
- PCollection.class.getSimpleName(),
- PValue.class.getSimpleName(),
- value.getValue()));
- }
- for (Map.Entry<TupleTag<?>, PCollection<?>> valueComponent :
- fullyExpand(value.getValue()).entrySet()) {
- PCollection<?> previous = result.put(valueComponent.getKey(), valueComponent.getValue());
- checkArgument(
- previous == null,
- "Found conflicting %ss in flattened expansion of %s: %s maps to %s and %s",
- output,
- TupleTag.class.getSimpleName(),
- valueComponent.getKey(),
- previous,
- valueComponent.getValue());
- }
- }
- }
- return result;
- }
-
/**
* Provides internal tracking of transform relationships with helper methods
* for initialization and ordered visitation.
@@ -288,36 +253,25 @@ public class TransformHierarchy {
boolean finishedSpecifying = false;
/**
- * Creates the root-level node. The root level node has a null enclosing node, a null transform,
- * an empty map of inputs, and a name equal to the empty string.
- */
- private Node() {
- this.enclosingNode = null;
- this.transform = null;
- this.fullName = "";
- this.inputs = Collections.emptyMap();
- }
-
- /**
* Creates a new Node with the given parent and transform.
*
+ * <p>EnclosingNode and transform may both be null for a root-level node, which holds all other
+ * nodes.
+ *
* @param enclosingNode the composite node containing this node
* @param transform the PTransform tracked by this node
* @param fullName the fully qualified name of the transform
* @param input the unexpanded input to the transform
*/
private Node(
- Node enclosingNode,
- PTransform<?, ?> transform,
+ @Nullable Node enclosingNode,
+ @Nullable PTransform<?, ?> transform,
String fullName,
- PInput input) {
+ @Nullable PInput input) {
this.enclosingNode = enclosingNode;
this.transform = transform;
this.fullName = fullName;
- ImmutableMap.Builder<TupleTag<?>, PValue> inputs = ImmutableMap.builder();
- inputs.putAll(input.expand());
- inputs.putAll(transform.getAdditionalInputs());
- this.inputs = inputs.build();
+ this.inputs = input == null ? Collections.<TupleTag<?>, PValue>emptyMap() : input.expand();
}
/**
@@ -398,7 +352,7 @@ public class TransformHierarchy {
return fullName;
}
- /** Returns the transform input, in fully expanded form. */
+ /** Returns the transform input, in unexpanded form. */
public Map<TupleTag<?>, PValue> getInputs() {
return inputs == null ? Collections.<TupleTag<?>, PValue>emptyMap() : inputs;
}
@@ -505,60 +459,10 @@ public class TransformHierarchy {
/**
* Visit the transform node.
*
- * <p>The visit proceeds in the following order:
- *
- * <ul>
- * <li>Visit all input {@link PValue PValues} returned by the flattened expansion of {@link
- * Node#getInputs()}.
- * <li>If the node is a composite:
- * <ul>
- * <li>Enter the node via {@link PipelineVisitor#enterCompositeTransform(Node)}.
- * <li>If the result of {@link PipelineVisitor#enterCompositeTransform(Node)} was {@link
- * CompositeBehavior#ENTER_TRANSFORM}, visit each child node of this {@link Node}.
- * <li>Leave the node via {@link PipelineVisitor#leaveCompositeTransform(Node)}.
- * </ul>
- * <li>If the node is a primitive, visit it via {@link
- * PipelineVisitor#visitPrimitiveTransform(Node)}.
- * <li>Visit each {@link PValue} that was output by this node.
- * </ul>
- *
- * <p>Additionally, the following ordering restrictions are observed:
- *
- * <ul>
- * <li>A {@link Node} will be visited after its enclosing node has been entered and before its
- * enclosing node has been left
- * <li>A {@link Node} will not be visited if any enclosing {@link Node} has returned {@link
- * CompositeBehavior#DO_NOT_ENTER_TRANSFORM} from the call to {@link
- * PipelineVisitor#enterCompositeTransform(Node)}.
- * <li>A {@link PValue} will only be visited after the {@link Node} that originally produced
- * it has been visited.
- * </ul>
- *
* <p>Provides an ordered visit of the input values, the primitive transform (or child nodes for
* composite transforms), then the output values.
*/
- private void visit(
- PipelineVisitor visitor,
- Set<PValue> visitedValues,
- Set<Node> visitedNodes,
- Set<Node> skippedComposites) {
- if (getEnclosingNode() != null && !visitedNodes.contains(getEnclosingNode())) {
- // Recursively enter all enclosing nodes, as appropriate.
- getEnclosingNode().visit(visitor, visitedValues, visitedNodes, skippedComposites);
- }
- // These checks occur after visiting the enclosing node to ensure that if this node has been
- // visited while visiting the enclosing node the node is not revisited, or, if an enclosing
- // Node is skipped, this node is also skipped.
- if (!visitedNodes.add(this)) {
- LOG.debug("Not revisiting previously visited node {}", this);
- return;
- } else if (childNodeOf(skippedComposites)) {
- // This node is a child of a node that has been passed over via CompositeBehavior, and
- // should also be skipped. All child nodes of a skipped composite should always be skipped.
- LOG.debug("Not revisiting Node {} which is a child of a previously passed composite", this);
- return;
- }
-
+ private void visit(PipelineVisitor visitor, Set<PValue> visitedValues) {
if (!finishedSpecifying) {
finishSpecifying();
}
@@ -566,31 +470,22 @@ public class TransformHierarchy {
if (!isRootNode()) {
// Visit inputs.
for (PValue inputValue : inputs.values()) {
- Node valueProducer = getProducer(inputValue);
- if (!visitedNodes.contains(valueProducer)) {
- valueProducer.visit(visitor, visitedValues, visitedNodes, skippedComposites);
- }
if (visitedValues.add(inputValue)) {
- LOG.debug("Visiting input value {}", inputValue);
- visitor.visitValue(inputValue, valueProducer);
+ visitor.visitValue(inputValue, getProducer(inputValue));
}
}
}
if (isCompositeNode()) {
- LOG.debug("Visiting composite node {}", this);
PipelineVisitor.CompositeBehavior recurse = visitor.enterCompositeTransform(this);
if (recurse.equals(CompositeBehavior.ENTER_TRANSFORM)) {
for (Node child : parts) {
- child.visit(visitor, visitedValues, visitedNodes, skippedComposites);
+ child.visit(visitor, visitedValues);
}
- } else {
- skippedComposites.add(this);
}
visitor.leaveCompositeTransform(this);
} else {
- LOG.debug("Visiting primitive node {}", this);
visitor.visitPrimitiveTransform(this);
}
@@ -599,24 +494,12 @@ public class TransformHierarchy {
// Visit outputs.
for (PValue pValue : outputs.values()) {
if (visitedValues.add(pValue)) {
- LOG.debug("Visiting output value {}", pValue);
visitor.visitValue(pValue, this);
}
}
}
}
- private boolean childNodeOf(Set<Node> nodes) {
- if (isRootNode()) {
- return false;
- }
- Node parent = this.getEnclosingNode();
- while (!parent.isRootNode() && !nodes.contains(parent)) {
- parent = parent.getEnclosingNode();
- }
- return nodes.contains(parent);
- }
-
/**
* Finish specifying a transform.
*
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/StaticWindows.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/StaticWindows.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/StaticWindows.java
index eba6978..c11057a 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/StaticWindows.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/StaticWindows.java
@@ -126,9 +126,4 @@ final class StaticWindows extends NonMergingWindowFn<Object, BoundedWindow> {
}
};
}
-
- @Override
- public boolean assignsToOneWindow() {
- return true;
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/TestStream.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/TestStream.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/TestStream.java
index d13fcf1..9ad8fd8 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/TestStream.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/TestStream.java
@@ -271,18 +271,6 @@ public final class TestStream<T> extends PTransform<PBegin, PCollection<T>> {
return events;
}
- /**
- * <b>For internal use only. No backwards-compatibility guarantees.</b>
- *
- * <p>Builder a test stream directly from events. No validation is performed on
- * watermark monotonicity, etc. This is assumed to be a previously-serialized
- * {@link TestStream} transform that is correct by construction.
- */
- @Internal
- public static <T> TestStream<T> fromRawEvents(Coder<T> coder, List<Event<T>> events) {
- return new TestStream<>(coder, events);
- }
-
@Override
public boolean equals(Object other) {
if (!(other instanceof TestStream)) {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Combine.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Combine.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Combine.java
index d7effb5..9e1cc71 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Combine.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Combine.java
@@ -20,6 +20,7 @@ package org.apache.beam.sdk.transforms;
import static com.google.common.base.Preconditions.checkState;
import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import java.io.IOException;
import java.io.InputStream;
@@ -1121,7 +1122,11 @@ public class Combine {
*/
@Override
public Map<TupleTag<?>, PValue> getAdditionalInputs() {
- return PCollectionViews.toAdditionalInputs(sideInputs);
+ ImmutableMap.Builder<TupleTag<?>, PValue> additionalInputs = ImmutableMap.builder();
+ for (PCollectionView<?> sideInput : sideInputs) {
+ additionalInputs.put(sideInput.getTagInternal(), sideInput.getPCollection());
+ }
+ return additionalInputs.build();
}
/**
@@ -1272,15 +1277,14 @@ public class Combine {
public PCollectionView<OutputT> expand(PCollection<InputT> input) {
PCollection<OutputT> combined =
input.apply(Combine.<InputT, OutputT>globally(fn).withoutDefaults().withFanout(fanout));
- PCollectionView<OutputT> view =
- PCollectionViews.singletonView(
- combined,
- input.getWindowingStrategy(),
- insertDefault,
- insertDefault ? fn.defaultValue() : null,
- combined.getCoder());
- combined.apply(CreatePCollectionView.<OutputT, OutputT>of(view));
- return view;
+ return combined.apply(
+ CreatePCollectionView.<OutputT, OutputT>of(
+ PCollectionViews.singletonView(
+ combined,
+ input.getWindowingStrategy(),
+ insertDefault,
+ insertDefault ? fn.defaultValue() : null,
+ combined.getCoder())));
}
public int getFanout() {
@@ -1573,7 +1577,11 @@ public class Combine {
*/
@Override
public Map<TupleTag<?>, PValue> getAdditionalInputs() {
- return PCollectionViews.toAdditionalInputs(sideInputs);
+ ImmutableMap.Builder<TupleTag<?>, PValue> additionalInputs = ImmutableMap.builder();
+ for (PCollectionView<?> sideInput : sideInputs) {
+ additionalInputs.put(sideInput.getTagInternal(), sideInput.getPCollection());
+ }
+ return additionalInputs.build();
}
@Override
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java
index 1b809c2..e711ac2 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java
@@ -17,7 +17,6 @@
*/
package org.apache.beam.sdk.transforms;
-import com.google.auto.value.AutoValue;
import java.io.Serializable;
import java.lang.annotation.Documented;
import java.lang.annotation.ElementType;
@@ -386,7 +385,7 @@ public abstract class DoFn<InputT, OutputT> implements Serializable, HasDisplayD
* <pre><code>{@literal new DoFn<KV<Key, Foo>, Baz>()} {
*
* {@literal @StateId("my-state-id")}
- * {@literal private final StateSpec<ValueState<MyState>>} myStateSpec =
+ * {@literal private final StateSpec<K, ValueState<MyState>>} myStateSpec =
* StateSpecs.value(new MyStateCoder());
*
* {@literal @ProcessElement}
@@ -546,15 +545,11 @@ public abstract class DoFn<InputT, OutputT> implements Serializable, HasDisplayD
* returned by {@link GetInitialRestriction} implements {@link HasDefaultTracker}.
* <li>It <i>may</i> define a {@link GetRestrictionCoder} method.
* <li>The type of restrictions used by all of these methods must be the same.
- * <li>Its {@link ProcessElement} method <i>may</i> return a {@link ProcessContinuation} to
- * indicate whether there is more work to be done for the current element.
* <li>Its {@link ProcessElement} method <i>must not</i> use any extra context parameters, such as
* {@link BoundedWindow}.
* <li>The {@link DoFn} itself <i>may</i> be annotated with {@link BoundedPerElement} or
* {@link UnboundedPerElement}, but not both at the same time. If it's not annotated with
- * either of these, it's assumed to be {@link BoundedPerElement} if its {@link
- * ProcessElement} method returns {@code void} and {@link UnboundedPerElement} if it
- * returns a {@link ProcessContinuation}.
+ * either of these, it's assumed to be {@link BoundedPerElement}.
* </ul>
*
* <p>A non-splittable {@link DoFn} <i>must not</i> define any of these methods.
@@ -682,49 +677,6 @@ public abstract class DoFn<InputT, OutputT> implements Serializable, HasDisplayD
@Experimental(Kind.SPLITTABLE_DO_FN)
public @interface UnboundedPerElement {}
- // This can't be put into ProcessContinuation itself due to the following problem:
- // http://ternarysearch.blogspot.com/2013/07/static-initialization-deadlock.html
- private static final ProcessContinuation PROCESS_CONTINUATION_STOP =
- new AutoValue_DoFn_ProcessContinuation(false, Duration.ZERO);
-
- /**
- * When used as a return value of {@link ProcessElement}, indicates whether there is more work to
- * be done for the current element.
- *
- * <p>If the {@link ProcessElement} call completes because of a failed {@code tryClaim()} call
- * on the {@link RestrictionTracker}, then the call MUST return {@link #stop()}.
- */
- @Experimental(Kind.SPLITTABLE_DO_FN)
- @AutoValue
- public abstract static class ProcessContinuation {
- /** Indicates that there is no more work to be done for the current element. */
- public static ProcessContinuation stop() {
- return PROCESS_CONTINUATION_STOP;
- }
-
- /** Indicates that there is more work to be done for the current element. */
- public static ProcessContinuation resume() {
- return new AutoValue_DoFn_ProcessContinuation(true, Duration.ZERO);
- }
-
- /**
- * If false, the {@link DoFn} promises that there is no more work remaining for the current
- * element, so the runner should not resume the {@link ProcessElement} call.
- */
- public abstract boolean shouldResume();
-
- /**
- * A minimum duration that should elapse between the end of this {@link ProcessElement} call and
- * the {@link ProcessElement} call continuing processing of the same element. By default, zero.
- */
- public abstract Duration resumeDelay();
-
- /** Builder method to set the value of {@link #resumeDelay()}. */
- public ProcessContinuation withResumeDelay(Duration resumeDelay) {
- return new AutoValue_DoFn_ProcessContinuation(shouldResume(), resumeDelay);
- }
- }
-
/**
* Finalize the {@link DoFn} construction to prepare for processing.
* This method should be called by runners before any processing methods.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java
index b2377dd..8a03f3c 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java
@@ -290,11 +290,6 @@ public class DoFnTester<InputT, OutputT> implements AutoCloseable {
}
@Override
- public PipelineOptions pipelineOptions() {
- return getPipelineOptions();
- }
-
- @Override
public DoFn<InputT, OutputT>.StartBundleContext startBundleContext(
DoFn<InputT, OutputT> doFn) {
throw new UnsupportedOperationException(
@@ -551,6 +546,11 @@ public class DoFnTester<InputT, OutputT> implements AutoCloseable {
fn.super();
}
+ private void throwUnsupportedOutputFromBundleMethods() {
+ throw new UnsupportedOperationException(
+ "DoFnTester doesn't support output from bundle methods");
+ }
+
@Override
public PipelineOptions getPipelineOptions() {
return options;
@@ -559,13 +559,12 @@ public class DoFnTester<InputT, OutputT> implements AutoCloseable {
@Override
public void output(
OutputT output, Instant timestamp, BoundedWindow window) {
- output(mainOutputTag, output, timestamp, window);
+ throwUnsupportedOutputFromBundleMethods();
}
@Override
public <T> void output(TupleTag<T> tag, T output, Instant timestamp, BoundedWindow window) {
- getMutableOutput(tag)
- .add(ValueInSingleWindow.of(output, timestamp, window, PaneInfo.NO_FIRING));
+ throwUnsupportedOutputFromBundleMethods();
}
}
@@ -643,6 +642,12 @@ public class DoFnTester<InputT, OutputT> implements AutoCloseable {
getMutableOutput(tag)
.add(ValueInSingleWindow.of(output, timestamp, element.getWindow(), element.getPane()));
}
+
+ private void throwUnsupportedOutputFromBundleMethods() {
+ throw new UnsupportedOperationException(
+ "DoFnTester doesn't support output from bundle methods");
+ }
+
}
@Override
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java
index 0d03835..edf1419 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java
@@ -20,6 +20,7 @@ package org.apache.beam.sdk.transforms;
import static com.google.common.base.Preconditions.checkArgument;
import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
import java.io.Serializable;
import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Type;
@@ -32,7 +33,6 @@ import org.apache.beam.sdk.PipelineRunner;
import org.apache.beam.sdk.coders.CannotProvideCoderException;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderRegistry;
-import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.transforms.DoFn.WindowedContext;
import org.apache.beam.sdk.transforms.display.DisplayData;
@@ -50,7 +50,6 @@ import org.apache.beam.sdk.util.SerializableUtils;
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.PCollectionViews;
import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
@@ -456,27 +455,6 @@ public class ParDo {
}
}
- private static void validateStateApplicableForInput(
- DoFn<?, ?> fn,
- PCollection<?> input) {
- Coder<?> inputCoder = input.getCoder();
- checkArgument(
- inputCoder instanceof KvCoder,
- "%s requires its input to use %s in order to use state and timers.",
- ParDo.class.getSimpleName(),
- KvCoder.class.getSimpleName());
-
- KvCoder<?, ?> kvCoder = (KvCoder<?, ?>) inputCoder;
- try {
- kvCoder.getKeyCoder().verifyDeterministic();
- } catch (Coder.NonDeterministicException exc) {
- throw new IllegalArgumentException(
- String.format(
- "%s requires a deterministic key coder in order to use state and timers",
- ParDo.class.getSimpleName()));
- }
- }
-
/**
* Try to provide coders for as many of the type arguments of given
* {@link DoFnSignature.StateDeclaration} as possible.
@@ -684,7 +662,11 @@ public class ParDo {
*/
@Override
public Map<TupleTag<?>, PValue> getAdditionalInputs() {
- return PCollectionViews.toAdditionalInputs(sideInputs);
+ ImmutableMap.Builder<TupleTag<?>, PValue> additionalInputs = ImmutableMap.builder();
+ for (PCollectionView<?> sideInput : sideInputs) {
+ additionalInputs.put(sideInput.getTagInternal(), sideInput.getPCollection());
+ }
+ return additionalInputs.build();
}
}
@@ -759,11 +741,6 @@ public class ParDo {
// Use coder registry to determine coders for all StateSpec defined in the fn signature.
finishSpecifyingStateSpecs(fn, input.getPipeline().getCoderRegistry(), input.getCoder());
- DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());
- if (signature.usesState() || signature.usesTimers()) {
- validateStateApplicableForInput(fn, input);
- }
-
PCollectionTuple outputs = PCollectionTuple.ofPrimitiveOutputsInternal(
input.getPipeline(),
TupleTagList.of(mainOutputTag).and(additionalOutputTags.getAll()),
@@ -830,7 +807,11 @@ public class ParDo {
*/
@Override
public Map<TupleTag<?>, PValue> getAdditionalInputs() {
- return PCollectionViews.toAdditionalInputs(sideInputs);
+ ImmutableMap.Builder<TupleTag<?>, PValue> additionalInputs = ImmutableMap.builder();
+ for (PCollectionView<?> sideInput : sideInputs) {
+ additionalInputs.put(sideInput.getTagInternal(), sideInput.getPCollection());
+ }
+ return additionalInputs.build();
}
}
[13/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/GroupByKeyTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/GroupByKeyTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/GroupByKeyTest.java
index 4b5d5f5..171171f 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/GroupByKeyTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/GroupByKeyTest.java
@@ -45,19 +45,14 @@ import org.apache.beam.sdk.coders.CoderProviders;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.MapCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.testing.LargeKeys;
import org.apache.beam.sdk.testing.NeedsRunner;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
-import org.apache.beam.sdk.testing.TestStream;
-import org.apache.beam.sdk.testing.UsesTestStream;
import org.apache.beam.sdk.testing.ValidatesRunner;
import org.apache.beam.sdk.transforms.display.DisplayData;
-import org.apache.beam.sdk.transforms.windowing.AfterProcessingTime;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
import org.apache.beam.sdk.transforms.windowing.InvalidWindows;
-import org.apache.beam.sdk.transforms.windowing.Repeatedly;
import org.apache.beam.sdk.transforms.windowing.Sessions;
import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
import org.apache.beam.sdk.transforms.windowing.Window;
@@ -189,40 +184,6 @@ public class GroupByKeyTest {
p.run();
}
- /**
- * Tests that when a processing time timers comes in after a window is expired it does not cause a
- * spurious output.
- */
- @Test
- @Category({ValidatesRunner.class, UsesTestStream.class})
- public void testCombiningAccumulatingProcessingTime() throws Exception {
- PCollection<Integer> triggeredSums =
- p.apply(
- TestStream.create(VarIntCoder.of())
- .advanceWatermarkTo(new Instant(0))
- .addElements(
- TimestampedValue.of(2, new Instant(2)),
- TimestampedValue.of(5, new Instant(5)))
- .advanceWatermarkTo(new Instant(100))
- .advanceProcessingTime(Duration.millis(10))
- .advanceWatermarkToInfinity())
- .apply(
- Window.<Integer>into(FixedWindows.of(Duration.millis(100)))
- .withTimestampCombiner(TimestampCombiner.EARLIEST)
- .accumulatingFiredPanes()
- .withAllowedLateness(Duration.ZERO)
- .triggering(
- Repeatedly.forever(
- AfterProcessingTime.pastFirstElementInPane()
- .plusDelayOf(Duration.millis(10)))))
- .apply(Sum.integersGlobally().withoutDefaults());
-
- PAssert.that(triggeredSums)
- .containsInAnyOrder(7);
-
- p.run();
- }
-
@Test
public void testGroupByKeyNonDeterministic() throws Exception {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java
index fa4949e..c67cf2a 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java
@@ -62,8 +62,6 @@ import org.apache.beam.sdk.coders.SetCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.io.GenerateSequence;
-import org.apache.beam.sdk.options.Default;
-import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.state.BagState;
import org.apache.beam.sdk.state.CombiningState;
import org.apache.beam.sdk.state.MapState;
@@ -1593,108 +1591,6 @@ public class ParDoTest implements Serializable {
}
@Test
- public void testStateNotKeyed() {
- final String stateId = "foo";
-
- DoFn<String, Integer> fn =
- new DoFn<String, Integer>() {
-
- @StateId(stateId)
- private final StateSpec<ValueState<Integer>> intState =
- StateSpecs.value();
-
- @ProcessElement
- public void processElement(
- ProcessContext c, @StateId(stateId) ValueState<Integer> state) {}
- };
-
- thrown.expect(IllegalArgumentException.class);
- thrown.expectMessage("state");
- thrown.expectMessage("KvCoder");
-
- pipeline.apply(Create.of("hello", "goodbye", "hello again")).apply(ParDo.of(fn));
- }
-
- @Test
- public void testStateNotDeterministic() {
- final String stateId = "foo";
-
- // DoubleCoder is not deterministic, so this should crash
- DoFn<KV<Double, String>, Integer> fn =
- new DoFn<KV<Double, String>, Integer>() {
-
- @StateId(stateId)
- private final StateSpec<ValueState<Integer>> intState =
- StateSpecs.value();
-
- @ProcessElement
- public void processElement(
- ProcessContext c, @StateId(stateId) ValueState<Integer> state) {}
- };
-
- thrown.expect(IllegalArgumentException.class);
- thrown.expectMessage("state");
- thrown.expectMessage("deterministic");
-
- pipeline
- .apply(Create.of(KV.of(1.0, "hello"), KV.of(5.4, "goodbye"), KV.of(7.2, "hello again")))
- .apply(ParDo.of(fn));
- }
-
- @Test
- public void testTimerNotKeyed() {
- final String timerId = "foo";
-
- DoFn<String, Integer> fn =
- new DoFn<String, Integer>() {
-
- @TimerId(timerId)
- private final TimerSpec timer = TimerSpecs.timer(TimeDomain.EVENT_TIME);
-
- @ProcessElement
- public void processElement(
- ProcessContext c, @TimerId(timerId) Timer timer) {}
-
- @OnTimer(timerId)
- public void onTimer() {}
- };
-
- thrown.expect(IllegalArgumentException.class);
- thrown.expectMessage("timer");
- thrown.expectMessage("KvCoder");
-
- pipeline.apply(Create.of("hello", "goodbye", "hello again")).apply(ParDo.of(fn));
- }
-
- @Test
- public void testTimerNotDeterministic() {
- final String timerId = "foo";
-
- // DoubleCoder is not deterministic, so this should crash
- DoFn<KV<Double, String>, Integer> fn =
- new DoFn<KV<Double, String>, Integer>() {
-
- @TimerId(timerId)
- private final TimerSpec timer = TimerSpecs.timer(TimeDomain.EVENT_TIME);
-
- @ProcessElement
- public void processElement(
- ProcessContext c, @TimerId(timerId) Timer timer) {}
-
- @OnTimer(timerId)
- public void onTimer() {}
- };
-
- thrown.expect(IllegalArgumentException.class);
- thrown.expectMessage("timer");
- thrown.expectMessage("deterministic");
-
- pipeline
- .apply(Create.of(KV.of(1.0, "hello"), KV.of(5.4, "goodbye"), KV.of(7.2, "hello again")))
- .apply(ParDo.of(fn));
- }
-
- @Test
@Category({ValidatesRunner.class, UsesStatefulParDo.class})
public void testValueStateCoderInference() {
final String stateId = "foo";
@@ -3046,65 +2942,4 @@ public class ParDoTest implements Serializable {
// If it doesn't crash, we made it!
}
-
- /** A {@link PipelineOptions} subclass for testing passing to a {@link DoFn}. */
- public interface MyOptions extends PipelineOptions {
- @Default.String("fake option")
- String getFakeOption();
- void setFakeOption(String value);
- }
-
- @Test
- @Category(ValidatesRunner.class)
- public void testPipelineOptionsParameter() {
- PCollection<String> results = pipeline
- .apply(Create.of(1))
- .apply(
- ParDo.of(
- new DoFn<Integer, String>() {
- @ProcessElement
- public void process(ProcessContext c, PipelineOptions options) {
- c.output(options.as(MyOptions.class).getFakeOption());
- }
- }));
-
- String testOptionValue = "not fake anymore";
- pipeline.getOptions().as(MyOptions.class).setFakeOption(testOptionValue);
- PAssert.that(results).containsInAnyOrder("not fake anymore");
-
- pipeline.run();
- }
-
- @Test
- @Category({ValidatesRunner.class, UsesTimersInParDo.class})
- public void testPipelineOptionsParameterOnTimer() {
- final String timerId = "thisTimer";
-
- PCollection<String> results =
- pipeline
- .apply(Create.of(KV.of(0, 0)))
- .apply(
- ParDo.of(
- new DoFn<KV<Integer, Integer>, String>() {
- @TimerId(timerId)
- private final TimerSpec spec = TimerSpecs.timer(TimeDomain.EVENT_TIME);
-
- @ProcessElement
- public void process(
- ProcessContext c, BoundedWindow w, @TimerId(timerId) Timer timer) {
- timer.set(w.maxTimestamp());
- }
-
- @OnTimer(timerId)
- public void onTimer(OnTimerContext c, PipelineOptions options) {
- c.output(options.as(MyOptions.class).getFakeOption());
- }
- }));
-
- String testOptionValue = "not fake anymore";
- pipeline.getOptions().as(MyOptions.class).setFakeOption(testOptionValue);
- PAssert.that(results).containsInAnyOrder("not fake anymore");
-
- pipeline.run();
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/SplittableDoFnTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/SplittableDoFnTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/SplittableDoFnTest.java
index d2d2529..02a44d2 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/SplittableDoFnTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/SplittableDoFnTest.java
@@ -18,14 +18,10 @@
package org.apache.beam.sdk.transforms;
import static com.google.common.base.Preconditions.checkState;
-import static org.apache.beam.sdk.testing.TestPipeline.testingPipelineOptions;
-import static org.apache.beam.sdk.transforms.DoFn.ProcessContinuation.resume;
-import static org.apache.beam.sdk.transforms.DoFn.ProcessContinuation.stop;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
-import com.google.common.collect.Ordering;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.Arrays;
@@ -33,9 +29,6 @@ import java.util.List;
import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.sdk.io.range.OffsetRange;
-import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.options.StreamingOptions;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.testing.TestStream;
@@ -44,6 +37,7 @@ import org.apache.beam.sdk.testing.UsesSplittableParDoWithWindowedSideInputs;
import org.apache.beam.sdk.testing.UsesTestStream;
import org.apache.beam.sdk.testing.ValidatesRunner;
import org.apache.beam.sdk.transforms.DoFn.BoundedPerElement;
+import org.apache.beam.sdk.transforms.splittabledofn.OffsetRange;
import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
@@ -73,16 +67,10 @@ public class SplittableDoFnTest implements Serializable {
static class PairStringWithIndexToLength extends DoFn<String, KV<String, Integer>> {
@ProcessElement
- public ProcessContinuation process(ProcessContext c, OffsetRangeTracker tracker) {
- for (long i = tracker.currentRestriction().getFrom(), numIterations = 0;
- tracker.tryClaim(i);
- ++i, ++numIterations) {
+ public void process(ProcessContext c, OffsetRangeTracker tracker) {
+ for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) {
c.output(KV.of(c.element(), (int) i));
- if (numIterations % 3 == 0) {
- return resume();
- }
}
- return stop();
}
@GetInitialRestriction
@@ -105,25 +93,8 @@ public class SplittableDoFnTest implements Serializable {
}
}
- private static PipelineOptions streamingTestPipelineOptions() {
- // Using testing options with streaming=true makes it possible to enable UsesSplittableParDo
- // tests in Dataflow runner, because as of writing, it can run Splittable DoFn only in
- // streaming mode.
- // This is a no-op for other runners currently (Direct runner doesn't care, and other
- // runners don't implement SDF at all yet).
- //
- // This is a workaround until https://issues.apache.org/jira/browse/BEAM-1620
- // is properly implemented and supports marking tests as streaming-only.
- //
- // https://issues.apache.org/jira/browse/BEAM-2483 specifically tracks the removal of the
- // current workaround.
- PipelineOptions options = testingPipelineOptions();
- options.as(StreamingOptions.class).setStreaming(true);
- return options;
- }
-
@Rule
- public final transient TestPipeline p = TestPipeline.fromOptions(streamingTestPipelineOptions());
+ public final transient TestPipeline p = TestPipeline.create();
@Test
@Category({ValidatesRunner.class, UsesSplittableParDo.class})
@@ -211,12 +182,6 @@ public class SplittableDoFnTest implements Serializable {
private static class SDFWithMultipleOutputsPerBlock extends DoFn<String, Integer> {
private static final int MAX_INDEX = 98765;
- private final int numClaimsPerCall;
-
- private SDFWithMultipleOutputsPerBlock(int numClaimsPerCall) {
- this.numClaimsPerCall = numClaimsPerCall;
- }
-
private static int snapToNextBlock(int index, int[] blockStarts) {
for (int i = 1; i < blockStarts.length; ++i) {
if (index > blockStarts[i - 1] && index <= blockStarts[i]) {
@@ -227,20 +192,14 @@ public class SplittableDoFnTest implements Serializable {
}
@ProcessElement
- public ProcessContinuation processElement(ProcessContext c, OffsetRangeTracker tracker) {
+ public void processElement(ProcessContext c, OffsetRangeTracker tracker) {
int[] blockStarts = {-1, 0, 12, 123, 1234, 12345, 34567, MAX_INDEX};
int trueStart = snapToNextBlock((int) tracker.currentRestriction().getFrom(), blockStarts);
- for (int i = trueStart, numIterations = 1;
- tracker.tryClaim(blockStarts[i]);
- ++i, ++numIterations) {
+ for (int i = trueStart; tracker.tryClaim(blockStarts[i]); ++i) {
for (int index = blockStarts[i]; index < blockStarts[i + 1]; ++index) {
c.output(index);
}
- if (numIterations == numClaimsPerCall) {
- return resume();
- }
}
- return stop();
}
@GetInitialRestriction
@@ -253,7 +212,7 @@ public class SplittableDoFnTest implements Serializable {
@Category({ValidatesRunner.class, UsesSplittableParDo.class})
public void testOutputAfterCheckpoint() throws Exception {
PCollection<Integer> outputs = p.apply(Create.of("foo"))
- .apply(ParDo.of(new SDFWithMultipleOutputsPerBlock(3)));
+ .apply(ParDo.of(new SDFWithMultipleOutputsPerBlock()));
PAssert.thatSingleton(outputs.apply(Count.<Integer>globally()))
.isEqualTo((long) SDFWithMultipleOutputsPerBlock.MAX_INDEX);
p.run();
@@ -328,105 +287,9 @@ public class SplittableDoFnTest implements Serializable {
PAssert.that(res).containsInAnyOrder("a:0", "a:1", "a:2", "a:3", "b:4", "b:5", "b:6", "b:7");
p.run();
- }
-
- @BoundedPerElement
- private static class SDFWithMultipleOutputsPerBlockAndSideInput
- extends DoFn<Integer, KV<String, Integer>> {
- private static final int MAX_INDEX = 98765;
- private final PCollectionView<String> sideInput;
- private final int numClaimsPerCall;
-
- public SDFWithMultipleOutputsPerBlockAndSideInput(
- PCollectionView<String> sideInput, int numClaimsPerCall) {
- this.sideInput = sideInput;
- this.numClaimsPerCall = numClaimsPerCall;
- }
-
- private static int snapToNextBlock(int index, int[] blockStarts) {
- for (int i = 1; i < blockStarts.length; ++i) {
- if (index > blockStarts[i - 1] && index <= blockStarts[i]) {
- return i;
- }
- }
- throw new IllegalStateException("Shouldn't get here");
- }
-
- @ProcessElement
- public ProcessContinuation processElement(ProcessContext c, OffsetRangeTracker tracker) {
- int[] blockStarts = {-1, 0, 12, 123, 1234, 12345, 34567, MAX_INDEX};
- int trueStart = snapToNextBlock((int) tracker.currentRestriction().getFrom(), blockStarts);
- for (int i = trueStart, numIterations = 1;
- tracker.tryClaim(blockStarts[i]);
- ++i, ++numIterations) {
- for (int index = blockStarts[i]; index < blockStarts[i + 1]; ++index) {
- c.output(KV.of(c.sideInput(sideInput) + ":" + c.element(), index));
- }
- if (numIterations == numClaimsPerCall) {
- return resume();
- }
- }
- return stop();
- }
-
- @GetInitialRestriction
- public OffsetRange getInitialRange(Integer element) {
- return new OffsetRange(0, MAX_INDEX);
- }
- }
-
- @Test
- @Category({
- ValidatesRunner.class,
- UsesSplittableParDo.class,
- UsesSplittableParDoWithWindowedSideInputs.class
- })
- public void testWindowedSideInputWithCheckpoints() throws Exception {
- PCollection<Integer> mainInput =
- p.apply("main",
- Create.timestamped(
- TimestampedValue.of(0, new Instant(0)),
- TimestampedValue.of(1, new Instant(1)),
- TimestampedValue.of(2, new Instant(2)),
- TimestampedValue.of(3, new Instant(3))))
- .apply("window 1", Window.<Integer>into(FixedWindows.of(Duration.millis(1))));
-
- PCollectionView<String> sideInput =
- p.apply("side",
- Create.timestamped(
- TimestampedValue.of("a", new Instant(0)),
- TimestampedValue.of("b", new Instant(2))))
- .apply("window 2", Window.<String>into(FixedWindows.of(Duration.millis(2))))
- .apply("singleton", View.<String>asSingleton());
-
- PCollection<KV<String, Integer>> res =
- mainInput.apply(
- ParDo.of(
- new SDFWithMultipleOutputsPerBlockAndSideInput(
- sideInput, 3 /* numClaimsPerCall */))
- .withSideInputs(sideInput));
- PCollection<KV<String, Iterable<Integer>>> grouped =
- res.apply(GroupByKey.<String, Integer>create());
-
- PAssert.that(grouped.apply(Keys.<String>create()))
- .containsInAnyOrder("a:0", "a:1", "b:2", "b:3");
- PAssert.that(grouped)
- .satisfies(
- new SerializableFunction<Iterable<KV<String, Iterable<Integer>>>, Void>() {
- @Override
- public Void apply(Iterable<KV<String, Iterable<Integer>>> input) {
- List<Integer> expected = new ArrayList<>();
- for (int i = 0; i < SDFWithMultipleOutputsPerBlockAndSideInput.MAX_INDEX; ++i) {
- expected.add(i);
- }
- for (KV<String, Iterable<Integer>> kv : input) {
- assertEquals(expected, Ordering.<Integer>natural().sortedCopy(kv.getValue()));
- }
- return null;
- }
- });
- p.run();
+ // TODO: also add test coverage when the SDF checkpoints - the resumed call should also
+ // properly access side inputs.
// TODO: also test coverage when some of the windows of the side input are not ready.
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java
index 2098c66..3edb194 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnInvokersTest.java
@@ -17,8 +17,6 @@
*/
package org.apache.beam.sdk.transforms.reflect;
-import static org.apache.beam.sdk.transforms.DoFn.ProcessContinuation.resume;
-import static org.apache.beam.sdk.transforms.DoFn.ProcessContinuation.stop;
import static org.hamcrest.CoreMatchers.instanceOf;
import static org.hamcrest.Matchers.equalTo;
import static org.junit.Assert.assertEquals;
@@ -91,8 +89,8 @@ public class DoFnInvokersTest {
when(mockArgumentProvider.processContext(Matchers.<DoFn>any())).thenReturn(mockProcessContext);
}
- private DoFn.ProcessContinuation invokeProcessElement(DoFn<String, String> fn) {
- return DoFnInvokers.invokerFor(fn).invokeProcessElement(mockArgumentProvider);
+ private void invokeProcessElement(DoFn<String, String> fn) {
+ DoFnInvokers.invokerFor(fn).invokeProcessElement(mockArgumentProvider);
}
private void invokeOnTimer(String timerId, DoFn<String, String> fn) {
@@ -121,7 +119,7 @@ public class DoFnInvokersTest {
public void processElement(ProcessContext c) throws Exception {}
}
MockFn mockFn = mock(MockFn.class);
- assertEquals(stop(), invokeProcessElement(mockFn));
+ invokeProcessElement(mockFn);
verify(mockFn).processElement(mockProcessContext);
}
@@ -142,7 +140,7 @@ public class DoFnInvokersTest {
public void testDoFnWithProcessElementInterface() throws Exception {
IdentityUsingInterfaceWithProcessElement fn =
mock(IdentityUsingInterfaceWithProcessElement.class);
- assertEquals(stop(), invokeProcessElement(fn));
+ invokeProcessElement(fn);
verify(fn).processElement(mockProcessContext);
}
@@ -163,14 +161,14 @@ public class DoFnInvokersTest {
@Test
public void testDoFnWithMethodInSuperclass() throws Exception {
IdentityChildWithoutOverride fn = mock(IdentityChildWithoutOverride.class);
- assertEquals(stop(), invokeProcessElement(fn));
+ invokeProcessElement(fn);
verify(fn).process(mockProcessContext);
}
@Test
public void testDoFnWithMethodInSubclass() throws Exception {
IdentityChildWithOverride fn = mock(IdentityChildWithOverride.class);
- assertEquals(stop(), invokeProcessElement(fn));
+ invokeProcessElement(fn);
verify(fn).process(mockProcessContext);
}
@@ -181,7 +179,7 @@ public class DoFnInvokersTest {
public void processElement(ProcessContext c, IntervalWindow w) throws Exception {}
}
MockFn fn = mock(MockFn.class);
- assertEquals(stop(), invokeProcessElement(fn));
+ invokeProcessElement(fn);
verify(fn).processElement(mockProcessContext, mockWindow);
}
@@ -205,7 +203,7 @@ public class DoFnInvokersTest {
throws Exception {}
}
MockFn fn = mock(MockFn.class);
- assertEquals(stop(), invokeProcessElement(fn));
+ invokeProcessElement(fn);
verify(fn).processElement(mockProcessContext, mockState);
}
@@ -231,35 +229,11 @@ public class DoFnInvokersTest {
public void onTimer() {}
}
MockFn fn = mock(MockFn.class);
- assertEquals(stop(), invokeProcessElement(fn));
+ invokeProcessElement(fn);
verify(fn).processElement(mockProcessContext, mockTimer);
}
@Test
- public void testDoFnWithReturn() throws Exception {
- class MockFn extends DoFn<String, String> {
- @DoFn.ProcessElement
- public ProcessContinuation processElement(ProcessContext c, SomeRestrictionTracker tracker)
- throws Exception {
- return null;
- }
-
- @GetInitialRestriction
- public SomeRestriction getInitialRestriction(String element) {
- return null;
- }
-
- @NewTracker
- public SomeRestrictionTracker newTracker(SomeRestriction restriction) {
- return null;
- }
- }
- MockFn fn = mock(MockFn.class);
- when(fn.processElement(mockProcessContext, null)).thenReturn(resume());
- assertEquals(resume(), invokeProcessElement(fn));
- }
-
- @Test
public void testDoFnWithStartBundleSetupTeardown() throws Exception {
class MockFn extends DoFn<String, String> {
@ProcessElement
@@ -314,9 +288,7 @@ public class DoFnInvokersTest {
/** Public so Mockito can do "delegatesTo()" in the test below. */
public static class MockFn extends DoFn<String, String> {
@ProcessElement
- public ProcessContinuation processElement(ProcessContext c, SomeRestrictionTracker tracker) {
- return null;
- }
+ public void processElement(ProcessContext c, SomeRestrictionTracker tracker) {}
@GetInitialRestriction
public SomeRestriction getInitialRestriction(String element) {
@@ -368,7 +340,7 @@ public class DoFnInvokersTest {
.splitRestriction(
eq("blah"), same(restriction), Mockito.<DoFn.OutputReceiver<SomeRestriction>>any());
when(fn.newTracker(restriction)).thenReturn(tracker);
- when(fn.processElement(mockProcessContext, tracker)).thenReturn(resume());
+ fn.processElement(mockProcessContext, tracker);
assertEquals(coder, invoker.invokeGetRestrictionCoder(CoderRegistry.createDefault()));
assertEquals(restriction, invoker.invokeGetInitialRestriction("blah"));
@@ -384,8 +356,6 @@ public class DoFnInvokersTest {
});
assertEquals(Arrays.asList(part1, part2, part3), outputs);
assertEquals(tracker, invoker.invokeNewTracker(restriction));
- assertEquals(
- resume(),
invoker.invokeProcessElement(
new FakeArgumentProvider<String, String>() {
@Override
@@ -397,7 +367,7 @@ public class DoFnInvokersTest {
public RestrictionTracker<?> restrictionTracker() {
return tracker;
}
- }));
+ });
}
private static class RestrictionWithDefaultTracker
@@ -471,7 +441,7 @@ public class DoFnInvokersTest {
assertEquals("foo", output);
}
});
- assertEquals(stop(), invoker.invokeProcessElement(mockArgumentProvider));
+ invoker.invokeProcessElement(mockArgumentProvider);
assertThat(
invoker.invokeNewTracker(new RestrictionWithDefaultTracker()),
instanceOf(DefaultTracker.class));
@@ -561,14 +531,14 @@ public class DoFnInvokersTest {
@Test
public void testLocalPrivateDoFnClass() throws Exception {
PrivateDoFnClass fn = mock(PrivateDoFnClass.class);
- assertEquals(stop(), invokeProcessElement(fn));
+ invokeProcessElement(fn);
verify(fn).processThis(mockProcessContext);
}
@Test
public void testStaticPackagePrivateDoFnClass() throws Exception {
DoFn<String, String> fn = mock(DoFnInvokersTestHelper.newStaticPackagePrivateDoFn().getClass());
- assertEquals(stop(), invokeProcessElement(fn));
+ invokeProcessElement(fn);
DoFnInvokersTestHelper.verifyStaticPackagePrivateDoFn(fn, mockProcessContext);
}
@@ -576,28 +546,28 @@ public class DoFnInvokersTest {
public void testInnerPackagePrivateDoFnClass() throws Exception {
DoFn<String, String> fn =
mock(new DoFnInvokersTestHelper().newInnerPackagePrivateDoFn().getClass());
- assertEquals(stop(), invokeProcessElement(fn));
+ invokeProcessElement(fn);
DoFnInvokersTestHelper.verifyInnerPackagePrivateDoFn(fn, mockProcessContext);
}
@Test
public void testStaticPrivateDoFnClass() throws Exception {
DoFn<String, String> fn = mock(DoFnInvokersTestHelper.newStaticPrivateDoFn().getClass());
- assertEquals(stop(), invokeProcessElement(fn));
+ invokeProcessElement(fn);
DoFnInvokersTestHelper.verifyStaticPrivateDoFn(fn, mockProcessContext);
}
@Test
public void testInnerPrivateDoFnClass() throws Exception {
DoFn<String, String> fn = mock(new DoFnInvokersTestHelper().newInnerPrivateDoFn().getClass());
- assertEquals(stop(), invokeProcessElement(fn));
+ invokeProcessElement(fn);
DoFnInvokersTestHelper.verifyInnerPrivateDoFn(fn, mockProcessContext);
}
@Test
public void testAnonymousInnerDoFn() throws Exception {
DoFn<String, String> fn = mock(new DoFnInvokersTestHelper().newInnerAnonymousDoFn().getClass());
- assertEquals(stop(), invokeProcessElement(fn));
+ invokeProcessElement(fn);
DoFnInvokersTestHelper.verifyInnerAnonymousDoFn(fn, mockProcessContext);
}
@@ -634,31 +604,6 @@ public class DoFnInvokersTest {
}
@Test
- public void testProcessElementExceptionWithReturn() throws Exception {
- thrown.expect(UserCodeException.class);
- thrown.expectMessage("bogus");
- DoFnInvokers.invokerFor(
- new DoFn<Integer, Integer>() {
- @ProcessElement
- public ProcessContinuation processElement(
- @SuppressWarnings("unused") ProcessContext c, SomeRestrictionTracker tracker) {
- throw new IllegalArgumentException("bogus");
- }
-
- @GetInitialRestriction
- public SomeRestriction getInitialRestriction(Integer element) {
- return null;
- }
-
- @NewTracker
- public SomeRestrictionTracker newTracker(SomeRestriction restriction) {
- return null;
- }
- })
- .invokeProcessElement(new FakeArgumentProvider<Integer, Integer>());
- }
-
- @Test
public void testStartBundleException() throws Exception {
DoFnInvoker<Integer, Integer> invoker =
DoFnInvokers.invokerFor(
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesProcessElementTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesProcessElementTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesProcessElementTest.java
index 44ae5c4..d321f54 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesProcessElementTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesProcessElementTest.java
@@ -50,7 +50,7 @@ public class DoFnSignaturesProcessElementTest {
@Test
public void testBadReturnType() throws Exception {
thrown.expect(IllegalArgumentException.class);
- thrown.expectMessage("Must return void or ProcessContinuation");
+ thrown.expectMessage("Must return void");
analyzeProcessElementMethod(
new AnonymousMethod() {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesSplittableDoFnTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesSplittableDoFnTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesSplittableDoFnTest.java
index 08af65e..07b3348 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesSplittableDoFnTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesSplittableDoFnTest.java
@@ -52,8 +52,7 @@ import org.junit.runners.JUnit4;
public class DoFnSignaturesSplittableDoFnTest {
@Rule public ExpectedException thrown = ExpectedException.none();
- private abstract static class SomeRestriction
- implements HasDefaultTracker<SomeRestriction, SomeRestrictionTracker> {}
+ private static class SomeRestriction {}
private abstract static class SomeRestrictionTracker
implements RestrictionTracker<SomeRestriction> {}
@@ -61,20 +60,6 @@ public class DoFnSignaturesSplittableDoFnTest {
private abstract static class SomeRestrictionCoder extends StructuredCoder<SomeRestriction> {}
@Test
- public void testReturnsProcessContinuation() throws Exception {
- DoFnSignature.ProcessElementMethod signature =
- analyzeProcessElementMethod(
- new AnonymousMethod() {
- private DoFn.ProcessContinuation method(
- DoFn<Integer, String>.ProcessContext context) {
- return null;
- }
- });
-
- assertTrue(signature.hasReturnValue());
- }
-
- @Test
public void testHasRestrictionTracker() throws Exception {
DoFnSignature.ProcessElementMethod signature =
analyzeProcessElementMethod(
@@ -115,6 +100,11 @@ public class DoFnSignaturesSplittableDoFnTest {
public SomeRestriction getInitialRestriction(Integer element) {
return null;
}
+
+ @NewTracker
+ public SomeRestrictionTracker newTracker(SomeRestriction restriction) {
+ return null;
+ }
}
@BoundedPerElement
@@ -140,55 +130,6 @@ public class DoFnSignaturesSplittableDoFnTest {
.isBoundedPerElement());
}
- private static class BaseFnWithoutContinuation extends DoFn<Integer, String> {
- @ProcessElement
- public void processElement(ProcessContext context, SomeRestrictionTracker tracker) {}
-
- @GetInitialRestriction
- public SomeRestriction getInitialRestriction(Integer element) {
- return null;
- }
- }
-
- private static class BaseFnWithContinuation extends DoFn<Integer, String> {
- @ProcessElement
- public ProcessContinuation processElement(
- ProcessContext context, SomeRestrictionTracker tracker) {
- return null;
- }
-
- @GetInitialRestriction
- public SomeRestriction getInitialRestriction(Integer element) {
- return null;
- }
- }
-
- @Test
- public void testSplittableBoundednessInferredFromReturnValue() throws Exception {
- assertEquals(
- PCollection.IsBounded.BOUNDED,
- DoFnSignatures.getSignature(BaseFnWithoutContinuation.class).isBoundedPerElement());
- assertEquals(
- PCollection.IsBounded.UNBOUNDED,
- DoFnSignatures.getSignature(BaseFnWithContinuation.class).isBoundedPerElement());
- }
-
- @Test
- public void testSplittableRespectsBoundednessAnnotation() throws Exception {
- @BoundedPerElement
- class BoundedFnWithContinuation extends BaseFnWithContinuation {}
-
- assertEquals(
- PCollection.IsBounded.BOUNDED,
- DoFnSignatures.getSignature(BoundedFnWithContinuation.class).isBoundedPerElement());
-
- @UnboundedPerElement
- class UnboundedFnWithContinuation extends BaseFnWithContinuation {}
-
- assertEquals(
- PCollection.IsBounded.UNBOUNDED,
- DoFnSignatures.getSignature(UnboundedFnWithContinuation.class).isBoundedPerElement());
- }
@Test
public void testUnsplittableIsBounded() throws Exception {
class UnsplittableFn extends DoFn<Integer, String> {
@@ -231,10 +172,8 @@ public class DoFnSignaturesSplittableDoFnTest {
public void testSplittableWithAllFunctions() throws Exception {
class GoodSplittableDoFn extends DoFn<Integer, String> {
@ProcessElement
- public ProcessContinuation processElement(
- ProcessContext context, SomeRestrictionTracker tracker) {
- return null;
- }
+ public void processElement(
+ ProcessContext context, SomeRestrictionTracker tracker) {}
@GetInitialRestriction
public SomeRestriction getInitialRestriction(Integer element) {
@@ -259,7 +198,6 @@ public class DoFnSignaturesSplittableDoFnTest {
DoFnSignature signature = DoFnSignatures.getSignature(GoodSplittableDoFn.class);
assertEquals(SomeRestrictionTracker.class, signature.processElement().trackerT().getRawType());
assertTrue(signature.processElement().isSplittable());
- assertTrue(signature.processElement().hasReturnValue());
assertEquals(
SomeRestriction.class, signature.getInitialRestriction().restrictionT().getRawType());
assertEquals(SomeRestriction.class, signature.splitRestriction().restrictionT().getRawType());
@@ -276,9 +214,7 @@ public class DoFnSignaturesSplittableDoFnTest {
public void testSplittableWithAllFunctionsGeneric() throws Exception {
class GoodGenericSplittableDoFn<RestrictionT, TrackerT, CoderT> extends DoFn<Integer, String> {
@ProcessElement
- public ProcessContinuation processElement(ProcessContext context, TrackerT tracker) {
- return null;
- }
+ public void processElement(ProcessContext context, TrackerT tracker) {}
@GetInitialRestriction
public RestrictionT getInitialRestriction(Integer element) {
@@ -306,7 +242,6 @@ public class DoFnSignaturesSplittableDoFnTest {
SomeRestriction, SomeRestrictionTracker, SomeRestrictionCoder>() {}.getClass());
assertEquals(SomeRestrictionTracker.class, signature.processElement().trackerT().getRawType());
assertTrue(signature.processElement().isSplittable());
- assertTrue(signature.processElement().hasReturnValue());
assertEquals(
SomeRestriction.class, signature.getInitialRestriction().restrictionT().getRawType());
assertEquals(SomeRestriction.class, signature.splitRestriction().restrictionT().getRawType());
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java
index 70c8dfd..cffb0ad 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/reflect/DoFnSignaturesTest.java
@@ -29,7 +29,6 @@ import static org.junit.Assert.fail;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.coders.VarLongCoder;
-import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.state.StateSpecs;
import org.apache.beam.sdk.state.TimeDomain;
@@ -330,19 +329,6 @@ public class DoFnSignaturesTest {
}
@Test
- public void testPipelineOptionsParameter() throws Exception {
- DoFnSignature sig =
- DoFnSignatures.getSignature(new DoFn<String, String>() {
- @ProcessElement
- public void process(ProcessContext c, PipelineOptions options) {}
- }.getClass());
-
- assertThat(
- sig.processElement().extraParameters(),
- Matchers.<Parameter>hasItem(instanceOf(Parameter.PipelineOptionsParameter.class)));
- }
-
- @Test
public void testDeclAndUsageOfTimerInSuperclass() throws Exception {
DoFnSignature sig =
DoFnSignatures.getSignature(new DoFnOverridingAbstractTimerUse().getClass());
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/splittabledofn/OffsetRangeTrackerTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/splittabledofn/OffsetRangeTrackerTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/splittabledofn/OffsetRangeTrackerTest.java
index 8aed6b9..831894c 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/splittabledofn/OffsetRangeTrackerTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/splittabledofn/OffsetRangeTrackerTest.java
@@ -21,7 +21,6 @@ import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
-import org.apache.beam.sdk.io.range.OffsetRange;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/windowing/SlidingWindowsTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/windowing/SlidingWindowsTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/windowing/SlidingWindowsTest.java
index bfd01f0..b14e221 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/windowing/SlidingWindowsTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/windowing/SlidingWindowsTest.java
@@ -21,7 +21,6 @@ import static org.apache.beam.sdk.testing.WindowFnTestUtils.runWindowFn;
import static org.apache.beam.sdk.testing.WindowFnTestUtils.set;
import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem;
import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.is;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertThat;
@@ -56,12 +55,11 @@ public class SlidingWindowsTest {
expected.put(new IntervalWindow(new Instant(0), new Instant(10)), set(1, 2, 5, 9));
expected.put(new IntervalWindow(new Instant(5), new Instant(15)), set(5, 9, 10, 11));
expected.put(new IntervalWindow(new Instant(10), new Instant(20)), set(10, 11));
- SlidingWindows windowFn = SlidingWindows.of(new Duration(10)).every(new Duration(5));
assertEquals(
expected,
- runWindowFn(windowFn,
+ runWindowFn(
+ SlidingWindows.of(new Duration(10)).every(new Duration(5)),
Arrays.asList(1L, 2L, 5L, 9L, 10L, 11L)));
- assertThat(windowFn.assignsToOneWindow(), is(false));
}
@Test
@@ -71,27 +69,11 @@ public class SlidingWindowsTest {
expected.put(new IntervalWindow(new Instant(0), new Instant(7)), set(1, 2, 5));
expected.put(new IntervalWindow(new Instant(5), new Instant(12)), set(5, 9, 10, 11));
expected.put(new IntervalWindow(new Instant(10), new Instant(17)), set(10, 11));
- SlidingWindows windowFn = SlidingWindows.of(new Duration(7)).every(new Duration(5));
- assertEquals(
- expected,
- runWindowFn(windowFn,
- Arrays.asList(1L, 2L, 5L, 9L, 10L, 11L)));
- assertThat(windowFn.assignsToOneWindow(), is(false));
- }
-
- @Test
- public void testEqualSize() throws Exception {
- Map<IntervalWindow, Set<String>> expected = new HashMap<>();
- expected.put(new IntervalWindow(new Instant(0), new Instant(3)), set(1, 2));
- expected.put(new IntervalWindow(new Instant(3), new Instant(6)), set(3, 4, 5));
- expected.put(new IntervalWindow(new Instant(6), new Instant(9)), set(6, 7));
- SlidingWindows windowFn = SlidingWindows.of(new Duration(3)).every(new Duration(3));
assertEquals(
expected,
runWindowFn(
- windowFn,
- Arrays.asList(1L, 2L, 3L, 4L, 5L, 6L, 7L)));
- assertThat(windowFn.assignsToOneWindow(), is(true));
+ SlidingWindows.of(new Duration(7)).every(new Duration(5)),
+ Arrays.asList(1L, 2L, 5L, 9L, 10L, 11L)));
}
@Test
@@ -100,14 +82,12 @@ public class SlidingWindowsTest {
expected.put(new IntervalWindow(new Instant(0), new Instant(3)), set(1, 2));
expected.put(new IntervalWindow(new Instant(10), new Instant(13)), set(10, 11));
expected.put(new IntervalWindow(new Instant(100), new Instant(103)), set(100));
- SlidingWindows windowFn = SlidingWindows.of(new Duration(3)).every(new Duration(10));
assertEquals(
expected,
runWindowFn(
// Only look at the first 3 millisecs of every 10-millisec interval.
- windowFn,
+ SlidingWindows.of(new Duration(3)).every(new Duration(10)),
Arrays.asList(1L, 2L, 3L, 5L, 9L, 10L, 11L, 100L)));
- assertThat(windowFn.assignsToOneWindow(), is(true));
}
@Test
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/extensions/google-cloud-platform-core/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/extensions/google-cloud-platform-core/pom.xml b/sdks/java/extensions/google-cloud-platform-core/pom.xml
index 7d54990..e4e951b 100644
--- a/sdks/java/extensions/google-cloud-platform-core/pom.xml
+++ b/sdks/java/extensions/google-cloud-platform-core/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-extensions-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/util/GcsUtil.java
----------------------------------------------------------------------
diff --git a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/util/GcsUtil.java b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/util/GcsUtil.java
index d7205bf..8d1fe74 100644
--- a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/util/GcsUtil.java
+++ b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/util/GcsUtil.java
@@ -135,7 +135,7 @@ public class GcsUtil {
private static final int MAX_CONCURRENT_BATCHES = 256;
private static final FluentBackoff BACKOFF_FACTORY =
- FluentBackoff.DEFAULT.withMaxRetries(10).withInitialBackoff(Duration.standardSeconds(1));
+ FluentBackoff.DEFAULT.withMaxRetries(3).withInitialBackoff(Duration.millis(200));
/////////////////////////////////////////////////////////////////////////////
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/util/RetryHttpRequestInitializer.java
----------------------------------------------------------------------
diff --git a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/util/RetryHttpRequestInitializer.java b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/util/RetryHttpRequestInitializer.java
index fd908cf..e5b48d3 100644
--- a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/util/RetryHttpRequestInitializer.java
+++ b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/util/RetryHttpRequestInitializer.java
@@ -17,9 +17,8 @@
*/
package org.apache.beam.sdk.util;
-import static com.google.api.client.util.BackOffUtils.next;
-
-import com.google.api.client.http.HttpIOExceptionHandler;
+import com.google.api.client.http.HttpBackOffIOExceptionHandler;
+import com.google.api.client.http.HttpBackOffUnsuccessfulResponseHandler;
import com.google.api.client.http.HttpRequest;
import com.google.api.client.http.HttpRequestInitializer;
import com.google.api.client.http.HttpResponse;
@@ -61,106 +60,64 @@ public class RetryHttpRequestInitializer implements HttpRequestInitializer {
*/
private static final int HANGING_GET_TIMEOUT_SEC = 80;
- /** Handlers used to provide additional logging information on unsuccessful HTTP requests. */
- private static class LoggingHttpBackOffHandler
- implements HttpIOExceptionHandler, HttpUnsuccessfulResponseHandler {
-
- private final Sleeper sleeper;
- private final BackOff ioExceptionBackOff;
- private final BackOff unsuccessfulResponseBackOff;
- private final Set<Integer> ignoredResponseCodes;
- private int ioExceptionRetries;
- private int unsuccessfulResponseRetries;
-
- private LoggingHttpBackOffHandler(
- Sleeper sleeper,
- BackOff ioExceptionBackOff,
- BackOff unsucessfulResponseBackOff,
- Set<Integer> ignoredResponseCodes) {
- this.sleeper = sleeper;
- this.ioExceptionBackOff = ioExceptionBackOff;
- this.unsuccessfulResponseBackOff = unsucessfulResponseBackOff;
- this.ignoredResponseCodes = ignoredResponseCodes;
+ private static class LoggingHttpBackOffIOExceptionHandler
+ extends HttpBackOffIOExceptionHandler {
+ public LoggingHttpBackOffIOExceptionHandler(BackOff backOff) {
+ super(backOff);
}
@Override
public boolean handleIOException(HttpRequest request, boolean supportsRetry)
throws IOException {
- // We will retry if the request supports retry or the backoff was successful.
- // Note that the order of these checks is important since
- // backOffWasSuccessful will perform a sleep.
- boolean willRetry = supportsRetry && backOffWasSuccessful(ioExceptionBackOff);
+ boolean willRetry = super.handleIOException(request, supportsRetry);
if (willRetry) {
- ioExceptionRetries += 1;
LOG.debug("Request failed with IOException, will retry: {}", request.getUrl());
} else {
- String message = "Request failed with IOException, "
- + "performed {} retries due to IOExceptions, "
- + "performed {} retries due to unsuccessful status codes, "
- + "HTTP framework says request {} be retried, "
- + "(caller responsible for retrying): {}";
- LOG.warn(message,
- ioExceptionRetries,
- unsuccessfulResponseRetries,
- supportsRetry ? "can" : "cannot",
+ LOG.warn(
+ "Request failed with IOException (caller responsible for retrying): {}",
request.getUrl());
}
return willRetry;
}
+ }
+
+ private static class LoggingHttpBackoffUnsuccessfulResponseHandler
+ implements HttpUnsuccessfulResponseHandler {
+ private final HttpBackOffUnsuccessfulResponseHandler handler;
+ private final Set<Integer> ignoredResponseCodes;
+
+ public LoggingHttpBackoffUnsuccessfulResponseHandler(BackOff backoff,
+ Sleeper sleeper, Set<Integer> ignoredResponseCodes) {
+ this.ignoredResponseCodes = ignoredResponseCodes;
+ handler = new HttpBackOffUnsuccessfulResponseHandler(backoff);
+ handler.setSleeper(sleeper);
+ handler.setBackOffRequired(
+ new HttpBackOffUnsuccessfulResponseHandler.BackOffRequired() {
+ @Override
+ public boolean isRequired(HttpResponse response) {
+ int statusCode = response.getStatusCode();
+ return (statusCode / 100 == 5) || // 5xx: server error
+ statusCode == 429; // 429: Too many requests
+ }
+ });
+ }
@Override
- public boolean handleResponse(HttpRequest request, HttpResponse response, boolean supportsRetry)
- throws IOException {
- // We will retry if the request supports retry and the status code requires a backoff
- // and the backoff was successful. Note that the order of these checks is important since
- // backOffWasSuccessful will perform a sleep.
- boolean willRetry = supportsRetry
- && retryOnStatusCode(response.getStatusCode())
- && backOffWasSuccessful(unsuccessfulResponseBackOff);
- if (willRetry) {
- unsuccessfulResponseRetries += 1;
+ public boolean handleResponse(HttpRequest request, HttpResponse response,
+ boolean supportsRetry) throws IOException {
+ boolean retry = handler.handleResponse(request, response, supportsRetry);
+ if (retry) {
LOG.debug("Request failed with code {}, will retry: {}",
response.getStatusCode(), request.getUrl());
- } else {
- String message = "Request failed with code {}, "
- + "performed {} retries due to IOExceptions, "
- + "performed {} retries due to unsuccessful status codes, "
- + "HTTP framework says request {} be retried, "
- + "(caller responsible for retrying): {}";
- if (ignoredResponseCodes.contains(response.getStatusCode())) {
- // Log ignored response codes at a lower level
- LOG.debug(message,
- response.getStatusCode(),
- ioExceptionRetries,
- unsuccessfulResponseRetries,
- supportsRetry ? "can" : "cannot",
- request.getUrl());
- } else {
- LOG.warn(message,
- response.getStatusCode(),
- ioExceptionRetries,
- unsuccessfulResponseRetries,
- supportsRetry ? "can" : "cannot",
- request.getUrl());
- }
- }
- return willRetry;
- }
- /** Returns true iff performing the backoff was successful. */
- private boolean backOffWasSuccessful(BackOff backOff) {
- try {
- return next(sleeper, backOff);
- } catch (InterruptedException | IOException e) {
- return false;
+ } else if (!ignoredResponseCodes.contains(response.getStatusCode())) {
+ LOG.warn(
+ "Request failed with code {} (caller responsible for retrying): {}",
+ response.getStatusCode(),
+ request.getUrl());
}
- }
- /** Returns true iff the {@code statusCode} represents an error that should be retried. */
- private boolean retryOnStatusCode(int statusCode) {
- return (statusCode == 0) // Code 0 usually means no response / network error
- || (statusCode / 100 == 5) // 5xx: server error
- || statusCode == 429; // 429: Too many requests
+ return retry;
}
}
@@ -216,20 +173,20 @@ public class RetryHttpRequestInitializer implements HttpRequestInitializer {
// TODO: Do this exclusively for work requests.
request.setReadTimeout(HANGING_GET_TIMEOUT_SEC * 1000);
- LoggingHttpBackOffHandler loggingHttpBackOffHandler = new LoggingHttpBackOffHandler(
- sleeper,
- // Retry immediately on IOExceptions.
- BackOff.ZERO_BACKOFF,
- // Back off on retryable http errors.
+ // Back off on retryable http errors.
+ request.setUnsuccessfulResponseHandler(
// A back-off multiplier of 2 raises the maximum request retrying time
// to approximately 5 minutes (keeping other back-off parameters to
// their default values).
- new ExponentialBackOff.Builder().setNanoClock(nanoClock).setMultiplier(2).build(),
- ignoredResponseCodes
- );
-
- request.setUnsuccessfulResponseHandler(loggingHttpBackOffHandler);
- request.setIOExceptionHandler(loggingHttpBackOffHandler);
+ new LoggingHttpBackoffUnsuccessfulResponseHandler(
+ new ExponentialBackOff.Builder().setNanoClock(nanoClock)
+ .setMultiplier(2).build(),
+ sleeper, ignoredResponseCodes));
+
+ // Retry immediately on IOExceptions.
+ LoggingHttpBackOffIOExceptionHandler loggingBackoffHandler =
+ new LoggingHttpBackOffIOExceptionHandler(BackOff.ZERO_BACKOFF);
+ request.setIOExceptionHandler(loggingBackoffHandler);
// Set response initializer
if (responseInterceptor != null) {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/GcpCoreApiSurfaceTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/GcpCoreApiSurfaceTest.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/GcpCoreApiSurfaceTest.java
index a0d9e4b..625c248 100644
--- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/GcpCoreApiSurfaceTest.java
+++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/GcpCoreApiSurfaceTest.java
@@ -15,16 +15,14 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.beam.sdk.extensions.gcp;
+package org.apache.beam;
-import static org.apache.beam.sdk.util.ApiSurface.classesInPackage;
-import static org.apache.beam.sdk.util.ApiSurface.containsOnlyClassesMatching;
+import static org.apache.beam.sdk.util.ApiSurface.containsOnlyPackages;
import static org.hamcrest.MatcherAssert.assertThat;
import com.google.common.collect.ImmutableSet;
import java.util.Set;
import org.apache.beam.sdk.util.ApiSurface;
-import org.hamcrest.Matcher;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
@@ -34,32 +32,28 @@ import org.junit.runners.JUnit4;
public class GcpCoreApiSurfaceTest {
@Test
- public void testGcpCoreApiSurface() throws Exception {
- final Package thisPackage = getClass().getPackage();
- final ClassLoader thisClassLoader = getClass().getClassLoader();
- final ApiSurface apiSurface =
- ApiSurface.ofPackage(thisPackage, thisClassLoader)
- .pruningPattern("org[.]apache[.]beam[.].*Test.*")
- .pruningPattern("org[.]apache[.]beam[.].*IT")
- .pruningPattern("java[.]lang.*")
- .pruningPattern("java[.]util.*");
+ public void testApiSurface() throws Exception {
@SuppressWarnings("unchecked")
- final Set<Matcher<Class<?>>> allowedClasses =
+ final Set<String> allowed =
ImmutableSet.of(
- classesInPackage("com.google.api.client.googleapis"),
- classesInPackage("com.google.api.client.http"),
- classesInPackage("com.google.api.client.json"),
- classesInPackage("com.google.api.client.util"),
- classesInPackage("com.google.api.services.storage"),
- classesInPackage("com.google.auth"),
- classesInPackage("com.fasterxml.jackson.annotation"),
- classesInPackage("java"),
- classesInPackage("javax"),
- classesInPackage("org.apache.beam.sdk"),
- classesInPackage("org.joda.time")
- );
+ "org.apache.beam",
+ "com.google.api.client",
+ "com.google.api.services.storage",
+ "com.google.auth",
+ "com.fasterxml.jackson.annotation",
+ "com.fasterxml.jackson.core",
+ "com.fasterxml.jackson.databind",
+ "org.apache.avro",
+ "org.hamcrest",
+ // via DataflowMatchers
+ "org.codehaus.jackson",
+ // via Avro
+ "org.joda.time",
+ "org.junit",
+ "sun.reflect");
- assertThat(apiSurface, containsOnlyClassesMatching(allowedClasses));
+ assertThat(
+ ApiSurface.getSdkApiSurface(getClass().getClassLoader()), containsOnlyPackages(allowed));
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/util/RetryHttpRequestInitializerTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/util/RetryHttpRequestInitializerTest.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/util/RetryHttpRequestInitializerTest.java
index 13a9309..37551a4 100644
--- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/util/RetryHttpRequestInitializerTest.java
+++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/util/RetryHttpRequestInitializerTest.java
@@ -49,11 +49,10 @@ import java.net.SocketTimeoutException;
import java.security.PrivateKey;
import java.util.Arrays;
import java.util.concurrent.atomic.AtomicLong;
-import org.apache.beam.sdk.testing.ExpectedLogs;
import org.hamcrest.Matchers;
import org.junit.After;
+import org.junit.Assert;
import org.junit.Before;
-import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
@@ -68,8 +67,6 @@ import org.mockito.stubbing.Answer;
@RunWith(JUnit4.class)
public class RetryHttpRequestInitializerTest {
- @Rule public ExpectedLogs expectedLogs = ExpectedLogs.none(RetryHttpRequestInitializer.class);
-
@Mock private PrivateKey mockPrivateKey;
@Mock private LowLevelHttpRequest mockLowLevelRequest;
@Mock private LowLevelHttpResponse mockLowLevelResponse;
@@ -138,7 +135,6 @@ public class RetryHttpRequestInitializerTest {
verify(mockLowLevelRequest).setTimeout(anyInt(), anyInt());
verify(mockLowLevelRequest).execute();
verify(mockLowLevelResponse).getStatusCode();
- expectedLogs.verifyNotLogged("Request failed");
}
/**
@@ -157,7 +153,7 @@ public class RetryHttpRequestInitializerTest {
HttpResponse response = result.executeUnparsed();
assertNotNull(response);
} catch (HttpResponseException e) {
- assertThat(e.getMessage(), Matchers.containsString("403"));
+ Assert.assertThat(e.getMessage(), Matchers.containsString("403"));
}
verify(mockHttpResponseInterceptor).interceptResponse(any(HttpResponse.class));
@@ -166,7 +162,6 @@ public class RetryHttpRequestInitializerTest {
verify(mockLowLevelRequest).setTimeout(anyInt(), anyInt());
verify(mockLowLevelRequest).execute();
verify(mockLowLevelResponse).getStatusCode();
- expectedLogs.verifyWarn("Request failed with code 403");
}
/**
@@ -193,7 +188,6 @@ public class RetryHttpRequestInitializerTest {
verify(mockLowLevelRequest, times(3)).setTimeout(anyInt(), anyInt());
verify(mockLowLevelRequest, times(3)).execute();
verify(mockLowLevelResponse, times(3)).getStatusCode();
- expectedLogs.verifyDebug("Request failed with code 503");
}
/**
@@ -217,7 +211,6 @@ public class RetryHttpRequestInitializerTest {
verify(mockLowLevelRequest, times(2)).setTimeout(anyInt(), anyInt());
verify(mockLowLevelRequest, times(2)).execute();
verify(mockLowLevelResponse).getStatusCode();
- expectedLogs.verifyDebug("Request failed with IOException");
}
/**
@@ -231,22 +224,19 @@ public class RetryHttpRequestInitializerTest {
int n = 0;
@Override
public Integer answer(InvocationOnMock invocation) {
- return n++ < retries ? 503 : 9999;
+ return (n++ < retries - 1) ? 503 : 200;
}});
Storage.Buckets.Get result = storage.buckets().get("test");
- try {
- result.executeUnparsed();
- fail();
- } catch (Throwable t) {
- }
+ HttpResponse response = result.executeUnparsed();
+ assertNotNull(response);
verify(mockHttpResponseInterceptor).interceptResponse(any(HttpResponse.class));
- verify(mockLowLevelRequest, atLeastOnce()).addHeader(anyString(), anyString());
- verify(mockLowLevelRequest, times(retries + 1)).setTimeout(anyInt(), anyInt());
- verify(mockLowLevelRequest, times(retries + 1)).execute();
- verify(mockLowLevelResponse, times(retries + 1)).getStatusCode();
- expectedLogs.verifyWarn("performed 10 retries due to unsuccessful status codes");
+ verify(mockLowLevelRequest, atLeastOnce()).addHeader(anyString(),
+ anyString());
+ verify(mockLowLevelRequest, times(retries)).setTimeout(anyInt(), anyInt());
+ verify(mockLowLevelRequest, times(retries)).execute();
+ verify(mockLowLevelResponse, times(retries)).getStatusCode();
}
/**
@@ -286,7 +276,6 @@ public class RetryHttpRequestInitializerTest {
} catch (Throwable e) {
assertThat(e, Matchers.<Throwable>instanceOf(SocketTimeoutException.class));
assertEquals(1 + defaultNumberOfRetries, executeCount.get());
- expectedLogs.verifyWarn("performed 10 retries due to IOExceptions");
}
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/extensions/jackson/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/extensions/jackson/pom.xml b/sdks/java/extensions/jackson/pom.xml
index 7fd38e0..4b09c11 100644
--- a/sdks/java/extensions/jackson/pom.xml
+++ b/sdks/java/extensions/jackson/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-extensions-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/extensions/join-library/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/extensions/join-library/pom.xml b/sdks/java/extensions/join-library/pom.xml
index ea24b75..556ec40 100644
--- a/sdks/java/extensions/join-library/pom.xml
+++ b/sdks/java/extensions/join-library/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-extensions-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/extensions/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/extensions/pom.xml b/sdks/java/extensions/pom.xml
index 1222476..3d63626 100644
--- a/sdks/java/extensions/pom.xml
+++ b/sdks/java/extensions/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/extensions/protobuf/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/extensions/protobuf/pom.xml b/sdks/java/extensions/protobuf/pom.xml
index 63855f8..ae909ab 100644
--- a/sdks/java/extensions/protobuf/pom.xml
+++ b/sdks/java/extensions/protobuf/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-extensions-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/extensions/sorter/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/extensions/sorter/pom.xml b/sdks/java/extensions/sorter/pom.xml
index 395c73f..9d25f9d 100644
--- a/sdks/java/extensions/sorter/pom.xml
+++ b/sdks/java/extensions/sorter/pom.xml
@@ -22,13 +22,17 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-extensions-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
<artifactId>beam-sdks-java-extensions-sorter</artifactId>
<name>Apache Beam :: SDKs :: Java :: Extensions :: Sorter</name>
+ <properties>
+ <hadoop.version>2.7.1</hadoop.version>
+ </properties>
+
<dependencies>
<dependency>
<groupId>org.apache.beam</groupId>
@@ -38,12 +42,14 @@
<dependency>
<groupId>org.apache.hadoop</groupId>
<artifactId>hadoop-mapreduce-client-core</artifactId>
+ <version>${hadoop.version}</version>
<scope>provided</scope>
</dependency>
<dependency>
<groupId>org.apache.hadoop</groupId>
<artifactId>hadoop-common</artifactId>
+ <version>${hadoop.version}</version>
<scope>provided</scope>
</dependency>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/harness/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/harness/pom.xml b/sdks/java/harness/pom.xml
index fe5c2f1..61a170a 100644
--- a/sdks/java/harness/pom.xml
+++ b/sdks/java/harness/pom.xml
@@ -23,7 +23,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-sdks-java-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
@@ -83,11 +83,6 @@
<dependency>
<groupId>org.apache.beam</groupId>
- <artifactId>beam-runners-core-construction-java</artifactId>
- </dependency>
-
- <dependency>
- <groupId>org.apache.beam</groupId>
<artifactId>beam-runners-google-cloud-dataflow-java</artifactId>
</dependency>
@@ -155,21 +150,10 @@
</dependency>
<dependency>
- <groupId>joda-time</groupId>
- <artifactId>joda-time</artifactId>
- </dependency>
-
- <dependency>
<groupId>org.slf4j</groupId>
<artifactId>slf4j-api</artifactId>
</dependency>
- <dependency>
- <groupId>com.google.auto.service</groupId>
- <artifactId>auto-service</artifactId>
- <optional>true</optional>
- </dependency>
-
<!-- test dependencies -->
<dependency>
<groupId>org.hamcrest</groupId>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
----------------------------------------------------------------------
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
index 2a9cef8..e33277a 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java
@@ -18,111 +18,82 @@
package org.apache.beam.fn.harness.control;
-import com.google.common.annotations.VisibleForTesting;
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.collect.Iterables.getOnlyElement;
+
+import com.google.common.collect.Collections2;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableMultimap;
import com.google.common.collect.Lists;
import com.google.common.collect.Multimap;
-import com.google.common.collect.Sets;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.BytesValue;
+import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.Message;
import java.io.IOException;
import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashSet;
import java.util.List;
import java.util.Map;
-import java.util.ServiceLoader;
-import java.util.Set;
+import java.util.Objects;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;
import org.apache.beam.fn.harness.data.BeamFnDataClient;
+import org.apache.beam.fn.harness.fake.FakeStepContext;
import org.apache.beam.fn.harness.fn.ThrowingConsumer;
import org.apache.beam.fn.harness.fn.ThrowingRunnable;
import org.apache.beam.fn.v1.BeamFnApi;
-import org.apache.beam.runners.core.PTransformRunnerFactory;
-import org.apache.beam.runners.core.PTransformRunnerFactory.Registrar;
+import org.apache.beam.runners.core.BeamFnDataReadRunner;
+import org.apache.beam.runners.core.BeamFnDataWriteRunner;
+import org.apache.beam.runners.core.BoundedSourceRunner;
+import org.apache.beam.runners.core.DoFnRunner;
+import org.apache.beam.runners.core.DoFnRunners;
+import org.apache.beam.runners.core.DoFnRunners.OutputManager;
+import org.apache.beam.runners.core.NullSideInputReader;
+import org.apache.beam.runners.dataflow.util.DoFnInfo;
import org.apache.beam.sdk.common.runner.v1.RunnerApi;
+import org.apache.beam.sdk.io.BoundedSource;
import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.sdk.util.common.ReflectHelpers;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.WindowingStrategy;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* Processes {@link org.apache.beam.fn.v1.BeamFnApi.ProcessBundleRequest}s by materializing
- * the set of required runners for each {@link RunnerApi.FunctionSpec},
+ * the set of required runners for each {@link org.apache.beam.fn.v1.BeamFnApi.FunctionSpec},
* wiring them together based upon the {@code input} and {@code output} map definitions.
*
* <p>Finally executes the DAG based graph by starting all runners in reverse topological order,
* and finishing all runners in forward topological order.
*/
public class ProcessBundleHandler {
-
// TODO: What should the initial set of URNs be?
private static final String DATA_INPUT_URN = "urn:org.apache.beam:source:runner:0.1";
- public static final String JAVA_SOURCE_URN = "urn:org.apache.beam:source:java:0.1";
+ private static final String DATA_OUTPUT_URN = "urn:org.apache.beam:sink:runner:0.1";
+ private static final String JAVA_DO_FN_URN = "urn:org.apache.beam:dofn:java:0.1";
+ private static final String JAVA_SOURCE_URN = "urn:org.apache.beam:source:java:0.1";
private static final Logger LOG = LoggerFactory.getLogger(ProcessBundleHandler.class);
- private static final Map<String, PTransformRunnerFactory> REGISTERED_RUNNER_FACTORIES;
-
- static {
- Set<Registrar> pipelineRunnerRegistrars =
- Sets.newTreeSet(ReflectHelpers.ObjectsClassComparator.INSTANCE);
- pipelineRunnerRegistrars.addAll(
- Lists.newArrayList(ServiceLoader.load(Registrar.class,
- ReflectHelpers.findClassLoader())));
-
- // Load all registered PTransform runner factories.
- ImmutableMap.Builder<String, PTransformRunnerFactory> builder =
- ImmutableMap.builder();
- for (Registrar registrar : pipelineRunnerRegistrars) {
- builder.putAll(registrar.getPTransformRunnerFactories());
- }
- REGISTERED_RUNNER_FACTORIES = builder.build();
- }
private final PipelineOptions options;
private final Function<String, Message> fnApiRegistry;
private final BeamFnDataClient beamFnDataClient;
- private final Map<String, PTransformRunnerFactory> urnToPTransformRunnerFactoryMap;
- private final PTransformRunnerFactory defaultPTransformRunnerFactory;
-
public ProcessBundleHandler(
PipelineOptions options,
Function<String, Message> fnApiRegistry,
BeamFnDataClient beamFnDataClient) {
- this(options, fnApiRegistry, beamFnDataClient, REGISTERED_RUNNER_FACTORIES);
- }
-
- @VisibleForTesting
- ProcessBundleHandler(
- PipelineOptions options,
- Function<String, Message> fnApiRegistry,
- BeamFnDataClient beamFnDataClient,
- Map<String, PTransformRunnerFactory> urnToPTransformRunnerFactoryMap) {
this.options = options;
this.fnApiRegistry = fnApiRegistry;
this.beamFnDataClient = beamFnDataClient;
- this.urnToPTransformRunnerFactoryMap = urnToPTransformRunnerFactoryMap;
- this.defaultPTransformRunnerFactory = new PTransformRunnerFactory<Object>() {
- @Override
- public Object createRunnerForPTransform(
- PipelineOptions pipelineOptions,
- BeamFnDataClient beamFnDataClient,
- String pTransformId,
- RunnerApi.PTransform pTransform,
- Supplier<String> processBundleInstructionId,
- Map<String, RunnerApi.PCollection> pCollections,
- Map<String, RunnerApi.Coder> coders,
- Multimap<String, ThrowingConsumer<WindowedValue<?>>> pCollectionIdsToConsumers,
- Consumer<ThrowingRunnable> addStartFunction,
- Consumer<ThrowingRunnable> addFinishFunction) {
- throw new IllegalStateException(String.format(
- "No factory registered for %s, known factories %s",
- pTransform.getSpec().getUrn(),
- urnToPTransformRunnerFactoryMap.keySet()));
- }
- };
}
private void createRunnerAndConsumersForPTransformRecursively(
@@ -157,19 +128,115 @@ public class ProcessBundleHandler {
}
}
- urnToPTransformRunnerFactoryMap.getOrDefault(
- pTransform.getSpec().getUrn(), defaultPTransformRunnerFactory)
- .createRunnerForPTransform(
- options,
- beamFnDataClient,
- pTransformId,
- pTransform,
- processBundleInstructionId,
- processBundleDescriptor.getPcollectionsMap(),
- processBundleDescriptor.getCodersMap(),
- pCollectionIdsToConsumers,
- addStartFunction,
- addFinishFunction);
+ createRunnerForPTransform(
+ pTransformId,
+ pTransform,
+ processBundleInstructionId,
+ processBundleDescriptor.getPcollectionsMap(),
+ pCollectionIdsToConsumers,
+ addStartFunction,
+ addFinishFunction);
+ }
+
+ protected void createRunnerForPTransform(
+ String pTransformId,
+ RunnerApi.PTransform pTransform,
+ Supplier<String> processBundleInstructionId,
+ Map<String, RunnerApi.PCollection> pCollections,
+ Multimap<String, ThrowingConsumer<WindowedValue<?>>> pCollectionIdsToConsumers,
+ Consumer<ThrowingRunnable> addStartFunction,
+ Consumer<ThrowingRunnable> addFinishFunction) throws IOException {
+
+
+ // For every output PCollection, create a map from output name to Consumer
+ ImmutableMap.Builder<String, Collection<ThrowingConsumer<WindowedValue<?>>>>
+ outputMapBuilder = ImmutableMap.builder();
+ for (Map.Entry<String, String> entry : pTransform.getOutputsMap().entrySet()) {
+ outputMapBuilder.put(
+ entry.getKey(),
+ pCollectionIdsToConsumers.get(entry.getValue()));
+ }
+ ImmutableMap<String, Collection<ThrowingConsumer<WindowedValue<?>>>> outputMap =
+ outputMapBuilder.build();
+
+
+ // Based upon the function spec, populate the start/finish/consumer information.
+ RunnerApi.FunctionSpec functionSpec = pTransform.getSpec();
+ ThrowingConsumer<WindowedValue<?>> consumer;
+ switch (functionSpec.getUrn()) {
+ default:
+ BeamFnApi.Target target;
+ RunnerApi.Coder coderSpec;
+ throw new IllegalArgumentException(
+ String.format("Unknown FunctionSpec %s", functionSpec));
+
+ case DATA_OUTPUT_URN:
+ target = BeamFnApi.Target.newBuilder()
+ .setPrimitiveTransformReference(pTransformId)
+ .setName(getOnlyElement(pTransform.getInputsMap().keySet()))
+ .build();
+ coderSpec = (RunnerApi.Coder) fnApiRegistry.apply(
+ pCollections.get(getOnlyElement(pTransform.getInputsMap().values())).getCoderId());
+ BeamFnDataWriteRunner<Object> remoteGrpcWriteRunner =
+ new BeamFnDataWriteRunner<Object>(
+ functionSpec,
+ processBundleInstructionId,
+ target,
+ coderSpec,
+ beamFnDataClient);
+ addStartFunction.accept(remoteGrpcWriteRunner::registerForOutput);
+ consumer = (ThrowingConsumer)
+ (ThrowingConsumer<WindowedValue<Object>>) remoteGrpcWriteRunner::consume;
+ addFinishFunction.accept(remoteGrpcWriteRunner::close);
+ break;
+
+ case DATA_INPUT_URN:
+ target = BeamFnApi.Target.newBuilder()
+ .setPrimitiveTransformReference(pTransformId)
+ .setName(getOnlyElement(pTransform.getOutputsMap().keySet()))
+ .build();
+ coderSpec = (RunnerApi.Coder) fnApiRegistry.apply(
+ pCollections.get(getOnlyElement(pTransform.getOutputsMap().values())).getCoderId());
+ BeamFnDataReadRunner<?> remoteGrpcReadRunner =
+ new BeamFnDataReadRunner<Object>(
+ functionSpec,
+ processBundleInstructionId,
+ target,
+ coderSpec,
+ beamFnDataClient,
+ (Map) outputMap);
+ addStartFunction.accept(remoteGrpcReadRunner::registerInputLocation);
+ consumer = null;
+ addFinishFunction.accept(remoteGrpcReadRunner::blockTillReadFinishes);
+ break;
+
+ case JAVA_DO_FN_URN:
+ DoFnRunner<Object, Object> doFnRunner = createDoFnRunner(functionSpec, (Map) outputMap);
+ addStartFunction.accept(doFnRunner::startBundle);
+ consumer = (ThrowingConsumer)
+ (ThrowingConsumer<WindowedValue<Object>>) doFnRunner::processElement;
+ addFinishFunction.accept(doFnRunner::finishBundle);
+ break;
+
+ case JAVA_SOURCE_URN:
+ @SuppressWarnings({"unchecked", "rawtypes"})
+ BoundedSourceRunner<BoundedSource<Object>, Object> sourceRunner =
+ createBoundedSourceRunner(functionSpec, (Map) outputMap);
+ // TODO: Remove and replace with source being sent across gRPC port
+ addStartFunction.accept(sourceRunner::start);
+ consumer = (ThrowingConsumer)
+ (ThrowingConsumer<WindowedValue<BoundedSource<Object>>>)
+ sourceRunner::runReadLoop;
+ break;
+ }
+
+ // If we created a consumer, add it to the map containing PCollection ids to consumers
+ if (consumer != null) {
+ for (String inputPCollectionId :
+ pTransform.getInputsMap().values()) {
+ pCollectionIdsToConsumers.put(inputPCollectionId, consumer);
+ }
+ }
}
public BeamFnApi.InstructionResponse.Builder processBundle(BeamFnApi.InstructionRequest request)
@@ -232,4 +299,88 @@ public class ProcessBundleHandler {
return response;
}
+
+ /**
+ * Converts a {@link org.apache.beam.fn.v1.BeamFnApi.FunctionSpec} into a {@link DoFnRunner}.
+ */
+ private <InputT, OutputT> DoFnRunner<InputT, OutputT> createDoFnRunner(
+ RunnerApi.FunctionSpec functionSpec,
+ Map<String, Collection<ThrowingConsumer<WindowedValue<OutputT>>>> outputMap) {
+ ByteString serializedFn;
+ try {
+ serializedFn = functionSpec.getParameter().unpack(BytesValue.class).getValue();
+ } catch (InvalidProtocolBufferException e) {
+ throw new IllegalArgumentException(
+ String.format("Unable to unwrap DoFn %s", functionSpec), e);
+ }
+ DoFnInfo<?, ?> doFnInfo =
+ (DoFnInfo<?, ?>)
+ SerializableUtils.deserializeFromByteArray(serializedFn.toByteArray(), "DoFnInfo");
+
+ checkArgument(
+ Objects.equals(
+ new HashSet<>(Collections2.transform(outputMap.keySet(), Long::parseLong)),
+ doFnInfo.getOutputMap().keySet()),
+ "Unexpected mismatch between transform output map %s and DoFnInfo output map %s.",
+ outputMap.keySet(),
+ doFnInfo.getOutputMap());
+
+ ImmutableMultimap.Builder<TupleTag<?>,
+ ThrowingConsumer<WindowedValue<OutputT>>> tagToOutput =
+ ImmutableMultimap.builder();
+ for (Map.Entry<Long, TupleTag<?>> entry : doFnInfo.getOutputMap().entrySet()) {
+ tagToOutput.putAll(entry.getValue(), outputMap.get(Long.toString(entry.getKey())));
+ }
+ @SuppressWarnings({"unchecked", "rawtypes"})
+ final Map<TupleTag<?>, Collection<ThrowingConsumer<WindowedValue<?>>>> tagBasedOutputMap =
+ (Map) tagToOutput.build().asMap();
+
+ OutputManager outputManager =
+ new OutputManager() {
+ Map<TupleTag<?>, Collection<ThrowingConsumer<WindowedValue<?>>>> tupleTagToOutput =
+ tagBasedOutputMap;
+
+ @Override
+ public <T> void output(TupleTag<T> tag, WindowedValue<T> output) {
+ try {
+ Collection<ThrowingConsumer<WindowedValue<?>>> consumers =
+ tupleTagToOutput.get(tag);
+ if (consumers == null) {
+ /* This is a normal case, e.g., if a DoFn has output but that output is not
+ * consumed. Drop the output. */
+ return;
+ }
+ for (ThrowingConsumer<WindowedValue<?>> consumer : consumers) {
+ consumer.accept(output);
+ }
+ } catch (Throwable t) {
+ throw new RuntimeException(t);
+ }
+ }
+ };
+
+ @SuppressWarnings({"unchecked", "rawtypes", "deprecation"})
+ DoFnRunner<InputT, OutputT> runner =
+ DoFnRunners.simpleRunner(
+ PipelineOptionsFactory.create(), /* TODO */
+ (DoFn) doFnInfo.getDoFn(),
+ NullSideInputReader.empty(), /* TODO */
+ outputManager,
+ (TupleTag) doFnInfo.getOutputMap().get(doFnInfo.getMainOutput()),
+ new ArrayList<>(doFnInfo.getOutputMap().values()),
+ new FakeStepContext(),
+ (WindowingStrategy) doFnInfo.getWindowingStrategy());
+ return runner;
+ }
+
+ private <InputT extends BoundedSource<OutputT>, OutputT>
+ BoundedSourceRunner<InputT, OutputT> createBoundedSourceRunner(
+ RunnerApi.FunctionSpec functionSpec,
+ Map<String, Collection<ThrowingConsumer<WindowedValue<OutputT>>>> outputMap) {
+
+ @SuppressWarnings({"rawtypes", "unchecked"})
+ BoundedSourceRunner<InputT, OutputT> runner =
+ new BoundedSourceRunner(options, functionSpec, outputMap);
+ return runner;
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/RegisterHandler.java
----------------------------------------------------------------------
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/RegisterHandler.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/RegisterHandler.java
index 0e738ac..276a120 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/RegisterHandler.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/RegisterHandler.java
@@ -79,7 +79,7 @@ public class RegisterHandler {
processBundleDescriptor.getClass());
computeIfAbsent(processBundleDescriptor.getId()).complete(processBundleDescriptor);
for (Map.Entry<String, RunnerApi.Coder> entry
- : processBundleDescriptor.getCodersMap().entrySet()) {
+ : processBundleDescriptor.getCodersyyyMap().entrySet()) {
LOG.debug("Registering {} with type {}",
entry.getKey(),
entry.getValue().getClass());
[10/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/ShardedKeyCoder.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/ShardedKeyCoder.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/ShardedKeyCoder.java
new file mode 100644
index 0000000..c2b62b7
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/ShardedKeyCoder.java
@@ -0,0 +1,74 @@
+/*
+ * 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.sdk.io.gcp.bigquery;
+
+import com.google.common.annotations.VisibleForTesting;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.Arrays;
+import java.util.List;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.StructuredCoder;
+import org.apache.beam.sdk.coders.VarIntCoder;
+
+
+/**
+ * A {@link Coder} for {@link ShardedKey}, using a wrapped key {@link Coder}.
+ */
+@VisibleForTesting
+class ShardedKeyCoder<KeyT>
+ extends StructuredCoder<ShardedKey<KeyT>> {
+ public static <KeyT> ShardedKeyCoder<KeyT> of(Coder<KeyT> keyCoder) {
+ return new ShardedKeyCoder<>(keyCoder);
+ }
+
+ private final Coder<KeyT> keyCoder;
+ private final VarIntCoder shardNumberCoder;
+
+ protected ShardedKeyCoder(Coder<KeyT> keyCoder) {
+ this.keyCoder = keyCoder;
+ this.shardNumberCoder = VarIntCoder.of();
+ }
+
+ @Override
+ public List<? extends Coder<?>> getCoderArguments() {
+ return Arrays.asList(keyCoder);
+ }
+
+ @Override
+ public void encode(ShardedKey<KeyT> key, OutputStream outStream)
+ throws IOException {
+ keyCoder.encode(key.getKey(), outStream);
+ shardNumberCoder.encode(key.getShardNumber(), outStream);
+ }
+
+ @Override
+ public ShardedKey<KeyT> decode(InputStream inStream)
+ throws IOException {
+ return new ShardedKey<>(
+ keyCoder.decode(inStream),
+ shardNumberCoder.decode(inStream));
+ }
+
+ @Override
+ public void verifyDeterministic() throws NonDeterministicException {
+ keyCoder.verifyDeterministic();
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteFn.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteFn.java
index a210858..63e5bc1 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteFn.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteFn.java
@@ -33,7 +33,6 @@ import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.SystemDoFnInternal;
import org.apache.beam.sdk.values.KV;
-import org.apache.beam.sdk.values.ShardedKey;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.ValueInSingleWindow;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java
index fa5b3ce..18b2033 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java
@@ -19,7 +19,6 @@ package org.apache.beam.sdk.io.gcp.bigquery;
import com.google.api.services.bigquery.model.TableRow;
import org.apache.beam.sdk.coders.KvCoder;
-import org.apache.beam.sdk.coders.ShardedKeyCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
@@ -30,7 +29,6 @@ import org.apache.beam.sdk.transforms.windowing.Window;
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.ShardedKey;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TagWithUniqueIds.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TagWithUniqueIds.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TagWithUniqueIds.java
index 51b9375..cd88222 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TagWithUniqueIds.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TagWithUniqueIds.java
@@ -26,7 +26,6 @@ import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.values.KV;
-import org.apache.beam.sdk.values.ShardedKey;
/**
* Fn that tags each table row with a unique id and destination table. To avoid calling
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteBundlesToFiles.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteBundlesToFiles.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteBundlesToFiles.java
index e1ed746..d68779a 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteBundlesToFiles.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteBundlesToFiles.java
@@ -19,7 +19,6 @@
package org.apache.beam.sdk.io.gcp.bigquery;
import static com.google.common.base.Preconditions.checkNotNull;
-
import com.google.api.services.bigquery.model.TableRow;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
@@ -41,7 +40,6 @@ import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.ShardedKey;
import org.apache.beam.sdk.values.TupleTag;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteGroupedRecordsToFiles.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteGroupedRecordsToFiles.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteGroupedRecordsToFiles.java
index 887cb93..45dc2a8 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteGroupedRecordsToFiles.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteGroupedRecordsToFiles.java
@@ -22,7 +22,6 @@ import com.google.api.services.bigquery.model.TableRow;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.ShardedKey;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WritePartition.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WritePartition.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WritePartition.java
index 451d1bd..acd1132 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WritePartition.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WritePartition.java
@@ -26,7 +26,6 @@ import org.apache.beam.sdk.io.gcp.bigquery.WriteBundlesToFiles.Result;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.ShardedKey;
import org.apache.beam.sdk.values.TupleTag;
/**
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java
index 9ed2916..c5494d8 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java
@@ -42,7 +42,6 @@ import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.JobService;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.ShardedKey;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java
index 0a90dde..1692cda 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java
@@ -175,7 +175,7 @@ import org.slf4j.LoggerFactory;
* pipeline. Please refer to the documentation of corresponding
* {@link PipelineRunner PipelineRunners} for more details.
*/
-@Experimental(Experimental.Kind.SOURCE_SINK)
+@Experimental
public class BigtableIO {
private static final Logger LOG = LoggerFactory.getLogger(BigtableIO.class);
@@ -211,7 +211,7 @@ public class BigtableIO {
*
* @see BigtableIO
*/
- @Experimental(Experimental.Kind.SOURCE_SINK)
+ @Experimental
@AutoValue
public abstract static class Read extends PTransform<PBegin, PCollection<Row>> {
@@ -415,7 +415,7 @@ public class BigtableIO {
*
* @see BigtableIO
*/
- @Experimental(Experimental.Kind.SOURCE_SINK)
+ @Experimental
@AutoValue
public abstract static class Write
extends PTransform<PCollection<KV<ByteString, Iterable<Mutation>>>, PDone> {
@@ -1027,7 +1027,7 @@ public class BigtableIO {
"{}: Failed to interpolate key for fraction {}.", rangeTracker.getRange(), fraction, e);
return null;
}
- LOG.info(
+ LOG.debug(
"Proposing to split {} at fraction {} (key {})", rangeTracker, fraction, splitKey);
BigtableSource primary;
BigtableSource residual;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImpl.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImpl.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImpl.java
index 07476e2..d1a17fe 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImpl.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImpl.java
@@ -168,13 +168,14 @@ class BigtableServiceImpl implements BigtableService {
private BigtableSession session;
private AsyncExecutor executor;
private BulkMutation bulkMutation;
- private final String tableName;
+ private final MutateRowRequest.Builder partialBuilder;
public BigtableWriterImpl(BigtableSession session, BigtableTableName tableName) {
this.session = session;
executor = session.createAsyncExecutor();
bulkMutation = session.createBulkMutation(tableName, executor);
- this.tableName = tableName.toString();
+
+ partialBuilder = MutateRowRequest.newBuilder().setTableName(tableName.toString());
}
@Override
@@ -207,8 +208,8 @@ class BigtableServiceImpl implements BigtableService {
KV<ByteString, Iterable<Mutation>> record)
throws IOException {
MutateRowRequest r =
- MutateRowRequest.newBuilder()
- .setTableName(tableName)
+ partialBuilder
+ .clone()
.setRowKey(record.getKey())
.addAllMutations(record.getValue())
.build();
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottler.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottler.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottler.java
deleted file mode 100644
index ce6ebe6..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottler.java
+++ /dev/null
@@ -1,103 +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.sdk.io.gcp.datastore;
-
-import com.google.common.annotations.VisibleForTesting;
-import java.util.Random;
-import org.apache.beam.sdk.transforms.Sum;
-import org.apache.beam.sdk.util.MovingFunction;
-
-
-/**
- * An implementation of client-side adaptive throttling. See
- * https://landing.google.com/sre/book/chapters/handling-overload.html#client-side-throttling-a7sYUg
- * for a full discussion of the use case and algorithm applied.
- */
-class AdaptiveThrottler {
- private final MovingFunction successfulRequests;
- private final MovingFunction allRequests;
-
- /** The target ratio between requests sent and successful requests. This is "K" in the formula in
- * https://landing.google.com/sre/book/chapters/handling-overload.html */
- private final double overloadRatio;
-
- /** The target minimum number of requests per samplePeriodMs, even if no requests succeed. Must be
- * greater than 0, else we could throttle to zero. Because every decision is probabilistic, there
- * is no guarantee that the request rate in any given interval will not be zero. (This is the +1
- * from the formula in https://landing.google.com/sre/book/chapters/handling-overload.html */
- private static final double MIN_REQUESTS = 1;
- private final Random random;
-
- /**
- * @param samplePeriodMs the time window to keep of request history to inform throttling
- * decisions.
- * @param sampleUpdateMs the length of buckets within this time window.
- * @param overloadRatio the target ratio between requests sent and successful requests. You should
- * always set this to more than 1, otherwise the client would never try to send more requests than
- * succeeded in the past - so it could never recover from temporary setbacks.
- */
- public AdaptiveThrottler(long samplePeriodMs, long sampleUpdateMs,
- double overloadRatio) {
- this(samplePeriodMs, sampleUpdateMs, overloadRatio, new Random());
- }
-
- @VisibleForTesting
- AdaptiveThrottler(long samplePeriodMs, long sampleUpdateMs,
- double overloadRatio, Random random) {
- allRequests =
- new MovingFunction(samplePeriodMs, sampleUpdateMs,
- 1 /* numSignificantBuckets */, 1 /* numSignificantSamples */, Sum.ofLongs());
- successfulRequests =
- new MovingFunction(samplePeriodMs, sampleUpdateMs,
- 1 /* numSignificantBuckets */, 1 /* numSignificantSamples */, Sum.ofLongs());
- this.overloadRatio = overloadRatio;
- this.random = random;
- }
-
- @VisibleForTesting
- double throttlingProbability(long nowMsSinceEpoch) {
- if (!allRequests.isSignificant()) {
- return 0;
- }
- long allRequestsNow = allRequests.get(nowMsSinceEpoch);
- long successfulRequestsNow = successfulRequests.get(nowMsSinceEpoch);
- return Math.max(0,
- (allRequestsNow - overloadRatio * successfulRequestsNow) / (allRequestsNow + MIN_REQUESTS));
- }
-
- /**
- * Call this before sending a request to the remote service; if this returns true, drop the
- * request (treating it as a failure or trying it again at a later time).
- */
- public boolean throttleRequest(long nowMsSinceEpoch) {
- double delayProbability = throttlingProbability(nowMsSinceEpoch);
- // Note that we increment the count of all requests here, even if we return true - so even if we
- // tell the client not to send a request at all, it still counts as a failed request.
- allRequests.add(nowMsSinceEpoch, 1);
-
- return (random.nextDouble() < delayProbability);
- }
-
- /**
- * Call this after {@link throttleRequest} if your request was successful.
- */
- public void successfulRequest(long nowMsSinceEpoch) {
- successfulRequests.add(nowMsSinceEpoch, 1);
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
index 5f65428..b198a6f 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
@@ -71,8 +71,6 @@ import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.annotations.Experimental.Kind;
import org.apache.beam.sdk.coders.SerializableCoder;
import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
-import org.apache.beam.sdk.metrics.Counter;
-import org.apache.beam.sdk.metrics.Metrics;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
@@ -203,31 +201,11 @@ public class DatastoreV1 {
DatastoreV1() {}
/**
- * The number of entity updates written per RPC, initially. We buffer updates in the connector and
- * write a batch to Datastore once we have collected a certain number. This is the initial batch
- * size; it is adjusted at runtime based on the performance of previous writes (see {@link
- * DatastoreV1.WriteBatcher}).
- *
- * <p>Testing has found that a batch of 200 entities will generally finish within the timeout even
- * in adverse conditions.
- */
- @VisibleForTesting
- static final int DATASTORE_BATCH_UPDATE_ENTITIES_START = 200;
-
- /**
- * When choosing the number of updates in a single RPC, never exceed the maximum allowed by the
- * API.
+ * Cloud Datastore has a limit of 500 mutations per batch operation, so we flush
+ * changes to Datastore every 500 entities.
*/
@VisibleForTesting
- static final int DATASTORE_BATCH_UPDATE_ENTITIES_LIMIT = 500;
-
- /**
- * When choosing the number of updates in a single RPC, do not go below this value. The actual
- * number of entities per request may be lower when we flush for the end of a bundle or if we hit
- * {@link DatastoreV1.DATASTORE_BATCH_UPDATE_BYTES_LIMIT}.
- */
- @VisibleForTesting
- static final int DATASTORE_BATCH_UPDATE_ENTITIES_MIN = 10;
+ static final int DATASTORE_BATCH_UPDATE_LIMIT = 500;
/**
* Cloud Datastore has a limit of 10MB per RPC, so we also flush if the total size of mutations
@@ -235,7 +213,7 @@ public class DatastoreV1 {
* the mutations themselves and not the CommitRequest wrapper around them.
*/
@VisibleForTesting
- static final int DATASTORE_BATCH_UPDATE_BYTES_LIMIT = 9_000_000;
+ static final int DATASTORE_BATCH_UPDATE_BYTES_LIMIT = 5_000_000;
/**
* Returns an empty {@link DatastoreV1.Read} builder. Configure the source {@code projectId},
@@ -1129,74 +1107,18 @@ public class DatastoreV1 {
}
}
- /** Determines batch sizes for commit RPCs. */
- @VisibleForTesting
- interface WriteBatcher {
- /** Call before using this WriteBatcher. */
- void start();
-
- /**
- * Reports the latency of a previous commit RPC, and the number of mutations that it contained.
- */
- void addRequestLatency(long timeSinceEpochMillis, long latencyMillis, int numMutations);
-
- /** Returns the number of entities to include in the next CommitRequest. */
- int nextBatchSize(long timeSinceEpochMillis);
- }
-
- /**
- * Determines batch sizes for commit RPCs based on past performance.
- *
- * <p>It aims for a target response time per RPC: it uses the response times for previous RPCs
- * and the number of entities contained in them, calculates a rolling average time-per-entity, and
- * chooses the number of entities for future writes to hit the target time.
- *
- * <p>This enables us to send large batches without sending over-large requests in the case of
- * expensive entity writes that may timeout before the server can apply them all.
- */
- @VisibleForTesting
- static class WriteBatcherImpl implements WriteBatcher, Serializable {
- /** Target time per RPC for writes. */
- static final int DATASTORE_BATCH_TARGET_LATENCY_MS = 5000;
-
- @Override
- public void start() {
- meanLatencyPerEntityMs = new MovingAverage(
- 120000 /* sample period 2 minutes */, 10000 /* sample interval 10s */,
- 1 /* numSignificantBuckets */, 1 /* numSignificantSamples */);
- }
-
- @Override
- public void addRequestLatency(long timeSinceEpochMillis, long latencyMillis, int numMutations) {
- meanLatencyPerEntityMs.add(timeSinceEpochMillis, latencyMillis / numMutations);
- }
-
- @Override
- public int nextBatchSize(long timeSinceEpochMillis) {
- if (!meanLatencyPerEntityMs.hasValue(timeSinceEpochMillis)) {
- return DATASTORE_BATCH_UPDATE_ENTITIES_START;
- }
- long recentMeanLatency = Math.max(meanLatencyPerEntityMs.get(timeSinceEpochMillis), 1);
- return (int) Math.max(DATASTORE_BATCH_UPDATE_ENTITIES_MIN,
- Math.min(DATASTORE_BATCH_UPDATE_ENTITIES_LIMIT,
- DATASTORE_BATCH_TARGET_LATENCY_MS / recentMeanLatency));
- }
-
- private transient MovingAverage meanLatencyPerEntityMs;
- }
-
/**
* {@link DoFn} that writes {@link Mutation}s to Cloud Datastore. Mutations are written in
- * batches; see {@link DatastoreV1.WriteBatcherImpl}.
+ * batches, where the maximum batch size is {@link DatastoreV1#DATASTORE_BATCH_UPDATE_LIMIT}.
*
* <p>See <a
* href="https://cloud.google.com/datastore/docs/concepts/entities">
* Datastore: Entities, Properties, and Keys</a> for information about entity keys and mutations.
*
* <p>Commits are non-transactional. If a commit fails because of a conflict over an entity
- * group, the commit will be retried (up to {@link DatastoreV1.DatastoreWriterFn#MAX_RETRIES}
+ * group, the commit will be retried (up to {@link DatastoreV1#DATASTORE_BATCH_UPDATE_LIMIT}
* times). This means that the mutation operation should be idempotent. Thus, the writer should
- * only be used for {@code upsert} and {@code delete} mutation operations, as these are the only
+ * only be used for {code upsert} and {@code delete} mutation operations, as these are the only
* two Cloud Datastore mutations that are idempotent.
*/
@VisibleForTesting
@@ -1210,14 +1132,6 @@ public class DatastoreV1 {
// Current batch of mutations to be written.
private final List<Mutation> mutations = new ArrayList<>();
private int mutationsSize = 0; // Accumulated size of protos in mutations.
- private WriteBatcher writeBatcher;
- private transient AdaptiveThrottler throttler;
- private final Counter throttledSeconds =
- Metrics.counter(DatastoreWriterFn.class, "cumulativeThrottlingSeconds");
- private final Counter rpcErrors =
- Metrics.counter(DatastoreWriterFn.class, "datastoreRpcErrors");
- private final Counter rpcSuccesses =
- Metrics.counter(DatastoreWriterFn.class, "datastoreRpcSuccesses");
private static final int MAX_RETRIES = 5;
private static final FluentBackoff BUNDLE_WRITE_BACKOFF =
@@ -1225,31 +1139,24 @@ public class DatastoreV1 {
.withMaxRetries(MAX_RETRIES).withInitialBackoff(Duration.standardSeconds(5));
DatastoreWriterFn(String projectId, @Nullable String localhost) {
- this(StaticValueProvider.of(projectId), localhost, new V1DatastoreFactory(),
- new WriteBatcherImpl());
+ this(StaticValueProvider.of(projectId), localhost, new V1DatastoreFactory());
}
DatastoreWriterFn(ValueProvider<String> projectId, @Nullable String localhost) {
- this(projectId, localhost, new V1DatastoreFactory(), new WriteBatcherImpl());
+ this(projectId, localhost, new V1DatastoreFactory());
}
@VisibleForTesting
DatastoreWriterFn(ValueProvider<String> projectId, @Nullable String localhost,
- V1DatastoreFactory datastoreFactory, WriteBatcher writeBatcher) {
+ V1DatastoreFactory datastoreFactory) {
this.projectId = checkNotNull(projectId, "projectId");
this.localhost = localhost;
this.datastoreFactory = datastoreFactory;
- this.writeBatcher = writeBatcher;
}
@StartBundle
public void startBundle(StartBundleContext c) {
datastore = datastoreFactory.getDatastore(c.getPipelineOptions(), projectId.get(), localhost);
- writeBatcher.start();
- if (throttler == null) {
- // Initialize throttler at first use, because it is not serializable.
- throttler = new AdaptiveThrottler(120000, 10000, 1.25);
- }
}
@ProcessElement
@@ -1262,7 +1169,7 @@ public class DatastoreV1 {
}
mutations.add(c.element());
mutationsSize += size;
- if (mutations.size() >= writeBatcher.nextBatchSize(System.currentTimeMillis())) {
+ if (mutations.size() >= DatastoreV1.DATASTORE_BATCH_UPDATE_LIMIT) {
flushBatch();
}
}
@@ -1292,42 +1199,18 @@ public class DatastoreV1 {
while (true) {
// Batch upsert entities.
- CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
- commitRequest.addAllMutations(mutations);
- commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
- long startTime = System.currentTimeMillis(), endTime;
-
- if (throttler.throttleRequest(startTime)) {
- LOG.info("Delaying request due to previous failures");
- throttledSeconds.inc(WriteBatcherImpl.DATASTORE_BATCH_TARGET_LATENCY_MS / 1000);
- sleeper.sleep(WriteBatcherImpl.DATASTORE_BATCH_TARGET_LATENCY_MS);
- continue;
- }
-
try {
+ CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
+ commitRequest.addAllMutations(mutations);
+ commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
datastore.commit(commitRequest.build());
- endTime = System.currentTimeMillis();
-
- writeBatcher.addRequestLatency(endTime, endTime - startTime, mutations.size());
- throttler.successfulRequest(startTime);
- rpcSuccesses.inc();
-
// Break if the commit threw no exception.
break;
} catch (DatastoreException exception) {
- if (exception.getCode() == Code.DEADLINE_EXCEEDED) {
- /* Most errors are not related to request size, and should not change our expectation of
- * the latency of successful requests. DEADLINE_EXCEEDED can be taken into
- * consideration, though. */
- endTime = System.currentTimeMillis();
- writeBatcher.addRequestLatency(endTime, endTime - startTime, mutations.size());
- }
// Only log the code and message for potentially-transient errors. The entire exception
// will be propagated upon the last retry.
- LOG.error("Error writing batch of {} mutations to Datastore ({}): {}", mutations.size(),
- exception.getCode(), exception.getMessage());
- rpcErrors.inc();
-
+ LOG.error("Error writing to the Datastore ({}): {}", exception.getCode(),
+ exception.getMessage());
if (!BackOffUtils.next(sleeper, backoff)) {
LOG.error("Aborting after {} retries.", MAX_RETRIES);
throw exception;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/MovingAverage.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/MovingAverage.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/MovingAverage.java
deleted file mode 100644
index 0890e79..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/MovingAverage.java
+++ /dev/null
@@ -1,50 +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.sdk.io.gcp.datastore;
-
-import org.apache.beam.sdk.transforms.Sum;
-import org.apache.beam.sdk.util.MovingFunction;
-
-
-class MovingAverage {
- private final MovingFunction sum;
- private final MovingFunction count;
-
- public MovingAverage(long samplePeriodMs, long sampleUpdateMs,
- int numSignificantBuckets, int numSignificantSamples) {
- sum = new MovingFunction(samplePeriodMs, sampleUpdateMs,
- numSignificantBuckets, numSignificantSamples, Sum.ofLongs());
- count = new MovingFunction(samplePeriodMs, sampleUpdateMs,
- numSignificantBuckets, numSignificantSamples, Sum.ofLongs());
- }
-
- public void add(long nowMsSinceEpoch, long value) {
- sum.add(nowMsSinceEpoch, value);
- count.add(nowMsSinceEpoch, 1);
- }
-
- public long get(long nowMsSinceEpoch) {
- return sum.get(nowMsSinceEpoch) / count.get(nowMsSinceEpoch);
- }
-
- public boolean hasValue(long nowMsSinceEpoch) {
- return sum.isSignificant() && count.isSignificant()
- && count.get(nowMsSinceEpoch) > 0;
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/AbstractSpannerFn.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/AbstractSpannerFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/AbstractSpannerFn.java
deleted file mode 100644
index 00008f1..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/AbstractSpannerFn.java
+++ /dev/null
@@ -1,58 +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.sdk.io.gcp.spanner;
-
-import com.google.cloud.spanner.DatabaseClient;
-import com.google.cloud.spanner.DatabaseId;
-import com.google.cloud.spanner.Spanner;
-import com.google.cloud.spanner.SpannerOptions;
-import org.apache.beam.sdk.transforms.DoFn;
-
-/**
- * Abstract {@link DoFn} that manages {@link Spanner} lifecycle. Use {@link
- * AbstractSpannerFn#databaseClient} to access the Cloud Spanner database client.
- */
-abstract class AbstractSpannerFn<InputT, OutputT> extends DoFn<InputT, OutputT> {
- private transient Spanner spanner;
- private transient DatabaseClient databaseClient;
-
- abstract SpannerConfig getSpannerConfig();
-
- @Setup
- public void setup() throws Exception {
- SpannerConfig spannerConfig = getSpannerConfig();
- SpannerOptions options = spannerConfig.buildSpannerOptions();
- spanner = options.getService();
- databaseClient = spanner.getDatabaseClient(DatabaseId
- .of(options.getProjectId(), spannerConfig.getInstanceId().get(),
- spannerConfig.getDatabaseId().get()));
- }
-
- @Teardown
- public void teardown() throws Exception {
- if (spanner == null) {
- return;
- }
- spanner.close();
- spanner = null;
- }
-
- protected DatabaseClient databaseClient() {
- return databaseClient;
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/CreateTransactionFn.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/CreateTransactionFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/CreateTransactionFn.java
deleted file mode 100644
index da8e8b1..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/CreateTransactionFn.java
+++ /dev/null
@@ -1,51 +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.sdk.io.gcp.spanner;
-
-import com.google.cloud.spanner.ReadOnlyTransaction;
-import com.google.cloud.spanner.ResultSet;
-import com.google.cloud.spanner.Statement;
-
-/** Creates a batch transaction. */
-class CreateTransactionFn extends AbstractSpannerFn<Object, Transaction> {
-
- private final SpannerIO.CreateTransaction config;
-
- CreateTransactionFn(SpannerIO.CreateTransaction config) {
- this.config = config;
- }
-
- @ProcessElement
- public void processElement(ProcessContext c) throws Exception {
- try (ReadOnlyTransaction readOnlyTransaction =
- databaseClient().readOnlyTransaction(config.getTimestampBound())) {
- // Run a dummy sql statement to force the RPC and obtain the timestamp from the server.
- ResultSet resultSet = readOnlyTransaction.executeQuery(Statement.of("SELECT 1"));
- while (resultSet.next()) {
- // do nothing
- }
- Transaction tx = Transaction.create(readOnlyTransaction.getReadTimestamp());
- c.output(tx);
- }
- }
-
- @Override
- SpannerConfig getSpannerConfig() {
- return config.getSpannerConfig();
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationGroup.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationGroup.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationGroup.java
deleted file mode 100644
index 5b08da2..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationGroup.java
+++ /dev/null
@@ -1,67 +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.sdk.io.gcp.spanner;
-
-import com.google.cloud.spanner.Mutation;
-import com.google.common.collect.ImmutableList;
-
-import java.io.Serializable;
-import java.util.Arrays;
-import java.util.Iterator;
-import java.util.List;
-
-/**
- * A bundle of mutations that must be submitted atomically.
- *
- * <p>One of the mutations is chosen to be "primary", and can be used to determine partitions.
- */
-public final class MutationGroup implements Serializable, Iterable<Mutation> {
- private final ImmutableList<Mutation> mutations;
-
- /**
- * Creates a new group.
- *
- * @param primary a primary mutation.
- * @param other other mutations, usually interleaved in parent.
- * @return new mutation group.
- */
- public static MutationGroup create(Mutation primary, Mutation... other) {
- return create(primary, Arrays.asList(other));
- }
-
- public static MutationGroup create(Mutation primary, Iterable<Mutation> other) {
- return new MutationGroup(ImmutableList.<Mutation>builder().add(primary).addAll(other).build());
- }
-
- @Override
- public Iterator<Mutation> iterator() {
- return mutations.iterator();
- }
-
- private MutationGroup(ImmutableList<Mutation> mutations) {
- this.mutations = mutations;
- }
-
- public Mutation primary() {
- return mutations.get(0);
- }
-
- public List<Mutation> attached() {
- return mutations.subList(1, mutations.size());
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimator.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimator.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimator.java
index 2418816..61652e7 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimator.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimator.java
@@ -44,15 +44,6 @@ class MutationSizeEstimator {
return result;
}
- /** Estimates a size of the mutation group in bytes. */
- public static long sizeOf(MutationGroup group) {
- long result = 0;
- for (Mutation m : group) {
- result += sizeOf(m);
- }
- return result;
- }
-
private static long estimatePrimitiveValue(Value v) {
switch (v.getType().getCode()) {
case BOOL:
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/NaiveSpannerReadFn.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/NaiveSpannerReadFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/NaiveSpannerReadFn.java
deleted file mode 100644
index d193b95..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/NaiveSpannerReadFn.java
+++ /dev/null
@@ -1,65 +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.sdk.io.gcp.spanner;
-
-import com.google.cloud.spanner.ReadOnlyTransaction;
-import com.google.cloud.spanner.ResultSet;
-import com.google.cloud.spanner.Struct;
-import com.google.cloud.spanner.TimestampBound;
-import com.google.common.annotations.VisibleForTesting;
-
-/** A simplest read function implementation. Parallelism support is coming. */
-@VisibleForTesting
-class NaiveSpannerReadFn extends AbstractSpannerFn<Object, Struct> {
- private final SpannerIO.Read config;
-
- NaiveSpannerReadFn(SpannerIO.Read config) {
- this.config = config;
- }
-
- SpannerConfig getSpannerConfig() {
- return config.getSpannerConfig();
- }
-
- @ProcessElement
- public void processElement(ProcessContext c) throws Exception {
- TimestampBound timestampBound = TimestampBound.strong();
- if (config.getTransaction() != null) {
- Transaction transaction = c.sideInput(config.getTransaction());
- timestampBound = TimestampBound.ofReadTimestamp(transaction.timestamp());
- }
- try (ReadOnlyTransaction readOnlyTransaction =
- databaseClient().readOnlyTransaction(timestampBound)) {
- ResultSet resultSet = execute(readOnlyTransaction);
- while (resultSet.next()) {
- c.output(resultSet.getCurrentRowAsStruct());
- }
- }
- }
-
- private ResultSet execute(ReadOnlyTransaction readOnlyTransaction) {
- if (config.getQuery() != null) {
- return readOnlyTransaction.executeQuery(config.getQuery());
- }
- if (config.getIndex() != null) {
- return readOnlyTransaction.readUsingIndex(
- config.getTable(), config.getIndex(), config.getKeySet(), config.getColumns());
- }
- return readOnlyTransaction.read(config.getTable(), config.getKeySet(), config.getColumns());
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java
deleted file mode 100644
index 02716fb..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java
+++ /dev/null
@@ -1,137 +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.sdk.io.gcp.spanner;
-
-import static com.google.common.base.Preconditions.checkNotNull;
-
-import com.google.auto.value.AutoValue;
-import com.google.cloud.ServiceFactory;
-import com.google.cloud.spanner.Spanner;
-import com.google.cloud.spanner.SpannerOptions;
-import com.google.common.annotations.VisibleForTesting;
-import java.io.Serializable;
-import javax.annotation.Nullable;
-import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.options.ValueProvider;
-import org.apache.beam.sdk.transforms.display.DisplayData;
-
-/** Configuration for a Cloud Spanner client. */
-@AutoValue
-public abstract class SpannerConfig implements Serializable {
-
- @Nullable
- abstract ValueProvider<String> getProjectId();
-
- @Nullable
- abstract ValueProvider<String> getInstanceId();
-
- @Nullable
- abstract ValueProvider<String> getDatabaseId();
-
- @Nullable
- @VisibleForTesting
- abstract ServiceFactory<Spanner, SpannerOptions> getServiceFactory();
-
- abstract Builder toBuilder();
-
- SpannerOptions buildSpannerOptions() {
- SpannerOptions.Builder builder = SpannerOptions.newBuilder();
- if (getProjectId() != null) {
- builder.setProjectId(getProjectId().get());
- }
- if (getServiceFactory() != null) {
- builder.setServiceFactory(getServiceFactory());
- }
- return builder.build();
- }
-
- public static SpannerConfig create() {
- return builder().build();
- }
-
- static Builder builder() {
- return new AutoValue_SpannerConfig.Builder();
- }
-
- public void validate(PipelineOptions options) {
- checkNotNull(
- getInstanceId(),
- "SpannerIO.read() requires instance id to be set with withInstanceId method");
- checkNotNull(
- getDatabaseId(),
- "SpannerIO.read() requires database id to be set with withDatabaseId method");
- }
-
- public void populateDisplayData(DisplayData.Builder builder) {
- builder
- .addIfNotNull(DisplayData.item("projectId", getProjectId()).withLabel("Output Project"))
- .addIfNotNull(DisplayData.item("instanceId", getInstanceId()).withLabel("Output Instance"))
- .addIfNotNull(DisplayData.item("databaseId", getDatabaseId()).withLabel("Output Database"));
-
- if (getServiceFactory() != null) {
- builder.addIfNotNull(
- DisplayData.item("serviceFactory", getServiceFactory().getClass().getName())
- .withLabel("Service Factory"));
- }
- }
-
- /** Builder for {@link SpannerConfig}. */
- @AutoValue.Builder
- public abstract static class Builder {
-
- abstract Builder setProjectId(ValueProvider<String> projectId);
-
- abstract Builder setInstanceId(ValueProvider<String> instanceId);
-
- abstract Builder setDatabaseId(ValueProvider<String> databaseId);
-
- abstract Builder setServiceFactory(ServiceFactory<Spanner, SpannerOptions> serviceFactory);
-
- public abstract SpannerConfig build();
- }
-
- public SpannerConfig withProjectId(ValueProvider<String> projectId) {
- return toBuilder().setProjectId(projectId).build();
- }
-
- public SpannerConfig withProjectId(String projectId) {
- return withProjectId(ValueProvider.StaticValueProvider.of(projectId));
- }
-
- public SpannerConfig withInstanceId(ValueProvider<String> instanceId) {
- return toBuilder().setInstanceId(instanceId).build();
- }
-
- public SpannerConfig withInstanceId(String instanceId) {
- return withInstanceId(ValueProvider.StaticValueProvider.of(instanceId));
- }
-
- public SpannerConfig withDatabaseId(ValueProvider<String> databaseId) {
- return toBuilder().setDatabaseId(databaseId).build();
- }
-
- public SpannerConfig withDatabaseId(String databaseId) {
- return withDatabaseId(ValueProvider.StaticValueProvider.of(databaseId));
- }
-
- @VisibleForTesting
- SpannerConfig withServiceFactory(ServiceFactory<Spanner, SpannerOptions> serviceFactory) {
- return toBuilder().setServiceFactory(serviceFactory).build();
- }
-
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
index a247d4c..5058d13 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
@@ -17,39 +17,38 @@
*/
package org.apache.beam.sdk.io.gcp.spanner;
-import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import com.google.auto.value.AutoValue;
import com.google.cloud.ServiceFactory;
-import com.google.cloud.Timestamp;
-import com.google.cloud.spanner.KeySet;
+import com.google.cloud.ServiceOptions;
+import com.google.cloud.spanner.AbortedException;
+import com.google.cloud.spanner.DatabaseClient;
+import com.google.cloud.spanner.DatabaseId;
import com.google.cloud.spanner.Mutation;
import com.google.cloud.spanner.Spanner;
import com.google.cloud.spanner.SpannerOptions;
-import com.google.cloud.spanner.Statement;
-import com.google.cloud.spanner.Struct;
-import com.google.cloud.spanner.TimestampBound;
import com.google.common.annotations.VisibleForTesting;
-
-import java.util.Arrays;
-import java.util.Collections;
+import java.io.IOException;
+import java.util.ArrayList;
import java.util.List;
import javax.annotation.Nullable;
-
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.options.ValueProvider;
-import org.apache.beam.sdk.transforms.Create;
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.View;
import org.apache.beam.sdk.transforms.display.DisplayData;
-import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.transforms.display.DisplayData.Builder;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.BackOffUtils;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.sdk.util.Sleeper;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PDone;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
* Experimental {@link PTransform Transforms} for reading from and writing to <a
@@ -57,69 +56,7 @@ import org.apache.beam.sdk.values.PDone;
*
* <h3>Reading from Cloud Spanner</h3>
*
- * <p>To read from Cloud Spanner, apply {@link SpannerIO.Read} transformation. It will return a
- * {@link PCollection} of {@link Struct Structs}, where each element represents
- * an individual row returned from the read operation. Both Query and Read APIs are supported.
- * See more information about <a href="https://cloud.google.com/spanner/docs/reads">reading from
- * Cloud Spanner</a>
- *
- * <p>To execute a <strong>query</strong>, specify a {@link SpannerIO.Read#withQuery(Statement)} or
- * {@link SpannerIO.Read#withQuery(String)} during the construction of the transform.
- *
- * <pre>{@code
- * PCollection<Struct> rows = p.apply(
- * SpannerIO.read()
- * .withInstanceId(instanceId)
- * .withDatabaseId(dbId)
- * .withQuery("SELECT id, name, email FROM users"));
- * }</pre>
- *
- * <p>To use the Read API, specify a {@link SpannerIO.Read#withTable(String) table name} and
- * a {@link SpannerIO.Read#withColumns(List) list of columns}.
- *
- * <pre>{@code
- * PCollection<Struct> rows = p.apply(
- * SpannerIO.read()
- * .withInstanceId(instanceId)
- * .withDatabaseId(dbId)
- * .withTable("users")
- * .withColumns("id", "name", "email"));
- * }</pre>
- *
- * <p>To optimally read using index, specify the index name using {@link SpannerIO.Read#withIndex}.
- *
- * <p>The transform is guaranteed to be executed on a consistent snapshot of data, utilizing the
- * power of read only transactions. Staleness of data can be controlled using
- * {@link SpannerIO.Read#withTimestampBound} or {@link SpannerIO.Read#withTimestamp(Timestamp)}
- * methods. <a href="https://cloud.google.com/spanner/docs/transactions">Read more</a> about
- * transactions in Cloud Spanner.
- *
- * <p>It is possible to read several {@link PCollection PCollections} within a single transaction.
- * Apply {@link SpannerIO#createTransaction()} transform, that lazily creates a transaction. The
- * result of this transformation can be passed to read operation using
- * {@link SpannerIO.Read#withTransaction(PCollectionView)}.
- *
- * <pre>{@code
- * SpannerConfig spannerConfig = ...
- *
- * PCollectionView<Transaction> tx =
- * p.apply(
- * SpannerIO.createTransaction()
- * .withSpannerConfig(spannerConfig)
- * .withTimestampBound(TimestampBound.strong()));
- *
- * PCollection<Struct> users = p.apply(
- * SpannerIO.read()
- * .withSpannerConfig(spannerConfig)
- * .withQuery("SELECT name, email FROM users")
- * .withTransaction(tx));
- *
- * PCollection<Struct> tweets = p.apply(
- * SpannerIO.read()
- * .withSpannerConfig(spannerConfig)
- * .withQuery("SELECT user, tweet, date FROM tweets")
- * .withTransaction(tx));
- * }</pre>
+ * <p>This functionality is not yet implemented.
*
* <h3>Writing to Cloud Spanner</h3>
*
@@ -151,11 +88,6 @@ import org.apache.beam.sdk.values.PDone;
* <li>If the pipeline was unexpectedly stopped, mutations that were already applied will not get
* rolled back.
* </ul>
- *
- * <p>Use {@link MutationGroup} to ensure that a small set mutations is bundled together. It is
- * guaranteed that mutations in a group are submitted in the same transaction. Build
- * {@link SpannerIO.Write} transform, and call {@link Write#grouped()} method. It will return a
- * transformation that can be applied to a PCollection of MutationGroup.
*/
@Experimental(Experimental.Kind.SOURCE_SINK)
public class SpannerIO {
@@ -163,33 +95,6 @@ public class SpannerIO {
private static final long DEFAULT_BATCH_SIZE_BYTES = 1024 * 1024; // 1 MB
/**
- * Creates an uninitialized instance of {@link Read}. Before use, the {@link Read} must be
- * configured with a {@link Read#withInstanceId} and {@link Read#withDatabaseId} that identify the
- * Cloud Spanner database.
- */
- @Experimental(Experimental.Kind.SOURCE_SINK)
- public static Read read() {
- return new AutoValue_SpannerIO_Read.Builder()
- .setSpannerConfig(SpannerConfig.create())
- .setTimestampBound(TimestampBound.strong())
- .setKeySet(KeySet.all())
- .build();
- }
-
- /**
- * Returns a transform that creates a batch transaction. By default,
- * {@link TimestampBound#strong()} transaction is created, to override this use
- * {@link CreateTransaction#withTimestampBound(TimestampBound)}.
- */
- @Experimental
- public static CreateTransaction createTransaction() {
- return new AutoValue_SpannerIO_CreateTransaction.Builder()
- .setSpannerConfig(SpannerConfig.create())
- .setTimestampBound(TimestampBound.strong())
- .build();
- }
-
- /**
* Creates an uninitialized instance of {@link Write}. Before use, the {@link Write} must be
* configured with a {@link Write#withInstanceId} and {@link Write#withDatabaseId} that identify
* the Cloud Spanner database being written.
@@ -197,408 +102,247 @@ public class SpannerIO {
@Experimental
public static Write write() {
return new AutoValue_SpannerIO_Write.Builder()
- .setSpannerConfig(SpannerConfig.create())
.setBatchSizeBytes(DEFAULT_BATCH_SIZE_BYTES)
.build();
}
/**
- * A {@link PTransform} that reads data from Google Cloud Spanner.
+ * A {@link PTransform} that writes {@link Mutation} objects to Google Cloud Spanner.
*
* @see SpannerIO
*/
@Experimental(Experimental.Kind.SOURCE_SINK)
@AutoValue
- public abstract static class Read extends PTransform<PBegin, PCollection<Struct>> {
-
- abstract SpannerConfig getSpannerConfig();
-
- @Nullable
- abstract TimestampBound getTimestampBound();
-
- @Nullable
- abstract Statement getQuery();
+ public abstract static class Write extends PTransform<PCollection<Mutation>, PDone> {
@Nullable
- abstract String getTable();
+ abstract String getProjectId();
@Nullable
- abstract String getIndex();
+ abstract String getInstanceId();
@Nullable
- abstract List<String> getColumns();
+ abstract String getDatabaseId();
- @Nullable
- abstract KeySet getKeySet();
+ abstract long getBatchSizeBytes();
@Nullable
- abstract PCollectionView<Transaction> getTransaction();
+ @VisibleForTesting
+ abstract ServiceFactory<Spanner, SpannerOptions> getServiceFactory();
abstract Builder toBuilder();
@AutoValue.Builder
abstract static class Builder {
- abstract Builder setSpannerConfig(SpannerConfig spannerConfig);
-
- abstract Builder setTimestampBound(TimestampBound timestampBound);
-
- abstract Builder setQuery(Statement statement);
-
- abstract Builder setTable(String table);
-
- abstract Builder setIndex(String index);
+ abstract Builder setProjectId(String projectId);
- abstract Builder setColumns(List<String> columns);
+ abstract Builder setInstanceId(String instanceId);
- abstract Builder setKeySet(KeySet keySet);
+ abstract Builder setDatabaseId(String databaseId);
- abstract Builder setTransaction(PCollectionView<Transaction> transaction);
-
- abstract Read build();
- }
-
- /** Specifies the Cloud Spanner configuration. */
- public Read withSpannerConfig(SpannerConfig spannerConfig) {
- return toBuilder().setSpannerConfig(spannerConfig).build();
- }
+ abstract Builder setBatchSizeBytes(long batchSizeBytes);
- /** Specifies the Cloud Spanner project. */
- public Read withProjectId(String projectId) {
- return withProjectId(ValueProvider.StaticValueProvider.of(projectId));
- }
+ @VisibleForTesting
+ abstract Builder setServiceFactory(ServiceFactory<Spanner, SpannerOptions> serviceFactory);
- /** Specifies the Cloud Spanner project. */
- public Read withProjectId(ValueProvider<String> projectId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withProjectId(projectId));
+ abstract Write build();
}
- /** Specifies the Cloud Spanner instance. */
- public Read withInstanceId(String instanceId) {
- return withInstanceId(ValueProvider.StaticValueProvider.of(instanceId));
+ /**
+ * Returns a new {@link SpannerIO.Write} that will write to the specified Cloud Spanner project.
+ *
+ * <p>Does not modify this object.
+ */
+ public Write withProjectId(String projectId) {
+ return toBuilder().setProjectId(projectId).build();
}
- /** Specifies the Cloud Spanner instance. */
- public Read withInstanceId(ValueProvider<String> instanceId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withInstanceId(instanceId));
+ /**
+ * Returns a new {@link SpannerIO.Write} that will write to the specified Cloud Spanner
+ * instance.
+ *
+ * <p>Does not modify this object.
+ */
+ public Write withInstanceId(String instanceId) {
+ return toBuilder().setInstanceId(instanceId).build();
}
- /** Specifies the Cloud Spanner database. */
- public Read withDatabaseId(String databaseId) {
- return withDatabaseId(ValueProvider.StaticValueProvider.of(databaseId));
+ /**
+ * Returns a new {@link SpannerIO.Write} with a new batch size limit.
+ *
+ * <p>Does not modify this object.
+ */
+ public Write withBatchSizeBytes(long batchSizeBytes) {
+ return toBuilder().setBatchSizeBytes(batchSizeBytes).build();
}
- /** Specifies the Cloud Spanner database. */
- public Read withDatabaseId(ValueProvider<String> databaseId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withDatabaseId(databaseId));
+ /**
+ * Returns a new {@link SpannerIO.Write} that will write to the specified Cloud Spanner
+ * database.
+ *
+ * <p>Does not modify this object.
+ */
+ public Write withDatabaseId(String databaseId) {
+ return toBuilder().setDatabaseId(databaseId).build();
}
@VisibleForTesting
- Read withServiceFactory(ServiceFactory<Spanner, SpannerOptions> serviceFactory) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withServiceFactory(serviceFactory));
- }
-
- public Read withTransaction(PCollectionView<Transaction> transaction) {
- return toBuilder().setTransaction(transaction).build();
- }
-
- public Read withTimestamp(Timestamp timestamp) {
- return withTimestampBound(TimestampBound.ofReadTimestamp(timestamp));
- }
-
- public Read withTimestampBound(TimestampBound timestampBound) {
- return toBuilder().setTimestampBound(timestampBound).build();
- }
-
- public Read withTable(String table) {
- return toBuilder().setTable(table).build();
- }
-
- public Read withColumns(String... columns) {
- return withColumns(Arrays.asList(columns));
- }
-
- public Read withColumns(List<String> columns) {
- return toBuilder().setColumns(columns).build();
- }
-
- public Read withQuery(Statement statement) {
- return toBuilder().setQuery(statement).build();
- }
-
- public Read withQuery(String sql) {
- return withQuery(Statement.of(sql));
- }
-
- public Read withKeySet(KeySet keySet) {
- return toBuilder().setKeySet(keySet).build();
- }
-
- public Read withIndex(String index) {
- return toBuilder().setIndex(index).build();
+ Write withServiceFactory(ServiceFactory<Spanner, SpannerOptions> serviceFactory) {
+ return toBuilder().setServiceFactory(serviceFactory).build();
}
-
@Override
public void validate(PipelineOptions options) {
- getSpannerConfig().validate(options);
checkNotNull(
- getTimestampBound(),
- "SpannerIO.read() runs in a read only transaction and requires timestamp to be set "
- + "with withTimestampBound or withTimestamp method");
-
- if (getQuery() != null) {
- // TODO: validate query?
- } else if (getTable() != null) {
- // Assume read
- checkNotNull(
- getColumns(),
- "For a read operation SpannerIO.read() requires a list of "
- + "columns to set with withColumns method");
- checkArgument(
- !getColumns().isEmpty(),
- "For a read operation SpannerIO.read() requires a"
- + " list of columns to set with withColumns method");
- } else {
- throw new IllegalArgumentException(
- "SpannerIO.read() requires configuring query or read operation.");
- }
+ getInstanceId(),
+ "SpannerIO.write() requires instance id to be set with withInstanceId method");
+ checkNotNull(
+ getDatabaseId(),
+ "SpannerIO.write() requires database id to be set with withDatabaseId method");
}
@Override
- public PCollection<Struct> expand(PBegin input) {
- Read config = this;
- List<PCollectionView<Transaction>> sideInputs = Collections.emptyList();
- if (getTimestampBound() != null) {
- PCollectionView<Transaction> transaction =
- input.apply(createTransaction().withSpannerConfig(getSpannerConfig()));
- config = config.withTransaction(transaction);
- sideInputs = Collections.singletonList(transaction);
- }
- return input
- .apply(Create.of(1))
- .apply(
- "Execute query", ParDo.of(new NaiveSpannerReadFn(config)).withSideInputs(sideInputs));
- }
- }
-
- /**
- * A {@link PTransform} that create a transaction.
- *
- * @see SpannerIO
- */
- @Experimental(Experimental.Kind.SOURCE_SINK)
- @AutoValue
- public abstract static class CreateTransaction
- extends PTransform<PBegin, PCollectionView<Transaction>> {
-
- abstract SpannerConfig getSpannerConfig();
-
- @Nullable
- abstract TimestampBound getTimestampBound();
-
- abstract Builder toBuilder();
-
- @Override
- public PCollectionView<Transaction> expand(PBegin input) {
- return input.apply(Create.of(1))
- .apply("Create transaction", ParDo.of(new CreateTransactionFn(this)))
- .apply("As PCollectionView", View.<Transaction>asSingleton());
- }
-
- /** Specifies the Cloud Spanner configuration. */
- public CreateTransaction withSpannerConfig(SpannerConfig spannerConfig) {
- return toBuilder().setSpannerConfig(spannerConfig).build();
- }
-
- /** Specifies the Cloud Spanner project. */
- public CreateTransaction withProjectId(String projectId) {
- return withProjectId(ValueProvider.StaticValueProvider.of(projectId));
- }
-
- /** Specifies the Cloud Spanner project. */
- public CreateTransaction withProjectId(ValueProvider<String> projectId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withProjectId(projectId));
- }
-
- /** Specifies the Cloud Spanner instance. */
- public CreateTransaction withInstanceId(String instanceId) {
- return withInstanceId(ValueProvider.StaticValueProvider.of(instanceId));
- }
-
- /** Specifies the Cloud Spanner instance. */
- public CreateTransaction withInstanceId(ValueProvider<String> instanceId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withInstanceId(instanceId));
- }
-
- /** Specifies the Cloud Spanner database. */
- public CreateTransaction withDatabaseId(String databaseId) {
- return withDatabaseId(ValueProvider.StaticValueProvider.of(databaseId));
- }
-
- /** Specifies the Cloud Spanner database. */
- public CreateTransaction withDatabaseId(ValueProvider<String> databaseId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withDatabaseId(databaseId));
- }
-
- @VisibleForTesting
- CreateTransaction withServiceFactory(
- ServiceFactory<Spanner, SpannerOptions> serviceFactory) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withServiceFactory(serviceFactory));
- }
-
- public CreateTransaction withTimestampBound(TimestampBound timestampBound) {
- return toBuilder().setTimestampBound(timestampBound).build();
+ public PDone expand(PCollection<Mutation> input) {
+ input.apply("Write mutations to Cloud Spanner", ParDo.of(new SpannerWriteFn(this)));
+ return PDone.in(input.getPipeline());
}
@Override
- public void validate(PipelineOptions options) {
- getSpannerConfig().validate(options);
- }
-
- /** A builder for {@link CreateTransaction}. */
- @AutoValue.Builder public abstract static class Builder {
-
- public abstract Builder setSpannerConfig(SpannerConfig spannerConfig);
-
- public abstract Builder setTimestampBound(TimestampBound newTimestampBound);
-
- public abstract CreateTransaction build();
+ public void populateDisplayData(DisplayData.Builder builder) {
+ super.populateDisplayData(builder);
+ builder
+ .addIfNotNull(DisplayData.item("projectId", getProjectId()).withLabel("Output Project"))
+ .addIfNotNull(
+ DisplayData.item("instanceId", getInstanceId()).withLabel("Output Instance"))
+ .addIfNotNull(
+ DisplayData.item("databaseId", getDatabaseId()).withLabel("Output Database"))
+ .add(DisplayData.item("batchSizeBytes", getBatchSizeBytes())
+ .withLabel("Batch Size in Bytes"));
+ if (getServiceFactory() != null) {
+ builder.addIfNotNull(
+ DisplayData.item("serviceFactory", getServiceFactory().getClass().getName())
+ .withLabel("Service Factory"));
+ }
}
}
+ /** Batches together and writes mutations to Google Cloud Spanner. */
+ @VisibleForTesting
+ static class SpannerWriteFn extends DoFn<Mutation, Void> {
+ private static final Logger LOG = LoggerFactory.getLogger(SpannerWriteFn.class);
+ private final Write spec;
+ private transient Spanner spanner;
+ private transient DatabaseClient dbClient;
+ // Current batch of mutations to be written.
+ private List<Mutation> mutations;
+ private long batchSizeBytes = 0;
+
+ private static final int MAX_RETRIES = 5;
+ private static final FluentBackoff BUNDLE_WRITE_BACKOFF =
+ FluentBackoff.DEFAULT
+ .withMaxRetries(MAX_RETRIES)
+ .withInitialBackoff(Duration.standardSeconds(5));
- /**
- * A {@link PTransform} that writes {@link Mutation} objects to Google Cloud Spanner.
- *
- * @see SpannerIO
- */
- @Experimental(Experimental.Kind.SOURCE_SINK)
- @AutoValue
- public abstract static class Write extends PTransform<PCollection<Mutation>, PDone> {
-
- abstract SpannerConfig getSpannerConfig();
-
- abstract long getBatchSizeBytes();
-
- abstract Builder toBuilder();
-
- @AutoValue.Builder
- abstract static class Builder {
-
- abstract Builder setSpannerConfig(SpannerConfig spannerConfig);
-
- abstract Builder setBatchSizeBytes(long batchSizeBytes);
-
- abstract Write build();
- }
-
- /** Specifies the Cloud Spanner configuration. */
- public Write withSpannerConfig(SpannerConfig spannerConfig) {
- return toBuilder().setSpannerConfig(spannerConfig).build();
- }
-
- /** Specifies the Cloud Spanner project. */
- public Write withProjectId(String projectId) {
- return withProjectId(ValueProvider.StaticValueProvider.of(projectId));
+ @VisibleForTesting
+ SpannerWriteFn(Write spec) {
+ this.spec = spec;
}
- /** Specifies the Cloud Spanner project. */
- public Write withProjectId(ValueProvider<String> projectId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withProjectId(projectId));
+ @Setup
+ public void setup() throws Exception {
+ SpannerOptions spannerOptions = getSpannerOptions();
+ spanner = spannerOptions.getService();
+ dbClient = spanner.getDatabaseClient(
+ DatabaseId.of(projectId(), spec.getInstanceId(), spec.getDatabaseId()));
+ mutations = new ArrayList<>();
+ batchSizeBytes = 0;
}
- /** Specifies the Cloud Spanner instance. */
- public Write withInstanceId(String instanceId) {
- return withInstanceId(ValueProvider.StaticValueProvider.of(instanceId));
+ @ProcessElement
+ public void processElement(ProcessContext c) throws Exception {
+ Mutation m = c.element();
+ mutations.add(m);
+ batchSizeBytes += MutationSizeEstimator.sizeOf(m);
+ if (batchSizeBytes >= spec.getBatchSizeBytes()) {
+ flushBatch();
+ }
}
- /** Specifies the Cloud Spanner instance. */
- public Write withInstanceId(ValueProvider<String> instanceId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withInstanceId(instanceId));
+ private String projectId() {
+ return spec.getProjectId() == null
+ ? ServiceOptions.getDefaultProjectId()
+ : spec.getProjectId();
}
- /** Specifies the Cloud Spanner database. */
- public Write withDatabaseId(String databaseId) {
- return withDatabaseId(ValueProvider.StaticValueProvider.of(databaseId));
+ @FinishBundle
+ public void finishBundle() throws Exception {
+ if (!mutations.isEmpty()) {
+ flushBatch();
+ }
}
- /** Specifies the Cloud Spanner database. */
- public Write withDatabaseId(ValueProvider<String> databaseId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withDatabaseId(databaseId));
+ @Teardown
+ public void teardown() throws Exception {
+ if (spanner == null) {
+ return;
+ }
+ spanner.closeAsync().get();
+ spanner = null;
}
- @VisibleForTesting
- Write withServiceFactory(ServiceFactory<Spanner, SpannerOptions> serviceFactory) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withServiceFactory(serviceFactory));
+ private SpannerOptions getSpannerOptions() {
+ SpannerOptions.Builder spannerOptionsBuider = SpannerOptions.newBuilder();
+ if (spec.getServiceFactory() != null) {
+ spannerOptionsBuider.setServiceFactory(spec.getServiceFactory());
+ }
+ if (spec.getProjectId() != null) {
+ spannerOptionsBuider.setProjectId(spec.getProjectId());
+ }
+ return spannerOptionsBuider.build();
}
/**
- * Same transform but can be applied to {@link PCollection} of {@link MutationGroup}.
+ * Writes a batch of mutations to Cloud Spanner.
+ *
+ * <p>If a commit fails, it will be retried up to {@link #MAX_RETRIES} times. If the retry limit
+ * is exceeded, the last exception from Cloud Spanner will be thrown.
+ *
+ * @throws AbortedException if the commit fails or IOException or InterruptedException if
+ * backing off between retries fails.
*/
- public WriteGrouped grouped() {
- return new WriteGrouped(this);
- }
-
- /** Specifies the batch size limit. */
- public Write withBatchSizeBytes(long batchSizeBytes) {
- return toBuilder().setBatchSizeBytes(batchSizeBytes).build();
- }
-
- @Override
- public void validate(PipelineOptions options) {
- getSpannerConfig().validate(options);
- }
-
- @Override
- public PDone expand(PCollection<Mutation> input) {
- input
- .apply("To mutation group", ParDo.of(new ToMutationGroupFn()))
- .apply("Write mutations to Cloud Spanner", ParDo.of(new SpannerWriteGroupFn(this)));
- return PDone.in(input.getPipeline());
+ private void flushBatch() throws AbortedException, IOException, InterruptedException {
+ LOG.debug("Writing batch of {} mutations", mutations.size());
+ Sleeper sleeper = Sleeper.DEFAULT;
+ BackOff backoff = BUNDLE_WRITE_BACKOFF.backoff();
+
+ while (true) {
+ // Batch upsert rows.
+ try {
+ dbClient.writeAtLeastOnce(mutations);
+
+ // Break if the commit threw no exception.
+ break;
+ } catch (AbortedException exception) {
+ // Only log the code and message for potentially-transient errors. The entire exception
+ // will be propagated upon the last retry.
+ LOG.error(
+ "Error writing to Spanner ({}): {}", exception.getCode(), exception.getMessage());
+ if (!BackOffUtils.next(sleeper, backoff)) {
+ LOG.error("Aborting after {} retries.", MAX_RETRIES);
+ throw exception;
+ }
+ }
+ }
+ LOG.debug("Successfully wrote {} mutations", mutations.size());
+ mutations = new ArrayList<>();
+ batchSizeBytes = 0;
}
-
@Override
- public void populateDisplayData(DisplayData.Builder builder) {
+ public void populateDisplayData(Builder builder) {
super.populateDisplayData(builder);
- getSpannerConfig().populateDisplayData(builder);
- builder.add(
- DisplayData.item("batchSizeBytes", getBatchSizeBytes()).withLabel("Batch Size in Bytes"));
- }
- }
-
- /** Same as {@link Write} but supports grouped mutations. */
- public static class WriteGrouped extends PTransform<PCollection<MutationGroup>, PDone> {
- private final Write spec;
-
- public WriteGrouped(Write spec) {
- this.spec = spec;
- }
-
- @Override public PDone expand(PCollection<MutationGroup> input) {
- input.apply("Write mutations to Cloud Spanner", ParDo.of(new SpannerWriteGroupFn(spec)));
- return PDone.in(input.getPipeline());
- }
- }
-
- private static class ToMutationGroupFn extends DoFn<Mutation, MutationGroup> {
- @ProcessElement
- public void processElement(ProcessContext c) throws Exception {
- Mutation value = c.element();
- c.output(MutationGroup.create(value));
+ spec.populateDisplayData(builder);
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteGroupFn.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteGroupFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteGroupFn.java
deleted file mode 100644
index 34a11da..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteGroupFn.java
+++ /dev/null
@@ -1,125 +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.sdk.io.gcp.spanner;
-
-import com.google.cloud.spanner.AbortedException;
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.Iterables;
-
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
-
-import org.apache.beam.sdk.transforms.display.DisplayData;
-import org.apache.beam.sdk.util.BackOff;
-import org.apache.beam.sdk.util.BackOffUtils;
-import org.apache.beam.sdk.util.FluentBackoff;
-import org.apache.beam.sdk.util.Sleeper;
-import org.joda.time.Duration;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/** Batches together and writes mutations to Google Cloud Spanner. */
-@VisibleForTesting class SpannerWriteGroupFn extends AbstractSpannerFn<MutationGroup, Void> {
- private static final Logger LOG = LoggerFactory.getLogger(SpannerWriteGroupFn.class);
- private final SpannerIO.Write spec;
- // Current batch of mutations to be written.
- private List<MutationGroup> mutations;
- private long batchSizeBytes = 0;
-
- private static final int MAX_RETRIES = 5;
- private static final FluentBackoff BUNDLE_WRITE_BACKOFF =
- FluentBackoff.DEFAULT
- .withMaxRetries(MAX_RETRIES)
- .withInitialBackoff(Duration.standardSeconds(5));
-
- @VisibleForTesting SpannerWriteGroupFn(SpannerIO.Write spec) {
- this.spec = spec;
- }
-
- @Override SpannerConfig getSpannerConfig() {
- return spec.getSpannerConfig();
- }
-
- @Setup
- public void setup() throws Exception {
- super.setup();
- mutations = new ArrayList<>();
- batchSizeBytes = 0;
- }
-
- @ProcessElement
- public void processElement(ProcessContext c) throws Exception {
- MutationGroup m = c.element();
- mutations.add(m);
- batchSizeBytes += MutationSizeEstimator.sizeOf(m);
- if (batchSizeBytes >= spec.getBatchSizeBytes()) {
- flushBatch();
- }
- }
-
- @FinishBundle
- public void finishBundle() throws Exception {
- if (!mutations.isEmpty()) {
- flushBatch();
- }
- }
-
- /**
- * Writes a batch of mutations to Cloud Spanner.
- *
- * <p>If a commit fails, it will be retried up to {@link #MAX_RETRIES} times. If the retry limit
- * is exceeded, the last exception from Cloud Spanner will be thrown.
- *
- * @throws AbortedException if the commit fails or IOException or InterruptedException if
- * backing off between retries fails.
- */
- private void flushBatch() throws AbortedException, IOException, InterruptedException {
- LOG.debug("Writing batch of {} mutations", mutations.size());
- Sleeper sleeper = Sleeper.DEFAULT;
- BackOff backoff = BUNDLE_WRITE_BACKOFF.backoff();
-
- while (true) {
- // Batch upsert rows.
- try {
- databaseClient().writeAtLeastOnce(Iterables.concat(mutations));
-
- // Break if the commit threw no exception.
- break;
- } catch (AbortedException exception) {
- // Only log the code and message for potentially-transient errors. The entire exception
- // will be propagated upon the last retry.
- LOG.error(
- "Error writing to Spanner ({}): {}", exception.getCode(), exception.getMessage());
- if (!BackOffUtils.next(sleeper, backoff)) {
- LOG.error("Aborting after {} retries.", MAX_RETRIES);
- throw exception;
- }
- }
- }
- LOG.debug("Successfully wrote {} mutations", mutations.size());
- mutations = new ArrayList<>();
- batchSizeBytes = 0;
- }
-
- @Override
- public void populateDisplayData(DisplayData.Builder builder) {
- super.populateDisplayData(builder);
- spec.populateDisplayData(builder);
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/Transaction.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/Transaction.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/Transaction.java
deleted file mode 100644
index 22af3b8..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/Transaction.java
+++ /dev/null
@@ -1,33 +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.sdk.io.gcp.spanner;
-
-import com.google.auto.value.AutoValue;
-import com.google.cloud.Timestamp;
-import java.io.Serializable;
-
-/** A transaction object. */
-@AutoValue
-public abstract class Transaction implements Serializable {
-
- abstract Timestamp timestamp();
-
- public static Transaction create(Timestamp timestamp) {
- return new AutoValue_Transaction(timestamp);
- }
-}
[22/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutabilityEnforcementFactoryTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutabilityEnforcementFactoryTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutabilityEnforcementFactoryTest.java
index 365b6c4..c0919b9 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutabilityEnforcementFactoryTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ImmutabilityEnforcementFactoryTest.java
@@ -64,9 +64,7 @@ public class ImmutabilityEnforcementFactoryTest implements Serializable {
c.element()[0] = 'b';
}
}));
- PCollection<Long> consumer = pcollection.apply(Count.<byte[]>globally());
- DirectGraphs.performDirectOverrides(p);
- this.consumer = DirectGraphs.getProducer(consumer);
+ consumer = DirectGraphs.getProducer(pcollection.apply(Count.<byte[]>globally()));
}
@Test
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ParDoEvaluatorTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ParDoEvaluatorTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ParDoEvaluatorTest.java
index 7912538..09a21ac 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ParDoEvaluatorTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ParDoEvaluatorTest.java
@@ -98,7 +98,7 @@ public class ParDoEvaluatorTest {
when(evaluationContext.createBundle(output)).thenReturn(outputBundle);
ParDoEvaluator<Integer> evaluator =
- createEvaluator(singletonView, fn, inputPc, output);
+ createEvaluator(singletonView, fn, output);
IntervalWindow nonGlobalWindow = new IntervalWindow(new Instant(0), new Instant(10_000L));
WindowedValue<Integer> first = WindowedValue.valueInGlobalWindow(3);
@@ -132,7 +132,6 @@ public class ParDoEvaluatorTest {
private ParDoEvaluator<Integer> createEvaluator(
PCollectionView<Integer> singletonView,
RecorderFn fn,
- PCollection<Integer> input,
PCollection<Integer> output) {
when(
evaluationContext.createSideInputReader(
@@ -150,7 +149,6 @@ public class ParDoEvaluatorTest {
Mockito.any(AppliedPTransform.class), Mockito.any(StructuralKey.class)))
.thenReturn(executionContext);
- DirectGraphs.performDirectOverrides(p);
@SuppressWarnings("unchecked")
AppliedPTransform<PCollection<Integer>, ?, ?> transform =
(AppliedPTransform<PCollection<Integer>, ?, ?>) DirectGraphs.getProducer(output);
@@ -158,7 +156,8 @@ public class ParDoEvaluatorTest {
evaluationContext,
stepContext,
transform,
- input.getWindowingStrategy(),
+ ((PCollection<?>) Iterables.getOnlyElement(transform.getInputs().values()))
+ .getWindowingStrategy(),
fn,
null /* key */,
ImmutableList.<PCollectionView<?>>of(singletonView),
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactoryTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactoryTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactoryTest.java
index fe0b743..9366b7c 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactoryTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/StatefulParDoEvaluatorFactoryTest.java
@@ -41,7 +41,6 @@ import org.apache.beam.runners.core.StateNamespace;
import org.apache.beam.runners.core.StateNamespaces;
import org.apache.beam.runners.core.StateTag;
import org.apache.beam.runners.core.StateTags;
-import org.apache.beam.runners.core.construction.TransformInputs;
import org.apache.beam.runners.direct.ParDoMultiOverrideFactory.StatefulParDo;
import org.apache.beam.runners.direct.WatermarkManager.TimerUpdate;
import org.apache.beam.sdk.coders.StringUtf8Coder;
@@ -53,6 +52,7 @@ import org.apache.beam.sdk.state.ValueState;
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.ParDo;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
@@ -128,17 +128,16 @@ public class StatefulParDoEvaluatorFactoryTest implements Serializable {
input
.apply(
new ParDoMultiOverrideFactory.GbkThenStatefulParDo<>(
- new DoFn<KV<String, Integer>, Integer>() {
- @StateId(stateId)
- private final StateSpec<ValueState<String>> spec =
- StateSpecs.value(StringUtf8Coder.of());
-
- @ProcessElement
- public void process(ProcessContext c) {}
- },
- mainOutput,
- TupleTagList.empty(),
- Collections.<PCollectionView<?>>emptyList()))
+ ParDo.of(
+ new DoFn<KV<String, Integer>, Integer>() {
+ @StateId(stateId)
+ private final StateSpec<ValueState<String>> spec =
+ StateSpecs.value(StringUtf8Coder.of());
+
+ @ProcessElement
+ public void process(ProcessContext c) {}
+ })
+ .withOutputTags(mainOutput, TupleTagList.empty())))
.get(mainOutput)
.setCoder(VarIntCoder.of());
@@ -154,7 +153,8 @@ public class StatefulParDoEvaluatorFactoryTest implements Serializable {
when(mockEvaluationContext.getExecutionContext(
eq(producingTransform), Mockito.<StructuralKey>any()))
.thenReturn(mockExecutionContext);
- when(mockExecutionContext.getStepContext(anyString())).thenReturn(mockStepContext);
+ when(mockExecutionContext.getStepContext(anyString()))
+ .thenReturn(mockStepContext);
IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(9));
IntervalWindow secondWindow = new IntervalWindow(new Instant(10), new Instant(19));
@@ -241,17 +241,18 @@ public class StatefulParDoEvaluatorFactoryTest implements Serializable {
mainInput
.apply(
new ParDoMultiOverrideFactory.GbkThenStatefulParDo<>(
- new DoFn<KV<String, Integer>, Integer>() {
- @StateId(stateId)
- private final StateSpec<ValueState<String>> spec =
- StateSpecs.value(StringUtf8Coder.of());
-
- @ProcessElement
- public void process(ProcessContext c) {}
- },
- mainOutput,
- TupleTagList.empty(),
- Collections.<PCollectionView<?>>singletonList(sideInput)))
+ ParDo
+ .of(
+ new DoFn<KV<String, Integer>, Integer>() {
+ @StateId(stateId)
+ private final StateSpec<ValueState<String>> spec =
+ StateSpecs.value(StringUtf8Coder.of());
+
+ @ProcessElement
+ public void process(ProcessContext c) {}
+ })
+ .withSideInputs(sideInput)
+ .withOutputTags(mainOutput, TupleTagList.empty())))
.get(mainOutput)
.setCoder(VarIntCoder.of());
@@ -268,7 +269,8 @@ public class StatefulParDoEvaluatorFactoryTest implements Serializable {
when(mockEvaluationContext.getExecutionContext(
eq(producingTransform), Mockito.<StructuralKey>any()))
.thenReturn(mockExecutionContext);
- when(mockExecutionContext.getStepContext(anyString())).thenReturn(mockStepContext);
+ when(mockExecutionContext.getStepContext(anyString()))
+ .thenReturn(mockStepContext);
when(mockEvaluationContext.createBundle(Matchers.<PCollection<Integer>>any()))
.thenReturn(mockUncommittedBundle);
when(mockStepContext.getTimerUpdate()).thenReturn(TimerUpdate.empty());
@@ -285,8 +287,11 @@ public class StatefulParDoEvaluatorFactoryTest implements Serializable {
// global window state merely by having the evaluator created. The cleanup logic does not
// depend on the window.
String key = "hello";
- WindowedValue<KV<String, Integer>> firstKv =
- WindowedValue.of(KV.of(key, 1), new Instant(3), firstWindow, PaneInfo.NO_FIRING);
+ WindowedValue<KV<String, Integer>> firstKv = WindowedValue.of(
+ KV.of(key, 1),
+ new Instant(3),
+ firstWindow,
+ PaneInfo.NO_FIRING);
WindowedValue<KeyedWorkItem<String, KV<String, Integer>>> gbkOutputElement =
firstKv.withValue(
@@ -301,8 +306,7 @@ public class StatefulParDoEvaluatorFactoryTest implements Serializable {
BUNDLE_FACTORY
.createBundle(
(PCollection<KeyedWorkItem<String, KV<String, Integer>>>)
- Iterables.getOnlyElement(
- TransformInputs.nonAdditionalInputs(producingTransform)))
+ Iterables.getOnlyElement(producingTransform.getInputs().values()))
.add(gbkOutputElement)
.commit(Instant.now());
TransformEvaluator<KeyedWorkItem<String, KV<String, Integer>>> evaluator =
@@ -312,7 +316,8 @@ public class StatefulParDoEvaluatorFactoryTest implements Serializable {
// This should push back every element as a KV<String, Iterable<Integer>>
// in the appropriate window. Since the keys are equal they are single-threaded
- TransformResult<KeyedWorkItem<String, KV<String, Integer>>> result = evaluator.finishBundle();
+ TransformResult<KeyedWorkItem<String, KV<String, Integer>>> result =
+ evaluator.finishBundle();
List<Integer> pushedBackInts = new ArrayList<>();
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/TransformExecutorTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/TransformExecutorTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/TransformExecutorTest.java
index b7f5a7c..86412a0 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/TransformExecutorTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/TransformExecutorTest.java
@@ -25,8 +25,6 @@ import static org.hamcrest.Matchers.nullValue;
import static org.junit.Assert.assertThat;
import static org.mockito.Mockito.when;
-import com.google.common.base.Optional;
-import com.google.common.collect.Iterables;
import com.google.common.util.concurrent.MoreExecutors;
import java.util.ArrayList;
import java.util.Collection;
@@ -92,7 +90,6 @@ public class TransformExecutorTest {
created = p.apply(Create.of("foo", "spam", "third"));
PCollection<KV<Integer, String>> downstream = created.apply(WithKeys.<Integer, String>of(3));
- DirectGraphs.performDirectOverrides(p);
DirectGraph graph = DirectGraphs.getGraph(p);
createdProducer = graph.getProducer(created);
downstreamProducer = graph.getProducer(downstream);
@@ -417,13 +414,8 @@ public class TransformExecutorTest {
? Collections.emptyList()
: result.getUnprocessedElements();
- Optional<? extends CommittedBundle<?>> unprocessedBundle;
- if (inputBundle == null || Iterables.isEmpty(unprocessedElements)) {
- unprocessedBundle = Optional.absent();
- } else {
- unprocessedBundle =
- Optional.<CommittedBundle<?>>of(inputBundle.withElements(unprocessedElements));
- }
+ CommittedBundle<?> unprocessedBundle =
+ inputBundle == null ? null : inputBundle.withElements(unprocessedElements);
return CommittedResult.create(
result,
unprocessedBundle,
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ViewEvaluatorFactoryTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ViewEvaluatorFactoryTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ViewEvaluatorFactoryTest.java
index 5bc48b7..419698e 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ViewEvaluatorFactoryTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ViewEvaluatorFactoryTest.java
@@ -36,6 +36,7 @@ import org.apache.beam.sdk.transforms.View.CreatePCollectionView;
import org.apache.beam.sdk.transforms.WithKeys;
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.PCollectionViews;
import org.joda.time.Instant;
import org.junit.Rule;
@@ -65,13 +66,12 @@ public class ViewEvaluatorFactoryTest {
.setCoder(KvCoder.of(VoidCoder.of(), StringUtf8Coder.of()))
.apply(GroupByKey.<Void, String>create())
.apply(Values.<Iterable<String>>create());
- PCollection<Iterable<String>> view =
- concat.apply(
- new ViewOverrideFactory.WriteView<String, Iterable<String>>(createView.getView()));
+ PCollectionView<Iterable<String>> view =
+ concat.apply(new ViewOverrideFactory.WriteView<>(createView));
EvaluationContext context = mock(EvaluationContext.class);
TestViewWriter<String, Iterable<String>> viewWriter = new TestViewWriter<>();
- when(context.createPCollectionViewWriter(concat, createView.getView())).thenReturn(viewWriter);
+ when(context.createPCollectionViewWriter(concat, view)).thenReturn(viewWriter);
CommittedBundle<String> inputBundle = bundleFactory.createBundle(input).commit(Instant.now());
AppliedPTransform<?, ?, ?> producer = DirectGraphs.getProducer(view);
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ViewOverrideFactoryTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ViewOverrideFactoryTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ViewOverrideFactoryTest.java
index 6af9273..024e15c 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ViewOverrideFactoryTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/ViewOverrideFactoryTest.java
@@ -37,11 +37,8 @@ import org.apache.beam.sdk.testing.PAssert;
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.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.View.CreatePCollectionView;
-import org.apache.beam.sdk.transforms.ViewFn;
-import org.apache.beam.sdk.transforms.windowing.WindowMappingFn;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PCollectionViews;
@@ -66,19 +63,24 @@ public class ViewOverrideFactoryTest implements Serializable {
PCollection<Integer> ints = p.apply("CreateContents", Create.of(1, 2, 3));
final PCollectionView<List<Integer>> view =
PCollectionViews.listView(ints, WindowingStrategy.globalDefault(), ints.getCoder());
- PTransformReplacement<PCollection<Integer>, PCollection<Integer>>
+ PTransformReplacement<PCollection<Integer>, PCollectionView<List<Integer>>>
replacementTransform =
factory.getReplacementTransform(
AppliedPTransform
- .<PCollection<Integer>, PCollection<Integer>,
- PTransform<PCollection<Integer>, PCollection<Integer>>>
+ .<PCollection<Integer>, PCollectionView<List<Integer>>,
+ CreatePCollectionView<Integer, List<Integer>>>
of(
"foo",
ints.expand(),
view.expand(),
CreatePCollectionView.<Integer, List<Integer>>of(view),
p));
- ints.apply(replacementTransform.getTransform());
+ PCollectionView<List<Integer>> afterReplacement =
+ ints.apply(replacementTransform.getTransform());
+ assertThat(
+ "The CreatePCollectionView replacement should return the same View",
+ afterReplacement,
+ equalTo(view));
PCollection<Set<Integer>> outputViewContents =
p.apply("CreateSingleton", Create.of(0))
@@ -102,11 +104,11 @@ public class ViewOverrideFactoryTest implements Serializable {
final PCollection<Integer> ints = p.apply("CreateContents", Create.of(1, 2, 3));
final PCollectionView<List<Integer>> view =
PCollectionViews.listView(ints, WindowingStrategy.globalDefault(), ints.getCoder());
- PTransformReplacement<PCollection<Integer>, PCollection<Integer>> replacement =
+ PTransformReplacement<PCollection<Integer>, PCollectionView<List<Integer>>> replacement =
factory.getReplacementTransform(
AppliedPTransform
- .<PCollection<Integer>, PCollection<Integer>,
- PTransform<PCollection<Integer>, PCollection<Integer>>>
+ .<PCollection<Integer>, PCollectionView<List<Integer>>,
+ CreatePCollectionView<Integer, List<Integer>>>
of(
"foo",
ints.expand(),
@@ -124,19 +126,8 @@ public class ViewOverrideFactoryTest implements Serializable {
"There should only be one WriteView primitive in the graph",
writeViewVisited.getAndSet(true),
is(false));
- PCollectionView<?> replacementView = ((WriteView) node.getTransform()).getView();
-
- // replacementView.getPCollection() is null, but that is not a requirement
- // so not asserted one way or the other
- assertThat(
- replacementView.getTagInternal(),
- equalTo(view.getTagInternal()));
- assertThat(
- replacementView.getViewFn(),
- Matchers.<ViewFn<?, ?>>equalTo(view.getViewFn()));
- assertThat(
- replacementView.getWindowMappingFn(),
- Matchers.<WindowMappingFn<?>>equalTo(view.getWindowMappingFn()));
+ PCollectionView replacementView = ((WriteView) node.getTransform()).getView();
+ assertThat(replacementView, Matchers.<PCollectionView>theInstance(view));
assertThat(node.getInputs().entrySet(), hasSize(1));
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WatermarkCallbackExecutorTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WatermarkCallbackExecutorTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WatermarkCallbackExecutorTest.java
index 1d8aac1..b667346 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WatermarkCallbackExecutorTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WatermarkCallbackExecutorTest.java
@@ -59,7 +59,6 @@ public class WatermarkCallbackExecutorTest {
public void setup() {
PCollection<Integer> created = p.apply(Create.of(1, 2, 3));
PCollection<Integer> summed = created.apply(Sum.integersGlobally());
- DirectGraphs.performDirectOverrides(p);
DirectGraph graph = DirectGraphs.getGraph(p);
create = graph.getProducer(created);
sum = graph.getProducer(summed);
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WatermarkManagerTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WatermarkManagerTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WatermarkManagerTest.java
index e3f6215..9528ac9 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WatermarkManagerTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WatermarkManagerTest.java
@@ -24,7 +24,6 @@ import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.not;
import static org.junit.Assert.assertThat;
-import com.google.common.base.Optional;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
@@ -122,7 +121,6 @@ public class WatermarkManagerTest implements Serializable {
flattened = preFlatten.apply("flattened", Flatten.<Integer>pCollections());
clock = MockClock.fromInstant(new Instant(1000));
- DirectGraphs.performDirectOverrides(p);
graph = DirectGraphs.getGraph(p);
manager = WatermarkManager.create(clock, graph);
@@ -319,7 +317,7 @@ public class WatermarkManagerTest implements Serializable {
TimerUpdate.empty(),
CommittedResult.create(
StepTransformResult.withoutHold(graph.getProducer(created)).build(),
- Optional.<CommittedBundle<?>>absent(),
+ root.withElements(Collections.<WindowedValue<Void>>emptyList()),
Collections.singleton(createBundle),
EnumSet.allOf(OutputType.class)),
BoundedWindow.TIMESTAMP_MAX_VALUE);
@@ -333,7 +331,7 @@ public class WatermarkManagerTest implements Serializable {
TimerUpdate.empty(),
CommittedResult.create(
StepTransformResult.withoutHold(theFlatten).build(),
- Optional.<CommittedBundle<?>>absent(),
+ createBundle.withElements(Collections.<WindowedValue<Integer>>emptyList()),
Collections.<CommittedBundle<?>>emptyList(),
EnumSet.allOf(OutputType.class)),
BoundedWindow.TIMESTAMP_MAX_VALUE);
@@ -346,7 +344,7 @@ public class WatermarkManagerTest implements Serializable {
TimerUpdate.empty(),
CommittedResult.create(
StepTransformResult.withoutHold(theFlatten).build(),
- Optional.<CommittedBundle<?>>absent(),
+ createBundle.withElements(Collections.<WindowedValue<Integer>>emptyList()),
Collections.<CommittedBundle<?>>emptyList(),
EnumSet.allOf(OutputType.class)),
BoundedWindow.TIMESTAMP_MAX_VALUE);
@@ -1502,15 +1500,9 @@ public class WatermarkManagerTest implements Serializable {
AppliedPTransform<?, ?, ?> transform,
@Nullable CommittedBundle<?> unprocessedBundle,
Iterable<? extends CommittedBundle<?>> bundles) {
- Optional<? extends CommittedBundle<?>> unprocessedElements;
- if (unprocessedBundle == null || Iterables.isEmpty(unprocessedBundle.getElements())) {
- unprocessedElements = Optional.absent();
- } else {
- unprocessedElements = Optional.of(unprocessedBundle);
- }
return CommittedResult.create(
StepTransformResult.withoutHold(transform).build(),
- unprocessedElements,
+ unprocessedBundle,
bundles,
Iterables.isEmpty(bundles)
? EnumSet.noneOf(OutputType.class)
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WriteWithShardingFactoryTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WriteWithShardingFactoryTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WriteWithShardingFactoryTest.java
index 546a181..a88d95e 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WriteWithShardingFactoryTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/WriteWithShardingFactoryTest.java
@@ -30,7 +30,6 @@ import static org.junit.Assert.assertThat;
import java.io.File;
import java.io.FileReader;
import java.io.Reader;
-import java.io.Serializable;
import java.nio.CharBuffer;
import java.util.ArrayList;
import java.util.Collections;
@@ -39,8 +38,9 @@ import java.util.UUID;
import org.apache.beam.runners.direct.WriteWithShardingFactory.CalculateShardsFn;
import org.apache.beam.sdk.coders.VarLongCoder;
import org.apache.beam.sdk.coders.VoidCoder;
-import org.apache.beam.sdk.io.DynamicFileDestinations;
+import org.apache.beam.sdk.io.DefaultFilenamePolicy;
import org.apache.beam.sdk.io.FileBasedSink;
+import org.apache.beam.sdk.io.FileBasedSink.FilenamePolicy;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.LocalResources;
import org.apache.beam.sdk.io.TextIO;
@@ -53,8 +53,6 @@ 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.DoFnTester;
-import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.transforms.SerializableFunctions;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;
@@ -73,17 +71,11 @@ import org.junit.runners.JUnit4;
* Tests for {@link WriteWithShardingFactory}.
*/
@RunWith(JUnit4.class)
-public class WriteWithShardingFactoryTest implements Serializable {
-
+public class WriteWithShardingFactoryTest {
private static final int INPUT_SIZE = 10000;
-
- @Rule public transient TemporaryFolder tmp = new TemporaryFolder();
-
- private transient WriteWithShardingFactory<Object> factory = new WriteWithShardingFactory<>();
-
- @Rule
- public final transient TestPipeline p =
- TestPipeline.create().enableAbandonedNodeEnforcement(false);
+ @Rule public TemporaryFolder tmp = new TemporaryFolder();
+ private WriteWithShardingFactory<Object> factory = new WriteWithShardingFactory<>();
+ @Rule public final TestPipeline p = TestPipeline.create().enableAbandonedNodeEnforcement(false);
@Test
public void dynamicallyReshardedWrite() throws Exception {
@@ -137,24 +129,26 @@ public class WriteWithShardingFactoryTest implements Serializable {
@Test
public void withNoShardingSpecifiedReturnsNewTransform() {
ResourceId outputDirectory = LocalResources.fromString("/foo", true /* isDirectory */);
-
- PTransform<PCollection<Object>, PDone> original =
+ FilenamePolicy policy =
+ DefaultFilenamePolicy.constructUsingStandardParameters(
+ StaticValueProvider.of(outputDirectory),
+ DefaultFilenamePolicy.DEFAULT_UNWINDOWED_SHARD_TEMPLATE,
+ "",
+ false);
+ WriteFiles<Object> original =
WriteFiles.to(
- new FileBasedSink<Object, Void>(
- StaticValueProvider.of(outputDirectory), DynamicFileDestinations.constant(null)) {
+ new FileBasedSink<Object>(StaticValueProvider.of(outputDirectory), policy) {
@Override
- public WriteOperation<Object, Void> createWriteOperation() {
+ public WriteOperation<Object> createWriteOperation() {
throw new IllegalArgumentException("Should not be used");
}
- },
- SerializableFunctions.identity());
+ });
@SuppressWarnings("unchecked")
PCollection<Object> objs = (PCollection) p.apply(Create.empty(VoidCoder.of()));
- AppliedPTransform<PCollection<Object>, PDone, PTransform<PCollection<Object>, PDone>>
- originalApplication =
- AppliedPTransform.of(
- "write", objs.expand(), Collections.<TupleTag<?>, PValue>emptyMap(), original, p);
+ AppliedPTransform<PCollection<Object>, PDone, WriteFiles<Object>> originalApplication =
+ AppliedPTransform.of(
+ "write", objs.expand(), Collections.<TupleTag<?>, PValue>emptyMap(), original, p);
assertThat(
factory.getReplacementTransform(originalApplication).getTransform(),
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/flink/pom.xml
----------------------------------------------------------------------
diff --git a/runners/flink/pom.xml b/runners/flink/pom.xml
index c063a2d..c4c6b55 100644
--- a/runners/flink/pom.xml
+++ b/runners/flink/pom.xml
@@ -22,7 +22,7 @@
<parent>
<groupId>org.apache.beam</groupId>
<artifactId>beam-runners-parent</artifactId>
- <version>2.2.0-SNAPSHOT</version>
+ <version>2.1.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>
@@ -91,6 +91,7 @@
<excludedGroups>
org.apache.beam.sdk.testing.FlattenWithHeterogeneousCoders,
org.apache.beam.sdk.testing.LargeKeys$Above100MB,
+ org.apache.beam.sdk.testing.UsesSetState,
org.apache.beam.sdk.testing.UsesCommittedMetrics,
org.apache.beam.sdk.testing.UsesTestStream,
org.apache.beam.sdk.testing.UsesSplittableParDo
@@ -380,13 +381,5 @@
<type>test-jar</type>
<scope>test</scope>
</dependency>
-
- <dependency>
- <groupId>org.apache.beam</groupId>
- <artifactId>beam-runners-core-java</artifactId>
- <type>test-jar</type>
- <scope>test</scope>
- </dependency>
-
</dependencies>
</project>
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/flink/src/main/java/org/apache/beam/runners/flink/CreateStreamingFlinkView.java
----------------------------------------------------------------------
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/CreateStreamingFlinkView.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/CreateStreamingFlinkView.java
deleted file mode 100644
index 0cc3aec..0000000
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/CreateStreamingFlinkView.java
+++ /dev/null
@@ -1,154 +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.flink;
-
-import com.google.common.collect.Iterables;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-import org.apache.beam.runners.core.construction.ReplacementOutputs;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.CoderRegistry;
-import org.apache.beam.sdk.coders.ListCoder;
-import org.apache.beam.sdk.runners.AppliedPTransform;
-import org.apache.beam.sdk.runners.PTransformOverrideFactory;
-import org.apache.beam.sdk.transforms.Combine;
-import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.transforms.View.CreatePCollectionView;
-import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.PValue;
-import org.apache.beam.sdk.values.TupleTag;
-
-/** Flink streaming overrides for various view (side input) transforms. */
-class CreateStreamingFlinkView<ElemT, ViewT>
- extends PTransform<PCollection<ElemT>, PCollection<ElemT>> {
- private final PCollectionView<ViewT> view;
-
- public CreateStreamingFlinkView(PCollectionView<ViewT> view) {
- this.view = view;
- }
-
- @Override
- public PCollection<ElemT> expand(PCollection<ElemT> input) {
- input
- .apply(Combine.globally(new Concatenate<ElemT>()).withoutDefaults())
- .apply(CreateFlinkPCollectionView.<ElemT, ViewT>of(view));
- return input;
- }
-
- /**
- * Combiner that combines {@code T}s into a single {@code List<T>} containing all inputs.
- *
- * <p>For internal use by {@link CreateStreamingFlinkView}. This combiner requires that the input
- * {@link PCollection} fits in memory. For a large {@link PCollection} this is expected to crash!
- *
- * @param <T> the type of elements to concatenate.
- */
- private static class Concatenate<T> extends Combine.CombineFn<T, List<T>, List<T>> {
- @Override
- public List<T> createAccumulator() {
- return new ArrayList<T>();
- }
-
- @Override
- public List<T> addInput(List<T> accumulator, T input) {
- accumulator.add(input);
- return accumulator;
- }
-
- @Override
- public List<T> mergeAccumulators(Iterable<List<T>> accumulators) {
- List<T> result = createAccumulator();
- for (List<T> accumulator : accumulators) {
- result.addAll(accumulator);
- }
- return result;
- }
-
- @Override
- public List<T> extractOutput(List<T> accumulator) {
- return accumulator;
- }
-
- @Override
- public Coder<List<T>> getAccumulatorCoder(CoderRegistry registry, Coder<T> inputCoder) {
- return ListCoder.of(inputCoder);
- }
-
- @Override
- public Coder<List<T>> getDefaultOutputCoder(CoderRegistry registry, Coder<T> inputCoder) {
- return ListCoder.of(inputCoder);
- }
- }
-
- /**
- * Creates a primitive {@link PCollectionView}.
- *
- * <p>For internal use only by runner implementors.
- *
- * @param <ElemT> The type of the elements of the input PCollection
- * @param <ViewT> The type associated with the {@link PCollectionView} used as a side input
- */
- public static class CreateFlinkPCollectionView<ElemT, ViewT>
- extends PTransform<PCollection<List<ElemT>>, PCollection<List<ElemT>>> {
- private PCollectionView<ViewT> view;
-
- private CreateFlinkPCollectionView(PCollectionView<ViewT> view) {
- this.view = view;
- }
-
- public static <ElemT, ViewT> CreateFlinkPCollectionView<ElemT, ViewT> of(
- PCollectionView<ViewT> view) {
- return new CreateFlinkPCollectionView<>(view);
- }
-
- @Override
- public PCollection<List<ElemT>> expand(PCollection<List<ElemT>> input) {
- return PCollection.<List<ElemT>>createPrimitiveOutputInternal(
- input.getPipeline(), input.getWindowingStrategy(), input.isBounded())
- .setCoder(input.getCoder());
- }
-
- public PCollectionView<ViewT> getView() {
- return view;
- }
- }
-
- public static class Factory<ElemT, ViewT>
- implements PTransformOverrideFactory<
- PCollection<ElemT>, PCollection<ElemT>, CreatePCollectionView<ElemT, ViewT>> {
- public Factory() {}
-
- @Override
- public PTransformReplacement<PCollection<ElemT>, PCollection<ElemT>> getReplacementTransform(
- AppliedPTransform<
- PCollection<ElemT>, PCollection<ElemT>, CreatePCollectionView<ElemT, ViewT>>
- transform) {
- return PTransformReplacement.of(
- (PCollection<ElemT>) Iterables.getOnlyElement(transform.getInputs().values()),
- new CreateStreamingFlinkView<ElemT, ViewT>(transform.getTransform().getView()));
- }
-
- @Override
- public Map<PValue, ReplacementOutput> mapOutputs(
- Map<TupleTag<?>, PValue> outputs, PCollection<ElemT> newOutput) {
- return ReplacementOutputs.singleton(outputs, newOutput);
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkBatchTranslationContext.java
----------------------------------------------------------------------
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkBatchTranslationContext.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkBatchTranslationContext.java
index 6e70198..0439119 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkBatchTranslationContext.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkBatchTranslationContext.java
@@ -20,7 +20,6 @@ package org.apache.beam.runners.flink;
import com.google.common.collect.Iterables;
import java.util.HashMap;
import java.util.Map;
-import org.apache.beam.runners.core.construction.TransformInputs;
import org.apache.beam.runners.flink.translation.types.CoderTypeInformation;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.options.PipelineOptions;
@@ -144,7 +143,7 @@ class FlinkBatchTranslationContext {
@SuppressWarnings("unchecked")
<T extends PValue> T getInput(PTransform<T, ?> transform) {
- return (T) Iterables.getOnlyElement(TransformInputs.nonAdditionalInputs(currentTransform));
+ return (T) Iterables.getOnlyElement(currentTransform.getInputs().values());
}
Map<TupleTag<?>, PValue> getOutputs(PTransform<?, ?> transform) {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineExecutionEnvironment.java
----------------------------------------------------------------------
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineExecutionEnvironment.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineExecutionEnvironment.java
index d2a2016..fe5dd87 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineExecutionEnvironment.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineExecutionEnvironment.java
@@ -84,8 +84,6 @@ class FlinkPipelineExecutionEnvironment {
this.flinkBatchEnv = null;
this.flinkStreamEnv = null;
- pipeline.replaceAll(FlinkTransformOverrides.getDefaultOverrides(options.isStreaming()));
-
PipelineTranslationOptimizer optimizer =
new PipelineTranslationOptimizer(TranslationMode.BATCH, options);
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/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 f733e2e..8da68c5 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
@@ -17,18 +17,27 @@
*/
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.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;
import org.apache.beam.sdk.runners.AppliedPTransform;
+import org.apache.beam.sdk.runners.PTransformOverride;
import org.apache.beam.sdk.runners.PTransformOverrideFactory;
import org.apache.beam.sdk.runners.TransformHierarchy;
+import org.apache.beam.sdk.transforms.Combine;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo.MultiOutput;
+import org.apache.beam.sdk.transforms.View;
+import org.apache.beam.sdk.util.InstanceBuilder;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.PValue;
@@ -64,8 +73,54 @@ class FlinkStreamingPipelineTranslator extends FlinkPipelineTranslator {
@Override
public void translate(Pipeline pipeline) {
+ List<PTransformOverride> transformOverrides =
+ ImmutableList.<PTransformOverride>builder()
+ .add(
+ PTransformOverride.of(
+ PTransformMatchers.splittableParDoMulti(),
+ 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)))
+ .add(
+ PTransformOverride.of(
+ PTransformMatchers.classEqualTo(View.AsList.class),
+ new ReflectiveOneToOneOverrideFactory(
+ FlinkStreamingViewOverrides.StreamingViewAsList.class, flinkRunner)))
+ .add(
+ PTransformOverride.of(
+ PTransformMatchers.classEqualTo(View.AsMap.class),
+ new ReflectiveOneToOneOverrideFactory(
+ FlinkStreamingViewOverrides.StreamingViewAsMap.class, flinkRunner)))
+ .add(
+ PTransformOverride.of(
+ PTransformMatchers.classEqualTo(View.AsMultimap.class),
+ new ReflectiveOneToOneOverrideFactory(
+ FlinkStreamingViewOverrides.StreamingViewAsMultimap.class, flinkRunner)))
+ .add(
+ PTransformOverride.of(
+ PTransformMatchers.classEqualTo(View.AsSingleton.class),
+ new ReflectiveOneToOneOverrideFactory(
+ FlinkStreamingViewOverrides.StreamingViewAsSingleton.class, flinkRunner)))
+ // this has to be last since the ViewAsSingleton override
+ // can expand to a Combine.GloballyAsSingletonView
+ .add(
+ PTransformOverride.of(
+ PTransformMatchers.classEqualTo(Combine.GloballyAsSingletonView.class),
+ new ReflectiveOneToOneOverrideFactory(
+ FlinkStreamingViewOverrides.StreamingCombineGloballyAsSingletonView.class,
+ flinkRunner)))
+ .build();
+
// Ensure all outputs of all reads are consumed.
UnconsumedReads.ensureAllReadsConsumed(pipeline);
+ pipeline.replaceAll(transformOverrides);
super.translate(pipeline);
}
@@ -173,6 +228,35 @@ class FlinkStreamingPipelineTranslator extends FlinkPipelineTranslator {
}
}
+ private static class ReflectiveOneToOneOverrideFactory<
+ InputT, OutputT, TransformT extends PTransform<PCollection<InputT>, PCollection<OutputT>>>
+ extends SingleInputOutputOverrideFactory<
+ PCollection<InputT>, PCollection<OutputT>, TransformT> {
+ private final Class<PTransform<PCollection<InputT>, PCollection<OutputT>>> replacement;
+ private final FlinkRunner runner;
+
+ private ReflectiveOneToOneOverrideFactory(
+ Class<PTransform<PCollection<InputT>, PCollection<OutputT>>> replacement,
+ FlinkRunner runner) {
+ this.replacement = replacement;
+ this.runner = runner;
+ }
+
+ @Override
+ public PTransformReplacement<PCollection<InputT>, PCollection<OutputT>> getReplacementTransform(
+ AppliedPTransform<PCollection<InputT>, PCollection<OutputT>, TransformT> transform) {
+ return PTransformReplacement.of(
+ PTransformReplacements.getSingletonMainInput(transform),
+ InstanceBuilder.ofType(replacement)
+ .withArg(FlinkRunner.class, runner)
+ .withArg(
+ (Class<PTransform<PCollection<InputT>, PCollection<OutputT>>>)
+ transform.getTransform().getClass(),
+ transform.getTransform())
+ .build());
+ }
+ }
+
/**
* A {@link PTransformOverrideFactory} that overrides a <a
* href="https://s.apache.org/splittable-do-fn">Splittable DoFn</a> with {@link SplittableParDo}.
@@ -188,7 +272,7 @@ class FlinkStreamingPipelineTranslator extends FlinkPipelineTranslator {
transform) {
return PTransformReplacement.of(
PTransformReplacements.getSingletonMainInput(transform),
- (SplittableParDo<InputT, OutputT, ?>) SplittableParDo.forAppliedParDo(transform));
+ new SplittableParDo<>(transform.getTransform()));
}
@Override
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/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 3d7e81f..2a7c5d6 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
@@ -29,6 +29,7 @@ import java.util.Map;
import org.apache.beam.runners.core.KeyedWorkItem;
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;
@@ -123,7 +124,7 @@ class FlinkStreamingTransformTranslators {
TRANSLATORS.put(Window.Assign.class, new WindowAssignTranslator());
TRANSLATORS.put(Flatten.PCollections.class, new FlattenPCollectionTranslator());
TRANSLATORS.put(
- CreateStreamingFlinkView.CreateFlinkPCollectionView.class,
+ FlinkStreamingViewOverrides.CreateFlinkPCollectionView.class,
new CreateViewStreamingTranslator());
TRANSLATORS.put(Reshuffle.class, new ReshuffleTranslatorStreaming());
@@ -362,13 +363,8 @@ class FlinkStreamingTransformTranslators {
Map<TupleTag<?>, OutputTag<WindowedValue<?>>> tagsToOutputTags = Maps.newHashMap();
for (Map.Entry<TupleTag<?>, PValue> entry : outputs.entrySet()) {
if (!tagsToOutputTags.containsKey(entry.getKey())) {
- tagsToOutputTags.put(
- entry.getKey(),
- new OutputTag<WindowedValue<?>>(
- entry.getKey().getId(),
- (TypeInformation) context.getTypeInfo((PCollection<?>) entry.getValue())
- )
- );
+ tagsToOutputTags.put(entry.getKey(), new OutputTag<>(entry.getKey().getId(),
+ (TypeInformation) context.getTypeInfo((PCollection<?>) entry.getValue())));
}
}
@@ -547,11 +543,14 @@ class FlinkStreamingTransformTranslators {
transform.getAdditionalOutputTags().getAll(),
context,
new ParDoTranslationHelper.DoFnOperatorFactory<
- KeyedWorkItem<String, KV<InputT, RestrictionT>>, OutputT>() {
+ KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT>() {
@Override
- public DoFnOperator<KeyedWorkItem<String, KV<InputT, RestrictionT>>, OutputT>
- createDoFnOperator(
- DoFn<KeyedWorkItem<String, KV<InputT, RestrictionT>>, OutputT> doFn,
+ public DoFnOperator<
+ KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>,
+ OutputT> createDoFnOperator(
+ DoFn<
+ KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>,
+ OutputT> doFn,
String stepName,
List<PCollectionView<?>> sideInputs,
TupleTag<OutputT> mainOutputTag,
@@ -559,8 +558,11 @@ class FlinkStreamingTransformTranslators {
FlinkStreamingTranslationContext context,
WindowingStrategy<?, ?> windowingStrategy,
Map<TupleTag<?>, OutputTag<WindowedValue<?>>> tagsToOutputTags,
- Coder<WindowedValue<KeyedWorkItem<String, KV<InputT, RestrictionT>>>>
- inputCoder,
+ Coder<
+ WindowedValue<
+ KeyedWorkItem<
+ String,
+ ElementAndRestriction<InputT, RestrictionT>>>> inputCoder,
Coder keyCoder,
Map<Integer, PCollectionView<?>> transformedSideInputs) {
return new SplittableDoFnOperator<>(
@@ -582,17 +584,17 @@ class FlinkStreamingTransformTranslators {
private static class CreateViewStreamingTranslator<ElemT, ViewT>
extends FlinkStreamingPipelineTranslator.StreamTransformTranslator<
- CreateStreamingFlinkView.CreateFlinkPCollectionView<ElemT, ViewT>> {
+ FlinkStreamingViewOverrides.CreateFlinkPCollectionView<ElemT, ViewT>> {
@Override
public void translateNode(
- CreateStreamingFlinkView.CreateFlinkPCollectionView<ElemT, ViewT> transform,
+ FlinkStreamingViewOverrides.CreateFlinkPCollectionView<ElemT, ViewT> transform,
FlinkStreamingTranslationContext context) {
// just forward
DataStream<WindowedValue<List<ElemT>>> inputDataSet =
context.getInputDataStream(context.getInput(transform));
- PCollectionView<ViewT> view = transform.getView();
+ PCollectionView<ViewT> view = context.getOutput(transform);
context.setOutputDataStream(view, inputDataSet);
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTranslationContext.java
----------------------------------------------------------------------
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTranslationContext.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTranslationContext.java
index 74a5fb9..ea5f6b3 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTranslationContext.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTranslationContext.java
@@ -22,7 +22,6 @@ import static com.google.common.base.Preconditions.checkNotNull;
import com.google.common.collect.Iterables;
import java.util.HashMap;
import java.util.Map;
-import org.apache.beam.runners.core.construction.TransformInputs;
import org.apache.beam.runners.flink.translation.types.CoderTypeInformation;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.options.PipelineOptions;
@@ -114,7 +113,7 @@ class FlinkStreamingTranslationContext {
@SuppressWarnings("unchecked")
public <T extends PValue> T getInput(PTransform<T, ?> transform) {
- return (T) Iterables.getOnlyElement(TransformInputs.nonAdditionalInputs(currentTransform));
+ return (T) Iterables.getOnlyElement(currentTransform.getInputs().values());
}
public <T extends PInput> Map<TupleTag<?>, PValue> getInputs(PTransform<T, ?> transform) {
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingViewOverrides.java
----------------------------------------------------------------------
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingViewOverrides.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingViewOverrides.java
new file mode 100644
index 0000000..ce1c895
--- /dev/null
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingViewOverrides.java
@@ -0,0 +1,372 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.flink;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.coders.ListCoder;
+import org.apache.beam.sdk.transforms.Combine;
+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.View;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.PCollectionViews;
+
+/**
+ * Flink streaming overrides for various view (side input) transforms.
+ */
+class FlinkStreamingViewOverrides {
+
+ /**
+ * Specialized implementation for
+ * {@link org.apache.beam.sdk.transforms.View.AsMap View.AsMap}
+ * for the Flink runner in streaming mode.
+ */
+ static class StreamingViewAsMap<K, V>
+ extends PTransform<PCollection<KV<K, V>>, PCollectionView<Map<K, V>>> {
+
+ private final transient FlinkRunner runner;
+
+ @SuppressWarnings("unused") // used via reflection in FlinkRunner#apply()
+ public StreamingViewAsMap(FlinkRunner runner, View.AsMap<K, V> transform) {
+ this.runner = runner;
+ }
+
+ @Override
+ public PCollectionView<Map<K, V>> expand(PCollection<KV<K, V>> input) {
+ PCollectionView<Map<K, V>> view =
+ PCollectionViews.mapView(
+ input,
+ input.getWindowingStrategy(),
+ input.getCoder());
+
+ @SuppressWarnings({"rawtypes", "unchecked"})
+ KvCoder<K, V> inputCoder = (KvCoder) input.getCoder();
+ try {
+ inputCoder.getKeyCoder().verifyDeterministic();
+ } catch (Coder.NonDeterministicException e) {
+ runner.recordViewUsesNonDeterministicKeyCoder(this);
+ }
+
+ return input
+ .apply(Combine.globally(new Concatenate<KV<K, V>>()).withoutDefaults())
+ .apply(CreateFlinkPCollectionView.<KV<K, V>, Map<K, V>>of(view));
+ }
+
+ @Override
+ protected String getKindString() {
+ return "StreamingViewAsMap";
+ }
+ }
+
+ /**
+ * Specialized expansion for {@link
+ * View.AsMultimap View.AsMultimap} for the
+ * Flink runner in streaming mode.
+ */
+ static class StreamingViewAsMultimap<K, V>
+ extends PTransform<PCollection<KV<K, V>>, PCollectionView<Map<K, Iterable<V>>>> {
+
+ private final transient FlinkRunner runner;
+
+ /**
+ * Builds an instance of this class from the overridden transform.
+ */
+ @SuppressWarnings("unused") // used via reflection in FlinkRunner#apply()
+ public StreamingViewAsMultimap(FlinkRunner runner, View.AsMultimap<K, V> transform) {
+ this.runner = runner;
+ }
+
+ @Override
+ public PCollectionView<Map<K, Iterable<V>>> expand(PCollection<KV<K, V>> input) {
+ PCollectionView<Map<K, Iterable<V>>> view =
+ PCollectionViews.multimapView(
+ input,
+ input.getWindowingStrategy(),
+ input.getCoder());
+
+ @SuppressWarnings({"rawtypes", "unchecked"})
+ KvCoder<K, V> inputCoder = (KvCoder) input.getCoder();
+ try {
+ inputCoder.getKeyCoder().verifyDeterministic();
+ } catch (Coder.NonDeterministicException e) {
+ runner.recordViewUsesNonDeterministicKeyCoder(this);
+ }
+
+ return input
+ .apply(Combine.globally(new Concatenate<KV<K, V>>()).withoutDefaults())
+ .apply(CreateFlinkPCollectionView.<KV<K, V>, Map<K, Iterable<V>>>of(view));
+ }
+
+ @Override
+ protected String getKindString() {
+ return "StreamingViewAsMultimap";
+ }
+ }
+
+ /**
+ * Specialized implementation for
+ * {@link View.AsList View.AsList} for the
+ * Flink runner in streaming mode.
+ */
+ static class StreamingViewAsList<T>
+ extends PTransform<PCollection<T>, PCollectionView<List<T>>> {
+ /**
+ * Builds an instance of this class from the overridden transform.
+ */
+ @SuppressWarnings("unused") // used via reflection in FlinkRunner#apply()
+ public StreamingViewAsList(FlinkRunner runner, View.AsList<T> transform) {}
+
+ @Override
+ public PCollectionView<List<T>> expand(PCollection<T> input) {
+ PCollectionView<List<T>> view =
+ PCollectionViews.listView(
+ input,
+ input.getWindowingStrategy(),
+ input.getCoder());
+
+ return input.apply(Combine.globally(new Concatenate<T>()).withoutDefaults())
+ .apply(CreateFlinkPCollectionView.<T, List<T>>of(view));
+ }
+
+ @Override
+ protected String getKindString() {
+ return "StreamingViewAsList";
+ }
+ }
+
+ /**
+ * Specialized implementation for
+ * {@link View.AsIterable View.AsIterable} for the
+ * Flink runner in streaming mode.
+ */
+ static class StreamingViewAsIterable<T>
+ extends PTransform<PCollection<T>, PCollectionView<Iterable<T>>> {
+ /**
+ * Builds an instance of this class from the overridden transform.
+ */
+ @SuppressWarnings("unused") // used via reflection in FlinkRunner#apply()
+ public StreamingViewAsIterable(FlinkRunner runner, View.AsIterable<T> transform) { }
+
+ @Override
+ public PCollectionView<Iterable<T>> expand(PCollection<T> input) {
+ PCollectionView<Iterable<T>> view =
+ PCollectionViews.iterableView(
+ input,
+ input.getWindowingStrategy(),
+ input.getCoder());
+
+ return input.apply(Combine.globally(new Concatenate<T>()).withoutDefaults())
+ .apply(CreateFlinkPCollectionView.<T, Iterable<T>>of(view));
+ }
+
+ @Override
+ protected String getKindString() {
+ return "StreamingViewAsIterable";
+ }
+ }
+
+ /**
+ * Specialized expansion for
+ * {@link View.AsSingleton View.AsSingleton} for the
+ * Flink runner in streaming mode.
+ */
+ static class StreamingViewAsSingleton<T>
+ extends PTransform<PCollection<T>, PCollectionView<T>> {
+ private View.AsSingleton<T> transform;
+
+ /**
+ * Builds an instance of this class from the overridden transform.
+ */
+ @SuppressWarnings("unused") // used via reflection in FlinkRunner#apply()
+ public StreamingViewAsSingleton(FlinkRunner runner, View.AsSingleton<T> transform) {
+ this.transform = transform;
+ }
+
+ @Override
+ public PCollectionView<T> expand(PCollection<T> input) {
+ Combine.Globally<T, T> combine = Combine.globally(
+ new SingletonCombine<>(transform.hasDefaultValue(), transform.defaultValue()));
+ if (!transform.hasDefaultValue()) {
+ combine = combine.withoutDefaults();
+ }
+ return input.apply(combine.asSingletonView());
+ }
+
+ @Override
+ protected String getKindString() {
+ return "StreamingViewAsSingleton";
+ }
+
+ private static class SingletonCombine<T> extends Combine.BinaryCombineFn<T> {
+ private boolean hasDefaultValue;
+ private T defaultValue;
+
+ SingletonCombine(boolean hasDefaultValue, T defaultValue) {
+ this.hasDefaultValue = hasDefaultValue;
+ this.defaultValue = defaultValue;
+ }
+
+ @Override
+ public T apply(T left, T right) {
+ throw new IllegalArgumentException("PCollection with more than one element "
+ + "accessed as a singleton view. Consider using Combine.globally().asSingleton() to "
+ + "combine the PCollection into a single value");
+ }
+
+ @Override
+ public T identity() {
+ if (hasDefaultValue) {
+ return defaultValue;
+ } else {
+ throw new IllegalArgumentException(
+ "Empty PCollection accessed as a singleton view. "
+ + "Consider setting withDefault to provide a default value");
+ }
+ }
+ }
+ }
+
+ static class StreamingCombineGloballyAsSingletonView<InputT, OutputT>
+ extends PTransform<PCollection<InputT>, PCollectionView<OutputT>> {
+ Combine.GloballyAsSingletonView<InputT, OutputT> transform;
+
+ /**
+ * Builds an instance of this class from the overridden transform.
+ */
+ @SuppressWarnings("unused") // used via reflection in FlinkRunner#apply()
+ public StreamingCombineGloballyAsSingletonView(
+ FlinkRunner runner,
+ Combine.GloballyAsSingletonView<InputT, OutputT> transform) {
+ this.transform = transform;
+ }
+
+ @Override
+ public PCollectionView<OutputT> expand(PCollection<InputT> input) {
+ PCollection<OutputT> combined =
+ input.apply(Combine.globally(transform.getCombineFn())
+ .withoutDefaults()
+ .withFanout(transform.getFanout()));
+
+ PCollectionView<OutputT> view = PCollectionViews.singletonView(
+ combined,
+ combined.getWindowingStrategy(),
+ transform.getInsertDefault(),
+ transform.getInsertDefault()
+ ? transform.getCombineFn().defaultValue() : null,
+ combined.getCoder());
+ return combined
+ .apply(ParDo.of(new WrapAsList<OutputT>()))
+ .apply(CreateFlinkPCollectionView.<OutputT, OutputT>of(view));
+ }
+
+ @Override
+ protected String getKindString() {
+ return "StreamingCombineGloballyAsSingletonView";
+ }
+ }
+
+ private static class WrapAsList<T> extends DoFn<T, List<T>> {
+ @ProcessElement
+ public void processElement(ProcessContext c) {
+ c.output(Collections.singletonList(c.element()));
+ }
+ }
+
+ /**
+ * Combiner that combines {@code T}s into a single {@code List<T>} containing all inputs.
+ *
+ * <p>For internal use by {@link StreamingViewAsMap}, {@link StreamingViewAsMultimap},
+ * {@link StreamingViewAsList}, {@link StreamingViewAsIterable}.
+ * They require the input {@link PCollection} fits in memory.
+ * For a large {@link PCollection} this is expected to crash!
+ *
+ * @param <T> the type of elements to concatenate.
+ */
+ private static class Concatenate<T> extends Combine.CombineFn<T, List<T>, List<T>> {
+ @Override
+ public List<T> createAccumulator() {
+ return new ArrayList<T>();
+ }
+
+ @Override
+ public List<T> addInput(List<T> accumulator, T input) {
+ accumulator.add(input);
+ return accumulator;
+ }
+
+ @Override
+ public List<T> mergeAccumulators(Iterable<List<T>> accumulators) {
+ List<T> result = createAccumulator();
+ for (List<T> accumulator : accumulators) {
+ result.addAll(accumulator);
+ }
+ return result;
+ }
+
+ @Override
+ public List<T> extractOutput(List<T> accumulator) {
+ return accumulator;
+ }
+
+ @Override
+ public Coder<List<T>> getAccumulatorCoder(CoderRegistry registry, Coder<T> inputCoder) {
+ return ListCoder.of(inputCoder);
+ }
+
+ @Override
+ public Coder<List<T>> getDefaultOutputCoder(CoderRegistry registry, Coder<T> inputCoder) {
+ return ListCoder.of(inputCoder);
+ }
+ }
+
+ /**
+ * Creates a primitive {@link PCollectionView}.
+ *
+ * <p>For internal use only by runner implementors.
+ *
+ * @param <ElemT> The type of the elements of the input PCollection
+ * @param <ViewT> The type associated with the {@link PCollectionView} used as a side input
+ */
+ public static class CreateFlinkPCollectionView<ElemT, ViewT>
+ extends PTransform<PCollection<List<ElemT>>, PCollectionView<ViewT>> {
+ private PCollectionView<ViewT> view;
+
+ private CreateFlinkPCollectionView(PCollectionView<ViewT> view) {
+ this.view = view;
+ }
+
+ public static <ElemT, ViewT> CreateFlinkPCollectionView<ElemT, ViewT> of(
+ PCollectionView<ViewT> view) {
+ return new CreateFlinkPCollectionView<>(view);
+ }
+
+ @Override
+ public PCollectionView<ViewT> expand(PCollection<List<ElemT>> input) {
+ return view;
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkTransformOverrides.java
----------------------------------------------------------------------
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkTransformOverrides.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkTransformOverrides.java
deleted file mode 100644
index 1dc8de9..0000000
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkTransformOverrides.java
+++ /dev/null
@@ -1,53 +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.flink;
-
-import com.google.common.collect.ImmutableList;
-import java.util.List;
-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.sdk.runners.PTransformOverride;
-import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.transforms.View;
-
-/**
- * {@link PTransform} overrides for Flink runner.
- */
-public class FlinkTransformOverrides {
- public static List<PTransformOverride> getDefaultOverrides(boolean streaming) {
- if (streaming) {
- return ImmutableList.<PTransformOverride>builder()
- .add(
- PTransformOverride.of(
- PTransformMatchers.splittableParDoMulti(),
- new FlinkStreamingPipelineTranslator.SplittableParDoOverrideFactory()))
- .add(
- PTransformOverride.of(
- PTransformMatchers.classEqualTo(SplittableParDo.ProcessKeyedElements.class),
- new SplittableParDoViaKeyedWorkItems.OverrideFactory()))
- .add(
- PTransformOverride.of(
- PTransformMatchers.classEqualTo(View.CreatePCollectionView.class),
- new CreateStreamingFlinkView.Factory()))
- .build();
- } else {
- return ImmutableList.of();
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/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 2f095d4..5d08eba 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
@@ -35,6 +35,7 @@ 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;
@@ -42,7 +43,6 @@ import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.util.WindowedValue;
-import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
@@ -55,15 +55,18 @@ import org.joda.time.Instant;
* the {@code @ProcessElement} method of a splittable {@link DoFn}.
*/
public class SplittableDoFnOperator<
- InputT, OutputT, RestrictionT, TrackerT extends RestrictionTracker<RestrictionT>>
- extends DoFnOperator<KeyedWorkItem<String, KV<InputT, RestrictionT>>, OutputT> {
+ InputT, OutputT, RestrictionT, TrackerT extends RestrictionTracker<RestrictionT>>
+ extends DoFnOperator<
+ KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT> {
private transient ScheduledExecutorService executorService;
public SplittableDoFnOperator(
- DoFn<KeyedWorkItem<String, KV<InputT, RestrictionT>>, OutputT> doFn,
+ DoFn<KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>, OutputT> doFn,
String stepName,
- Coder<WindowedValue<KeyedWorkItem<String, KV<InputT, RestrictionT>>>> inputCoder,
+ Coder<
+ WindowedValue<
+ KeyedWorkItem<String, ElementAndRestriction<InputT, RestrictionT>>>> inputCoder,
TupleTag<OutputT> mainOutputTag,
List<TupleTag<?>> additionalOutputTags,
OutputManagerFactory<OutputT> outputManagerFactory,
@@ -84,6 +87,7 @@ public class SplittableDoFnOperator<
sideInputs,
options,
keyCoder);
+
}
@Override
@@ -147,7 +151,7 @@ public class SplittableDoFnOperator<
@Override
public void fireTimer(InternalTimer<?, TimerInternals.TimerData> timer) {
doFnRunner.processElement(WindowedValue.valueInGlobalWindow(
- KeyedWorkItems.<String, KV<InputT, RestrictionT>>timersWorkItem(
+ KeyedWorkItems.<String, ElementAndRestriction<InputT, RestrictionT>>timersWorkItem(
(String) stateInternals.getKey(),
Collections.singletonList(timer.getNamespace()))));
}