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/04/12 19:56:32 UTC
[01/50] [abbrv] beam git commit: This closes #2487 [Forced Update!]
Repository: beam
Updated Branches:
refs/heads/DSL_SQL 870e42dec -> 3625dbd9e (forced update)
This closes #2487
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/8beea73c
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/8beea73c
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/8beea73c
Branch: refs/heads/DSL_SQL
Commit: 8beea73c1b57414dfb2b0a3f44f5bee29f088b4e
Parents: 986fcef 01e5a8d
Author: Thomas Groh <tg...@google.com>
Authored: Tue Apr 11 09:58:44 2017 -0700
Committer: Thomas Groh <tg...@google.com>
Committed: Tue Apr 11 09:58:44 2017 -0700
----------------------------------------------------------------------
.../beam/runners/core/construction/Coders.java | 162 +++++++++++++++++++
.../core/construction/SdkComponents.java | 5 +-
.../runners/core/construction/CodersTest.java | 99 ++++++++++++
.../core/construction/SdkComponentsTest.java | 25 ++-
4 files changed, 288 insertions(+), 3 deletions(-)
----------------------------------------------------------------------
[36/50] [abbrv] beam git commit: Use SdkComponents in
WindowingStrategy.toProto
Posted by jb...@apache.org.
Use SdkComponents in WindowingStrategy.toProto
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/e0fd0a22
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/e0fd0a22
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/e0fd0a22
Branch: refs/heads/DSL_SQL
Commit: e0fd0a222510b733b98d6c33c694431139bbc40d
Parents: 7e603d5
Author: Thomas Groh <tg...@google.com>
Authored: Fri Apr 7 13:41:29 2017 -0700
Committer: Dan Halperin <dh...@google.com>
Committed: Wed Apr 12 09:20:43 2017 -0700
----------------------------------------------------------------------
.../beam/runners/core/construction/Coders.java | 6 +-
.../core/construction/SdkComponents.java | 5 +-
.../core/construction/WindowingStrategies.java | 94 ++++++++------------
.../core/construction/SdkComponentsTest.java | 6 +-
.../construction/WindowingStrategiesTest.java | 18 ++++
5 files changed, 66 insertions(+), 63 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/e0fd0a22/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/Coders.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/Coders.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/Coders.java
index d890de7..7b96240 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/Coders.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/Coders.java
@@ -76,8 +76,10 @@ public class Coders {
private static RunnerApi.Coder toKnownCoder(Coder<?> coder, SdkComponents components)
throws IOException {
List<String> componentIds = new ArrayList<>();
- for (Coder<?> componentCoder : coder.getCoderArguments()) {
- componentIds.add(components.registerCoder(componentCoder));
+ if (coder.getCoderArguments() != null) {
+ for (Coder<?> componentCoder : coder.getCoderArguments()) {
+ componentIds.add(components.registerCoder(componentCoder));
+ }
}
return RunnerApi.Coder.newBuilder()
.addAllComponentCoderIds(componentIds)
http://git-wip-us.apache.org/repos/asf/beam/blob/e0fd0a22/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SdkComponents.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SdkComponents.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SdkComponents.java
index 5cb0a00..03f3a03 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SdkComponents.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SdkComponents.java
@@ -96,7 +96,7 @@ class SdkComponents {
* unique ID for the {@link WindowingStrategy}. Multiple registrations of the same {@link
* WindowingStrategy} will return the same unique ID.
*/
- String registerWindowingStrategy(WindowingStrategy<?, ?> windowingStrategy) {
+ String registerWindowingStrategy(WindowingStrategy<?, ?> windowingStrategy) throws IOException {
String existing = windowingStrategyIds.get(windowingStrategy);
if (existing != null) {
return existing;
@@ -108,6 +108,9 @@ class SdkComponents {
NameUtils.approximateSimpleName(windowingStrategy.getWindowFn()));
String name = uniqify(baseName, windowingStrategyIds.values());
windowingStrategyIds.put(windowingStrategy, name);
+ RunnerApi.WindowingStrategy windowingStrategyProto =
+ WindowingStrategies.toProto(windowingStrategy, this);
+ componentsBuilder.putWindowingStrategies(name, windowingStrategyProto);
return name;
}
http://git-wip-us.apache.org/repos/asf/beam/blob/e0fd0a22/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowingStrategies.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowingStrategies.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowingStrategies.java
index 353be05..6d721b0 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowingStrategies.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/WindowingStrategies.java
@@ -19,15 +19,12 @@ package org.apache.beam.runners.core.construction;
import static com.google.common.base.Preconditions.checkArgument;
-import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.protobuf.Any;
import com.google.protobuf.ByteString;
import com.google.protobuf.BytesValue;
import com.google.protobuf.InvalidProtocolBufferException;
import java.io.IOException;
import java.io.Serializable;
-import java.util.UUID;
-import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.common.runner.v1.RunnerApi;
import org.apache.beam.sdk.common.runner.v1.RunnerApi.Components;
import org.apache.beam.sdk.common.runner.v1.RunnerApi.FunctionSpec;
@@ -127,62 +124,30 @@ public class WindowingStrategies implements Serializable {
}
}
- // This URN says that the coder is just a UDF blob the indicated SDK understands
- // TODO: standardize such things
- public static final String CUSTOM_CODER_URN = "urn:beam:coders:javasdk:0.1";
-
// This URN says that the WindowFn is just a UDF blob the indicated SDK understands
// TODO: standardize such things
public static final String CUSTOM_WINDOWFN_URN = "urn:beam:windowfn:javasdk:0.1";
- private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
-
/**
- * Converts a {@link WindowFn} into a {@link RunnerApi.MessageWithComponents} where
- * {@link RunnerApi.MessageWithComponents#getFunctionSpec()} is a {@link FunctionSpec}
- * for the input {@link WindowFn}.
+ * Converts a {@link WindowFn} into a {@link RunnerApi.MessageWithComponents} where {@link
+ * RunnerApi.MessageWithComponents#getFunctionSpec()} is a {@link RunnerApi.FunctionSpec} for the
+ * input {@link WindowFn}.
*/
- public static RunnerApi.MessageWithComponents toProto(WindowFn<?, ?> windowFn)
+ public static SdkFunctionSpec toProto(
+ WindowFn<?, ?> windowFn, @SuppressWarnings("unused") SdkComponents components)
throws IOException {
- Coder<?> windowCoder = windowFn.windowCoder();
-
- // TODO: re-use components
- String windowCoderId = UUID.randomUUID().toString();
-
- SdkFunctionSpec windowFnSpec =
- SdkFunctionSpec.newBuilder()
- .setSpec(
- FunctionSpec.newBuilder()
- .setUrn(CUSTOM_WINDOWFN_URN)
- .setParameter(
- Any.pack(
- BytesValue.newBuilder()
- .setValue(
- ByteString.copyFrom(
- SerializableUtils.serializeToByteArray(windowFn)))
- .build())))
- .build();
-
- RunnerApi.Coder windowCoderProto =
- RunnerApi.Coder.newBuilder()
- .setSpec(
- SdkFunctionSpec.newBuilder()
- .setSpec(
- FunctionSpec.newBuilder()
- .setUrn(CUSTOM_CODER_URN)
- .setParameter(
- Any.pack(
- BytesValue.newBuilder()
- .setValue(
- ByteString.copyFrom(
- OBJECT_MAPPER.writeValueAsBytes(
- windowCoder.asCloudObject())))
- .build()))))
- .build();
-
- return RunnerApi.MessageWithComponents.newBuilder()
- .setSdkFunctionSpec(windowFnSpec)
- .setComponents(Components.newBuilder().putCoders(windowCoderId, windowCoderProto))
+ return SdkFunctionSpec.newBuilder()
+ // TODO: Set environment ID
+ .setSpec(
+ FunctionSpec.newBuilder()
+ .setUrn(CUSTOM_WINDOWFN_URN)
+ .setParameter(
+ Any.pack(
+ BytesValue.newBuilder()
+ .setValue(
+ ByteString.copyFrom(
+ SerializableUtils.serializeToByteArray(windowFn)))
+ .build())))
.build();
}
@@ -194,9 +159,22 @@ public class WindowingStrategies implements Serializable {
*/
public static RunnerApi.MessageWithComponents toProto(WindowingStrategy<?, ?> windowingStrategy)
throws IOException {
+ SdkComponents components = SdkComponents.create();
+ RunnerApi.WindowingStrategy windowingStrategyProto = toProto(windowingStrategy, components);
- RunnerApi.MessageWithComponents windowFnWithComponents =
- toProto(windowingStrategy.getWindowFn());
+ return RunnerApi.MessageWithComponents.newBuilder()
+ .setWindowingStrategy(windowingStrategyProto)
+ .setComponents(components.toComponents())
+ .build();
+ }
+
+ /**
+ * Converts a {@link WindowingStrategy} into a {@link RunnerApi.WindowingStrategy}, registering
+ * any components in the provided {@link SdkComponents}.
+ */
+ public static RunnerApi.WindowingStrategy toProto(
+ WindowingStrategy<?, ?> windowingStrategy, SdkComponents components) throws IOException {
+ SdkFunctionSpec windowFnSpec = toProto(windowingStrategy.getWindowFn(), components);
RunnerApi.WindowingStrategy.Builder windowingStrategyProto =
RunnerApi.WindowingStrategy.newBuilder()
@@ -205,11 +183,11 @@ public class WindowingStrategies implements Serializable {
.setClosingBehavior(toProto(windowingStrategy.getClosingBehavior()))
.setAllowedLateness(windowingStrategy.getAllowedLateness().getMillis())
.setTrigger(Triggers.toProto(windowingStrategy.getTrigger()))
- .setWindowFn(windowFnWithComponents.getSdkFunctionSpec());
+ .setWindowFn(windowFnSpec)
+ .setWindowCoderId(
+ components.registerCoder(windowingStrategy.getWindowFn().windowCoder()));
- return RunnerApi.MessageWithComponents.newBuilder()
- .setWindowingStrategy(windowingStrategyProto)
- .setComponents(windowFnWithComponents.getComponents()).build();
+ return windowingStrategyProto.build();
}
/**
http://git-wip-us.apache.org/repos/asf/beam/blob/e0fd0a22/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SdkComponentsTest.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SdkComponentsTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SdkComponentsTest.java
index 28b4911..ef4b16b 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SdkComponentsTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SdkComponentsTest.java
@@ -132,15 +132,17 @@ public class SdkComponentsTest {
}
@Test
- public void registerWindowingStrategy() {
+ public void registerWindowingStrategy() throws IOException {
WindowingStrategy<?, ?> strategy =
WindowingStrategy.globalDefault().withMode(AccumulationMode.ACCUMULATING_FIRED_PANES);
String name = components.registerWindowingStrategy(strategy);
assertThat(name, not(isEmptyOrNullString()));
+
+ components.toComponents().getWindowingStrategiesOrThrow(name);
}
@Test
- public void registerWindowingStrategyIdEqualStrategies() {
+ public void registerWindowingStrategyIdEqualStrategies() throws IOException {
WindowingStrategy<?, ?> strategy =
WindowingStrategy.globalDefault().withMode(AccumulationMode.ACCUMULATING_FIRED_PANES);
String name = components.registerWindowingStrategy(strategy);
http://git-wip-us.apache.org/repos/asf/beam/blob/e0fd0a22/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WindowingStrategiesTest.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WindowingStrategiesTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WindowingStrategiesTest.java
index b603d65..62bba8e 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WindowingStrategiesTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/WindowingStrategiesTest.java
@@ -22,6 +22,7 @@ import static org.junit.Assert.assertThat;
import com.google.auto.value.AutoValue;
import com.google.common.collect.ImmutableList;
+import org.apache.beam.sdk.common.runner.v1.RunnerApi;
import org.apache.beam.sdk.transforms.windowing.AfterWatermark;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
import org.apache.beam.sdk.transforms.windowing.OutputTimeFns;
@@ -30,6 +31,7 @@ import org.apache.beam.sdk.transforms.windowing.Window.ClosingBehavior;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.util.WindowingStrategy;
import org.apache.beam.sdk.util.WindowingStrategy.AccumulationMode;
+import org.hamcrest.Matchers;
import org.joda.time.Duration;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -89,4 +91,20 @@ public class WindowingStrategiesTest {
toProtoAndBackWindowingStrategy,
equalTo((WindowingStrategy) windowingStrategy.fixDefaults()));
}
+
+ @Test
+ public void testToProtoAndBackWithComponents() throws Exception {
+ WindowingStrategy<?, ?> windowingStrategy = toProtoAndBackSpec.getWindowingStrategy();
+ SdkComponents components = SdkComponents.create();
+ RunnerApi.WindowingStrategy proto =
+ WindowingStrategies.toProto(windowingStrategy, components);
+ RunnerApi.Components protoComponents = components.toComponents();
+
+ assertThat(
+ WindowingStrategies.fromProto(proto, protoComponents).fixDefaults(),
+ Matchers.<WindowingStrategy<?, ?>>equalTo(windowingStrategy.fixDefaults()));
+
+ protoComponents.getCodersOrThrow(
+ components.registerCoder(windowingStrategy.getWindowFn().windowCoder()));
+ }
}
[20/50] [abbrv] beam git commit: Throw specialized exception in value
providers
Posted by jb...@apache.org.
Throw specialized exception in value providers
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/ee92b964
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/ee92b964
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/ee92b964
Branch: refs/heads/DSL_SQL
Commit: ee92b9642bb6b6e42bb701ab638c55539163bb69
Parents: 17a41ab
Author: Sourabh Bajaj <so...@google.com>
Authored: Tue Apr 11 13:38:26 2017 -0700
Committer: Ahmet Altay <al...@google.com>
Committed: Tue Apr 11 15:32:14 2017 -0700
----------------------------------------------------------------------
.../apache_beam/runners/dataflow/dataflow_runner.py | 6 ++++++
sdks/python/apache_beam/utils/value_provider.py | 14 ++++++++++++--
2 files changed, 18 insertions(+), 2 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/ee92b964/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 1a935c1..4e81788 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
@@ -43,6 +43,7 @@ from apache_beam.runners.runner import PipelineState
from apache_beam.transforms.display import DisplayData
from apache_beam.typehints import typehints
from apache_beam.utils.pipeline_options import StandardOptions
+from apache_beam.utils.value_provider import RuntimeValueProviderError
class DataflowRunner(PipelineRunner):
@@ -476,6 +477,11 @@ class DataflowRunner(PipelineRunner):
'estimated_size_bytes': json_value.get_typed_value_descriptor(
transform.source.estimate_size())
}
+ except RuntimeValueProviderError:
+ # Size estimation is best effort, and this error is by value provider.
+ logging.info(
+ 'Could not estimate size of source %r due to ' + \
+ 'RuntimeValueProviderError', transform.source)
except Exception: # pylint: disable=broad-except
# Size estimation is best effort. So we log the error and continue.
logging.info(
http://git-wip-us.apache.org/repos/asf/beam/blob/ee92b964/sdks/python/apache_beam/utils/value_provider.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/utils/value_provider.py b/sdks/python/apache_beam/utils/value_provider.py
index a72fc4c..271202d 100644
--- a/sdks/python/apache_beam/utils/value_provider.py
+++ b/sdks/python/apache_beam/utils/value_provider.py
@@ -22,6 +22,15 @@ and dynamically provided values.
from functools import wraps
+class RuntimeValueProviderError(RuntimeError):
+ def __init__(self, msg):
+ """Class representing the errors thrown during runtime by the valueprovider
+ Args:
+ msg: Message string for the exception thrown
+ """
+ super(RuntimeValueProviderError, self).__init__(msg)
+
+
class ValueProvider(object):
def is_accessible(self):
raise NotImplementedError(
@@ -67,7 +76,8 @@ class RuntimeValueProvider(ValueProvider):
runtime_options = (
RuntimeValueProvider.runtime_options_map.get(self.options_id))
if runtime_options is None:
- raise RuntimeError('%s.get() not called from a runtime context' % self)
+ raise RuntimeValueProviderError(
+ '%s.get() not called from a runtime context' % self)
candidate = runtime_options.get(self.option_name)
if candidate:
@@ -104,7 +114,7 @@ def check_accessible(value_provider_list):
def _f(self, *args, **kwargs):
for obj in [getattr(self, vp) for vp in value_provider_list]:
if not obj.is_accessible():
- raise RuntimeError('%s not accessible' % obj)
+ raise RuntimeValueProviderError('%s not accessible' % obj)
return fnc(self, *args, **kwargs)
return _f
return _check_accessible
[43/50] [abbrv] beam git commit: Removes
FlatMapElements.MissingOutputTypeDescriptor
Posted by jb...@apache.org.
Removes FlatMapElements.MissingOutputTypeDescriptor
This comes from changing FlatMapElements.via(fn).withOutputType(td)
to FlatMapElements.into(td).via(fn) which is also shorter.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/831b11fb
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/831b11fb
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/831b11fb
Branch: refs/heads/DSL_SQL
Commit: 831b11fb712049ce78965ec2e677f9cf9ea66fc4
Parents: 3bf6c6b
Author: Eugene Kirpichov <ki...@google.com>
Authored: Wed Mar 29 13:56:00 2017 -0700
Committer: Eugene Kirpichov <ki...@google.com>
Committed: Wed Apr 12 11:35:05 2017 -0700
----------------------------------------------------------------------
.../beam/examples/MinimalWordCountJava8.java | 5 +-
.../examples/MinimalWordCountJava8Test.java | 5 +-
.../beam/sdk/transforms/FlatMapElements.java | 113 +++++++++----------
.../transforms/FlatMapElementsJava8Test.java | 10 +-
4 files changed, 63 insertions(+), 70 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/831b11fb/examples/java8/src/main/java/org/apache/beam/examples/MinimalWordCountJava8.java
----------------------------------------------------------------------
diff --git a/examples/java8/src/main/java/org/apache/beam/examples/MinimalWordCountJava8.java b/examples/java8/src/main/java/org/apache/beam/examples/MinimalWordCountJava8.java
index 0072886..f424a7b 100644
--- a/examples/java8/src/main/java/org/apache/beam/examples/MinimalWordCountJava8.java
+++ b/examples/java8/src/main/java/org/apache/beam/examples/MinimalWordCountJava8.java
@@ -56,8 +56,9 @@ public class MinimalWordCountJava8 {
Pipeline p = Pipeline.create(options);
p.apply(TextIO.Read.from("gs://apache-beam-samples/shakespeare/*"))
- .apply(FlatMapElements.via((String word) -> Arrays.asList(word.split("[^a-zA-Z']+")))
- .withOutputType(TypeDescriptors.strings()))
+ .apply(FlatMapElements
+ .into(TypeDescriptors.strings())
+ .via((String word) -> Arrays.asList(word.split("[^a-zA-Z']+"))))
.apply(Filter.by((String word) -> !word.isEmpty()))
.apply(Count.<String>perElement())
.apply(MapElements
http://git-wip-us.apache.org/repos/asf/beam/blob/831b11fb/examples/java8/src/test/java/org/apache/beam/examples/MinimalWordCountJava8Test.java
----------------------------------------------------------------------
diff --git a/examples/java8/src/test/java/org/apache/beam/examples/MinimalWordCountJava8Test.java b/examples/java8/src/test/java/org/apache/beam/examples/MinimalWordCountJava8Test.java
index 911ccf6..6c66d8f 100644
--- a/examples/java8/src/test/java/org/apache/beam/examples/MinimalWordCountJava8Test.java
+++ b/examples/java8/src/test/java/org/apache/beam/examples/MinimalWordCountJava8Test.java
@@ -63,8 +63,9 @@ public class MinimalWordCountJava8Test implements Serializable {
p.getOptions().as(GcsOptions.class).setGcsUtil(buildMockGcsUtil());
p.apply(TextIO.Read.from("gs://apache-beam-samples/shakespeare/*"))
- .apply(FlatMapElements.via((String word) -> Arrays.asList(word.split("[^a-zA-Z']+")))
- .withOutputType(TypeDescriptors.strings()))
+ .apply(FlatMapElements
+ .into(TypeDescriptors.strings())
+ .via((String word) -> Arrays.asList(word.split("[^a-zA-Z']+"))))
.apply(Filter.by((String word) -> !word.isEmpty()))
.apply(Count.<String>perElement())
.apply(MapElements
http://git-wip-us.apache.org/repos/asf/beam/blob/831b11fb/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/FlatMapElements.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/FlatMapElements.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/FlatMapElements.java
index c165f7f..0983165 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/FlatMapElements.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/FlatMapElements.java
@@ -17,7 +17,10 @@
*/
package org.apache.beam.sdk.transforms;
+import static com.google.common.base.Preconditions.checkNotNull;
+
import java.lang.reflect.ParameterizedType;
+import javax.annotation.Nullable;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.TypeDescriptor;
@@ -30,32 +33,29 @@ import org.apache.beam.sdk.values.TypeDescriptors;
public class FlatMapElements<InputT, OutputT>
extends PTransform<PCollection<? extends InputT>, PCollection<OutputT>> {
/**
- * For a {@code SerializableFunction<InputT, ? extends Iterable<OutputT>>} {@code fn},
- * returns a {@link PTransform} that applies {@code fn} to every element of the input
- * {@code PCollection<InputT>} and outputs all of the elements to the output
- * {@code PCollection<OutputT>}.
- *
- * <p>Example of use in Java 8:
- * <pre>{@code
- * PCollection<String> words = lines.apply(
- * FlatMapElements.via((String line) -> Arrays.asList(line.split(" ")))
- * .withOutputType(new TypeDescriptor<String>(){});
- * }</pre>
- *
- * <p>In Java 7, the overload {@link #via(SimpleFunction)} is more concise as the output type
- * descriptor need not be provided.
+ * Temporarily stores the argument of {@link #into(TypeDescriptor)} until combined with the
+ * argument of {@link #via(SerializableFunction)} into the fully-specified {@link #fn}. Stays null
+ * if constructed using {@link #via(SimpleFunction)} directly.
+ */
+ @Nullable
+ private final transient TypeDescriptor<Iterable<OutputT>> outputType;
+
+ /**
+ * Non-null on a fully specified transform - is null only when constructed using {@link
+ * #into(TypeDescriptor)}, until the fn is specified using {@link #via(SerializableFunction)}.
*/
- public static <InputT, OutputT> MissingOutputTypeDescriptor<InputT, OutputT>
- via(SerializableFunction<? super InputT, ? extends Iterable<OutputT>> fn) {
+ @Nullable
+ private final SimpleFunction<InputT, Iterable<OutputT>> fn;
+ private final DisplayData.ItemSpec<?> fnClassDisplayData;
- // TypeDescriptor interacts poorly with the wildcards needed to correctly express
- // covariance and contravariance in Java, so instead we cast it to an invariant
- // function here.
- @SuppressWarnings("unchecked") // safe covariant cast
- SerializableFunction<InputT, Iterable<OutputT>> simplerFn =
- (SerializableFunction<InputT, Iterable<OutputT>>) fn;
+ private FlatMapElements(
+ @Nullable SimpleFunction<InputT, Iterable<OutputT>> fn,
+ @Nullable TypeDescriptor<Iterable<OutputT>> outputType,
+ @Nullable Class<?> fnClass) {
+ this.fn = fn;
+ this.outputType = outputType;
+ this.fnClassDisplayData = DisplayData.item("flatMapFn", fnClass).withLabel("FlatMap Function");
- return new MissingOutputTypeDescriptor<>(simplerFn);
}
/**
@@ -82,54 +82,45 @@ extends PTransform<PCollection<? extends InputT>, PCollection<OutputT>> {
*/
public static <InputT, OutputT> FlatMapElements<InputT, OutputT>
via(SimpleFunction<? super InputT, ? extends Iterable<OutputT>> fn) {
- // TypeDescriptor interacts poorly with the wildcards needed to correctly express
- // covariance and contravariance in Java, so instead we cast it to an invariant
- // function here.
- @SuppressWarnings("unchecked") // safe covariant cast
- SimpleFunction<InputT, Iterable<OutputT>> simplerFn =
- (SimpleFunction<InputT, Iterable<OutputT>>) fn;
-
- return new FlatMapElements<>(simplerFn, fn.getClass());
+ return new FlatMapElements(fn, null, fn.getClass());
}
/**
- * An intermediate builder for a {@link FlatMapElements} transform. To complete the transform,
- * provide an output type descriptor to {@link MissingOutputTypeDescriptor#withOutputType}. See
- * {@link #via(SerializableFunction)} for a full example of use.
+ * Returns a new {@link FlatMapElements} transform with the given type descriptor for the output
+ * type, but the mapping function yet to be specified using {@link #via(SerializableFunction)}.
*/
- public static final class MissingOutputTypeDescriptor<InputT, OutputT> {
-
- private final SerializableFunction<InputT, Iterable<OutputT>> fn;
-
- private MissingOutputTypeDescriptor(
- SerializableFunction<InputT, Iterable<OutputT>> fn) {
- this.fn = fn;
- }
-
- public FlatMapElements<InputT, OutputT> withOutputType(TypeDescriptor<OutputT> outputType) {
- TypeDescriptor<Iterable<OutputT>> iterableOutputType = TypeDescriptors.iterables(outputType);
-
- return new FlatMapElements<>(
- SimpleFunction.fromSerializableFunctionWithOutputType(fn,
- iterableOutputType),
- fn.getClass());
- }
+ public static <OutputT> FlatMapElements<?, OutputT>
+ into(final TypeDescriptor<OutputT> outputType) {
+ return new FlatMapElements<>(null, TypeDescriptors.iterables(outputType), null);
}
- //////////////////////////////////////////////////////////////////////////////////////////////////
-
- private final SimpleFunction<InputT, ? extends Iterable<OutputT>> fn;
- private final DisplayData.ItemSpec<?> fnClassDisplayData;
-
- private FlatMapElements(
- SimpleFunction<InputT, ? extends Iterable<OutputT>> fn,
- Class<?> fnClass) {
- this.fn = fn;
- this.fnClassDisplayData = DisplayData.item("flatMapFn", fnClass).withLabel("FlatMap Function");
+ /**
+ * For a {@code SerializableFunction<InputT, ? extends Iterable<OutputT>>} {@code fn},
+ * returns a {@link PTransform} that applies {@code fn} to every element of the input
+ * {@code PCollection<InputT>} and outputs all of the elements to the output
+ * {@code PCollection<OutputT>}.
+ *
+ * <p>Example of use in Java 8:
+ * <pre>{@code
+ * PCollection<String> words = lines.apply(
+ * FlatMapElements.via((String line) -> Arrays.asList(line.split(" ")))
+ * .withOutputType(new TypeDescriptor<String>(){});
+ * }</pre>
+ *
+ * <p>In Java 7, the overload {@link #via(SimpleFunction)} is more concise as the output type
+ * descriptor need not be provided.
+ */
+ public <NewInputT> FlatMapElements<NewInputT, OutputT>
+ via(SerializableFunction<NewInputT, ? extends Iterable<OutputT>> fn) {
+ return new FlatMapElements(
+ SimpleFunction.fromSerializableFunctionWithOutputType(fn, (TypeDescriptor) outputType),
+ null,
+ fn.getClass());
}
@Override
public PCollection<OutputT> expand(PCollection<? extends InputT> input) {
+ checkNotNull(fn, "Must specify a function on FlatMapElements using .via()");
return input.apply(
"FlatMap",
ParDo.of(
http://git-wip-us.apache.org/repos/asf/beam/blob/831b11fb/sdks/java/java8tests/src/test/java/org/apache/beam/sdk/transforms/FlatMapElementsJava8Test.java
----------------------------------------------------------------------
diff --git a/sdks/java/java8tests/src/test/java/org/apache/beam/sdk/transforms/FlatMapElementsJava8Test.java b/sdks/java/java8tests/src/test/java/org/apache/beam/sdk/transforms/FlatMapElementsJava8Test.java
index 471724d..501b0d1 100644
--- a/sdks/java/java8tests/src/test/java/org/apache/beam/sdk/transforms/FlatMapElementsJava8Test.java
+++ b/sdks/java/java8tests/src/test/java/org/apache/beam/sdk/transforms/FlatMapElementsJava8Test.java
@@ -23,7 +23,7 @@ import java.util.List;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.sdk.values.TypeDescriptors;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
@@ -52,8 +52,8 @@ public class FlatMapElementsJava8Test implements Serializable {
.apply(Create.of(1, 2, 3))
.apply(FlatMapElements
// Note that the input type annotation is required.
- .via((Integer i) -> ImmutableList.of(i, -i))
- .withOutputType(new TypeDescriptor<Integer>() {}));
+ .into(TypeDescriptors.integers())
+ .via((Integer i) -> ImmutableList.of(i, -i)));
PAssert.that(output).containsInAnyOrder(1, 3, -1, -3, 2, -2);
pipeline.run();
@@ -69,8 +69,8 @@ public class FlatMapElementsJava8Test implements Serializable {
.apply(Create.of(1, 2, 3))
.apply(FlatMapElements
// Note that the input type annotation is required.
- .via(new Negater()::numAndNegation)
- .withOutputType(new TypeDescriptor<Integer>() {}));
+ .into(TypeDescriptors.integers())
+ .via(new Negater()::numAndNegation));
PAssert.that(output).containsInAnyOrder(1, 3, -1, -3, 2, -2);
pipeline.run();
[32/50] [abbrv] beam git commit: [BEAM-1722] Move PubsubIO into the
google-cloud-platform module
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/src/main/java/org/apache/beam/sdk/util/PubsubGrpcClient.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/PubsubGrpcClient.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/PubsubGrpcClient.java
deleted file mode 100644
index 4a6ddac..0000000
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/PubsubGrpcClient.java
+++ /dev/null
@@ -1,424 +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.util;
-
-import static com.google.common.base.Preconditions.checkState;
-
-import com.google.auth.Credentials;
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Strings;
-import com.google.common.collect.ImmutableList;
-import com.google.protobuf.ByteString;
-import com.google.protobuf.Timestamp;
-import com.google.pubsub.v1.AcknowledgeRequest;
-import com.google.pubsub.v1.DeleteSubscriptionRequest;
-import com.google.pubsub.v1.DeleteTopicRequest;
-import com.google.pubsub.v1.GetSubscriptionRequest;
-import com.google.pubsub.v1.ListSubscriptionsRequest;
-import com.google.pubsub.v1.ListSubscriptionsResponse;
-import com.google.pubsub.v1.ListTopicsRequest;
-import com.google.pubsub.v1.ListTopicsResponse;
-import com.google.pubsub.v1.ModifyAckDeadlineRequest;
-import com.google.pubsub.v1.PublishRequest;
-import com.google.pubsub.v1.PublishResponse;
-import com.google.pubsub.v1.PublisherGrpc;
-import com.google.pubsub.v1.PublisherGrpc.PublisherBlockingStub;
-import com.google.pubsub.v1.PubsubMessage;
-import com.google.pubsub.v1.PullRequest;
-import com.google.pubsub.v1.PullResponse;
-import com.google.pubsub.v1.ReceivedMessage;
-import com.google.pubsub.v1.SubscriberGrpc;
-import com.google.pubsub.v1.SubscriberGrpc.SubscriberBlockingStub;
-import com.google.pubsub.v1.Subscription;
-import com.google.pubsub.v1.Topic;
-import io.grpc.Channel;
-import io.grpc.ClientInterceptors;
-import io.grpc.ManagedChannel;
-import io.grpc.auth.ClientAuthInterceptor;
-import io.grpc.netty.GrpcSslContexts;
-import io.grpc.netty.NegotiationType;
-import io.grpc.netty.NettyChannelBuilder;
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-import java.util.concurrent.Executors;
-import java.util.concurrent.TimeUnit;
-import javax.annotation.Nullable;
-import org.apache.beam.sdk.options.GcpOptions;
-import org.apache.beam.sdk.options.PubsubOptions;
-
-/**
- * A helper class for talking to Pubsub via grpc.
- *
- * <p>CAUTION: Currently uses the application default credentials and does not respect any
- * credentials-related arguments in {@link GcpOptions}.
- */
-public class PubsubGrpcClient extends PubsubClient {
- private static final String PUBSUB_ADDRESS = "pubsub.googleapis.com";
- private static final int PUBSUB_PORT = 443;
- private static final int LIST_BATCH_SIZE = 1000;
-
- private static final int DEFAULT_TIMEOUT_S = 15;
-
- private static class PubsubGrpcClientFactory implements PubsubClientFactory {
- @Override
- public PubsubClient newClient(
- @Nullable String timestampLabel, @Nullable String idLabel, PubsubOptions options)
- throws IOException {
- ManagedChannel channel = NettyChannelBuilder
- .forAddress(PUBSUB_ADDRESS, PUBSUB_PORT)
- .negotiationType(NegotiationType.TLS)
- .sslContext(GrpcSslContexts.forClient().ciphers(null).build())
- .build();
-
- return new PubsubGrpcClient(timestampLabel,
- idLabel,
- DEFAULT_TIMEOUT_S,
- channel,
- options.getGcpCredential());
- }
-
- @Override
- public String getKind() {
- return "Grpc";
- }
- }
-
- /**
- * Factory for creating Pubsub clients using gRCP transport.
- */
- public static final PubsubClientFactory FACTORY = new PubsubGrpcClientFactory();
-
- /**
- * Timeout for grpc calls (in s).
- */
- private final int timeoutSec;
-
- /**
- * Underlying netty channel, or {@literal null} if closed.
- */
- @Nullable
- private ManagedChannel publisherChannel;
-
- /**
- * Credentials determined from options and environment.
- */
- private final Credentials credentials;
-
- /**
- * Label to use for custom timestamps, or {@literal null} if should use Pubsub publish time
- * instead.
- */
- @Nullable
- private final String timestampLabel;
-
- /**
- * Label to use for custom ids, or {@literal null} if should use Pubsub provided ids.
- */
- @Nullable
- private final String idLabel;
-
-
- /**
- * Cached stubs, or null if not cached.
- */
- @Nullable
- private PublisherGrpc.PublisherBlockingStub cachedPublisherStub;
- private SubscriberGrpc.SubscriberBlockingStub cachedSubscriberStub;
-
- @VisibleForTesting
- PubsubGrpcClient(
- @Nullable String timestampLabel,
- @Nullable String idLabel,
- int timeoutSec,
- ManagedChannel publisherChannel,
- Credentials credentials) {
- this.timestampLabel = timestampLabel;
- this.idLabel = idLabel;
- this.timeoutSec = timeoutSec;
- this.publisherChannel = publisherChannel;
- this.credentials = credentials;
- }
-
- /**
- * Gracefully close the underlying netty channel.
- */
- @Override
- public void close() {
- if (publisherChannel == null) {
- // Already closed.
- return;
- }
- // Can gc the underlying stubs.
- cachedPublisherStub = null;
- cachedSubscriberStub = null;
- // Mark the client as having been closed before going further
- // in case we have an exception from the channel.
- ManagedChannel publisherChannel = this.publisherChannel;
- this.publisherChannel = null;
- // Gracefully shutdown the channel.
- publisherChannel.shutdown();
- try {
- publisherChannel.awaitTermination(timeoutSec, TimeUnit.SECONDS);
- } catch (InterruptedException e) {
- // Ignore.
- Thread.currentThread().interrupt();
- }
- }
-
- /**
- * Return channel with interceptor for returning credentials.
- */
- private Channel newChannel() throws IOException {
- checkState(publisherChannel != null, "PubsubGrpcClient has been closed");
- ClientAuthInterceptor interceptor =
- new ClientAuthInterceptor(credentials, Executors.newSingleThreadExecutor());
- return ClientInterceptors.intercept(publisherChannel, interceptor);
- }
-
- /**
- * Return a stub for making a publish request with a timeout.
- */
- private PublisherBlockingStub publisherStub() throws IOException {
- if (cachedPublisherStub == null) {
- cachedPublisherStub = PublisherGrpc.newBlockingStub(newChannel());
- }
- return cachedPublisherStub.withDeadlineAfter(timeoutSec, TimeUnit.SECONDS);
- }
-
- /**
- * Return a stub for making a subscribe request with a timeout.
- */
- private SubscriberBlockingStub subscriberStub() throws IOException {
- if (cachedSubscriberStub == null) {
- cachedSubscriberStub = SubscriberGrpc.newBlockingStub(newChannel());
- }
- return cachedSubscriberStub.withDeadlineAfter(timeoutSec, TimeUnit.SECONDS);
- }
-
- @Override
- public int publish(TopicPath topic, List<OutgoingMessage> outgoingMessages)
- throws IOException {
- PublishRequest.Builder request = PublishRequest.newBuilder()
- .setTopic(topic.getPath());
- for (OutgoingMessage outgoingMessage : outgoingMessages) {
- PubsubMessage.Builder message =
- PubsubMessage.newBuilder()
- .setData(ByteString.copyFrom(outgoingMessage.elementBytes));
-
- if (outgoingMessage.attributes != null) {
- message.putAllAttributes(outgoingMessage.attributes);
- }
-
- if (timestampLabel != null) {
- message.getMutableAttributes()
- .put(timestampLabel, String.valueOf(outgoingMessage.timestampMsSinceEpoch));
- }
-
- if (idLabel != null && !Strings.isNullOrEmpty(outgoingMessage.recordId)) {
- message.getMutableAttributes().put(idLabel, outgoingMessage.recordId);
- }
-
- request.addMessages(message);
- }
-
- PublishResponse response = publisherStub().publish(request.build());
- return response.getMessageIdsCount();
- }
-
- @Override
- public List<IncomingMessage> pull(
- long requestTimeMsSinceEpoch,
- SubscriptionPath subscription,
- int batchSize,
- boolean returnImmediately) throws IOException {
- PullRequest request = PullRequest.newBuilder()
- .setSubscription(subscription.getPath())
- .setReturnImmediately(returnImmediately)
- .setMaxMessages(batchSize)
- .build();
- PullResponse response = subscriberStub().pull(request);
- if (response.getReceivedMessagesCount() == 0) {
- return ImmutableList.of();
- }
- List<IncomingMessage> incomingMessages = new ArrayList<>(response.getReceivedMessagesCount());
- for (ReceivedMessage message : response.getReceivedMessagesList()) {
- PubsubMessage pubsubMessage = message.getMessage();
- @Nullable Map<String, String> attributes = pubsubMessage.getAttributes();
-
- // Payload.
- byte[] elementBytes = pubsubMessage.getData().toByteArray();
-
- // Timestamp.
- String pubsubTimestampString = null;
- Timestamp timestampProto = pubsubMessage.getPublishTime();
- if (timestampProto != null) {
- pubsubTimestampString = String.valueOf(timestampProto.getSeconds()
- + timestampProto.getNanos() / 1000L);
- }
- long timestampMsSinceEpoch =
- extractTimestamp(timestampLabel, pubsubTimestampString, attributes);
-
- // Ack id.
- String ackId = message.getAckId();
- checkState(!Strings.isNullOrEmpty(ackId));
-
- // Record id, if any.
- @Nullable String recordId = null;
- if (idLabel != null && attributes != null) {
- recordId = attributes.get(idLabel);
- }
- if (Strings.isNullOrEmpty(recordId)) {
- // Fall back to the Pubsub provided message id.
- recordId = pubsubMessage.getMessageId();
- }
-
- incomingMessages.add(new IncomingMessage(elementBytes, attributes, timestampMsSinceEpoch,
- requestTimeMsSinceEpoch, ackId, recordId));
- }
- return incomingMessages;
- }
-
- @Override
- public void acknowledge(SubscriptionPath subscription, List<String> ackIds)
- throws IOException {
- AcknowledgeRequest request = AcknowledgeRequest.newBuilder()
- .setSubscription(subscription.getPath())
- .addAllAckIds(ackIds)
- .build();
- subscriberStub().acknowledge(request); // ignore Empty result.
- }
-
- @Override
- public void modifyAckDeadline(
- SubscriptionPath subscription, List<String> ackIds, int deadlineSeconds)
- throws IOException {
- ModifyAckDeadlineRequest request =
- ModifyAckDeadlineRequest.newBuilder()
- .setSubscription(subscription.getPath())
- .addAllAckIds(ackIds)
- .setAckDeadlineSeconds(deadlineSeconds)
- .build();
- subscriberStub().modifyAckDeadline(request); // ignore Empty result.
- }
-
- @Override
- public void createTopic(TopicPath topic) throws IOException {
- Topic request = Topic.newBuilder()
- .setName(topic.getPath())
- .build();
- publisherStub().createTopic(request); // ignore Topic result.
- }
-
- @Override
- public void deleteTopic(TopicPath topic) throws IOException {
- DeleteTopicRequest request = DeleteTopicRequest.newBuilder()
- .setTopic(topic.getPath())
- .build();
- publisherStub().deleteTopic(request); // ignore Empty result.
- }
-
- @Override
- public List<TopicPath> listTopics(ProjectPath project) throws IOException {
- ListTopicsRequest.Builder request =
- ListTopicsRequest.newBuilder()
- .setProject(project.getPath())
- .setPageSize(LIST_BATCH_SIZE);
- ListTopicsResponse response = publisherStub().listTopics(request.build());
- if (response.getTopicsCount() == 0) {
- return ImmutableList.of();
- }
- List<TopicPath> topics = new ArrayList<>(response.getTopicsCount());
- while (true) {
- for (Topic topic : response.getTopicsList()) {
- topics.add(topicPathFromPath(topic.getName()));
- }
- if (response.getNextPageToken().isEmpty()) {
- break;
- }
- request.setPageToken(response.getNextPageToken());
- response = publisherStub().listTopics(request.build());
- }
- return topics;
- }
-
- @Override
- public void createSubscription(
- TopicPath topic, SubscriptionPath subscription,
- int ackDeadlineSeconds) throws IOException {
- Subscription request = Subscription.newBuilder()
- .setTopic(topic.getPath())
- .setName(subscription.getPath())
- .setAckDeadlineSeconds(ackDeadlineSeconds)
- .build();
- subscriberStub().createSubscription(request); // ignore Subscription result.
- }
-
- @Override
- public void deleteSubscription(SubscriptionPath subscription) throws IOException {
- DeleteSubscriptionRequest request =
- DeleteSubscriptionRequest.newBuilder()
- .setSubscription(subscription.getPath())
- .build();
- subscriberStub().deleteSubscription(request); // ignore Empty result.
- }
-
- @Override
- public List<SubscriptionPath> listSubscriptions(ProjectPath project, TopicPath topic)
- throws IOException {
- ListSubscriptionsRequest.Builder request =
- ListSubscriptionsRequest.newBuilder()
- .setProject(project.getPath())
- .setPageSize(LIST_BATCH_SIZE);
- ListSubscriptionsResponse response = subscriberStub().listSubscriptions(request.build());
- if (response.getSubscriptionsCount() == 0) {
- return ImmutableList.of();
- }
- List<SubscriptionPath> subscriptions = new ArrayList<>(response.getSubscriptionsCount());
- while (true) {
- for (Subscription subscription : response.getSubscriptionsList()) {
- if (subscription.getTopic().equals(topic.getPath())) {
- subscriptions.add(subscriptionPathFromPath(subscription.getName()));
- }
- }
- if (response.getNextPageToken().isEmpty()) {
- break;
- }
- request.setPageToken(response.getNextPageToken());
- response = subscriberStub().listSubscriptions(request.build());
- }
- return subscriptions;
- }
-
- @Override
- public int ackDeadlineSeconds(SubscriptionPath subscription) throws IOException {
- GetSubscriptionRequest request =
- GetSubscriptionRequest.newBuilder()
- .setSubscription(subscription.getPath())
- .build();
- Subscription response = subscriberStub().getSubscription(request);
- return response.getAckDeadlineSeconds();
- }
-
- @Override
- public boolean isEOF() {
- return false;
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/src/main/java/org/apache/beam/sdk/util/PubsubJsonClient.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/PubsubJsonClient.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/PubsubJsonClient.java
deleted file mode 100644
index ef8abfd..0000000
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/PubsubJsonClient.java
+++ /dev/null
@@ -1,317 +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.util;
-
-import static com.google.common.base.Preconditions.checkState;
-
-import com.google.api.client.http.HttpRequestInitializer;
-import com.google.api.services.pubsub.Pubsub;
-import com.google.api.services.pubsub.Pubsub.Builder;
-import com.google.api.services.pubsub.model.AcknowledgeRequest;
-import com.google.api.services.pubsub.model.ListSubscriptionsResponse;
-import com.google.api.services.pubsub.model.ListTopicsResponse;
-import com.google.api.services.pubsub.model.ModifyAckDeadlineRequest;
-import com.google.api.services.pubsub.model.PublishRequest;
-import com.google.api.services.pubsub.model.PublishResponse;
-import com.google.api.services.pubsub.model.PubsubMessage;
-import com.google.api.services.pubsub.model.PullRequest;
-import com.google.api.services.pubsub.model.PullResponse;
-import com.google.api.services.pubsub.model.ReceivedMessage;
-import com.google.api.services.pubsub.model.Subscription;
-import com.google.api.services.pubsub.model.Topic;
-import com.google.auth.Credentials;
-import com.google.auth.http.HttpCredentialsAdapter;
-import com.google.cloud.hadoop.util.ChainingHttpRequestInitializer;
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Strings;
-import com.google.common.collect.ImmutableList;
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-import java.util.TreeMap;
-import javax.annotation.Nullable;
-import org.apache.beam.sdk.options.PubsubOptions;
-
-/**
- * A Pubsub client using JSON transport.
- */
-public class PubsubJsonClient extends PubsubClient {
-
- private static class PubsubJsonClientFactory implements PubsubClientFactory {
- private static HttpRequestInitializer chainHttpRequestInitializer(
- Credentials credential, HttpRequestInitializer httpRequestInitializer) {
- if (credential == null) {
- return httpRequestInitializer;
- } else {
- return new ChainingHttpRequestInitializer(
- new HttpCredentialsAdapter(credential),
- httpRequestInitializer);
- }
- }
-
- @Override
- public PubsubClient newClient(
- @Nullable String timestampLabel, @Nullable String idLabel, PubsubOptions options)
- throws IOException {
- Pubsub pubsub = new Builder(
- Transport.getTransport(),
- Transport.getJsonFactory(),
- chainHttpRequestInitializer(
- options.getGcpCredential(),
- // Do not log 404. It clutters the output and is possibly even required by the caller.
- new RetryHttpRequestInitializer(ImmutableList.of(404))))
- .setRootUrl(options.getPubsubRootUrl())
- .setApplicationName(options.getAppName())
- .setGoogleClientRequestInitializer(options.getGoogleApiTrace())
- .build();
- return new PubsubJsonClient(timestampLabel, idLabel, pubsub);
- }
-
- @Override
- public String getKind() {
- return "Json";
- }
- }
-
- /**
- * Factory for creating Pubsub clients using Json transport.
- */
- public static final PubsubClientFactory FACTORY = new PubsubJsonClientFactory();
-
- /**
- * Label to use for custom timestamps, or {@literal null} if should use Pubsub publish time
- * instead.
- */
- @Nullable
- private final String timestampLabel;
-
- /**
- * Label to use for custom ids, or {@literal null} if should use Pubsub provided ids.
- */
- @Nullable
- private final String idLabel;
-
- /**
- * Underlying JSON transport.
- */
- private Pubsub pubsub;
-
- @VisibleForTesting
- PubsubJsonClient(
- @Nullable String timestampLabel,
- @Nullable String idLabel,
- Pubsub pubsub) {
- this.timestampLabel = timestampLabel;
- this.idLabel = idLabel;
- this.pubsub = pubsub;
- }
-
- @Override
- public void close() {
- // Nothing to close.
- }
-
- @Override
- public int publish(TopicPath topic, List<OutgoingMessage> outgoingMessages)
- throws IOException {
- List<PubsubMessage> pubsubMessages = new ArrayList<>(outgoingMessages.size());
- for (OutgoingMessage outgoingMessage : outgoingMessages) {
- PubsubMessage pubsubMessage = new PubsubMessage().encodeData(outgoingMessage.elementBytes);
-
- Map<String, String> attributes = outgoingMessage.attributes;
- if ((timestampLabel != null || idLabel != null) && attributes == null) {
- attributes = new TreeMap<>();
- }
- if (attributes != null) {
- pubsubMessage.setAttributes(attributes);
- }
-
- if (timestampLabel != null) {
- attributes.put(timestampLabel, String.valueOf(outgoingMessage.timestampMsSinceEpoch));
- }
-
- if (idLabel != null && !Strings.isNullOrEmpty(outgoingMessage.recordId)) {
- attributes.put(idLabel, outgoingMessage.recordId);
- }
-
- pubsubMessages.add(pubsubMessage);
- }
- PublishRequest request = new PublishRequest().setMessages(pubsubMessages);
- PublishResponse response = pubsub.projects()
- .topics()
- .publish(topic.getPath(), request)
- .execute();
- return response.getMessageIds().size();
- }
-
- @Override
- public List<IncomingMessage> pull(
- long requestTimeMsSinceEpoch,
- SubscriptionPath subscription,
- int batchSize,
- boolean returnImmediately) throws IOException {
- PullRequest request = new PullRequest()
- .setReturnImmediately(returnImmediately)
- .setMaxMessages(batchSize);
- PullResponse response = pubsub.projects()
- .subscriptions()
- .pull(subscription.getPath(), request)
- .execute();
- if (response.getReceivedMessages() == null || response.getReceivedMessages().size() == 0) {
- return ImmutableList.of();
- }
- List<IncomingMessage> incomingMessages = new ArrayList<>(response.getReceivedMessages().size());
- for (ReceivedMessage message : response.getReceivedMessages()) {
- PubsubMessage pubsubMessage = message.getMessage();
- @Nullable Map<String, String> attributes = pubsubMessage.getAttributes();
-
- // Payload.
- byte[] elementBytes = pubsubMessage.decodeData();
-
- // Timestamp.
- long timestampMsSinceEpoch =
- extractTimestamp(timestampLabel, message.getMessage().getPublishTime(), attributes);
-
- // Ack id.
- String ackId = message.getAckId();
- checkState(!Strings.isNullOrEmpty(ackId));
-
- // Record id, if any.
- @Nullable String recordId = null;
- if (idLabel != null && attributes != null) {
- recordId = attributes.get(idLabel);
- }
- if (Strings.isNullOrEmpty(recordId)) {
- // Fall back to the Pubsub provided message id.
- recordId = pubsubMessage.getMessageId();
- }
-
- incomingMessages.add(new IncomingMessage(elementBytes, attributes, timestampMsSinceEpoch,
- requestTimeMsSinceEpoch, ackId, recordId));
- }
-
- return incomingMessages;
- }
-
- @Override
- public void acknowledge(SubscriptionPath subscription, List<String> ackIds) throws IOException {
- AcknowledgeRequest request = new AcknowledgeRequest().setAckIds(ackIds);
- pubsub.projects()
- .subscriptions()
- .acknowledge(subscription.getPath(), request)
- .execute(); // ignore Empty result.
- }
-
- @Override
- public void modifyAckDeadline(
- SubscriptionPath subscription, List<String> ackIds, int deadlineSeconds)
- throws IOException {
- ModifyAckDeadlineRequest request =
- new ModifyAckDeadlineRequest().setAckIds(ackIds)
- .setAckDeadlineSeconds(deadlineSeconds);
- pubsub.projects()
- .subscriptions()
- .modifyAckDeadline(subscription.getPath(), request)
- .execute(); // ignore Empty result.
- }
-
- @Override
- public void createTopic(TopicPath topic) throws IOException {
- pubsub.projects()
- .topics()
- .create(topic.getPath(), new Topic())
- .execute(); // ignore Topic result.
- }
-
- @Override
- public void deleteTopic(TopicPath topic) throws IOException {
- pubsub.projects()
- .topics()
- .delete(topic.getPath())
- .execute(); // ignore Empty result.
- }
-
- @Override
- public List<TopicPath> listTopics(ProjectPath project) throws IOException {
- ListTopicsResponse response = pubsub.projects()
- .topics()
- .list(project.getPath())
- .execute();
- if (response.getTopics() == null || response.getTopics().isEmpty()) {
- return ImmutableList.of();
- }
- List<TopicPath> topics = new ArrayList<>(response.getTopics().size());
- for (Topic topic : response.getTopics()) {
- topics.add(topicPathFromPath(topic.getName()));
- }
- return topics;
- }
-
- @Override
- public void createSubscription(
- TopicPath topic, SubscriptionPath subscription,
- int ackDeadlineSeconds) throws IOException {
- Subscription request = new Subscription()
- .setTopic(topic.getPath())
- .setAckDeadlineSeconds(ackDeadlineSeconds);
- pubsub.projects()
- .subscriptions()
- .create(subscription.getPath(), request)
- .execute(); // ignore Subscription result.
- }
-
- @Override
- public void deleteSubscription(SubscriptionPath subscription) throws IOException {
- pubsub.projects()
- .subscriptions()
- .delete(subscription.getPath())
- .execute(); // ignore Empty result.
- }
-
- @Override
- public List<SubscriptionPath> listSubscriptions(ProjectPath project, TopicPath topic)
- throws IOException {
- ListSubscriptionsResponse response = pubsub.projects()
- .subscriptions()
- .list(project.getPath())
- .execute();
- if (response.getSubscriptions() == null || response.getSubscriptions().isEmpty()) {
- return ImmutableList.of();
- }
- List<SubscriptionPath> subscriptions = new ArrayList<>(response.getSubscriptions().size());
- for (Subscription subscription : response.getSubscriptions()) {
- if (subscription.getTopic().equals(topic.getPath())) {
- subscriptions.add(subscriptionPathFromPath(subscription.getName()));
- }
- }
- return subscriptions;
- }
-
- @Override
- public int ackDeadlineSeconds(SubscriptionPath subscription) throws IOException {
- Subscription response = pubsub.projects().subscriptions().get(subscription.getPath()).execute();
- return response.getAckDeadlineSeconds();
- }
-
- @Override
- public boolean isEOF() {
- return false;
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/src/main/java/org/apache/beam/sdk/util/PubsubTestClient.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/PubsubTestClient.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/PubsubTestClient.java
deleted file mode 100644
index 61479f9..0000000
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/PubsubTestClient.java
+++ /dev/null
@@ -1,436 +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.util;
-
-import static com.google.common.base.Preconditions.checkState;
-
-import com.google.api.client.util.Clock;
-import com.google.common.collect.Lists;
-import com.google.common.collect.Sets;
-import java.io.Closeable;
-import java.io.IOException;
-import java.io.Serializable;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import javax.annotation.Nullable;
-import org.apache.beam.sdk.options.PubsubOptions;
-
-/**
- * A (partial) implementation of {@link PubsubClient} for use by unit tests. Only suitable for
- * testing {@link #publish}, {@link #pull}, {@link #acknowledge} and {@link #modifyAckDeadline}
- * methods. Relies on statics to mimic the Pubsub service, though we try to hide that.
- */
-public class PubsubTestClient extends PubsubClient implements Serializable {
- /**
- * Mimic the state of the simulated Pubsub 'service'.
- *
- * <p>Note that the {@link PubsubTestClientFactory} is serialized/deserialized even when running
- * test pipelines. Meanwhile it is valid for multiple {@link PubsubTestClient}s to be created
- * from the same client factory and run in parallel. Thus we can't enforce aliasing of the
- * following data structures over all clients and must resort to a static.
- */
- private static class State {
- /**
- * True if has been primed for a test but not yet validated.
- */
- boolean isActive;
-
- /**
- * Publish mode only: Only publish calls for this topic are allowed.
- */
- @Nullable
- TopicPath expectedTopic;
-
- /**
- * Publish mode only: Messages yet to seen in a {@link #publish} call.
- */
- @Nullable
- Set<OutgoingMessage> remainingExpectedOutgoingMessages;
-
- /**
- * Publish mode only: Messages which should throw when first sent to simulate transient publish
- * failure.
- */
- @Nullable
- Set<OutgoingMessage> remainingFailingOutgoingMessages;
-
- /**
- * Pull mode only: Clock from which to get current time.
- */
- @Nullable
- Clock clock;
-
- /**
- * Pull mode only: Only pull calls for this subscription are allowed.
- */
- @Nullable
- SubscriptionPath expectedSubscription;
-
- /**
- * Pull mode only: Timeout to simulate.
- */
- int ackTimeoutSec;
-
- /**
- * Pull mode only: Messages waiting to be received by a {@link #pull} call.
- */
- @Nullable
- List<IncomingMessage> remainingPendingIncomingMessages;
-
- /**
- * Pull mode only: Messages which have been returned from a {@link #pull} call and
- * not yet ACKed by an {@link #acknowledge} call.
- */
- @Nullable
- Map<String, IncomingMessage> pendingAckIncomingMessages;
-
- /**
- * Pull mode only: When above messages are due to have their ACK deadlines expire.
- */
- @Nullable
- Map<String, Long> ackDeadline;
- }
-
- private static final State STATE = new State();
-
- /** Closing the factory will validate all expected messages were processed. */
- public interface PubsubTestClientFactory
- extends PubsubClientFactory, Closeable, Serializable {
- }
-
- /**
- * Return a factory for testing publishers. Only one factory may be in-flight at a time.
- * The factory must be closed when the test is complete, at which point final validation will
- * occur.
- */
- public static PubsubTestClientFactory createFactoryForPublish(
- final TopicPath expectedTopic,
- final Iterable<OutgoingMessage> expectedOutgoingMessages,
- final Iterable<OutgoingMessage> failingOutgoingMessages) {
- synchronized (STATE) {
- checkState(!STATE.isActive, "Test still in flight");
- STATE.expectedTopic = expectedTopic;
- STATE.remainingExpectedOutgoingMessages = Sets.newHashSet(expectedOutgoingMessages);
- STATE.remainingFailingOutgoingMessages = Sets.newHashSet(failingOutgoingMessages);
- STATE.isActive = true;
- }
- return new PubsubTestClientFactory() {
- @Override
- public PubsubClient newClient(
- @Nullable String timestampLabel, @Nullable String idLabel, PubsubOptions options)
- throws IOException {
- return new PubsubTestClient();
- }
-
- @Override
- public String getKind() {
- return "PublishTest";
- }
-
- @Override
- public void close() {
- synchronized (STATE) {
- checkState(STATE.isActive, "No test still in flight");
- checkState(STATE.remainingExpectedOutgoingMessages.isEmpty(),
- "Still waiting for %s messages to be published",
- STATE.remainingExpectedOutgoingMessages.size());
- STATE.isActive = false;
- STATE.remainingExpectedOutgoingMessages = null;
- }
- }
- };
- }
-
- /**
- * Return a factory for testing subscribers. Only one factory may be in-flight at a time.
- * The factory must be closed when the test in complete
- */
- public static PubsubTestClientFactory createFactoryForPull(
- final Clock clock,
- final SubscriptionPath expectedSubscription,
- final int ackTimeoutSec,
- final Iterable<IncomingMessage> expectedIncomingMessages) {
- synchronized (STATE) {
- checkState(!STATE.isActive, "Test still in flight");
- STATE.clock = clock;
- STATE.expectedSubscription = expectedSubscription;
- STATE.ackTimeoutSec = ackTimeoutSec;
- STATE.remainingPendingIncomingMessages = Lists.newArrayList(expectedIncomingMessages);
- STATE.pendingAckIncomingMessages = new HashMap<>();
- STATE.ackDeadline = new HashMap<>();
- STATE.isActive = true;
- }
- return new PubsubTestClientFactory() {
- @Override
- public PubsubClient newClient(
- @Nullable String timestampLabel, @Nullable String idLabel, PubsubOptions options)
- throws IOException {
- return new PubsubTestClient();
- }
-
- @Override
- public String getKind() {
- return "PullTest";
- }
-
- @Override
- public void close() {
- synchronized (STATE) {
- checkState(STATE.isActive, "No test still in flight");
- checkState(STATE.remainingPendingIncomingMessages.isEmpty(),
- "Still waiting for %s messages to be pulled",
- STATE.remainingPendingIncomingMessages.size());
- checkState(STATE.pendingAckIncomingMessages.isEmpty(),
- "Still waiting for %s messages to be ACKed",
- STATE.pendingAckIncomingMessages.size());
- checkState(STATE.ackDeadline.isEmpty(),
- "Still waiting for %s messages to be ACKed",
- STATE.ackDeadline.size());
- STATE.isActive = false;
- STATE.remainingPendingIncomingMessages = null;
- STATE.pendingAckIncomingMessages = null;
- STATE.ackDeadline = null;
- }
- }
- };
- }
-
- public static PubsubTestClientFactory createFactoryForCreateSubscription() {
- return new PubsubTestClientFactory() {
- int numCalls = 0;
-
- @Override
- public void close() throws IOException {
- checkState(
- numCalls == 1, "Expected exactly one subscription to be created, got %s", numCalls);
- }
-
- @Override
- public PubsubClient newClient(
- @Nullable String timestampLabel, @Nullable String idLabel, PubsubOptions options)
- throws IOException {
- return new PubsubTestClient() {
- @Override
- public void createSubscription(
- TopicPath topic, SubscriptionPath subscription, int ackDeadlineSeconds)
- throws IOException {
- checkState(numCalls == 0, "Expected at most one subscription to be created");
- numCalls++;
- }
- };
- }
-
- @Override
- public String getKind() {
- return "CreateSubscriptionTest";
- }
- };
- }
-
- /**
- * Return true if in pull mode.
- */
- private boolean inPullMode() {
- checkState(STATE.isActive, "No test is active");
- return STATE.expectedSubscription != null;
- }
-
- /**
- * Return true if in publish mode.
- */
- private boolean inPublishMode() {
- checkState(STATE.isActive, "No test is active");
- return STATE.expectedTopic != null;
- }
-
- /**
- * For subscription mode only:
- * Track progression of time according to the {@link Clock} passed . This will simulate Pubsub
- * expiring
- * outstanding ACKs.
- */
- public void advance() {
- synchronized (STATE) {
- checkState(inPullMode(), "Can only advance in pull mode");
- // Any messages who's ACKs timed out are available for re-pulling.
- Iterator<Map.Entry<String, Long>> deadlineItr = STATE.ackDeadline.entrySet().iterator();
- while (deadlineItr.hasNext()) {
- Map.Entry<String, Long> entry = deadlineItr.next();
- if (entry.getValue() <= STATE.clock.currentTimeMillis()) {
- STATE.remainingPendingIncomingMessages.add(
- STATE.pendingAckIncomingMessages.remove(entry.getKey()));
- deadlineItr.remove();
- }
- }
- }
- }
-
- @Override
- public void close() {
- }
-
- @Override
- public int publish(
- TopicPath topic, List<OutgoingMessage> outgoingMessages) throws IOException {
- synchronized (STATE) {
- checkState(inPublishMode(), "Can only publish in publish mode");
- checkState(topic.equals(STATE.expectedTopic), "Topic %s does not match expected %s", topic,
- STATE.expectedTopic);
- for (OutgoingMessage outgoingMessage : outgoingMessages) {
- if (STATE.remainingFailingOutgoingMessages.remove(outgoingMessage)) {
- throw new RuntimeException("Simulating failure for " + outgoingMessage);
- }
- checkState(STATE.remainingExpectedOutgoingMessages.remove(outgoingMessage),
- "Unexpected outgoing message %s", outgoingMessage);
- }
- return outgoingMessages.size();
- }
- }
-
- @Override
- public List<IncomingMessage> pull(
- long requestTimeMsSinceEpoch, SubscriptionPath subscription, int batchSize,
- boolean returnImmediately) throws IOException {
- synchronized (STATE) {
- checkState(inPullMode(), "Can only pull in pull mode");
- long now = STATE.clock.currentTimeMillis();
- checkState(requestTimeMsSinceEpoch == now,
- "Simulated time %s does not match request time %s", now, requestTimeMsSinceEpoch);
- checkState(subscription.equals(STATE.expectedSubscription),
- "Subscription %s does not match expected %s", subscription,
- STATE.expectedSubscription);
- checkState(returnImmediately, "Pull only supported if returning immediately");
-
- List<IncomingMessage> incomingMessages = new ArrayList<>();
- Iterator<IncomingMessage> pendItr = STATE.remainingPendingIncomingMessages.iterator();
- while (pendItr.hasNext()) {
- IncomingMessage incomingMessage = pendItr.next();
- pendItr.remove();
- IncomingMessage incomingMessageWithRequestTime =
- incomingMessage.withRequestTime(requestTimeMsSinceEpoch);
- incomingMessages.add(incomingMessageWithRequestTime);
- STATE.pendingAckIncomingMessages.put(incomingMessageWithRequestTime.ackId,
- incomingMessageWithRequestTime);
- STATE.ackDeadline.put(incomingMessageWithRequestTime.ackId,
- requestTimeMsSinceEpoch + STATE.ackTimeoutSec * 1000);
- if (incomingMessages.size() >= batchSize) {
- break;
- }
- }
- return incomingMessages;
- }
- }
-
- @Override
- public void acknowledge(
- SubscriptionPath subscription,
- List<String> ackIds) throws IOException {
- synchronized (STATE) {
- checkState(inPullMode(), "Can only acknowledge in pull mode");
- checkState(subscription.equals(STATE.expectedSubscription),
- "Subscription %s does not match expected %s", subscription,
- STATE.expectedSubscription);
-
- for (String ackId : ackIds) {
- checkState(STATE.ackDeadline.remove(ackId) != null,
- "No message with ACK id %s is waiting for an ACK", ackId);
- checkState(STATE.pendingAckIncomingMessages.remove(ackId) != null,
- "No message with ACK id %s is waiting for an ACK", ackId);
- }
- }
- }
-
- @Override
- public void modifyAckDeadline(
- SubscriptionPath subscription, List<String> ackIds, int deadlineSeconds) throws IOException {
- synchronized (STATE) {
- checkState(inPullMode(), "Can only modify ack deadline in pull mode");
- checkState(subscription.equals(STATE.expectedSubscription),
- "Subscription %s does not match expected %s", subscription,
- STATE.expectedSubscription);
-
- for (String ackId : ackIds) {
- if (deadlineSeconds > 0) {
- checkState(STATE.ackDeadline.remove(ackId) != null,
- "No message with ACK id %s is waiting for an ACK", ackId);
- checkState(STATE.pendingAckIncomingMessages.containsKey(ackId),
- "No message with ACK id %s is waiting for an ACK", ackId);
- STATE.ackDeadline.put(ackId, STATE.clock.currentTimeMillis() + deadlineSeconds * 1000);
- } else {
- checkState(STATE.ackDeadline.remove(ackId) != null,
- "No message with ACK id %s is waiting for an ACK", ackId);
- IncomingMessage message = STATE.pendingAckIncomingMessages.remove(ackId);
- checkState(message != null, "No message with ACK id %s is waiting for an ACK", ackId);
- STATE.remainingPendingIncomingMessages.add(message);
- }
- }
- }
- }
-
- @Override
- public void createTopic(TopicPath topic) throws IOException {
- throw new UnsupportedOperationException();
- }
-
- @Override
- public void deleteTopic(TopicPath topic) throws IOException {
- throw new UnsupportedOperationException();
- }
-
- @Override
- public List<TopicPath> listTopics(ProjectPath project) throws IOException {
- throw new UnsupportedOperationException();
- }
-
- @Override
- public void createSubscription(
- TopicPath topic, SubscriptionPath subscription, int ackDeadlineSeconds) throws IOException {
- throw new UnsupportedOperationException();
- }
-
- @Override
- public void deleteSubscription(SubscriptionPath subscription) throws IOException {
- throw new UnsupportedOperationException();
- }
-
- @Override
- public List<SubscriptionPath> listSubscriptions(
- ProjectPath project, TopicPath topic) throws IOException {
- throw new UnsupportedOperationException();
- }
-
- @Override
- public int ackDeadlineSeconds(SubscriptionPath subscription) throws IOException {
- synchronized (STATE) {
- return STATE.ackTimeoutSec;
- }
- }
-
- @Override
- public boolean isEOF() {
- synchronized (STATE) {
- checkState(inPullMode(), "Can only check EOF in pull mode");
- return STATE.remainingPendingIncomingMessages.isEmpty();
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/src/main/java/org/apache/beam/sdk/util/Transport.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/Transport.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/Transport.java
index 1edfa1d..80c093b 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/Transport.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/Transport.java
@@ -107,8 +107,7 @@ public class Transport {
/**
* Returns a Pubsub client builder using the specified {@link PubsubOptions}.
*
- * @deprecated Use an appropriate
- * {@link org.apache.beam.sdk.util.PubsubClient.PubsubClientFactory}
+ * @deprecated Use an appropriate org.apache.beam.sdk.util.PubsubClient.PubsubClientFactory
*/
@Deprecated
public static Pubsub.Builder
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/src/test/java/org/apache/beam/sdk/io/PubsubIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/PubsubIOTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/PubsubIOTest.java
deleted file mode 100644
index c996409..0000000
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/PubsubIOTest.java
+++ /dev/null
@@ -1,189 +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 static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem;
-import static org.hamcrest.Matchers.hasItem;
-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 java.util.Set;
-import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
-import org.apache.beam.sdk.testing.UsesUnboundedPCollections;
-import org.apache.beam.sdk.testing.ValidatesRunner;
-import org.apache.beam.sdk.transforms.display.DisplayData;
-import org.apache.beam.sdk.transforms.display.DisplayDataEvaluator;
-import org.joda.time.Duration;
-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;
-
-/**
- * Tests for PubsubIO Read and Write transforms.
- */
-@RunWith(JUnit4.class)
-public class PubsubIOTest {
- @Rule
- public ExpectedException thrown = ExpectedException.none();
-
- @Test
- public void testPubsubIOGetName() {
- assertEquals("PubsubIO.Read",
- PubsubIO.<String>read().topic("projects/myproject/topics/mytopic").getName());
- assertEquals("PubsubIO.Write",
- PubsubIO.<String>write().topic("projects/myproject/topics/mytopic").getName());
- }
-
- @Test
- public void testTopicValidationSuccess() throws Exception {
- PubsubIO.<String>read().topic("projects/my-project/topics/abc");
- PubsubIO.<String>read().topic("projects/my-project/topics/ABC");
- PubsubIO.<String>read().topic("projects/my-project/topics/AbC-DeF");
- PubsubIO.<String>read().topic("projects/my-project/topics/AbC-1234");
- PubsubIO.<String>read().topic("projects/my-project/topics/AbC-1234-_.~%+-_.~%+-_.~%+-abc");
- PubsubIO.<String>read().topic(new StringBuilder()
- .append("projects/my-project/topics/A-really-long-one-")
- .append("111111111111111111111111111111111111111111111111111111111111111111111111111111111")
- .append("111111111111111111111111111111111111111111111111111111111111111111111111111111111")
- .append("11111111111111111111111111111111111111111111111111111111111111111111111111")
- .toString());
- }
-
- @Test
- public void testTopicValidationBadCharacter() throws Exception {
- thrown.expect(IllegalArgumentException.class);
- PubsubIO.<String>read().topic("projects/my-project/topics/abc-*-abc");
- }
-
- @Test
- public void testTopicValidationTooLong() throws Exception {
- thrown.expect(IllegalArgumentException.class);
- PubsubIO.<String>read().topic(new StringBuilder().append
- ("projects/my-project/topics/A-really-long-one-")
- .append("111111111111111111111111111111111111111111111111111111111111111111111111111111111")
- .append("111111111111111111111111111111111111111111111111111111111111111111111111111111111")
- .append("1111111111111111111111111111111111111111111111111111111111111111111111111111")
- .toString());
- }
-
- @Test
- public void testReadTopicDisplayData() {
- String topic = "projects/project/topics/topic";
- String subscription = "projects/project/subscriptions/subscription";
- Duration maxReadTime = Duration.standardMinutes(5);
- PubsubIO.Read<String> read = PubsubIO.<String>read()
- .topic(StaticValueProvider.of(topic))
- .timestampLabel("myTimestamp")
- .idLabel("myId");
-
- DisplayData displayData = DisplayData.from(read);
-
- assertThat(displayData, hasDisplayItem("topic", topic));
- assertThat(displayData, hasDisplayItem("timestampLabel", "myTimestamp"));
- assertThat(displayData, hasDisplayItem("idLabel", "myId"));
- }
-
- @Test
- public void testReadSubscriptionDisplayData() {
- String topic = "projects/project/topics/topic";
- String subscription = "projects/project/subscriptions/subscription";
- Duration maxReadTime = Duration.standardMinutes(5);
- PubsubIO.Read<String> read = PubsubIO.<String>read()
- .subscription(StaticValueProvider.of(subscription))
- .timestampLabel("myTimestamp")
- .idLabel("myId");
-
- DisplayData displayData = DisplayData.from(read);
-
- assertThat(displayData, hasDisplayItem("subscription", subscription));
- assertThat(displayData, hasDisplayItem("timestampLabel", "myTimestamp"));
- assertThat(displayData, hasDisplayItem("idLabel", "myId"));
- }
-
- @Test
- public void testNullTopic() {
- String subscription = "projects/project/subscriptions/subscription";
- PubsubIO.Read<String> read = PubsubIO.<String>read()
- .subscription(StaticValueProvider.of(subscription));
- assertNull(read.getTopic());
- assertNotNull(read.getSubscription());
- assertNotNull(DisplayData.from(read));
- }
-
- @Test
- public void testNullSubscription() {
- String topic = "projects/project/topics/topic";
- PubsubIO.Read<String> read = PubsubIO.<String>read()
- .topic(StaticValueProvider.of(topic));
- assertNotNull(read.getTopic());
- assertNull(read.getSubscription());
- assertNotNull(DisplayData.from(read));
- }
-
- @Test
- @Category({ValidatesRunner.class, UsesUnboundedPCollections.class})
- public void testPrimitiveReadDisplayData() {
- DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
- Set<DisplayData> displayData;
- PubsubIO.Read<String> read = PubsubIO.<String>read().withCoder(StringUtf8Coder.of());
-
- // Reading from a subscription.
- read = read.subscription("projects/project/subscriptions/subscription");
- displayData = evaluator.displayDataForPrimitiveSourceTransforms(read);
- assertThat("PubsubIO.Read should include the subscription in its primitive display data",
- displayData, hasItem(hasDisplayItem("subscription")));
-
- // Reading from a topic.
- read = read.topic("projects/project/topics/topic");
- displayData = evaluator.displayDataForPrimitiveSourceTransforms(read);
- assertThat("PubsubIO.Read should include the topic in its primitive display data",
- displayData, hasItem(hasDisplayItem("topic")));
- }
-
- @Test
- public void testWriteDisplayData() {
- String topic = "projects/project/topics/topic";
- PubsubIO.Write<?> write = PubsubIO.<String>write()
- .topic(topic)
- .timestampLabel("myTimestamp")
- .idLabel("myId");
-
- DisplayData displayData = DisplayData.from(write);
-
- assertThat(displayData, hasDisplayItem("topic", topic));
- assertThat(displayData, hasDisplayItem("timestampLabel", "myTimestamp"));
- assertThat(displayData, hasDisplayItem("idLabel", "myId"));
- }
-
- @Test
- @Category(ValidatesRunner.class)
- public void testPrimitiveWriteDisplayData() {
- DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
- PubsubIO.Write<?> write = PubsubIO.<String>write().topic("projects/project/topics/topic");
-
- Set<DisplayData> displayData = evaluator.displayDataForPrimitiveTransforms(write);
- assertThat("PubsubIO.Write should include the topic in its primitive display data",
- displayData, hasItem(hasDisplayItem("topic")));
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/src/test/java/org/apache/beam/sdk/io/PubsubUnboundedSinkTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/PubsubUnboundedSinkTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/PubsubUnboundedSinkTest.java
deleted file mode 100644
index 7a4be62..0000000
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/PubsubUnboundedSinkTest.java
+++ /dev/null
@@ -1,190 +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 com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.hash.Hashing;
-import java.io.IOException;
-import java.io.Serializable;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-
-import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.sdk.io.PubsubUnboundedSink.RecordIdMethod;
-import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
-import org.apache.beam.sdk.testing.CoderProperties;
-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.DoFn;
-import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.transforms.SimpleFunction;
-import org.apache.beam.sdk.util.PubsubClient;
-import org.apache.beam.sdk.util.PubsubClient.OutgoingMessage;
-import org.apache.beam.sdk.util.PubsubClient.TopicPath;
-import org.apache.beam.sdk.util.PubsubTestClient;
-import org.apache.beam.sdk.util.PubsubTestClient.PubsubTestClientFactory;
-import org.joda.time.Duration;
-import org.joda.time.Instant;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.experimental.categories.Category;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-/**
- * Test PubsubUnboundedSink.
- */
-@RunWith(JUnit4.class)
-public class PubsubUnboundedSinkTest implements Serializable {
- private static final TopicPath TOPIC = PubsubClient.topicPathFromName("testProject", "testTopic");
- private static final String DATA = "testData";
- private static final Map<String, String> ATTRIBUTES =
- ImmutableMap.<String, String>builder().put("a", "b").put("c", "d").build();
- private static final long TIMESTAMP = 1234L;
- private static final String TIMESTAMP_LABEL = "timestamp";
- private static final String ID_LABEL = "id";
- private static final int NUM_SHARDS = 10;
-
- private static class Stamp extends DoFn<String, String> {
- @ProcessElement
- public void processElement(ProcessContext c) {
- c.outputWithTimestamp(c.element(), new Instant(TIMESTAMP));
- }
- }
-
- private String getRecordId(String data) {
- return Hashing.murmur3_128().hashBytes(data.getBytes()).toString();
- }
-
- @Rule
- public transient TestPipeline p = TestPipeline.create();
-
- @Test
- public void saneCoder() throws Exception {
- OutgoingMessage message = new OutgoingMessage(
- DATA.getBytes(), ImmutableMap.<String, String>of(), TIMESTAMP, getRecordId(DATA));
- CoderProperties.coderDecodeEncodeEqual(PubsubUnboundedSink.CODER, message);
- CoderProperties.coderSerializable(PubsubUnboundedSink.CODER);
- }
-
- @Test
- @Category(NeedsRunner.class)
- public void sendOneMessage() throws IOException {
- List<OutgoingMessage> outgoing =
- ImmutableList.of(new OutgoingMessage(
- DATA.getBytes(),
- ATTRIBUTES,
- TIMESTAMP, getRecordId(DATA)));
- int batchSize = 1;
- int batchBytes = 1;
- try (PubsubTestClientFactory factory =
- PubsubTestClient.createFactoryForPublish(TOPIC, outgoing,
- ImmutableList.<OutgoingMessage>of())) {
- PubsubUnboundedSink<String> sink =
- new PubsubUnboundedSink<>(factory, StaticValueProvider.of(TOPIC), StringUtf8Coder.of(),
- TIMESTAMP_LABEL, ID_LABEL, NUM_SHARDS, batchSize, batchBytes,
- Duration.standardSeconds(2),
- new SimpleFunction<String, PubsubIO.PubsubMessage>() {
- @Override
- public PubsubIO.PubsubMessage apply(String input) {
- return new PubsubIO.PubsubMessage(input.getBytes(), ATTRIBUTES);
- }
- },
- RecordIdMethod.DETERMINISTIC);
- p.apply(Create.of(ImmutableList.of(DATA)))
- .apply(ParDo.of(new Stamp()))
- .apply(sink);
- p.run();
- }
- // The PubsubTestClientFactory will assert fail on close if the actual published
- // message does not match the expected publish message.
- }
-
- @Test
- @Category(NeedsRunner.class)
- public void sendMoreThanOneBatchByNumMessages() throws IOException {
- List<OutgoingMessage> outgoing = new ArrayList<>();
- List<String> data = new ArrayList<>();
- int batchSize = 2;
- int batchBytes = 1000;
- for (int i = 0; i < batchSize * 10; i++) {
- String str = String.valueOf(i);
- outgoing.add(new OutgoingMessage(
- str.getBytes(), ImmutableMap.<String, String>of(), TIMESTAMP, getRecordId(str)));
- data.add(str);
- }
- try (PubsubTestClientFactory factory =
- PubsubTestClient.createFactoryForPublish(TOPIC, outgoing,
- ImmutableList.<OutgoingMessage>of())) {
- PubsubUnboundedSink<String> sink =
- new PubsubUnboundedSink<>(factory, StaticValueProvider.of(TOPIC), StringUtf8Coder.of(),
- TIMESTAMP_LABEL, ID_LABEL, NUM_SHARDS, batchSize, batchBytes,
- Duration.standardSeconds(2), null, RecordIdMethod.DETERMINISTIC);
- p.apply(Create.of(data))
- .apply(ParDo.of(new Stamp()))
- .apply(sink);
- p.run();
- }
- // The PubsubTestClientFactory will assert fail on close if the actual published
- // message does not match the expected publish message.
- }
-
- @Test
- @Category(NeedsRunner.class)
- public void sendMoreThanOneBatchByByteSize() throws IOException {
- List<OutgoingMessage> outgoing = new ArrayList<>();
- List<String> data = new ArrayList<>();
- int batchSize = 100;
- int batchBytes = 10;
- int n = 0;
- while (n < batchBytes * 10) {
- StringBuilder sb = new StringBuilder();
- for (int i = 0; i < batchBytes; i++) {
- sb.append(String.valueOf(n));
- }
- String str = sb.toString();
- outgoing.add(new OutgoingMessage(
- str.getBytes(), ImmutableMap.<String, String>of(), TIMESTAMP, getRecordId(str)));
- data.add(str);
- n += str.length();
- }
- try (PubsubTestClientFactory factory =
- PubsubTestClient.createFactoryForPublish(TOPIC, outgoing,
- ImmutableList.<OutgoingMessage>of())) {
- PubsubUnboundedSink<String> sink =
- new PubsubUnboundedSink<>(factory, StaticValueProvider.of(TOPIC),
- StringUtf8Coder.of(), TIMESTAMP_LABEL, ID_LABEL,
- NUM_SHARDS, batchSize, batchBytes, Duration.standardSeconds(2),
- null, RecordIdMethod.DETERMINISTIC);
- p.apply(Create.of(data))
- .apply(ParDo.of(new Stamp()))
- .apply(sink);
- p.run();
- }
- // The PubsubTestClientFactory will assert fail on close if the actual published
- // message does not match the expected publish message.
- }
-
- // TODO: We would like to test that failed Pubsub publish calls cause the already assigned
- // (and random) record ids to be reused. However that can't be done without the test runnner
- // supporting retrying bundles.
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/src/test/java/org/apache/beam/sdk/io/PubsubUnboundedSourceTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/PubsubUnboundedSourceTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/PubsubUnboundedSourceTest.java
deleted file mode 100644
index d9df2ca..0000000
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/PubsubUnboundedSourceTest.java
+++ /dev/null
@@ -1,411 +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 static junit.framework.TestCase.assertFalse;
-import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.greaterThan;
-import static org.hamcrest.Matchers.hasSize;
-import static org.hamcrest.Matchers.lessThanOrEqualTo;
-import static org.hamcrest.Matchers.not;
-import static org.hamcrest.Matchers.nullValue;
-import static org.junit.Assert.assertArrayEquals;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertThat;
-import static org.junit.Assert.assertTrue;
-
-import com.google.api.client.util.Clock;
-import com.google.common.collect.ImmutableList;
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.concurrent.atomic.AtomicLong;
-import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.sdk.io.PubsubUnboundedSource.PubsubCheckpoint;
-import org.apache.beam.sdk.io.PubsubUnboundedSource.PubsubReader;
-import org.apache.beam.sdk.io.PubsubUnboundedSource.PubsubSource;
-import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.options.PipelineOptionsFactory;
-import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
-import org.apache.beam.sdk.testing.CoderProperties;
-import org.apache.beam.sdk.testing.TestPipeline;
-import org.apache.beam.sdk.util.CoderUtils;
-import org.apache.beam.sdk.util.PubsubClient;
-import org.apache.beam.sdk.util.PubsubClient.IncomingMessage;
-import org.apache.beam.sdk.util.PubsubClient.SubscriptionPath;
-import org.apache.beam.sdk.util.PubsubClient.TopicPath;
-import org.apache.beam.sdk.util.PubsubTestClient;
-import org.apache.beam.sdk.util.PubsubTestClient.PubsubTestClientFactory;
-import org.joda.time.Instant;
-import org.junit.After;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-/**
- * Test PubsubUnboundedSource.
- */
-@RunWith(JUnit4.class)
-public class PubsubUnboundedSourceTest {
- private static final SubscriptionPath SUBSCRIPTION =
- PubsubClient.subscriptionPathFromName("testProject", "testSubscription");
- private static final String DATA = "testData";
- private static final long TIMESTAMP = 1234L;
- private static final long REQ_TIME = 6373L;
- private static final String TIMESTAMP_LABEL = "timestamp";
- private static final String ID_LABEL = "id";
- private static final String ACK_ID = "testAckId";
- private static final String RECORD_ID = "testRecordId";
- private static final int ACK_TIMEOUT_S = 60;
-
- private AtomicLong now;
- private Clock clock;
- private PubsubTestClientFactory factory;
- private PubsubSource<String> primSource;
-
- @Rule
- public TestPipeline p = TestPipeline.create();
-
- private void setupOneMessage(Iterable<IncomingMessage> incoming) {
- now = new AtomicLong(REQ_TIME);
- clock = new Clock() {
- @Override
- public long currentTimeMillis() {
- return now.get();
- }
- };
- factory = PubsubTestClient.createFactoryForPull(clock, SUBSCRIPTION, ACK_TIMEOUT_S, incoming);
- PubsubUnboundedSource<String> source =
- new PubsubUnboundedSource<>(
- clock, factory, null, null, StaticValueProvider.of(SUBSCRIPTION),
- StringUtf8Coder.of(), TIMESTAMP_LABEL, ID_LABEL, null);
- primSource = new PubsubSource<>(source);
- }
-
- private void setupOneMessage() {
- setupOneMessage(ImmutableList.of(
- new IncomingMessage(DATA.getBytes(), null, TIMESTAMP, 0, ACK_ID, RECORD_ID)));
- }
-
- @After
- public void after() throws IOException {
- factory.close();
- now = null;
- clock = null;
- primSource = null;
- factory = null;
- }
-
- @Test
- public void checkpointCoderIsSane() throws Exception {
- setupOneMessage(ImmutableList.<IncomingMessage>of());
- CoderProperties.coderSerializable(primSource.getCheckpointMarkCoder());
- // Since we only serialize/deserialize the 'notYetReadIds', and we don't want to make
- // equals on checkpoints ignore those fields, we'll test serialization and deserialization
- // of checkpoints in multipleReaders below.
- }
-
- @Test
- public void readOneMessage() throws IOException {
- setupOneMessage();
- PubsubReader<String> reader = primSource.createReader(p.getOptions(), null);
- // Read one message.
- assertTrue(reader.start());
- assertEquals(DATA, reader.getCurrent());
- assertFalse(reader.advance());
- // ACK the message.
- PubsubCheckpoint<String> checkpoint = reader.getCheckpointMark();
- checkpoint.finalizeCheckpoint();
- reader.close();
- }
-
- @Test
- public void timeoutAckAndRereadOneMessage() throws IOException {
- setupOneMessage();
- PubsubReader<String> reader = primSource.createReader(p.getOptions(), null);
- PubsubTestClient pubsubClient = (PubsubTestClient) reader.getPubsubClient();
- assertTrue(reader.start());
- assertEquals(DATA, reader.getCurrent());
- // Let the ACK deadline for the above expire.
- now.addAndGet(65 * 1000);
- pubsubClient.advance();
- // We'll now receive the same message again.
- assertTrue(reader.advance());
- assertEquals(DATA, reader.getCurrent());
- assertFalse(reader.advance());
- // Now ACK the message.
- PubsubCheckpoint<String> checkpoint = reader.getCheckpointMark();
- checkpoint.finalizeCheckpoint();
- reader.close();
- }
-
- @Test
- public void extendAck() throws IOException {
- setupOneMessage();
- PubsubReader<String> reader = primSource.createReader(p.getOptions(), null);
- PubsubTestClient pubsubClient = (PubsubTestClient) reader.getPubsubClient();
- // Pull the first message but don't take a checkpoint for it.
- assertTrue(reader.start());
- assertEquals(DATA, reader.getCurrent());
- // Extend the ack
- now.addAndGet(55 * 1000);
- pubsubClient.advance();
- assertFalse(reader.advance());
- // Extend the ack again
- now.addAndGet(25 * 1000);
- pubsubClient.advance();
- assertFalse(reader.advance());
- // Now ACK the message.
- PubsubCheckpoint<String> checkpoint = reader.getCheckpointMark();
- checkpoint.finalizeCheckpoint();
- reader.close();
- }
-
- @Test
- public void timeoutAckExtensions() throws IOException {
- setupOneMessage();
- PubsubReader<String> reader = primSource.createReader(p.getOptions(), null);
- PubsubTestClient pubsubClient = (PubsubTestClient) reader.getPubsubClient();
- // Pull the first message but don't take a checkpoint for it.
- assertTrue(reader.start());
- assertEquals(DATA, reader.getCurrent());
- // Extend the ack.
- now.addAndGet(55 * 1000);
- pubsubClient.advance();
- assertFalse(reader.advance());
- // Let the ack expire.
- for (int i = 0; i < 3; i++) {
- now.addAndGet(25 * 1000);
- pubsubClient.advance();
- assertFalse(reader.advance());
- }
- // Wait for resend.
- now.addAndGet(25 * 1000);
- pubsubClient.advance();
- // Reread the same message.
- assertTrue(reader.advance());
- assertEquals(DATA, reader.getCurrent());
- // Now ACK the message.
- PubsubCheckpoint<String> checkpoint = reader.getCheckpointMark();
- checkpoint.finalizeCheckpoint();
- reader.close();
- }
-
- @Test
- public void multipleReaders() throws IOException {
- List<IncomingMessage> incoming = new ArrayList<>();
- for (int i = 0; i < 2; i++) {
- String data = String.format("data_%d", i);
- String ackid = String.format("ackid_%d", i);
- incoming.add(new IncomingMessage(data.getBytes(), null, TIMESTAMP, 0, ackid, RECORD_ID));
- }
- setupOneMessage(incoming);
- PubsubReader<String> reader = primSource.createReader(p.getOptions(), null);
- // Consume two messages, only read one.
- assertTrue(reader.start());
- assertEquals("data_0", reader.getCurrent());
-
- // Grab checkpoint.
- PubsubCheckpoint<String> checkpoint = reader.getCheckpointMark();
- checkpoint.finalizeCheckpoint();
- assertEquals(1, checkpoint.notYetReadIds.size());
- assertEquals("ackid_1", checkpoint.notYetReadIds.get(0));
-
- // Read second message.
- assertTrue(reader.advance());
- assertEquals("data_1", reader.getCurrent());
-
- // Restore from checkpoint.
- byte[] checkpointBytes =
- CoderUtils.encodeToByteArray(primSource.getCheckpointMarkCoder(), checkpoint);
- checkpoint = CoderUtils.decodeFromByteArray(primSource.getCheckpointMarkCoder(),
- checkpointBytes);
- assertEquals(1, checkpoint.notYetReadIds.size());
- assertEquals("ackid_1", checkpoint.notYetReadIds.get(0));
-
- // Re-read second message.
- reader = primSource.createReader(p.getOptions(), checkpoint);
- assertTrue(reader.start());
- assertEquals("data_1", reader.getCurrent());
-
- // We are done.
- assertFalse(reader.advance());
-
- // ACK final message.
- checkpoint = reader.getCheckpointMark();
- checkpoint.finalizeCheckpoint();
- reader.close();
- }
-
- private long messageNumToTimestamp(int messageNum) {
- return TIMESTAMP + messageNum * 100;
- }
-
- @Test
- public void readManyMessages() throws IOException {
- Map<String, Integer> dataToMessageNum = new HashMap<>();
-
- final int m = 97;
- final int n = 10000;
- List<IncomingMessage> incoming = new ArrayList<>();
- for (int i = 0; i < n; i++) {
- // Make the messages timestamps slightly out of order.
- int messageNum = ((i / m) * m) + (m - 1) - (i % m);
- String data = String.format("data_%d", messageNum);
- dataToMessageNum.put(data, messageNum);
- String recid = String.format("recordid_%d", messageNum);
- String ackId = String.format("ackid_%d", messageNum);
- incoming.add(new IncomingMessage(data.getBytes(), null, messageNumToTimestamp(messageNum), 0,
- ackId, recid));
- }
- setupOneMessage(incoming);
-
- PubsubReader<String> reader = primSource.createReader(p.getOptions(), null);
- PubsubTestClient pubsubClient = (PubsubTestClient) reader.getPubsubClient();
-
- for (int i = 0; i < n; i++) {
- if (i == 0) {
- assertTrue(reader.start());
- } else {
- assertTrue(reader.advance());
- }
- // We'll checkpoint and ack within the 2min limit.
- now.addAndGet(30);
- pubsubClient.advance();
- String data = reader.getCurrent();
- Integer messageNum = dataToMessageNum.remove(data);
- // No duplicate messages.
- assertNotNull(messageNum);
- // Preserve timestamp.
- assertEquals(new Instant(messageNumToTimestamp(messageNum)), reader.getCurrentTimestamp());
- // Preserve record id.
- String recid = String.format("recordid_%d", messageNum);
- assertArrayEquals(recid.getBytes(), reader.getCurrentRecordId());
-
- if (i % 1000 == 999) {
- // Estimated watermark can never get ahead of actual outstanding messages.
- long watermark = reader.getWatermark().getMillis();
- long minOutstandingTimestamp = Long.MAX_VALUE;
- for (Integer outstandingMessageNum : dataToMessageNum.values()) {
- minOutstandingTimestamp =
- Math.min(minOutstandingTimestamp, messageNumToTimestamp(outstandingMessageNum));
- }
- assertThat(watermark, lessThanOrEqualTo(minOutstandingTimestamp));
- // Ack messages, but only every other finalization.
- PubsubCheckpoint<String> checkpoint = reader.getCheckpointMark();
- if (i % 2000 == 1999) {
- checkpoint.finalizeCheckpoint();
- }
- }
- }
- // We are done.
- assertFalse(reader.advance());
- // We saw each message exactly once.
- assertTrue(dataToMessageNum.isEmpty());
- reader.close();
- }
-
- @Test
- public void noSubscriptionSplitIntoBundlesGeneratesSubscription() throws Exception {
- TopicPath topicPath = PubsubClient.topicPathFromName("my_project", "my_topic");
- factory = PubsubTestClient.createFactoryForCreateSubscription();
- PubsubUnboundedSource<String> source =
- new PubsubUnboundedSource<>(
- factory,
- StaticValueProvider.of(PubsubClient.projectPathFromId("my_project")),
- StaticValueProvider.of(topicPath),
- null,
- StringUtf8Coder.of(),
- null,
- null,
- null);
- assertThat(source.getSubscription(), nullValue());
-
- assertThat(source.getSubscription(), nullValue());
-
- PipelineOptions options = PipelineOptionsFactory.create();
- List<PubsubSource<String>> splits =
- (new PubsubSource<>(source)).generateInitialSplits(3, options);
- // We have at least one returned split
- assertThat(splits, hasSize(greaterThan(0)));
- for (PubsubSource<String> split : splits) {
- // Each split is equal
- assertThat(split, equalTo(splits.get(0)));
- }
-
- assertThat(splits.get(0).subscriptionPath, not(nullValue()));
- }
-
- @Test
- public void noSubscriptionNoSplitGeneratesSubscription() throws Exception {
- TopicPath topicPath = PubsubClient.topicPathFromName("my_project", "my_topic");
- factory = PubsubTestClient.createFactoryForCreateSubscription();
- PubsubUnboundedSource<String> source =
- new PubsubUnboundedSource<>(
- factory,
- StaticValueProvider.of(PubsubClient.projectPathFromId("my_project")),
- StaticValueProvider.of(topicPath),
- null,
- StringUtf8Coder.of(),
- null,
- null,
- null);
- assertThat(source.getSubscription(), nullValue());
-
- assertThat(source.getSubscription(), nullValue());
-
- PipelineOptions options = PipelineOptionsFactory.create();
- PubsubSource<String> actualSource = new PubsubSource<>(source);
- PubsubReader<String> reader = actualSource.createReader(options, null);
- SubscriptionPath createdSubscription = reader.subscription;
- assertThat(createdSubscription, not(nullValue()));
-
- PubsubCheckpoint<String> checkpoint = reader.getCheckpointMark();
- assertThat(checkpoint.subscriptionPath, equalTo(createdSubscription.getPath()));
-
- checkpoint.finalizeCheckpoint();
- PubsubCheckpoint<String> deserCheckpoint =
- CoderUtils.clone(actualSource.getCheckpointMarkCoder(), checkpoint);
- assertThat(checkpoint.subscriptionPath, not(nullValue()));
- assertThat(checkpoint.subscriptionPath, equalTo(deserCheckpoint.subscriptionPath));
-
- PubsubReader<String> readerFromOriginal = actualSource.createReader(options, checkpoint);
- PubsubReader<String> readerFromDeser = actualSource.createReader(options, deserCheckpoint);
-
- assertThat(readerFromOriginal.subscription, equalTo(createdSubscription));
- assertThat(readerFromDeser.subscription, equalTo(createdSubscription));
- }
-
- /**
- * Tests that checkpoints finalized after the reader is closed succeed.
- */
- @Test
- public void closeWithActiveCheckpoints() throws Exception {
- setupOneMessage();
- PubsubReader<String> reader = primSource.createReader(p.getOptions(), null);
- reader.start();
- PubsubCheckpoint<String> checkpoint = reader.getCheckpointMark();
- reader.close();
- checkpoint.finalizeCheckpoint();
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/src/test/java/org/apache/beam/sdk/util/PubsubClientTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/PubsubClientTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/PubsubClientTest.java
deleted file mode 100644
index 1a99d38..0000000
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/PubsubClientTest.java
+++ /dev/null
@@ -1,189 +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.util;
-
-import static org.junit.Assert.assertEquals;
-
-import com.google.common.collect.ImmutableMap;
-import java.util.Map;
-import org.apache.beam.sdk.util.PubsubClient.ProjectPath;
-import org.apache.beam.sdk.util.PubsubClient.SubscriptionPath;
-import org.apache.beam.sdk.util.PubsubClient.TopicPath;
-import org.joda.time.Instant;
-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 helper classes and methods in PubsubClient.
- */
-@RunWith(JUnit4.class)
-public class PubsubClientTest {
- @Rule
- public ExpectedException thrown = ExpectedException.none();
-
- //
- // Timestamp handling
- //
-
- private long parse(String timestamp) {
- Map<String, String> map = ImmutableMap.of("myLabel", timestamp);
- return PubsubClient.extractTimestamp("myLabel", null, map);
- }
-
- private void roundTripRfc339(String timestamp) {
- assertEquals(Instant.parse(timestamp).getMillis(), parse(timestamp));
- }
-
- private void truncatedRfc339(String timestamp, String truncatedTimestmap) {
- assertEquals(Instant.parse(truncatedTimestmap).getMillis(), parse(timestamp));
- }
-
- @Test
- public void noTimestampLabelReturnsPubsubPublish() {
- final long time = 987654321L;
- long timestamp = PubsubClient.extractTimestamp(null, String.valueOf(time), null);
- assertEquals(time, timestamp);
- }
-
- @Test
- public void noTimestampLabelAndInvalidPubsubPublishThrowsError() {
- thrown.expect(NumberFormatException.class);
- PubsubClient.extractTimestamp(null, "not-a-date", null);
- }
-
- @Test
- public void timestampLabelWithNullAttributesThrowsError() {
- thrown.expect(RuntimeException.class);
- thrown.expectMessage("PubSub message is missing a value for timestamp label myLabel");
- PubsubClient.extractTimestamp("myLabel", null, null);
- }
-
- @Test
- public void timestampLabelSetWithMissingAttributeThrowsError() {
- thrown.expect(RuntimeException.class);
- thrown.expectMessage("PubSub message is missing a value for timestamp label myLabel");
- Map<String, String> map = ImmutableMap.of("otherLabel", "whatever");
- PubsubClient.extractTimestamp("myLabel", null, map);
- }
-
- @Test
- public void timestampLabelParsesMillisecondsSinceEpoch() {
- long time = 1446162101123L;
- Map<String, String> map = ImmutableMap.of("myLabel", String.valueOf(time));
- long timestamp = PubsubClient.extractTimestamp("myLabel", null, map);
- assertEquals(time, timestamp);
- }
-
- @Test
- public void timestampLabelParsesRfc3339Seconds() {
- roundTripRfc339("2015-10-29T23:41:41Z");
- }
-
- @Test
- public void timestampLabelParsesRfc3339Tenths() {
- roundTripRfc339("2015-10-29T23:41:41.1Z");
- }
-
- @Test
- public void timestampLabelParsesRfc3339Hundredths() {
- roundTripRfc339("2015-10-29T23:41:41.12Z");
- }
-
- @Test
- public void timestampLabelParsesRfc3339Millis() {
- roundTripRfc339("2015-10-29T23:41:41.123Z");
- }
-
- @Test
- public void timestampLabelParsesRfc3339Micros() {
- // Note: micros part 456/1000 is dropped.
- truncatedRfc339("2015-10-29T23:41:41.123456Z", "2015-10-29T23:41:41.123Z");
- }
-
- @Test
- public void timestampLabelParsesRfc3339MicrosRounding() {
- // Note: micros part 999/1000 is dropped, not rounded up.
- truncatedRfc339("2015-10-29T23:41:41.123999Z", "2015-10-29T23:41:41.123Z");
- }
-
- @Test
- public void timestampLabelWithInvalidFormatThrowsError() {
- thrown.expect(NumberFormatException.class);
- parse("not-a-timestamp");
- }
-
- @Test
- public void timestampLabelWithInvalidFormat2ThrowsError() {
- thrown.expect(NumberFormatException.class);
- parse("null");
- }
-
- @Test
- public void timestampLabelWithInvalidFormat3ThrowsError() {
- thrown.expect(NumberFormatException.class);
- parse("2015-10");
- }
-
- @Test
- public void timestampLabelParsesRfc3339WithSmallYear() {
- // Google and JodaTime agree on dates after 1582-10-15, when the Gregorian Calendar was adopted
- // This is therefore a "small year" until this difference is reconciled.
- roundTripRfc339("1582-10-15T01:23:45.123Z");
- }
-
- @Test
- public void timestampLabelParsesRfc3339WithLargeYear() {
- // Year 9999 in range.
- roundTripRfc339("9999-10-29T23:41:41.123999Z");
- }
-
- @Test
- public void timestampLabelRfc3339WithTooLargeYearThrowsError() {
- thrown.expect(NumberFormatException.class);
- // Year 10000 out of range.
- parse("10000-10-29T23:41:41.123999Z");
- }
-
- //
- // Paths
- //
-
- @Test
- public void projectPathFromIdWellFormed() {
- ProjectPath path = PubsubClient.projectPathFromId("test");
- assertEquals("projects/test", path.getPath());
- }
-
- @Test
- public void subscriptionPathFromNameWellFormed() {
- SubscriptionPath path = PubsubClient.subscriptionPathFromName("test", "something");
- assertEquals("projects/test/subscriptions/something", path.getPath());
- assertEquals("/subscriptions/test/something", path.getV1Beta1Path());
- }
-
- @Test
- public void topicPathFromNameWellFormed() {
- TopicPath path = PubsubClient.topicPathFromName("test", "something");
- assertEquals("projects/test/topics/something", path.getPath());
- assertEquals("/topics/test/something", path.getV1Beta1Path());
- }
-}
[04/50] [abbrv] beam git commit: Remove Jdk1.8-tests/.toDelete
Posted by jb...@apache.org.
Remove Jdk1.8-tests/.toDelete
Breaks RAT. File not useful.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/ee7965c9
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/ee7965c9
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/ee7965c9
Branch: refs/heads/DSL_SQL
Commit: ee7965c92af8f4d7cc1066a466a479a7599a5a57
Parents: d255fa2
Author: Thomas Groh <tg...@google.com>
Authored: Tue Apr 11 10:39:13 2017 -0700
Committer: Thomas Groh <tg...@google.com>
Committed: Tue Apr 11 10:39:13 2017 -0700
----------------------------------------------------------------------
sdks/java/io/hadoop/jdk1.8-tests/.toDelete | 0
1 file changed, 0 insertions(+), 0 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/ee7965c9/sdks/java/io/hadoop/jdk1.8-tests/.toDelete
----------------------------------------------------------------------
diff --git a/sdks/java/io/hadoop/jdk1.8-tests/.toDelete b/sdks/java/io/hadoop/jdk1.8-tests/.toDelete
deleted file mode 100644
index e69de29..0000000
[17/50] [abbrv] beam git commit: This closes #2497
Posted by jb...@apache.org.
This closes #2497
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/bef2d373
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/bef2d373
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/bef2d373
Branch: refs/heads/DSL_SQL
Commit: bef2d3738227cdca1d81123f93ea0f4028711517
Parents: 7587d29 3fb3a0c
Author: Dan Halperin <dh...@google.com>
Authored: Tue Apr 11 15:01:37 2017 -0700
Committer: Dan Halperin <dh...@google.com>
Committed: Tue Apr 11 15:01:37 2017 -0700
----------------------------------------------------------------------
.../beam/runners/dataflow/DataflowPipelineTranslator.java | 3 ---
.../java/org/apache/beam/runners/dataflow/DataflowRunner.java | 7 ++++---
2 files changed, 4 insertions(+), 6 deletions(-)
----------------------------------------------------------------------
[13/50] [abbrv] beam git commit: Remove Orderdness of Input,
Output expansions
Posted by jb...@apache.org.
Remove Orderdness of Input, Output expansions
This brings the PInput and POutput expansion signatures back in line
with the Runner API.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/0e5737fd
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/0e5737fd
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/0e5737fd
Branch: refs/heads/DSL_SQL
Commit: 0e5737fdbee5478ee7f39c4b1a1ac95353ec7b08
Parents: a5a5bf9
Author: Thomas Groh <tg...@google.com>
Authored: Tue Apr 4 16:51:55 2017 -0700
Committer: Thomas Groh <tg...@google.com>
Committed: Tue Apr 11 12:40:46 2017 -0700
----------------------------------------------------------------------
.../FlattenPCollectionTranslator.java | 13 +-
.../apex/translation/ParDoTranslator.java | 13 +-
.../apex/translation/TranslationContext.java | 10 +-
.../DeduplicatedFlattenFactory.java | 18 ++-
.../EmptyFlattenAsCreateFactory.java | 7 +-
.../core/construction/PTransformMatchers.java | 5 +-
.../core/construction/PrimitiveCreate.java | 7 +-
.../core/construction/ReplacementOutputs.java | 63 ++++-----
.../SingleInputOutputOverrideFactory.java | 11 +-
.../UnsupportedOverrideFactory.java | 8 +-
.../DeduplicatedFlattenFactoryTest.java | 6 +-
.../EmptyFlattenAsCreateFactoryTest.java | 8 +-
.../construction/PTransformMatchersTest.java | 132 +++++++++----------
.../construction/ReplacementOutputsTest.java | 109 ++-------------
.../SingleInputOutputOverrideFactoryTest.java | 6 +-
.../UnsupportedOverrideFactoryTest.java | 7 +-
.../direct/BoundedReadEvaluatorFactory.java | 2 +-
.../beam/runners/direct/DirectGraphVisitor.java | 5 +-
.../direct/ExecutorServiceParallelExecutor.java | 4 +-
.../runners/direct/FlattenEvaluatorFactory.java | 2 +-
.../GroupAlsoByWindowEvaluatorFactory.java | 2 +-
.../direct/GroupByKeyOnlyEvaluatorFactory.java | 4 +-
.../direct/KeyedPValueTrackingVisitor.java | 14 +-
.../runners/direct/ParDoEvaluatorFactory.java | 10 +-
.../direct/ParDoMultiOverrideFactory.java | 9 +-
.../direct/StatefulParDoEvaluatorFactory.java | 8 +-
.../direct/TestStreamEvaluatorFactory.java | 8 +-
.../direct/UnboundedReadEvaluatorFactory.java | 4 +-
.../runners/direct/ViewEvaluatorFactory.java | 4 +-
.../runners/direct/ViewOverrideFactory.java | 9 +-
.../beam/runners/direct/WatermarkManager.java | 19 +--
.../runners/direct/WindowEvaluatorFactory.java | 2 +-
.../direct/WriteWithShardingFactory.java | 10 +-
.../runners/direct/DirectGraphVisitorTest.java | 7 +-
.../beam/runners/direct/ParDoEvaluatorTest.java | 2 +-
.../StatefulParDoEvaluatorFactoryTest.java | 2 +-
.../direct/TestStreamEvaluatorFactoryTest.java | 5 +-
.../runners/direct/ViewOverrideFactoryTest.java | 2 +-
.../flink/FlinkBatchTransformTranslators.java | 36 ++---
.../flink/FlinkBatchTranslationContext.java | 11 +-
.../flink/FlinkStreamingPipelineTranslator.java | 9 +-
.../FlinkStreamingTransformTranslators.java | 32 ++---
.../flink/FlinkStreamingTranslationContext.java | 12 +-
.../dataflow/BatchStatefulParDoOverrides.java | 15 +--
.../dataflow/DataflowPipelineTranslator.java | 20 +--
.../beam/runners/dataflow/DataflowRunner.java | 27 ++--
.../runners/dataflow/TransformTranslator.java | 6 +-
.../dataflow/DataflowPipelineJobTest.java | 7 +-
.../apache/beam/runners/spark/SparkRunner.java | 20 +--
.../beam/runners/spark/TestSparkRunner.java | 7 +-
.../spark/translation/EvaluationContext.java | 11 +-
.../spark/translation/TransformTranslator.java | 25 ++--
.../streaming/StreamingTransformTranslator.java | 20 +--
.../sdk/runners/PTransformOverrideFactory.java | 10 +-
.../beam/sdk/runners/TransformHierarchy.java | 80 +++++------
.../beam/sdk/transforms/AppliedPTransform.java | 17 ++-
.../transforms/join/KeyedPCollectionTuple.java | 12 +-
.../java/org/apache/beam/sdk/values/PBegin.java | 6 +-
.../apache/beam/sdk/values/PCollectionList.java | 27 ++--
.../beam/sdk/values/PCollectionTuple.java | 10 +-
.../java/org/apache/beam/sdk/values/PDone.java | 6 +-
.../java/org/apache/beam/sdk/values/PInput.java | 4 +-
.../org/apache/beam/sdk/values/POutput.java | 4 +-
.../java/org/apache/beam/sdk/values/PValue.java | 4 +-
.../org/apache/beam/sdk/values/PValueBase.java | 6 +-
.../apache/beam/sdk/values/TaggedPValue.java | 5 +
.../java/org/apache/beam/sdk/PipelineTest.java | 20 ++-
.../sdk/runners/TransformHierarchyTest.java | 72 +++++-----
.../beam/sdk/values/PCollectionListTest.java | 70 ++++------
.../beam/sdk/values/PCollectionTupleTest.java | 5 +-
.../beam/sdk/io/gcp/bigquery/WriteResult.java | 11 +-
71 files changed, 526 insertions(+), 658 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/FlattenPCollectionTranslator.java
----------------------------------------------------------------------
diff --git a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/FlattenPCollectionTranslator.java b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/FlattenPCollectionTranslator.java
index 080c5e9..440b801 100644
--- a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/FlattenPCollectionTranslator.java
+++ b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/FlattenPCollectionTranslator.java
@@ -32,7 +32,8 @@ import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.transforms.Flatten;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.PValue;
+import org.apache.beam.sdk.values.TupleTag;
/**
* {@link Flatten.PCollections} translation to Apex operator.
@@ -63,15 +64,15 @@ class FlattenPCollectionTranslator<T> implements
}
}
- private List<PCollection<T>> extractPCollections(List<TaggedPValue> inputs) {
+ private List<PCollection<T>> extractPCollections(Map<TupleTag<?>, PValue> inputs) {
List<PCollection<T>> collections = Lists.newArrayList();
- for (TaggedPValue pv : inputs) {
+ for (PValue pv : inputs.values()) {
checkArgument(
- pv.getValue() instanceof PCollection,
+ pv instanceof PCollection,
"Non-PCollection provided as input to flatten: %s of type %s",
- pv.getValue(),
+ pv,
pv.getClass().getSimpleName());
- collections.add((PCollection<T>) pv.getValue());
+ collections.add((PCollection<T>) pv);
}
return collections;
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/ParDoTranslator.java
----------------------------------------------------------------------
diff --git a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/ParDoTranslator.java b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/ParDoTranslator.java
index fa9d21d..9213c1f 100644
--- a/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/ParDoTranslator.java
+++ b/runners/apex/src/main/java/org/apache/beam/runners/apex/translation/ParDoTranslator.java
@@ -27,6 +27,7 @@ import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.Map.Entry;
import org.apache.beam.runners.apex.ApexRunner;
import org.apache.beam.runners.apex.translation.operators.ApexParDoOperator;
import org.apache.beam.sdk.coders.Coder;
@@ -38,7 +39,7 @@ import org.apache.beam.sdk.util.WindowedValue.FullWindowedValueCoder;
import org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TupleTag;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -81,8 +82,8 @@ class ParDoTranslator<InputT, OutputT>
ApexRunner.class.getSimpleName()));
}
- List<TaggedPValue> outputs = context.getOutputs();
- PCollection<InputT> input = context.getInput();
+ Map<TupleTag<?>, PValue> outputs = context.getOutputs();
+ PCollection<InputT> input = (PCollection<InputT>) context.getInput();
List<PCollectionView<?>> sideInputs = transform.getSideInputs();
Coder<InputT> inputCoder = input.getCoder();
WindowedValueCoder<InputT> wvInputCoder =
@@ -100,7 +101,7 @@ class ParDoTranslator<InputT, OutputT>
context.getStateBackend());
Map<PCollection<?>, OutputPort<?>> ports = Maps.newHashMapWithExpectedSize(outputs.size());
- for (TaggedPValue output : outputs) {
+ for (Entry<TupleTag<?>, PValue> output : outputs.entrySet()) {
checkArgument(
output.getValue() instanceof PCollection,
"%s %s outputs non-PCollection %s of type %s",
@@ -109,12 +110,12 @@ class ParDoTranslator<InputT, OutputT>
output.getValue(),
output.getValue().getClass().getSimpleName());
PCollection<?> pc = (PCollection<?>) output.getValue();
- if (output.getTag().equals(transform.getMainOutputTag())) {
+ if (output.getKey().equals(transform.getMainOutputTag())) {
ports.put(pc, operator.output);
} else {
int portIndex = 0;
for (TupleTag<?> tag : transform.getSideOutputTags().getAll()) {
- if (tag.equals(output.getTag())) {
+ if (tag.equals(output.getKey())) {
ports.put(pc, operator.sideOutputPorts[portIndex]);
break;
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 81507ef..c78028e 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
@@ -42,7 +42,7 @@ 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.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.apache.commons.lang3.tuple.ImmutablePair;
import org.apache.commons.lang3.tuple.Pair;
@@ -85,20 +85,20 @@ class TranslationContext {
return getCurrentTransform().getFullName();
}
- public List<TaggedPValue> getInputs() {
+ public Map<TupleTag<?>, PValue> getInputs() {
return getCurrentTransform().getInputs();
}
public <InputT extends PValue> InputT getInput() {
- return (InputT) Iterables.getOnlyElement(getCurrentTransform().getInputs()).getValue();
+ return (InputT) Iterables.getOnlyElement(getCurrentTransform().getInputs().values());
}
- public List<TaggedPValue> getOutputs() {
+ public Map<TupleTag<?>, PValue> getOutputs() {
return getCurrentTransform().getOutputs();
}
public <OutputT extends PValue> OutputT getOutput() {
- return (OutputT) Iterables.getOnlyElement(getCurrentTransform().getOutputs()).getValue();
+ return (OutputT) Iterables.getOnlyElement(getCurrentTransform().getOutputs().values());
}
private AppliedPTransform<?, ?, ?> getCurrentTransform() {
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/DeduplicatedFlattenFactory.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/DeduplicatedFlattenFactory.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/DeduplicatedFlattenFactory.java
index 093385e..c12c548 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/DeduplicatedFlattenFactory.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/DeduplicatedFlattenFactory.java
@@ -19,7 +19,6 @@
package org.apache.beam.runners.core.construction;
import java.util.HashMap;
-import java.util.List;
import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.runners.PTransformOverrideFactory;
@@ -31,7 +30,7 @@ import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PValue;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
/**
* A {@link PTransformOverrideFactory} that will apply a flatten where no element appears in the
@@ -46,6 +45,7 @@ public class DeduplicatedFlattenFactory<T>
}
private DeduplicatedFlattenFactory() {}
+
@Override
public PTransform<PCollectionList<T>, PCollection<T>> getReplacementTransform(
PCollections<T> transform) {
@@ -75,12 +75,16 @@ public class DeduplicatedFlattenFactory<T>
};
}
+ /**
+ * {@inheritDoc}.
+ *
+ * <p>The input {@link PCollectionList} that is constructed will have the same values in the same
+ */
@Override
- public PCollectionList<T> getInput(
- List<TaggedPValue> inputs, Pipeline p) {
+ public PCollectionList<T> getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
PCollectionList<T> pCollections = PCollectionList.empty(p);
- for (TaggedPValue input : inputs) {
- PCollection<T> pcollection = (PCollection<T>) input.getValue();
+ for (PValue input : inputs.values()) {
+ PCollection<T> pcollection = (PCollection<T>) input;
pCollections = pCollections.and(pcollection);
}
return pCollections;
@@ -88,7 +92,7 @@ public class DeduplicatedFlattenFactory<T>
@Override
public Map<PValue, ReplacementOutput> mapOutputs(
- List<TaggedPValue> outputs, PCollection<T> newOutput) {
+ Map<TupleTag<?>, PValue> outputs, PCollection<T> newOutput) {
return ReplacementOutputs.singleton(outputs, newOutput);
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/EmptyFlattenAsCreateFactory.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/EmptyFlattenAsCreateFactory.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/EmptyFlattenAsCreateFactory.java
index 4328cf3..936bc08 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/EmptyFlattenAsCreateFactory.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/EmptyFlattenAsCreateFactory.java
@@ -20,7 +20,6 @@ package org.apache.beam.runners.core.construction;
import static com.google.common.base.Preconditions.checkArgument;
-import java.util.List;
import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.VoidCoder;
@@ -31,7 +30,7 @@ import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PValue;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
/**
* A {@link PTransformOverrideFactory} that provides an empty {@link Create} to replace a {@link
@@ -57,7 +56,7 @@ public class EmptyFlattenAsCreateFactory<T>
@Override
public PCollectionList<T> getInput(
- List<TaggedPValue> inputs, Pipeline p) {
+ Map<TupleTag<?>, PValue> inputs, Pipeline p) {
checkArgument(
inputs.isEmpty(),
"Unexpected nonempty input %s for %s",
@@ -68,7 +67,7 @@ public class EmptyFlattenAsCreateFactory<T>
@Override
public Map<PValue, ReplacementOutput> mapOutputs(
- List<TaggedPValue> outputs, PCollection<T> newOutput) {
+ Map<TupleTag<?>, PValue> outputs, PCollection<T> newOutput) {
return ReplacementOutputs.singleton(outputs, newOutput);
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 6437f7e..94ec38c 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
@@ -37,7 +37,6 @@ import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PValue;
-import org.apache.beam.sdk.values.TaggedPValue;
/**
* A {@link PTransformMatcher} that matches {@link PTransform PTransforms} based on the class of the
@@ -250,8 +249,8 @@ public class PTransformMatchers {
public boolean matches(AppliedPTransform<?, ?, ?> application) {
if (application.getTransform() instanceof Flatten.PCollections) {
Set<PValue> observed = new HashSet<>();
- for (TaggedPValue pvalue : application.getInputs()) {
- boolean firstInstance = observed.add(pvalue.getValue());
+ for (PValue pvalue : application.getInputs().values()) {
+ boolean firstInstance = observed.add(pvalue);
if (!firstInstance) {
return true;
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PrimitiveCreate.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PrimitiveCreate.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PrimitiveCreate.java
index 7bd38b0..9335f3a 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PrimitiveCreate.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PrimitiveCreate.java
@@ -18,7 +18,6 @@
package org.apache.beam.runners.core.construction;
-import java.util.List;
import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.runners.PTransformOverrideFactory;
@@ -30,7 +29,7 @@ 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.PValue;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
/**
* An implementation of {@link Create} that returns a primitive {@link PCollection}.
@@ -63,13 +62,13 @@ public class PrimitiveCreate<T> extends PTransform<PBegin, PCollection<T>> {
}
@Override
- public PBegin getInput(List<TaggedPValue> inputs, Pipeline p) {
+ public PBegin getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
return p.begin();
}
@Override
public Map<PValue, ReplacementOutput> mapOutputs(
- List<TaggedPValue> outputs, PCollection<T> newOutput) {
+ Map<TupleTag<?>, PValue> outputs, PCollection<T> newOutput) {
return ReplacementOutputs.singleton(outputs, newOutput);
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ReplacementOutputs.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ReplacementOutputs.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ReplacementOutputs.java
index 11b4449..3d485ae 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ReplacementOutputs.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ReplacementOutputs.java
@@ -21,10 +21,11 @@ import static com.google.common.base.Preconditions.checkArgument;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
+import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
-import java.util.List;
import java.util.Map;
+import java.util.Map.Entry;
import java.util.Set;
import org.apache.beam.sdk.runners.PTransformOverrideFactory.ReplacementOutput;
import org.apache.beam.sdk.values.POutput;
@@ -39,60 +40,40 @@ public class ReplacementOutputs {
private ReplacementOutputs() {}
public static Map<PValue, ReplacementOutput> singleton(
- List<TaggedPValue> original, PValue replacement) {
- TaggedPValue taggedReplacement = Iterables.getOnlyElement(replacement.expand());
- return ImmutableMap.<PValue, ReplacementOutput>builder()
- .put(
- taggedReplacement.getValue(),
- ReplacementOutput.of(Iterables.getOnlyElement(original), taggedReplacement))
- .build();
- }
-
- public static Map<PValue, ReplacementOutput> ordered(
- List<TaggedPValue> original, POutput replacement) {
- ImmutableMap.Builder<PValue, ReplacementOutput> result = ImmutableMap.builder();
- List<TaggedPValue> replacements = replacement.expand();
- checkArgument(
- original.size() == replacements.size(),
- "Original and Replacements must be the same size. Original: %s Replacement: %s",
- original.size(),
- replacements.size());
- int i = 0;
- for (TaggedPValue replacementPvalue : replacements) {
- result.put(
- replacementPvalue.getValue(), ReplacementOutput.of(original.get(i), replacementPvalue));
- i++;
- }
- return result.build();
+ Map<TupleTag<?>, PValue> original, PValue replacement) {
+ Entry<TupleTag<?>, PValue> originalElement = Iterables.getOnlyElement(original.entrySet());
+ TupleTag<?> replacementTag = Iterables.getOnlyElement(replacement.expand().entrySet()).getKey();
+ return Collections.singletonMap(
+ replacement,
+ ReplacementOutput.of(
+ TaggedPValue.of(originalElement.getKey(), originalElement.getValue()),
+ TaggedPValue.of(replacementTag, replacement)));
}
public static Map<PValue, ReplacementOutput> tagged(
- List<TaggedPValue> original, POutput replacement) {
+ Map<TupleTag<?>, PValue> original, POutput replacement) {
Map<TupleTag<?>, TaggedPValue> originalTags = new HashMap<>();
- for (TaggedPValue value : original) {
- TaggedPValue former = originalTags.put(value.getTag(), value);
- checkArgument(
- former == null || former.equals(value),
- "Found two tags in an expanded output which map to different values: output: %s "
- + "Values: %s and %s",
- original,
- former,
- value);
+ for (Map.Entry<TupleTag<?>, PValue> originalValue : original.entrySet()) {
+ originalTags.put(
+ originalValue.getKey(),
+ TaggedPValue.of(originalValue.getKey(), originalValue.getValue()));
}
ImmutableMap.Builder<PValue, ReplacementOutput> resultBuilder = ImmutableMap.builder();
Set<TupleTag<?>> missingTags = new HashSet<>(originalTags.keySet());
- for (TaggedPValue replacementValue : replacement.expand()) {
- TaggedPValue mapped = originalTags.get(replacementValue.getTag());
+ for (Map.Entry<TupleTag<?>, PValue> replacementValue : replacement.expand().entrySet()) {
+ TaggedPValue mapped = originalTags.get(replacementValue.getKey());
checkArgument(
mapped != null,
"Missing original output for Tag %s and Value %s Between original %s and replacement %s",
- replacementValue.getTag(),
+ replacementValue.getKey(),
replacementValue.getValue(),
original,
replacement.expand());
resultBuilder.put(
- replacementValue.getValue(), ReplacementOutput.of(mapped, replacementValue));
- missingTags.remove(replacementValue.getTag());
+ replacementValue.getValue(),
+ ReplacementOutput.of(
+ mapped, TaggedPValue.of(replacementValue.getKey(), replacementValue.getValue())));
+ missingTags.remove(replacementValue.getKey());
}
ImmutableMap<PValue, ReplacementOutput> result = resultBuilder.build();
checkArgument(
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SingleInputOutputOverrideFactory.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SingleInputOutputOverrideFactory.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SingleInputOutputOverrideFactory.java
index 43bf556..6d0d571 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SingleInputOutputOverrideFactory.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SingleInputOutputOverrideFactory.java
@@ -19,17 +19,16 @@
package org.apache.beam.runners.core.construction;
import com.google.common.collect.Iterables;
-import java.util.List;
import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.runners.PTransformOverrideFactory;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.values.PValue;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
/**
* A {@link PTransformOverrideFactory} which consumes from a {@link PValue} and produces a
- * {@link PValue}. {@link #getInput(List, Pipeline)} and {@link #mapOutputs(List, PValue)} are
+ * {@link PValue}. {@link #getInput(Map, Pipeline)} and {@link #mapOutputs(Map, PValue)} are
* implemented.
*/
public abstract class SingleInputOutputOverrideFactory<
@@ -38,13 +37,13 @@ public abstract class SingleInputOutputOverrideFactory<
TransformT extends PTransform<InputT, OutputT>>
implements PTransformOverrideFactory<InputT, OutputT, TransformT> {
@Override
- public final InputT getInput(List<TaggedPValue> inputs, Pipeline p) {
- return (InputT) Iterables.getOnlyElement(inputs).getValue();
+ public final InputT getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
+ return (InputT) Iterables.getOnlyElement(inputs.values());
}
@Override
public final Map<PValue, ReplacementOutput> mapOutputs(
- List<TaggedPValue> outputs, OutputT newOutput) {
+ Map<TupleTag<?>, PValue> outputs, OutputT newOutput) {
return ReplacementOutputs.singleton(outputs, newOutput);
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/UnsupportedOverrideFactory.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/UnsupportedOverrideFactory.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/UnsupportedOverrideFactory.java
index 38cbd2a..7b9d704 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/UnsupportedOverrideFactory.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/UnsupportedOverrideFactory.java
@@ -18,7 +18,6 @@
package org.apache.beam.runners.core.construction;
-import java.util.List;
import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.runners.PTransformOverrideFactory;
@@ -26,7 +25,7 @@ import org.apache.beam.sdk.transforms.PTransform;
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.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
/**
* A {@link PTransformOverrideFactory} that throws an exception when a call to
@@ -60,12 +59,13 @@ public final class UnsupportedOverrideFactory<
}
@Override
- public InputT getInput(List<TaggedPValue> inputs, Pipeline p) {
+ public InputT getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
throw new UnsupportedOperationException(message);
}
@Override
- public Map<PValue, ReplacementOutput> mapOutputs(List<TaggedPValue> outputs, OutputT newOutput) {
+ public Map<PValue, ReplacementOutput> mapOutputs(
+ Map<TupleTag<?>, PValue> outputs, OutputT newOutput) {
throw new UnsupportedOperationException(message);
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/DeduplicatedFlattenFactoryTest.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/DeduplicatedFlattenFactoryTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/DeduplicatedFlattenFactoryTest.java
index a251f5a..14aa1e6 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/DeduplicatedFlattenFactoryTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/DeduplicatedFlattenFactoryTest.java
@@ -22,7 +22,6 @@ import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.not;
import static org.junit.Assert.assertThat;
-import com.google.common.collect.Iterables;
import org.apache.beam.sdk.Pipeline.PipelineVisitor.Defaults;
import org.apache.beam.sdk.runners.PTransformOverrideFactory.ReplacementOutput;
import org.apache.beam.sdk.runners.TransformHierarchy;
@@ -35,6 +34,7 @@ import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PValue;
+import org.apache.beam.sdk.values.TaggedPValue;
import org.hamcrest.Matchers;
import org.junit.Rule;
import org.junit.Test;
@@ -106,7 +106,7 @@ public class DeduplicatedFlattenFactoryTest {
Matchers.<PValue, ReplacementOutput>hasEntry(
replacement,
ReplacementOutput.of(
- Iterables.getOnlyElement(original.expand()),
- Iterables.getOnlyElement(replacement.expand()))));
+ TaggedPValue.ofExpandedValue(original),
+ TaggedPValue.ofExpandedValue(replacement))));
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/EmptyFlattenAsCreateFactoryTest.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/EmptyFlattenAsCreateFactoryTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/EmptyFlattenAsCreateFactoryTest.java
index ad9d908..90bbee7 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/EmptyFlattenAsCreateFactoryTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/EmptyFlattenAsCreateFactoryTest.java
@@ -21,7 +21,6 @@ package org.apache.beam.runners.core.construction;
import static org.hamcrest.Matchers.equalTo;
import static org.junit.Assert.assertThat;
-import com.google.common.collect.Iterables;
import java.util.Collections;
import java.util.Map;
import org.apache.beam.sdk.io.CountingInput;
@@ -34,6 +33,7 @@ import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.hamcrest.Matchers;
import org.junit.Rule;
import org.junit.Test;
@@ -55,7 +55,7 @@ public class EmptyFlattenAsCreateFactoryTest {
@Test
public void getInputEmptySucceeds() {
assertThat(
- factory.getInput(Collections.<TaggedPValue>emptyList(), pipeline).size(), equalTo(0));
+ factory.getInput(Collections.<TupleTag<?>, PValue>emptyMap(), pipeline).size(), equalTo(0));
}
@Test
@@ -80,8 +80,8 @@ public class EmptyFlattenAsCreateFactoryTest {
Matchers.<PValue, ReplacementOutput>hasEntry(
replacement,
ReplacementOutput.of(
- Iterables.getOnlyElement(original.expand()),
- Iterables.getOnlyElement(replacement.expand()))));
+ TaggedPValue.ofExpandedValue(original),
+ TaggedPValue.ofExpandedValue(replacement))));
}
@Test
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 d053f62..4084cdc 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
@@ -24,7 +24,7 @@ import static org.hamcrest.Matchers.not;
import static org.junit.Assert.assertThat;
import com.google.common.base.MoreObjects;
-import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
import java.io.Serializable;
import java.util.Collections;
import org.apache.beam.sdk.coders.VarIntCoder;
@@ -65,7 +65,7 @@ import org.apache.beam.sdk.values.PCollection.IsBounded;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PDone;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
import org.hamcrest.Matchers;
@@ -385,17 +385,15 @@ public class PTransformMatchersTest implements Serializable {
public void emptyFlattenWithEmptyFlatten() {
AppliedPTransform application =
AppliedPTransform
- .<PCollectionList<Object>, PCollection<Object>, Flatten.PCollections<Object>>
- of(
- "EmptyFlatten",
- Collections.<TaggedPValue>emptyList(),
- Collections.singletonList(
- TaggedPValue.of(
- new TupleTag<Object>(),
- PCollection.createPrimitiveOutputInternal(
- p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED))),
- Flatten.pCollections(),
- p);
+ .<PCollectionList<Object>, PCollection<Object>, Flatten.PCollections<Object>>of(
+ "EmptyFlatten",
+ Collections.<TupleTag<?>, PValue>emptyMap(),
+ Collections.<TupleTag<?>, PValue>singletonMap(
+ new TupleTag<Object>(),
+ PCollection.createPrimitiveOutputInternal(
+ p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)),
+ Flatten.pCollections(),
+ p);
assertThat(PTransformMatchers.emptyFlatten().matches(application), is(true));
}
@@ -404,21 +402,18 @@ public class PTransformMatchersTest implements Serializable {
public void emptyFlattenWithNonEmptyFlatten() {
AppliedPTransform application =
AppliedPTransform
- .<PCollectionList<Object>, PCollection<Object>, Flatten.PCollections<Object>>
- of(
- "Flatten",
- Collections.singletonList(
- TaggedPValue.of(
- new TupleTag<Object>(),
- PCollection.createPrimitiveOutputInternal(
- p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED))),
- Collections.singletonList(
- TaggedPValue.of(
- new TupleTag<Object>(),
- PCollection.createPrimitiveOutputInternal(
- p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED))),
- Flatten.pCollections(),
- p);
+ .<PCollectionList<Object>, PCollection<Object>, Flatten.PCollections<Object>>of(
+ "Flatten",
+ Collections.<TupleTag<?>, PValue>singletonMap(
+ new TupleTag<Object>(),
+ PCollection.createPrimitiveOutputInternal(
+ p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)),
+ Collections.<TupleTag<?>, PValue>singletonMap(
+ new TupleTag<Object>(),
+ PCollection.createPrimitiveOutputInternal(
+ p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)),
+ Flatten.pCollections(),
+ p);
assertThat(PTransformMatchers.emptyFlatten().matches(application), is(false));
}
@@ -427,18 +422,16 @@ public class PTransformMatchersTest implements Serializable {
public void emptyFlattenWithNonFlatten() {
AppliedPTransform application =
AppliedPTransform
- .<PCollection<Iterable<Object>>, PCollection<Object>, Flatten.Iterables<Object>>
- of(
- "EmptyFlatten",
- Collections.<TaggedPValue>emptyList(),
- Collections.singletonList(
- TaggedPValue.of(
- new TupleTag<Object>(),
- PCollection.createPrimitiveOutputInternal(
- p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED))),
- Flatten.iterables() /* This isn't actually possible to construct,
+ .<PCollection<Iterable<Object>>, PCollection<Object>, Flatten.Iterables<Object>>of(
+ "EmptyFlatten",
+ Collections.<TupleTag<?>, PValue>emptyMap(),
+ Collections.<TupleTag<?>, PValue>singletonMap(
+ new TupleTag<Object>(),
+ PCollection.createPrimitiveOutputInternal(
+ p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)),
+ Flatten.iterables() /* This isn't actually possible to construct,
* but for the sake of example */,
- p);
+ p);
assertThat(PTransformMatchers.emptyFlatten().matches(application), is(false));
}
@@ -447,19 +440,16 @@ public class PTransformMatchersTest implements Serializable {
public void flattenWithDuplicateInputsWithoutDuplicates() {
AppliedPTransform application =
AppliedPTransform
- .<PCollectionList<Object>, PCollection<Object>, Flatten.PCollections<Object>>
- of(
+ .<PCollectionList<Object>, PCollection<Object>, Flatten.PCollections<Object>>of(
"Flatten",
- Collections.singletonList(
- TaggedPValue.of(
- new TupleTag<Object>(),
- PCollection.createPrimitiveOutputInternal(
- p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED))),
- Collections.singletonList(
- TaggedPValue.of(
- new TupleTag<Object>(),
- PCollection.createPrimitiveOutputInternal(
- p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED))),
+ Collections.<TupleTag<?>, PValue>singletonMap(
+ new TupleTag<Object>(),
+ PCollection.createPrimitiveOutputInternal(
+ p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)),
+ Collections.<TupleTag<?>, PValue>singletonMap(
+ new TupleTag<Object>(),
+ PCollection.createPrimitiveOutputInternal(
+ p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)),
Flatten.pCollections(),
p);
@@ -468,21 +458,21 @@ public class PTransformMatchersTest implements Serializable {
@Test
public void flattenWithDuplicateInputsWithDuplicates() {
- PCollection<Object> duplicate = PCollection.createPrimitiveOutputInternal(p,
- WindowingStrategy.globalDefault(),
- IsBounded.BOUNDED);
+ PCollection<Object> duplicate =
+ PCollection.createPrimitiveOutputInternal(
+ p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED);
AppliedPTransform application =
AppliedPTransform
.<PCollectionList<Object>, PCollection<Object>, Flatten.PCollections<Object>>of(
"Flatten",
- ImmutableList.of(
- TaggedPValue.of(new TupleTag<Object>(), duplicate),
- TaggedPValue.of(new TupleTag<Object>(), duplicate)),
- Collections.singletonList(
- TaggedPValue.of(
- new TupleTag<Object>(),
- PCollection.createPrimitiveOutputInternal(
- p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED))),
+ ImmutableMap.<TupleTag<?>, PValue>builder()
+ .put(new TupleTag<Object>(), duplicate)
+ .put(new TupleTag<Object>(), duplicate)
+ .build(),
+ Collections.<TupleTag<?>, PValue>singletonMap(
+ new TupleTag<Object>(),
+ PCollection.createPrimitiveOutputInternal(
+ p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)),
Flatten.pCollections(),
p);
@@ -493,15 +483,13 @@ public class PTransformMatchersTest implements Serializable {
public void flattenWithDuplicateInputsNonFlatten() {
AppliedPTransform application =
AppliedPTransform
- .<PCollection<Iterable<Object>>, PCollection<Object>, Flatten.Iterables<Object>>
- of(
+ .<PCollection<Iterable<Object>>, PCollection<Object>, Flatten.Iterables<Object>>of(
"EmptyFlatten",
- Collections.<TaggedPValue>emptyList(),
- Collections.singletonList(
- TaggedPValue.of(
- new TupleTag<Object>(),
- PCollection.createPrimitiveOutputInternal(
- p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED))),
+ Collections.<TupleTag<?>, PValue>emptyMap(),
+ Collections.<TupleTag<?>, PValue>singletonMap(
+ new TupleTag<Object>(),
+ PCollection.createPrimitiveOutputInternal(
+ p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)),
Flatten.iterables() /* This isn't actually possible to construct,
* but for the sake of example */,
p);
@@ -541,8 +529,8 @@ public class PTransformMatchersTest implements Serializable {
private AppliedPTransform<?, ?, ?> appliedWrite(Write<Integer> write) {
return AppliedPTransform.<PCollection<Integer>, PDone, Write<Integer>>of(
"Write",
- Collections.<TaggedPValue>emptyList(),
- Collections.<TaggedPValue>emptyList(),
+ Collections.<TupleTag<?>, PValue>emptyMap(),
+ Collections.<TupleTag<?>, PValue>emptyMap(),
write,
p);
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ReplacementOutputsTest.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ReplacementOutputsTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ReplacementOutputsTest.java
index abfdeef..00c436d 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ReplacementOutputsTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ReplacementOutputsTest.java
@@ -21,18 +21,15 @@ package org.apache.beam.runners.core.construction;
import static org.hamcrest.Matchers.equalTo;
import static org.junit.Assert.assertThat;
-import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
-import java.util.List;
import java.util.Map;
import org.apache.beam.sdk.runners.PTransformOverrideFactory.ReplacementOutput;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.util.WindowingStrategy;
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.POutput;
import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TaggedPValue;
import org.apache.beam.sdk.values.TupleTag;
@@ -79,8 +76,10 @@ public class ReplacementOutputsTest {
assertThat(replacements, Matchers.<PValue>hasKey(replacementInts));
ReplacementOutput replacement = replacements.get(replacementInts);
- TaggedPValue taggedInts = Iterables.getOnlyElement(ints.expand());
- assertThat(replacement.getOriginal(), equalTo(taggedInts));
+ Map.Entry<TupleTag<?>, PValue> taggedInts = Iterables.getOnlyElement(ints.expand().entrySet());
+ assertThat(
+ replacement.getOriginal().getTag(), Matchers.<TupleTag<?>>equalTo(taggedInts.getKey()));
+ assertThat(replacement.getOriginal().getValue(), equalTo(taggedInts.getValue()));
assertThat(replacement.getReplacement().getValue(), Matchers.<PValue>equalTo(replacementInts));
}
@@ -88,44 +87,11 @@ public class ReplacementOutputsTest {
public void singletonMultipleOriginalsThrows() {
thrown.expect(IllegalArgumentException.class);
ReplacementOutputs.singleton(
- ImmutableList.copyOf(Iterables.concat(ints.expand(), moreInts.expand())), replacementInts);
- }
-
- @Test
- public void orderedSucceeds() {
- List<TaggedPValue> originals = PCollectionList.of(ints).and(moreInts).expand();
- Map<PValue, ReplacementOutput> replacements =
- ReplacementOutputs.ordered(
- originals, PCollectionList.of(replacementInts).and(moreReplacementInts));
- assertThat(
- replacements.keySet(),
- Matchers.<PValue>containsInAnyOrder(replacementInts, moreReplacementInts));
-
- ReplacementOutput intsMapping = replacements.get(replacementInts);
- assertThat(intsMapping.getOriginal().getValue(), Matchers.<PValue>equalTo(ints));
- assertThat(intsMapping.getReplacement().getValue(), Matchers.<PValue>equalTo(replacementInts));
-
- ReplacementOutput moreIntsMapping = replacements.get(moreReplacementInts);
- assertThat(moreIntsMapping.getOriginal().getValue(), Matchers.<PValue>equalTo(moreInts));
- assertThat(
- moreIntsMapping.getReplacement().getValue(), Matchers.<PValue>equalTo(moreReplacementInts));
- }
-
- @Test
- public void orderedTooManyReplacements() {
- thrown.expect(IllegalArgumentException.class);
- thrown.expectMessage("same size");
- ReplacementOutputs.ordered(
- PCollectionList.of(ints).expand(),
- PCollectionList.of(replacementInts).and(moreReplacementInts));
- }
-
- @Test
- public void orderedTooFewReplacements() {
- thrown.expect(IllegalArgumentException.class);
- thrown.expectMessage("same size");
- ReplacementOutputs.ordered(
- PCollectionList.of(ints).and(moreInts).expand(), PCollectionList.of(moreReplacementInts));
+ ImmutableMap.<TupleTag<?>, PValue>builder()
+ .putAll(ints.expand())
+ .putAll(moreInts.expand())
+ .build(),
+ replacementInts);
}
private TupleTag<Integer> intsTag = new TupleTag<>();
@@ -168,61 +134,6 @@ public class ReplacementOutputsTest {
TaggedPValue.of(moreIntsTag, moreReplacementInts))));
}
- /**
- * When a call to {@link ReplacementOutputs#tagged(List, POutput)} is made where the first
- * argument contains multiple copies of the same {@link TaggedPValue}, the call succeeds using
- * that mapping.
- */
- @Test
- public void taggedMultipleInstances() {
- List<TaggedPValue> original =
- ImmutableList.of(
- TaggedPValue.of(intsTag, ints),
- TaggedPValue.of(strsTag, strs),
- TaggedPValue.of(intsTag, ints));
-
- Map<PValue, ReplacementOutput> replacements =
- ReplacementOutputs.tagged(
- original, PCollectionTuple.of(strsTag, replacementStrs).and(intsTag, replacementInts));
- assertThat(
- replacements.keySet(),
- Matchers.<PValue>containsInAnyOrder(replacementStrs, replacementInts));
- ReplacementOutput intsReplacement = replacements.get(replacementInts);
- ReplacementOutput strsReplacement = replacements.get(replacementStrs);
-
- assertThat(
- intsReplacement,
- equalTo(
- ReplacementOutput.of(
- TaggedPValue.of(intsTag, ints), TaggedPValue.of(intsTag, replacementInts))));
- assertThat(
- strsReplacement,
- equalTo(
- ReplacementOutput.of(
- TaggedPValue.of(strsTag, strs), TaggedPValue.of(strsTag, replacementStrs))));
- }
-
- /**
- * When a call to {@link ReplacementOutputs#tagged(List, POutput)} is made where a single tag
- * has multiple {@link PValue PValues} mapped to it, the call fails.
- */
- @Test
- public void taggedMultipleConflictingInstancesThrows() {
- List<TaggedPValue> original =
- ImmutableList.of(
- TaggedPValue.of(intsTag, ints), TaggedPValue.of(intsTag, moreReplacementInts));
- thrown.expect(IllegalArgumentException.class);
- thrown.expectMessage("different values");
- thrown.expectMessage(intsTag.toString());
- thrown.expectMessage(ints.toString());
- thrown.expectMessage(moreReplacementInts.toString());
- ReplacementOutputs.tagged(
- original,
- PCollectionTuple.of(strsTag, replacementStrs)
- .and(moreIntsTag, moreReplacementInts)
- .and(intsTag, replacementInts));
- }
-
@Test
public void taggedMissingReplacementThrows() {
PCollectionTuple original =
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SingleInputOutputOverrideFactoryTest.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SingleInputOutputOverrideFactoryTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SingleInputOutputOverrideFactoryTest.java
index b4cdd1f..07352f5 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SingleInputOutputOverrideFactoryTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/SingleInputOutputOverrideFactoryTest.java
@@ -20,7 +20,6 @@ package org.apache.beam.runners.core.construction;
import static org.junit.Assert.assertThat;
-import com.google.common.collect.Iterables;
import java.io.Serializable;
import java.util.Map;
import org.apache.beam.sdk.runners.PTransformOverrideFactory.ReplacementOutput;
@@ -32,6 +31,7 @@ import org.apache.beam.sdk.transforms.SimpleFunction;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PValue;
+import org.apache.beam.sdk.values.TaggedPValue;
import org.hamcrest.Matchers;
import org.junit.Rule;
import org.junit.Test;
@@ -97,8 +97,8 @@ public class SingleInputOutputOverrideFactoryTest implements Serializable {
Matchers.<PValue, ReplacementOutput>hasEntry(
reappliedOutput,
ReplacementOutput.of(
- Iterables.getOnlyElement(output.expand()),
- Iterables.getOnlyElement(reappliedOutput.expand()))));
+ TaggedPValue.ofExpandedValue(output),
+ TaggedPValue.ofExpandedValue(reappliedOutput))));
}
@Test
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/UnsupportedOverrideFactoryTest.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/UnsupportedOverrideFactoryTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/UnsupportedOverrideFactoryTest.java
index f33d0f9..81ce00d 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/UnsupportedOverrideFactoryTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/UnsupportedOverrideFactoryTest.java
@@ -23,7 +23,8 @@ import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.values.PDone;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.PValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
@@ -53,13 +54,13 @@ public class UnsupportedOverrideFactoryTest {
public void getInputThrows() {
thrown.expect(UnsupportedOperationException.class);
thrown.expectMessage(message);
- factory.getInput(Collections.<TaggedPValue>emptyList(), pipeline);
+ factory.getInput(Collections.<TupleTag<?>, PValue>emptyMap(), pipeline);
}
@Test
public void mapOutputThrows() {
thrown.expect(UnsupportedOperationException.class);
thrown.expectMessage(message);
- factory.mapOutputs(Collections.<TaggedPValue>emptyList(), PDone.in(pipeline));
+ factory.mapOutputs(Collections.<TupleTag<?>, PValue>emptyMap(), PDone.in(pipeline));
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/direct-java/src/main/java/org/apache/beam/runners/direct/BoundedReadEvaluatorFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/BoundedReadEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/BoundedReadEvaluatorFactory.java
index 57735e7..5bd6f7e 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/BoundedReadEvaluatorFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/BoundedReadEvaluatorFactory.java
@@ -117,7 +117,7 @@ final class BoundedReadEvaluatorFactory implements TransformEvaluatorFactory {
ExecutorService executor) {
this.evaluationContext = evaluationContext;
this.outputPCollection =
- (PCollection<OutputT>) Iterables.getOnlyElement(transform.getOutputs()).getValue();
+ (PCollection<OutputT>) Iterables.getOnlyElement(transform.getOutputs().values());
this.resultBuilder = StepTransformResult.withoutHold(transform);
this.minimumDynamicSplitSize = minimumDynamicSplitSize;
this.produceSplitExecutor = executor;
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 7e6845d..c342136 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
@@ -34,7 +34,6 @@ import org.apache.beam.sdk.values.PCollectionView;
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.TaggedPValue;
/**
* Tracks the {@link AppliedPTransform AppliedPTransforms} that consume each {@link PValue} in the
@@ -83,8 +82,8 @@ class DirectGraphVisitor extends PipelineVisitor.Defaults {
if (node.getInputs().isEmpty()) {
rootTransforms.add(appliedTransform);
} else {
- for (TaggedPValue value : node.getInputs()) {
- primitiveConsumers.put(value.getValue(), appliedTransform);
+ for (PValue value : node.getInputs().values()) {
+ primitiveConsumers.put(value, appliedTransform);
}
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 8b9f995..c802c58 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
@@ -534,8 +534,8 @@ final class ExecutorServiceParallelExecutor implements PipelineExecutor {
.createKeyedBundle(
transformTimers.getKey(),
(PCollection)
- Iterables.getOnlyElement(transformTimers.getTransform().getInputs())
- .getValue())
+ Iterables.getOnlyElement(
+ transformTimers.getTransform().getInputs().values()))
.add(WindowedValue.valueInGlobalWindow(work))
.commit(evaluationContext.now());
scheduleConsumption(
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/direct-java/src/main/java/org/apache/beam/runners/direct/FlattenEvaluatorFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/FlattenEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/FlattenEvaluatorFactory.java
index 8528905..7c6d2a1 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/FlattenEvaluatorFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/FlattenEvaluatorFactory.java
@@ -57,7 +57,7 @@ class FlattenEvaluatorFactory implements TransformEvaluatorFactory {
application) {
final UncommittedBundle<InputT> outputBundle =
evaluationContext.createBundle(
- (PCollection<InputT>) Iterables.getOnlyElement(application.getOutputs()).getValue());
+ (PCollection<InputT>) Iterables.getOnlyElement(application.getOutputs().values()));
final TransformResult<InputT> result =
StepTransformResult.<InputT>withoutHold(application).addOutput(outputBundle).build();
return new FlattenEvaluator<>(outputBundle, result);
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java
index 04e5aaa..f7fd4cf 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupAlsoByWindowEvaluatorFactory.java
@@ -162,7 +162,7 @@ class GroupAlsoByWindowEvaluatorFactory implements TransformEvaluatorFactory {
evaluationContext.createKeyedBundle(
structuralKey,
(PCollection<KV<K, Iterable<V>>>)
- Iterables.getOnlyElement(application.getOutputs()).getValue());
+ Iterables.getOnlyElement(application.getOutputs().values()));
outputBundles.add(bundle);
CopyOnAccessInMemoryStateInternals<K> stateInternals =
(CopyOnAccessInMemoryStateInternals<K>) stepContext.stateInternals();
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupByKeyOnlyEvaluatorFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupByKeyOnlyEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupByKeyOnlyEvaluatorFactory.java
index ef96a92..ac0b14f 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupByKeyOnlyEvaluatorFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/GroupByKeyOnlyEvaluatorFactory.java
@@ -105,7 +105,7 @@ class GroupByKeyOnlyEvaluatorFactory implements TransformEvaluatorFactory {
this.application = application;
this.keyCoder =
getKeyCoder(
- ((PCollection<KV<K, V>>) Iterables.getOnlyElement(application.getInputs()).getValue())
+ ((PCollection<KV<K, V>>) Iterables.getOnlyElement(application.getInputs().values()))
.getCoder());
this.groupingMap = new HashMap<>();
}
@@ -158,7 +158,7 @@ class GroupByKeyOnlyEvaluatorFactory implements TransformEvaluatorFactory {
evaluationContext.createKeyedBundle(
StructuralKey.of(key, keyCoder),
(PCollection<KeyedWorkItem<K, V>>)
- Iterables.getOnlyElement(application.getOutputs()).getValue());
+ Iterables.getOnlyElement(application.getOutputs().values()));
bundle.add(WindowedValue.valueInGlobalWindow(groupedKv));
resultBuilder.addOutput(bundle);
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java
index 02b1bed..f9b6eba 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java
@@ -21,7 +21,7 @@ import static com.google.common.base.Preconditions.checkState;
import com.google.common.collect.ImmutableSet;
import java.util.HashSet;
-import java.util.List;
+import java.util.Map;
import java.util.Set;
import org.apache.beam.runners.core.SplittableParDo;
import org.apache.beam.runners.direct.DirectGroupByKey.DirectGroupAlsoByWindow;
@@ -32,7 +32,7 @@ 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.values.PValue;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
/**
* A pipeline visitor that tracks all keyed {@link PValue PValues}. A {@link PValue} is keyed if it
@@ -83,9 +83,9 @@ class KeyedPValueTrackingVisitor implements PipelineVisitor {
if (node.isRootNode()) {
finalized = true;
} else if (PRODUCES_KEYED_OUTPUTS.contains(node.getTransform().getClass())) {
- List<TaggedPValue> outputs = node.getOutputs();
- for (TaggedPValue output : outputs) {
- keyedValues.add(output.getValue());
+ Map<TupleTag<?>, PValue> outputs = node.getOutputs();
+ for (PValue output : outputs.values()) {
+ keyedValues.add(output);
}
}
}
@@ -96,8 +96,8 @@ class KeyedPValueTrackingVisitor implements PipelineVisitor {
@Override
public void visitValue(PValue value, TransformHierarchy.Node producer) {
boolean inputsAreKeyed = true;
- for (TaggedPValue input : producer.getInputs()) {
- inputsAreKeyed = inputsAreKeyed && keyedValues.contains(input.getValue());
+ for (PValue input : producer.getInputs().values()) {
+ inputsAreKeyed = inputsAreKeyed && keyedValues.contains(input);
}
if (PRODUCES_KEYED_OUTPUTS.contains(producer.getTransform().getClass())
|| (isKeyPreserving(producer.getTransform()) && inputsAreKeyed)) {
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 b0e97fb..b8a13e2 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
@@ -32,7 +32,7 @@ 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.PCollectionView;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TupleTag;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -141,7 +141,7 @@ final class ParDoEvaluatorFactory<InputT, OutputT> implements TransformEvaluator
evaluationContext,
stepContext,
application,
- ((PCollection<InputT>) Iterables.getOnlyElement(application.getInputs()).getValue())
+ ((PCollection<InputT>) Iterables.getOnlyElement(application.getInputs().values()))
.getWindowingStrategy(),
fn,
key,
@@ -162,10 +162,10 @@ final class ParDoEvaluatorFactory<InputT, OutputT> implements TransformEvaluator
}
}
- private Map<TupleTag<?>, PCollection<?>> pcollections(List<TaggedPValue> outputs) {
+ private Map<TupleTag<?>, PCollection<?>> pcollections(Map<TupleTag<?>, PValue> outputs) {
Map<TupleTag<?>, PCollection<?>> pcs = new HashMap<>();
- for (TaggedPValue output : outputs) {
- pcs.put(output.getTag(), (PCollection<?>) output.getValue());
+ for (Map.Entry<TupleTag<?>, PValue> output : outputs.entrySet()) {
+ pcs.put(output.getKey(), (PCollection<?>) output.getValue());
}
return pcs;
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 056a0c3..00c0d6a 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
@@ -20,7 +20,6 @@ package org.apache.beam.runners.direct;
import static com.google.common.base.Preconditions.checkState;
import com.google.common.collect.Iterables;
-import java.util.List;
import java.util.Map;
import org.apache.beam.runners.core.KeyedWorkItem;
import org.apache.beam.runners.core.KeyedWorkItemCoder;
@@ -50,7 +49,7 @@ 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.PValue;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
import org.apache.beam.sdk.values.TypedPValue;
@@ -86,13 +85,13 @@ class ParDoMultiOverrideFactory<InputT, OutputT>
@Override
public PCollection<? extends InputT> getInput(
- List<TaggedPValue> inputs, Pipeline p) {
- return (PCollection<? extends InputT>) Iterables.getOnlyElement(inputs).getValue();
+ Map<TupleTag<?>, PValue> inputs, Pipeline p) {
+ return (PCollection<? extends InputT>) Iterables.getOnlyElement(inputs.values());
}
@Override
public Map<PValue, ReplacementOutput> mapOutputs(
- List<TaggedPValue> outputs, PCollectionTuple newOutput) {
+ Map<TupleTag<?>, PValue> outputs, PCollectionTuple newOutput) {
return ReplacementOutputs.tagged(outputs, newOutput);
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 77bebb2..f8fe3d6 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
@@ -27,6 +27,7 @@ import com.google.common.collect.Lists;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
+import java.util.Map.Entry;
import org.apache.beam.runners.core.KeyedWorkItem;
import org.apache.beam.runners.core.KeyedWorkItems;
import org.apache.beam.runners.core.StateNamespace;
@@ -52,7 +53,7 @@ import org.apache.beam.sdk.util.state.StateSpec;
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.TaggedPValue;
+import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TupleTag;
/** A {@link TransformEvaluatorFactory} for stateful {@link ParDo}. */
@@ -139,8 +140,9 @@ final class StatefulParDoEvaluatorFactory<K, InputT, OutputT> implements Transfo
String stepName = evaluationContext.getStepName(transformOutputWindow.getTransform());
Map<TupleTag<?>, PCollection<?>> taggedValues = new HashMap<>();
- for (TaggedPValue pv : transformOutputWindow.getTransform().getOutputs()) {
- taggedValues.put(pv.getTag(), (PCollection<?>) pv.getValue());
+ for (Entry<TupleTag<?>, PValue> pv :
+ transformOutputWindow.getTransform().getOutputs().entrySet()) {
+ taggedValues.put(pv.getKey(), (PCollection<?>) pv.getValue());
}
PCollection<?> pc =
taggedValues
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 0dd8919..6e0a4fc 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
@@ -48,8 +48,8 @@ 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.PValue;
-import org.apache.beam.sdk.values.TaggedPValue;
import org.apache.beam.sdk.values.TimestampedValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.joda.time.Duration;
import org.joda.time.Instant;
@@ -108,7 +108,7 @@ class TestStreamEvaluatorFactory implements TransformEvaluatorFactory {
if (event.getType().equals(EventType.ELEMENT)) {
UncommittedBundle<T> bundle =
context.createBundle(
- (PCollection<T>) Iterables.getOnlyElement(application.getOutputs()).getValue());
+ (PCollection<T>) Iterables.getOnlyElement(application.getOutputs().values()));
for (TimestampedValue<T> elem : ((ElementEvent<T>) event).getElements()) {
bundle.add(
WindowedValue.timestampedValueInGlobalWindow(elem.getValue(), elem.getTimestamp()));
@@ -176,13 +176,13 @@ class TestStreamEvaluatorFactory implements TransformEvaluatorFactory {
}
@Override
- public PBegin getInput(List<TaggedPValue> inputs, Pipeline p) {
+ public PBegin getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
return p.begin();
}
@Override
public Map<PValue, ReplacementOutput> mapOutputs(
- List<TaggedPValue> outputs, PCollection<T> newOutput) {
+ Map<TupleTag<?>, PValue> outputs, PCollection<T> newOutput) {
return ReplacementOutputs.singleton(outputs, newOutput);
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/direct-java/src/main/java/org/apache/beam/runners/direct/UnboundedReadEvaluatorFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/UnboundedReadEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/UnboundedReadEvaluatorFactory.java
index f0eef58..91e7248 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/UnboundedReadEvaluatorFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/UnboundedReadEvaluatorFactory.java
@@ -121,7 +121,7 @@ class UnboundedReadEvaluatorFactory implements TransformEvaluatorFactory {
WindowedValue<UnboundedSourceShard<OutputT, CheckpointMarkT>> element) throws IOException {
UncommittedBundle<OutputT> output =
evaluationContext.createBundle(
- (PCollection<OutputT>) getOnlyElement(transform.getOutputs()).getValue());
+ (PCollection<OutputT>) getOnlyElement(transform.getOutputs().values()));
UnboundedSourceShard<OutputT, CheckpointMarkT> shard = element.getValue();
UnboundedReader<OutputT> reader = null;
try {
@@ -227,7 +227,7 @@ class UnboundedReadEvaluatorFactory implements TransformEvaluatorFactory {
// committing the output.
if (!reader.getWatermark().isBefore(BoundedWindow.TIMESTAMP_MAX_VALUE)) {
PCollection<OutputT> outputPc =
- (PCollection<OutputT>) Iterables.getOnlyElement(transform.getOutputs()).getValue();
+ (PCollection<OutputT>) Iterables.getOnlyElement(transform.getOutputs().values());
evaluationContext.scheduleAfterOutputWouldBeProduced(
outputPc,
GlobalWindow.INSTANCE,
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 dc74d3e..8cbe8fc 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
@@ -64,9 +64,9 @@ class ViewEvaluatorFactory implements TransformEvaluatorFactory {
final AppliedPTransform<PCollection<Iterable<InT>>, PCollectionView<OuT>, WriteView<InT, OuT>>
application) {
PCollection<Iterable<InT>> input =
- (PCollection<Iterable<InT>>) Iterables.getOnlyElement(application.getInputs()).getValue();
+ (PCollection<Iterable<InT>>) Iterables.getOnlyElement(application.getInputs().values());
final PCollectionViewWriter<InT, OuT> writer = context.createPCollectionViewWriter(input,
- (PCollectionView<OuT>) Iterables.getOnlyElement(application.getOutputs()).getValue());
+ (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/0e5737fd/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 64e1218..52dc329 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
@@ -20,7 +20,6 @@ package org.apache.beam.runners.direct;
import com.google.common.collect.Iterables;
import java.util.Collections;
-import java.util.List;
import java.util.Map;
import org.apache.beam.runners.core.construction.ForwardingPTransform;
import org.apache.beam.sdk.Pipeline;
@@ -35,7 +34,7 @@ import org.apache.beam.sdk.transforms.WithKeys;
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.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
/**
* A {@link PTransformOverrideFactory} that provides overrides for the {@link CreatePCollectionView}
@@ -51,13 +50,13 @@ class ViewOverrideFactory<ElemT, ViewT>
}
@Override
- public PCollection<ElemT> getInput(List<TaggedPValue> inputs, Pipeline p) {
- return (PCollection<ElemT>) Iterables.getOnlyElement(inputs).getValue();
+ public PCollection<ElemT> getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
+ return (PCollection<ElemT>) Iterables.getOnlyElement(inputs.values());
}
@Override
public Map<PValue, ReplacementOutput> mapOutputs(
- List<TaggedPValue> outputs, PCollectionView<ViewT> newOutput) {
+ Map<TupleTag<?>, PValue> outputs, PCollectionView<ViewT> newOutput) {
return Collections.emptyMap();
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 942e67c..8c04362 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
@@ -61,7 +61,8 @@ import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.TimeDomain;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.PValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.joda.time.Instant;
/**
@@ -818,13 +819,13 @@ public class WatermarkManager {
private Collection<Watermark> getInputProcessingWatermarks(AppliedPTransform<?, ?, ?> transform) {
ImmutableList.Builder<Watermark> inputWmsBuilder = ImmutableList.builder();
- List<TaggedPValue> inputs = transform.getInputs();
+ Map<TupleTag<?>, PValue> inputs = transform.getInputs();
if (inputs.isEmpty()) {
inputWmsBuilder.add(THE_END_OF_TIME);
}
- for (TaggedPValue pvalue : inputs) {
+ for (PValue pvalue : inputs.values()) {
Watermark producerOutputWatermark =
- getTransformWatermark(graph.getProducer(pvalue.getValue()))
+ getTransformWatermark(graph.getProducer(pvalue))
.synchronizedProcessingOutputWatermark;
inputWmsBuilder.add(producerOutputWatermark);
}
@@ -833,13 +834,13 @@ public class WatermarkManager {
private List<Watermark> getInputWatermarks(AppliedPTransform<?, ?, ?> transform) {
ImmutableList.Builder<Watermark> inputWatermarksBuilder = ImmutableList.builder();
- List<TaggedPValue> inputs = transform.getInputs();
+ Map<TupleTag<?>, PValue> inputs = transform.getInputs();
if (inputs.isEmpty()) {
inputWatermarksBuilder.add(THE_END_OF_TIME);
}
- for (TaggedPValue pvalue : inputs) {
+ for (PValue pvalue : inputs.values()) {
Watermark producerOutputWatermark =
- getTransformWatermark(graph.getProducer(pvalue.getValue())).outputWatermark;
+ getTransformWatermark(graph.getProducer(pvalue)).outputWatermark;
inputWatermarksBuilder.add(producerOutputWatermark);
}
List<Watermark> inputCollectionWatermarks = inputWatermarksBuilder.build();
@@ -1023,8 +1024,8 @@ public class WatermarkManager {
WatermarkUpdate updateResult = myWatermarks.refresh();
if (updateResult.isAdvanced()) {
Set<AppliedPTransform<?, ?, ?>> additionalRefreshes = new HashSet<>();
- for (TaggedPValue outputPValue : toRefresh.getOutputs()) {
- additionalRefreshes.addAll(graph.getPrimitiveConsumers(outputPValue.getValue()));
+ for (PValue outputPValue : toRefresh.getOutputs().values()) {
+ additionalRefreshes.addAll(graph.getPrimitiveConsumers(outputPValue));
}
return additionalRefreshes;
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WindowEvaluatorFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WindowEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WindowEvaluatorFactory.java
index 8974c67..2550924 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WindowEvaluatorFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/WindowEvaluatorFactory.java
@@ -57,7 +57,7 @@ class WindowEvaluatorFactory implements TransformEvaluatorFactory {
WindowFn<? super InputT, ?> fn = transform.getTransform().getWindowFn();
UncommittedBundle<InputT> outputBundle =
evaluationContext.createBundle(
- (PCollection<InputT>) Iterables.getOnlyElement(transform.getOutputs()).getValue());
+ (PCollection<InputT>) Iterables.getOnlyElement(transform.getOutputs().values()));
if (fn == null) {
return PassthroughTransformEvaluator.create(transform, outputBundle);
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 1bf5839..b3f92ab 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
@@ -24,7 +24,6 @@ import com.google.common.base.Suppliers;
import com.google.common.collect.Iterables;
import java.io.Serializable;
import java.util.Collections;
-import java.util.List;
import java.util.Map;
import java.util.concurrent.ThreadLocalRandom;
import org.apache.beam.sdk.Pipeline;
@@ -41,7 +40,7 @@ import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PDone;
import org.apache.beam.sdk.values.PValue;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
/**
* A {@link PTransformOverrideFactory} that overrides {@link Write} {@link PTransform PTransforms}
@@ -60,12 +59,13 @@ class WriteWithShardingFactory<InputT>
}
@Override
- public PCollection<InputT> getInput(List<TaggedPValue> inputs, Pipeline p) {
- return (PCollection<InputT>) Iterables.getOnlyElement(inputs).getValue();
+ public PCollection<InputT> getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
+ return (PCollection<InputT>) Iterables.getOnlyElement(inputs.values());
}
@Override
- public Map<PValue, ReplacementOutput> mapOutputs(List<TaggedPValue> outputs, PDone newOutput) {
+ public Map<PValue, ReplacementOutput> mapOutputs(
+ Map<TupleTag<?>, PValue> outputs, PDone newOutput) {
return Collections.emptyMap();
}
[34/50] [abbrv] beam git commit: [BEAM-1722] Move PubsubIO into the
google-cloud-platform module
Posted by jb...@apache.org.
[BEAM-1722] Move PubsubIO into the google-cloud-platform module
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/5bcb8c57
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/5bcb8c57
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/5bcb8c57
Branch: refs/heads/DSL_SQL
Commit: 5bcb8c579656bc2d6e4d2a8dd5dcb2a46875812f
Parents: 82f2f2c
Author: Isma�l Mej�a <ie...@apache.org>
Authored: Wed Apr 12 14:49:58 2017 +0200
Committer: Isma�l Mej�a <ie...@apache.org>
Committed: Wed Apr 12 14:49:58 2017 +0200
----------------------------------------------------------------------
.../beam/examples/complete/game/GameStats.java | 2 +-
.../examples/complete/game/LeaderBoard.java | 2 +-
.../triggers/AfterWatermarkStateMachine.java | 14 +-
runners/google-cloud-dataflow-java/pom.xml | 5 +-
.../beam/runners/dataflow/DataflowRunner.java | 4 +-
sdks/java/core/pom.xml | 44 -
.../java/org/apache/beam/sdk/io/PubsubIO.java | 1016 ------------
.../apache/beam/sdk/io/PubsubUnboundedSink.java | 494 ------
.../beam/sdk/io/PubsubUnboundedSource.java | 1463 ------------------
.../apache/beam/sdk/transforms/GroupByKey.java | 4 +-
.../transforms/windowing/AfterWatermark.java | 14 +-
.../org/apache/beam/sdk/util/PubsubClient.java | 544 -------
.../apache/beam/sdk/util/PubsubGrpcClient.java | 424 -----
.../apache/beam/sdk/util/PubsubJsonClient.java | 317 ----
.../apache/beam/sdk/util/PubsubTestClient.java | 436 ------
.../org/apache/beam/sdk/util/Transport.java | 3 +-
.../org/apache/beam/sdk/io/PubsubIOTest.java | 189 ---
.../beam/sdk/io/PubsubUnboundedSinkTest.java | 190 ---
.../beam/sdk/io/PubsubUnboundedSourceTest.java | 411 -----
.../apache/beam/sdk/util/PubsubClientTest.java | 189 ---
.../beam/sdk/util/PubsubGrpcClientTest.java | 207 ---
.../beam/sdk/util/PubsubJsonClientTest.java | 140 --
.../beam/sdk/util/PubsubTestClientTest.java | 114 --
sdks/java/io/google-cloud-platform/pom.xml | 44 +
.../beam/sdk/io/gcp/pubsub/PubsubClient.java | 544 +++++++
.../sdk/io/gcp/pubsub/PubsubGrpcClient.java | 424 +++++
.../apache/beam/sdk/io/gcp/pubsub/PubsubIO.java | 1014 ++++++++++++
.../sdk/io/gcp/pubsub/PubsubJsonClient.java | 319 ++++
.../sdk/io/gcp/pubsub/PubsubTestClient.java | 436 ++++++
.../sdk/io/gcp/pubsub/PubsubUnboundedSink.java | 490 ++++++
.../io/gcp/pubsub/PubsubUnboundedSource.java | 1463 ++++++++++++++++++
.../beam/sdk/io/gcp/pubsub/package-info.java | 24 +
.../beam/sdk/io/gcp/GcpApiSurfaceTest.java | 5 +-
.../sdk/io/gcp/pubsub/PubsubClientTest.java | 189 +++
.../sdk/io/gcp/pubsub/PubsubGrpcClientTest.java | 208 +++
.../beam/sdk/io/gcp/pubsub/PubsubIOTest.java | 189 +++
.../sdk/io/gcp/pubsub/PubsubJsonClientTest.java | 139 ++
.../sdk/io/gcp/pubsub/PubsubTestClientTest.java | 114 ++
.../io/gcp/pubsub/PubsubUnboundedSinkTest.java | 188 +++
.../gcp/pubsub/PubsubUnboundedSourceTest.java | 409 +++++
40 files changed, 6218 insertions(+), 6207 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/examples/java8/src/main/java/org/apache/beam/examples/complete/game/GameStats.java
----------------------------------------------------------------------
diff --git a/examples/java8/src/main/java/org/apache/beam/examples/complete/game/GameStats.java b/examples/java8/src/main/java/org/apache/beam/examples/complete/game/GameStats.java
index 93e8254..6874953 100644
--- a/examples/java8/src/main/java/org/apache/beam/examples/complete/game/GameStats.java
+++ b/examples/java8/src/main/java/org/apache/beam/examples/complete/game/GameStats.java
@@ -25,7 +25,7 @@ import org.apache.beam.examples.complete.game.utils.WriteWindowedToBigQuery;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.sdk.io.PubsubIO;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubIO;
import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.Description;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/examples/java8/src/main/java/org/apache/beam/examples/complete/game/LeaderBoard.java
----------------------------------------------------------------------
diff --git a/examples/java8/src/main/java/org/apache/beam/examples/complete/game/LeaderBoard.java b/examples/java8/src/main/java/org/apache/beam/examples/complete/game/LeaderBoard.java
index 35b586b..96f4291 100644
--- a/examples/java8/src/main/java/org/apache/beam/examples/complete/game/LeaderBoard.java
+++ b/examples/java8/src/main/java/org/apache/beam/examples/complete/game/LeaderBoard.java
@@ -28,7 +28,7 @@ import org.apache.beam.examples.complete.game.utils.WriteWindowedToBigQuery;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.sdk.io.PubsubIO;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubIO;
import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.Description;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/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 0b12005..1b117d2 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
@@ -31,18 +31,16 @@ import org.apache.beam.sdk.util.TimeDomain;
* lower-bound, sometimes heuristically established, on event times that have been fully processed
* by the pipeline.
*
- * <p>For sources that provide non-heuristic watermarks (e.g.
- * {@link org.apache.beam.sdk.io.PubsubIO} when using arrival times as event times), the
- * watermark is a strict guarantee that no data with an event time earlier than
+ * <p>For sources that provide non-heuristic watermarks (e.g. PubsubIO when using arrival times as
+ * event times), the watermark is a strict guarantee that no data with an event time earlier than
* that watermark will ever be observed in the pipeline. In this case, it's safe to assume that any
* pane triggered by an {@code AfterWatermark} trigger with a reference point at or beyond the end
* of the window will be the last pane ever for that window.
*
- * <p>For sources that provide heuristic watermarks (e.g.
- * {@link org.apache.beam.sdk.io.PubsubIO} when using user-supplied event times), the
- * watermark itself becomes an <i>estimate</i> that no data with an event time earlier than that
- * watermark (i.e. "late data") will ever be observed in the pipeline. These heuristics can
- * often be quite accurate, but the chance of seeing late data for any given window is non-zero.
+ * <p>For sources that provide heuristic watermarks (e.g. PubsubIO when using user-supplied event
+ * times), the watermark itself becomes an <i>estimate</i> that no data with an event time earlier
+ * than that watermark (i.e. "late data") will ever be observed in the pipeline. These heuristics
+ * can often be quite accurate, but the chance of seeing late data for any given window is non-zero.
* Thus, if absolute correctness over time is important to your use case, you may want to consider
* using a trigger that accounts for late data. The default trigger,
* {@code Repeatedly.forever(AfterWatermark.pastEndOfWindow())}, which fires
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/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 a57744c..96eced8 100644
--- a/runners/google-cloud-dataflow-java/pom.xml
+++ b/runners/google-cloud-dataflow-java/pom.xml
@@ -188,13 +188,12 @@
<dependency>
<groupId>org.apache.beam</groupId>
- <artifactId>beam-runners-core-construction-java</artifactId>
+ <artifactId>beam-sdks-java-io-google-cloud-platform</artifactId>
</dependency>
<dependency>
<groupId>org.apache.beam</groupId>
- <artifactId>beam-sdks-java-io-google-cloud-platform</artifactId>
- <scope>test</scope>
+ <artifactId>beam-runners-core-construction-java</artifactId>
</dependency>
<dependency>
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/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 8726635..684dc14 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
@@ -82,11 +82,11 @@ import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.io.BoundedSource;
import org.apache.beam.sdk.io.FileBasedSink;
-import org.apache.beam.sdk.io.PubsubUnboundedSink;
-import org.apache.beam.sdk.io.PubsubUnboundedSource;
import org.apache.beam.sdk.io.Read;
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.io.Write;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSink;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsValidator;
import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider;
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/pom.xml
----------------------------------------------------------------------
diff --git a/sdks/java/core/pom.xml b/sdks/java/core/pom.xml
index d117d5a..0ac40f4 100644
--- a/sdks/java/core/pom.xml
+++ b/sdks/java/core/pom.xml
@@ -140,40 +140,6 @@
</dependency>
<dependency>
- <groupId>io.grpc</groupId>
- <artifactId>grpc-auth</artifactId>
- </dependency>
-
- <dependency>
- <groupId>io.grpc</groupId>
- <artifactId>grpc-core</artifactId>
- </dependency>
-
- <dependency>
- <groupId>io.grpc</groupId>
- <artifactId>grpc-netty</artifactId>
- </dependency>
-
- <dependency>
- <groupId>io.grpc</groupId>
- <artifactId>grpc-stub</artifactId>
- </dependency>
-
- <!-- grpc-all does not obey IWYU, so we need to exclude from compile
- scope and depend on it at runtime. -->
- <dependency>
- <groupId>io.grpc</groupId>
- <artifactId>grpc-all</artifactId>
- <scope>runtime</scope>
- </dependency>
-
- <dependency>
- <groupId>io.grpc</groupId>
- <artifactId>grpc-protobuf</artifactId>
- <scope>runtime</scope>
- </dependency>
-
- <dependency>
<groupId>com.google.auth</groupId>
<artifactId>google-auth-library-credentials</artifactId>
</dependency>
@@ -184,16 +150,6 @@
</dependency>
<dependency>
- <groupId>io.netty</groupId>
- <artifactId>netty-handler</artifactId>
- </dependency>
-
- <dependency>
- <groupId>com.google.api.grpc</groupId>
- <artifactId>grpc-google-pubsub-v1</artifactId>
- </dependency>
-
- <dependency>
<groupId>com.google.api-client</groupId>
<artifactId>google-api-client</artifactId>
</dependency>
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/src/main/java/org/apache/beam/sdk/io/PubsubIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/PubsubIO.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/PubsubIO.java
deleted file mode 100644
index 67ab2ec..0000000
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/PubsubIO.java
+++ /dev/null
@@ -1,1016 +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 static com.google.common.base.Preconditions.checkNotNull;
-import static com.google.common.base.Preconditions.checkState;
-
-import java.io.IOException;
-import java.io.Serializable;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
-import javax.annotation.Nullable;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.VoidCoder;
-import org.apache.beam.sdk.options.PubsubOptions;
-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.runners.PipelineRunner;
-import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.transforms.SerializableFunction;
-import org.apache.beam.sdk.transforms.SimpleFunction;
-import org.apache.beam.sdk.transforms.display.DisplayData;
-import org.apache.beam.sdk.transforms.windowing.AfterWatermark;
-import org.apache.beam.sdk.util.CoderUtils;
-import org.apache.beam.sdk.util.PubsubClient;
-import org.apache.beam.sdk.util.PubsubClient.OutgoingMessage;
-import org.apache.beam.sdk.util.PubsubClient.ProjectPath;
-import org.apache.beam.sdk.util.PubsubClient.SubscriptionPath;
-import org.apache.beam.sdk.util.PubsubClient.TopicPath;
-import org.apache.beam.sdk.util.PubsubJsonClient;
-import org.apache.beam.sdk.values.PBegin;
-import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PDone;
-import org.joda.time.Instant;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/**
- * Read and Write {@link PTransform}s for Cloud Pub/Sub streams. These transforms create
- * and consume unbounded {@link PCollection PCollections}.
- *
- * <h3>Permissions</h3>
- *
- * <p>Permission requirements depend on the {@link PipelineRunner} that is used to execute the
- * Beam pipeline. Please refer to the documentation of corresponding
- * {@link PipelineRunner PipelineRunners} for more details.
- */
-public class PubsubIO {
-
- private static final Logger LOG = LoggerFactory.getLogger(PubsubIO.class);
-
- /** Factory for creating pubsub client to manage transport. */
- private static final PubsubClient.PubsubClientFactory FACTORY = PubsubJsonClient.FACTORY;
-
- /**
- * Project IDs must contain 6-63 lowercase letters, digits, or dashes.
- * IDs must start with a letter and may not end with a dash.
- * This regex isn't exact - this allows for patterns that would be rejected by
- * the service, but this is sufficient for basic parsing of table references.
- */
- private static final Pattern PROJECT_ID_REGEXP =
- Pattern.compile("[a-z][-a-z0-9:.]{4,61}[a-z0-9]");
-
- private static final Pattern SUBSCRIPTION_REGEXP =
- Pattern.compile("projects/([^/]+)/subscriptions/(.+)");
-
- private static final Pattern TOPIC_REGEXP = Pattern.compile("projects/([^/]+)/topics/(.+)");
-
- private static final Pattern V1BETA1_SUBSCRIPTION_REGEXP =
- Pattern.compile("/subscriptions/([^/]+)/(.+)");
-
- private static final Pattern V1BETA1_TOPIC_REGEXP = Pattern.compile("/topics/([^/]+)/(.+)");
-
- private static final Pattern PUBSUB_NAME_REGEXP = Pattern.compile("[a-zA-Z][-._~%+a-zA-Z0-9]+");
-
- private static final int PUBSUB_NAME_MIN_LENGTH = 3;
- private static final int PUBSUB_NAME_MAX_LENGTH = 255;
-
- private static final String SUBSCRIPTION_RANDOM_TEST_PREFIX = "_random/";
- private static final String SUBSCRIPTION_STARTING_SIGNAL = "_starting_signal/";
- private static final String TOPIC_DEV_NULL_TEST_NAME = "/topics/dev/null";
-
- private static void validateProjectName(String project) {
- Matcher match = PROJECT_ID_REGEXP.matcher(project);
- if (!match.matches()) {
- throw new IllegalArgumentException(
- "Illegal project name specified in Pubsub subscription: " + project);
- }
- }
-
- private static void validatePubsubName(String name) {
- if (name.length() < PUBSUB_NAME_MIN_LENGTH) {
- throw new IllegalArgumentException(
- "Pubsub object name is shorter than 3 characters: " + name);
- }
- if (name.length() > PUBSUB_NAME_MAX_LENGTH) {
- throw new IllegalArgumentException(
- "Pubsub object name is longer than 255 characters: " + name);
- }
-
- if (name.startsWith("goog")) {
- throw new IllegalArgumentException("Pubsub object name cannot start with goog: " + name);
- }
-
- Matcher match = PUBSUB_NAME_REGEXP.matcher(name);
- if (!match.matches()) {
- throw new IllegalArgumentException("Illegal Pubsub object name specified: " + name
- + " Please see Javadoc for naming rules.");
- }
- }
-
- /**
- * Populate common {@link DisplayData} between Pubsub source and sink.
- */
- private static void populateCommonDisplayData(DisplayData.Builder builder,
- String timestampLabel, String idLabel, ValueProvider<PubsubTopic> topic) {
- builder
- .addIfNotNull(DisplayData.item("timestampLabel", timestampLabel)
- .withLabel("Timestamp Label Attribute"))
- .addIfNotNull(DisplayData.item("idLabel", idLabel)
- .withLabel("ID Label Attribute"));
-
- if (topic != null) {
- String topicString = topic.isAccessible() ? topic.get().asPath()
- : topic.toString();
- builder.add(DisplayData.item("topic", topicString)
- .withLabel("Pubsub Topic"));
- }
- }
-
- /**
- * Class representing a Pub/Sub message. Each message contains a single message payload and
- * a map of attached attributes.
- */
- public static class PubsubMessage {
-
- private byte[] message;
- private Map<String, String> attributes;
-
- public PubsubMessage(byte[] message, Map<String, String> attributes) {
- this.message = message;
- this.attributes = attributes;
- }
-
- /**
- * Returns the main PubSub message.
- */
- public byte[] getMessage() {
- return message;
- }
-
- /**
- * Returns the given attribute value. If not such attribute exists, returns null.
- */
- @Nullable
- public String getAttribute(String attribute) {
- checkNotNull(attribute, "attribute");
- return attributes.get(attribute);
- }
-
- /**
- * Returns the full map of attributes. This is an unmodifiable map.
- */
- public Map<String, String> getAttributeMap() {
- return attributes;
- }
- }
-
- /**
- * Class representing a Cloud Pub/Sub Subscription.
- */
- public static class PubsubSubscription implements Serializable {
-
- private enum Type {NORMAL, FAKE}
-
- private final Type type;
- private final String project;
- private final String subscription;
-
- private PubsubSubscription(Type type, String project, String subscription) {
- this.type = type;
- this.project = project;
- this.subscription = subscription;
- }
-
- /**
- * Creates a class representing a Pub/Sub subscription from the specified subscription path.
- *
- * <p>Cloud Pub/Sub subscription names should be of the form
- * {@code projects/<project>/subscriptions/<subscription>}, where {@code <project>} is the name
- * of the project the subscription belongs to. The {@code <subscription>} component must comply
- * with the following requirements:
- *
- * <ul>
- * <li>Can only contain lowercase letters, numbers, dashes ('-'), underscores ('_') and periods
- * ('.').</li>
- * <li>Must be between 3 and 255 characters.</li>
- * <li>Must begin with a letter.</li>
- * <li>Must end with a letter or a number.</li>
- * <li>Cannot begin with {@code 'goog'} prefix.</li>
- * </ul>
- */
- public static PubsubSubscription fromPath(String path) {
- if (path.startsWith(SUBSCRIPTION_RANDOM_TEST_PREFIX)
- || path.startsWith(SUBSCRIPTION_STARTING_SIGNAL)) {
- return new PubsubSubscription(Type.FAKE, "", path);
- }
-
- String projectName, subscriptionName;
-
- Matcher v1beta1Match = V1BETA1_SUBSCRIPTION_REGEXP.matcher(path);
- if (v1beta1Match.matches()) {
- LOG.warn("Saw subscription in v1beta1 format. Subscriptions should be in the format "
- + "projects/<project_id>/subscriptions/<subscription_name>");
- projectName = v1beta1Match.group(1);
- subscriptionName = v1beta1Match.group(2);
- } else {
- Matcher match = SUBSCRIPTION_REGEXP.matcher(path);
- if (!match.matches()) {
- throw new IllegalArgumentException("Pubsub subscription is not in "
- + "projects/<project_id>/subscriptions/<subscription_name> format: " + path);
- }
- projectName = match.group(1);
- subscriptionName = match.group(2);
- }
-
- validateProjectName(projectName);
- validatePubsubName(subscriptionName);
- return new PubsubSubscription(Type.NORMAL, projectName, subscriptionName);
- }
-
- /**
- * Returns the string representation of this subscription as a path used in the Cloud Pub/Sub
- * v1beta1 API.
- *
- * @deprecated the v1beta1 API for Cloud Pub/Sub is deprecated.
- */
- @Deprecated
- public String asV1Beta1Path() {
- if (type == Type.NORMAL) {
- return "/subscriptions/" + project + "/" + subscription;
- } else {
- return subscription;
- }
- }
-
- /**
- * Returns the string representation of this subscription as a path used in the Cloud Pub/Sub
- * v1beta2 API.
- *
- * @deprecated the v1beta2 API for Cloud Pub/Sub is deprecated.
- */
- @Deprecated
- public String asV1Beta2Path() {
- if (type == Type.NORMAL) {
- return "projects/" + project + "/subscriptions/" + subscription;
- } else {
- return subscription;
- }
- }
-
- /**
- * Returns the string representation of this subscription as a path used in the Cloud Pub/Sub
- * API.
- */
- public String asPath() {
- if (type == Type.NORMAL) {
- return "projects/" + project + "/subscriptions/" + subscription;
- } else {
- return subscription;
- }
- }
- }
-
- /**
- * Used to build a {@link ValueProvider} for {@link PubsubSubscription}.
- */
- private static class SubscriptionTranslator
- implements SerializableFunction<String, PubsubSubscription> {
-
- @Override
- public PubsubSubscription apply(String from) {
- return PubsubSubscription.fromPath(from);
- }
- }
-
- /**
- * Used to build a {@link ValueProvider} for {@link SubscriptionPath}.
- */
- private static class SubscriptionPathTranslator
- implements SerializableFunction<PubsubSubscription, SubscriptionPath> {
-
- @Override
- public SubscriptionPath apply(PubsubSubscription from) {
- return PubsubClient.subscriptionPathFromName(from.project, from.subscription);
- }
- }
-
- /**
- * Used to build a {@link ValueProvider} for {@link PubsubTopic}.
- */
- private static class TopicTranslator
- implements SerializableFunction<String, PubsubTopic> {
-
- @Override
- public PubsubTopic apply(String from) {
- return PubsubTopic.fromPath(from);
- }
- }
-
- /**
- * Used to build a {@link ValueProvider} for {@link TopicPath}.
- */
- private static class TopicPathTranslator
- implements SerializableFunction<PubsubTopic, TopicPath> {
-
- @Override
- public TopicPath apply(PubsubTopic from) {
- return PubsubClient.topicPathFromName(from.project, from.topic);
- }
- }
-
- /**
- * Used to build a {@link ValueProvider} for {@link ProjectPath}.
- */
- private static class ProjectPathTranslator
- implements SerializableFunction<PubsubTopic, ProjectPath> {
-
- @Override
- public ProjectPath apply(PubsubTopic from) {
- return PubsubClient.projectPathFromId(from.project);
- }
- }
-
- /**
- * Class representing a Cloud Pub/Sub Topic.
- */
- public static class PubsubTopic implements Serializable {
-
- private enum Type {NORMAL, FAKE}
-
- private final Type type;
- private final String project;
- private final String topic;
-
- private PubsubTopic(Type type, String project, String topic) {
- this.type = type;
- this.project = project;
- this.topic = topic;
- }
-
- /**
- * Creates a class representing a Cloud Pub/Sub topic from the specified topic path.
- *
- * <p>Cloud Pub/Sub topic names should be of the form
- * {@code /topics/<project>/<topic>}, where {@code <project>} is the name of
- * the publishing project. The {@code <topic>} component must comply with
- * the following requirements:
- *
- * <ul>
- * <li>Can only contain lowercase letters, numbers, dashes ('-'), underscores ('_') and periods
- * ('.').</li>
- * <li>Must be between 3 and 255 characters.</li>
- * <li>Must begin with a letter.</li>
- * <li>Must end with a letter or a number.</li>
- * <li>Cannot begin with 'goog' prefix.</li>
- * </ul>
- */
- public static PubsubTopic fromPath(String path) {
- if (path.equals(TOPIC_DEV_NULL_TEST_NAME)) {
- return new PubsubTopic(Type.FAKE, "", path);
- }
-
- String projectName, topicName;
-
- Matcher v1beta1Match = V1BETA1_TOPIC_REGEXP.matcher(path);
- if (v1beta1Match.matches()) {
- LOG.warn("Saw topic in v1beta1 format. Topics should be in the format "
- + "projects/<project_id>/topics/<topic_name>");
- projectName = v1beta1Match.group(1);
- topicName = v1beta1Match.group(2);
- } else {
- Matcher match = TOPIC_REGEXP.matcher(path);
- if (!match.matches()) {
- throw new IllegalArgumentException(
- "Pubsub topic is not in projects/<project_id>/topics/<topic_name> format: " + path);
- }
- projectName = match.group(1);
- topicName = match.group(2);
- }
-
- validateProjectName(projectName);
- validatePubsubName(topicName);
- return new PubsubTopic(Type.NORMAL, projectName, topicName);
- }
-
- /**
- * Returns the string representation of this topic as a path used in the Cloud Pub/Sub
- * v1beta1 API.
- *
- * @deprecated the v1beta1 API for Cloud Pub/Sub is deprecated.
- */
- @Deprecated
- public String asV1Beta1Path() {
- if (type == Type.NORMAL) {
- return "/topics/" + project + "/" + topic;
- } else {
- return topic;
- }
- }
-
- /**
- * Returns the string representation of this topic as a path used in the Cloud Pub/Sub
- * v1beta2 API.
- *
- * @deprecated the v1beta2 API for Cloud Pub/Sub is deprecated.
- */
- @Deprecated
- public String asV1Beta2Path() {
- if (type == Type.NORMAL) {
- return "projects/" + project + "/topics/" + topic;
- } else {
- return topic;
- }
- }
-
- /**
- * Returns the string representation of this topic as a path used in the Cloud Pub/Sub
- * API.
- */
- public String asPath() {
- if (type == Type.NORMAL) {
- return "projects/" + project + "/topics/" + topic;
- } else {
- return topic;
- }
- }
- }
-
- public static <T> Read<T> read() {
- return new Read<>();
- }
-
- public static <T> Write<T> write() {
- return new Write<>();
- }
-
- /**
- * A {@link PTransform} that continuously reads from a Google Cloud Pub/Sub stream and
- * returns a {@link PCollection} of {@link String Strings} containing the items from
- * the stream.
- */
- public static class Read<T> extends PTransform<PBegin, PCollection<T>> {
-
- /** The Cloud Pub/Sub topic to read from. */
- @Nullable
- private final ValueProvider<PubsubTopic> topic;
-
- /** The Cloud Pub/Sub subscription to read from. */
- @Nullable
- private final ValueProvider<PubsubSubscription> subscription;
-
- /** The name of the message attribute to read timestamps from. */
- @Nullable
- private final String timestampLabel;
-
- /** The name of the message attribute to read unique message IDs from. */
- @Nullable
- private final String idLabel;
-
- /** The coder used to decode each record. */
- @Nullable
- private final Coder<T> coder;
-
- /** User function for parsing PubsubMessage object. */
- SimpleFunction<PubsubMessage, T> parseFn;
-
- private Read() {
- this(null, null, null, null, null, null, null);
- }
-
- private Read(String name, ValueProvider<PubsubSubscription> subscription,
- ValueProvider<PubsubTopic> topic, String timestampLabel, Coder<T> coder,
- String idLabel,
- SimpleFunction<PubsubMessage, T> parseFn) {
- super(name);
- this.subscription = subscription;
- this.topic = topic;
- this.timestampLabel = timestampLabel;
- this.coder = coder;
- this.idLabel = idLabel;
- this.parseFn = parseFn;
- }
-
- /**
- * Returns a transform that's like this one but reading from the
- * given subscription.
- *
- * <p>See {@link PubsubIO.PubsubSubscription#fromPath(String)} for more details on the format
- * of the {@code subscription} string.
- *
- * <p>Multiple readers reading from the same subscription will each receive
- * some arbitrary portion of the data. Most likely, separate readers should
- * use their own subscriptions.
- *
- * <p>Does not modify this object.
- */
- public Read<T> subscription(String subscription) {
- return subscription(StaticValueProvider.of(subscription));
- }
-
- /**
- * Like {@code subscription()} but with a {@link ValueProvider}.
- */
- public Read<T> subscription(ValueProvider<String> subscription) {
- if (subscription.isAccessible()) {
- // Validate.
- PubsubSubscription.fromPath(subscription.get());
- }
- return new Read<>(
- name, NestedValueProvider.of(subscription, new SubscriptionTranslator()),
- null /* reset topic to null */, timestampLabel, coder, idLabel, parseFn);
- }
-
- /**
- * Creates and returns a transform for reading from a Cloud Pub/Sub topic. Mutually exclusive
- * with {@link #subscription(String)}.
- *
- * <p>See {@link PubsubIO.PubsubTopic#fromPath(String)} for more details on the format
- * of the {@code topic} string.
- *
- * <p>The Beam runner will start reading data published on this topic from the time the
- * pipeline is started. Any data published on the topic before the pipeline is started will
- * not be read by the runner.
- */
- public Read<T> topic(String topic) {
- return topic(StaticValueProvider.of(topic));
- }
-
- /**
- * Like {@code topic()} but with a {@link ValueProvider}.
- */
- public Read<T> topic(ValueProvider<String> topic) {
- if (topic.isAccessible()) {
- // Validate.
- PubsubTopic.fromPath(topic.get());
- }
- return new Read<>(name, null /* reset subscription to null */,
- NestedValueProvider.of(topic, new TopicTranslator()),
- timestampLabel, coder, idLabel, parseFn);
- }
-
- /**
- * Creates and returns a transform reading from Cloud Pub/Sub where record timestamps are
- * expected to be provided as Pub/Sub message attributes. The {@code timestampLabel}
- * parameter specifies the name of the attribute that contains the timestamp.
- *
- * <p>The timestamp value is expected to be represented in the attribute as either:
- *
- * <ul>
- * <li>a numerical value representing the number of milliseconds since the Unix epoch. For
- * example, if using the Joda time classes, {@link Instant#getMillis()} returns the correct
- * value for this attribute.
- * <li>a String in RFC 3339 format. For example, {@code 2015-10-29T23:41:41.123Z}. The
- * sub-second component of the timestamp is optional, and digits beyond the first three
- * (i.e., time units smaller than milliseconds) will be ignored.
- * </ul>
- *
- * <p>If {@code timestampLabel} is not provided, the system will generate record timestamps
- * the first time it sees each record. All windowing will be done relative to these
- * timestamps.
- *
- * <p>By default, windows are emitted based on an estimate of when this source is likely
- * done producing data for a given timestamp (referred to as the Watermark; see
- * {@link AfterWatermark} for more details). Any late data will be handled by the trigger
- * specified with the windowing strategy – by default it will be output immediately.
- *
- * <p>Note that the system can guarantee that no late data will ever be seen when it assigns
- * timestamps by arrival time (i.e. {@code timestampLabel} is not provided).
- *
- * @see <a href="https://www.ietf.org/rfc/rfc3339.txt">RFC 3339</a>
- */
- public Read<T> timestampLabel(String timestampLabel) {
- return new Read<>(
- name, subscription, topic, timestampLabel, coder, idLabel,
- parseFn);
- }
-
- /**
- * Creates and returns a transform for reading from Cloud Pub/Sub where unique record
- * identifiers are expected to be provided as Pub/Sub message attributes. The {@code idLabel}
- * parameter specifies the attribute name. The value of the attribute can be any string
- * that uniquely identifies this record.
- *
- * <p>Pub/Sub cannot guarantee that no duplicate data will be delivered on the Pub/Sub stream.
- * If {@code idLabel} is not provided, Beam cannot guarantee that no duplicate data will
- * be delivered, and deduplication of the stream will be strictly best effort.
- */
- public Read<T> idLabel(String idLabel) {
- return new Read<>(
- name, subscription, topic, timestampLabel, coder, idLabel,
- parseFn);
- }
-
- /**
- * Returns a transform that's like this one but that uses the given
- * {@link Coder} to decode each record into a value of type {@code T}.
- *
- * <p>Does not modify this object.
- */
- public Read<T> withCoder(Coder<T> coder) {
- return new Read<>(
- name, subscription, topic, timestampLabel, coder, idLabel,
- parseFn);
- }
-
- /**
- * Causes the source to return a PubsubMessage that includes Pubsub attributes.
- * The user must supply a parsing function to transform the PubsubMessage into an output type.
- * A Coder for the output type T must be registered or set on the output via
- * {@link PCollection#setCoder(Coder)}.
- */
- public Read<T> withAttributes(SimpleFunction<PubsubMessage, T> parseFn) {
- return new Read<T>(
- name, subscription, topic, timestampLabel, coder, idLabel,
- parseFn);
- }
-
- @Override
- public PCollection<T> expand(PBegin input) {
- if (topic == null && subscription == null) {
- throw new IllegalStateException("Need to set either the topic or the subscription for "
- + "a PubsubIO.Read transform");
- }
- if (topic != null && subscription != null) {
- throw new IllegalStateException("Can't set both the topic and the subscription for "
- + "a PubsubIO.Read transform");
- }
- if (coder == null) {
- throw new IllegalStateException("PubsubIO.Read requires that a coder be set using "
- + "the withCoder method.");
- }
-
- @Nullable ValueProvider<ProjectPath> projectPath =
- topic == null ? null : NestedValueProvider.of(topic, new ProjectPathTranslator());
- @Nullable ValueProvider<TopicPath> topicPath =
- topic == null ? null : NestedValueProvider.of(topic, new TopicPathTranslator());
- @Nullable ValueProvider<SubscriptionPath> subscriptionPath =
- subscription == null
- ? null
- : NestedValueProvider.of(subscription, new SubscriptionPathTranslator());
- PubsubUnboundedSource<T> source = new PubsubUnboundedSource<T>(
- FACTORY, projectPath, topicPath, subscriptionPath,
- coder, timestampLabel, idLabel, parseFn);
- return input.getPipeline().apply(source);
- }
-
- @Override
- public void populateDisplayData(DisplayData.Builder builder) {
- super.populateDisplayData(builder);
- populateCommonDisplayData(builder, timestampLabel, idLabel, topic);
-
- if (subscription != null) {
- String subscriptionString = subscription.isAccessible()
- ? subscription.get().asPath() : subscription.toString();
- builder.add(DisplayData.item("subscription", subscriptionString)
- .withLabel("Pubsub Subscription"));
- }
- }
-
- @Override
- protected Coder<T> getDefaultOutputCoder() {
- return coder;
- }
-
- /**
- * Get the topic being read from.
- */
- @Nullable
- public PubsubTopic getTopic() {
- return topic == null ? null : topic.get();
- }
-
- /**
- * Get the {@link ValueProvider} for the topic being read from.
- */
- public ValueProvider<PubsubTopic> getTopicProvider() {
- return topic;
- }
-
- /**
- * Get the subscription being read from.
- */
- @Nullable
- public PubsubSubscription getSubscription() {
- return subscription == null ? null : subscription.get();
- }
-
- /**
- * Get the {@link ValueProvider} for the subscription being read from.
- */
- public ValueProvider<PubsubSubscription> getSubscriptionProvider() {
- return subscription;
- }
-
- /**
- * Get the timestamp label.
- */
- @Nullable
- public String getTimestampLabel() {
- return timestampLabel;
- }
-
- /**
- * Get the id label.
- */
- @Nullable
- public String getIdLabel() {
- return idLabel;
- }
-
-
- /**
- * Get the {@link Coder} used for the transform's output.
- */
- @Nullable
- public Coder<T> getCoder() {
- return coder;
- }
-
- /**
- * Get the parse function used for PubSub attributes.
- */
- @Nullable
- public SimpleFunction<PubsubMessage, T> getPubSubMessageParseFn() {
- return parseFn;
- }
-
- }
-
- /////////////////////////////////////////////////////////////////////////////
-
- /** Disallow construction of utility class. */
- private PubsubIO() {}
-
-
- /**
- * A {@link PTransform} that writes an unbounded {@link PCollection} of {@link String Strings}
- * to a Cloud Pub/Sub stream.
- */
- public static class Write<T> extends PTransform<PCollection<T>, PDone> {
-
- /** The Cloud Pub/Sub topic to publish to. */
- @Nullable
- private final ValueProvider<PubsubTopic> topic;
- /** The name of the message attribute to publish message timestamps in. */
- @Nullable
- private final String timestampLabel;
- /** The name of the message attribute to publish unique message IDs in. */
- @Nullable
- private final String idLabel;
- /** The input type Coder. */
- private final Coder<T> coder;
- /** The format function for input PubsubMessage objects. */
- SimpleFunction<T, PubsubMessage> formatFn;
-
- private Write() {
- this(null, null, null, null, null, null);
- }
-
- private Write(
- String name, ValueProvider<PubsubTopic> topic, String timestampLabel,
- String idLabel, Coder<T> coder, SimpleFunction<T, PubsubMessage> formatFn) {
- super(name);
- this.topic = topic;
- this.timestampLabel = timestampLabel;
- this.idLabel = idLabel;
- this.coder = coder;
- this.formatFn = formatFn;
- }
-
- /**
- * Creates a transform that publishes to the specified topic.
- *
- * <p>See {@link PubsubIO.PubsubTopic#fromPath(String)} for more details on the format of the
- * {@code topic} string.
- */
- public Write<T> topic(String topic) {
- return topic(StaticValueProvider.of(topic));
- }
-
- /**
- * Like {@code topic()} but with a {@link ValueProvider}.
- */
- public Write<T> topic(ValueProvider<String> topic) {
- return new Write<>(name, NestedValueProvider.of(topic, new TopicTranslator()),
- timestampLabel, idLabel, coder, formatFn);
- }
-
- /**
- * Creates a transform that writes to Pub/Sub, adds each record's timestamp to the published
- * messages in an attribute with the specified name. The value of the attribute will be a number
- * representing the number of milliseconds since the Unix epoch. For example, if using the Joda
- * time classes, {@link Instant#Instant(long)} can be used to parse this value.
- *
- * <p>If the output from this sink is being read by another Beam pipeline, then
- * {@link PubsubIO.Read#timestampLabel(String)} can be used to ensure the other source reads
- * these timestamps from the appropriate attribute.
- */
- public Write<T> timestampLabel(String timestampLabel) {
- return new Write<>(name, topic, timestampLabel, idLabel, coder, formatFn);
- }
-
- /**
- * Creates a transform that writes to Pub/Sub, adding each record's unique identifier to the
- * published messages in an attribute with the specified name. The value of the attribute is an
- * opaque string.
- *
- * <p>If the the output from this sink is being read by another Beam pipeline, then
- * {@link PubsubIO.Read#idLabel(String)} can be used to ensure that* the other source reads
- * these unique identifiers from the appropriate attribute.
- */
- public Write<T> idLabel(String idLabel) {
- return new Write<>(name, topic, timestampLabel, idLabel, coder, formatFn);
- }
-
- /**
- * Returns a new transform that's like this one
- * but that uses the given {@link Coder} to encode each of
- * the elements of the input {@link PCollection} into an
- * output record.
- *
- * <p>Does not modify this object.
- */
- public Write<T> withCoder(Coder<T> coder) {
- return new Write<>(name, topic, timestampLabel, idLabel, coder, formatFn);
- }
-
- /**
- * Used to write a PubSub message together with PubSub attributes. The user-supplied format
- * function translates the input type T to a PubsubMessage object, which is used by the sink
- * to separately set the PubSub message's payload and attributes.
- */
- public Write<T> withAttributes(SimpleFunction<T, PubsubMessage> formatFn) {
- return new Write<T>(name, topic, timestampLabel, idLabel, coder, formatFn);
- }
-
- @Override
- public PDone expand(PCollection<T> input) {
- if (topic == null) {
- throw new IllegalStateException("need to set the topic of a PubsubIO.Write transform");
- }
- switch (input.isBounded()) {
- case BOUNDED:
- input.apply(ParDo.of(new PubsubBoundedWriter()));
- return PDone.in(input.getPipeline());
- case UNBOUNDED:
- return input.apply(new PubsubUnboundedSink<T>(
- FACTORY,
- NestedValueProvider.of(topic, new TopicPathTranslator()),
- coder,
- timestampLabel,
- idLabel,
- formatFn,
- 100 /* numShards */));
- }
- throw new RuntimeException(); // cases are exhaustive.
- }
-
- @Override
- public void populateDisplayData(DisplayData.Builder builder) {
- super.populateDisplayData(builder);
- populateCommonDisplayData(builder, timestampLabel, idLabel, topic);
- }
-
- @Override
- protected Coder<Void> getDefaultOutputCoder() {
- return VoidCoder.of();
- }
-
- /**
- * Returns the PubSub topic being written to.
- */
- @Nullable
- public PubsubTopic getTopic() {
- return (topic == null) ? null : topic.get();
- }
-
- /**
- * Returns the {@link ValueProvider} for the topic being written to.
- */
- @Nullable
- public ValueProvider<PubsubTopic> getTopicProvider() {
- return topic;
- }
-
- /**
- * Returns the timestamp label.
- */
- @Nullable
- public String getTimestampLabel() {
- return timestampLabel;
- }
-
- /**
- * Returns the id label.
- */
- @Nullable
- public String getIdLabel() {
- return idLabel;
- }
-
- /**
- * Returns the output coder.
- */
- @Nullable
- public Coder<T> getCoder() {
- return coder;
- }
-
- /**
- * Returns the formatting function used if publishing attributes.
- */
- @Nullable
- public SimpleFunction<T, PubsubMessage> getFormatFn() {
- return formatFn;
- }
-
- /**
- * Writer to Pubsub which batches messages from bounded collections.
- *
- * <p>Public so can be suppressed by runners.
- */
- public class PubsubBoundedWriter extends DoFn<T, Void> {
-
- private static final int MAX_PUBLISH_BATCH_SIZE = 100;
- private transient List<OutgoingMessage> output;
- private transient PubsubClient pubsubClient;
-
- @StartBundle
- public void startBundle(Context c) throws IOException {
- this.output = new ArrayList<>();
- // NOTE: idLabel is ignored.
- this.pubsubClient =
- FACTORY.newClient(timestampLabel, null,
- c.getPipelineOptions().as(PubsubOptions.class));
- }
-
- @ProcessElement
- public void processElement(ProcessContext c) throws IOException {
- byte[] payload = null;
- Map<String, String> attributes = null;
- if (formatFn != null) {
- PubsubMessage message = formatFn.apply(c.element());
- payload = message.getMessage();
- attributes = message.getAttributeMap();
- } else {
- payload = CoderUtils.encodeToByteArray(getCoder(), c.element());
- }
- // NOTE: The record id is always null.
- OutgoingMessage message =
- new OutgoingMessage(payload, attributes, c.timestamp().getMillis(), null);
- output.add(message);
-
- if (output.size() >= MAX_PUBLISH_BATCH_SIZE) {
- publish();
- }
- }
-
- @FinishBundle
- public void finishBundle(Context c) throws IOException {
- if (!output.isEmpty()) {
- publish();
- }
- output = null;
- pubsubClient.close();
- pubsubClient = null;
- }
-
- private void publish() throws IOException {
- int n = pubsubClient.publish(
- PubsubClient.topicPathFromName(getTopic().project, getTopic().topic),
- output);
- checkState(n == output.size());
- output.clear();
- }
-
- @Override
- public void populateDisplayData(DisplayData.Builder builder) {
- super.populateDisplayData(builder);
- builder.delegate(Write.this);
- }
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/src/main/java/org/apache/beam/sdk/io/PubsubUnboundedSink.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/PubsubUnboundedSink.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/PubsubUnboundedSink.java
deleted file mode 100644
index 55605b3..0000000
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/PubsubUnboundedSink.java
+++ /dev/null
@@ -1,494 +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 static com.google.common.base.Preconditions.checkState;
-
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.hash.Hashing;
-
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-import java.util.UUID;
-import java.util.concurrent.ThreadLocalRandom;
-import javax.annotation.Nullable;
-
-import org.apache.beam.sdk.coders.AtomicCoder;
-import org.apache.beam.sdk.coders.BigEndianLongCoder;
-import org.apache.beam.sdk.coders.ByteArrayCoder;
-import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.coders.CoderException;
-import org.apache.beam.sdk.coders.KvCoder;
-import org.apache.beam.sdk.coders.MapCoder;
-import org.apache.beam.sdk.coders.NullableCoder;
-import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.sdk.coders.VarIntCoder;
-import org.apache.beam.sdk.io.PubsubIO.PubsubMessage;
-import org.apache.beam.sdk.metrics.Counter;
-import org.apache.beam.sdk.metrics.Metrics;
-import org.apache.beam.sdk.options.PubsubOptions;
-import org.apache.beam.sdk.options.ValueProvider;
-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.SimpleFunction;
-import org.apache.beam.sdk.transforms.display.DisplayData;
-import org.apache.beam.sdk.transforms.display.DisplayData.Builder;
-import org.apache.beam.sdk.transforms.windowing.AfterFirst;
-import org.apache.beam.sdk.transforms.windowing.AfterPane;
-import org.apache.beam.sdk.transforms.windowing.AfterProcessingTime;
-import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
-import org.apache.beam.sdk.transforms.windowing.Repeatedly;
-import org.apache.beam.sdk.transforms.windowing.Window;
-import org.apache.beam.sdk.util.CoderUtils;
-import org.apache.beam.sdk.util.PubsubClient;
-import org.apache.beam.sdk.util.PubsubClient.OutgoingMessage;
-import org.apache.beam.sdk.util.PubsubClient.PubsubClientFactory;
-import org.apache.beam.sdk.util.PubsubClient.TopicPath;
-import org.apache.beam.sdk.values.KV;
-import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PDone;
-import org.joda.time.Duration;
-
-/**
- * A PTransform which streams messages to Pubsub.
- * <ul>
- * <li>The underlying implementation is just a {@link GroupByKey} followed by a {@link ParDo} which
- * publishes as a side effect. (In the future we want to design and switch to a custom
- * {@code UnboundedSink} implementation so as to gain access to system watermark and
- * end-of-pipeline cleanup.)
- * <li>We try to send messages in batches while also limiting send latency.
- * <li>No stats are logged. Rather some counters are used to keep track of elements and batches.
- * <li>Though some background threads are used by the underlying netty system all actual Pubsub
- * calls are blocking. We rely on the underlying runner to allow multiple {@link DoFn} instances
- * to execute concurrently and hide latency.
- * <li>A failed bundle will cause messages to be resent. Thus we rely on the Pubsub consumer
- * to dedup messages.
- * </ul>
- */
-public class PubsubUnboundedSink<T> extends PTransform<PCollection<T>, PDone> {
- /**
- * Default maximum number of messages per publish.
- */
- private static final int DEFAULT_PUBLISH_BATCH_SIZE = 1000;
-
- /**
- * Default maximum size of a publish batch, in bytes.
- */
- private static final int DEFAULT_PUBLISH_BATCH_BYTES = 400000;
-
- /**
- * Default longest delay between receiving a message and pushing it to Pubsub.
- */
- private static final Duration DEFAULT_MAX_LATENCY = Duration.standardSeconds(2);
-
- /**
- * Coder for conveying outgoing messages between internal stages.
- */
- private static class OutgoingMessageCoder extends AtomicCoder<OutgoingMessage> {
- private static final NullableCoder<String> RECORD_ID_CODER =
- NullableCoder.of(StringUtf8Coder.of());
- private static final NullableCoder<Map<String, String>> ATTRIBUTES_CODER =
- NullableCoder.of(MapCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()));
-
- @Override
- public void encode(
- OutgoingMessage value, OutputStream outStream, Context context)
- throws CoderException, IOException {
- ByteArrayCoder.of().encode(value.elementBytes, outStream, context.nested());
- ATTRIBUTES_CODER.encode(value.attributes, outStream, context.nested());
- BigEndianLongCoder.of().encode(value.timestampMsSinceEpoch, outStream, context.nested());
- RECORD_ID_CODER.encode(value.recordId, outStream, context.nested());
- }
-
- @Override
- public OutgoingMessage decode(
- InputStream inStream, Context context) throws CoderException, IOException {
- byte[] elementBytes = ByteArrayCoder.of().decode(inStream, context.nested());
- Map<String, String> attributes = ATTRIBUTES_CODER.decode(inStream, context.nested());
- long timestampMsSinceEpoch = BigEndianLongCoder.of().decode(inStream, context.nested());
- @Nullable String recordId = RECORD_ID_CODER.decode(inStream, context.nested());
- return new OutgoingMessage(elementBytes, attributes, timestampMsSinceEpoch, recordId);
- }
- }
-
- @VisibleForTesting
- static final Coder<OutgoingMessage> CODER = new OutgoingMessageCoder();
-
- // ================================================================================
- // RecordIdMethod
- // ================================================================================
-
- /**
- * Specify how record ids are to be generated.
- */
- @VisibleForTesting
- enum RecordIdMethod {
- /** Leave null. */
- NONE,
- /** Generate randomly. */
- RANDOM,
- /** Generate deterministically. For testing only. */
- DETERMINISTIC
- }
-
- // ================================================================================
- // ShardFn
- // ================================================================================
-
- /**
- * Convert elements to messages and shard them.
- */
- private static class ShardFn<T> extends DoFn<T, KV<Integer, OutgoingMessage>> {
- private final Counter elementCounter = Metrics.counter(ShardFn.class, "elements");
- private final Coder<T> elementCoder;
- private final int numShards;
- private final RecordIdMethod recordIdMethod;
- private final SimpleFunction<T, PubsubMessage> formatFn;
-
- ShardFn(Coder<T> elementCoder, int numShards,
- SimpleFunction<T, PubsubIO.PubsubMessage> formatFn, RecordIdMethod recordIdMethod) {
- this.elementCoder = elementCoder;
- this.numShards = numShards;
- this.formatFn = formatFn;
- this.recordIdMethod = recordIdMethod;
- }
-
- @ProcessElement
- public void processElement(ProcessContext c) throws Exception {
- elementCounter.inc();
- byte[] elementBytes = null;
- Map<String, String> attributes = ImmutableMap.<String, String>of();
- if (formatFn != null) {
- PubsubIO.PubsubMessage message = formatFn.apply(c.element());
- elementBytes = message.getMessage();
- attributes = message.getAttributeMap();
- } else {
- elementBytes = CoderUtils.encodeToByteArray(elementCoder, c.element());
- }
-
- long timestampMsSinceEpoch = c.timestamp().getMillis();
- @Nullable String recordId = null;
- switch (recordIdMethod) {
- case NONE:
- break;
- case DETERMINISTIC:
- recordId = Hashing.murmur3_128().hashBytes(elementBytes).toString();
- break;
- case RANDOM:
- // Since these elements go through a GroupByKey, any failures while sending to
- // Pubsub will be retried without falling back and generating a new record id.
- // Thus even though we may send the same message to Pubsub twice, it is guaranteed
- // to have the same record id.
- recordId = UUID.randomUUID().toString();
- break;
- }
- c.output(KV.of(ThreadLocalRandom.current().nextInt(numShards),
- new OutgoingMessage(elementBytes, attributes, timestampMsSinceEpoch,
- recordId)));
- }
-
- @Override
- public void populateDisplayData(Builder builder) {
- super.populateDisplayData(builder);
- builder.add(DisplayData.item("numShards", numShards));
- }
- }
-
- // ================================================================================
- // WriterFn
- // ================================================================================
-
- /**
- * Publish messages to Pubsub in batches.
- */
- private static class WriterFn
- extends DoFn<KV<Integer, Iterable<OutgoingMessage>>, Void> {
- private final PubsubClientFactory pubsubFactory;
- private final ValueProvider<TopicPath> topic;
- private final String timestampLabel;
- private final String idLabel;
- private final int publishBatchSize;
- private final int publishBatchBytes;
-
- /**
- * Client on which to talk to Pubsub. Null until created by {@link #startBundle}.
- */
- @Nullable
- private transient PubsubClient pubsubClient;
-
- private final Counter batchCounter = Metrics.counter(WriterFn.class, "batches");
- private final Counter elementCounter = Metrics.counter(WriterFn.class, "elements");
- private final Counter byteCounter = Metrics.counter(WriterFn.class, "bytes");
-
- WriterFn(
- PubsubClientFactory pubsubFactory, ValueProvider<TopicPath> topic,
- String timestampLabel, String idLabel, int publishBatchSize, int publishBatchBytes) {
- this.pubsubFactory = pubsubFactory;
- this.topic = topic;
- this.timestampLabel = timestampLabel;
- this.idLabel = idLabel;
- this.publishBatchSize = publishBatchSize;
- this.publishBatchBytes = publishBatchBytes;
- }
-
- /**
- * BLOCKING
- * Send {@code messages} as a batch to Pubsub.
- */
- private void publishBatch(List<OutgoingMessage> messages, int bytes)
- throws IOException {
- int n = pubsubClient.publish(topic.get(), messages);
- checkState(n == messages.size(), "Attempted to publish %s messages but %s were successful",
- messages.size(), n);
- batchCounter.inc();
- elementCounter.inc(messages.size());
- byteCounter.inc(bytes);
- }
-
- @StartBundle
- public void startBundle(Context c) throws Exception {
- checkState(pubsubClient == null, "startBundle invoked without prior finishBundle");
- pubsubClient = pubsubFactory.newClient(timestampLabel, idLabel,
- c.getPipelineOptions().as(PubsubOptions.class));
- }
-
- @ProcessElement
- public void processElement(ProcessContext c) throws Exception {
- List<OutgoingMessage> pubsubMessages = new ArrayList<>(publishBatchSize);
- int bytes = 0;
- for (OutgoingMessage message : c.element().getValue()) {
- if (!pubsubMessages.isEmpty()
- && bytes + message.elementBytes.length > publishBatchBytes) {
- // Break large (in bytes) batches into smaller.
- // (We've already broken by batch size using the trigger below, though that may
- // run slightly over the actual PUBLISH_BATCH_SIZE. We'll consider that ok since
- // the hard limit from Pubsub is by bytes rather than number of messages.)
- // BLOCKS until published.
- publishBatch(pubsubMessages, bytes);
- pubsubMessages.clear();
- bytes = 0;
- }
- pubsubMessages.add(message);
- bytes += message.elementBytes.length;
- }
- if (!pubsubMessages.isEmpty()) {
- // BLOCKS until published.
- publishBatch(pubsubMessages, bytes);
- }
- }
-
- @FinishBundle
- public void finishBundle(Context c) throws Exception {
- pubsubClient.close();
- pubsubClient = null;
- }
-
- @Override
- public void populateDisplayData(Builder builder) {
- super.populateDisplayData(builder);
- String topicString =
- topic == null ? null
- : topic.isAccessible() ? topic.get().getPath()
- : topic.toString();
- builder.add(DisplayData.item("topic", topicString));
- builder.add(DisplayData.item("transport", pubsubFactory.getKind()));
- builder.addIfNotNull(DisplayData.item("timestampLabel", timestampLabel));
- builder.addIfNotNull(DisplayData.item("idLabel", idLabel));
- }
- }
-
- // ================================================================================
- // PubsubUnboundedSink
- // ================================================================================
-
- /**
- * Which factory to use for creating Pubsub transport.
- */
- private final PubsubClientFactory pubsubFactory;
-
- /**
- * Pubsub topic to publish to.
- */
- private final ValueProvider<TopicPath> topic;
-
- /**
- * Coder for elements. It is the responsibility of the underlying Pubsub transport to
- * re-encode element bytes if necessary, eg as Base64 strings.
- */
- private final Coder<T> elementCoder;
-
- /**
- * Pubsub metadata field holding timestamp of each element, or {@literal null} if should use
- * Pubsub message publish timestamp instead.
- */
- @Nullable
- private final String timestampLabel;
-
- /**
- * Pubsub metadata field holding id for each element, or {@literal null} if need to generate
- * a unique id ourselves.
- */
- @Nullable
- private final String idLabel;
-
- /**
- * Number of 'shards' to use so that latency in Pubsub publish can be hidden. Generally this
- * should be a small multiple of the number of available cores. Too smoll a number results
- * in too much time lost to blocking Pubsub calls. To large a number results in too many
- * single-element batches being sent to Pubsub with high per-batch overhead.
- */
- private final int numShards;
-
- /**
- * Maximum number of messages per publish.
- */
- private final int publishBatchSize;
-
- /**
- * Maximum size of a publish batch, in bytes.
- */
- private final int publishBatchBytes;
-
- /**
- * Longest delay between receiving a message and pushing it to Pubsub.
- */
- private final Duration maxLatency;
-
- /**
- * How record ids should be generated for each record (if {@link #idLabel} is non-{@literal
- * null}).
- */
- private final RecordIdMethod recordIdMethod;
-
- /**
- * In order to publish attributes, a formatting function is used to format the output into
- * a {@link PubsubIO.PubsubMessage}.
- */
- private final SimpleFunction<T, PubsubIO.PubsubMessage> formatFn;
-
- @VisibleForTesting
- PubsubUnboundedSink(
- PubsubClientFactory pubsubFactory,
- ValueProvider<TopicPath> topic,
- Coder<T> elementCoder,
- String timestampLabel,
- String idLabel,
- int numShards,
- int publishBatchSize,
- int publishBatchBytes,
- Duration maxLatency,
- SimpleFunction<T, PubsubIO.PubsubMessage> formatFn,
- RecordIdMethod recordIdMethod) {
- this.pubsubFactory = pubsubFactory;
- this.topic = topic;
- this.elementCoder = elementCoder;
- this.timestampLabel = timestampLabel;
- this.idLabel = idLabel;
- this.numShards = numShards;
- this.publishBatchSize = publishBatchSize;
- this.publishBatchBytes = publishBatchBytes;
- this.maxLatency = maxLatency;
- this.formatFn = formatFn;
- this.recordIdMethod = idLabel == null ? RecordIdMethod.NONE : recordIdMethod;
- }
-
- public PubsubUnboundedSink(
- PubsubClientFactory pubsubFactory,
- ValueProvider<TopicPath> topic,
- Coder<T> elementCoder,
- String timestampLabel,
- String idLabel,
- SimpleFunction<T, PubsubIO.PubsubMessage> formatFn,
- int numShards) {
- this(pubsubFactory, topic, elementCoder, timestampLabel, idLabel, numShards,
- DEFAULT_PUBLISH_BATCH_SIZE, DEFAULT_PUBLISH_BATCH_BYTES, DEFAULT_MAX_LATENCY,
- formatFn, RecordIdMethod.RANDOM);
- }
-
- /**
- * Get the topic being written to.
- */
- public TopicPath getTopic() {
- return topic.get();
- }
-
- /**
- * Get the {@link ValueProvider} for the topic being written to.
- */
- public ValueProvider<TopicPath> getTopicProvider() {
- return topic;
- }
-
- /**
- * Get the timestamp label.
- */
- @Nullable
- public String getTimestampLabel() {
- return timestampLabel;
- }
-
- /**
- * Get the id label.
- */
- @Nullable
- public String getIdLabel() {
- return idLabel;
- }
-
- /**
- * Get the format function used for PubSub attributes.
- */
- @Nullable
- public SimpleFunction<T, PubsubIO.PubsubMessage> getFormatFn() {
- return formatFn;
- }
-
- /**
- * Get the Coder used to encode output elements.
- */
- public Coder<T> getElementCoder() {
- return elementCoder;
- }
-
- @Override
- public PDone expand(PCollection<T> input) {
- input.apply("PubsubUnboundedSink.Window", Window.<T>into(new GlobalWindows())
- .triggering(
- Repeatedly.forever(
- AfterFirst.of(AfterPane.elementCountAtLeast(publishBatchSize),
- AfterProcessingTime.pastFirstElementInPane()
- .plusDelayOf(maxLatency))))
- .discardingFiredPanes())
- .apply("PubsubUnboundedSink.Shard",
- ParDo.of(new ShardFn<T>(elementCoder, numShards, formatFn, recordIdMethod)))
- .setCoder(KvCoder.of(VarIntCoder.of(), CODER))
- .apply(GroupByKey.<Integer, OutgoingMessage>create())
- .apply("PubsubUnboundedSink.Writer",
- ParDo.of(new WriterFn(pubsubFactory, topic, timestampLabel, idLabel,
- publishBatchSize, publishBatchBytes)));
- return PDone.in(input.getPipeline());
- }
-}
[44/50] [abbrv] beam git commit: This closes #2363
Posted by jb...@apache.org.
This closes #2363
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/b7ddab28
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/b7ddab28
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/b7ddab28
Branch: refs/heads/DSL_SQL
Commit: b7ddab2818cebb702ecb61a36fdb6c79c4a4956f
Parents: f314354 831b11f
Author: Eugene Kirpichov <ki...@google.com>
Authored: Wed Apr 12 11:38:32 2017 -0700
Committer: Eugene Kirpichov <ki...@google.com>
Committed: Wed Apr 12 11:38:32 2017 -0700
----------------------------------------------------------------------
.../beam/examples/MinimalWordCountJava8.java | 9 +-
.../beam/examples/complete/game/GameStats.java | 6 +-
.../beam/examples/complete/game/UserScore.java | 5 +-
.../examples/MinimalWordCountJava8Test.java | 9 +-
.../complete/game/HourlyTeamScoreTest.java | 5 +-
.../examples/complete/game/UserScoreTest.java | 6 +-
.../beam/runners/direct/DirectRunnerTest.java | 5 +-
.../beam/sdk/transforms/FlatMapElements.java | 113 +++++++++----------
.../apache/beam/sdk/transforms/MapElements.java | 99 ++++++++--------
.../beam/sdk/transforms/MapElementsTest.java | 25 ++--
.../sdk/io/gcp/bigquery/BatchLoadBigQuery.java | 7 +-
.../transforms/FlatMapElementsJava8Test.java | 10 +-
.../sdk/transforms/MapElementsJava8Test.java | 10 +-
13 files changed, 153 insertions(+), 156 deletions(-)
----------------------------------------------------------------------
[42/50] [abbrv] beam git commit: Removes
MapElements.MissingOutputTypeDescriptor
Posted by jb...@apache.org.
Removes MapElements.MissingOutputTypeDescriptor
This comes from changing MapElements.via(fn).withOutputType(td)
to MapElements.into(td).via(fn) which is also shorter.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/3bf6c6b3
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/3bf6c6b3
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/3bf6c6b3
Branch: refs/heads/DSL_SQL
Commit: 3bf6c6b3f81c643a1107346674088f554aca29a8
Parents: f314354
Author: Eugene Kirpichov <ki...@google.com>
Authored: Wed Mar 29 13:42:29 2017 -0700
Committer: Eugene Kirpichov <ki...@google.com>
Committed: Wed Apr 12 11:35:05 2017 -0700
----------------------------------------------------------------------
.../beam/examples/MinimalWordCountJava8.java | 4 +-
.../beam/examples/complete/game/GameStats.java | 6 +-
.../beam/examples/complete/game/UserScore.java | 5 +-
.../examples/MinimalWordCountJava8Test.java | 4 +-
.../complete/game/HourlyTeamScoreTest.java | 5 +-
.../examples/complete/game/UserScoreTest.java | 6 +-
.../beam/runners/direct/DirectRunnerTest.java | 5 +-
.../apache/beam/sdk/transforms/MapElements.java | 99 ++++++++++----------
.../beam/sdk/transforms/MapElementsTest.java | 25 +++--
.../sdk/io/gcp/bigquery/BatchLoadBigQuery.java | 7 +-
.../sdk/transforms/MapElementsJava8Test.java | 10 +-
11 files changed, 90 insertions(+), 86 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/3bf6c6b3/examples/java8/src/main/java/org/apache/beam/examples/MinimalWordCountJava8.java
----------------------------------------------------------------------
diff --git a/examples/java8/src/main/java/org/apache/beam/examples/MinimalWordCountJava8.java b/examples/java8/src/main/java/org/apache/beam/examples/MinimalWordCountJava8.java
index 738b64d..0072886 100644
--- a/examples/java8/src/main/java/org/apache/beam/examples/MinimalWordCountJava8.java
+++ b/examples/java8/src/main/java/org/apache/beam/examples/MinimalWordCountJava8.java
@@ -61,8 +61,8 @@ public class MinimalWordCountJava8 {
.apply(Filter.by((String word) -> !word.isEmpty()))
.apply(Count.<String>perElement())
.apply(MapElements
- .via((KV<String, Long> wordCount) -> wordCount.getKey() + ": " + wordCount.getValue())
- .withOutputType(TypeDescriptors.strings()))
+ .into(TypeDescriptors.strings())
+ .via((KV<String, Long> wordCount) -> wordCount.getKey() + ": " + wordCount.getValue()))
// CHANGE 3/3: The Google Cloud Storage path is required for outputting the results to.
.apply(TextIO.Write.to("gs://YOUR_OUTPUT_BUCKET/AND_OUTPUT_PREFIX"));
http://git-wip-us.apache.org/repos/asf/beam/blob/3bf6c6b3/examples/java8/src/main/java/org/apache/beam/examples/complete/game/GameStats.java
----------------------------------------------------------------------
diff --git a/examples/java8/src/main/java/org/apache/beam/examples/complete/game/GameStats.java b/examples/java8/src/main/java/org/apache/beam/examples/complete/game/GameStats.java
index 6874953..9c79fad 100644
--- a/examples/java8/src/main/java/org/apache/beam/examples/complete/game/GameStats.java
+++ b/examples/java8/src/main/java/org/apache/beam/examples/complete/game/GameStats.java
@@ -260,9 +260,9 @@ public class GameStats extends LeaderBoard {
// Extract username/score pairs from the event stream
PCollection<KV<String, Integer>> userEvents =
rawEvents.apply("ExtractUserScore",
- MapElements.via((GameActionInfo gInfo) -> KV.of(gInfo.getUser(), gInfo.getScore()))
- .withOutputType(
- TypeDescriptors.kvs(TypeDescriptors.strings(), TypeDescriptors.integers())));
+ MapElements
+ .into(TypeDescriptors.kvs(TypeDescriptors.strings(), TypeDescriptors.integers()))
+ .via((GameActionInfo gInfo) -> KV.of(gInfo.getUser(), gInfo.getScore())));
// Calculate the total score per user over fixed windows, and
// cumulative updates for late data.
http://git-wip-us.apache.org/repos/asf/beam/blob/3bf6c6b3/examples/java8/src/main/java/org/apache/beam/examples/complete/game/UserScore.java
----------------------------------------------------------------------
diff --git a/examples/java8/src/main/java/org/apache/beam/examples/complete/game/UserScore.java b/examples/java8/src/main/java/org/apache/beam/examples/complete/game/UserScore.java
index 7dd5a8e..b4b023f 100644
--- a/examples/java8/src/main/java/org/apache/beam/examples/complete/game/UserScore.java
+++ b/examples/java8/src/main/java/org/apache/beam/examples/complete/game/UserScore.java
@@ -165,9 +165,8 @@ public class UserScore {
return gameInfo
.apply(MapElements
- .via((GameActionInfo gInfo) -> KV.of(gInfo.getKey(field), gInfo.getScore()))
- .withOutputType(
- TypeDescriptors.kvs(TypeDescriptors.strings(), TypeDescriptors.integers())))
+ .into(TypeDescriptors.kvs(TypeDescriptors.strings(), TypeDescriptors.integers()))
+ .via((GameActionInfo gInfo) -> KV.of(gInfo.getKey(field), gInfo.getScore())))
.apply(Sum.<String>integersPerKey());
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/3bf6c6b3/examples/java8/src/test/java/org/apache/beam/examples/MinimalWordCountJava8Test.java
----------------------------------------------------------------------
diff --git a/examples/java8/src/test/java/org/apache/beam/examples/MinimalWordCountJava8Test.java b/examples/java8/src/test/java/org/apache/beam/examples/MinimalWordCountJava8Test.java
index c2f3efe..911ccf6 100644
--- a/examples/java8/src/test/java/org/apache/beam/examples/MinimalWordCountJava8Test.java
+++ b/examples/java8/src/test/java/org/apache/beam/examples/MinimalWordCountJava8Test.java
@@ -68,8 +68,8 @@ public class MinimalWordCountJava8Test implements Serializable {
.apply(Filter.by((String word) -> !word.isEmpty()))
.apply(Count.<String>perElement())
.apply(MapElements
- .via((KV<String, Long> wordCount) -> wordCount.getKey() + ": " + wordCount.getValue())
- .withOutputType(TypeDescriptors.strings()))
+ .into(TypeDescriptors.strings())
+ .via((KV<String, Long> wordCount) -> wordCount.getKey() + ": " + wordCount.getValue()))
.apply(TextIO.Write.to("gs://your-output-bucket/and-output-prefix"));
}
http://git-wip-us.apache.org/repos/asf/beam/blob/3bf6c6b3/examples/java8/src/test/java/org/apache/beam/examples/complete/game/HourlyTeamScoreTest.java
----------------------------------------------------------------------
diff --git a/examples/java8/src/test/java/org/apache/beam/examples/complete/game/HourlyTeamScoreTest.java b/examples/java8/src/test/java/org/apache/beam/examples/complete/game/HourlyTeamScoreTest.java
index 40bbfdb..409fc92 100644
--- a/examples/java8/src/test/java/org/apache/beam/examples/complete/game/HourlyTeamScoreTest.java
+++ b/examples/java8/src/test/java/org/apache/beam/examples/complete/game/HourlyTeamScoreTest.java
@@ -101,9 +101,8 @@ public class HourlyTeamScoreTest implements Serializable {
-> gInfo.getTimestamp() > startMinTimestamp.getMillis()))
// run a map to access the fields in the result.
.apply(MapElements
- .via((GameActionInfo gInfo) -> KV.of(gInfo.getUser(), gInfo.getScore()))
- .withOutputType(
- TypeDescriptors.kvs(TypeDescriptors.strings(), TypeDescriptors.integers())));
+ .into(TypeDescriptors.kvs(TypeDescriptors.strings(), TypeDescriptors.integers()))
+ .via((GameActionInfo gInfo) -> KV.of(gInfo.getUser(), gInfo.getScore())));
PAssert.that(output).containsInAnyOrder(FILTERED_EVENTS);
http://git-wip-us.apache.org/repos/asf/beam/blob/3bf6c6b3/examples/java8/src/test/java/org/apache/beam/examples/complete/game/UserScoreTest.java
----------------------------------------------------------------------
diff --git a/examples/java8/src/test/java/org/apache/beam/examples/complete/game/UserScoreTest.java b/examples/java8/src/test/java/org/apache/beam/examples/complete/game/UserScoreTest.java
index f0c28ab..2eb63aa 100644
--- a/examples/java8/src/test/java/org/apache/beam/examples/complete/game/UserScoreTest.java
+++ b/examples/java8/src/test/java/org/apache/beam/examples/complete/game/UserScoreTest.java
@@ -143,9 +143,9 @@ public class UserScoreTest implements Serializable {
PCollection<KV<String, Integer>> extract = input
.apply(ParDo.of(new ParseEventFn()))
.apply(
- MapElements.via((GameActionInfo gInfo) -> KV.of(gInfo.getUser(), gInfo.getScore()))
- .withOutputType(
- TypeDescriptors.kvs(TypeDescriptors.strings(), TypeDescriptors.integers())));
+ MapElements
+ .into(TypeDescriptors.kvs(TypeDescriptors.strings(), TypeDescriptors.integers()))
+ .via((GameActionInfo gInfo) -> KV.of(gInfo.getUser(), gInfo.getScore())));
PAssert.that(extract).empty();
http://git-wip-us.apache.org/repos/asf/beam/blob/3bf6c6b3/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectRunnerTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectRunnerTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectRunnerTest.java
index 28c24ad..3b81f4d 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectRunnerTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectRunnerTest.java
@@ -192,9 +192,10 @@ public class DirectRunnerTest implements Serializable {
TypeDescriptor<byte[]> td = new TypeDescriptor<byte[]>() {
};
PCollection<byte[]> foos =
- p.apply(Create.of(1, 1, 1, 2, 2, 3)).apply(MapElements.via(getBytes).withOutputType(td));
+ p.apply(Create.of(1, 1, 1, 2, 2, 3))
+ .apply(MapElements.into(td).via(getBytes));
PCollection<byte[]> msync =
- p.apply(Create.of(1, -2, -8, -16)).apply(MapElements.via(getBytes).withOutputType(td));
+ p.apply(Create.of(1, -2, -8, -16)).apply(MapElements.into(td).via(getBytes));
PCollection<byte[]> bytes =
PCollectionList.of(foos).and(msync).apply(Flatten.<byte[]>pCollections());
PCollection<KV<byte[], Long>> counts = bytes.apply(Count.<byte[]>perElement());
http://git-wip-us.apache.org/repos/asf/beam/blob/3bf6c6b3/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/MapElements.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/MapElements.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/MapElements.java
index 421b2ab..82cf753 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/MapElements.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/MapElements.java
@@ -17,6 +17,9 @@
*/
package org.apache.beam.sdk.transforms;
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import javax.annotation.Nullable;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.TypeDescriptor;
@@ -26,34 +29,27 @@ import org.apache.beam.sdk.values.TypeDescriptor;
*/
public class MapElements<InputT, OutputT>
extends PTransform<PCollection<? extends InputT>, PCollection<OutputT>> {
-
/**
- * For a {@code SerializableFunction<InputT, OutputT>} {@code fn} and output type descriptor,
- * returns a {@code PTransform} that takes an input {@code PCollection<InputT>} and returns
- * a {@code PCollection<OutputT>} containing {@code fn.apply(v)} for every element {@code v} in
- * the input.
- *
- * <p>Example of use in Java 8:
- * <pre>{@code
- * PCollection<Integer> wordLengths = words.apply(
- * MapElements.via((String word) -> word.length())
- * .withOutputType(new TypeDescriptor<Integer>() {});
- * }</pre>
- *
- * <p>In Java 7, the overload {@link #via(SimpleFunction)} is more concise as the output type
- * descriptor need not be provided.
+ * Temporarily stores the argument of {@link #into(TypeDescriptor)} until combined with the
+ * argument of {@link #via(SerializableFunction)} into the fully-specified {@link #fn}. Stays null
+ * if constructed using {@link #via(SimpleFunction)} directly.
*/
- public static <InputT, OutputT> MissingOutputTypeDescriptor<InputT, OutputT>
- via(SerializableFunction<? super InputT, OutputT> fn) {
+ @Nullable private final transient TypeDescriptor<OutputT> outputType;
- // TypeDescriptor interacts poorly with the wildcards needed to correctly express
- // covariance and contravariance in Java, so instead we cast it to an invariant
- // function here.
- @SuppressWarnings("unchecked") // safe covariant cast
- SerializableFunction<InputT, OutputT> simplerFn =
- (SerializableFunction<InputT, OutputT>) fn;
+ /**
+ * Non-null on a fully specified transform - is null only when constructed using {@link
+ * #into(TypeDescriptor)}, until the fn is specified using {@link #via(SerializableFunction)}.
+ */
+ @Nullable private final SimpleFunction<InputT, OutputT> fn;
+ private final DisplayData.ItemSpec<?> fnClassDisplayData;
- return new MissingOutputTypeDescriptor<>(simplerFn);
+ private MapElements(
+ @Nullable SimpleFunction<InputT, OutputT> fn,
+ @Nullable TypeDescriptor<OutputT> outputType,
+ @Nullable Class<?> fnClass) {
+ this.fn = fn;
+ this.outputType = outputType;
+ this.fnClassDisplayData = DisplayData.item("mapFn", fnClass).withLabel("Map Function");
}
/**
@@ -77,41 +73,46 @@ extends PTransform<PCollection<? extends InputT>, PCollection<OutputT>> {
*/
public static <InputT, OutputT> MapElements<InputT, OutputT> via(
final SimpleFunction<InputT, OutputT> fn) {
- return new MapElements<>(fn, fn.getClass());
+ return new MapElements<>(fn, null, fn.getClass());
}
/**
- * An intermediate builder for a {@link MapElements} transform. To complete the transform, provide
- * an output type descriptor to {@link MissingOutputTypeDescriptor#withOutputType}. See
- * {@link #via(SerializableFunction)} for a full example of use.
+ * Returns a new {@link MapElements} transform with the given type descriptor for the output
+ * type, but the mapping function yet to be specified using {@link #via(SerializableFunction)}.
*/
- public static final class MissingOutputTypeDescriptor<InputT, OutputT> {
-
- private final SerializableFunction<InputT, OutputT> fn;
-
- private MissingOutputTypeDescriptor(SerializableFunction<InputT, OutputT> fn) {
- this.fn = fn;
- }
-
- public MapElements<InputT, OutputT> withOutputType(final TypeDescriptor<OutputT> outputType) {
- return new MapElements<>(
- SimpleFunction.fromSerializableFunctionWithOutputType(fn, outputType), fn.getClass());
- }
-
+ public static <OutputT> MapElements<?, OutputT>
+ into(final TypeDescriptor<OutputT> outputType) {
+ return new MapElements<>(null, outputType, null);
}
- ///////////////////////////////////////////////////////////////////
-
- private final SimpleFunction<InputT, OutputT> fn;
- private final DisplayData.ItemSpec<?> fnClassDisplayData;
-
- private MapElements(SimpleFunction<InputT, OutputT> fn, Class<?> fnClass) {
- this.fn = fn;
- this.fnClassDisplayData = DisplayData.item("mapFn", fnClass).withLabel("Map Function");
+ /**
+ * For a {@code SerializableFunction<InputT, OutputT>} {@code fn} and output type descriptor,
+ * returns a {@code PTransform} that takes an input {@code PCollection<InputT>} and returns a
+ * {@code PCollection<OutputT>} containing {@code fn.apply(v)} for every element {@code v} in the
+ * input.
+ *
+ * <p>Example of use in Java 8:
+ *
+ * <pre>{@code
+ * PCollection<Integer> wordLengths = words.apply(
+ * MapElements.via((String word) -> word.length())
+ * .withOutputType(new TypeDescriptor<Integer>() {});
+ * }</pre>
+ *
+ * <p>In Java 7, the overload {@link #via(SimpleFunction)} is more concise as the output type
+ * descriptor need not be provided.
+ */
+ public <NewInputT> MapElements<NewInputT, OutputT> via(
+ SerializableFunction<NewInputT, OutputT> fn) {
+ return new MapElements<>(
+ SimpleFunction.fromSerializableFunctionWithOutputType(fn, outputType),
+ null,
+ fn.getClass());
}
@Override
public PCollection<OutputT> expand(PCollection<? extends InputT> input) {
+ checkNotNull(fn, "Must specify a function on MapElements using .via()");
return input.apply(
"Map",
ParDo.of(
http://git-wip-us.apache.org/repos/asf/beam/blob/3bf6c6b3/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/MapElementsTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/MapElementsTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/MapElementsTest.java
index 82e856e..7bf94a0 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/MapElementsTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/MapElementsTest.java
@@ -33,6 +33,7 @@ import org.apache.beam.sdk.transforms.display.DisplayDataEvaluator;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.sdk.values.TypeDescriptors;
import org.junit.Rule;
import org.junit.Test;
import org.junit.experimental.categories.Category;
@@ -152,14 +153,18 @@ public class MapElementsTest implements Serializable {
@Test
@Category(NeedsRunner.class)
public void testMapBasicSerializableFunction() throws Exception {
- PCollection<Integer> output = pipeline
- .apply(Create.of(1, 2, 3))
- .apply(MapElements.via(new SerializableFunction<Integer, Integer>() {
- @Override
- public Integer apply(Integer input) {
- return -input;
- }
- }).withOutputType(new TypeDescriptor<Integer>() {}));
+ PCollection<Integer> output =
+ pipeline
+ .apply(Create.of(1, 2, 3))
+ .apply(
+ MapElements.into(TypeDescriptors.integers())
+ .via(
+ new SerializableFunction<Integer, Integer>() {
+ @Override
+ public Integer apply(Integer input) {
+ return -input;
+ }
+ }));
PAssert.that(output).containsInAnyOrder(-2, -1, -3);
pipeline.run();
@@ -210,8 +215,8 @@ public class MapElementsTest implements Serializable {
}
};
- MapElements<?, ?> serializableMap = MapElements.via(serializableFn)
- .withOutputType(TypeDescriptor.of(Integer.class));
+ MapElements<?, ?> serializableMap =
+ MapElements.into(TypeDescriptors.integers()).via(serializableFn);
assertThat(DisplayData.from(serializableMap),
hasDisplayItem("mapFn", serializableFn.getClass()));
}
http://git-wip-us.apache.org/repos/asf/beam/blob/3bf6c6b3/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoadBigQuery.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoadBigQuery.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoadBigQuery.java
index c323858..160b231 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoadBigQuery.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoadBigQuery.java
@@ -103,10 +103,9 @@ class BatchLoadBigQuery<T> extends PTransform<PCollection<T>, WriteResult> {
if (write.getFormatFunction() == BigQueryIO.IDENTITY_FORMATTER) {
inputInGlobalWindow = (PCollection<TableRow>) typedInputInGlobalWindow;
} else {
- inputInGlobalWindow = typedInputInGlobalWindow
- .apply(MapElements.via(write.getFormatFunction())
- .withOutputType(new TypeDescriptor<TableRow>() {
- }));
+ inputInGlobalWindow =
+ typedInputInGlobalWindow.apply(
+ MapElements.into(new TypeDescriptor<TableRow>() {}).via(write.getFormatFunction()));
}
// PCollection of filename, file byte size.
http://git-wip-us.apache.org/repos/asf/beam/blob/3bf6c6b3/sdks/java/java8tests/src/test/java/org/apache/beam/sdk/transforms/MapElementsJava8Test.java
----------------------------------------------------------------------
diff --git a/sdks/java/java8tests/src/test/java/org/apache/beam/sdk/transforms/MapElementsJava8Test.java b/sdks/java/java8tests/src/test/java/org/apache/beam/sdk/transforms/MapElementsJava8Test.java
index e0e9d9b4..dbd5ef3 100644
--- a/sdks/java/java8tests/src/test/java/org/apache/beam/sdk/transforms/MapElementsJava8Test.java
+++ b/sdks/java/java8tests/src/test/java/org/apache/beam/sdk/transforms/MapElementsJava8Test.java
@@ -21,7 +21,7 @@ import java.io.Serializable;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.beam.sdk.values.TypeDescriptors;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -47,8 +47,8 @@ public class MapElementsJava8Test implements Serializable {
.apply(Create.of(1, 2, 3))
.apply(MapElements
// Note that the type annotation is required.
- .via((Integer i) -> i * 2)
- .withOutputType(new TypeDescriptor<Integer>() {}));
+ .into(TypeDescriptors.integers())
+ .via((Integer i) -> i * 2));
PAssert.that(output).containsInAnyOrder(6, 2, 4);
pipeline.run();
@@ -82,8 +82,8 @@ public class MapElementsJava8Test implements Serializable {
.apply(Create.of(1, 2, 3))
.apply(MapElements
// Note that the type annotation is required.
- .via(new Doubler()::doubleIt)
- .withOutputType(new TypeDescriptor<Integer>() {}));
+ .into(TypeDescriptors.integers())
+ .via(new Doubler()::doubleIt));
PAssert.that(output).containsInAnyOrder(6, 2, 4);
pipeline.run();
[18/50] [abbrv] beam git commit: DataflowRunner: send windowing
strategy using Runner API proto
Posted by jb...@apache.org.
DataflowRunner: send windowing strategy using Runner API proto
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/69f412dc
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/69f412dc
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/69f412dc
Branch: refs/heads/DSL_SQL
Commit: 69f412dc34f36df40b034c2160b8b0cdad815011
Parents: bef2d37
Author: Dan Halperin <dh...@google.com>
Authored: Tue Apr 11 13:32:11 2017 -0700
Committer: Dan Halperin <dh...@google.com>
Committed: Tue Apr 11 15:04:38 2017 -0700
----------------------------------------------------------------------
runners/google-cloud-dataflow-java/pom.xml | 12 +++++++++++-
.../runners/dataflow/DataflowPipelineTranslator.java | 14 ++++++++++++--
2 files changed, 23 insertions(+), 3 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/69f412dc/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 d0d86e6..a57744c 100644
--- a/runners/google-cloud-dataflow-java/pom.xml
+++ b/runners/google-cloud-dataflow-java/pom.xml
@@ -33,7 +33,7 @@
<packaging>jar</packaging>
<properties>
- <dataflow.container_version>beam-master-20170410</dataflow.container_version>
+ <dataflow.container_version>beam-master-20170411</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>
@@ -114,6 +114,7 @@
<includes>
<include>com.google.cloud.bigtable:bigtable-client-core</include>
<include>com.google.guava:guava</include>
+ <include>org.apache.beam:beam-runners-core-construction-java</include>
</includes>
</artifactSet>
<filters>
@@ -153,6 +154,10 @@
<exclude>com.google.cloud.bigtable.grpc.BigtableTableName</exclude>
</excludes>
</relocation>
+ <relocation>
+ <pattern>org.apache.beam.runners.core</pattern>
+ <shadedPattern>org.apache.beam.runners.dataflow.repackaged.org.apache.beam.runners.core</shadedPattern>
+ </relocation>
</relocations>
<transformers>
<transformer implementation="org.apache.maven.plugins.shade.resource.ServicesResourceTransformer"/>
@@ -178,6 +183,11 @@
<dependency>
<groupId>org.apache.beam</groupId>
+ <artifactId>beam-sdks-common-runner-api</artifactId>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.beam</groupId>
<artifactId>beam-runners-core-construction-java</artifactId>
</dependency>
http://git-wip-us.apache.org/repos/asf/beam/blob/69f412dc/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 34da996..abeca4d 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
@@ -55,6 +55,7 @@ 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.WindowingStrategies;
import org.apache.beam.runners.dataflow.BatchViewOverrides.GroupByKeyAndSortValuesOnly;
import org.apache.beam.runners.dataflow.DataflowRunner.CombineGroupedValues;
import org.apache.beam.runners.dataflow.PrimitiveParDoSingleFactory.ParDoSingle;
@@ -111,6 +112,15 @@ public class DataflowPipelineTranslator {
private static final Logger LOG = LoggerFactory.getLogger(DataflowPipelineTranslator.class);
private static final ObjectMapper MAPPER = new ObjectMapper();
+ private static byte[] serializeWindowingStrategy(WindowingStrategy<?, ?> windowingStrategy) {
+ try {
+ return WindowingStrategies.toProto(windowingStrategy).toByteArray();
+ } catch (Exception e) {
+ throw new RuntimeException(
+ String.format("Unable to format windowing strategy %s as bytes", windowingStrategy), e);
+ }
+ }
+
/**
* A map from {@link PTransform} subclass to the corresponding
* {@link TransformTranslator} to use to translate that transform.
@@ -813,7 +823,7 @@ public class DataflowPipelineTranslator {
stepContext.addInput(PropertyNames.DISALLOW_COMBINER_LIFTING, disallowCombinerLifting);
stepContext.addInput(
PropertyNames.SERIALIZED_FN,
- byteArrayToJsonString(serializeToByteArray(windowingStrategy)));
+ byteArrayToJsonString(serializeWindowingStrategy(windowingStrategy)));
stepContext.addInput(
PropertyNames.IS_MERGING_WINDOW_FN,
!windowingStrategy.getWindowFn().isNonMerging());
@@ -891,7 +901,7 @@ public class DataflowPipelineTranslator {
stepContext.addOutput(context.getOutput(transform));
WindowingStrategy<?, ?> strategy = context.getOutput(transform).getWindowingStrategy();
- byte[] serializedBytes = serializeToByteArray(strategy);
+ byte[] serializedBytes = serializeWindowingStrategy(strategy);
String serializedJson = byteArrayToJsonString(serializedBytes);
stepContext.addInput(PropertyNames.SERIALIZED_FN, serializedJson);
}
[19/50] [abbrv] beam git commit: This closes #2495
Posted by jb...@apache.org.
This closes #2495
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/17a41ab1
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/17a41ab1
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/17a41ab1
Branch: refs/heads/DSL_SQL
Commit: 17a41ab10cd841e495b8f9f634119f5e10be4ca9
Parents: bef2d37 69f412d
Author: Dan Halperin <dh...@google.com>
Authored: Tue Apr 11 15:06:29 2017 -0700
Committer: Dan Halperin <dh...@google.com>
Committed: Tue Apr 11 15:06:29 2017 -0700
----------------------------------------------------------------------
runners/google-cloud-dataflow-java/pom.xml | 12 +++++++++++-
.../runners/dataflow/DataflowPipelineTranslator.java | 14 ++++++++++++--
2 files changed, 23 insertions(+), 3 deletions(-)
----------------------------------------------------------------------
[29/50] [abbrv] beam git commit: [BEAM-1722] Move PubsubIO into the
google-cloud-platform module
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSource.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSource.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSource.java
new file mode 100644
index 0000000..9d8763b
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSource.java
@@ -0,0 +1,1463 @@
+/*
+ * 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.pubsub;
+
+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.api.client.util.Clock;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Charsets;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.security.GeneralSecurityException;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.NoSuchElementException;
+import java.util.Queue;
+import java.util.Set;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+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.ListCoder;
+import org.apache.beam.sdk.coders.NullableCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.io.Read;
+import org.apache.beam.sdk.io.UnboundedSource;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.ProjectPath;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.PubsubClientFactory;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath;
+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.PubsubOptions;
+import org.apache.beam.sdk.options.ValueProvider;
+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.SimpleFunction;
+import org.apache.beam.sdk.transforms.Sum;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.DisplayData.Builder;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.util.BucketingFunction;
+import org.apache.beam.sdk.util.CoderUtils;
+import org.apache.beam.sdk.util.MovingFunction;
+import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.values.PCollection;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A PTransform which streams messages from Pubsub.
+ * <ul>
+ * <li>The underlying implementation in an {@link UnboundedSource} which receives messages
+ * in batches and hands them out one at a time.
+ * <li>The watermark (either in Pubsub processing time or custom timestamp time) is estimated
+ * by keeping track of the minimum of the last minutes worth of messages. This assumes Pubsub
+ * delivers the oldest (in Pubsub processing time) available message at least once a minute,
+ * and that custom timestamps are 'mostly' monotonic with Pubsub processing time. Unfortunately
+ * both of those assumptions are fragile. Thus the estimated watermark may get ahead of
+ * the 'true' watermark and cause some messages to be late.
+ * <li>Checkpoints are used both to ACK received messages back to Pubsub (so that they may
+ * be retired on the Pubsub end), and to NACK already consumed messages should a checkpoint
+ * need to be restored (so that Pubsub will resend those messages promptly).
+ * <li>The backlog is determined by each reader using the messages which have been pulled from
+ * Pubsub but not yet consumed downstream. The backlog does not take account of any messages queued
+ * by Pubsub for the subscription. Unfortunately there is currently no API to determine
+ * the size of the Pubsub queue's backlog.
+ * <li>The subscription must already exist.
+ * <li>The subscription timeout is read whenever a reader is started. However it is not
+ * checked thereafter despite the timeout being user-changeable on-the-fly.
+ * <li>We log vital stats every 30 seconds.
+ * <li>Though some background threads may be used by the underlying transport all Pubsub calls
+ * are blocking. We rely on the underlying runner to allow multiple
+ * {@link UnboundedSource.UnboundedReader} instances to execute concurrently and thus hide latency.
+ * </ul>
+ */
+public class PubsubUnboundedSource<T> extends PTransform<PBegin, PCollection<T>> {
+ private static final Logger LOG = LoggerFactory.getLogger(PubsubUnboundedSource.class);
+
+ /**
+ * Default ACK timeout for created subscriptions.
+ */
+ private static final int DEAULT_ACK_TIMEOUT_SEC = 60;
+
+ /**
+ * Coder for checkpoints.
+ */
+ private static final PubsubCheckpointCoder<?> CHECKPOINT_CODER = new PubsubCheckpointCoder<>();
+
+ /**
+ * Maximum number of messages per pull.
+ */
+ private static final int PULL_BATCH_SIZE = 1000;
+
+ /**
+ * Maximum number of ACK ids per ACK or ACK extension call.
+ */
+ private static final int ACK_BATCH_SIZE = 2000;
+
+ /**
+ * Maximum number of messages in flight.
+ */
+ private static final int MAX_IN_FLIGHT = 20000;
+
+ /**
+ * Timeout for round trip from receiving a message to finally ACKing it back to Pubsub.
+ */
+ private static final Duration PROCESSING_TIMEOUT = Duration.standardSeconds(120);
+
+ /**
+ * Percentage of ack timeout by which to extend acks when they are near timeout.
+ */
+ private static final int ACK_EXTENSION_PCT = 50;
+
+ /**
+ * Percentage of ack timeout we should use as a safety margin. We'll try to extend acks
+ * by this margin before the ack actually expires.
+ */
+ private static final int ACK_SAFETY_PCT = 20;
+
+ /**
+ * For stats only: How close we can get to an ack deadline before we risk it being already
+ * considered passed by Pubsub.
+ */
+ private static final Duration ACK_TOO_LATE = Duration.standardSeconds(2);
+
+ /**
+ * Period of samples to determine watermark and other stats.
+ */
+ private static final Duration SAMPLE_PERIOD = Duration.standardMinutes(1);
+
+ /**
+ * Period of updates to determine watermark and other stats.
+ */
+ private static final Duration SAMPLE_UPDATE = Duration.standardSeconds(5);
+
+ /**
+ * Period for logging stats.
+ */
+ private static final Duration LOG_PERIOD = Duration.standardSeconds(30);
+
+ /**
+ * Minimum number of unread messages required before considering updating watermark.
+ */
+ private static final int MIN_WATERMARK_MESSAGES = 10;
+
+ /**
+ * Minimum number of SAMPLE_UPDATE periods over which unread messages shoud be spread
+ * before considering updating watermark.
+ */
+ private static final int MIN_WATERMARK_SPREAD = 2;
+
+ /**
+ * Additional sharding so that we can hide read message latency.
+ */
+ private static final int SCALE_OUT = 4;
+
+ // TODO: Would prefer to use MinLongFn but it is a BinaryCombineFn<Long> rather
+ // than a BinaryCombineLongFn. [BEAM-285]
+ private static final Combine.BinaryCombineLongFn MIN =
+ new Combine.BinaryCombineLongFn() {
+ @Override
+ public long apply(long left, long right) {
+ return Math.min(left, right);
+ }
+
+ @Override
+ public long identity() {
+ return Long.MAX_VALUE;
+ }
+ };
+
+ private static final Combine.BinaryCombineLongFn MAX =
+ new Combine.BinaryCombineLongFn() {
+ @Override
+ public long apply(long left, long right) {
+ return Math.max(left, right);
+ }
+
+ @Override
+ public long identity() {
+ return Long.MIN_VALUE;
+ }
+ };
+
+ private static final Combine.BinaryCombineLongFn SUM = Sum.ofLongs();
+
+ // ================================================================================
+ // Checkpoint
+ // ================================================================================
+
+ /**
+ * Which messages have been durably committed and thus can now be ACKed.
+ * Which messages have been read but not yet committed, in which case they should be NACKed if
+ * we need to restore.
+ */
+ @VisibleForTesting
+ static class PubsubCheckpoint<T> implements UnboundedSource.CheckpointMark {
+ /**
+ * The {@link SubscriptionPath} to the subscription the reader is reading from. May be
+ * {@code null} if the {@link PubsubUnboundedSource} contains the subscription.
+ */
+ @VisibleForTesting
+ @Nullable String subscriptionPath;
+
+ /**
+ * If the checkpoint is for persisting: the reader who's snapshotted state we are persisting.
+ * If the checkpoint is for restoring: {@literal null}.
+ * Not persisted in durable checkpoint.
+ * CAUTION: Between a checkpoint being taken and {@link #finalizeCheckpoint()} being called
+ * the 'true' active reader may have changed.
+ */
+ @Nullable
+ private PubsubReader<T> reader;
+
+ /**
+ * If the checkpoint is for persisting: The ACK ids of messages which have been passed
+ * downstream since the last checkpoint.
+ * If the checkpoint is for restoring: {@literal null}.
+ * Not persisted in durable checkpoint.
+ */
+ @Nullable
+ private List<String> safeToAckIds;
+
+ /**
+ * If the checkpoint is for persisting: The ACK ids of messages which have been received
+ * from Pubsub but not yet passed downstream at the time of the snapshot.
+ * If the checkpoint is for restoring: Same, but recovered from durable storage.
+ */
+ @VisibleForTesting
+ final List<String> notYetReadIds;
+
+ public PubsubCheckpoint(
+ @Nullable String subscriptionPath,
+ @Nullable PubsubReader<T> reader,
+ @Nullable List<String> safeToAckIds,
+ List<String> notYetReadIds) {
+ this.subscriptionPath = subscriptionPath;
+ this.reader = reader;
+ this.safeToAckIds = safeToAckIds;
+ this.notYetReadIds = notYetReadIds;
+ }
+
+ @Nullable
+ private SubscriptionPath getSubscription() {
+ return subscriptionPath == null
+ ? null
+ : PubsubClient.subscriptionPathFromPath(subscriptionPath);
+ }
+
+ /**
+ * BLOCKING
+ * All messages which have been passed downstream have now been durably committed.
+ * We can ACK them upstream.
+ * CAUTION: This may never be called.
+ */
+ @Override
+ public void finalizeCheckpoint() throws IOException {
+ checkState(reader != null && safeToAckIds != null, "Cannot finalize a restored checkpoint");
+ // Even if the 'true' active reader has changed since the checkpoint was taken we are
+ // fine:
+ // - The underlying Pubsub topic will not have changed, so the following ACKs will still
+ // go to the right place.
+ // - We'll delete the ACK ids from the readers in-flight state, but that only effects
+ // flow control and stats, neither of which are relevant anymore.
+ try {
+ int n = safeToAckIds.size();
+ List<String> batchSafeToAckIds = new ArrayList<>(Math.min(n, ACK_BATCH_SIZE));
+ for (String ackId : safeToAckIds) {
+ batchSafeToAckIds.add(ackId);
+ if (batchSafeToAckIds.size() >= ACK_BATCH_SIZE) {
+ reader.ackBatch(batchSafeToAckIds);
+ n -= batchSafeToAckIds.size();
+ // CAUTION: Don't reuse the same list since ackBatch holds on to it.
+ batchSafeToAckIds = new ArrayList<>(Math.min(n, ACK_BATCH_SIZE));
+ }
+ }
+ if (!batchSafeToAckIds.isEmpty()) {
+ reader.ackBatch(batchSafeToAckIds);
+ }
+ } finally {
+ int remainingInFlight = reader.numInFlightCheckpoints.decrementAndGet();
+ checkState(remainingInFlight >= 0,
+ "Miscounted in-flight checkpoints");
+ reader.maybeCloseClient();
+ reader = null;
+ safeToAckIds = null;
+ }
+ }
+
+ /**
+ * Return current time according to {@code reader}.
+ */
+ private static long now(PubsubReader<?> reader) {
+ if (reader.outer.outer.clock == null) {
+ return System.currentTimeMillis();
+ } else {
+ return reader.outer.outer.clock.currentTimeMillis();
+ }
+ }
+
+ /**
+ * BLOCKING
+ * NACK all messages which have been read from Pubsub but not passed downstream.
+ * This way Pubsub will send them again promptly.
+ */
+ public void nackAll(PubsubReader<T> reader) throws IOException {
+ checkState(this.reader == null, "Cannot nackAll on persisting checkpoint");
+ List<String> batchYetToAckIds =
+ new ArrayList<>(Math.min(notYetReadIds.size(), ACK_BATCH_SIZE));
+ for (String ackId : notYetReadIds) {
+ batchYetToAckIds.add(ackId);
+ if (batchYetToAckIds.size() >= ACK_BATCH_SIZE) {
+ long nowMsSinceEpoch = now(reader);
+ reader.nackBatch(nowMsSinceEpoch, batchYetToAckIds);
+ batchYetToAckIds.clear();
+ }
+ }
+ if (!batchYetToAckIds.isEmpty()) {
+ long nowMsSinceEpoch = now(reader);
+ reader.nackBatch(nowMsSinceEpoch, batchYetToAckIds);
+ }
+ }
+ }
+
+ /** The coder for our checkpoints. */
+ private static class PubsubCheckpointCoder<T> extends AtomicCoder<PubsubCheckpoint<T>> {
+ private static final Coder<String> SUBSCRIPTION_PATH_CODER =
+ NullableCoder.of(StringUtf8Coder.of());
+ private static final Coder<List<String>> LIST_CODER = ListCoder.of(StringUtf8Coder.of());
+
+ @Override
+ public void encode(PubsubCheckpoint<T> value, OutputStream outStream, Context context)
+ throws IOException {
+ SUBSCRIPTION_PATH_CODER.encode(
+ value.subscriptionPath,
+ outStream,
+ context.nested());
+ LIST_CODER.encode(value.notYetReadIds, outStream, context);
+ }
+
+ @Override
+ public PubsubCheckpoint<T> decode(InputStream inStream, Context context) throws IOException {
+ String path = SUBSCRIPTION_PATH_CODER.decode(inStream, context.nested());
+ List<String> notYetReadIds = LIST_CODER.decode(inStream, context);
+ return new PubsubCheckpoint<>(path, null, null, notYetReadIds);
+ }
+ }
+
+ // ================================================================================
+ // Reader
+ // ================================================================================
+
+ /**
+ * A reader which keeps track of which messages have been received from Pubsub
+ * but not yet consumed downstream and/or ACKed back to Pubsub.
+ */
+ @VisibleForTesting
+ static class PubsubReader<T> extends UnboundedSource.UnboundedReader<T> {
+ /**
+ * For access to topic and checkpointCoder.
+ */
+ private final PubsubSource<T> outer;
+ @VisibleForTesting
+ final SubscriptionPath subscription;
+
+ private final SimpleFunction<PubsubIO.PubsubMessage, T> parseFn;
+
+ /**
+ * Client on which to talk to Pubsub. Contains a null value if the client has been closed.
+ */
+ private AtomicReference<PubsubClient> pubsubClient;
+
+ /**
+ * The closed state of this {@link PubsubReader}. If true, the reader has not yet been closed,
+ * and it will have a non-null value within {@link #pubsubClient}.
+ */
+ private AtomicBoolean active = new AtomicBoolean(true);
+
+ /**
+ * Ack timeout, in ms, as set on subscription when we first start reading. Not
+ * updated thereafter. -1 if not yet determined.
+ */
+ private int ackTimeoutMs;
+
+ /**
+ * ACK ids of messages we have delivered downstream but not yet ACKed.
+ */
+ private Set<String> safeToAckIds;
+
+ /**
+ * Messages we have received from Pubsub and not yet delivered downstream.
+ * We preserve their order.
+ */
+ private final Queue<PubsubClient.IncomingMessage> notYetRead;
+
+ private static class InFlightState {
+ /**
+ * When request which yielded message was issues.
+ */
+ long requestTimeMsSinceEpoch;
+
+ /**
+ * When Pubsub will consider this message's ACK to timeout and thus it needs to be
+ * extended.
+ */
+ long ackDeadlineMsSinceEpoch;
+
+ public InFlightState(long requestTimeMsSinceEpoch, long ackDeadlineMsSinceEpoch) {
+ this.requestTimeMsSinceEpoch = requestTimeMsSinceEpoch;
+ this.ackDeadlineMsSinceEpoch = ackDeadlineMsSinceEpoch;
+ }
+ }
+
+ /**
+ * Map from ACK ids of messages we have received from Pubsub but not yet ACKed to their
+ * in flight state. Ordered from earliest to latest ACK deadline.
+ */
+ private final LinkedHashMap<String, InFlightState> inFlight;
+
+ /**
+ * Batches of successfully ACKed ids which need to be pruned from the above.
+ * CAUTION: Accessed by both reader and checkpointing threads.
+ */
+ private final Queue<List<String>> ackedIds;
+
+ /**
+ * Byte size of undecoded elements in {@link #notYetRead}.
+ */
+ private long notYetReadBytes;
+
+ /**
+ * Bucketed map from received time (as system time, ms since epoch) to message
+ * timestamps (mssince epoch) of all received but not-yet read messages.
+ * Used to estimate watermark.
+ */
+ private BucketingFunction minUnreadTimestampMsSinceEpoch;
+
+ /**
+ * Minimum of timestamps (ms since epoch) of all recently read messages.
+ * Used to estimate watermark.
+ */
+ private MovingFunction minReadTimestampMsSinceEpoch;
+
+ /**
+ * System time (ms since epoch) we last received a message from Pubsub, or -1 if
+ * not yet received any messages.
+ */
+ private long lastReceivedMsSinceEpoch;
+
+ /**
+ * The last reported watermark (ms since epoch), or beginning of time if none yet reported.
+ */
+ private long lastWatermarkMsSinceEpoch;
+
+ /**
+ * The current message, or {@literal null} if none.
+ */
+ @Nullable
+ private PubsubClient.IncomingMessage current;
+
+ /**
+ * Stats only: System time (ms since epoch) we last logs stats, or -1 if never.
+ */
+ private long lastLogTimestampMsSinceEpoch;
+
+ /**
+ * Stats only: Total number of messages received.
+ */
+ private long numReceived;
+
+ /**
+ * Stats only: Number of messages which have recently been received.
+ */
+ private MovingFunction numReceivedRecently;
+
+ /**
+ * Stats only: Number of messages which have recently had their deadline extended.
+ */
+ private MovingFunction numExtendedDeadlines;
+
+ /**
+ * Stats only: Number of messages which have recenttly had their deadline extended even
+ * though it may be too late to do so.
+ */
+ private MovingFunction numLateDeadlines;
+
+
+ /**
+ * Stats only: Number of messages which have recently been ACKed.
+ */
+ private MovingFunction numAcked;
+
+ /**
+ * Stats only: Number of messages which have recently expired (ACKs were extended for too
+ * long).
+ */
+ private MovingFunction numExpired;
+
+ /**
+ * Stats only: Number of messages which have recently been NACKed.
+ */
+ private MovingFunction numNacked;
+
+ /**
+ * Stats only: Number of message bytes which have recently been read by downstream consumer.
+ */
+ private MovingFunction numReadBytes;
+
+ /**
+ * Stats only: Minimum of timestamp (ms since epoch) of all recently received messages.
+ * Used to estimate timestamp skew. Does not contribute to watermark estimator.
+ */
+ private MovingFunction minReceivedTimestampMsSinceEpoch;
+
+ /**
+ * Stats only: Maximum of timestamp (ms since epoch) of all recently received messages.
+ * Used to estimate timestamp skew.
+ */
+ private MovingFunction maxReceivedTimestampMsSinceEpoch;
+
+ /**
+ * Stats only: Minimum of recent estimated watermarks (ms since epoch).
+ */
+ private MovingFunction minWatermarkMsSinceEpoch;
+
+ /**
+ * Stats ony: Maximum of recent estimated watermarks (ms since epoch).
+ */
+ private MovingFunction maxWatermarkMsSinceEpoch;
+
+ /**
+ * Stats only: Number of messages with timestamps strictly behind the estimated watermark
+ * at the time they are received. These may be considered 'late' by downstream computations.
+ */
+ private MovingFunction numLateMessages;
+
+ /**
+ * Stats only: Current number of checkpoints in flight.
+ * CAUTION: Accessed by both checkpointing and reader threads.
+ */
+ private AtomicInteger numInFlightCheckpoints;
+
+ /**
+ * Stats only: Maximum number of checkpoints in flight at any time.
+ */
+ private int maxInFlightCheckpoints;
+
+ private static MovingFunction newFun(Combine.BinaryCombineLongFn function) {
+ return new MovingFunction(SAMPLE_PERIOD.getMillis(),
+ SAMPLE_UPDATE.getMillis(),
+ MIN_WATERMARK_SPREAD,
+ MIN_WATERMARK_MESSAGES,
+ function);
+ }
+
+ /**
+ * Construct a reader.
+ */
+ public PubsubReader(PubsubOptions options, PubsubSource<T> outer, SubscriptionPath subscription,
+ SimpleFunction<PubsubIO.PubsubMessage, T> parseFn)
+ throws IOException, GeneralSecurityException {
+ this.outer = outer;
+ this.subscription = subscription;
+ this.parseFn = parseFn;
+ pubsubClient =
+ new AtomicReference<>(
+ outer.outer.pubsubFactory.newClient(
+ outer.outer.timestampLabel, outer.outer.idLabel, options));
+ ackTimeoutMs = -1;
+ safeToAckIds = new HashSet<>();
+ notYetRead = new ArrayDeque<>();
+ inFlight = new LinkedHashMap<>();
+ ackedIds = new ConcurrentLinkedQueue<>();
+ notYetReadBytes = 0;
+ minUnreadTimestampMsSinceEpoch = new BucketingFunction(SAMPLE_UPDATE.getMillis(),
+ MIN_WATERMARK_SPREAD,
+ MIN_WATERMARK_MESSAGES,
+ MIN);
+ minReadTimestampMsSinceEpoch = newFun(MIN);
+ lastReceivedMsSinceEpoch = -1;
+ lastWatermarkMsSinceEpoch = BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis();
+ current = null;
+ lastLogTimestampMsSinceEpoch = -1;
+ numReceived = 0L;
+ numReceivedRecently = newFun(SUM);
+ numExtendedDeadlines = newFun(SUM);
+ numLateDeadlines = newFun(SUM);
+ numAcked = newFun(SUM);
+ numExpired = newFun(SUM);
+ numNacked = newFun(SUM);
+ numReadBytes = newFun(SUM);
+ minReceivedTimestampMsSinceEpoch = newFun(MIN);
+ maxReceivedTimestampMsSinceEpoch = newFun(MAX);
+ minWatermarkMsSinceEpoch = newFun(MIN);
+ maxWatermarkMsSinceEpoch = newFun(MAX);
+ numLateMessages = newFun(SUM);
+ numInFlightCheckpoints = new AtomicInteger();
+ maxInFlightCheckpoints = 0;
+ }
+
+ @VisibleForTesting
+ PubsubClient getPubsubClient() {
+ return pubsubClient.get();
+ }
+
+ /**
+ * Acks the provided {@code ackIds} back to Pubsub, blocking until all of the messages are
+ * ACKed.
+ *
+ * <p>CAUTION: May be invoked from a separate thread.
+ *
+ * <p>CAUTION: Retains {@code ackIds}.
+ */
+ void ackBatch(List<String> ackIds) throws IOException {
+ pubsubClient.get().acknowledge(subscription, ackIds);
+ ackedIds.add(ackIds);
+ }
+
+ /**
+ * BLOCKING
+ * NACK (ie request deadline extension of 0) receipt of messages from Pubsub
+ * with the given {@code ockIds}. Does not retain {@code ackIds}.
+ */
+ public void nackBatch(long nowMsSinceEpoch, List<String> ackIds) throws IOException {
+ pubsubClient.get().modifyAckDeadline(subscription, ackIds, 0);
+ numNacked.add(nowMsSinceEpoch, ackIds.size());
+ }
+
+ /**
+ * BLOCKING
+ * Extend the processing deadline for messages from Pubsub with the given {@code ackIds}.
+ * Does not retain {@code ackIds}.
+ */
+ private void extendBatch(long nowMsSinceEpoch, List<String> ackIds) throws IOException {
+ int extensionSec = (ackTimeoutMs * ACK_EXTENSION_PCT) / (100 * 1000);
+ pubsubClient.get().modifyAckDeadline(subscription, ackIds, extensionSec);
+ numExtendedDeadlines.add(nowMsSinceEpoch, ackIds.size());
+ }
+
+ /**
+ * Return the current time, in ms since epoch.
+ */
+ private long now() {
+ if (outer.outer.clock == null) {
+ return System.currentTimeMillis();
+ } else {
+ return outer.outer.clock.currentTimeMillis();
+ }
+ }
+
+ /**
+ * Messages which have been ACKed (via the checkpoint finalize) are no longer in flight.
+ * This is only used for flow control and stats.
+ */
+ private void retire() throws IOException {
+ long nowMsSinceEpoch = now();
+ while (true) {
+ List<String> ackIds = ackedIds.poll();
+ if (ackIds == null) {
+ return;
+ }
+ numAcked.add(nowMsSinceEpoch, ackIds.size());
+ for (String ackId : ackIds) {
+ inFlight.remove(ackId);
+ safeToAckIds.remove(ackId);
+ }
+ }
+ }
+
+ /**
+ * BLOCKING
+ * Extend deadline for all messages which need it.
+ * CAUTION: If extensions can't keep up with wallclock then we'll never return.
+ */
+ private void extend() throws IOException {
+ while (true) {
+ long nowMsSinceEpoch = now();
+ List<String> assumeExpired = new ArrayList<>();
+ List<String> toBeExtended = new ArrayList<>();
+ List<String> toBeExpired = new ArrayList<>();
+ // Messages will be in increasing deadline order.
+ for (Map.Entry<String, InFlightState> entry : inFlight.entrySet()) {
+ if (entry.getValue().ackDeadlineMsSinceEpoch - (ackTimeoutMs * ACK_SAFETY_PCT) / 100
+ > nowMsSinceEpoch) {
+ // All remaining messages don't need their ACKs to be extended.
+ break;
+ }
+
+ if (entry.getValue().ackDeadlineMsSinceEpoch - ACK_TOO_LATE.getMillis()
+ < nowMsSinceEpoch) {
+ // Pubsub may have already considered this message to have expired.
+ // If so it will (eventually) be made available on a future pull request.
+ // If this message ends up being committed then it will be considered a duplicate
+ // when re-pulled.
+ assumeExpired.add(entry.getKey());
+ continue;
+ }
+
+ if (entry.getValue().requestTimeMsSinceEpoch + PROCESSING_TIMEOUT.getMillis()
+ < nowMsSinceEpoch) {
+ // This message has been in-flight for too long.
+ // Give up on it, otherwise we risk extending its ACK indefinitely.
+ toBeExpired.add(entry.getKey());
+ continue;
+ }
+
+ // Extend the ACK for this message.
+ toBeExtended.add(entry.getKey());
+ if (toBeExtended.size() >= ACK_BATCH_SIZE) {
+ // Enough for one batch.
+ break;
+ }
+ }
+
+ if (assumeExpired.isEmpty() && toBeExtended.isEmpty() && toBeExpired.isEmpty()) {
+ // Nothing to be done.
+ return;
+ }
+
+ if (!assumeExpired.isEmpty()) {
+ // If we didn't make the ACK deadline assume expired and no longer in flight.
+ numLateDeadlines.add(nowMsSinceEpoch, assumeExpired.size());
+ for (String ackId : assumeExpired) {
+ inFlight.remove(ackId);
+ }
+ }
+
+ if (!toBeExpired.isEmpty()) {
+ // Expired messages are no longer considered in flight.
+ numExpired.add(nowMsSinceEpoch, toBeExpired.size());
+ for (String ackId : toBeExpired) {
+ inFlight.remove(ackId);
+ }
+ }
+
+ if (!toBeExtended.isEmpty()) {
+ // Pubsub extends acks from it's notion of current time.
+ // We'll try to track that on our side, but note the deadlines won't necessarily agree.
+ long newDeadlineMsSinceEpoch = nowMsSinceEpoch + (ackTimeoutMs * ACK_EXTENSION_PCT) / 100;
+ for (String ackId : toBeExtended) {
+ // Maintain increasing ack deadline order.
+ InFlightState state = inFlight.remove(ackId);
+ inFlight.put(ackId,
+ new InFlightState(state.requestTimeMsSinceEpoch, newDeadlineMsSinceEpoch));
+ }
+ // BLOCKs until extended.
+ extendBatch(nowMsSinceEpoch, toBeExtended);
+ }
+ }
+ }
+
+ /**
+ * BLOCKING
+ * Fetch another batch of messages from Pubsub.
+ */
+ private void pull() throws IOException {
+ if (inFlight.size() >= MAX_IN_FLIGHT) {
+ // Wait for checkpoint to be finalized before pulling anymore.
+ // There may be lag while checkpoints are persisted and the finalizeCheckpoint method
+ // is invoked. By limiting the in-flight messages we can ensure we don't end up consuming
+ // messages faster than we can checkpoint them.
+ return;
+ }
+
+ long requestTimeMsSinceEpoch = now();
+ long deadlineMsSinceEpoch = requestTimeMsSinceEpoch + ackTimeoutMs;
+
+ // Pull the next batch.
+ // BLOCKs until received.
+ Collection<PubsubClient.IncomingMessage> receivedMessages =
+ pubsubClient.get().pull(requestTimeMsSinceEpoch, subscription, PULL_BATCH_SIZE, true);
+ if (receivedMessages.isEmpty()) {
+ // Nothing available yet. Try again later.
+ return;
+ }
+
+ lastReceivedMsSinceEpoch = requestTimeMsSinceEpoch;
+
+ // Capture the received messages.
+ for (PubsubClient.IncomingMessage incomingMessage : receivedMessages) {
+ notYetRead.add(incomingMessage);
+ notYetReadBytes += incomingMessage.elementBytes.length;
+ inFlight.put(incomingMessage.ackId,
+ new InFlightState(requestTimeMsSinceEpoch, deadlineMsSinceEpoch));
+ numReceived++;
+ numReceivedRecently.add(requestTimeMsSinceEpoch, 1L);
+ minReceivedTimestampMsSinceEpoch.add(requestTimeMsSinceEpoch,
+ incomingMessage.timestampMsSinceEpoch);
+ maxReceivedTimestampMsSinceEpoch.add(requestTimeMsSinceEpoch,
+ incomingMessage.timestampMsSinceEpoch);
+ minUnreadTimestampMsSinceEpoch.add(requestTimeMsSinceEpoch,
+ incomingMessage.timestampMsSinceEpoch);
+ }
+ }
+
+ /**
+ * Log stats if time to do so.
+ */
+ private void stats() {
+ long nowMsSinceEpoch = now();
+ if (lastLogTimestampMsSinceEpoch < 0) {
+ lastLogTimestampMsSinceEpoch = nowMsSinceEpoch;
+ return;
+ }
+ long deltaMs = nowMsSinceEpoch - lastLogTimestampMsSinceEpoch;
+ if (deltaMs < LOG_PERIOD.getMillis()) {
+ return;
+ }
+
+ String messageSkew = "unknown";
+ long minTimestamp = minReceivedTimestampMsSinceEpoch.get(nowMsSinceEpoch);
+ long maxTimestamp = maxReceivedTimestampMsSinceEpoch.get(nowMsSinceEpoch);
+ if (minTimestamp < Long.MAX_VALUE && maxTimestamp > Long.MIN_VALUE) {
+ messageSkew = (maxTimestamp - minTimestamp) + "ms";
+ }
+
+ String watermarkSkew = "unknown";
+ long minWatermark = minWatermarkMsSinceEpoch.get(nowMsSinceEpoch);
+ long maxWatermark = maxWatermarkMsSinceEpoch.get(nowMsSinceEpoch);
+ if (minWatermark < Long.MAX_VALUE && maxWatermark > Long.MIN_VALUE) {
+ watermarkSkew = (maxWatermark - minWatermark) + "ms";
+ }
+
+ String oldestInFlight = "no";
+ String oldestAckId = Iterables.getFirst(inFlight.keySet(), null);
+ if (oldestAckId != null) {
+ oldestInFlight =
+ (nowMsSinceEpoch - inFlight.get(oldestAckId).requestTimeMsSinceEpoch) + "ms";
+ }
+
+ LOG.info("Pubsub {} has "
+ + "{} received messages, "
+ + "{} current unread messages, "
+ + "{} current unread bytes, "
+ + "{} current in-flight msgs, "
+ + "{} oldest in-flight, "
+ + "{} current in-flight checkpoints, "
+ + "{} max in-flight checkpoints, "
+ + "{}B/s recent read, "
+ + "{} recent received, "
+ + "{} recent extended, "
+ + "{} recent late extended, "
+ + "{} recent ACKed, "
+ + "{} recent NACKed, "
+ + "{} recent expired, "
+ + "{} recent message timestamp skew, "
+ + "{} recent watermark skew, "
+ + "{} recent late messages, "
+ + "{} last reported watermark",
+ subscription,
+ numReceived,
+ notYetRead.size(),
+ notYetReadBytes,
+ inFlight.size(),
+ oldestInFlight,
+ numInFlightCheckpoints.get(),
+ maxInFlightCheckpoints,
+ numReadBytes.get(nowMsSinceEpoch) / (SAMPLE_PERIOD.getMillis() / 1000L),
+ numReceivedRecently.get(nowMsSinceEpoch),
+ numExtendedDeadlines.get(nowMsSinceEpoch),
+ numLateDeadlines.get(nowMsSinceEpoch),
+ numAcked.get(nowMsSinceEpoch),
+ numNacked.get(nowMsSinceEpoch),
+ numExpired.get(nowMsSinceEpoch),
+ messageSkew,
+ watermarkSkew,
+ numLateMessages.get(nowMsSinceEpoch),
+ new Instant(lastWatermarkMsSinceEpoch));
+
+ lastLogTimestampMsSinceEpoch = nowMsSinceEpoch;
+ }
+
+ @Override
+ public boolean start() throws IOException {
+ // Determine the ack timeout.
+ ackTimeoutMs = pubsubClient.get().ackDeadlineSeconds(subscription) * 1000;
+ return advance();
+ }
+
+ /**
+ * BLOCKING
+ * Return {@literal true} if a Pubsub messaage is available, {@literal false} if
+ * none is available at this time or we are over-subscribed. May BLOCK while extending
+ * ACKs or fetching available messages. Will not block waiting for messages.
+ */
+ @Override
+ public boolean advance() throws IOException {
+ // Emit stats.
+ stats();
+
+ if (current != null) {
+ // Current is consumed. It can no longer contribute to holding back the watermark.
+ minUnreadTimestampMsSinceEpoch.remove(current.requestTimeMsSinceEpoch);
+ current = null;
+ }
+
+ // Retire state associated with ACKed messages.
+ retire();
+
+ // Extend all pressing deadlines.
+ // Will BLOCK until done.
+ // If the system is pulling messages only to let them sit in a downsteam queue then
+ // this will have the effect of slowing down the pull rate.
+ // However, if the system is genuinely taking longer to process each message then
+ // the work to extend ACKs would be better done in the background.
+ extend();
+
+ if (notYetRead.isEmpty()) {
+ // Pull another batch.
+ // Will BLOCK until fetch returns, but will not block until a message is available.
+ pull();
+ }
+
+ // Take one message from queue.
+ current = notYetRead.poll();
+ if (current == null) {
+ // Try again later.
+ return false;
+ }
+ notYetReadBytes -= current.elementBytes.length;
+ checkState(notYetReadBytes >= 0);
+ long nowMsSinceEpoch = now();
+ numReadBytes.add(nowMsSinceEpoch, current.elementBytes.length);
+ minReadTimestampMsSinceEpoch.add(nowMsSinceEpoch, current.timestampMsSinceEpoch);
+ if (current.timestampMsSinceEpoch < lastWatermarkMsSinceEpoch) {
+ numLateMessages.add(nowMsSinceEpoch, 1L);
+ }
+
+ // Current message can be considered 'read' and will be persisted by the next
+ // checkpoint. So it is now safe to ACK back to Pubsub.
+ safeToAckIds.add(current.ackId);
+ return true;
+ }
+
+ @Override
+ public T getCurrent() throws NoSuchElementException {
+ if (current == null) {
+ throw new NoSuchElementException();
+ }
+ try {
+ if (parseFn != null) {
+ return parseFn.apply(new PubsubIO.PubsubMessage(
+ current.elementBytes, current.attributes));
+ } else {
+ return CoderUtils.decodeFromByteArray(outer.outer.elementCoder, current.elementBytes);
+ }
+ } catch (CoderException e) {
+ throw new RuntimeException("Unable to decode element from Pubsub message: ", e);
+ }
+ }
+
+ @Override
+ public Instant getCurrentTimestamp() throws NoSuchElementException {
+ if (current == null) {
+ throw new NoSuchElementException();
+ }
+ return new Instant(current.timestampMsSinceEpoch);
+ }
+
+ @Override
+ public byte[] getCurrentRecordId() throws NoSuchElementException {
+ if (current == null) {
+ throw new NoSuchElementException();
+ }
+ return current.recordId.getBytes(Charsets.UTF_8);
+ }
+
+ /**
+ * {@inheritDoc}.
+ *
+ * <p>Marks this {@link PubsubReader} as no longer active. The {@link PubsubClient}
+ * continue to exist and be active beyond the life of this call if there are any in-flight
+ * checkpoints. When no in-flight checkpoints remain, the reader will be closed.
+ */
+ @Override
+ public void close() throws IOException {
+ active.set(false);
+ maybeCloseClient();
+ }
+
+ /**
+ * Close this reader's underlying {@link PubsubClient} if the reader has been closed and there
+ * are no outstanding checkpoints.
+ */
+ private void maybeCloseClient() throws IOException {
+ if (!active.get() && numInFlightCheckpoints.get() == 0) {
+ // The reader has been closed and it has no more outstanding checkpoints. The client
+ // must be closed so it doesn't leak
+ PubsubClient client = pubsubClient.getAndSet(null);
+ if (client != null) {
+ client.close();
+ }
+ }
+ }
+
+ @Override
+ public PubsubSource<T> getCurrentSource() {
+ return outer;
+ }
+
+ @Override
+ public Instant getWatermark() {
+ if (pubsubClient.get().isEOF() && notYetRead.isEmpty()) {
+ // For testing only: Advance the watermark to the end of time to signal
+ // the test is complete.
+ return BoundedWindow.TIMESTAMP_MAX_VALUE;
+ }
+
+ // NOTE: We'll allow the watermark to go backwards. The underlying runner is responsible
+ // for aggregating all reported watermarks and ensuring the aggregate is latched.
+ // If we attempt to latch locally then it is possible a temporary starvation of one reader
+ // could cause its estimated watermark to fast forward to current system time. Then when
+ // the reader resumes its watermark would be unable to resume tracking.
+ // By letting the underlying runner latch we avoid any problems due to localized starvation.
+ long nowMsSinceEpoch = now();
+ long readMin = minReadTimestampMsSinceEpoch.get(nowMsSinceEpoch);
+ long unreadMin = minUnreadTimestampMsSinceEpoch.get();
+ if (readMin == Long.MAX_VALUE
+ && unreadMin == Long.MAX_VALUE
+ && lastReceivedMsSinceEpoch >= 0
+ && nowMsSinceEpoch > lastReceivedMsSinceEpoch + SAMPLE_PERIOD.getMillis()) {
+ // We don't currently have any unread messages pending, we have not had any messages
+ // read for a while, and we have not received any new messages from Pubsub for a while.
+ // Advance watermark to current time.
+ // TODO: Estimate a timestamp lag.
+ lastWatermarkMsSinceEpoch = nowMsSinceEpoch;
+ } else if (minReadTimestampMsSinceEpoch.isSignificant()
+ || minUnreadTimestampMsSinceEpoch.isSignificant()) {
+ // Take minimum of the timestamps in all unread messages and recently read messages.
+ lastWatermarkMsSinceEpoch = Math.min(readMin, unreadMin);
+ }
+ // else: We're not confident enough to estimate a new watermark. Stick with the old one.
+ minWatermarkMsSinceEpoch.add(nowMsSinceEpoch, lastWatermarkMsSinceEpoch);
+ maxWatermarkMsSinceEpoch.add(nowMsSinceEpoch, lastWatermarkMsSinceEpoch);
+ return new Instant(lastWatermarkMsSinceEpoch);
+ }
+
+ @Override
+ public PubsubCheckpoint<T> getCheckpointMark() {
+ int cur = numInFlightCheckpoints.incrementAndGet();
+ maxInFlightCheckpoints = Math.max(maxInFlightCheckpoints, cur);
+ // It's possible for a checkpoint to be taken but never finalized.
+ // So we simply copy whatever safeToAckIds we currently have.
+ // It's possible a later checkpoint will be taken before an earlier one is finalized,
+ // in which case we'll double ACK messages to Pubsub. However Pubsub is fine with that.
+ List<String> snapshotSafeToAckIds = Lists.newArrayList(safeToAckIds);
+ List<String> snapshotNotYetReadIds = new ArrayList<>(notYetRead.size());
+ for (PubsubClient.IncomingMessage incomingMessage : notYetRead) {
+ snapshotNotYetReadIds.add(incomingMessage.ackId);
+ }
+ if (outer.subscriptionPath == null) {
+ // need to include the subscription in case we resume, as it's not stored in the source.
+ return new PubsubCheckpoint<>(
+ subscription.getPath(), this, snapshotSafeToAckIds, snapshotNotYetReadIds);
+ }
+ return new PubsubCheckpoint<>(null, this, snapshotSafeToAckIds, snapshotNotYetReadIds);
+ }
+
+ @Override
+ public long getSplitBacklogBytes() {
+ return notYetReadBytes;
+ }
+ }
+
+ // ================================================================================
+ // Source
+ // ================================================================================
+
+ @VisibleForTesting
+ static class PubsubSource<T> extends UnboundedSource<T, PubsubCheckpoint<T>> {
+ public final PubsubUnboundedSource<T> outer;
+ // The subscription to read from.
+ @VisibleForTesting
+ final SubscriptionPath subscriptionPath;
+
+ public PubsubSource(PubsubUnboundedSource<T> outer) {
+ this(outer, outer.getSubscription());
+ }
+
+ private PubsubSource(PubsubUnboundedSource<T> outer, SubscriptionPath subscriptionPath) {
+ this.outer = outer;
+ this.subscriptionPath = subscriptionPath;
+ }
+
+ @Override
+ public List<PubsubSource<T>> generateInitialSplits(
+ int desiredNumSplits, PipelineOptions options) throws Exception {
+ List<PubsubSource<T>> result = new ArrayList<>(desiredNumSplits);
+ PubsubSource<T> splitSource = this;
+ if (subscriptionPath == null) {
+ splitSource = new PubsubSource<>(outer, outer.createRandomSubscription(options));
+ }
+ for (int i = 0; i < desiredNumSplits * SCALE_OUT; i++) {
+ // Since the source is immutable and Pubsub automatically shards we simply
+ // replicate ourselves the requested number of times
+ result.add(splitSource);
+ }
+ return result;
+ }
+
+ @Override
+ public PubsubReader<T> createReader(
+ PipelineOptions options,
+ @Nullable PubsubCheckpoint<T> checkpoint) {
+ PubsubReader<T> reader;
+ SubscriptionPath subscription = subscriptionPath;
+ if (subscription == null) {
+ if (checkpoint == null) {
+ // This reader has never been started and there was no call to #splitIntoBundles; create
+ // a single random subscription, which will be kept in the checkpoint.
+ subscription = outer.createRandomSubscription(options);
+ } else {
+ subscription = checkpoint.getSubscription();
+ }
+ }
+ try {
+ reader = new PubsubReader<>(options.as(PubsubOptions.class), this, subscription,
+ outer.parseFn);
+ } catch (GeneralSecurityException | IOException e) {
+ throw new RuntimeException("Unable to subscribe to " + subscriptionPath + ": ", e);
+ }
+ if (checkpoint != null) {
+ // NACK all messages we may have lost.
+ try {
+ // Will BLOCK until NACKed.
+ checkpoint.nackAll(reader);
+ } catch (IOException e) {
+ LOG.error("Pubsub {} cannot have {} lost messages NACKed, ignoring: {}",
+ subscriptionPath, checkpoint.notYetReadIds.size(), e);
+ }
+ }
+ return reader;
+ }
+
+ @Nullable
+ @Override
+ public Coder<PubsubCheckpoint<T>> getCheckpointMarkCoder() {
+ @SuppressWarnings("unchecked") PubsubCheckpointCoder<T> typedCoder =
+ (PubsubCheckpointCoder<T>) CHECKPOINT_CODER;
+ return typedCoder;
+ }
+
+ @Override
+ public Coder<T> getDefaultOutputCoder() {
+ return outer.elementCoder;
+ }
+
+ @Override
+ public void validate() {
+ // Nothing to validate.
+ }
+
+ @Override
+ public boolean requiresDeduping() {
+ // We cannot prevent re-offering already read messages after a restore from checkpoint.
+ return true;
+ }
+ }
+
+ // ================================================================================
+ // StatsFn
+ // ================================================================================
+
+ private static class StatsFn<T> extends DoFn<T, T> {
+ private final Counter elementCounter = Metrics.counter(StatsFn.class, "elements");
+
+ private final PubsubClientFactory pubsubFactory;
+ @Nullable
+ private final ValueProvider<SubscriptionPath> subscription;
+ @Nullable
+ private final ValueProvider<TopicPath> topic;
+ @Nullable
+ private final String timestampLabel;
+ @Nullable
+ private final String idLabel;
+
+ public StatsFn(
+ PubsubClientFactory pubsubFactory,
+ @Nullable ValueProvider<SubscriptionPath> subscription,
+ @Nullable ValueProvider<TopicPath> topic,
+ @Nullable String timestampLabel,
+ @Nullable String idLabel) {
+ checkArgument(pubsubFactory != null, "pubsubFactory should not be null");
+ this.pubsubFactory = pubsubFactory;
+ this.subscription = subscription;
+ this.topic = topic;
+ this.timestampLabel = timestampLabel;
+ this.idLabel = idLabel;
+ }
+
+ @ProcessElement
+ public void processElement(ProcessContext c) throws Exception {
+ elementCounter.inc();
+ c.output(c.element());
+ }
+
+ @Override
+ public void populateDisplayData(Builder builder) {
+ super.populateDisplayData(builder);
+ if (subscription != null) {
+ String subscriptionString = subscription.isAccessible()
+ ? subscription.get().getPath()
+ : subscription.toString();
+ builder.add(DisplayData.item("subscription", subscriptionString));
+ }
+ if (topic != null) {
+ String topicString = topic.isAccessible()
+ ? topic.get().getPath()
+ : topic.toString();
+ builder.add(DisplayData.item("topic", topicString));
+ }
+ builder.add(DisplayData.item("transport", pubsubFactory.getKind()));
+ builder.addIfNotNull(DisplayData.item("timestampLabel", timestampLabel));
+ builder.addIfNotNull(DisplayData.item("idLabel", idLabel));
+ }
+ }
+
+ // ================================================================================
+ // PubsubUnboundedSource
+ // ================================================================================
+
+ /**
+ * For testing only: Clock to use for all timekeeping. If {@literal null} use system clock.
+ */
+ @Nullable
+ private Clock clock;
+
+ /**
+ * Factory for creating underlying Pubsub transport.
+ */
+ private final PubsubClientFactory pubsubFactory;
+
+ /**
+ * Project under which to create a subscription if only the {@link #topic} was given.
+ */
+ @Nullable
+ private final ValueProvider<ProjectPath> project;
+
+ /**
+ * Topic to read from. If {@literal null}, then {@link #subscription} must be given.
+ * Otherwise {@link #subscription} must be null.
+ */
+ @Nullable
+ private final ValueProvider<TopicPath> topic;
+
+ /**
+ * Subscription to read from. If {@literal null} then {@link #topic} must be given.
+ * Otherwise {@link #topic} must be null.
+ *
+ * <p>If no subscription is given a random one will be created when the transorm is
+ * applied. This field will be update with that subscription's path. The created
+ * subscription is never deleted.
+ */
+ @Nullable
+ private ValueProvider<SubscriptionPath> subscription;
+
+ /**
+ * Coder for elements. Elements are effectively double-encoded: first to a byte array
+ * using this checkpointCoder, then to a base-64 string to conform to Pubsub's payload
+ * conventions.
+ */
+ private final Coder<T> elementCoder;
+
+ /**
+ * Pubsub metadata field holding timestamp of each element, or {@literal null} if should use
+ * Pubsub message publish timestamp instead.
+ */
+ @Nullable
+ private final String timestampLabel;
+
+ /**
+ * Pubsub metadata field holding id for each element, or {@literal null} if need to generate
+ * a unique id ourselves.
+ */
+ @Nullable
+ private final String idLabel;
+
+ /**
+ * If not {@literal null}, the user is asking for PubSub attributes. This parse function will be
+ * used to parse {@link PubsubIO.PubsubMessage}s containing a payload and attributes.
+ */
+ @Nullable
+ SimpleFunction<PubsubIO.PubsubMessage, T> parseFn;
+
+ @VisibleForTesting
+ PubsubUnboundedSource(
+ Clock clock,
+ PubsubClientFactory pubsubFactory,
+ @Nullable ValueProvider<ProjectPath> project,
+ @Nullable ValueProvider<TopicPath> topic,
+ @Nullable ValueProvider<SubscriptionPath> subscription,
+ Coder<T> elementCoder,
+ @Nullable String timestampLabel,
+ @Nullable String idLabel,
+ @Nullable SimpleFunction<PubsubIO.PubsubMessage, T> parseFn) {
+ checkArgument((topic == null) != (subscription == null),
+ "Exactly one of topic and subscription must be given");
+ checkArgument((topic == null) == (project == null),
+ "Project must be given if topic is given");
+ this.clock = clock;
+ this.pubsubFactory = checkNotNull(pubsubFactory);
+ this.project = project;
+ this.topic = topic;
+ this.subscription = subscription;
+ this.elementCoder = checkNotNull(elementCoder);
+ this.timestampLabel = timestampLabel;
+ this.idLabel = idLabel;
+ this.parseFn = parseFn;
+ }
+
+ /**
+ * Construct an unbounded source to consume from the Pubsub {@code subscription}.
+ */
+ public PubsubUnboundedSource(
+ PubsubClientFactory pubsubFactory,
+ @Nullable ValueProvider<ProjectPath> project,
+ @Nullable ValueProvider<TopicPath> topic,
+ @Nullable ValueProvider<SubscriptionPath> subscription,
+ Coder<T> elementCoder,
+ @Nullable String timestampLabel,
+ @Nullable String idLabel,
+ @Nullable SimpleFunction<PubsubIO.PubsubMessage, T> parseFn) {
+ this(null, pubsubFactory, project, topic, subscription, elementCoder, timestampLabel, idLabel,
+ parseFn);
+ }
+
+ /**
+ * Get the coder used for elements.
+ */
+ public Coder<T> getElementCoder() {
+ return elementCoder;
+ }
+
+ /**
+ * Get the project path.
+ */
+ @Nullable
+ public ProjectPath getProject() {
+ return project == null ? null : project.get();
+ }
+
+ /**
+ * Get the topic being read from.
+ */
+ @Nullable
+ public TopicPath getTopic() {
+ return topic == null ? null : topic.get();
+ }
+
+ /**
+ * Get the {@link ValueProvider} for the topic being read from.
+ */
+ @Nullable
+ public ValueProvider<TopicPath> getTopicProvider() {
+ return topic;
+ }
+
+ /**
+ * Get the subscription being read from.
+ */
+ @Nullable
+ public SubscriptionPath getSubscription() {
+ return subscription == null ? null : subscription.get();
+ }
+
+ /**
+ * Get the {@link ValueProvider} for the subscription being read from.
+ */
+ @Nullable
+ public ValueProvider<SubscriptionPath> getSubscriptionProvider() {
+ return subscription;
+ }
+
+ /**
+ * Get the timestamp label.
+ */
+ @Nullable
+ public String getTimestampLabel() {
+ return timestampLabel;
+ }
+
+ /**
+ * Get the id label.
+ */
+ @Nullable
+ public String getIdLabel() {
+ return idLabel;
+ }
+
+ /**
+ * Get the parsing function for PubSub attributes.
+ */
+ @Nullable
+ public SimpleFunction<PubsubIO.PubsubMessage, T> getWithAttributesParseFn() {
+ return parseFn;
+ }
+
+ @Override
+ public PCollection<T> expand(PBegin input) {
+ return input.getPipeline().begin()
+ .apply(Read.from(new PubsubSource<T>(this)))
+ .apply("PubsubUnboundedSource.Stats",
+ ParDo.of(new StatsFn<T>(
+ pubsubFactory, subscription, topic, timestampLabel, idLabel)));
+ }
+
+ private SubscriptionPath createRandomSubscription(PipelineOptions options) {
+ try {
+ try (PubsubClient pubsubClient =
+ pubsubFactory.newClient(timestampLabel, idLabel, options.as(PubsubOptions.class))) {
+ checkState(project.isAccessible(), "createRandomSubscription must be called at runtime.");
+ checkState(topic.isAccessible(), "createRandomSubscription must be called at runtime.");
+ SubscriptionPath subscriptionPath =
+ pubsubClient.createRandomSubscription(
+ project.get(), topic.get(), DEAULT_ACK_TIMEOUT_SEC);
+ LOG.warn(
+ "Created subscription {} to topic {}."
+ + " Note this subscription WILL NOT be deleted when the pipeline terminates",
+ subscriptionPath,
+ topic);
+ return subscriptionPath;
+ }
+ } catch (Exception e) {
+ throw new RuntimeException("Failed to create subscription: ", e);
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/package-info.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/package-info.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/package-info.java
new file mode 100644
index 0000000..55befba
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/package-info.java
@@ -0,0 +1,24 @@
+/*
+ * 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.
+ */
+
+/**
+ * Defines transforms for reading and writing from
+ * <a href="https://cloud.google.com/pubsub/">Google Cloud Pub/Sub</a>.
+ * @see org.apache.beam.sdk.io.gcp.pubsub.PubsubIO
+ */
+package org.apache.beam.sdk.io.gcp.pubsub;
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/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 0987140..f468ec0 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
@@ -44,11 +44,14 @@ public class GcpApiSurfaceTest {
ApiSurface.ofPackage(thisPackage, thisClassLoader)
.pruningPattern("org[.]apache[.]beam[.].*Test.*")
.pruningPattern("org[.]apache[.]beam[.].*IT")
- .pruningPattern("java[.]lang.*");
+ .pruningPattern("java[.]lang.*")
+ .pruningPattern("java[.]util.*");
@SuppressWarnings("unchecked")
final Set<Matcher<Class<?>>> allowedClasses =
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.bigquery.model"),
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubClientTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubClientTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubClientTest.java
new file mode 100644
index 0000000..14c36f9
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubClientTest.java
@@ -0,0 +1,189 @@
+/*
+ * 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.pubsub;
+
+import static org.junit.Assert.assertEquals;
+
+import com.google.common.collect.ImmutableMap;
+import java.util.Map;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.ProjectPath;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath;
+import org.joda.time.Instant;
+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 helper classes and methods in PubsubClient.
+ */
+@RunWith(JUnit4.class)
+public class PubsubClientTest {
+ @Rule
+ public ExpectedException thrown = ExpectedException.none();
+
+ //
+ // Timestamp handling
+ //
+
+ private long parse(String timestamp) {
+ Map<String, String> map = ImmutableMap.of("myLabel", timestamp);
+ return PubsubClient.extractTimestamp("myLabel", null, map);
+ }
+
+ private void roundTripRfc339(String timestamp) {
+ assertEquals(Instant.parse(timestamp).getMillis(), parse(timestamp));
+ }
+
+ private void truncatedRfc339(String timestamp, String truncatedTimestmap) {
+ assertEquals(Instant.parse(truncatedTimestmap).getMillis(), parse(timestamp));
+ }
+
+ @Test
+ public void noTimestampLabelReturnsPubsubPublish() {
+ final long time = 987654321L;
+ long timestamp = PubsubClient.extractTimestamp(null, String.valueOf(time), null);
+ assertEquals(time, timestamp);
+ }
+
+ @Test
+ public void noTimestampLabelAndInvalidPubsubPublishThrowsError() {
+ thrown.expect(NumberFormatException.class);
+ PubsubClient.extractTimestamp(null, "not-a-date", null);
+ }
+
+ @Test
+ public void timestampLabelWithNullAttributesThrowsError() {
+ thrown.expect(RuntimeException.class);
+ thrown.expectMessage("PubSub message is missing a value for timestamp label myLabel");
+ PubsubClient.extractTimestamp("myLabel", null, null);
+ }
+
+ @Test
+ public void timestampLabelSetWithMissingAttributeThrowsError() {
+ thrown.expect(RuntimeException.class);
+ thrown.expectMessage("PubSub message is missing a value for timestamp label myLabel");
+ Map<String, String> map = ImmutableMap.of("otherLabel", "whatever");
+ PubsubClient.extractTimestamp("myLabel", null, map);
+ }
+
+ @Test
+ public void timestampLabelParsesMillisecondsSinceEpoch() {
+ long time = 1446162101123L;
+ Map<String, String> map = ImmutableMap.of("myLabel", String.valueOf(time));
+ long timestamp = PubsubClient.extractTimestamp("myLabel", null, map);
+ assertEquals(time, timestamp);
+ }
+
+ @Test
+ public void timestampLabelParsesRfc3339Seconds() {
+ roundTripRfc339("2015-10-29T23:41:41Z");
+ }
+
+ @Test
+ public void timestampLabelParsesRfc3339Tenths() {
+ roundTripRfc339("2015-10-29T23:41:41.1Z");
+ }
+
+ @Test
+ public void timestampLabelParsesRfc3339Hundredths() {
+ roundTripRfc339("2015-10-29T23:41:41.12Z");
+ }
+
+ @Test
+ public void timestampLabelParsesRfc3339Millis() {
+ roundTripRfc339("2015-10-29T23:41:41.123Z");
+ }
+
+ @Test
+ public void timestampLabelParsesRfc3339Micros() {
+ // Note: micros part 456/1000 is dropped.
+ truncatedRfc339("2015-10-29T23:41:41.123456Z", "2015-10-29T23:41:41.123Z");
+ }
+
+ @Test
+ public void timestampLabelParsesRfc3339MicrosRounding() {
+ // Note: micros part 999/1000 is dropped, not rounded up.
+ truncatedRfc339("2015-10-29T23:41:41.123999Z", "2015-10-29T23:41:41.123Z");
+ }
+
+ @Test
+ public void timestampLabelWithInvalidFormatThrowsError() {
+ thrown.expect(NumberFormatException.class);
+ parse("not-a-timestamp");
+ }
+
+ @Test
+ public void timestampLabelWithInvalidFormat2ThrowsError() {
+ thrown.expect(NumberFormatException.class);
+ parse("null");
+ }
+
+ @Test
+ public void timestampLabelWithInvalidFormat3ThrowsError() {
+ thrown.expect(NumberFormatException.class);
+ parse("2015-10");
+ }
+
+ @Test
+ public void timestampLabelParsesRfc3339WithSmallYear() {
+ // Google and JodaTime agree on dates after 1582-10-15, when the Gregorian Calendar was adopted
+ // This is therefore a "small year" until this difference is reconciled.
+ roundTripRfc339("1582-10-15T01:23:45.123Z");
+ }
+
+ @Test
+ public void timestampLabelParsesRfc3339WithLargeYear() {
+ // Year 9999 in range.
+ roundTripRfc339("9999-10-29T23:41:41.123999Z");
+ }
+
+ @Test
+ public void timestampLabelRfc3339WithTooLargeYearThrowsError() {
+ thrown.expect(NumberFormatException.class);
+ // Year 10000 out of range.
+ parse("10000-10-29T23:41:41.123999Z");
+ }
+
+ //
+ // Paths
+ //
+
+ @Test
+ public void projectPathFromIdWellFormed() {
+ ProjectPath path = PubsubClient.projectPathFromId("test");
+ assertEquals("projects/test", path.getPath());
+ }
+
+ @Test
+ public void subscriptionPathFromNameWellFormed() {
+ SubscriptionPath path = PubsubClient.subscriptionPathFromName("test", "something");
+ assertEquals("projects/test/subscriptions/something", path.getPath());
+ assertEquals("/subscriptions/test/something", path.getV1Beta1Path());
+ }
+
+ @Test
+ public void topicPathFromNameWellFormed() {
+ TopicPath path = PubsubClient.topicPathFromName("test", "something");
+ assertEquals("projects/test/topics/something", path.getPath());
+ assertEquals("/topics/test/something", path.getV1Beta1Path());
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubGrpcClientTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubGrpcClientTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubGrpcClientTest.java
new file mode 100644
index 0000000..28e07e2
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubGrpcClientTest.java
@@ -0,0 +1,208 @@
+/*
+ * 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.pubsub;
+
+import static org.junit.Assert.assertEquals;
+
+import com.google.auth.Credentials;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Iterables;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Timestamp;
+import com.google.pubsub.v1.PublishRequest;
+import com.google.pubsub.v1.PublishResponse;
+import com.google.pubsub.v1.PublisherGrpc.PublisherImplBase;
+import com.google.pubsub.v1.PubsubMessage;
+import com.google.pubsub.v1.PullRequest;
+import com.google.pubsub.v1.PullResponse;
+import com.google.pubsub.v1.ReceivedMessage;
+import com.google.pubsub.v1.SubscriberGrpc.SubscriberImplBase;
+import io.grpc.ManagedChannel;
+import io.grpc.Server;
+import io.grpc.inprocess.InProcessChannelBuilder;
+import io.grpc.inprocess.InProcessServerBuilder;
+import io.grpc.stub.StreamObserver;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ThreadLocalRandom;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.OutgoingMessage;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath;
+import org.apache.beam.sdk.util.TestCredential;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Tests for PubsubGrpcClient.
+ */
+@RunWith(JUnit4.class)
+public class PubsubGrpcClientTest {
+ private ManagedChannel inProcessChannel;
+ private Credentials testCredentials;
+
+ private PubsubClient client;
+ private String channelName;
+
+ private static final TopicPath TOPIC = PubsubClient.topicPathFromName("testProject", "testTopic");
+ private static final SubscriptionPath SUBSCRIPTION =
+ PubsubClient.subscriptionPathFromName("testProject", "testSubscription");
+ private static final long REQ_TIME = 1234L;
+ private static final long PUB_TIME = 3456L;
+ private static final long MESSAGE_TIME = 6789L;
+ private static final String TIMESTAMP_LABEL = "timestamp";
+ private static final String ID_LABEL = "id";
+ private static final String MESSAGE_ID = "testMessageId";
+ private static final String DATA = "testData";
+ private static final String RECORD_ID = "testRecordId";
+ private static final String ACK_ID = "testAckId";
+ private static final Map<String, String> ATTRIBUTES =
+ ImmutableMap.<String, String>builder().put("a", "b").put("c", "d").build();
+
+ @Before
+ public void setup() {
+ channelName = String.format("%s-%s",
+ PubsubGrpcClientTest.class.getName(), ThreadLocalRandom.current().nextInt());
+ inProcessChannel = InProcessChannelBuilder.forName(channelName).directExecutor().build();
+ testCredentials = new TestCredential();
+ client = new PubsubGrpcClient(TIMESTAMP_LABEL, ID_LABEL, 10, inProcessChannel, testCredentials);
+ }
+
+ @After
+ public void teardown() throws IOException {
+ client.close();
+ inProcessChannel.shutdownNow();
+ }
+
+ @Test
+ public void pullOneMessage() throws IOException {
+ String expectedSubscription = SUBSCRIPTION.getPath();
+ final PullRequest expectedRequest =
+ PullRequest.newBuilder()
+ .setSubscription(expectedSubscription)
+ .setReturnImmediately(true)
+ .setMaxMessages(10)
+ .build();
+ Timestamp timestamp = Timestamp.newBuilder()
+ .setSeconds(PUB_TIME / 1000)
+ .setNanos((int) (PUB_TIME % 1000) * 1000)
+ .build();
+ PubsubMessage expectedPubsubMessage =
+ PubsubMessage.newBuilder()
+ .setMessageId(MESSAGE_ID)
+ .setData(
+ ByteString.copyFrom(DATA.getBytes()))
+ .setPublishTime(timestamp)
+ .putAllAttributes(ATTRIBUTES)
+ .putAllAttributes(
+ ImmutableMap.of(TIMESTAMP_LABEL,
+ String.valueOf(MESSAGE_TIME),
+ ID_LABEL, RECORD_ID))
+ .build();
+ ReceivedMessage expectedReceivedMessage =
+ ReceivedMessage.newBuilder()
+ .setMessage(expectedPubsubMessage)
+ .setAckId(ACK_ID)
+ .build();
+ final PullResponse response =
+ PullResponse.newBuilder()
+ .addAllReceivedMessages(ImmutableList.of(expectedReceivedMessage))
+ .build();
+
+ final List<PullRequest> requestsReceived = new ArrayList<>();
+ SubscriberImplBase subscriberImplBase = new SubscriberImplBase() {
+ @Override
+ public void pull(PullRequest request, StreamObserver<PullResponse> responseObserver) {
+ requestsReceived.add(request);
+ responseObserver.onNext(response);
+ responseObserver.onCompleted();
+ }
+ };
+ Server server = InProcessServerBuilder.forName(channelName)
+ .addService(subscriberImplBase)
+ .build()
+ .start();
+ try {
+ List<IncomingMessage> acutalMessages = client.pull(REQ_TIME, SUBSCRIPTION, 10, true);
+ assertEquals(1, acutalMessages.size());
+ IncomingMessage actualMessage = acutalMessages.get(0);
+ assertEquals(ACK_ID, actualMessage.ackId);
+ assertEquals(DATA, new String(actualMessage.elementBytes));
+ assertEquals(RECORD_ID, actualMessage.recordId);
+ assertEquals(REQ_TIME, actualMessage.requestTimeMsSinceEpoch);
+ assertEquals(MESSAGE_TIME, actualMessage.timestampMsSinceEpoch);
+ assertEquals(expectedRequest, Iterables.getOnlyElement(requestsReceived));
+ } finally {
+ server.shutdownNow();
+ }
+ }
+
+ @Test
+ public void publishOneMessage() throws IOException {
+ String expectedTopic = TOPIC.getPath();
+ PubsubMessage expectedPubsubMessage =
+ PubsubMessage.newBuilder()
+ .setData(ByteString.copyFrom(DATA.getBytes()))
+ .putAllAttributes(ATTRIBUTES)
+ .putAllAttributes(
+ ImmutableMap.of(TIMESTAMP_LABEL, String.valueOf(MESSAGE_TIME),
+ ID_LABEL, RECORD_ID))
+ .build();
+ final PublishRequest expectedRequest =
+ PublishRequest.newBuilder()
+ .setTopic(expectedTopic)
+ .addAllMessages(
+ ImmutableList.of(expectedPubsubMessage))
+ .build();
+ final PublishResponse response =
+ PublishResponse.newBuilder()
+ .addAllMessageIds(ImmutableList.of(MESSAGE_ID))
+ .build();
+
+ final List<PublishRequest> requestsReceived = new ArrayList<>();
+ PublisherImplBase publisherImplBase = new PublisherImplBase() {
+ @Override
+ public void publish(
+ PublishRequest request, StreamObserver<PublishResponse> responseObserver) {
+ requestsReceived.add(request);
+ responseObserver.onNext(response);
+ responseObserver.onCompleted();
+ }
+ };
+ Server server = InProcessServerBuilder.forName(channelName)
+ .addService(publisherImplBase)
+ .build()
+ .start();
+ try {
+ OutgoingMessage actualMessage = new OutgoingMessage(
+ DATA.getBytes(), ATTRIBUTES, MESSAGE_TIME, RECORD_ID);
+ int n = client.publish(TOPIC, ImmutableList.of(actualMessage));
+ assertEquals(1, n);
+ assertEquals(expectedRequest, Iterables.getOnlyElement(requestsReceived));
+ } finally {
+ server.shutdownNow();
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOTest.java
new file mode 100644
index 0000000..6e9922c
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOTest.java
@@ -0,0 +1,189 @@
+/*
+ * 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.pubsub;
+
+import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem;
+import static org.hamcrest.Matchers.hasItem;
+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 java.util.Set;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
+import org.apache.beam.sdk.testing.UsesUnboundedPCollections;
+import org.apache.beam.sdk.testing.ValidatesRunner;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.DisplayDataEvaluator;
+import org.joda.time.Duration;
+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;
+
+/**
+ * Tests for PubsubIO Read and Write transforms.
+ */
+@RunWith(JUnit4.class)
+public class PubsubIOTest {
+ @Rule
+ public ExpectedException thrown = ExpectedException.none();
+
+ @Test
+ public void testPubsubIOGetName() {
+ assertEquals("PubsubIO.Read",
+ PubsubIO.<String>read().topic("projects/myproject/topics/mytopic").getName());
+ assertEquals("PubsubIO.Write",
+ PubsubIO.<String>write().topic("projects/myproject/topics/mytopic").getName());
+ }
+
+ @Test
+ public void testTopicValidationSuccess() throws Exception {
+ PubsubIO.<String>read().topic("projects/my-project/topics/abc");
+ PubsubIO.<String>read().topic("projects/my-project/topics/ABC");
+ PubsubIO.<String>read().topic("projects/my-project/topics/AbC-DeF");
+ PubsubIO.<String>read().topic("projects/my-project/topics/AbC-1234");
+ PubsubIO.<String>read().topic("projects/my-project/topics/AbC-1234-_.~%+-_.~%+-_.~%+-abc");
+ PubsubIO.<String>read().topic(new StringBuilder()
+ .append("projects/my-project/topics/A-really-long-one-")
+ .append("111111111111111111111111111111111111111111111111111111111111111111111111111111111")
+ .append("111111111111111111111111111111111111111111111111111111111111111111111111111111111")
+ .append("11111111111111111111111111111111111111111111111111111111111111111111111111")
+ .toString());
+ }
+
+ @Test
+ public void testTopicValidationBadCharacter() throws Exception {
+ thrown.expect(IllegalArgumentException.class);
+ PubsubIO.<String>read().topic("projects/my-project/topics/abc-*-abc");
+ }
+
+ @Test
+ public void testTopicValidationTooLong() throws Exception {
+ thrown.expect(IllegalArgumentException.class);
+ PubsubIO.<String>read().topic(new StringBuilder().append
+ ("projects/my-project/topics/A-really-long-one-")
+ .append("111111111111111111111111111111111111111111111111111111111111111111111111111111111")
+ .append("111111111111111111111111111111111111111111111111111111111111111111111111111111111")
+ .append("1111111111111111111111111111111111111111111111111111111111111111111111111111")
+ .toString());
+ }
+
+ @Test
+ public void testReadTopicDisplayData() {
+ String topic = "projects/project/topics/topic";
+ String subscription = "projects/project/subscriptions/subscription";
+ Duration maxReadTime = Duration.standardMinutes(5);
+ PubsubIO.Read<String> read = PubsubIO.<String>read()
+ .topic(StaticValueProvider.of(topic))
+ .timestampLabel("myTimestamp")
+ .idLabel("myId");
+
+ DisplayData displayData = DisplayData.from(read);
+
+ assertThat(displayData, hasDisplayItem("topic", topic));
+ assertThat(displayData, hasDisplayItem("timestampLabel", "myTimestamp"));
+ assertThat(displayData, hasDisplayItem("idLabel", "myId"));
+ }
+
+ @Test
+ public void testReadSubscriptionDisplayData() {
+ String topic = "projects/project/topics/topic";
+ String subscription = "projects/project/subscriptions/subscription";
+ Duration maxReadTime = Duration.standardMinutes(5);
+ PubsubIO.Read<String> read = PubsubIO.<String>read()
+ .subscription(StaticValueProvider.of(subscription))
+ .timestampLabel("myTimestamp")
+ .idLabel("myId");
+
+ DisplayData displayData = DisplayData.from(read);
+
+ assertThat(displayData, hasDisplayItem("subscription", subscription));
+ assertThat(displayData, hasDisplayItem("timestampLabel", "myTimestamp"));
+ assertThat(displayData, hasDisplayItem("idLabel", "myId"));
+ }
+
+ @Test
+ public void testNullTopic() {
+ String subscription = "projects/project/subscriptions/subscription";
+ PubsubIO.Read<String> read = PubsubIO.<String>read()
+ .subscription(StaticValueProvider.of(subscription));
+ assertNull(read.getTopic());
+ assertNotNull(read.getSubscription());
+ assertNotNull(DisplayData.from(read));
+ }
+
+ @Test
+ public void testNullSubscription() {
+ String topic = "projects/project/topics/topic";
+ PubsubIO.Read<String> read = PubsubIO.<String>read()
+ .topic(StaticValueProvider.of(topic));
+ assertNotNull(read.getTopic());
+ assertNull(read.getSubscription());
+ assertNotNull(DisplayData.from(read));
+ }
+
+ @Test
+ @Category({ValidatesRunner.class, UsesUnboundedPCollections.class})
+ public void testPrimitiveReadDisplayData() {
+ DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
+ Set<DisplayData> displayData;
+ PubsubIO.Read<String> read = PubsubIO.<String>read().withCoder(StringUtf8Coder.of());
+
+ // Reading from a subscription.
+ read = read.subscription("projects/project/subscriptions/subscription");
+ displayData = evaluator.displayDataForPrimitiveSourceTransforms(read);
+ assertThat("PubsubIO.Read should include the subscription in its primitive display data",
+ displayData, hasItem(hasDisplayItem("subscription")));
+
+ // Reading from a topic.
+ read = read.topic("projects/project/topics/topic");
+ displayData = evaluator.displayDataForPrimitiveSourceTransforms(read);
+ assertThat("PubsubIO.Read should include the topic in its primitive display data",
+ displayData, hasItem(hasDisplayItem("topic")));
+ }
+
+ @Test
+ public void testWriteDisplayData() {
+ String topic = "projects/project/topics/topic";
+ PubsubIO.Write<?> write = PubsubIO.<String>write()
+ .topic(topic)
+ .timestampLabel("myTimestamp")
+ .idLabel("myId");
+
+ DisplayData displayData = DisplayData.from(write);
+
+ assertThat(displayData, hasDisplayItem("topic", topic));
+ assertThat(displayData, hasDisplayItem("timestampLabel", "myTimestamp"));
+ assertThat(displayData, hasDisplayItem("idLabel", "myId"));
+ }
+
+ @Test
+ @Category(ValidatesRunner.class)
+ public void testPrimitiveWriteDisplayData() {
+ DisplayDataEvaluator evaluator = DisplayDataEvaluator.create();
+ PubsubIO.Write<?> write = PubsubIO.<String>write().topic("projects/project/topics/topic");
+
+ Set<DisplayData> displayData = evaluator.displayDataForPrimitiveTransforms(write);
+ assertThat("PubsubIO.Write should include the topic in its primitive display data",
+ displayData, hasItem(hasDisplayItem("topic")));
+ }
+}
[27/50] [abbrv] beam git commit: This closes #2504
Posted by jb...@apache.org.
This closes #2504
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/82f2f2cf
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/82f2f2cf
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/82f2f2cf
Branch: refs/heads/DSL_SQL
Commit: 82f2f2cff0f026800df465f1bff5451a7891ee13
Parents: ccf247c e7f4337
Author: Dan Halperin <dh...@google.com>
Authored: Tue Apr 11 18:34:14 2017 -0700
Committer: Dan Halperin <dh...@google.com>
Committed: Tue Apr 11 18:34:14 2017 -0700
----------------------------------------------------------------------
.../java/org/apache/beam/sdk/transforms/windowing/Triggers.java | 2 ++
1 file changed, 2 insertions(+)
----------------------------------------------------------------------
[39/50] [abbrv] beam git commit: [BEAM-115] Update timer/state fields
on ParDoPayload to use a map field for consistent name usage
Posted by jb...@apache.org.
[BEAM-115] Update timer/state fields on ParDoPayload to use a map field for consistent name usage
This closes #2511
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/435c3b42
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/435c3b42
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/435c3b42
Branch: refs/heads/DSL_SQL
Commit: 435c3b4267f527543849db3777f021fecc5964df
Parents: 571631a b114844
Author: Luke Cwik <lc...@google.com>
Authored: Wed Apr 12 10:59:58 2017 -0700
Committer: Luke Cwik <lc...@google.com>
Committed: Wed Apr 12 10:59:58 2017 -0700
----------------------------------------------------------------------
.../runner-api/src/main/proto/beam_runner_api.proto | 10 +++++-----
1 file changed, 5 insertions(+), 5 deletions(-)
----------------------------------------------------------------------
[49/50] [abbrv] beam git commit: [BEAM-911] This closes #2510
Posted by jb...@apache.org.
[BEAM-911] This closes #2510
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/1788cef9
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/1788cef9
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/1788cef9
Branch: refs/heads/DSL_SQL
Commit: 1788cef96916ebfb588848922a5cbc2584586e9b
Parents: a405a2c 024b6e0
Author: Jean-Baptiste Onofr� <jb...@apache.org>
Authored: Wed Apr 12 21:53:15 2017 +0200
Committer: Jean-Baptiste Onofr� <jb...@apache.org>
Committed: Wed Apr 12 21:53:15 2017 +0200
----------------------------------------------------------------------
.../org/apache/beam/sdk/io/elasticsearch/ElasticsearchIO.java | 2 ++
.../beam/sdk/io/hadoop/inputformat/HadoopInputFormatIO.java | 3 ++-
.../src/main/java/org/apache/beam/sdk/io/hdfs/HDFSFileSink.java | 2 ++
.../src/main/java/org/apache/beam/sdk/io/hdfs/HDFSFileSource.java | 2 ++
.../io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java | 2 ++
.../io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java | 2 ++
.../kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java | 2 ++
.../src/main/java/org/apache/beam/sdk/io/kinesis/KinesisIO.java | 2 ++
.../main/java/org/apache/beam/sdk/io/mongodb/MongoDbGridFSIO.java | 2 ++
.../src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java | 2 ++
.../io/mqtt/src/main/java/org/apache/beam/sdk/io/mqtt/MqttIO.java | 2 ++
11 files changed, 22 insertions(+), 1 deletion(-)
----------------------------------------------------------------------
[12/50] [abbrv] beam git commit: Remove Orderdness of Input,
Output expansions
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 8b4573f..b44c890 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
@@ -45,7 +45,6 @@ import org.apache.beam.sdk.values.PCollectionView;
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.TaggedPValue;
import org.hamcrest.Matchers;
import org.junit.Rule;
import org.junit.Test;
@@ -101,9 +100,9 @@ public class DirectGraphVisitorTest implements Serializable {
graph.getProducer(created), graph.getProducer(counted), graph.getProducer(unCounted)));
for (AppliedPTransform<?, ?, ?> root : graph.getRootTransforms()) {
// Root transforms will have no inputs
- assertThat(root.getInputs(), emptyIterable());
+ assertThat(root.getInputs().entrySet(), emptyIterable());
assertThat(
- Iterables.getOnlyElement(root.getOutputs()).getValue(),
+ Iterables.getOnlyElement(root.getOutputs().values()),
Matchers.<POutput>isOneOf(created, counted, unCounted));
}
}
@@ -121,7 +120,7 @@ public class DirectGraphVisitorTest implements Serializable {
Matchers.<AppliedPTransform<?, ?, ?>>containsInAnyOrder(graph.getProducer(empty)));
AppliedPTransform<?, ?, ?> onlyRoot = Iterables.getOnlyElement(graph.getRootTransforms());
assertThat(onlyRoot.getTransform(), Matchers.<PTransform<?, ?>>equalTo(flatten));
- assertThat(onlyRoot.getInputs(), Matchers.<TaggedPValue>emptyIterable());
+ assertThat(onlyRoot.getInputs().entrySet(), emptyIterable());
assertThat(onlyRoot.getOutputs(), equalTo(empty.expand()));
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 c85b85e..2a94d48 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
@@ -162,7 +162,7 @@ public class ParDoEvaluatorTest {
evaluationContext,
stepContext,
transform,
- ((PCollection<?>) Iterables.getOnlyElement(transform.getInputs()).getValue())
+ ((PCollection<?>) Iterables.getOnlyElement(transform.getInputs().values()))
.getWindowingStrategy(),
fn,
null /* key */,
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 946cd69..ecb8130 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
@@ -308,7 +308,7 @@ public class StatefulParDoEvaluatorFactoryTest implements Serializable {
BUNDLE_FACTORY
.createBundle(
(PCollection<KeyedWorkItem<String, KV<String, Integer>>>)
- Iterables.getOnlyElement(producingTransform.getInputs()).getValue())
+ Iterables.getOnlyElement(producingTransform.getInputs().values()))
.add(gbkOutputElement)
.commit(Instant.now());
TransformEvaluator<KeyedWorkItem<String, KV<String, Integer>>> evaluator =
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/direct-java/src/test/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactoryTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactoryTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactoryTest.java
index fc689fe..0d909c2 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactoryTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactoryTest.java
@@ -40,8 +40,9 @@ import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TimestampedValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.hamcrest.Matchers;
import org.joda.time.Duration;
import org.joda.time.Instant;
@@ -183,7 +184,7 @@ public class TestStreamEvaluatorFactoryTest {
@Test
public void overrideFactoryGetInputSucceeds() {
DirectTestStreamFactory<?> factory = new DirectTestStreamFactory<>(runner);
- PBegin begin = factory.getInput(Collections.<TaggedPValue>emptyList(), p);
+ PBegin begin = factory.getInput(Collections.<TupleTag<?>, PValue>emptyMap(), p);
assertThat(begin.getPipeline(), Matchers.<Pipeline>equalTo(p));
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 6dcc13c..258cb46 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
@@ -107,7 +107,7 @@ public class ViewOverrideFactoryTest implements Serializable {
is(false));
PCollectionView replacementView = ((WriteView) node.getTransform()).getView();
assertThat(replacementView, Matchers.<PCollectionView>theInstance(view));
- assertThat(node.getInputs(), hasSize(1));
+ assertThat(node.getInputs().entrySet(), hasSize(1));
}
}
});
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkBatchTransformTranslators.java
----------------------------------------------------------------------
diff --git a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkBatchTransformTranslators.java b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkBatchTransformTranslators.java
index 1d6728b..ff9521c 100644
--- a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkBatchTransformTranslators.java
+++ b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkBatchTransformTranslators.java
@@ -27,6 +27,7 @@ import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.Map.Entry;
import org.apache.beam.runners.flink.translation.functions.FlinkAssignWindows;
import org.apache.beam.runners.flink.translation.functions.FlinkDoFnFunction;
import org.apache.beam.runners.flink.translation.functions.FlinkMergingNonShuffleReduceFunction;
@@ -71,7 +72,6 @@ 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.PValue;
-import org.apache.beam.sdk.values.TaggedPValue;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.flink.api.common.functions.FilterFunction;
import org.apache.flink.api.common.functions.FlatMapFunction;
@@ -511,15 +511,15 @@ class FlinkBatchTransformTranslators {
DataSet<WindowedValue<InputT>> inputDataSet =
context.getInputDataSet(context.getInput(transform));
- List<TaggedPValue> outputs = context.getOutputs(transform);
+ Map<TupleTag<?>, PValue> outputs = context.getOutputs(transform);
Map<TupleTag<?>, Integer> outputMap = Maps.newHashMap();
// put the main output at index 0, FlinkMultiOutputDoFnFunction expects this
outputMap.put(transform.getMainOutputTag(), 0);
int count = 1;
- for (TaggedPValue taggedValue : outputs) {
- if (!outputMap.containsKey(taggedValue.getTag())) {
- outputMap.put(taggedValue.getTag(), count++);
+ for (TupleTag<?> tag : outputs.keySet()) {
+ if (!outputMap.containsKey(tag)) {
+ outputMap.put(tag, count++);
}
}
@@ -528,13 +528,13 @@ class FlinkBatchTransformTranslators {
// collect all output Coders and create a UnionCoder for our tagged outputs
List<Coder<?>> outputCoders = Lists.newArrayList();
- for (TaggedPValue taggedValue : outputs) {
+ for (PValue taggedValue : outputs.values()) {
checkState(
- taggedValue.getValue() instanceof PCollection,
+ taggedValue instanceof PCollection,
"Within ParDo, got a non-PCollection output %s of type %s",
- taggedValue.getValue(),
- taggedValue.getValue().getClass().getSimpleName());
- PCollection<?> coll = (PCollection<?>) taggedValue.getValue();
+ taggedValue,
+ taggedValue.getClass().getSimpleName());
+ PCollection<?> coll = (PCollection<?>) taggedValue;
outputCoders.add(coll.getCoder());
windowingStrategy = coll.getWindowingStrategy();
}
@@ -599,11 +599,11 @@ class FlinkBatchTransformTranslators {
transformSideInputs(sideInputs, outputDataSet, context);
- for (TaggedPValue output : outputs) {
+ for (Entry<TupleTag<?>, PValue> output : outputs.entrySet()) {
pruneOutput(
outputDataSet,
context,
- outputMap.get(output.getTag()),
+ outputMap.get(output.getKey()),
(PCollection) output.getValue());
}
@@ -640,7 +640,7 @@ class FlinkBatchTransformTranslators {
Flatten.PCollections<T> transform,
FlinkBatchTranslationContext context) {
- List<TaggedPValue> allInputs = context.getInputs(transform);
+ Map<TupleTag<?>, PValue> allInputs = context.getInputs(transform);
DataSet<WindowedValue<T>> result = null;
if (allInputs.isEmpty()) {
@@ -661,13 +661,13 @@ class FlinkBatchTransformTranslators {
(Coder<T>) VoidCoder.of(),
GlobalWindow.Coder.INSTANCE)));
} else {
- for (TaggedPValue taggedPc : allInputs) {
+ for (PValue taggedPc : allInputs.values()) {
checkArgument(
- taggedPc.getValue() instanceof PCollection,
+ taggedPc instanceof PCollection,
"Got non-PCollection input to flatten: %s of type %s",
- taggedPc.getValue(),
- taggedPc.getValue().getClass().getSimpleName());
- PCollection<T> collection = (PCollection<T>) taggedPc.getValue();
+ taggedPc,
+ taggedPc.getClass().getSimpleName());
+ PCollection<T> collection = (PCollection<T>) taggedPc;
DataSet<WindowedValue<T>> current = context.getInputDataSet(collection);
if (result == null) {
result = current;
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkBatchTranslationContext.java
----------------------------------------------------------------------
diff --git a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkBatchTranslationContext.java b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkBatchTranslationContext.java
index cb69575..98dd0fb 100644
--- a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkBatchTranslationContext.java
+++ b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkBatchTranslationContext.java
@@ -19,7 +19,6 @@ package org.apache.beam.runners.flink;
import com.google.common.collect.Iterables;
import java.util.HashMap;
-import java.util.List;
import java.util.Map;
import org.apache.beam.runners.flink.translation.types.CoderTypeInformation;
import org.apache.beam.sdk.coders.Coder;
@@ -31,7 +30,7 @@ import org.apache.beam.sdk.util.WindowingStrategy;
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.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.DataSet;
import org.apache.flink.api.java.ExecutionEnvironment;
@@ -134,21 +133,21 @@ class FlinkBatchTranslationContext {
return new CoderTypeInformation<>(windowedValueCoder);
}
- List<TaggedPValue> getInputs(PTransform<?, ?> transform) {
+ Map<TupleTag<?>, PValue> getInputs(PTransform<?, ?> transform) {
return currentTransform.getInputs();
}
@SuppressWarnings("unchecked")
<T extends PValue> T getInput(PTransform<T, ?> transform) {
- return (T) Iterables.getOnlyElement(currentTransform.getInputs()).getValue();
+ return (T) Iterables.getOnlyElement(currentTransform.getInputs().values());
}
- List<TaggedPValue> getOutputs(PTransform<?, ?> transform) {
+ Map<TupleTag<?>, PValue> getOutputs(PTransform<?, ?> transform) {
return currentTransform.getOutputs();
}
@SuppressWarnings("unchecked")
<T extends PValue> T getOutput(PTransform<?, T> transform) {
- return (T) Iterables.getOnlyElement(currentTransform.getOutputs()).getValue();
+ return (T) Iterables.getOnlyElement(currentTransform.getOutputs().values());
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslator.java
----------------------------------------------------------------------
diff --git a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslator.java b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslator.java
index 8b5637e..70da2b3 100644
--- a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslator.java
+++ b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkStreamingPipelineTranslator.java
@@ -39,7 +39,7 @@ 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;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -259,14 +259,13 @@ class FlinkStreamingPipelineTranslator extends FlinkPipelineTranslator {
}
@Override
- public PCollection<? extends InputT> getInput(
- List<TaggedPValue> inputs, Pipeline p) {
- return (PCollection<? extends InputT>) Iterables.getOnlyElement(inputs).getValue();
+ public PCollection<? extends InputT> getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
+ return (PCollection<? extends InputT>) Iterables.getOnlyElement(inputs.values());
}
@Override
public Map<PValue, ReplacementOutput> mapOutputs(
- List<TaggedPValue> outputs, PCollectionTuple newOutput) {
+ Map<TupleTag<?>, PValue> outputs, PCollectionTuple newOutput) {
return ReplacementOutputs.tagged(outputs, newOutput);
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java
----------------------------------------------------------------------
diff --git a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java
index 5c29db2..af157f0 100644
--- a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java
+++ b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java
@@ -29,6 +29,7 @@ import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.Map.Entry;
import org.apache.beam.runners.core.ElementAndRestriction;
import org.apache.beam.runners.core.KeyedWorkItem;
import org.apache.beam.runners.core.SplittableParDo;
@@ -79,7 +80,6 @@ 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.PValue;
-import org.apache.beam.sdk.values.TaggedPValue;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.functions.MapFunction;
@@ -420,7 +420,7 @@ class FlinkStreamingTransformTranslators {
DoFn<InputT, OutputT> doFn,
PCollection<InputT> input,
List<PCollectionView<?>> sideInputs,
- List<TaggedPValue> outputs,
+ Map<TupleTag<?>, PValue> outputs,
TupleTag<OutputT> mainOutputTag,
List<TupleTag<?>> sideOutputTags,
FlinkStreamingTranslationContext context,
@@ -537,8 +537,8 @@ class FlinkStreamingTransformTranslators {
}
});
- for (TaggedPValue output : outputs) {
- final int outputTag = tagsToLabels.get(output.getTag());
+ for (Entry<TupleTag<?>, PValue> output : outputs.entrySet()) {
+ final int outputTag = tagsToLabels.get(output.getKey());
TypeInformation outputTypeInfo = context.getTypeInfo((PCollection<?>) output.getValue());
@@ -557,28 +557,28 @@ class FlinkStreamingTransformTranslators {
private static Map<TupleTag<?>, Integer> transformTupleTagsToLabels(
TupleTag<?> mainTag,
- List<TaggedPValue> allTaggedValues) {
+ Map<TupleTag<?>, PValue> allTaggedValues) {
Map<TupleTag<?>, Integer> tagToLabelMap = Maps.newHashMap();
int count = 0;
tagToLabelMap.put(mainTag, count++);
- for (TaggedPValue taggedPValue : allTaggedValues) {
- if (!tagToLabelMap.containsKey(taggedPValue.getTag())) {
- tagToLabelMap.put(taggedPValue.getTag(), count++);
+ for (TupleTag<?> key : allTaggedValues.keySet()) {
+ if (!tagToLabelMap.containsKey(key)) {
+ tagToLabelMap.put(key, count++);
}
}
return tagToLabelMap;
}
- private static UnionCoder createUnionCoder(Collection<TaggedPValue> taggedCollections) {
+ private static UnionCoder createUnionCoder(Map<TupleTag<?>, PValue> taggedCollections) {
List<Coder<?>> outputCoders = Lists.newArrayList();
- for (TaggedPValue taggedColl : taggedCollections) {
+ for (PValue taggedColl : taggedCollections.values()) {
checkArgument(
- taggedColl.getValue() instanceof PCollection,
+ taggedColl instanceof PCollection,
"A Union Coder can only be created for a Collection of Tagged %s. Got %s",
PCollection.class.getSimpleName(),
- taggedColl.getValue().getClass().getSimpleName());
- PCollection<?> coll = (PCollection<?>) taggedColl.getValue();
+ taggedColl.getClass().getSimpleName());
+ PCollection<?> coll = (PCollection<?>) taggedColl;
WindowedValue.FullWindowedValueCoder<?> windowedValueCoder =
WindowedValue.getFullCoder(
coll.getCoder(),
@@ -1042,7 +1042,7 @@ class FlinkStreamingTransformTranslators {
public void translateNode(
Flatten.PCollections<T> transform,
FlinkStreamingTranslationContext context) {
- List<TaggedPValue> allInputs = context.getInputs(transform);
+ Map<TupleTag<?>, PValue> allInputs = context.getInputs(transform);
if (allInputs.isEmpty()) {
@@ -1069,8 +1069,8 @@ class FlinkStreamingTransformTranslators {
} else {
DataStream<T> result = null;
- for (TaggedPValue input : allInputs) {
- DataStream<T> current = context.getInputDataStream(input.getValue());
+ for (PValue input : allInputs.values()) {
+ DataStream<T> current = context.getInputDataStream(input);
result = (result == null) ? current : result.union(current);
}
context.setOutputDataStream(context.getOutput(transform), result);
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTranslationContext.java
----------------------------------------------------------------------
diff --git a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTranslationContext.java b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTranslationContext.java
index 3d5b83f..1a943a3 100644
--- a/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTranslationContext.java
+++ b/runners/flink/runner/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTranslationContext.java
@@ -21,7 +21,6 @@ import static com.google.common.base.Preconditions.checkNotNull;
import com.google.common.collect.Iterables;
import java.util.HashMap;
-import java.util.List;
import java.util.Map;
import org.apache.beam.runners.flink.translation.types.CoderTypeInformation;
import org.apache.beam.sdk.coders.Coder;
@@ -33,7 +32,7 @@ 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;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
@@ -111,19 +110,20 @@ class FlinkStreamingTranslationContext {
@SuppressWarnings("unchecked")
public <T extends PValue> T getInput(PTransform<T, ?> transform) {
- return (T) Iterables.getOnlyElement(currentTransform.getInputs()).getValue();
+ return (T) Iterables.getOnlyElement(currentTransform.getInputs().values());
}
- public <T extends PInput> List<TaggedPValue> getInputs(PTransform<T, ?> transform) {
+ public <T extends PInput> Map<TupleTag<?>, PValue> getInputs(PTransform<T, ?> transform) {
return currentTransform.getInputs();
}
@SuppressWarnings("unchecked")
public <T extends PValue> T getOutput(PTransform<?, T> transform) {
- return (T) Iterables.getOnlyElement(currentTransform.getOutputs()).getValue();
+ return (T) Iterables.getOnlyElement(currentTransform.getOutputs().values());
}
- public <OutputT extends POutput> List<TaggedPValue> getOutputs(PTransform<?, OutputT> transform) {
+ public <OutputT extends POutput> Map<TupleTag<?>, PValue> getOutputs(
+ PTransform<?, OutputT> transform) {
return currentTransform.getOutputs();
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 1d19d64..3ded079 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
@@ -20,7 +20,6 @@ package org.apache.beam.runners.dataflow;
import static com.google.common.base.Preconditions.checkState;
import com.google.common.collect.Iterables;
-import java.util.List;
import java.util.Map;
import org.apache.beam.runners.core.construction.ReplacementOutputs;
import org.apache.beam.runners.dataflow.BatchViewOverrides.GroupByKeyAndSortValuesOnly;
@@ -42,7 +41,7 @@ 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.PValue;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TypeDescriptor;
import org.joda.time.Instant;
@@ -93,13 +92,13 @@ public class BatchStatefulParDoOverrides {
}
@Override
- public PCollection<KV<K, InputT>> getInput(List<TaggedPValue> inputs, Pipeline p) {
- return (PCollection<KV<K, InputT>>) Iterables.getOnlyElement(inputs).getValue();
+ public PCollection<KV<K, InputT>> getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
+ return (PCollection<KV<K, InputT>>) Iterables.getOnlyElement(inputs.values());
}
@Override
public Map<PValue, ReplacementOutput> mapOutputs(
- List<TaggedPValue> outputs, PCollection<OutputT> newOutput) {
+ Map<TupleTag<?>, PValue> outputs, PCollection<OutputT> newOutput) {
return ReplacementOutputs.singleton(outputs, newOutput);
}
}
@@ -116,13 +115,13 @@ public class BatchStatefulParDoOverrides {
}
@Override
- public PCollection<KV<K, InputT>> getInput(List<TaggedPValue> inputs, Pipeline p) {
- return (PCollection<KV<K, InputT>>) Iterables.getOnlyElement(inputs).getValue();
+ public PCollection<KV<K, InputT>> getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
+ return (PCollection<KV<K, InputT>>) Iterables.getOnlyElement(inputs.values());
}
@Override
public Map<PValue, ReplacementOutput> mapOutputs(
- List<TaggedPValue> outputs, PCollectionTuple newOutput) {
+ Map<TupleTag<?>, PValue> outputs, PCollectionTuple newOutput) {
return ReplacementOutputs.tagged(outputs, newOutput);
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 ab9df70..1a2e663 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
@@ -97,7 +97,6 @@ import org.apache.beam.sdk.values.PCollectionView;
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.TaggedPValue;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TypedPValue;
import org.slf4j.Logger;
@@ -371,24 +370,25 @@ public class DataflowPipelineTranslator {
}
@Override
- public <InputT extends PInput> List<TaggedPValue> getInputs(PTransform<InputT, ?> transform) {
+ public <InputT extends PInput> Map<TupleTag<?>, PValue> getInputs(
+ PTransform<InputT, ?> transform) {
return getCurrentTransform(transform).getInputs();
}
@Override
public <InputT extends PValue> InputT getInput(PTransform<InputT, ?> transform) {
- return (InputT) Iterables.getOnlyElement(getInputs(transform)).getValue();
+ return (InputT) Iterables.getOnlyElement(getInputs(transform).values());
}
@Override
- public <OutputT extends POutput> List<TaggedPValue> getOutputs(
+ public <OutputT extends POutput> Map<TupleTag<?>, PValue> getOutputs(
PTransform<?, OutputT> transform) {
return getCurrentTransform(transform).getOutputs();
}
@Override
public <OutputT extends PValue> OutputT getOutput(PTransform<?, OutputT> transform) {
- return (OutputT) Iterables.getOnlyElement(getOutputs(transform)).getValue();
+ return (OutputT) Iterables.getOnlyElement(getOutputs(transform).values());
}
@Override
@@ -758,10 +758,10 @@ public class DataflowPipelineTranslator {
StepTranslationContext stepContext = context.addStep(transform, "Flatten");
List<OutputReference> inputs = new LinkedList<>();
- for (TaggedPValue input : context.getInputs(transform)) {
+ for (PValue input : context.getInputs(transform).values()) {
inputs.add(
context.asOutputReference(
- input.getValue(), context.getProducer(input.getValue())));
+ input, context.getProducer(input)));
}
stepContext.addInput(PropertyNames.INPUTS, inputs);
stepContext.addOutput(context.getOutput(transform));
@@ -967,11 +967,11 @@ public class DataflowPipelineTranslator {
}
private static BiMap<Long, TupleTag<?>> translateOutputs(
- List<TaggedPValue> outputs,
+ Map<TupleTag<?>, PValue> outputs,
StepTranslationContext stepContext) {
ImmutableBiMap.Builder<Long, TupleTag<?>> mapBuilder = ImmutableBiMap.builder();
- for (TaggedPValue taggedOutput : outputs) {
- TupleTag<?> tag = taggedOutput.getTag();
+ for (Map.Entry<TupleTag<?>, PValue> taggedOutput : outputs.entrySet()) {
+ TupleTag<?> tag = taggedOutput.getKey();
checkArgument(taggedOutput.getValue() instanceof PCollection,
"Non %s returned from Multi-output %s",
PCollection.class.getSimpleName(),
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 f789769..9b993f4 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
@@ -127,7 +127,7 @@ import org.apache.beam.sdk.values.PCollectionView;
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.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.joda.time.DateTimeUtils;
import org.joda.time.DateTimeZone;
import org.joda.time.format.DateTimeFormat;
@@ -450,13 +450,13 @@ public class DataflowRunner extends PipelineRunner<DataflowPipelineJob> {
}
@Override
- public PBegin getInput(List<TaggedPValue> inputs, Pipeline p) {
+ public PBegin getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
return p.begin();
}
@Override
public Map<PValue, ReplacementOutput> mapOutputs(
- List<TaggedPValue> outputs, PCollection<T> newOutput) {
+ Map<TupleTag<?>, PValue> outputs, PCollection<T> newOutput) {
return ReplacementOutputs.singleton(outputs, newOutput);
}
}
@@ -760,7 +760,7 @@ public class DataflowRunner extends PipelineRunner<DataflowPipelineJob> {
if (node.getTransform() instanceof View.AsMap
|| node.getTransform() instanceof View.AsMultimap) {
PCollection<KV<?, ?>> input =
- (PCollection<KV<?, ?>>) Iterables.getOnlyElement(node.getInputs()).getValue();
+ (PCollection<KV<?, ?>>) Iterables.getOnlyElement(node.getInputs().values());
KvCoder<?, ?> inputCoder = (KvCoder) input.getCoder();
try {
inputCoder.getKeyCoder().verifyDeterministic();
@@ -825,13 +825,13 @@ public class DataflowRunner extends PipelineRunner<DataflowPipelineJob> {
}
@Override
- public PCollection<T> getInput(List<TaggedPValue> inputs, Pipeline p) {
- return (PCollection<T>) Iterables.getOnlyElement(inputs).getValue();
+ public PCollection<T> getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
+ return (PCollection<T>) Iterables.getOnlyElement(inputs.values());
}
@Override
public Map<PValue, ReplacementOutput> mapOutputs(
- List<TaggedPValue> outputs, PDone newOutput) {
+ Map<TupleTag<?>, PValue> outputs, PDone newOutput) {
return Collections.emptyMap();
}
}
@@ -1317,13 +1317,13 @@ public class DataflowRunner extends PipelineRunner<DataflowPipelineJob> {
@Override
public PCollection<KV<K, Iterable<InputT>>> getInput(
- List<TaggedPValue> inputs, Pipeline p) {
- return (PCollection<KV<K, Iterable<InputT>>>) Iterables.getOnlyElement(inputs).getValue();
+ Map<TupleTag<?>, PValue> inputs, Pipeline p) {
+ return (PCollection<KV<K, Iterable<InputT>>>) Iterables.getOnlyElement(inputs.values());
}
@Override
public Map<PValue, ReplacementOutput> mapOutputs(
- List<TaggedPValue> outputs, PCollection<KV<K, OutputT>> newOutput) {
+ Map<TupleTag<?>, PValue> outputs, PCollection<KV<K, OutputT>> newOutput) {
return ReplacementOutputs.singleton(outputs, newOutput);
}
}
@@ -1343,12 +1343,13 @@ public class DataflowRunner extends PipelineRunner<DataflowPipelineJob> {
}
@Override
- public PCollection<T> getInput(List<TaggedPValue> inputs, Pipeline p) {
- return (PCollection<T>) Iterables.getOnlyElement(inputs).getValue();
+ public PCollection<T> getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
+ return (PCollection<T>) Iterables.getOnlyElement(inputs.values());
}
@Override
- public Map<PValue, ReplacementOutput> mapOutputs(List<TaggedPValue> outputs, PDone newOutput) {
+ public Map<PValue, ReplacementOutput> mapOutputs(
+ Map<TupleTag<?>, PValue> outputs, PDone newOutput) {
return Collections.emptyMap();
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 e020e83..52b3a31 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
@@ -29,7 +29,7 @@ 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;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
/**
* A {@link TransformTranslator} knows how to translate a particular subclass of {@link PTransform}
@@ -47,12 +47,12 @@ interface TransformTranslator<TransformT extends PTransform> {
DataflowPipelineOptions getPipelineOptions();
/** Returns the input of the currently being translated transform. */
- <InputT extends PInput> List<TaggedPValue> getInputs(PTransform<InputT, ?> transform);
+ <InputT extends PInput> Map<TupleTag<?>, PValue> getInputs(PTransform<InputT, ?> transform);
<InputT extends PValue> InputT getInput(PTransform<InputT, ?> transform);
/** Returns the output of the currently being translated transform. */
- <OutputT extends POutput> List<TaggedPValue> getOutputs(PTransform<?, OutputT> transform);
+ <OutputT extends POutput> Map<TupleTag<?>, PValue> getOutputs(PTransform<?, OutputT> transform);
<OutputT extends PValue> OutputT getOutput(PTransform<?, OutputT> transform);
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineJobTest.java
----------------------------------------------------------------------
diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineJobTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineJobTest.java
index e3d2e4e..e7f2b48 100644
--- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineJobTest.java
+++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineJobTest.java
@@ -72,7 +72,8 @@ import org.apache.beam.sdk.util.NoopPathValidator;
import org.apache.beam.sdk.util.TestCredential;
import org.apache.beam.sdk.values.PInput;
import org.apache.beam.sdk.values.POutput;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.PValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.joda.time.Duration;
import org.junit.Before;
import org.junit.Rule;
@@ -689,8 +690,8 @@ public class DataflowPipelineJobTest {
when(input.getPipeline()).thenReturn(p);
return AppliedPTransform.of(
fullName,
- Collections.<TaggedPValue>emptyList(),
- Collections.<TaggedPValue>emptyList(),
+ Collections.<TupleTag<?>, PValue>emptyMap(),
+ Collections.<TupleTag<?>, PValue>emptyMap(),
transform,
p);
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 5b4f73e..97487f3 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
@@ -22,6 +22,7 @@ import com.google.common.collect.Iterables;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
+import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
@@ -55,7 +56,7 @@ import org.apache.beam.sdk.values.PCollectionView;
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.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.apache.spark.SparkEnv$;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.metrics.MetricsSystem;
@@ -315,8 +316,7 @@ public final class SparkRunner extends PipelineRunner<SparkPipelineResult> {
// The goal is to detect the PCollections accessed more than one time, and so enable cache
// on the underlying RDDs or DStreams.
- for (TaggedPValue input : node.getInputs()) {
- PValue value = input.getValue();
+ for (PValue value : node.getInputs().values()) {
if (value instanceof PCollection) {
long count = 1L;
if (ctxt.getCacheCandidates().get(value) != null) {
@@ -362,7 +362,7 @@ public final class SparkRunner extends PipelineRunner<SparkPipelineResult> {
if (node.getInputs().size() != 1) {
return false;
}
- PValue input = Iterables.getOnlyElement(node.getInputs()).getValue();
+ PValue input = Iterables.getOnlyElement(node.getInputs().values());
if (!(input instanceof PCollection)
|| ((PCollection) input).getWindowingStrategy().getWindowFn().isNonMerging()) {
return false;
@@ -420,14 +420,14 @@ public final class SparkRunner extends PipelineRunner<SparkPipelineResult> {
//--- determine if node is bounded/unbounded.
// usually, the input determines if the PCollection to apply the next transformation to
// is BOUNDED or UNBOUNDED, meaning RDD/DStream.
- Collection<TaggedPValue> pValues;
+ Map<TupleTag<?>, PValue> pValues;
if (node.getInputs().isEmpty()) {
// in case of a PBegin, it's the output.
pValues = node.getOutputs();
} else {
pValues = node.getInputs();
}
- PCollection.IsBounded isNodeBounded = isBoundedCollection(pValues);
+ PCollection.IsBounded isNodeBounded = isBoundedCollection(pValues.values());
// translate accordingly.
LOG.debug("Translating {} as {}", transform, isNodeBounded);
return isNodeBounded.equals(PCollection.IsBounded.BOUNDED)
@@ -435,15 +435,15 @@ public final class SparkRunner extends PipelineRunner<SparkPipelineResult> {
: translator.translateUnbounded(transformClass);
}
- protected PCollection.IsBounded isBoundedCollection(Collection<TaggedPValue> pValues) {
+ protected PCollection.IsBounded isBoundedCollection(Collection<PValue> pValues) {
// anything that is not a PCollection, is BOUNDED.
// For PCollections:
// BOUNDED behaves as the Identity Element, BOUNDED + BOUNDED = BOUNDED
// while BOUNDED + UNBOUNDED = UNBOUNDED.
PCollection.IsBounded isBounded = PCollection.IsBounded.BOUNDED;
- for (TaggedPValue pValue : pValues) {
- if (pValue.getValue() instanceof PCollection) {
- isBounded = isBounded.and(((PCollection) pValue.getValue()).isBounded());
+ for (PValue pValue : pValues) {
+ if (pValue instanceof PCollection) {
+ isBounded = isBounded.and(((PCollection) pValue).isBounded());
} else {
isBounded = isBounded.and(PCollection.IsBounded.BOUNDED);
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 fcc00f9..aacb942 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
@@ -28,7 +28,6 @@ import com.google.common.util.concurrent.Uninterruptibles;
import java.io.File;
import java.io.IOException;
import java.util.Collections;
-import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import org.apache.beam.runners.core.construction.PTransformMatchers;
@@ -53,7 +52,7 @@ import org.apache.beam.sdk.util.ValueWithRecordId;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PValue;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.apache.commons.io.FileUtils;
import org.joda.time.Duration;
import org.joda.time.Instant;
@@ -251,13 +250,13 @@ public final class TestSparkRunner extends PipelineRunner<SparkPipelineResult> {
}
@Override
- public PBegin getInput(List<TaggedPValue> inputs, Pipeline p) {
+ public PBegin getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
return p.begin();
}
@Override
public Map<PValue, ReplacementOutput> mapOutputs(
- List<TaggedPValue> outputs, PCollection<T> newOutput) {
+ Map<TupleTag<?>, PValue> outputs, PCollection<T> newOutput) {
return ReplacementOutputs.singleton(outputs, newOutput);
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 643749d..838c504 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
@@ -24,7 +24,6 @@ import com.google.common.collect.Iterables;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
-import java.util.List;
import java.util.Map;
import java.util.Set;
import org.apache.beam.runners.spark.SparkPipelineOptions;
@@ -37,7 +36,7 @@ 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.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.streaming.api.java.JavaStreamingContext;
@@ -96,11 +95,11 @@ public class EvaluationContext {
public <T extends PValue> T getInput(PTransform<T, ?> transform) {
@SuppressWarnings("unchecked")
- T input = (T) Iterables.getOnlyElement(getInputs(transform)).getValue();
+ T input = (T) Iterables.getOnlyElement(getInputs(transform).values());
return input;
}
- public <T> List<TaggedPValue> getInputs(PTransform<?, ?> transform) {
+ public <T> Map<TupleTag<?>, PValue> getInputs(PTransform<?, ?> transform) {
checkArgument(currentTransform != null && currentTransform.getTransform() == transform,
"can only be called with current transform");
return currentTransform.getInputs();
@@ -108,11 +107,11 @@ public class EvaluationContext {
public <T extends PValue> T getOutput(PTransform<?, T> transform) {
@SuppressWarnings("unchecked")
- T output = (T) Iterables.getOnlyElement(getOutputs(transform)).getValue();
+ T output = (T) Iterables.getOnlyElement(getOutputs(transform).values());
return output;
}
- public List<TaggedPValue> getOutputs(PTransform<?, ?> transform) {
+ public Map<TupleTag<?>, PValue> getOutputs(PTransform<?, ?> transform) {
checkArgument(currentTransform != null && currentTransform.getTransform() == transform,
"can only be called with current transform");
return currentTransform.getOutputs();
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 7894c4e..c2a8b06 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
@@ -26,9 +26,10 @@ import static org.apache.beam.runners.spark.translation.TranslationUtils.rejectS
import com.google.common.base.Optional;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
+import java.util.Collection;
import java.util.Collections;
-import java.util.List;
import java.util.Map;
+import java.util.Map.Entry;
import org.apache.beam.runners.core.SystemReduceFn;
import org.apache.beam.runners.spark.aggregators.AggregatorsAccumulator;
import org.apache.beam.runners.spark.aggregators.NamedAggregators;
@@ -61,7 +62,7 @@ import org.apache.beam.sdk.util.WindowingStrategy;
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.TaggedPValue;
+import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.spark.Accumulator;
import org.apache.spark.api.java.JavaPairRDD;
@@ -83,19 +84,21 @@ public final class TransformTranslator {
@SuppressWarnings("unchecked")
@Override
public void evaluate(Flatten.PCollections<T> transform, EvaluationContext context) {
- List<TaggedPValue> pcs = context.getInputs(transform);
+ Collection<PValue> pcs = context.getInputs(transform).values();
JavaRDD<WindowedValue<T>> unionRDD;
if (pcs.size() == 0) {
unionRDD = context.getSparkContext().emptyRDD();
} else {
JavaRDD<WindowedValue<T>>[] rdds = new JavaRDD[pcs.size()];
- for (int i = 0; i < rdds.length; i++) {
+ int index = 0;
+ for (PValue pc : pcs) {
checkArgument(
- pcs.get(i).getValue() instanceof PCollection,
+ pc instanceof PCollection,
"Flatten had non-PCollection value in input: %s of type %s",
- pcs.get(i).getValue(),
- pcs.get(i).getValue().getClass().getSimpleName());
- rdds[i] = ((BoundedDataset<T>) context.borrowDataset(pcs.get(i).getValue())).getRDD();
+ pc,
+ pc.getClass().getSimpleName());
+ rdds[index] = ((BoundedDataset<T>) context.borrowDataset(pc)).getRDD();
+ index++;
}
unionRDD = context.getSparkContext().union(rdds);
}
@@ -360,15 +363,15 @@ public final class TransformTranslator {
transform.getMainOutputTag(),
TranslationUtils.getSideInputs(transform.getSideInputs(), context),
windowingStrategy));
- List<TaggedPValue> outputs = context.getOutputs(transform);
+ Map<TupleTag<?>, PValue> outputs = context.getOutputs(transform);
if (outputs.size() > 1) {
// cache the RDD if we're going to filter it more than once.
all.cache();
}
- for (TaggedPValue output : outputs) {
+ for (Map.Entry<TupleTag<?>, PValue> output : outputs.entrySet()) {
@SuppressWarnings("unchecked")
JavaPairRDD<TupleTag<?>, WindowedValue<?>> filtered =
- all.filter(new TranslationUtils.TupleTagFilter(output.getTag()));
+ all.filter(new TranslationUtils.TupleTagFilter(output.getKey()));
@SuppressWarnings("unchecked")
// Object is the best we can do since different outputs can have different tags
JavaRDD<WindowedValue<Object>> values =
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/StreamingTransformTranslator.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/StreamingTransformTranslator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/StreamingTransformTranslator.java
index d4c6c9d..65892d2 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/StreamingTransformTranslator.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/StreamingTransformTranslator.java
@@ -77,7 +77,7 @@ import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.util.WindowingStrategy;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TimestampedValue;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.spark.Accumulator;
@@ -191,19 +191,19 @@ public final class StreamingTransformTranslator {
@SuppressWarnings("unchecked")
@Override
public void evaluate(Flatten.PCollections<T> transform, EvaluationContext context) {
- List<TaggedPValue> pcs = context.getInputs(transform);
+ Map<TupleTag<?>, PValue> pcs = context.getInputs(transform);
// since this is a streaming pipeline, at least one of the PCollections to "flatten" are
// unbounded, meaning it represents a DStream.
// So we could end up with an unbounded unified DStream.
final List<JavaDStream<WindowedValue<T>>> dStreams = new ArrayList<>();
final List<Integer> streamingSources = new ArrayList<>();
- for (TaggedPValue pv : pcs) {
+ for (PValue pv : pcs.values()) {
checkArgument(
- pv.getValue() instanceof PCollection,
+ pv instanceof PCollection,
"Flatten had non-PCollection value in input: %s of type %s",
- pv.getValue(),
- pv.getValue().getClass().getSimpleName());
- PCollection<T> pcol = (PCollection<T>) pv.getValue();
+ pv,
+ pv.getClass().getSimpleName());
+ PCollection<T> pcol = (PCollection<T>) pv;
Dataset dataset = context.borrowDataset(pcol);
if (dataset instanceof UnboundedDataset) {
UnboundedDataset<T> unboundedDataset = (UnboundedDataset<T>) dataset;
@@ -416,15 +416,15 @@ public final class StreamingTransformTranslator {
windowingStrategy));
}
});
- List<TaggedPValue> outputs = context.getOutputs(transform);
+ Map<TupleTag<?>, PValue> outputs = context.getOutputs(transform);
if (outputs.size() > 1) {
// cache the DStream if we're going to filter it more than once.
all.cache();
}
- for (TaggedPValue output : outputs) {
+ for (Map.Entry<TupleTag<?>, PValue> output : outputs.entrySet()) {
@SuppressWarnings("unchecked")
JavaPairDStream<TupleTag<?>, WindowedValue<?>> filtered =
- all.filter(new TranslationUtils.TupleTagFilter(output.getTag()));
+ all.filter(new TranslationUtils.TupleTagFilter(output.getKey()));
@SuppressWarnings("unchecked")
// Object is the best we can do since different outputs can have different tags
JavaDStream<WindowedValue<Object>> values =
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/PTransformOverrideFactory.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/PTransformOverrideFactory.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/PTransformOverrideFactory.java
index e2b6009..57cba50 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/PTransformOverrideFactory.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/runners/PTransformOverrideFactory.java
@@ -20,7 +20,6 @@
package org.apache.beam.sdk.runners;
import com.google.auto.value.AutoValue;
-import java.util.List;
import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.annotations.Experimental;
@@ -30,6 +29,7 @@ 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.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
/**
* Produces {@link PipelineRunner}-specific overrides of {@link PTransform PTransforms}, and
@@ -48,17 +48,15 @@ public interface PTransformOverrideFactory<
/**
* Returns the composite type that replacement transforms consumed from an equivalent expansion.
*/
- InputT getInput(List<TaggedPValue> inputs, Pipeline p);
+ InputT getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p);
/**
* Returns a {@link Map} from the expanded values in {@code newOutput} to the values produced by
* the original transform.
*/
- Map<PValue, ReplacementOutput> mapOutputs(List<TaggedPValue> outputs, OutputT newOutput);
+ Map<PValue, ReplacementOutput> mapOutputs(Map<TupleTag<?>, PValue> outputs, OutputT newOutput);
- /**
- * A mapping between original {@link TaggedPValue} outputs and their replacements.
- */
+ /** A mapping between original {@link TaggedPValue} outputs and their replacements. */
@AutoValue
abstract class ReplacementOutput {
public static ReplacementOutput of(TaggedPValue original, TaggedPValue replacement) {
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 972cb5b..18bf2e9 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
@@ -23,6 +23,7 @@ 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 java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
@@ -41,7 +42,7 @@ import org.apache.beam.sdk.transforms.PTransform;
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.TaggedPValue;
+import org.apache.beam.sdk.values.TupleTag;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -103,8 +104,8 @@ public class TransformHierarchy {
"Replacing a node when the graph has an unexpanded input. This is an SDK bug.");
Node replacement =
new Node(existing.getEnclosingNode(), transform, existing.getFullName(), input);
- for (TaggedPValue output : existing.getOutputs()) {
- Node producer = producers.get(output.getValue());
+ for (PValue output : existing.getOutputs().values()) {
+ Node producer = producers.get(output);
boolean producedInExisting = false;
do {
if (producer.equals(existing)) {
@@ -114,13 +115,13 @@ public class TransformHierarchy {
}
} while (!producedInExisting && !producer.isRootNode());
if (producedInExisting) {
- producers.remove(output.getValue());
+ producers.remove(output);
LOG.debug("Removed producer for value {} as it is part of a replaced composite {}",
- output.getValue(),
+ output,
existing.getFullName());
} else {
LOG.debug(
- "Value {} not produced in existing node {}", output.getValue(), existing.getFullName());
+ "Value {} not produced in existing node {}", output, existing.getFullName());
}
}
existing.getEnclosingNode().replaceChild(existing, replacement);
@@ -137,18 +138,18 @@ public class TransformHierarchy {
*/
public void finishSpecifyingInput() {
// Inputs must be completely specified before they are consumed by a transform.
- for (TaggedPValue inputValue : current.getInputs()) {
- Node producerNode = getProducer(inputValue.getValue());
- PInput input = producerInput.remove(inputValue.getValue());
- inputValue.getValue().finishSpecifying(input, producerNode.getTransform());
+ for (PValue inputValue : current.getInputs().values()) {
+ Node producerNode = getProducer(inputValue);
+ PInput input = producerInput.remove(inputValue);
+ inputValue.finishSpecifying(input, producerNode.getTransform());
checkState(
- producers.get(inputValue.getValue()) != null,
+ producers.get(inputValue) != null,
"Producer unknown for input %s",
inputValue);
checkState(
- producers.get(inputValue.getValue()) != null,
+ producers.get(inputValue) != null,
"Producer unknown for input %s",
- inputValue.getValue());
+ inputValue);
}
}
@@ -163,12 +164,12 @@ public class TransformHierarchy {
* nodes.
*/
public void setOutput(POutput output) {
- for (TaggedPValue value : output.expand()) {
- if (!producers.containsKey(value.getValue())) {
- producers.put(value.getValue(), current);
+ for (PValue value : output.expand().values()) {
+ if (!producers.containsKey(value)) {
+ producers.put(value, current);
}
- value.getValue().finishSpecifyingOutput(unexpandedInputs.get(current), current.transform);
- producerInput.put(value.getValue(), unexpandedInputs.get(current));
+ value.finishSpecifyingOutput(unexpandedInputs.get(current), current.transform);
+ producerInput.put(value, unexpandedInputs.get(current));
}
output.finishSpecifyingOutput(unexpandedInputs.get(current), current.transform);
current.setOutput(output);
@@ -241,11 +242,11 @@ public class TransformHierarchy {
private final List<Node> parts = new ArrayList<>();
// Input to the transform, in expanded form.
- private final List<TaggedPValue> inputs;
+ private final Map<TupleTag<?>, PValue> inputs;
// TODO: track which outputs need to be exported to parent.
// Output of the transform, in expanded form.
- private List<TaggedPValue> outputs;
+ private Map<TupleTag<?>, PValue> outputs;
@VisibleForTesting
boolean finishedSpecifying = false;
@@ -269,7 +270,7 @@ public class TransformHierarchy {
this.enclosingNode = enclosingNode;
this.transform = transform;
this.fullName = fullName;
- this.inputs = input == null ? Collections.<TaggedPValue>emptyList() : input.expand();
+ this.inputs = input == null ? Collections.<TupleTag<?>, PValue>emptyMap() : input.expand();
}
/**
@@ -333,8 +334,8 @@ public class TransformHierarchy {
private boolean returnsOthersOutput() {
PTransform<?, ?> transform = getTransform();
if (outputs != null) {
- for (TaggedPValue outputValue : outputs) {
- if (!getProducer(outputValue.getValue()).getTransform().equals(transform)) {
+ for (PValue outputValue : outputs.values()) {
+ if (!getProducer(outputValue).getTransform().equals(transform)) {
return true;
}
}
@@ -351,8 +352,8 @@ public class TransformHierarchy {
}
/** Returns the transform input, in unexpanded form. */
- public List<TaggedPValue> getInputs() {
- return inputs == null ? Collections.<TaggedPValue>emptyList() : inputs;
+ public Map<TupleTag<?>, PValue> getInputs() {
+ return inputs == null ? Collections.<TupleTag<?>, PValue>emptyMap() : inputs;
}
/**
@@ -368,8 +369,8 @@ public class TransformHierarchy {
// Validate that a primitive transform produces only primitive output, and a composite
// transform does not produce primitive output.
Set<Node> outputProducers = new HashSet<>();
- for (TaggedPValue outputValue : output.expand()) {
- outputProducers.add(getProducer(outputValue.getValue()));
+ for (PValue outputValue : output.expand().values()) {
+ outputProducers.add(getProducer(outputValue));
}
if (outputProducers.contains(this)) {
if (!parts.isEmpty() || outputProducers.size() > 1) {
@@ -412,8 +413,8 @@ public class TransformHierarchy {
// Replace the outputs of the component nodes
component.replaceOutputs(originalToReplacement);
}
- List<TaggedPValue> newOutputs = new ArrayList<>(outputs.size());
- for (TaggedPValue output : outputs) {
+ ImmutableMap.Builder<TupleTag<?>, PValue> newOutputsBuilder = ImmutableMap.builder();
+ for (Map.Entry<TupleTag<?>, PValue> output : outputs.entrySet()) {
ReplacementOutput mapping = originalToReplacement.get(output.getValue());
if (mapping != null) {
if (this.equals(producers.get(mapping.getReplacement().getValue()))) {
@@ -429,11 +430,12 @@ public class TransformHierarchy {
"Replacing output {} with original {}",
mapping.getReplacement(),
mapping.getOriginal());
- newOutputs.add(TaggedPValue.of(output.getTag(), mapping.getOriginal().getValue()));
+ newOutputsBuilder.put(output.getKey(), mapping.getOriginal().getValue());
} else {
- newOutputs.add(output);
+ newOutputsBuilder.put(output);
}
}
+ ImmutableMap<TupleTag<?>, PValue> newOutputs = newOutputsBuilder.build();
checkState(
outputs.size() == newOutputs.size(),
"Number of outputs must be stable across replacement");
@@ -441,8 +443,8 @@ public class TransformHierarchy {
}
/** Returns the transform output, in expanded form. */
- public List<TaggedPValue> getOutputs() {
- return outputs == null ? Collections.<TaggedPValue>emptyList() : outputs;
+ public Map<TupleTag<?>, PValue> getOutputs() {
+ return outputs == null ? Collections.<TupleTag<?>, PValue>emptyMap() : outputs;
}
/**
@@ -466,9 +468,9 @@ public class TransformHierarchy {
if (!isRootNode()) {
// Visit inputs.
- for (TaggedPValue inputValue : inputs) {
- if (visitedValues.add(inputValue.getValue())) {
- visitor.visitValue(inputValue.getValue(), getProducer(inputValue.getValue()));
+ for (PValue inputValue : inputs.values()) {
+ if (visitedValues.add(inputValue)) {
+ visitor.visitValue(inputValue, getProducer(inputValue));
}
}
}
@@ -489,9 +491,9 @@ public class TransformHierarchy {
if (!isRootNode()) {
checkNotNull(outputs, "Outputs for non-root node %s are null", getFullName());
// Visit outputs.
- for (TaggedPValue pValue : outputs) {
- if (visitedValues.add(pValue.getValue())) {
- visitor.visitValue(pValue.getValue(), this);
+ for (PValue pValue : outputs.values()) {
+ if (visitedValues.add(pValue)) {
+ visitor.visitValue(pValue, this);
}
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AppliedPTransform.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AppliedPTransform.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AppliedPTransform.java
index e78d795..8d99a62 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AppliedPTransform.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/AppliedPTransform.java
@@ -18,11 +18,12 @@
package org.apache.beam.sdk.transforms;
import com.google.auto.value.AutoValue;
-import java.util.List;
+import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.values.PInput;
import org.apache.beam.sdk.values.POutput;
-import org.apache.beam.sdk.values.TaggedPValue;
+import org.apache.beam.sdk.values.PValue;
+import org.apache.beam.sdk.values.TupleTag;
/**
* Represents the application of a {@link PTransform} to a specific input to produce
@@ -41,12 +42,14 @@ public abstract class AppliedPTransform<
// To prevent extension outside of this package.
AppliedPTransform() {}
- public static <InputT extends PInput, OutputT extends POutput,
+ public static <
+ InputT extends PInput,
+ OutputT extends POutput,
TransformT extends PTransform<? super InputT, OutputT>>
AppliedPTransform<InputT, OutputT, TransformT> of(
String fullName,
- List<TaggedPValue> input,
- List<TaggedPValue> output,
+ Map<TupleTag<?>, PValue> input,
+ Map<TupleTag<?>, PValue> output,
TransformT transform,
Pipeline p) {
return new AutoValue_AppliedPTransform<InputT, OutputT, TransformT>(
@@ -55,9 +58,9 @@ public abstract class AppliedPTransform<
public abstract String getFullName();
- public abstract List<TaggedPValue> getInputs();
+ public abstract Map<TupleTag<?>, PValue> getInputs();
- public abstract List<TaggedPValue> getOutputs();
+ public abstract Map<TupleTag<?>, PValue> getOutputs();
public abstract TransformT getTransform();
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/join/KeyedPCollectionTuple.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/join/KeyedPCollectionTuple.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/join/KeyedPCollectionTuple.java
index b373909..2e7dd01 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/join/KeyedPCollectionTuple.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/join/KeyedPCollectionTuple.java
@@ -17,8 +17,10 @@
*/
package org.apache.beam.sdk.transforms.join;
+import com.google.common.collect.ImmutableMap;
import java.util.ArrayList;
import java.util.List;
+import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
@@ -27,7 +29,7 @@ import org.apache.beam.sdk.values.KV;
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.TaggedPValue;
+import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
@@ -120,12 +122,12 @@ public class KeyedPCollectionTuple<K> implements PInput {
* any tag-specific information.
*/
@Override
- public List<TaggedPValue> expand() {
- List<TaggedPValue> retval = new ArrayList<>();
+ public Map<TupleTag<?>, PValue> expand() {
+ ImmutableMap.Builder<TupleTag<?>, PValue> retval = ImmutableMap.builder();
for (TaggedKeyedPCollection<K, ?> taggedPCollection : keyedCollections) {
- retval.add(TaggedPValue.of(taggedPCollection.tupleTag, taggedPCollection.pCollection));
+ retval.put(taggedPCollection.tupleTag, taggedPCollection.pCollection);
}
- return retval;
+ return retval.build();
}
/**
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PBegin.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PBegin.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PBegin.java
index 2ba0f1c..04d1bdb 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PBegin.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PBegin.java
@@ -18,7 +18,7 @@
package org.apache.beam.sdk.values;
import java.util.Collections;
-import java.util.List;
+import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.TextIO.Read;
import org.apache.beam.sdk.transforms.Create;
@@ -64,9 +64,9 @@ public class PBegin implements PInput {
}
@Override
- public List<TaggedPValue> expand() {
+ public Map<TupleTag<?>, PValue> expand() {
// A PBegin contains no PValues.
- return Collections.emptyList();
+ return Collections.emptyMap();
}
/////////////////////////////////////////////////////////////////////////////
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionList.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionList.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionList.java
index dcb64a8..7b45deb 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionList.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionList.java
@@ -18,9 +18,10 @@
package org.apache.beam.sdk.values;
import com.google.common.collect.ImmutableList;
-import com.google.common.collect.Iterables;
+import com.google.common.collect.ImmutableMap;
import java.util.Iterator;
import java.util.List;
+import java.util.Map;
import java.util.Objects;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.transforms.AppliedPTransform;
@@ -116,7 +117,7 @@ public class PCollectionList<T> implements PInput, POutput {
return new PCollectionList<>(pipeline,
ImmutableList.<TaggedPValue>builder()
.addAll(pcollections)
- .add(Iterables.getOnlyElement(pc.expand()))
+ .add(TaggedPValue.of(new TupleTag<T>(), pc))
.build());
}
@@ -133,10 +134,9 @@ public class PCollectionList<T> implements PInput, POutput {
builder.addAll(pcollections);
for (PCollection<T> pc : pcs) {
if (pc.getPipeline() != pipeline) {
- throw new IllegalArgumentException(
- "PCollections come from different Pipelines");
+ throw new IllegalArgumentException("PCollections come from different Pipelines");
}
- builder.add(Iterables.getOnlyElement(pc.expand()));
+ builder.add(TaggedPValue.of(new TupleTag<T>(), pc));
}
return new PCollectionList<>(pipeline, builder.build());
}
@@ -200,7 +200,10 @@ public class PCollectionList<T> implements PInput, POutput {
// Internal details below here.
final Pipeline pipeline;
- // ImmutableMap has a defined iteration order.
+ /**
+ * The {@link PCollection PCollections} contained by this {@link PCollectionList}, and an
+ * arbitrary tags associated with each.
+ */
final List<TaggedPValue> pcollections;
PCollectionList(Pipeline pipeline) {
@@ -218,8 +221,12 @@ public class PCollectionList<T> implements PInput, POutput {
}
@Override
- public List<TaggedPValue> expand() {
- return pcollections;
+ public Map<TupleTag<?>, PValue> expand() {
+ ImmutableMap.Builder<TupleTag<?>, PValue> expanded = ImmutableMap.builder();
+ for (TaggedPValue tagged : pcollections) {
+ expanded.put(tagged.getTag(), tagged.getValue());
+ }
+ return expanded.build();
}
@Override
@@ -244,11 +251,11 @@ public class PCollectionList<T> implements PInput, POutput {
return false;
}
PCollectionList that = (PCollectionList) other;
- return this.pipeline.equals(that.pipeline) && this.pcollections.equals(that.pcollections);
+ return this.pipeline.equals(that.pipeline) && this.getAll().equals(that.getAll());
}
@Override
public int hashCode() {
- return Objects.hash(this.pipeline, this.pcollections);
+ return Objects.hash(this.pipeline, this.getAll());
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionTuple.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionTuple.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionTuple.java
index d61db51..0ab26ca 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionTuple.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionTuple.java
@@ -17,11 +17,9 @@
*/
package org.apache.beam.sdk.values;
-import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import java.util.Collections;
import java.util.LinkedHashMap;
-import java.util.List;
import java.util.Map;
import java.util.Objects;
import org.apache.beam.sdk.Pipeline;
@@ -234,12 +232,8 @@ public class PCollectionTuple implements PInput, POutput {
}
@Override
- public List<TaggedPValue> expand() {
- ImmutableList.Builder<TaggedPValue> values = ImmutableList.builder();
- for (Map.Entry<TupleTag<?>, PCollection<?>> entry : pcollectionMap.entrySet()) {
- values.add(TaggedPValue.of(entry.getKey(), entry.getValue()));
- }
- return values.build();
+ public Map<TupleTag<?>, PValue> expand() {
+ return ImmutableMap.<TupleTag<?>, PValue>copyOf(pcollectionMap);
}
@Override
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PDone.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PDone.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PDone.java
index b4a3025..eb5db20 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PDone.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PDone.java
@@ -18,7 +18,7 @@
package org.apache.beam.sdk.values;
import java.util.Collections;
-import java.util.List;
+import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.transforms.PTransform;
@@ -36,9 +36,9 @@ public class PDone extends POutputValueBase {
}
@Override
- public List<TaggedPValue> expand() {
+ public Map<TupleTag<?>, PValue> expand() {
// A PDone contains no PValues.
- return Collections.emptyList();
+ return Collections.emptyMap();
}
private PDone(Pipeline pipeline) {
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PInput.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PInput.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PInput.java
index 30d4297..caf7812 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PInput.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PInput.java
@@ -17,7 +17,7 @@
*/
package org.apache.beam.sdk.values;
-import java.util.List;
+import java.util.Map;
import org.apache.beam.sdk.Pipeline;
/**
@@ -43,5 +43,5 @@ public interface PInput {
*
* <p>Not intended to be invoked directly by user code.
*/
- List<TaggedPValue> expand();
+ Map<TupleTag<?>, PValue> expand();
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/sdks/java/core/src/main/java/org/apache/beam/sdk/values/POutput.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/POutput.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/POutput.java
index 062f565..bb01beb 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/POutput.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/POutput.java
@@ -17,7 +17,7 @@
*/
package org.apache.beam.sdk.values;
-import java.util.List;
+import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.transforms.AppliedPTransform;
import org.apache.beam.sdk.transforms.PTransform;
@@ -45,7 +45,7 @@ public interface POutput {
*
* <p>Not intended to be invoked directly by user code.
*/
- List<TaggedPValue> expand();
+ Map<TupleTag<?>, PValue> expand();
/**
* Records that this {@code POutput} is an output of the given
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PValue.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PValue.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PValue.java
index 4c62972..06546aa 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PValue.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PValue.java
@@ -17,7 +17,7 @@
*/
package org.apache.beam.sdk.values;
-import java.util.List;
+import java.util.Map;
import org.apache.beam.sdk.transforms.PTransform;
/**
@@ -37,7 +37,7 @@ public interface PValue extends POutput, PInput {
* never appropriate.
*/
@Deprecated
- List<TaggedPValue> expand();
+ Map<TupleTag<?>, PValue> expand();
/**
* After building, finalizes this {@code PValue} to make it ready for being used as an input to a
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 8778597..91ee392 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
@@ -18,7 +18,7 @@
package org.apache.beam.sdk.values;
import java.util.Collections;
-import java.util.List;
+import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.transforms.AppliedPTransform;
import org.apache.beam.sdk.transforms.PTransform;
@@ -133,8 +133,8 @@ public abstract class PValueBase extends POutputValueBase implements PValue {
}
@Override
- public final List<TaggedPValue> expand() {
- return Collections.singletonList(TaggedPValue.of(tag, this));
+ public final Map<TupleTag<?>, PValue> expand() {
+ return Collections.<TupleTag<?>, PValue>singletonMap(tag, this);
}
@Override
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/sdks/java/core/src/main/java/org/apache/beam/sdk/values/TaggedPValue.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/TaggedPValue.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/TaggedPValue.java
index 458d16f..3b4d599 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/TaggedPValue.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/TaggedPValue.java
@@ -20,6 +20,7 @@
package org.apache.beam.sdk.values;
import com.google.auto.value.AutoValue;
+import com.google.common.collect.Iterables;
/**
* A (TupleTag, PValue) pair used in the expansion of a {@link PInput} or {@link POutput}.
@@ -30,6 +31,10 @@ public abstract class TaggedPValue {
return new AutoValue_TaggedPValue(tag, value);
}
+ public static TaggedPValue ofExpandedValue(PValue value) {
+ return of(Iterables.getOnlyElement(value.expand().keySet()), value);
+ }
+
/**
* Returns the local tag associated with the {@link PValue}.
*/
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/sdks/java/core/src/test/java/org/apache/beam/sdk/PipelineTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/PipelineTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/PipelineTest.java
index efe8db4..0a5746b 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/PipelineTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/PipelineTest.java
@@ -393,17 +393,21 @@ public class PipelineTest {
}
@Override
- public PBegin getInput(List<TaggedPValue> inputs, Pipeline p) {
+ public PBegin getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
return p.begin();
}
@Override
public Map<PValue, ReplacementOutput> mapOutputs(
- List<TaggedPValue> outputs, PCollection<Long> newOutput) {
+ Map<TupleTag<?>, PValue> outputs, PCollection<Long> newOutput) {
+ Map.Entry<TupleTag<?>, PValue> original = Iterables.getOnlyElement(outputs.entrySet());
+ Map.Entry<TupleTag<?>, PValue> replacement =
+ Iterables.getOnlyElement(newOutput.expand().entrySet());
return Collections.<PValue, ReplacementOutput>singletonMap(
newOutput,
ReplacementOutput.of(
- Iterables.getOnlyElement(outputs), Iterables.getOnlyElement(newOutput.expand())));
+ TaggedPValue.of(original.getKey(), original.getValue()),
+ TaggedPValue.of(replacement.getKey(), replacement.getValue())));
}
}
static class UnboundedCountingInputOverride
@@ -415,17 +419,21 @@ public class PipelineTest {
}
@Override
- public PBegin getInput(List<TaggedPValue> inputs, Pipeline p) {
+ public PBegin getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
return p.begin();
}
@Override
public Map<PValue, ReplacementOutput> mapOutputs(
- List<TaggedPValue> outputs, PCollection<Long> newOutput) {
+ Map<TupleTag<?>, PValue> outputs, PCollection<Long> newOutput) {
+ Map.Entry<TupleTag<?>, PValue> original = Iterables.getOnlyElement(outputs.entrySet());
+ Map.Entry<TupleTag<?>, PValue> replacement =
+ Iterables.getOnlyElement(newOutput.expand().entrySet());
return Collections.<PValue, ReplacementOutput>singletonMap(
newOutput,
ReplacementOutput.of(
- Iterables.getOnlyElement(outputs), Iterables.getOnlyElement(newOutput.expand())));
+ TaggedPValue.of(original.getKey(), original.getValue()),
+ TaggedPValue.of(replacement.getKey(), replacement.getValue())));
}
}
}
[30/50] [abbrv] beam git commit: [BEAM-1722] Move PubsubIO into the
google-cloud-platform module
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java
new file mode 100644
index 0000000..8fc1c19
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java
@@ -0,0 +1,1014 @@
+/*
+ * 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.pubsub;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.common.base.Preconditions.checkState;
+
+import java.io.IOException;
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.VoidCoder;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.OutgoingMessage;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.ProjectPath;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath;
+import org.apache.beam.sdk.options.PubsubOptions;
+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.runners.PipelineRunner;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.windowing.AfterWatermark;
+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.PDone;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Read and Write {@link PTransform}s for Cloud Pub/Sub streams. These transforms create
+ * and consume unbounded {@link PCollection PCollections}.
+ *
+ * <h3>Permissions</h3>
+ *
+ * <p>Permission requirements depend on the {@link PipelineRunner} that is used to execute the
+ * Beam pipeline. Please refer to the documentation of corresponding
+ * {@link PipelineRunner PipelineRunners} for more details.
+ */
+public class PubsubIO {
+
+ private static final Logger LOG = LoggerFactory.getLogger(PubsubIO.class);
+
+ /** Factory for creating pubsub client to manage transport. */
+ private static final PubsubClient.PubsubClientFactory FACTORY = PubsubJsonClient.FACTORY;
+
+ /**
+ * Project IDs must contain 6-63 lowercase letters, digits, or dashes.
+ * IDs must start with a letter and may not end with a dash.
+ * This regex isn't exact - this allows for patterns that would be rejected by
+ * the service, but this is sufficient for basic parsing of table references.
+ */
+ private static final Pattern PROJECT_ID_REGEXP =
+ Pattern.compile("[a-z][-a-z0-9:.]{4,61}[a-z0-9]");
+
+ private static final Pattern SUBSCRIPTION_REGEXP =
+ Pattern.compile("projects/([^/]+)/subscriptions/(.+)");
+
+ private static final Pattern TOPIC_REGEXP = Pattern.compile("projects/([^/]+)/topics/(.+)");
+
+ private static final Pattern V1BETA1_SUBSCRIPTION_REGEXP =
+ Pattern.compile("/subscriptions/([^/]+)/(.+)");
+
+ private static final Pattern V1BETA1_TOPIC_REGEXP = Pattern.compile("/topics/([^/]+)/(.+)");
+
+ private static final Pattern PUBSUB_NAME_REGEXP = Pattern.compile("[a-zA-Z][-._~%+a-zA-Z0-9]+");
+
+ private static final int PUBSUB_NAME_MIN_LENGTH = 3;
+ private static final int PUBSUB_NAME_MAX_LENGTH = 255;
+
+ private static final String SUBSCRIPTION_RANDOM_TEST_PREFIX = "_random/";
+ private static final String SUBSCRIPTION_STARTING_SIGNAL = "_starting_signal/";
+ private static final String TOPIC_DEV_NULL_TEST_NAME = "/topics/dev/null";
+
+ private static void validateProjectName(String project) {
+ Matcher match = PROJECT_ID_REGEXP.matcher(project);
+ if (!match.matches()) {
+ throw new IllegalArgumentException(
+ "Illegal project name specified in Pubsub subscription: " + project);
+ }
+ }
+
+ private static void validatePubsubName(String name) {
+ if (name.length() < PUBSUB_NAME_MIN_LENGTH) {
+ throw new IllegalArgumentException(
+ "Pubsub object name is shorter than 3 characters: " + name);
+ }
+ if (name.length() > PUBSUB_NAME_MAX_LENGTH) {
+ throw new IllegalArgumentException(
+ "Pubsub object name is longer than 255 characters: " + name);
+ }
+
+ if (name.startsWith("goog")) {
+ throw new IllegalArgumentException("Pubsub object name cannot start with goog: " + name);
+ }
+
+ Matcher match = PUBSUB_NAME_REGEXP.matcher(name);
+ if (!match.matches()) {
+ throw new IllegalArgumentException("Illegal Pubsub object name specified: " + name
+ + " Please see Javadoc for naming rules.");
+ }
+ }
+
+ /**
+ * Populate common {@link DisplayData} between Pubsub source and sink.
+ */
+ private static void populateCommonDisplayData(DisplayData.Builder builder,
+ String timestampLabel, String idLabel, ValueProvider<PubsubTopic> topic) {
+ builder
+ .addIfNotNull(DisplayData.item("timestampLabel", timestampLabel)
+ .withLabel("Timestamp Label Attribute"))
+ .addIfNotNull(DisplayData.item("idLabel", idLabel)
+ .withLabel("ID Label Attribute"));
+
+ if (topic != null) {
+ String topicString = topic.isAccessible() ? topic.get().asPath()
+ : topic.toString();
+ builder.add(DisplayData.item("topic", topicString)
+ .withLabel("Pubsub Topic"));
+ }
+ }
+
+ /**
+ * Class representing a Pub/Sub message. Each message contains a single message payload and
+ * a map of attached attributes.
+ */
+ public static class PubsubMessage {
+
+ private byte[] message;
+ private Map<String, String> attributes;
+
+ public PubsubMessage(byte[] message, Map<String, String> attributes) {
+ this.message = message;
+ this.attributes = attributes;
+ }
+
+ /**
+ * Returns the main PubSub message.
+ */
+ public byte[] getMessage() {
+ return message;
+ }
+
+ /**
+ * Returns the given attribute value. If not such attribute exists, returns null.
+ */
+ @Nullable
+ public String getAttribute(String attribute) {
+ checkNotNull(attribute, "attribute");
+ return attributes.get(attribute);
+ }
+
+ /**
+ * Returns the full map of attributes. This is an unmodifiable map.
+ */
+ public Map<String, String> getAttributeMap() {
+ return attributes;
+ }
+ }
+
+ /**
+ * Class representing a Cloud Pub/Sub Subscription.
+ */
+ public static class PubsubSubscription implements Serializable {
+
+ private enum Type {NORMAL, FAKE}
+
+ private final Type type;
+ private final String project;
+ private final String subscription;
+
+ private PubsubSubscription(Type type, String project, String subscription) {
+ this.type = type;
+ this.project = project;
+ this.subscription = subscription;
+ }
+
+ /**
+ * Creates a class representing a Pub/Sub subscription from the specified subscription path.
+ *
+ * <p>Cloud Pub/Sub subscription names should be of the form
+ * {@code projects/<project>/subscriptions/<subscription>}, where {@code <project>} is the name
+ * of the project the subscription belongs to. The {@code <subscription>} component must comply
+ * with the following requirements:
+ *
+ * <ul>
+ * <li>Can only contain lowercase letters, numbers, dashes ('-'), underscores ('_') and periods
+ * ('.').</li>
+ * <li>Must be between 3 and 255 characters.</li>
+ * <li>Must begin with a letter.</li>
+ * <li>Must end with a letter or a number.</li>
+ * <li>Cannot begin with {@code 'goog'} prefix.</li>
+ * </ul>
+ */
+ public static PubsubSubscription fromPath(String path) {
+ if (path.startsWith(SUBSCRIPTION_RANDOM_TEST_PREFIX)
+ || path.startsWith(SUBSCRIPTION_STARTING_SIGNAL)) {
+ return new PubsubSubscription(Type.FAKE, "", path);
+ }
+
+ String projectName, subscriptionName;
+
+ Matcher v1beta1Match = V1BETA1_SUBSCRIPTION_REGEXP.matcher(path);
+ if (v1beta1Match.matches()) {
+ LOG.warn("Saw subscription in v1beta1 format. Subscriptions should be in the format "
+ + "projects/<project_id>/subscriptions/<subscription_name>");
+ projectName = v1beta1Match.group(1);
+ subscriptionName = v1beta1Match.group(2);
+ } else {
+ Matcher match = SUBSCRIPTION_REGEXP.matcher(path);
+ if (!match.matches()) {
+ throw new IllegalArgumentException("Pubsub subscription is not in "
+ + "projects/<project_id>/subscriptions/<subscription_name> format: " + path);
+ }
+ projectName = match.group(1);
+ subscriptionName = match.group(2);
+ }
+
+ validateProjectName(projectName);
+ validatePubsubName(subscriptionName);
+ return new PubsubSubscription(Type.NORMAL, projectName, subscriptionName);
+ }
+
+ /**
+ * Returns the string representation of this subscription as a path used in the Cloud Pub/Sub
+ * v1beta1 API.
+ *
+ * @deprecated the v1beta1 API for Cloud Pub/Sub is deprecated.
+ */
+ @Deprecated
+ public String asV1Beta1Path() {
+ if (type == Type.NORMAL) {
+ return "/subscriptions/" + project + "/" + subscription;
+ } else {
+ return subscription;
+ }
+ }
+
+ /**
+ * Returns the string representation of this subscription as a path used in the Cloud Pub/Sub
+ * v1beta2 API.
+ *
+ * @deprecated the v1beta2 API for Cloud Pub/Sub is deprecated.
+ */
+ @Deprecated
+ public String asV1Beta2Path() {
+ if (type == Type.NORMAL) {
+ return "projects/" + project + "/subscriptions/" + subscription;
+ } else {
+ return subscription;
+ }
+ }
+
+ /**
+ * Returns the string representation of this subscription as a path used in the Cloud Pub/Sub
+ * API.
+ */
+ public String asPath() {
+ if (type == Type.NORMAL) {
+ return "projects/" + project + "/subscriptions/" + subscription;
+ } else {
+ return subscription;
+ }
+ }
+ }
+
+ /**
+ * Used to build a {@link ValueProvider} for {@link PubsubSubscription}.
+ */
+ private static class SubscriptionTranslator
+ implements SerializableFunction<String, PubsubSubscription> {
+
+ @Override
+ public PubsubSubscription apply(String from) {
+ return PubsubSubscription.fromPath(from);
+ }
+ }
+
+ /**
+ * Used to build a {@link ValueProvider} for {@link SubscriptionPath}.
+ */
+ private static class SubscriptionPathTranslator
+ implements SerializableFunction<PubsubSubscription, SubscriptionPath> {
+
+ @Override
+ public SubscriptionPath apply(PubsubSubscription from) {
+ return PubsubClient.subscriptionPathFromName(from.project, from.subscription);
+ }
+ }
+
+ /**
+ * Used to build a {@link ValueProvider} for {@link PubsubTopic}.
+ */
+ private static class TopicTranslator
+ implements SerializableFunction<String, PubsubTopic> {
+
+ @Override
+ public PubsubTopic apply(String from) {
+ return PubsubTopic.fromPath(from);
+ }
+ }
+
+ /**
+ * Used to build a {@link ValueProvider} for {@link TopicPath}.
+ */
+ private static class TopicPathTranslator
+ implements SerializableFunction<PubsubTopic, TopicPath> {
+
+ @Override
+ public TopicPath apply(PubsubTopic from) {
+ return PubsubClient.topicPathFromName(from.project, from.topic);
+ }
+ }
+
+ /**
+ * Used to build a {@link ValueProvider} for {@link ProjectPath}.
+ */
+ private static class ProjectPathTranslator
+ implements SerializableFunction<PubsubTopic, ProjectPath> {
+
+ @Override
+ public ProjectPath apply(PubsubTopic from) {
+ return PubsubClient.projectPathFromId(from.project);
+ }
+ }
+
+ /**
+ * Class representing a Cloud Pub/Sub Topic.
+ */
+ public static class PubsubTopic implements Serializable {
+
+ private enum Type {NORMAL, FAKE}
+
+ private final Type type;
+ private final String project;
+ private final String topic;
+
+ private PubsubTopic(Type type, String project, String topic) {
+ this.type = type;
+ this.project = project;
+ this.topic = topic;
+ }
+
+ /**
+ * Creates a class representing a Cloud Pub/Sub topic from the specified topic path.
+ *
+ * <p>Cloud Pub/Sub topic names should be of the form
+ * {@code /topics/<project>/<topic>}, where {@code <project>} is the name of
+ * the publishing project. The {@code <topic>} component must comply with
+ * the following requirements:
+ *
+ * <ul>
+ * <li>Can only contain lowercase letters, numbers, dashes ('-'), underscores ('_') and periods
+ * ('.').</li>
+ * <li>Must be between 3 and 255 characters.</li>
+ * <li>Must begin with a letter.</li>
+ * <li>Must end with a letter or a number.</li>
+ * <li>Cannot begin with 'goog' prefix.</li>
+ * </ul>
+ */
+ public static PubsubTopic fromPath(String path) {
+ if (path.equals(TOPIC_DEV_NULL_TEST_NAME)) {
+ return new PubsubTopic(Type.FAKE, "", path);
+ }
+
+ String projectName, topicName;
+
+ Matcher v1beta1Match = V1BETA1_TOPIC_REGEXP.matcher(path);
+ if (v1beta1Match.matches()) {
+ LOG.warn("Saw topic in v1beta1 format. Topics should be in the format "
+ + "projects/<project_id>/topics/<topic_name>");
+ projectName = v1beta1Match.group(1);
+ topicName = v1beta1Match.group(2);
+ } else {
+ Matcher match = TOPIC_REGEXP.matcher(path);
+ if (!match.matches()) {
+ throw new IllegalArgumentException(
+ "Pubsub topic is not in projects/<project_id>/topics/<topic_name> format: " + path);
+ }
+ projectName = match.group(1);
+ topicName = match.group(2);
+ }
+
+ validateProjectName(projectName);
+ validatePubsubName(topicName);
+ return new PubsubTopic(Type.NORMAL, projectName, topicName);
+ }
+
+ /**
+ * Returns the string representation of this topic as a path used in the Cloud Pub/Sub
+ * v1beta1 API.
+ *
+ * @deprecated the v1beta1 API for Cloud Pub/Sub is deprecated.
+ */
+ @Deprecated
+ public String asV1Beta1Path() {
+ if (type == Type.NORMAL) {
+ return "/topics/" + project + "/" + topic;
+ } else {
+ return topic;
+ }
+ }
+
+ /**
+ * Returns the string representation of this topic as a path used in the Cloud Pub/Sub
+ * v1beta2 API.
+ *
+ * @deprecated the v1beta2 API for Cloud Pub/Sub is deprecated.
+ */
+ @Deprecated
+ public String asV1Beta2Path() {
+ if (type == Type.NORMAL) {
+ return "projects/" + project + "/topics/" + topic;
+ } else {
+ return topic;
+ }
+ }
+
+ /**
+ * Returns the string representation of this topic as a path used in the Cloud Pub/Sub
+ * API.
+ */
+ public String asPath() {
+ if (type == Type.NORMAL) {
+ return "projects/" + project + "/topics/" + topic;
+ } else {
+ return topic;
+ }
+ }
+ }
+
+ public static <T> Read<T> read() {
+ return new Read<>();
+ }
+
+ public static <T> Write<T> write() {
+ return new Write<>();
+ }
+
+ /**
+ * A {@link PTransform} that continuously reads from a Google Cloud Pub/Sub stream and
+ * returns a {@link PCollection} of {@link String Strings} containing the items from
+ * the stream.
+ */
+ public static class Read<T> extends PTransform<PBegin, PCollection<T>> {
+
+ /** The Cloud Pub/Sub topic to read from. */
+ @Nullable
+ private final ValueProvider<PubsubTopic> topic;
+
+ /** The Cloud Pub/Sub subscription to read from. */
+ @Nullable
+ private final ValueProvider<PubsubSubscription> subscription;
+
+ /** The name of the message attribute to read timestamps from. */
+ @Nullable
+ private final String timestampLabel;
+
+ /** The name of the message attribute to read unique message IDs from. */
+ @Nullable
+ private final String idLabel;
+
+ /** The coder used to decode each record. */
+ @Nullable
+ private final Coder<T> coder;
+
+ /** User function for parsing PubsubMessage object. */
+ SimpleFunction<PubsubMessage, T> parseFn;
+
+ private Read() {
+ this(null, null, null, null, null, null, null);
+ }
+
+ private Read(String name, ValueProvider<PubsubSubscription> subscription,
+ ValueProvider<PubsubTopic> topic, String timestampLabel, Coder<T> coder,
+ String idLabel,
+ SimpleFunction<PubsubMessage, T> parseFn) {
+ super(name);
+ this.subscription = subscription;
+ this.topic = topic;
+ this.timestampLabel = timestampLabel;
+ this.coder = coder;
+ this.idLabel = idLabel;
+ this.parseFn = parseFn;
+ }
+
+ /**
+ * Returns a transform that's like this one but reading from the
+ * given subscription.
+ *
+ * <p>See {@link PubsubIO.PubsubSubscription#fromPath(String)} for more details on the format
+ * of the {@code subscription} string.
+ *
+ * <p>Multiple readers reading from the same subscription will each receive
+ * some arbitrary portion of the data. Most likely, separate readers should
+ * use their own subscriptions.
+ *
+ * <p>Does not modify this object.
+ */
+ public Read<T> subscription(String subscription) {
+ return subscription(StaticValueProvider.of(subscription));
+ }
+
+ /**
+ * Like {@code subscription()} but with a {@link ValueProvider}.
+ */
+ public Read<T> subscription(ValueProvider<String> subscription) {
+ if (subscription.isAccessible()) {
+ // Validate.
+ PubsubSubscription.fromPath(subscription.get());
+ }
+ return new Read<>(
+ name, NestedValueProvider.of(subscription, new SubscriptionTranslator()),
+ null /* reset topic to null */, timestampLabel, coder, idLabel, parseFn);
+ }
+
+ /**
+ * Creates and returns a transform for reading from a Cloud Pub/Sub topic. Mutually exclusive
+ * with {@link #subscription(String)}.
+ *
+ * <p>See {@link PubsubIO.PubsubTopic#fromPath(String)} for more details on the format
+ * of the {@code topic} string.
+ *
+ * <p>The Beam runner will start reading data published on this topic from the time the
+ * pipeline is started. Any data published on the topic before the pipeline is started will
+ * not be read by the runner.
+ */
+ public Read<T> topic(String topic) {
+ return topic(StaticValueProvider.of(topic));
+ }
+
+ /**
+ * Like {@code topic()} but with a {@link ValueProvider}.
+ */
+ public Read<T> topic(ValueProvider<String> topic) {
+ if (topic.isAccessible()) {
+ // Validate.
+ PubsubTopic.fromPath(topic.get());
+ }
+ return new Read<>(name, null /* reset subscription to null */,
+ NestedValueProvider.of(topic, new TopicTranslator()),
+ timestampLabel, coder, idLabel, parseFn);
+ }
+
+ /**
+ * Creates and returns a transform reading from Cloud Pub/Sub where record timestamps are
+ * expected to be provided as Pub/Sub message attributes. The {@code timestampLabel}
+ * parameter specifies the name of the attribute that contains the timestamp.
+ *
+ * <p>The timestamp value is expected to be represented in the attribute as either:
+ *
+ * <ul>
+ * <li>a numerical value representing the number of milliseconds since the Unix epoch. For
+ * example, if using the Joda time classes, {@link Instant#getMillis()} returns the correct
+ * value for this attribute.
+ * <li>a String in RFC 3339 format. For example, {@code 2015-10-29T23:41:41.123Z}. The
+ * sub-second component of the timestamp is optional, and digits beyond the first three
+ * (i.e., time units smaller than milliseconds) will be ignored.
+ * </ul>
+ *
+ * <p>If {@code timestampLabel} is not provided, the system will generate record timestamps
+ * the first time it sees each record. All windowing will be done relative to these
+ * timestamps.
+ *
+ * <p>By default, windows are emitted based on an estimate of when this source is likely
+ * done producing data for a given timestamp (referred to as the Watermark; see
+ * {@link AfterWatermark} for more details). Any late data will be handled by the trigger
+ * specified with the windowing strategy – by default it will be output immediately.
+ *
+ * <p>Note that the system can guarantee that no late data will ever be seen when it assigns
+ * timestamps by arrival time (i.e. {@code timestampLabel} is not provided).
+ *
+ * @see <a href="https://www.ietf.org/rfc/rfc3339.txt">RFC 3339</a>
+ */
+ public Read<T> timestampLabel(String timestampLabel) {
+ return new Read<>(
+ name, subscription, topic, timestampLabel, coder, idLabel,
+ parseFn);
+ }
+
+ /**
+ * Creates and returns a transform for reading from Cloud Pub/Sub where unique record
+ * identifiers are expected to be provided as Pub/Sub message attributes. The {@code idLabel}
+ * parameter specifies the attribute name. The value of the attribute can be any string
+ * that uniquely identifies this record.
+ *
+ * <p>Pub/Sub cannot guarantee that no duplicate data will be delivered on the Pub/Sub stream.
+ * If {@code idLabel} is not provided, Beam cannot guarantee that no duplicate data will
+ * be delivered, and deduplication of the stream will be strictly best effort.
+ */
+ public Read<T> idLabel(String idLabel) {
+ return new Read<>(
+ name, subscription, topic, timestampLabel, coder, idLabel,
+ parseFn);
+ }
+
+ /**
+ * Returns a transform that's like this one but that uses the given
+ * {@link Coder} to decode each record into a value of type {@code T}.
+ *
+ * <p>Does not modify this object.
+ */
+ public Read<T> withCoder(Coder<T> coder) {
+ return new Read<>(
+ name, subscription, topic, timestampLabel, coder, idLabel,
+ parseFn);
+ }
+
+ /**
+ * Causes the source to return a PubsubMessage that includes Pubsub attributes.
+ * The user must supply a parsing function to transform the PubsubMessage into an output type.
+ * A Coder for the output type T must be registered or set on the output via
+ * {@link PCollection#setCoder(Coder)}.
+ */
+ public Read<T> withAttributes(SimpleFunction<PubsubMessage, T> parseFn) {
+ return new Read<T>(
+ name, subscription, topic, timestampLabel, coder, idLabel,
+ parseFn);
+ }
+
+ @Override
+ public PCollection<T> expand(PBegin input) {
+ if (topic == null && subscription == null) {
+ throw new IllegalStateException("Need to set either the topic or the subscription for "
+ + "a PubsubIO.Read transform");
+ }
+ if (topic != null && subscription != null) {
+ throw new IllegalStateException("Can't set both the topic and the subscription for "
+ + "a PubsubIO.Read transform");
+ }
+ if (coder == null) {
+ throw new IllegalStateException("PubsubIO.Read requires that a coder be set using "
+ + "the withCoder method.");
+ }
+
+ @Nullable ValueProvider<ProjectPath> projectPath =
+ topic == null ? null : NestedValueProvider.of(topic, new ProjectPathTranslator());
+ @Nullable ValueProvider<TopicPath> topicPath =
+ topic == null ? null : NestedValueProvider.of(topic, new TopicPathTranslator());
+ @Nullable ValueProvider<SubscriptionPath> subscriptionPath =
+ subscription == null
+ ? null
+ : NestedValueProvider.of(subscription, new SubscriptionPathTranslator());
+ PubsubUnboundedSource<T> source = new PubsubUnboundedSource<T>(
+ FACTORY, projectPath, topicPath, subscriptionPath,
+ coder, timestampLabel, idLabel, parseFn);
+ return input.getPipeline().apply(source);
+ }
+
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ super.populateDisplayData(builder);
+ populateCommonDisplayData(builder, timestampLabel, idLabel, topic);
+
+ if (subscription != null) {
+ String subscriptionString = subscription.isAccessible()
+ ? subscription.get().asPath() : subscription.toString();
+ builder.add(DisplayData.item("subscription", subscriptionString)
+ .withLabel("Pubsub Subscription"));
+ }
+ }
+
+ @Override
+ protected Coder<T> getDefaultOutputCoder() {
+ return coder;
+ }
+
+ /**
+ * Get the topic being read from.
+ */
+ @Nullable
+ public PubsubTopic getTopic() {
+ return topic == null ? null : topic.get();
+ }
+
+ /**
+ * Get the {@link ValueProvider} for the topic being read from.
+ */
+ public ValueProvider<PubsubTopic> getTopicProvider() {
+ return topic;
+ }
+
+ /**
+ * Get the subscription being read from.
+ */
+ @Nullable
+ public PubsubSubscription getSubscription() {
+ return subscription == null ? null : subscription.get();
+ }
+
+ /**
+ * Get the {@link ValueProvider} for the subscription being read from.
+ */
+ public ValueProvider<PubsubSubscription> getSubscriptionProvider() {
+ return subscription;
+ }
+
+ /**
+ * Get the timestamp label.
+ */
+ @Nullable
+ public String getTimestampLabel() {
+ return timestampLabel;
+ }
+
+ /**
+ * Get the id label.
+ */
+ @Nullable
+ public String getIdLabel() {
+ return idLabel;
+ }
+
+
+ /**
+ * Get the {@link Coder} used for the transform's output.
+ */
+ @Nullable
+ public Coder<T> getCoder() {
+ return coder;
+ }
+
+ /**
+ * Get the parse function used for PubSub attributes.
+ */
+ @Nullable
+ public SimpleFunction<PubsubMessage, T> getPubSubMessageParseFn() {
+ return parseFn;
+ }
+
+ }
+
+ /////////////////////////////////////////////////////////////////////////////
+
+ /** Disallow construction of utility class. */
+ private PubsubIO() {}
+
+
+ /**
+ * A {@link PTransform} that writes an unbounded {@link PCollection} of {@link String Strings}
+ * to a Cloud Pub/Sub stream.
+ */
+ public static class Write<T> extends PTransform<PCollection<T>, PDone> {
+
+ /** The Cloud Pub/Sub topic to publish to. */
+ @Nullable
+ private final ValueProvider<PubsubTopic> topic;
+ /** The name of the message attribute to publish message timestamps in. */
+ @Nullable
+ private final String timestampLabel;
+ /** The name of the message attribute to publish unique message IDs in. */
+ @Nullable
+ private final String idLabel;
+ /** The input type Coder. */
+ private final Coder<T> coder;
+ /** The format function for input PubsubMessage objects. */
+ SimpleFunction<T, PubsubMessage> formatFn;
+
+ private Write() {
+ this(null, null, null, null, null, null);
+ }
+
+ private Write(
+ String name, ValueProvider<PubsubTopic> topic, String timestampLabel,
+ String idLabel, Coder<T> coder, SimpleFunction<T, PubsubMessage> formatFn) {
+ super(name);
+ this.topic = topic;
+ this.timestampLabel = timestampLabel;
+ this.idLabel = idLabel;
+ this.coder = coder;
+ this.formatFn = formatFn;
+ }
+
+ /**
+ * Creates a transform that publishes to the specified topic.
+ *
+ * <p>See {@link PubsubIO.PubsubTopic#fromPath(String)} for more details on the format of the
+ * {@code topic} string.
+ */
+ public Write<T> topic(String topic) {
+ return topic(StaticValueProvider.of(topic));
+ }
+
+ /**
+ * Like {@code topic()} but with a {@link ValueProvider}.
+ */
+ public Write<T> topic(ValueProvider<String> topic) {
+ return new Write<>(name, NestedValueProvider.of(topic, new TopicTranslator()),
+ timestampLabel, idLabel, coder, formatFn);
+ }
+
+ /**
+ * Creates a transform that writes to Pub/Sub, adds each record's timestamp to the published
+ * messages in an attribute with the specified name. The value of the attribute will be a number
+ * representing the number of milliseconds since the Unix epoch. For example, if using the Joda
+ * time classes, {@link Instant#Instant(long)} can be used to parse this value.
+ *
+ * <p>If the output from this sink is being read by another Beam pipeline, then
+ * {@link PubsubIO.Read#timestampLabel(String)} can be used to ensure the other source reads
+ * these timestamps from the appropriate attribute.
+ */
+ public Write<T> timestampLabel(String timestampLabel) {
+ return new Write<>(name, topic, timestampLabel, idLabel, coder, formatFn);
+ }
+
+ /**
+ * Creates a transform that writes to Pub/Sub, adding each record's unique identifier to the
+ * published messages in an attribute with the specified name. The value of the attribute is an
+ * opaque string.
+ *
+ * <p>If the the output from this sink is being read by another Beam pipeline, then
+ * {@link PubsubIO.Read#idLabel(String)} can be used to ensure that* the other source reads
+ * these unique identifiers from the appropriate attribute.
+ */
+ public Write<T> idLabel(String idLabel) {
+ return new Write<>(name, topic, timestampLabel, idLabel, coder, formatFn);
+ }
+
+ /**
+ * Returns a new transform that's like this one
+ * but that uses the given {@link Coder} to encode each of
+ * the elements of the input {@link PCollection} into an
+ * output record.
+ *
+ * <p>Does not modify this object.
+ */
+ public Write<T> withCoder(Coder<T> coder) {
+ return new Write<>(name, topic, timestampLabel, idLabel, coder, formatFn);
+ }
+
+ /**
+ * Used to write a PubSub message together with PubSub attributes. The user-supplied format
+ * function translates the input type T to a PubsubMessage object, which is used by the sink
+ * to separately set the PubSub message's payload and attributes.
+ */
+ public Write<T> withAttributes(SimpleFunction<T, PubsubMessage> formatFn) {
+ return new Write<T>(name, topic, timestampLabel, idLabel, coder, formatFn);
+ }
+
+ @Override
+ public PDone expand(PCollection<T> input) {
+ if (topic == null) {
+ throw new IllegalStateException("need to set the topic of a PubsubIO.Write transform");
+ }
+ switch (input.isBounded()) {
+ case BOUNDED:
+ input.apply(ParDo.of(new PubsubBoundedWriter()));
+ return PDone.in(input.getPipeline());
+ case UNBOUNDED:
+ return input.apply(new PubsubUnboundedSink<T>(
+ FACTORY,
+ NestedValueProvider.of(topic, new TopicPathTranslator()),
+ coder,
+ timestampLabel,
+ idLabel,
+ formatFn,
+ 100 /* numShards */));
+ }
+ throw new RuntimeException(); // cases are exhaustive.
+ }
+
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ super.populateDisplayData(builder);
+ populateCommonDisplayData(builder, timestampLabel, idLabel, topic);
+ }
+
+ @Override
+ protected Coder<Void> getDefaultOutputCoder() {
+ return VoidCoder.of();
+ }
+
+ /**
+ * Returns the PubSub topic being written to.
+ */
+ @Nullable
+ public PubsubTopic getTopic() {
+ return (topic == null) ? null : topic.get();
+ }
+
+ /**
+ * Returns the {@link ValueProvider} for the topic being written to.
+ */
+ @Nullable
+ public ValueProvider<PubsubTopic> getTopicProvider() {
+ return topic;
+ }
+
+ /**
+ * Returns the timestamp label.
+ */
+ @Nullable
+ public String getTimestampLabel() {
+ return timestampLabel;
+ }
+
+ /**
+ * Returns the id label.
+ */
+ @Nullable
+ public String getIdLabel() {
+ return idLabel;
+ }
+
+ /**
+ * Returns the output coder.
+ */
+ @Nullable
+ public Coder<T> getCoder() {
+ return coder;
+ }
+
+ /**
+ * Returns the formatting function used if publishing attributes.
+ */
+ @Nullable
+ public SimpleFunction<T, PubsubMessage> getFormatFn() {
+ return formatFn;
+ }
+
+ /**
+ * Writer to Pubsub which batches messages from bounded collections.
+ *
+ * <p>Public so can be suppressed by runners.
+ */
+ public class PubsubBoundedWriter extends DoFn<T, Void> {
+
+ private static final int MAX_PUBLISH_BATCH_SIZE = 100;
+ private transient List<OutgoingMessage> output;
+ private transient PubsubClient pubsubClient;
+
+ @StartBundle
+ public void startBundle(Context c) throws IOException {
+ this.output = new ArrayList<>();
+ // NOTE: idLabel is ignored.
+ this.pubsubClient =
+ FACTORY.newClient(timestampLabel, null,
+ c.getPipelineOptions().as(PubsubOptions.class));
+ }
+
+ @ProcessElement
+ public void processElement(ProcessContext c) throws IOException {
+ byte[] payload = null;
+ Map<String, String> attributes = null;
+ if (formatFn != null) {
+ PubsubMessage message = formatFn.apply(c.element());
+ payload = message.getMessage();
+ attributes = message.getAttributeMap();
+ } else {
+ payload = CoderUtils.encodeToByteArray(getCoder(), c.element());
+ }
+ // NOTE: The record id is always null.
+ OutgoingMessage message =
+ new OutgoingMessage(payload, attributes, c.timestamp().getMillis(), null);
+ output.add(message);
+
+ if (output.size() >= MAX_PUBLISH_BATCH_SIZE) {
+ publish();
+ }
+ }
+
+ @FinishBundle
+ public void finishBundle(Context c) throws IOException {
+ if (!output.isEmpty()) {
+ publish();
+ }
+ output = null;
+ pubsubClient.close();
+ pubsubClient = null;
+ }
+
+ private void publish() throws IOException {
+ int n = pubsubClient.publish(
+ PubsubClient.topicPathFromName(getTopic().project, getTopic().topic),
+ output);
+ checkState(n == output.size());
+ output.clear();
+ }
+
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ super.populateDisplayData(builder);
+ builder.delegate(Write.this);
+ }
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubJsonClient.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubJsonClient.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubJsonClient.java
new file mode 100644
index 0000000..e290a6b
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubJsonClient.java
@@ -0,0 +1,319 @@
+/*
+ * 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.pubsub;
+
+import static com.google.common.base.Preconditions.checkState;
+
+import com.google.api.client.http.HttpRequestInitializer;
+import com.google.api.services.pubsub.Pubsub;
+import com.google.api.services.pubsub.Pubsub.Builder;
+import com.google.api.services.pubsub.model.AcknowledgeRequest;
+import com.google.api.services.pubsub.model.ListSubscriptionsResponse;
+import com.google.api.services.pubsub.model.ListTopicsResponse;
+import com.google.api.services.pubsub.model.ModifyAckDeadlineRequest;
+import com.google.api.services.pubsub.model.PublishRequest;
+import com.google.api.services.pubsub.model.PublishResponse;
+import com.google.api.services.pubsub.model.PubsubMessage;
+import com.google.api.services.pubsub.model.PullRequest;
+import com.google.api.services.pubsub.model.PullResponse;
+import com.google.api.services.pubsub.model.ReceivedMessage;
+import com.google.api.services.pubsub.model.Subscription;
+import com.google.api.services.pubsub.model.Topic;
+import com.google.auth.Credentials;
+import com.google.auth.http.HttpCredentialsAdapter;
+import com.google.cloud.hadoop.util.ChainingHttpRequestInitializer;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Strings;
+import com.google.common.collect.ImmutableList;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.options.PubsubOptions;
+import org.apache.beam.sdk.util.RetryHttpRequestInitializer;
+import org.apache.beam.sdk.util.Transport;
+
+/**
+ * A Pubsub client using JSON transport.
+ */
+class PubsubJsonClient extends PubsubClient {
+
+ private static class PubsubJsonClientFactory implements PubsubClientFactory {
+ private static HttpRequestInitializer chainHttpRequestInitializer(
+ Credentials credential, HttpRequestInitializer httpRequestInitializer) {
+ if (credential == null) {
+ return httpRequestInitializer;
+ } else {
+ return new ChainingHttpRequestInitializer(
+ new HttpCredentialsAdapter(credential),
+ httpRequestInitializer);
+ }
+ }
+
+ @Override
+ public PubsubClient newClient(
+ @Nullable String timestampLabel, @Nullable String idLabel, PubsubOptions options)
+ throws IOException {
+ Pubsub pubsub = new Builder(
+ Transport.getTransport(),
+ Transport.getJsonFactory(),
+ chainHttpRequestInitializer(
+ options.getGcpCredential(),
+ // Do not log 404. It clutters the output and is possibly even required by the caller.
+ new RetryHttpRequestInitializer(ImmutableList.of(404))))
+ .setRootUrl(options.getPubsubRootUrl())
+ .setApplicationName(options.getAppName())
+ .setGoogleClientRequestInitializer(options.getGoogleApiTrace())
+ .build();
+ return new PubsubJsonClient(timestampLabel, idLabel, pubsub);
+ }
+
+ @Override
+ public String getKind() {
+ return "Json";
+ }
+ }
+
+ /**
+ * Factory for creating Pubsub clients using Json transport.
+ */
+ public static final PubsubClientFactory FACTORY = new PubsubJsonClientFactory();
+
+ /**
+ * Label to use for custom timestamps, or {@literal null} if should use Pubsub publish time
+ * instead.
+ */
+ @Nullable
+ private final String timestampLabel;
+
+ /**
+ * Label to use for custom ids, or {@literal null} if should use Pubsub provided ids.
+ */
+ @Nullable
+ private final String idLabel;
+
+ /**
+ * Underlying JSON transport.
+ */
+ private Pubsub pubsub;
+
+ @VisibleForTesting
+ PubsubJsonClient(
+ @Nullable String timestampLabel,
+ @Nullable String idLabel,
+ Pubsub pubsub) {
+ this.timestampLabel = timestampLabel;
+ this.idLabel = idLabel;
+ this.pubsub = pubsub;
+ }
+
+ @Override
+ public void close() {
+ // Nothing to close.
+ }
+
+ @Override
+ public int publish(TopicPath topic, List<OutgoingMessage> outgoingMessages)
+ throws IOException {
+ List<PubsubMessage> pubsubMessages = new ArrayList<>(outgoingMessages.size());
+ for (OutgoingMessage outgoingMessage : outgoingMessages) {
+ PubsubMessage pubsubMessage = new PubsubMessage().encodeData(outgoingMessage.elementBytes);
+
+ Map<String, String> attributes = outgoingMessage.attributes;
+ if ((timestampLabel != null || idLabel != null) && attributes == null) {
+ attributes = new TreeMap<>();
+ }
+ if (attributes != null) {
+ pubsubMessage.setAttributes(attributes);
+ }
+
+ if (timestampLabel != null) {
+ attributes.put(timestampLabel, String.valueOf(outgoingMessage.timestampMsSinceEpoch));
+ }
+
+ if (idLabel != null && !Strings.isNullOrEmpty(outgoingMessage.recordId)) {
+ attributes.put(idLabel, outgoingMessage.recordId);
+ }
+
+ pubsubMessages.add(pubsubMessage);
+ }
+ PublishRequest request = new PublishRequest().setMessages(pubsubMessages);
+ PublishResponse response = pubsub.projects()
+ .topics()
+ .publish(topic.getPath(), request)
+ .execute();
+ return response.getMessageIds().size();
+ }
+
+ @Override
+ public List<IncomingMessage> pull(
+ long requestTimeMsSinceEpoch,
+ SubscriptionPath subscription,
+ int batchSize,
+ boolean returnImmediately) throws IOException {
+ PullRequest request = new PullRequest()
+ .setReturnImmediately(returnImmediately)
+ .setMaxMessages(batchSize);
+ PullResponse response = pubsub.projects()
+ .subscriptions()
+ .pull(subscription.getPath(), request)
+ .execute();
+ if (response.getReceivedMessages() == null || response.getReceivedMessages().size() == 0) {
+ return ImmutableList.of();
+ }
+ List<IncomingMessage> incomingMessages = new ArrayList<>(response.getReceivedMessages().size());
+ for (ReceivedMessage message : response.getReceivedMessages()) {
+ PubsubMessage pubsubMessage = message.getMessage();
+ @Nullable Map<String, String> attributes = pubsubMessage.getAttributes();
+
+ // Payload.
+ byte[] elementBytes = pubsubMessage.decodeData();
+
+ // Timestamp.
+ long timestampMsSinceEpoch =
+ extractTimestamp(timestampLabel, message.getMessage().getPublishTime(), attributes);
+
+ // Ack id.
+ String ackId = message.getAckId();
+ checkState(!Strings.isNullOrEmpty(ackId));
+
+ // Record id, if any.
+ @Nullable String recordId = null;
+ if (idLabel != null && attributes != null) {
+ recordId = attributes.get(idLabel);
+ }
+ if (Strings.isNullOrEmpty(recordId)) {
+ // Fall back to the Pubsub provided message id.
+ recordId = pubsubMessage.getMessageId();
+ }
+
+ incomingMessages.add(new IncomingMessage(elementBytes, attributes, timestampMsSinceEpoch,
+ requestTimeMsSinceEpoch, ackId, recordId));
+ }
+
+ return incomingMessages;
+ }
+
+ @Override
+ public void acknowledge(SubscriptionPath subscription, List<String> ackIds) throws IOException {
+ AcknowledgeRequest request = new AcknowledgeRequest().setAckIds(ackIds);
+ pubsub.projects()
+ .subscriptions()
+ .acknowledge(subscription.getPath(), request)
+ .execute(); // ignore Empty result.
+ }
+
+ @Override
+ public void modifyAckDeadline(
+ SubscriptionPath subscription, List<String> ackIds, int deadlineSeconds)
+ throws IOException {
+ ModifyAckDeadlineRequest request =
+ new ModifyAckDeadlineRequest().setAckIds(ackIds)
+ .setAckDeadlineSeconds(deadlineSeconds);
+ pubsub.projects()
+ .subscriptions()
+ .modifyAckDeadline(subscription.getPath(), request)
+ .execute(); // ignore Empty result.
+ }
+
+ @Override
+ public void createTopic(TopicPath topic) throws IOException {
+ pubsub.projects()
+ .topics()
+ .create(topic.getPath(), new Topic())
+ .execute(); // ignore Topic result.
+ }
+
+ @Override
+ public void deleteTopic(TopicPath topic) throws IOException {
+ pubsub.projects()
+ .topics()
+ .delete(topic.getPath())
+ .execute(); // ignore Empty result.
+ }
+
+ @Override
+ public List<TopicPath> listTopics(ProjectPath project) throws IOException {
+ ListTopicsResponse response = pubsub.projects()
+ .topics()
+ .list(project.getPath())
+ .execute();
+ if (response.getTopics() == null || response.getTopics().isEmpty()) {
+ return ImmutableList.of();
+ }
+ List<TopicPath> topics = new ArrayList<>(response.getTopics().size());
+ for (Topic topic : response.getTopics()) {
+ topics.add(topicPathFromPath(topic.getName()));
+ }
+ return topics;
+ }
+
+ @Override
+ public void createSubscription(
+ TopicPath topic, SubscriptionPath subscription,
+ int ackDeadlineSeconds) throws IOException {
+ Subscription request = new Subscription()
+ .setTopic(topic.getPath())
+ .setAckDeadlineSeconds(ackDeadlineSeconds);
+ pubsub.projects()
+ .subscriptions()
+ .create(subscription.getPath(), request)
+ .execute(); // ignore Subscription result.
+ }
+
+ @Override
+ public void deleteSubscription(SubscriptionPath subscription) throws IOException {
+ pubsub.projects()
+ .subscriptions()
+ .delete(subscription.getPath())
+ .execute(); // ignore Empty result.
+ }
+
+ @Override
+ public List<SubscriptionPath> listSubscriptions(ProjectPath project, TopicPath topic)
+ throws IOException {
+ ListSubscriptionsResponse response = pubsub.projects()
+ .subscriptions()
+ .list(project.getPath())
+ .execute();
+ if (response.getSubscriptions() == null || response.getSubscriptions().isEmpty()) {
+ return ImmutableList.of();
+ }
+ List<SubscriptionPath> subscriptions = new ArrayList<>(response.getSubscriptions().size());
+ for (Subscription subscription : response.getSubscriptions()) {
+ if (subscription.getTopic().equals(topic.getPath())) {
+ subscriptions.add(subscriptionPathFromPath(subscription.getName()));
+ }
+ }
+ return subscriptions;
+ }
+
+ @Override
+ public int ackDeadlineSeconds(SubscriptionPath subscription) throws IOException {
+ Subscription response = pubsub.projects().subscriptions().get(subscription.getPath()).execute();
+ return response.getAckDeadlineSeconds();
+ }
+
+ @Override
+ public boolean isEOF() {
+ return false;
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubTestClient.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubTestClient.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubTestClient.java
new file mode 100644
index 0000000..c88576e
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubTestClient.java
@@ -0,0 +1,436 @@
+/*
+ * 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.pubsub;
+
+import static com.google.common.base.Preconditions.checkState;
+
+import com.google.api.client.util.Clock;
+import com.google.common.collect.Lists;
+import com.google.common.collect.Sets;
+import java.io.Closeable;
+import java.io.IOException;
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.options.PubsubOptions;
+
+/**
+ * A (partial) implementation of {@link PubsubClient} for use by unit tests. Only suitable for
+ * testing {@link #publish}, {@link #pull}, {@link #acknowledge} and {@link #modifyAckDeadline}
+ * methods. Relies on statics to mimic the Pubsub service, though we try to hide that.
+ */
+class PubsubTestClient extends PubsubClient implements Serializable {
+ /**
+ * Mimic the state of the simulated Pubsub 'service'.
+ *
+ * <p>Note that the {@link PubsubTestClientFactory} is serialized/deserialized even when running
+ * test pipelines. Meanwhile it is valid for multiple {@link PubsubTestClient}s to be created
+ * from the same client factory and run in parallel. Thus we can't enforce aliasing of the
+ * following data structures over all clients and must resort to a static.
+ */
+ private static class State {
+ /**
+ * True if has been primed for a test but not yet validated.
+ */
+ boolean isActive;
+
+ /**
+ * Publish mode only: Only publish calls for this topic are allowed.
+ */
+ @Nullable
+ TopicPath expectedTopic;
+
+ /**
+ * Publish mode only: Messages yet to seen in a {@link #publish} call.
+ */
+ @Nullable
+ Set<OutgoingMessage> remainingExpectedOutgoingMessages;
+
+ /**
+ * Publish mode only: Messages which should throw when first sent to simulate transient publish
+ * failure.
+ */
+ @Nullable
+ Set<OutgoingMessage> remainingFailingOutgoingMessages;
+
+ /**
+ * Pull mode only: Clock from which to get current time.
+ */
+ @Nullable
+ Clock clock;
+
+ /**
+ * Pull mode only: Only pull calls for this subscription are allowed.
+ */
+ @Nullable
+ SubscriptionPath expectedSubscription;
+
+ /**
+ * Pull mode only: Timeout to simulate.
+ */
+ int ackTimeoutSec;
+
+ /**
+ * Pull mode only: Messages waiting to be received by a {@link #pull} call.
+ */
+ @Nullable
+ List<IncomingMessage> remainingPendingIncomingMessages;
+
+ /**
+ * Pull mode only: Messages which have been returned from a {@link #pull} call and
+ * not yet ACKed by an {@link #acknowledge} call.
+ */
+ @Nullable
+ Map<String, IncomingMessage> pendingAckIncomingMessages;
+
+ /**
+ * Pull mode only: When above messages are due to have their ACK deadlines expire.
+ */
+ @Nullable
+ Map<String, Long> ackDeadline;
+ }
+
+ private static final State STATE = new State();
+
+ /** Closing the factory will validate all expected messages were processed. */
+ public interface PubsubTestClientFactory
+ extends PubsubClientFactory, Closeable, Serializable {
+ }
+
+ /**
+ * Return a factory for testing publishers. Only one factory may be in-flight at a time.
+ * The factory must be closed when the test is complete, at which point final validation will
+ * occur.
+ */
+ static PubsubTestClientFactory createFactoryForPublish(
+ final TopicPath expectedTopic,
+ final Iterable<OutgoingMessage> expectedOutgoingMessages,
+ final Iterable<OutgoingMessage> failingOutgoingMessages) {
+ synchronized (STATE) {
+ checkState(!STATE.isActive, "Test still in flight");
+ STATE.expectedTopic = expectedTopic;
+ STATE.remainingExpectedOutgoingMessages = Sets.newHashSet(expectedOutgoingMessages);
+ STATE.remainingFailingOutgoingMessages = Sets.newHashSet(failingOutgoingMessages);
+ STATE.isActive = true;
+ }
+ return new PubsubTestClientFactory() {
+ @Override
+ public PubsubClient newClient(
+ @Nullable String timestampLabel, @Nullable String idLabel, PubsubOptions options)
+ throws IOException {
+ return new PubsubTestClient();
+ }
+
+ @Override
+ public String getKind() {
+ return "PublishTest";
+ }
+
+ @Override
+ public void close() {
+ synchronized (STATE) {
+ checkState(STATE.isActive, "No test still in flight");
+ checkState(STATE.remainingExpectedOutgoingMessages.isEmpty(),
+ "Still waiting for %s messages to be published",
+ STATE.remainingExpectedOutgoingMessages.size());
+ STATE.isActive = false;
+ STATE.remainingExpectedOutgoingMessages = null;
+ }
+ }
+ };
+ }
+
+ /**
+ * Return a factory for testing subscribers. Only one factory may be in-flight at a time.
+ * The factory must be closed when the test in complete
+ */
+ public static PubsubTestClientFactory createFactoryForPull(
+ final Clock clock,
+ final SubscriptionPath expectedSubscription,
+ final int ackTimeoutSec,
+ final Iterable<IncomingMessage> expectedIncomingMessages) {
+ synchronized (STATE) {
+ checkState(!STATE.isActive, "Test still in flight");
+ STATE.clock = clock;
+ STATE.expectedSubscription = expectedSubscription;
+ STATE.ackTimeoutSec = ackTimeoutSec;
+ STATE.remainingPendingIncomingMessages = Lists.newArrayList(expectedIncomingMessages);
+ STATE.pendingAckIncomingMessages = new HashMap<>();
+ STATE.ackDeadline = new HashMap<>();
+ STATE.isActive = true;
+ }
+ return new PubsubTestClientFactory() {
+ @Override
+ public PubsubClient newClient(
+ @Nullable String timestampLabel, @Nullable String idLabel, PubsubOptions options)
+ throws IOException {
+ return new PubsubTestClient();
+ }
+
+ @Override
+ public String getKind() {
+ return "PullTest";
+ }
+
+ @Override
+ public void close() {
+ synchronized (STATE) {
+ checkState(STATE.isActive, "No test still in flight");
+ checkState(STATE.remainingPendingIncomingMessages.isEmpty(),
+ "Still waiting for %s messages to be pulled",
+ STATE.remainingPendingIncomingMessages.size());
+ checkState(STATE.pendingAckIncomingMessages.isEmpty(),
+ "Still waiting for %s messages to be ACKed",
+ STATE.pendingAckIncomingMessages.size());
+ checkState(STATE.ackDeadline.isEmpty(),
+ "Still waiting for %s messages to be ACKed",
+ STATE.ackDeadline.size());
+ STATE.isActive = false;
+ STATE.remainingPendingIncomingMessages = null;
+ STATE.pendingAckIncomingMessages = null;
+ STATE.ackDeadline = null;
+ }
+ }
+ };
+ }
+
+ public static PubsubTestClientFactory createFactoryForCreateSubscription() {
+ return new PubsubTestClientFactory() {
+ int numCalls = 0;
+
+ @Override
+ public void close() throws IOException {
+ checkState(
+ numCalls == 1, "Expected exactly one subscription to be created, got %s", numCalls);
+ }
+
+ @Override
+ public PubsubClient newClient(
+ @Nullable String timestampLabel, @Nullable String idLabel, PubsubOptions options)
+ throws IOException {
+ return new PubsubTestClient() {
+ @Override
+ public void createSubscription(
+ TopicPath topic, SubscriptionPath subscription, int ackDeadlineSeconds)
+ throws IOException {
+ checkState(numCalls == 0, "Expected at most one subscription to be created");
+ numCalls++;
+ }
+ };
+ }
+
+ @Override
+ public String getKind() {
+ return "CreateSubscriptionTest";
+ }
+ };
+ }
+
+ /**
+ * Return true if in pull mode.
+ */
+ private boolean inPullMode() {
+ checkState(STATE.isActive, "No test is active");
+ return STATE.expectedSubscription != null;
+ }
+
+ /**
+ * Return true if in publish mode.
+ */
+ private boolean inPublishMode() {
+ checkState(STATE.isActive, "No test is active");
+ return STATE.expectedTopic != null;
+ }
+
+ /**
+ * For subscription mode only:
+ * Track progression of time according to the {@link Clock} passed . This will simulate Pubsub
+ * expiring
+ * outstanding ACKs.
+ */
+ public void advance() {
+ synchronized (STATE) {
+ checkState(inPullMode(), "Can only advance in pull mode");
+ // Any messages who's ACKs timed out are available for re-pulling.
+ Iterator<Map.Entry<String, Long>> deadlineItr = STATE.ackDeadline.entrySet().iterator();
+ while (deadlineItr.hasNext()) {
+ Map.Entry<String, Long> entry = deadlineItr.next();
+ if (entry.getValue() <= STATE.clock.currentTimeMillis()) {
+ STATE.remainingPendingIncomingMessages.add(
+ STATE.pendingAckIncomingMessages.remove(entry.getKey()));
+ deadlineItr.remove();
+ }
+ }
+ }
+ }
+
+ @Override
+ public void close() {
+ }
+
+ @Override
+ public int publish(
+ TopicPath topic, List<OutgoingMessage> outgoingMessages) throws IOException {
+ synchronized (STATE) {
+ checkState(inPublishMode(), "Can only publish in publish mode");
+ checkState(topic.equals(STATE.expectedTopic), "Topic %s does not match expected %s", topic,
+ STATE.expectedTopic);
+ for (OutgoingMessage outgoingMessage : outgoingMessages) {
+ if (STATE.remainingFailingOutgoingMessages.remove(outgoingMessage)) {
+ throw new RuntimeException("Simulating failure for " + outgoingMessage);
+ }
+ checkState(STATE.remainingExpectedOutgoingMessages.remove(outgoingMessage),
+ "Unexpected outgoing message %s", outgoingMessage);
+ }
+ return outgoingMessages.size();
+ }
+ }
+
+ @Override
+ public List<IncomingMessage> pull(
+ long requestTimeMsSinceEpoch, SubscriptionPath subscription, int batchSize,
+ boolean returnImmediately) throws IOException {
+ synchronized (STATE) {
+ checkState(inPullMode(), "Can only pull in pull mode");
+ long now = STATE.clock.currentTimeMillis();
+ checkState(requestTimeMsSinceEpoch == now,
+ "Simulated time %s does not match request time %s", now, requestTimeMsSinceEpoch);
+ checkState(subscription.equals(STATE.expectedSubscription),
+ "Subscription %s does not match expected %s", subscription,
+ STATE.expectedSubscription);
+ checkState(returnImmediately, "Pull only supported if returning immediately");
+
+ List<IncomingMessage> incomingMessages = new ArrayList<>();
+ Iterator<IncomingMessage> pendItr = STATE.remainingPendingIncomingMessages.iterator();
+ while (pendItr.hasNext()) {
+ IncomingMessage incomingMessage = pendItr.next();
+ pendItr.remove();
+ IncomingMessage incomingMessageWithRequestTime =
+ incomingMessage.withRequestTime(requestTimeMsSinceEpoch);
+ incomingMessages.add(incomingMessageWithRequestTime);
+ STATE.pendingAckIncomingMessages.put(incomingMessageWithRequestTime.ackId,
+ incomingMessageWithRequestTime);
+ STATE.ackDeadline.put(incomingMessageWithRequestTime.ackId,
+ requestTimeMsSinceEpoch + STATE.ackTimeoutSec * 1000);
+ if (incomingMessages.size() >= batchSize) {
+ break;
+ }
+ }
+ return incomingMessages;
+ }
+ }
+
+ @Override
+ public void acknowledge(
+ SubscriptionPath subscription,
+ List<String> ackIds) throws IOException {
+ synchronized (STATE) {
+ checkState(inPullMode(), "Can only acknowledge in pull mode");
+ checkState(subscription.equals(STATE.expectedSubscription),
+ "Subscription %s does not match expected %s", subscription,
+ STATE.expectedSubscription);
+
+ for (String ackId : ackIds) {
+ checkState(STATE.ackDeadline.remove(ackId) != null,
+ "No message with ACK id %s is waiting for an ACK", ackId);
+ checkState(STATE.pendingAckIncomingMessages.remove(ackId) != null,
+ "No message with ACK id %s is waiting for an ACK", ackId);
+ }
+ }
+ }
+
+ @Override
+ public void modifyAckDeadline(
+ SubscriptionPath subscription, List<String> ackIds, int deadlineSeconds) throws IOException {
+ synchronized (STATE) {
+ checkState(inPullMode(), "Can only modify ack deadline in pull mode");
+ checkState(subscription.equals(STATE.expectedSubscription),
+ "Subscription %s does not match expected %s", subscription,
+ STATE.expectedSubscription);
+
+ for (String ackId : ackIds) {
+ if (deadlineSeconds > 0) {
+ checkState(STATE.ackDeadline.remove(ackId) != null,
+ "No message with ACK id %s is waiting for an ACK", ackId);
+ checkState(STATE.pendingAckIncomingMessages.containsKey(ackId),
+ "No message with ACK id %s is waiting for an ACK", ackId);
+ STATE.ackDeadline.put(ackId, STATE.clock.currentTimeMillis() + deadlineSeconds * 1000);
+ } else {
+ checkState(STATE.ackDeadline.remove(ackId) != null,
+ "No message with ACK id %s is waiting for an ACK", ackId);
+ IncomingMessage message = STATE.pendingAckIncomingMessages.remove(ackId);
+ checkState(message != null, "No message with ACK id %s is waiting for an ACK", ackId);
+ STATE.remainingPendingIncomingMessages.add(message);
+ }
+ }
+ }
+ }
+
+ @Override
+ public void createTopic(TopicPath topic) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public void deleteTopic(TopicPath topic) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public List<TopicPath> listTopics(ProjectPath project) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public void createSubscription(
+ TopicPath topic, SubscriptionPath subscription, int ackDeadlineSeconds) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public void deleteSubscription(SubscriptionPath subscription) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public List<SubscriptionPath> listSubscriptions(
+ ProjectPath project, TopicPath topic) throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public int ackDeadlineSeconds(SubscriptionPath subscription) throws IOException {
+ synchronized (STATE) {
+ return STATE.ackTimeoutSec;
+ }
+ }
+
+ @Override
+ public boolean isEOF() {
+ synchronized (STATE) {
+ checkState(inPullMode(), "Can only check EOF in pull mode");
+ return STATE.remainingPendingIncomingMessages.isEmpty();
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSink.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSink.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSink.java
new file mode 100644
index 0000000..3ce9224
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSink.java
@@ -0,0 +1,490 @@
+/*
+ * 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.pubsub;
+
+import static com.google.common.base.Preconditions.checkState;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.hash.Hashing;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.ThreadLocalRandom;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.AtomicCoder;
+import org.apache.beam.sdk.coders.BigEndianLongCoder;
+import org.apache.beam.sdk.coders.ByteArrayCoder;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderException;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.coders.MapCoder;
+import org.apache.beam.sdk.coders.NullableCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.OutgoingMessage;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.PubsubClientFactory;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath;
+import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.metrics.Metrics;
+import org.apache.beam.sdk.options.PubsubOptions;
+import org.apache.beam.sdk.options.ValueProvider;
+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.SimpleFunction;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.display.DisplayData.Builder;
+import org.apache.beam.sdk.transforms.windowing.AfterFirst;
+import org.apache.beam.sdk.transforms.windowing.AfterPane;
+import org.apache.beam.sdk.transforms.windowing.AfterProcessingTime;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
+import org.apache.beam.sdk.transforms.windowing.Repeatedly;
+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.PDone;
+import org.joda.time.Duration;
+
+/**
+ * A PTransform which streams messages to Pubsub.
+ * <ul>
+ * <li>The underlying implementation is just a {@link GroupByKey} followed by a {@link ParDo} which
+ * publishes as a side effect. (In the future we want to design and switch to a custom
+ * {@code UnboundedSink} implementation so as to gain access to system watermark and
+ * end-of-pipeline cleanup.)
+ * <li>We try to send messages in batches while also limiting send latency.
+ * <li>No stats are logged. Rather some counters are used to keep track of elements and batches.
+ * <li>Though some background threads are used by the underlying netty system all actual Pubsub
+ * calls are blocking. We rely on the underlying runner to allow multiple {@link DoFn} instances
+ * to execute concurrently and hide latency.
+ * <li>A failed bundle will cause messages to be resent. Thus we rely on the Pubsub consumer
+ * to dedup messages.
+ * </ul>
+ */
+public class PubsubUnboundedSink<T> extends PTransform<PCollection<T>, PDone> {
+ /**
+ * Default maximum number of messages per publish.
+ */
+ private static final int DEFAULT_PUBLISH_BATCH_SIZE = 1000;
+
+ /**
+ * Default maximum size of a publish batch, in bytes.
+ */
+ private static final int DEFAULT_PUBLISH_BATCH_BYTES = 400000;
+
+ /**
+ * Default longest delay between receiving a message and pushing it to Pubsub.
+ */
+ private static final Duration DEFAULT_MAX_LATENCY = Duration.standardSeconds(2);
+
+ /**
+ * Coder for conveying outgoing messages between internal stages.
+ */
+ private static class OutgoingMessageCoder extends AtomicCoder<OutgoingMessage> {
+ private static final NullableCoder<String> RECORD_ID_CODER =
+ NullableCoder.of(StringUtf8Coder.of());
+ private static final NullableCoder<Map<String, String>> ATTRIBUTES_CODER =
+ NullableCoder.of(MapCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()));
+
+ @Override
+ public void encode(
+ OutgoingMessage value, OutputStream outStream, Context context)
+ throws CoderException, IOException {
+ ByteArrayCoder.of().encode(value.elementBytes, outStream, context.nested());
+ ATTRIBUTES_CODER.encode(value.attributes, outStream, context.nested());
+ BigEndianLongCoder.of().encode(value.timestampMsSinceEpoch, outStream, context.nested());
+ RECORD_ID_CODER.encode(value.recordId, outStream, context.nested());
+ }
+
+ @Override
+ public OutgoingMessage decode(
+ InputStream inStream, Context context) throws CoderException, IOException {
+ byte[] elementBytes = ByteArrayCoder.of().decode(inStream, context.nested());
+ Map<String, String> attributes = ATTRIBUTES_CODER.decode(inStream, context.nested());
+ long timestampMsSinceEpoch = BigEndianLongCoder.of().decode(inStream, context.nested());
+ @Nullable String recordId = RECORD_ID_CODER.decode(inStream, context.nested());
+ return new OutgoingMessage(elementBytes, attributes, timestampMsSinceEpoch, recordId);
+ }
+ }
+
+ @VisibleForTesting
+ static final Coder<OutgoingMessage> CODER = new OutgoingMessageCoder();
+
+ // ================================================================================
+ // RecordIdMethod
+ // ================================================================================
+
+ /**
+ * Specify how record ids are to be generated.
+ */
+ @VisibleForTesting
+ enum RecordIdMethod {
+ /** Leave null. */
+ NONE,
+ /** Generate randomly. */
+ RANDOM,
+ /** Generate deterministically. For testing only. */
+ DETERMINISTIC
+ }
+
+ // ================================================================================
+ // ShardFn
+ // ================================================================================
+
+ /**
+ * Convert elements to messages and shard them.
+ */
+ private static class ShardFn<T> extends DoFn<T, KV<Integer, OutgoingMessage>> {
+ private final Counter elementCounter = Metrics.counter(ShardFn.class, "elements");
+ private final Coder<T> elementCoder;
+ private final int numShards;
+ private final RecordIdMethod recordIdMethod;
+ private final SimpleFunction<T, PubsubIO.PubsubMessage> formatFn;
+
+ ShardFn(Coder<T> elementCoder, int numShards,
+ SimpleFunction<T, PubsubIO.PubsubMessage> formatFn, RecordIdMethod recordIdMethod) {
+ this.elementCoder = elementCoder;
+ this.numShards = numShards;
+ this.formatFn = formatFn;
+ this.recordIdMethod = recordIdMethod;
+ }
+
+ @ProcessElement
+ public void processElement(ProcessContext c) throws Exception {
+ elementCounter.inc();
+ byte[] elementBytes = null;
+ Map<String, String> attributes = ImmutableMap.<String, String>of();
+ if (formatFn != null) {
+ PubsubIO.PubsubMessage message = formatFn.apply(c.element());
+ elementBytes = message.getMessage();
+ attributes = message.getAttributeMap();
+ } else {
+ elementBytes = CoderUtils.encodeToByteArray(elementCoder, c.element());
+ }
+
+ long timestampMsSinceEpoch = c.timestamp().getMillis();
+ @Nullable String recordId = null;
+ switch (recordIdMethod) {
+ case NONE:
+ break;
+ case DETERMINISTIC:
+ recordId = Hashing.murmur3_128().hashBytes(elementBytes).toString();
+ break;
+ case RANDOM:
+ // Since these elements go through a GroupByKey, any failures while sending to
+ // Pubsub will be retried without falling back and generating a new record id.
+ // Thus even though we may send the same message to Pubsub twice, it is guaranteed
+ // to have the same record id.
+ recordId = UUID.randomUUID().toString();
+ break;
+ }
+ c.output(KV.of(ThreadLocalRandom.current().nextInt(numShards),
+ new OutgoingMessage(elementBytes, attributes, timestampMsSinceEpoch,
+ recordId)));
+ }
+
+ @Override
+ public void populateDisplayData(Builder builder) {
+ super.populateDisplayData(builder);
+ builder.add(DisplayData.item("numShards", numShards));
+ }
+ }
+
+ // ================================================================================
+ // WriterFn
+ // ================================================================================
+
+ /**
+ * Publish messages to Pubsub in batches.
+ */
+ private static class WriterFn
+ extends DoFn<KV<Integer, Iterable<OutgoingMessage>>, Void> {
+ private final PubsubClientFactory pubsubFactory;
+ private final ValueProvider<TopicPath> topic;
+ private final String timestampLabel;
+ private final String idLabel;
+ private final int publishBatchSize;
+ private final int publishBatchBytes;
+
+ /**
+ * Client on which to talk to Pubsub. Null until created by {@link #startBundle}.
+ */
+ @Nullable
+ private transient PubsubClient pubsubClient;
+
+ private final Counter batchCounter = Metrics.counter(WriterFn.class, "batches");
+ private final Counter elementCounter = Metrics.counter(WriterFn.class, "elements");
+ private final Counter byteCounter = Metrics.counter(WriterFn.class, "bytes");
+
+ WriterFn(
+ PubsubClientFactory pubsubFactory, ValueProvider<TopicPath> topic,
+ String timestampLabel, String idLabel, int publishBatchSize, int publishBatchBytes) {
+ this.pubsubFactory = pubsubFactory;
+ this.topic = topic;
+ this.timestampLabel = timestampLabel;
+ this.idLabel = idLabel;
+ this.publishBatchSize = publishBatchSize;
+ this.publishBatchBytes = publishBatchBytes;
+ }
+
+ /**
+ * BLOCKING
+ * Send {@code messages} as a batch to Pubsub.
+ */
+ private void publishBatch(List<OutgoingMessage> messages, int bytes)
+ throws IOException {
+ int n = pubsubClient.publish(topic.get(), messages);
+ checkState(n == messages.size(), "Attempted to publish %s messages but %s were successful",
+ messages.size(), n);
+ batchCounter.inc();
+ elementCounter.inc(messages.size());
+ byteCounter.inc(bytes);
+ }
+
+ @StartBundle
+ public void startBundle(Context c) throws Exception {
+ checkState(pubsubClient == null, "startBundle invoked without prior finishBundle");
+ pubsubClient = pubsubFactory.newClient(timestampLabel, idLabel,
+ c.getPipelineOptions().as(PubsubOptions.class));
+ }
+
+ @ProcessElement
+ public void processElement(ProcessContext c) throws Exception {
+ List<OutgoingMessage> pubsubMessages = new ArrayList<>(publishBatchSize);
+ int bytes = 0;
+ for (OutgoingMessage message : c.element().getValue()) {
+ if (!pubsubMessages.isEmpty()
+ && bytes + message.elementBytes.length > publishBatchBytes) {
+ // Break large (in bytes) batches into smaller.
+ // (We've already broken by batch size using the trigger below, though that may
+ // run slightly over the actual PUBLISH_BATCH_SIZE. We'll consider that ok since
+ // the hard limit from Pubsub is by bytes rather than number of messages.)
+ // BLOCKS until published.
+ publishBatch(pubsubMessages, bytes);
+ pubsubMessages.clear();
+ bytes = 0;
+ }
+ pubsubMessages.add(message);
+ bytes += message.elementBytes.length;
+ }
+ if (!pubsubMessages.isEmpty()) {
+ // BLOCKS until published.
+ publishBatch(pubsubMessages, bytes);
+ }
+ }
+
+ @FinishBundle
+ public void finishBundle(Context c) throws Exception {
+ pubsubClient.close();
+ pubsubClient = null;
+ }
+
+ @Override
+ public void populateDisplayData(Builder builder) {
+ super.populateDisplayData(builder);
+ String topicString =
+ topic == null ? null
+ : topic.isAccessible() ? topic.get().getPath()
+ : topic.toString();
+ builder.add(DisplayData.item("topic", topicString));
+ builder.add(DisplayData.item("transport", pubsubFactory.getKind()));
+ builder.addIfNotNull(DisplayData.item("timestampLabel", timestampLabel));
+ builder.addIfNotNull(DisplayData.item("idLabel", idLabel));
+ }
+ }
+
+ // ================================================================================
+ // PubsubUnboundedSink
+ // ================================================================================
+
+ /**
+ * Which factory to use for creating Pubsub transport.
+ */
+ private final PubsubClientFactory pubsubFactory;
+
+ /**
+ * Pubsub topic to publish to.
+ */
+ private final ValueProvider<TopicPath> topic;
+
+ /**
+ * Coder for elements. It is the responsibility of the underlying Pubsub transport to
+ * re-encode element bytes if necessary, eg as Base64 strings.
+ */
+ private final Coder<T> elementCoder;
+
+ /**
+ * Pubsub metadata field holding timestamp of each element, or {@literal null} if should use
+ * Pubsub message publish timestamp instead.
+ */
+ @Nullable
+ private final String timestampLabel;
+
+ /**
+ * Pubsub metadata field holding id for each element, or {@literal null} if need to generate
+ * a unique id ourselves.
+ */
+ @Nullable
+ private final String idLabel;
+
+ /**
+ * Number of 'shards' to use so that latency in Pubsub publish can be hidden. Generally this
+ * should be a small multiple of the number of available cores. Too smoll a number results
+ * in too much time lost to blocking Pubsub calls. To large a number results in too many
+ * single-element batches being sent to Pubsub with high per-batch overhead.
+ */
+ private final int numShards;
+
+ /**
+ * Maximum number of messages per publish.
+ */
+ private final int publishBatchSize;
+
+ /**
+ * Maximum size of a publish batch, in bytes.
+ */
+ private final int publishBatchBytes;
+
+ /**
+ * Longest delay between receiving a message and pushing it to Pubsub.
+ */
+ private final Duration maxLatency;
+
+ /**
+ * How record ids should be generated for each record (if {@link #idLabel} is non-{@literal
+ * null}).
+ */
+ private final RecordIdMethod recordIdMethod;
+
+ /**
+ * In order to publish attributes, a formatting function is used to format the output into
+ * a {@link PubsubIO.PubsubMessage}.
+ */
+ private final SimpleFunction<T, PubsubIO.PubsubMessage> formatFn;
+
+ @VisibleForTesting
+ PubsubUnboundedSink(
+ PubsubClientFactory pubsubFactory,
+ ValueProvider<TopicPath> topic,
+ Coder<T> elementCoder,
+ String timestampLabel,
+ String idLabel,
+ int numShards,
+ int publishBatchSize,
+ int publishBatchBytes,
+ Duration maxLatency,
+ SimpleFunction<T, PubsubIO.PubsubMessage> formatFn,
+ RecordIdMethod recordIdMethod) {
+ this.pubsubFactory = pubsubFactory;
+ this.topic = topic;
+ this.elementCoder = elementCoder;
+ this.timestampLabel = timestampLabel;
+ this.idLabel = idLabel;
+ this.numShards = numShards;
+ this.publishBatchSize = publishBatchSize;
+ this.publishBatchBytes = publishBatchBytes;
+ this.maxLatency = maxLatency;
+ this.formatFn = formatFn;
+ this.recordIdMethod = idLabel == null ? RecordIdMethod.NONE : recordIdMethod;
+ }
+
+ public PubsubUnboundedSink(
+ PubsubClientFactory pubsubFactory,
+ ValueProvider<TopicPath> topic,
+ Coder<T> elementCoder,
+ String timestampLabel,
+ String idLabel,
+ SimpleFunction<T, PubsubIO.PubsubMessage> formatFn,
+ int numShards) {
+ this(pubsubFactory, topic, elementCoder, timestampLabel, idLabel, numShards,
+ DEFAULT_PUBLISH_BATCH_SIZE, DEFAULT_PUBLISH_BATCH_BYTES, DEFAULT_MAX_LATENCY,
+ formatFn, RecordIdMethod.RANDOM);
+ }
+
+ /**
+ * Get the topic being written to.
+ */
+ public TopicPath getTopic() {
+ return topic.get();
+ }
+
+ /**
+ * Get the {@link ValueProvider} for the topic being written to.
+ */
+ public ValueProvider<TopicPath> getTopicProvider() {
+ return topic;
+ }
+
+ /**
+ * Get the timestamp label.
+ */
+ @Nullable
+ public String getTimestampLabel() {
+ return timestampLabel;
+ }
+
+ /**
+ * Get the id label.
+ */
+ @Nullable
+ public String getIdLabel() {
+ return idLabel;
+ }
+
+ /**
+ * Get the format function used for PubSub attributes.
+ */
+ @Nullable
+ public SimpleFunction<T, PubsubIO.PubsubMessage> getFormatFn() {
+ return formatFn;
+ }
+
+ /**
+ * Get the Coder used to encode output elements.
+ */
+ public Coder<T> getElementCoder() {
+ return elementCoder;
+ }
+
+ @Override
+ public PDone expand(PCollection<T> input) {
+ input.apply("PubsubUnboundedSink.Window", Window.<T>into(new GlobalWindows())
+ .triggering(
+ Repeatedly.forever(
+ AfterFirst.of(AfterPane.elementCountAtLeast(publishBatchSize),
+ AfterProcessingTime.pastFirstElementInPane()
+ .plusDelayOf(maxLatency))))
+ .discardingFiredPanes())
+ .apply("PubsubUnboundedSink.Shard",
+ ParDo.of(new ShardFn<T>(elementCoder, numShards, formatFn, recordIdMethod)))
+ .setCoder(KvCoder.of(VarIntCoder.of(), CODER))
+ .apply(GroupByKey.<Integer, OutgoingMessage>create())
+ .apply("PubsubUnboundedSink.Writer",
+ ParDo.of(new WriterFn(pubsubFactory, topic, timestampLabel, idLabel,
+ publishBatchSize, publishBatchBytes)));
+ return PDone.in(input.getPipeline());
+ }
+}
[11/50] [abbrv] beam git commit: Remove Orderdness of Input,
Output expansions
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/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 f62b320..3638fc8 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
@@ -17,21 +17,18 @@
*/
package org.apache.beam.sdk.runners;
-import static org.hamcrest.Matchers.contains;
import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.not;
import static org.junit.Assert.assertThat;
-import com.google.common.base.Function;
import com.google.common.collect.Iterables;
-import com.google.common.collect.Lists;
import java.io.Serializable;
import java.util.Collections;
import java.util.HashSet;
-import java.util.List;
import java.util.Map;
+import java.util.Map.Entry;
import java.util.Set;
import org.apache.beam.sdk.Pipeline.PipelineVisitor;
import org.apache.beam.sdk.Pipeline.PipelineVisitor.Defaults;
@@ -224,11 +221,13 @@ public class TransformHierarchyTest implements Serializable {
assertThat(hierarchy.getCurrent(), equalTo(replacement));
hierarchy.setOutput(replacementOutput);
- TaggedPValue taggedOriginal = Iterables.getOnlyElement(originalOutput.expand());
- TaggedPValue taggedReplacement = Iterables.getOnlyElement(replacementOutput.expand());
+ TaggedPValue taggedReplacement = TaggedPValue.ofExpandedValue(replacementOutput);
Map<PValue, ReplacementOutput> replacementOutputs =
Collections.<PValue, ReplacementOutput>singletonMap(
- replacementOutput, ReplacementOutput.of(taggedOriginal, taggedReplacement));
+ replacementOutput,
+ ReplacementOutput.of(
+ TaggedPValue.ofExpandedValue(originalOutput),
+ taggedReplacement));
hierarchy.replaceOutputs(replacementOutputs);
assertThat(replacement.getInputs(), equalTo(original.getInputs()));
@@ -238,8 +237,9 @@ public class TransformHierarchyTest implements Serializable {
replacement.getTransform(), Matchers.<PTransform<?, ?>>equalTo(replacementTransform));
// THe tags of the replacement transform are matched to the appropriate PValues of the original
assertThat(
- replacement.getOutputs(),
- contains(TaggedPValue.of(taggedReplacement.getTag(), taggedOriginal.getValue())));
+ replacement.getOutputs().keySet(),
+ Matchers.<TupleTag<?>>contains(taggedReplacement.getTag()));
+ assertThat(replacement.getOutputs().values(), Matchers.<PValue>contains(originalOutput));
hierarchy.popNode();
}
@@ -294,21 +294,23 @@ public class TransformHierarchyTest implements Serializable {
hierarchy.popNode();
hierarchy.setOutput(replacementOutput.get(longs));
- TaggedPValue originalLongs = Iterables.getOnlyElement(output.expand());
- TaggedPValue replacementLongs = Iterables.getOnlyElement(replacementOutput.expand());
+ Entry<TupleTag<?>, PValue>
+ replacementLongs = Iterables.getOnlyElement(replacementOutput.expand().entrySet());
hierarchy.replaceOutputs(
Collections.<PValue, ReplacementOutput>singletonMap(
- replacementOutput.get(longs), ReplacementOutput.of(originalLongs, replacementLongs)));
+ replacementOutput.get(longs),
+ ReplacementOutput.of(
+ TaggedPValue.ofExpandedValue(output),
+ TaggedPValue.of(replacementLongs.getKey(), replacementLongs.getValue()))));
assertThat(
- replacementParNode.getOutputs(),
- contains(TaggedPValue.of(replacementLongs.getTag(), originalLongs.getValue())));
+ replacementParNode.getOutputs().keySet(),
+ Matchers.<TupleTag<?>>contains(replacementLongs.getKey()));
+ assertThat(replacementParNode.getOutputs().values(), Matchers.<PValue>contains(output));
assertThat(
- compositeNode.getOutputs(),
- contains(
- TaggedPValue.of(
- Iterables.getOnlyElement(replacementOutput.get(longs).expand()).getTag(),
- originalLongs.getValue())));
+ compositeNode.getOutputs().keySet(),
+ equalTo(replacementOutput.get(longs).expand().keySet()));
+ assertThat(compositeNode.getOutputs().values(), Matchers.<PValue>contains(output));
hierarchy.popNode();
}
@@ -340,10 +342,10 @@ public class TransformHierarchyTest implements Serializable {
TransformHierarchy.Node compositeNode = hierarchy.pushNode("Create", begin, create);
hierarchy.finishSpecifyingInput();
assertThat(hierarchy.getCurrent(), equalTo(compositeNode));
- assertThat(compositeNode.getInputs(), Matchers.emptyIterable());
+ assertThat(compositeNode.getInputs().entrySet(), Matchers.empty());
assertThat(compositeNode.getTransform(), Matchers.<PTransform<?, ?>>equalTo(create));
// Not yet set
- assertThat(compositeNode.getOutputs(), Matchers.emptyIterable());
+ assertThat(compositeNode.getOutputs().entrySet(), Matchers.emptyIterable());
assertThat(compositeNode.getEnclosingNode().isRootNode(), is(true));
TransformHierarchy.Node primitiveNode = hierarchy.pushNode("Create/Read", begin, read);
@@ -351,16 +353,14 @@ public class TransformHierarchyTest implements Serializable {
hierarchy.finishSpecifyingInput();
hierarchy.setOutput(created);
hierarchy.popNode();
- assertThat(
- fromTaggedValues(primitiveNode.getOutputs()), Matchers.<PValue>containsInAnyOrder(created));
- assertThat(primitiveNode.getInputs(), Matchers.<TaggedPValue>emptyIterable());
+ assertThat(primitiveNode.getOutputs().values(), Matchers.<PValue>containsInAnyOrder(created));
+ assertThat(primitiveNode.getInputs().entrySet(), Matchers.emptyIterable());
assertThat(primitiveNode.getTransform(), Matchers.<PTransform<?, ?>>equalTo(read));
assertThat(primitiveNode.getEnclosingNode(), equalTo(compositeNode));
hierarchy.setOutput(created);
// The composite is listed as outputting a PValue created by the contained primitive
- assertThat(
- fromTaggedValues(compositeNode.getOutputs()), Matchers.<PValue>containsInAnyOrder(created));
+ assertThat(compositeNode.getOutputs().values(), Matchers.<PValue>containsInAnyOrder(created));
// The producer of that PValue is still the primitive in which it is first output
assertThat(hierarchy.getProducer(created), equalTo(primitiveNode));
hierarchy.popNode();
@@ -457,11 +457,14 @@ public class TransformHierarchyTest implements Serializable {
hierarchy.popNode();
hierarchy.setOutput(replacementOutput.get(longs));
- TaggedPValue originalLongs = Iterables.getOnlyElement(output.expand());
- TaggedPValue replacementLongs = Iterables.getOnlyElement(replacementOutput.expand());
+ Entry<TupleTag<?>, PValue> replacementLongs =
+ Iterables.getOnlyElement(replacementOutput.expand().entrySet());
hierarchy.replaceOutputs(
Collections.<PValue, ReplacementOutput>singletonMap(
- replacementOutput.get(longs), ReplacementOutput.of(originalLongs, replacementLongs)));
+ replacementOutput.get(longs),
+ ReplacementOutput.of(
+ TaggedPValue.ofExpandedValue(output),
+ TaggedPValue.of(replacementLongs.getKey(), replacementLongs.getValue()))));
hierarchy.popNode();
final Set<Node> visitedCompositeNodes = new HashSet<>();
@@ -489,15 +492,4 @@ public class TransformHierarchyTest implements Serializable {
assertThat(visitedPrimitiveNodes, containsInAnyOrder(upstreamNode, replacementParNode));
assertThat(visitedValues, Matchers.<PValue>containsInAnyOrder(upstream, output));
}
-
- private static List<PValue> fromTaggedValues(List<TaggedPValue> taggedValues) {
- return Lists.transform(
- taggedValues,
- new Function<TaggedPValue, PValue>() {
- @Override
- public PValue apply(TaggedPValue input) {
- return input.getValue();
- }
- });
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/sdks/java/core/src/test/java/org/apache/beam/sdk/values/PCollectionListTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/values/PCollectionListTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/values/PCollectionListTest.java
index 2482f32..76cba01 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/values/PCollectionListTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/values/PCollectionListTest.java
@@ -19,8 +19,8 @@ package org.apache.beam.sdk.values;
import static org.hamcrest.CoreMatchers.containsString;
import static org.hamcrest.Matchers.contains;
+import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.hasSize;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.fail;
@@ -28,9 +28,9 @@ import com.google.common.collect.ImmutableList;
import com.google.common.testing.EqualsTester;
import java.util.Collections;
import java.util.List;
+import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.CountingInput;
-import org.apache.beam.sdk.io.CountingInput.BoundedCountingInput;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Create;
import org.hamcrest.Matchers;
@@ -90,27 +90,25 @@ public class PCollectionListTest {
assertThat(
fromEmpty.getAll(), contains(unboundedCount, createOne, boundedCount, maxRecordsCount));
- List<TaggedPValue> expansion = fromEmpty.expand();
- // TaggedPValues are stable between expansions
+ Map<TupleTag<?>, PValue> expansion = fromEmpty.expand();
+ // Tag->PValue mappings are stable between expansions. They don't need to be stable across
+ // different list instances, though
assertThat(expansion, equalTo(fromEmpty.expand()));
- // TaggedPValues are equivalent between equivalent lists
- assertThat(
- expansion,
- equalTo(
- PCollectionList.of(unboundedCount)
- .and(createOne)
- .and(boundedCount)
- .and(maxRecordsCount)
- .expand()));
List<PCollection<Long>> expectedList =
ImmutableList.of(unboundedCount, createOne, boundedCount, maxRecordsCount);
- for (int i = 0; i < expansion.size(); i++) {
- assertThat(
- "Index " + i + " should have equal PValue",
- expansion.get(i).getValue(),
- Matchers.<PValue>equalTo(expectedList.get(i)));
- }
+ assertThat(expansion.values(), containsInAnyOrder(expectedList.toArray()));
+ }
+
+ @Test
+ public void testExpandWithDuplicates() {
+ Pipeline p = TestPipeline.create();
+ PCollection<Long> createOne = p.apply("CreateOne", Create.of(1L, 2L, 3L));
+
+ PCollectionList<Long> list = PCollectionList.of(createOne).and(createOne).and(createOne);
+ assertThat(
+ list.expand().values(),
+ Matchers.<PValue>containsInAnyOrder(createOne, createOne, createOne));
}
@Test
@@ -121,15 +119,15 @@ public class PCollectionListTest {
PCollection<String> third = p.apply("Syntactic", Create.of("eggs", "baz"));
EqualsTester tester = new EqualsTester();
- tester.addEqualityGroup(PCollectionList.empty(p), PCollectionList.empty(p));
- tester.addEqualityGroup(PCollectionList.of(first).and(second));
+// tester.addEqualityGroup(PCollectionList.empty(p), PCollectionList.empty(p));
+// tester.addEqualityGroup(PCollectionList.of(first).and(second));
// Constructors should all produce equivalent
tester.addEqualityGroup(
PCollectionList.of(first).and(second).and(third),
PCollectionList.of(first).and(second).and(third),
- PCollectionList.<String>empty(p).and(first).and(second).and(third),
- PCollectionList.of(ImmutableList.of(first, second, third)),
- PCollectionList.of(first).and(ImmutableList.of(second, third)),
+// PCollectionList.<String>empty(p).and(first).and(second).and(third),
+// PCollectionList.of(ImmutableList.of(first, second, third)),
+// PCollectionList.of(first).and(ImmutableList.of(second, third)),
PCollectionList.of(ImmutableList.of(first, second)).and(third));
// Order is considered
tester.addEqualityGroup(PCollectionList.of(first).and(third).and(second));
@@ -137,28 +135,4 @@ public class PCollectionListTest {
tester.testEquals();
}
-
- @Test
- public void testExpansionOrderWithDuplicates() {
- TestPipeline p = TestPipeline.create();
- BoundedCountingInput count = CountingInput.upTo(10L);
- PCollection<Long> firstCount = p.apply("CountFirst", count);
- PCollection<Long> secondCount = p.apply("CountSecond", count);
-
- PCollectionList<Long> counts =
- PCollectionList.of(firstCount).and(secondCount).and(firstCount).and(firstCount);
-
- ImmutableList<PCollection<Long>> expectedOrder =
- ImmutableList.of(firstCount, secondCount, firstCount, firstCount);
- PCollectionList<Long> reconstructed = PCollectionList.empty(p);
- assertThat(counts.expand(), hasSize(4));
- for (int i = 0; i < 4; i++) {
- PValue value = counts.expand().get(i).getValue();
- assertThat(
- "Index " + i + " should be equal", value,
- Matchers.<PValue>equalTo(expectedOrder.get(i)));
- reconstructed = reconstructed.and((PCollection<Long>) value);
- }
- assertThat(reconstructed, equalTo(counts));
- }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/sdks/java/core/src/test/java/org/apache/beam/sdk/values/PCollectionTupleTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/values/PCollectionTupleTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/values/PCollectionTupleTest.java
index 010d726..0a0abd6 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/values/PCollectionTupleTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/values/PCollectionTupleTest.java
@@ -29,6 +29,7 @@ import java.io.Serializable;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
+import java.util.Map.Entry;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.CountingInput;
import org.apache.beam.sdk.testing.PAssert;
@@ -153,8 +154,8 @@ public final class PCollectionTupleTest implements Serializable {
PCollectionTuple.of(intTag, ints).and(longTag, longs).and(strTag, strs);
assertThat(tuple.getAll(), equalTo(pcsByTag));
PCollectionTuple reconstructed = PCollectionTuple.empty(p);
- for (TaggedPValue taggedValue : tuple.expand()) {
- TupleTag<?> tag = taggedValue.getTag();
+ for (Entry<TupleTag<?>, PValue> taggedValue : tuple.expand().entrySet()) {
+ TupleTag<?> tag = taggedValue.getKey();
PValue value = taggedValue.getValue();
assertThat("The tag should map back to the value", tuple.get(tag), equalTo(value));
assertThat(value, Matchers.<PValue>equalTo(pcsByTag.get(tag)));
http://git-wip-us.apache.org/repos/asf/beam/blob/0e5737fd/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteResult.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteResult.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteResult.java
index 07fbc68..3e0f51c 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteResult.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteResult.java
@@ -18,12 +18,11 @@
package org.apache.beam.sdk.io.gcp.bigquery;
import java.util.Collections;
-import java.util.List;
-
+import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.values.POutputValueBase;
-import org.apache.beam.sdk.values.TaggedPValue;
-
+import org.apache.beam.sdk.values.PValue;
+import org.apache.beam.sdk.values.TupleTag;
/**
* The result of a {@link BigQueryIO.Write} transform.
@@ -37,8 +36,8 @@ final class WriteResult extends POutputValueBase {
}
@Override
- public List<TaggedPValue> expand() {
- return Collections.emptyList();
+ public Map<TupleTag<?>, PValue> expand() {
+ return Collections.emptyMap();
}
private WriteResult(Pipeline pipeline) {
[38/50] [abbrv] beam git commit: [BEAM-115] Update timer/state fields
on ParDoPayload to use a map field for consistent tag usage
Posted by jb...@apache.org.
[BEAM-115] Update timer/state fields on ParDoPayload to use a map field for consistent tag usage
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/b1148442
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/b1148442
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/b1148442
Branch: refs/heads/DSL_SQL
Commit: b1148442eefebd198a8a876ab29dbc6e9820bffa
Parents: 571631a
Author: Luke Cwik <lc...@google.com>
Authored: Wed Apr 12 09:25:39 2017 -0700
Committer: Luke Cwik <lc...@google.com>
Committed: Wed Apr 12 10:59:22 2017 -0700
----------------------------------------------------------------------
.../runner-api/src/main/proto/beam_runner_api.proto | 10 +++++-----
1 file changed, 5 insertions(+), 5 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/b1148442/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 3b7c5bd..ebbd928 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
@@ -215,11 +215,11 @@ message ParDoPayload {
// the expected access pattern.
map<string, SideInput> side_inputs = 3;
- // (Optional) if the DoFn uses state, a list of the specs for cells.
- repeated StateSpec state_specs = 4;
+ // (Optional) A mapping of local state names to state specifications.
+ map<string, StateSpec> state_specs = 4;
- // (Optional) if the DoFn uses timers, a list of the specs for timers.
- repeated TimerSpec timer_specs = 5;
+ // (Optional) A mapping of local timer names to timer specifications.
+ map<string, TimerSpec> timer_specs = 5;
}
// Parameters that a UDF might require.
@@ -227,7 +227,7 @@ message ParDoPayload {
// The details of how a runner sends these parameters to the SDK harness
// are the subject of the Fn API.
//
-// The details ofo how an SDK harness delivers them to the UDF is entirely
+// The details of how an SDK harness delivers them to the UDF is entirely
// up to the SDK. (for some SDKs there may be parameters that are not
// represented here if the runner doesn't need to do anything)
//
[06/50] [abbrv] beam git commit: This closes #2492
Posted by jb...@apache.org.
This closes #2492
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/77712c93
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/77712c93
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/77712c93
Branch: refs/heads/DSL_SQL
Commit: 77712c9363937892208e791d6f6856639cbf0c67
Parents: 38cf23e 732610d
Author: Thomas Groh <tg...@google.com>
Authored: Tue Apr 11 11:12:13 2017 -0700
Committer: Thomas Groh <tg...@google.com>
Committed: Tue Apr 11 11:12:13 2017 -0700
----------------------------------------------------------------------
.../io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
----------------------------------------------------------------------
[21/50] [abbrv] beam git commit: This closes #2496
Posted by jb...@apache.org.
This closes #2496
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/132d4c55
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/132d4c55
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/132d4c55
Branch: refs/heads/DSL_SQL
Commit: 132d4c55dac66f63adab0d32ad7dc53ffda82df9
Parents: 17a41ab ee92b96
Author: Ahmet Altay <al...@google.com>
Authored: Tue Apr 11 15:32:21 2017 -0700
Committer: Ahmet Altay <al...@google.com>
Committed: Tue Apr 11 15:32:21 2017 -0700
----------------------------------------------------------------------
.../apache_beam/runners/dataflow/dataflow_runner.py | 6 ++++++
sdks/python/apache_beam/utils/value_provider.py | 14 ++++++++++++--
2 files changed, 18 insertions(+), 2 deletions(-)
----------------------------------------------------------------------
[50/50] [abbrv] beam git commit: Add DSLs module
Posted by jb...@apache.org.
Add DSLs module
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/3625dbd9
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/3625dbd9
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/3625dbd9
Branch: refs/heads/DSL_SQL
Commit: 3625dbd9e00db28d2175c6ac000792fd5c5d8f65
Parents: 1788cef
Author: Jean-Baptiste Onofr� <jb...@apache.org>
Authored: Tue Aug 9 11:17:32 2016 +0200
Committer: Jean-Baptiste Onofr� <jb...@apache.org>
Committed: Wed Apr 12 21:56:01 2017 +0200
----------------------------------------------------------------------
dsls/pom.xml | 56 +++++++++++++++++++++++++++++++++++++++++++++++++++++++
pom.xml | 1 +
2 files changed, 57 insertions(+)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/3625dbd9/dsls/pom.xml
----------------------------------------------------------------------
diff --git a/dsls/pom.xml b/dsls/pom.xml
new file mode 100644
index 0000000..6e00171
--- /dev/null
+++ b/dsls/pom.xml
@@ -0,0 +1,56 @@
+<?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-parent</artifactId>
+ <version>0.7.0-SNAPSHOT</version>
+ <relativePath>../pom.xml</relativePath>
+ </parent>
+
+ <artifactId>beam-dsls-parent</artifactId>
+ <name>Apache Beam :: DSLs</name>
+
+ <modules>
+ <!-- <module>sql</module> -->
+ </modules>
+
+ <build>
+ <pluginManagement>
+ <plugins>
+ <!-- DSLs will generally offer test suites for runners, as sdks/java does. -->
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-jar-plugin</artifactId>
+ <executions>
+ <execution>
+ <id>default-test-jar</id>
+ <goals>
+ <goal>test-jar</goal>
+ </goals>
+ </execution>
+ </executions>
+ </plugin>
+ </plugins>
+ </pluginManagement>
+ </build>
+
+</project>
\ No newline at end of file
http://git-wip-us.apache.org/repos/asf/beam/blob/3625dbd9/pom.xml
----------------------------------------------------------------------
diff --git a/pom.xml b/pom.xml
index 09f3985..ef312c1 100644
--- a/pom.xml
+++ b/pom.xml
@@ -152,6 +152,7 @@
<module>sdks/java/build-tools</module>
<module>sdks</module>
<module>runners</module>
+ <module>dsls</module>
<module>examples</module>
<!-- sdks/java/javadoc builds project-wide Javadoc. It has to run last. -->
<module>sdks/java/javadoc</module>
[47/50] [abbrv] beam git commit: This closes #2515
Posted by jb...@apache.org.
This closes #2515
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/a405a2c4
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/a405a2c4
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/a405a2c4
Branch: refs/heads/DSL_SQL
Commit: a405a2c46c92875e6a10e1d990cf4507a6ea7650
Parents: b7ddab2 0f5c363
Author: Ahmet Altay <al...@google.com>
Authored: Wed Apr 12 11:51:09 2017 -0700
Committer: Ahmet Altay <al...@google.com>
Committed: Wed Apr 12 11:51:09 2017 -0700
----------------------------------------------------------------------
sdks/python/apache_beam/examples/wordcount.py | 32 ++--
.../apache_beam/internal/gcp/json_value.py | 6 -
sdks/python/apache_beam/io/filebasedsource.py | 54 ++----
.../apache_beam/io/filebasedsource_test.py | 24 ---
sdks/python/apache_beam/io/fileio.py | 56 ++-----
sdks/python/apache_beam/io/fileio_test.py | 45 ++---
.../runners/dataflow/dataflow_runner.py | 6 -
.../runners/dataflow/internal/apiclient.py | 1 -
.../apache_beam/runners/direct/direct_runner.py | 9 -
sdks/python/apache_beam/transforms/display.py | 1 -
.../apache_beam/transforms/display_test.py | 36 ----
.../apache_beam/utils/pipeline_options.py | 92 +----------
.../apache_beam/utils/pipeline_options_test.py | 52 +-----
sdks/python/apache_beam/utils/value_provider.py | 120 --------------
.../apache_beam/utils/value_provider_test.py | 165 -------------------
15 files changed, 65 insertions(+), 634 deletions(-)
----------------------------------------------------------------------
[28/50] [abbrv] beam git commit: [BEAM-1722] Move PubsubIO into the
google-cloud-platform module
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubJsonClientTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubJsonClientTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubJsonClientTest.java
new file mode 100644
index 0000000..d290994
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubJsonClientTest.java
@@ -0,0 +1,139 @@
+/*
+ * 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.pubsub;
+
+import static org.junit.Assert.assertEquals;
+
+import com.google.api.services.pubsub.Pubsub;
+import com.google.api.services.pubsub.model.PublishRequest;
+import com.google.api.services.pubsub.model.PublishResponse;
+import com.google.api.services.pubsub.model.PubsubMessage;
+import com.google.api.services.pubsub.model.PullRequest;
+import com.google.api.services.pubsub.model.PullResponse;
+import com.google.api.services.pubsub.model.ReceivedMessage;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.OutgoingMessage;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+/**
+ * Tests for PubsubJsonClient.
+ */
+@RunWith(JUnit4.class)
+public class PubsubJsonClientTest {
+ private Pubsub mockPubsub;
+ private PubsubClient client;
+
+ private static final TopicPath TOPIC = PubsubClient.topicPathFromName("testProject", "testTopic");
+ private static final SubscriptionPath SUBSCRIPTION =
+ PubsubClient.subscriptionPathFromName("testProject", "testSubscription");
+ private static final long REQ_TIME = 1234L;
+ private static final long PUB_TIME = 3456L;
+ private static final long MESSAGE_TIME = 6789L;
+ private static final String TIMESTAMP_LABEL = "timestamp";
+ private static final String ID_LABEL = "id";
+ private static final String MESSAGE_ID = "testMessageId";
+ private static final String DATA = "testData";
+ private static final String RECORD_ID = "testRecordId";
+ private static final String ACK_ID = "testAckId";
+
+ @Before
+ public void setup() throws IOException {
+ mockPubsub = Mockito.mock(Pubsub.class, Mockito.RETURNS_DEEP_STUBS);
+ client = new PubsubJsonClient(TIMESTAMP_LABEL, ID_LABEL, mockPubsub);
+ }
+
+ @After
+ public void teardown() throws IOException {
+ client.close();
+ client = null;
+ mockPubsub = null;
+ }
+
+ @Test
+ public void pullOneMessage() throws IOException {
+ String expectedSubscription = SUBSCRIPTION.getPath();
+ PullRequest expectedRequest =
+ new PullRequest().setReturnImmediately(true).setMaxMessages(10);
+ PubsubMessage expectedPubsubMessage = new PubsubMessage()
+ .setMessageId(MESSAGE_ID)
+ .encodeData(DATA.getBytes())
+ .setPublishTime(String.valueOf(PUB_TIME))
+ .setAttributes(
+ ImmutableMap.of(TIMESTAMP_LABEL, String.valueOf(MESSAGE_TIME),
+ ID_LABEL, RECORD_ID));
+ ReceivedMessage expectedReceivedMessage =
+ new ReceivedMessage().setMessage(expectedPubsubMessage)
+ .setAckId(ACK_ID);
+ PullResponse expectedResponse =
+ new PullResponse().setReceivedMessages(ImmutableList.of(expectedReceivedMessage));
+ Mockito.when((Object) (mockPubsub.projects()
+ .subscriptions()
+ .pull(expectedSubscription, expectedRequest)
+ .execute()))
+ .thenReturn(expectedResponse);
+ List<IncomingMessage> acutalMessages = client.pull(REQ_TIME, SUBSCRIPTION, 10, true);
+ assertEquals(1, acutalMessages.size());
+ IncomingMessage actualMessage = acutalMessages.get(0);
+ assertEquals(ACK_ID, actualMessage.ackId);
+ assertEquals(DATA, new String(actualMessage.elementBytes));
+ assertEquals(RECORD_ID, actualMessage.recordId);
+ assertEquals(REQ_TIME, actualMessage.requestTimeMsSinceEpoch);
+ assertEquals(MESSAGE_TIME, actualMessage.timestampMsSinceEpoch);
+ }
+
+ @Test
+ public void publishOneMessage() throws IOException {
+ String expectedTopic = TOPIC.getPath();
+ PubsubMessage expectedPubsubMessage = new PubsubMessage()
+ .encodeData(DATA.getBytes())
+ .setAttributes(
+ ImmutableMap.<String, String> builder()
+ .put(TIMESTAMP_LABEL, String.valueOf(MESSAGE_TIME))
+ .put(ID_LABEL, RECORD_ID)
+ .put("k", "v").build());
+ PublishRequest expectedRequest = new PublishRequest()
+ .setMessages(ImmutableList.of(expectedPubsubMessage));
+ PublishResponse expectedResponse = new PublishResponse()
+ .setMessageIds(ImmutableList.of(MESSAGE_ID));
+ Mockito.when((Object) (mockPubsub.projects()
+ .topics()
+ .publish(expectedTopic, expectedRequest)
+ .execute()))
+ .thenReturn(expectedResponse);
+ Map<String, String> attrs = new HashMap<>();
+ attrs.put("k", "v");
+ OutgoingMessage actualMessage = new OutgoingMessage(
+ DATA.getBytes(), attrs, MESSAGE_TIME, RECORD_ID);
+ int n = client.publish(TOPIC, ImmutableList.of(actualMessage));
+ assertEquals(1, n);
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubTestClientTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubTestClientTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubTestClientTest.java
new file mode 100644
index 0000000..18180af
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubTestClientTest.java
@@ -0,0 +1,114 @@
+/*
+ * 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.pubsub;
+
+import static org.junit.Assert.assertEquals;
+
+import com.google.api.client.util.Clock;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Lists;
+import com.google.common.collect.Sets;
+import java.io.IOException;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicLong;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.OutgoingMessage;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubTestClient.PubsubTestClientFactory;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Tests for PubsubTestClient.
+ */
+@RunWith(JUnit4.class)
+public class PubsubTestClientTest {
+ private static final TopicPath TOPIC = PubsubClient.topicPathFromName("testProject", "testTopic");
+ private static final SubscriptionPath SUBSCRIPTION =
+ PubsubClient.subscriptionPathFromName("testProject", "testSubscription");
+ private static final long REQ_TIME = 1234L;
+ private static final long MESSAGE_TIME = 6789L;
+ private static final String MESSAGE_ID = "testMessageId";
+ private static final String DATA = "testData";
+ private static final String ACK_ID = "testAckId";
+ private static final int ACK_TIMEOUT_S = 60;
+
+ @Test
+ public void pullOneMessage() throws IOException {
+ final AtomicLong now = new AtomicLong();
+ Clock clock = new Clock() {
+ @Override
+ public long currentTimeMillis() {
+ return now.get();
+ }
+ };
+ IncomingMessage expectedIncomingMessage =
+ new IncomingMessage(DATA.getBytes(), null, MESSAGE_TIME, REQ_TIME, ACK_ID, MESSAGE_ID);
+ try (PubsubTestClientFactory factory =
+ PubsubTestClient.createFactoryForPull(clock, SUBSCRIPTION, ACK_TIMEOUT_S,
+ Lists.newArrayList(expectedIncomingMessage))) {
+ try (PubsubTestClient client = (PubsubTestClient) factory.newClient(null, null, null)) {
+ now.set(REQ_TIME);
+ client.advance();
+ List<IncomingMessage> incomingMessages = client.pull(now.get(), SUBSCRIPTION, 1, true);
+ assertEquals(1, incomingMessages.size());
+ assertEquals(expectedIncomingMessage, incomingMessages.get(0));
+ // Timeout on ACK.
+ now.addAndGet((ACK_TIMEOUT_S + 10) * 1000);
+ client.advance();
+ incomingMessages = client.pull(now.get(), SUBSCRIPTION, 1, true);
+ assertEquals(1, incomingMessages.size());
+ assertEquals(expectedIncomingMessage.withRequestTime(now.get()), incomingMessages.get(0));
+ now.addAndGet(10 * 1000);
+ client.advance();
+ // Extend ack
+ client.modifyAckDeadline(SUBSCRIPTION, ImmutableList.of(ACK_ID), 20);
+ // Timeout on extended ACK
+ now.addAndGet(30 * 1000);
+ client.advance();
+ incomingMessages = client.pull(now.get(), SUBSCRIPTION, 1, true);
+ assertEquals(1, incomingMessages.size());
+ assertEquals(expectedIncomingMessage.withRequestTime(now.get()), incomingMessages.get(0));
+ // Extend ack
+ client.modifyAckDeadline(SUBSCRIPTION, ImmutableList.of(ACK_ID), 20);
+ // Ack
+ now.addAndGet(15 * 1000);
+ client.advance();
+ client.acknowledge(SUBSCRIPTION, ImmutableList.of(ACK_ID));
+ }
+ }
+ }
+
+ @Test
+ public void publishOneMessage() throws IOException {
+ OutgoingMessage expectedOutgoingMessage =
+ new OutgoingMessage(DATA.getBytes(), null, MESSAGE_TIME, MESSAGE_ID);
+ try (PubsubTestClientFactory factory =
+ PubsubTestClient.createFactoryForPublish(
+ TOPIC,
+ Sets.newHashSet(expectedOutgoingMessage),
+ ImmutableList.<OutgoingMessage>of())) {
+ try (PubsubTestClient client = (PubsubTestClient) factory.newClient(null, null, null)) {
+ client.publish(TOPIC, ImmutableList.of(expectedOutgoingMessage));
+ }
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSinkTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSinkTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSinkTest.java
new file mode 100644
index 0000000..be425d4
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSinkTest.java
@@ -0,0 +1,188 @@
+/*
+ * 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.pubsub;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.hash.Hashing;
+import java.io.IOException;
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.OutgoingMessage;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubTestClient.PubsubTestClientFactory;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSink.RecordIdMethod;
+import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
+import org.apache.beam.sdk.testing.CoderProperties;
+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.DoFn;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.joda.time.Duration;
+import org.joda.time.Instant;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Test PubsubUnboundedSink.
+ */
+@RunWith(JUnit4.class)
+public class PubsubUnboundedSinkTest implements Serializable {
+ private static final TopicPath TOPIC = PubsubClient.topicPathFromName("testProject", "testTopic");
+ private static final String DATA = "testData";
+ private static final Map<String, String> ATTRIBUTES =
+ ImmutableMap.<String, String>builder().put("a", "b").put("c", "d").build();
+ private static final long TIMESTAMP = 1234L;
+ private static final String TIMESTAMP_LABEL = "timestamp";
+ private static final String ID_LABEL = "id";
+ private static final int NUM_SHARDS = 10;
+
+ private static class Stamp extends DoFn<String, String> {
+ @ProcessElement
+ public void processElement(ProcessContext c) {
+ c.outputWithTimestamp(c.element(), new Instant(TIMESTAMP));
+ }
+ }
+
+ private String getRecordId(String data) {
+ return Hashing.murmur3_128().hashBytes(data.getBytes()).toString();
+ }
+
+ @Rule
+ public transient TestPipeline p = TestPipeline.create();
+
+ @Test
+ public void saneCoder() throws Exception {
+ OutgoingMessage message = new OutgoingMessage(
+ DATA.getBytes(), ImmutableMap.<String, String>of(), TIMESTAMP, getRecordId(DATA));
+ CoderProperties.coderDecodeEncodeEqual(PubsubUnboundedSink.CODER, message);
+ CoderProperties.coderSerializable(PubsubUnboundedSink.CODER);
+ }
+
+ @Test
+ @Category(NeedsRunner.class)
+ public void sendOneMessage() throws IOException {
+ List<OutgoingMessage> outgoing =
+ ImmutableList.of(new OutgoingMessage(
+ DATA.getBytes(),
+ ATTRIBUTES,
+ TIMESTAMP, getRecordId(DATA)));
+ int batchSize = 1;
+ int batchBytes = 1;
+ try (PubsubTestClientFactory factory =
+ PubsubTestClient.createFactoryForPublish(TOPIC, outgoing,
+ ImmutableList.<OutgoingMessage>of())) {
+ PubsubUnboundedSink<String> sink =
+ new PubsubUnboundedSink<>(factory, StaticValueProvider.of(TOPIC), StringUtf8Coder.of(),
+ TIMESTAMP_LABEL, ID_LABEL, NUM_SHARDS, batchSize, batchBytes,
+ Duration.standardSeconds(2),
+ new SimpleFunction<String, PubsubIO.PubsubMessage>() {
+ @Override
+ public PubsubIO.PubsubMessage apply(String input) {
+ return new PubsubIO.PubsubMessage(input.getBytes(), ATTRIBUTES);
+ }
+ },
+ RecordIdMethod.DETERMINISTIC);
+ p.apply(Create.of(ImmutableList.of(DATA)))
+ .apply(ParDo.of(new Stamp()))
+ .apply(sink);
+ p.run();
+ }
+ // The PubsubTestClientFactory will assert fail on close if the actual published
+ // message does not match the expected publish message.
+ }
+
+ @Test
+ @Category(NeedsRunner.class)
+ public void sendMoreThanOneBatchByNumMessages() throws IOException {
+ List<OutgoingMessage> outgoing = new ArrayList<>();
+ List<String> data = new ArrayList<>();
+ int batchSize = 2;
+ int batchBytes = 1000;
+ for (int i = 0; i < batchSize * 10; i++) {
+ String str = String.valueOf(i);
+ outgoing.add(new OutgoingMessage(
+ str.getBytes(), ImmutableMap.<String, String>of(), TIMESTAMP, getRecordId(str)));
+ data.add(str);
+ }
+ try (PubsubTestClientFactory factory =
+ PubsubTestClient.createFactoryForPublish(TOPIC, outgoing,
+ ImmutableList.<OutgoingMessage>of())) {
+ PubsubUnboundedSink<String> sink =
+ new PubsubUnboundedSink<>(factory, StaticValueProvider.of(TOPIC), StringUtf8Coder.of(),
+ TIMESTAMP_LABEL, ID_LABEL, NUM_SHARDS, batchSize, batchBytes,
+ Duration.standardSeconds(2), null, RecordIdMethod.DETERMINISTIC);
+ p.apply(Create.of(data))
+ .apply(ParDo.of(new Stamp()))
+ .apply(sink);
+ p.run();
+ }
+ // The PubsubTestClientFactory will assert fail on close if the actual published
+ // message does not match the expected publish message.
+ }
+
+ @Test
+ @Category(NeedsRunner.class)
+ public void sendMoreThanOneBatchByByteSize() throws IOException {
+ List<OutgoingMessage> outgoing = new ArrayList<>();
+ List<String> data = new ArrayList<>();
+ int batchSize = 100;
+ int batchBytes = 10;
+ int n = 0;
+ while (n < batchBytes * 10) {
+ StringBuilder sb = new StringBuilder();
+ for (int i = 0; i < batchBytes; i++) {
+ sb.append(String.valueOf(n));
+ }
+ String str = sb.toString();
+ outgoing.add(new OutgoingMessage(
+ str.getBytes(), ImmutableMap.<String, String>of(), TIMESTAMP, getRecordId(str)));
+ data.add(str);
+ n += str.length();
+ }
+ try (PubsubTestClientFactory factory =
+ PubsubTestClient.createFactoryForPublish(TOPIC, outgoing,
+ ImmutableList.<OutgoingMessage>of())) {
+ PubsubUnboundedSink<String> sink =
+ new PubsubUnboundedSink<>(factory, StaticValueProvider.of(TOPIC),
+ StringUtf8Coder.of(), TIMESTAMP_LABEL, ID_LABEL,
+ NUM_SHARDS, batchSize, batchBytes, Duration.standardSeconds(2),
+ null, RecordIdMethod.DETERMINISTIC);
+ p.apply(Create.of(data))
+ .apply(ParDo.of(new Stamp()))
+ .apply(sink);
+ p.run();
+ }
+ // The PubsubTestClientFactory will assert fail on close if the actual published
+ // message does not match the expected publish message.
+ }
+
+ // TODO: We would like to test that failed Pubsub publish calls cause the already assigned
+ // (and random) record ids to be reused. However that can't be done without the test runnner
+ // supporting retrying bundles.
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSourceTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSourceTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSourceTest.java
new file mode 100644
index 0000000..d2e88c3
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubUnboundedSourceTest.java
@@ -0,0 +1,409 @@
+/*
+ * 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.pubsub;
+
+import static junit.framework.TestCase.assertFalse;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.greaterThan;
+import static org.hamcrest.Matchers.hasSize;
+import static org.hamcrest.Matchers.lessThanOrEqualTo;
+import static org.hamcrest.Matchers.not;
+import static org.hamcrest.Matchers.nullValue;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
+
+import com.google.api.client.util.Clock;
+import com.google.common.collect.ImmutableList;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicLong;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubTestClient.PubsubTestClientFactory;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubCheckpoint;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubReader;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubSource;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
+import org.apache.beam.sdk.testing.CoderProperties;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.util.CoderUtils;
+import org.joda.time.Instant;
+import org.junit.After;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Test PubsubUnboundedSource.
+ */
+@RunWith(JUnit4.class)
+public class PubsubUnboundedSourceTest {
+ private static final SubscriptionPath SUBSCRIPTION =
+ PubsubClient.subscriptionPathFromName("testProject", "testSubscription");
+ private static final String DATA = "testData";
+ private static final long TIMESTAMP = 1234L;
+ private static final long REQ_TIME = 6373L;
+ private static final String TIMESTAMP_LABEL = "timestamp";
+ private static final String ID_LABEL = "id";
+ private static final String ACK_ID = "testAckId";
+ private static final String RECORD_ID = "testRecordId";
+ private static final int ACK_TIMEOUT_S = 60;
+
+ private AtomicLong now;
+ private Clock clock;
+ private PubsubTestClientFactory factory;
+ private PubsubSource<String> primSource;
+
+ @Rule
+ public TestPipeline p = TestPipeline.create();
+
+ private void setupOneMessage(Iterable<IncomingMessage> incoming) {
+ now = new AtomicLong(REQ_TIME);
+ clock = new Clock() {
+ @Override
+ public long currentTimeMillis() {
+ return now.get();
+ }
+ };
+ factory = PubsubTestClient.createFactoryForPull(clock, SUBSCRIPTION, ACK_TIMEOUT_S, incoming);
+ PubsubUnboundedSource<String> source =
+ new PubsubUnboundedSource<>(
+ clock, factory, null, null, StaticValueProvider.of(SUBSCRIPTION),
+ StringUtf8Coder.of(), TIMESTAMP_LABEL, ID_LABEL, null);
+ primSource = new PubsubSource<>(source);
+ }
+
+ private void setupOneMessage() {
+ setupOneMessage(ImmutableList.of(
+ new IncomingMessage(DATA.getBytes(), null, TIMESTAMP, 0, ACK_ID, RECORD_ID)));
+ }
+
+ @After
+ public void after() throws IOException {
+ factory.close();
+ now = null;
+ clock = null;
+ primSource = null;
+ factory = null;
+ }
+
+ @Test
+ public void checkpointCoderIsSane() throws Exception {
+ setupOneMessage(ImmutableList.<IncomingMessage>of());
+ CoderProperties.coderSerializable(primSource.getCheckpointMarkCoder());
+ // Since we only serialize/deserialize the 'notYetReadIds', and we don't want to make
+ // equals on checkpoints ignore those fields, we'll test serialization and deserialization
+ // of checkpoints in multipleReaders below.
+ }
+
+ @Test
+ public void readOneMessage() throws IOException {
+ setupOneMessage();
+ PubsubReader<String> reader = primSource.createReader(p.getOptions(), null);
+ // Read one message.
+ assertTrue(reader.start());
+ assertEquals(DATA, reader.getCurrent());
+ assertFalse(reader.advance());
+ // ACK the message.
+ PubsubCheckpoint<String> checkpoint = reader.getCheckpointMark();
+ checkpoint.finalizeCheckpoint();
+ reader.close();
+ }
+
+ @Test
+ public void timeoutAckAndRereadOneMessage() throws IOException {
+ setupOneMessage();
+ PubsubReader<String> reader = primSource.createReader(p.getOptions(), null);
+ PubsubTestClient pubsubClient = (PubsubTestClient) reader.getPubsubClient();
+ assertTrue(reader.start());
+ assertEquals(DATA, reader.getCurrent());
+ // Let the ACK deadline for the above expire.
+ now.addAndGet(65 * 1000);
+ pubsubClient.advance();
+ // We'll now receive the same message again.
+ assertTrue(reader.advance());
+ assertEquals(DATA, reader.getCurrent());
+ assertFalse(reader.advance());
+ // Now ACK the message.
+ PubsubCheckpoint<String> checkpoint = reader.getCheckpointMark();
+ checkpoint.finalizeCheckpoint();
+ reader.close();
+ }
+
+ @Test
+ public void extendAck() throws IOException {
+ setupOneMessage();
+ PubsubReader<String> reader = primSource.createReader(p.getOptions(), null);
+ PubsubTestClient pubsubClient = (PubsubTestClient) reader.getPubsubClient();
+ // Pull the first message but don't take a checkpoint for it.
+ assertTrue(reader.start());
+ assertEquals(DATA, reader.getCurrent());
+ // Extend the ack
+ now.addAndGet(55 * 1000);
+ pubsubClient.advance();
+ assertFalse(reader.advance());
+ // Extend the ack again
+ now.addAndGet(25 * 1000);
+ pubsubClient.advance();
+ assertFalse(reader.advance());
+ // Now ACK the message.
+ PubsubCheckpoint<String> checkpoint = reader.getCheckpointMark();
+ checkpoint.finalizeCheckpoint();
+ reader.close();
+ }
+
+ @Test
+ public void timeoutAckExtensions() throws IOException {
+ setupOneMessage();
+ PubsubReader<String> reader = primSource.createReader(p.getOptions(), null);
+ PubsubTestClient pubsubClient = (PubsubTestClient) reader.getPubsubClient();
+ // Pull the first message but don't take a checkpoint for it.
+ assertTrue(reader.start());
+ assertEquals(DATA, reader.getCurrent());
+ // Extend the ack.
+ now.addAndGet(55 * 1000);
+ pubsubClient.advance();
+ assertFalse(reader.advance());
+ // Let the ack expire.
+ for (int i = 0; i < 3; i++) {
+ now.addAndGet(25 * 1000);
+ pubsubClient.advance();
+ assertFalse(reader.advance());
+ }
+ // Wait for resend.
+ now.addAndGet(25 * 1000);
+ pubsubClient.advance();
+ // Reread the same message.
+ assertTrue(reader.advance());
+ assertEquals(DATA, reader.getCurrent());
+ // Now ACK the message.
+ PubsubCheckpoint<String> checkpoint = reader.getCheckpointMark();
+ checkpoint.finalizeCheckpoint();
+ reader.close();
+ }
+
+ @Test
+ public void multipleReaders() throws IOException {
+ List<IncomingMessage> incoming = new ArrayList<>();
+ for (int i = 0; i < 2; i++) {
+ String data = String.format("data_%d", i);
+ String ackid = String.format("ackid_%d", i);
+ incoming.add(new IncomingMessage(data.getBytes(), null, TIMESTAMP, 0, ackid, RECORD_ID));
+ }
+ setupOneMessage(incoming);
+ PubsubReader<String> reader = primSource.createReader(p.getOptions(), null);
+ // Consume two messages, only read one.
+ assertTrue(reader.start());
+ assertEquals("data_0", reader.getCurrent());
+
+ // Grab checkpoint.
+ PubsubCheckpoint<String> checkpoint = reader.getCheckpointMark();
+ checkpoint.finalizeCheckpoint();
+ assertEquals(1, checkpoint.notYetReadIds.size());
+ assertEquals("ackid_1", checkpoint.notYetReadIds.get(0));
+
+ // Read second message.
+ assertTrue(reader.advance());
+ assertEquals("data_1", reader.getCurrent());
+
+ // Restore from checkpoint.
+ byte[] checkpointBytes =
+ CoderUtils.encodeToByteArray(primSource.getCheckpointMarkCoder(), checkpoint);
+ checkpoint = CoderUtils.decodeFromByteArray(primSource.getCheckpointMarkCoder(),
+ checkpointBytes);
+ assertEquals(1, checkpoint.notYetReadIds.size());
+ assertEquals("ackid_1", checkpoint.notYetReadIds.get(0));
+
+ // Re-read second message.
+ reader = primSource.createReader(p.getOptions(), checkpoint);
+ assertTrue(reader.start());
+ assertEquals("data_1", reader.getCurrent());
+
+ // We are done.
+ assertFalse(reader.advance());
+
+ // ACK final message.
+ checkpoint = reader.getCheckpointMark();
+ checkpoint.finalizeCheckpoint();
+ reader.close();
+ }
+
+ private long messageNumToTimestamp(int messageNum) {
+ return TIMESTAMP + messageNum * 100;
+ }
+
+ @Test
+ public void readManyMessages() throws IOException {
+ Map<String, Integer> dataToMessageNum = new HashMap<>();
+
+ final int m = 97;
+ final int n = 10000;
+ List<IncomingMessage> incoming = new ArrayList<>();
+ for (int i = 0; i < n; i++) {
+ // Make the messages timestamps slightly out of order.
+ int messageNum = ((i / m) * m) + (m - 1) - (i % m);
+ String data = String.format("data_%d", messageNum);
+ dataToMessageNum.put(data, messageNum);
+ String recid = String.format("recordid_%d", messageNum);
+ String ackId = String.format("ackid_%d", messageNum);
+ incoming.add(new IncomingMessage(data.getBytes(), null, messageNumToTimestamp(messageNum), 0,
+ ackId, recid));
+ }
+ setupOneMessage(incoming);
+
+ PubsubReader<String> reader = primSource.createReader(p.getOptions(), null);
+ PubsubTestClient pubsubClient = (PubsubTestClient) reader.getPubsubClient();
+
+ for (int i = 0; i < n; i++) {
+ if (i == 0) {
+ assertTrue(reader.start());
+ } else {
+ assertTrue(reader.advance());
+ }
+ // We'll checkpoint and ack within the 2min limit.
+ now.addAndGet(30);
+ pubsubClient.advance();
+ String data = reader.getCurrent();
+ Integer messageNum = dataToMessageNum.remove(data);
+ // No duplicate messages.
+ assertNotNull(messageNum);
+ // Preserve timestamp.
+ assertEquals(new Instant(messageNumToTimestamp(messageNum)), reader.getCurrentTimestamp());
+ // Preserve record id.
+ String recid = String.format("recordid_%d", messageNum);
+ assertArrayEquals(recid.getBytes(), reader.getCurrentRecordId());
+
+ if (i % 1000 == 999) {
+ // Estimated watermark can never get ahead of actual outstanding messages.
+ long watermark = reader.getWatermark().getMillis();
+ long minOutstandingTimestamp = Long.MAX_VALUE;
+ for (Integer outstandingMessageNum : dataToMessageNum.values()) {
+ minOutstandingTimestamp =
+ Math.min(minOutstandingTimestamp, messageNumToTimestamp(outstandingMessageNum));
+ }
+ assertThat(watermark, lessThanOrEqualTo(minOutstandingTimestamp));
+ // Ack messages, but only every other finalization.
+ PubsubCheckpoint<String> checkpoint = reader.getCheckpointMark();
+ if (i % 2000 == 1999) {
+ checkpoint.finalizeCheckpoint();
+ }
+ }
+ }
+ // We are done.
+ assertFalse(reader.advance());
+ // We saw each message exactly once.
+ assertTrue(dataToMessageNum.isEmpty());
+ reader.close();
+ }
+
+ @Test
+ public void noSubscriptionSplitIntoBundlesGeneratesSubscription() throws Exception {
+ TopicPath topicPath = PubsubClient.topicPathFromName("my_project", "my_topic");
+ factory = PubsubTestClient.createFactoryForCreateSubscription();
+ PubsubUnboundedSource<String> source =
+ new PubsubUnboundedSource<>(
+ factory,
+ StaticValueProvider.of(PubsubClient.projectPathFromId("my_project")),
+ StaticValueProvider.of(topicPath),
+ null,
+ StringUtf8Coder.of(),
+ null,
+ null,
+ null);
+ assertThat(source.getSubscription(), nullValue());
+
+ assertThat(source.getSubscription(), nullValue());
+
+ PipelineOptions options = PipelineOptionsFactory.create();
+ List<PubsubSource<String>> splits =
+ (new PubsubSource<>(source)).generateInitialSplits(3, options);
+ // We have at least one returned split
+ assertThat(splits, hasSize(greaterThan(0)));
+ for (PubsubSource<String> split : splits) {
+ // Each split is equal
+ assertThat(split, equalTo(splits.get(0)));
+ }
+
+ assertThat(splits.get(0).subscriptionPath, not(nullValue()));
+ }
+
+ @Test
+ public void noSubscriptionNoSplitGeneratesSubscription() throws Exception {
+ TopicPath topicPath = PubsubClient.topicPathFromName("my_project", "my_topic");
+ factory = PubsubTestClient.createFactoryForCreateSubscription();
+ PubsubUnboundedSource<String> source =
+ new PubsubUnboundedSource<>(
+ factory,
+ StaticValueProvider.of(PubsubClient.projectPathFromId("my_project")),
+ StaticValueProvider.of(topicPath),
+ null,
+ StringUtf8Coder.of(),
+ null,
+ null,
+ null);
+ assertThat(source.getSubscription(), nullValue());
+
+ assertThat(source.getSubscription(), nullValue());
+
+ PipelineOptions options = PipelineOptionsFactory.create();
+ PubsubSource<String> actualSource = new PubsubSource<>(source);
+ PubsubReader<String> reader = actualSource.createReader(options, null);
+ SubscriptionPath createdSubscription = reader.subscription;
+ assertThat(createdSubscription, not(nullValue()));
+
+ PubsubCheckpoint<String> checkpoint = reader.getCheckpointMark();
+ assertThat(checkpoint.subscriptionPath, equalTo(createdSubscription.getPath()));
+
+ checkpoint.finalizeCheckpoint();
+ PubsubCheckpoint<String> deserCheckpoint =
+ CoderUtils.clone(actualSource.getCheckpointMarkCoder(), checkpoint);
+ assertThat(checkpoint.subscriptionPath, not(nullValue()));
+ assertThat(checkpoint.subscriptionPath, equalTo(deserCheckpoint.subscriptionPath));
+
+ PubsubReader<String> readerFromOriginal = actualSource.createReader(options, checkpoint);
+ PubsubReader<String> readerFromDeser = actualSource.createReader(options, deserCheckpoint);
+
+ assertThat(readerFromOriginal.subscription, equalTo(createdSubscription));
+ assertThat(readerFromDeser.subscription, equalTo(createdSubscription));
+ }
+
+ /**
+ * Tests that checkpoints finalized after the reader is closed succeed.
+ */
+ @Test
+ public void closeWithActiveCheckpoints() throws Exception {
+ setupOneMessage();
+ PubsubReader<String> reader = primSource.createReader(p.getOptions(), null);
+ reader.start();
+ PubsubCheckpoint<String> checkpoint = reader.getCheckpointMark();
+ reader.close();
+ checkpoint.finalizeCheckpoint();
+ }
+}
[03/50] [abbrv] beam git commit: This closes #2460
Posted by jb...@apache.org.
This closes #2460
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/d255fa25
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/d255fa25
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/d255fa25
Branch: refs/heads/DSL_SQL
Commit: d255fa25f1a1a0bf61ea159e2922f014d36eda43
Parents: 8beea73 c08b7b1
Author: Dan Halperin <dh...@google.com>
Authored: Tue Apr 11 10:20:35 2017 -0700
Committer: Dan Halperin <dh...@google.com>
Committed: Tue Apr 11 10:20:35 2017 -0700
----------------------------------------------------------------------
.../hadoop/inputformat/HadoopInputFormatIO.java | 33 +--
.../inputformat/HadoopInputFormatIOTest.java | 99 ++++++---
sdks/java/io/hadoop/jdk1.8-tests/.toDelete | 0
.../HIFIOWithEmbeddedCassandraTest.java | 214 +++++++++++++++++++
4 files changed, 306 insertions(+), 40 deletions(-)
----------------------------------------------------------------------
[40/50] [abbrv] beam git commit: [BEAM-1882] Update postgres k8
scripts & add scripts for running local dev test
Posted by jb...@apache.org.
[BEAM-1882] Update postgres k8 scripts & add scripts for running local dev test
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/4d1feebf
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/4d1feebf
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/4d1feebf
Branch: refs/heads/DSL_SQL
Commit: 4d1feebfd6931882554cec8f459da85047064982
Parents: 435c3b4
Author: Stephen Sisk <si...@google.com>
Authored: Fri Mar 17 15:09:00 2017 -0700
Committer: Jean-Baptiste Onofr� <jb...@apache.org>
Committed: Wed Apr 12 20:07:08 2017 +0200
----------------------------------------------------------------------
.../kubernetes/postgres-pod-no-vol.yml | 32 -----------
.../postgres-service-for-local-dev.yml | 28 ++++++++++
.../kubernetes/postgres-service-public.yml | 28 ----------
.../src/test/resources/kubernetes/postgres.yml | 56 ++++++++++++++++++++
.../jdbc/src/test/resources/kubernetes/setup.sh | 3 +-
.../src/test/resources/kubernetes/teardown.sh | 3 +-
6 files changed, 86 insertions(+), 64 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/4d1feebf/sdks/java/io/jdbc/src/test/resources/kubernetes/postgres-pod-no-vol.yml
----------------------------------------------------------------------
diff --git a/sdks/java/io/jdbc/src/test/resources/kubernetes/postgres-pod-no-vol.yml b/sdks/java/io/jdbc/src/test/resources/kubernetes/postgres-pod-no-vol.yml
deleted file mode 100644
index 17cad93..0000000
--- a/sdks/java/io/jdbc/src/test/resources/kubernetes/postgres-pod-no-vol.yml
+++ /dev/null
@@ -1,32 +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.
-
-apiVersion: v1
-kind: Pod
-metadata:
- name: postgres-no-pv
- labels:
- name: postgres-no-pv
-spec:
- containers:
- - name: postgres
- image: postgres
- env:
- - name: POSTGRES_PASSWORD
- value: uuinkks
- - name: PGDATA
- value: /var/lib/postgresql/data/pgdata
- ports:
- - containerPort: 5432
http://git-wip-us.apache.org/repos/asf/beam/blob/4d1feebf/sdks/java/io/jdbc/src/test/resources/kubernetes/postgres-service-for-local-dev.yml
----------------------------------------------------------------------
diff --git a/sdks/java/io/jdbc/src/test/resources/kubernetes/postgres-service-for-local-dev.yml b/sdks/java/io/jdbc/src/test/resources/kubernetes/postgres-service-for-local-dev.yml
new file mode 100644
index 0000000..5d2c664
--- /dev/null
+++ b/sdks/java/io/jdbc/src/test/resources/kubernetes/postgres-service-for-local-dev.yml
@@ -0,0 +1,28 @@
+# 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.
+
+
+apiVersion: v1
+kind: Service
+metadata:
+ name: postgres-for-dev
+ labels:
+ name: postgres
+spec:
+ ports:
+ - port: 5432
+ selector:
+ name: postgres
+ type: LoadBalancer
http://git-wip-us.apache.org/repos/asf/beam/blob/4d1feebf/sdks/java/io/jdbc/src/test/resources/kubernetes/postgres-service-public.yml
----------------------------------------------------------------------
diff --git a/sdks/java/io/jdbc/src/test/resources/kubernetes/postgres-service-public.yml b/sdks/java/io/jdbc/src/test/resources/kubernetes/postgres-service-public.yml
deleted file mode 100644
index 24d8eee..0000000
--- a/sdks/java/io/jdbc/src/test/resources/kubernetes/postgres-service-public.yml
+++ /dev/null
@@ -1,28 +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.
-
-apiVersion: v1
-kind: Service
-metadata:
- name: postgres-no-pv
- labels:
- name: postgres-no-pv
-spec:
- ports:
- - port: 5432
- nodePort: 31234
- selector:
- name: postgres-no-pv
- type: NodePort
http://git-wip-us.apache.org/repos/asf/beam/blob/4d1feebf/sdks/java/io/jdbc/src/test/resources/kubernetes/postgres.yml
----------------------------------------------------------------------
diff --git a/sdks/java/io/jdbc/src/test/resources/kubernetes/postgres.yml b/sdks/java/io/jdbc/src/test/resources/kubernetes/postgres.yml
new file mode 100644
index 0000000..6244968
--- /dev/null
+++ b/sdks/java/io/jdbc/src/test/resources/kubernetes/postgres.yml
@@ -0,0 +1,56 @@
+# 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.
+
+
+apiVersion: v1
+kind: Service
+metadata:
+ name: postgres
+ labels:
+ name: postgres
+spec:
+ ports:
+ - port: 5432
+ nodePort: 31234
+ selector:
+ name: postgres
+ type: NodePort
+
+---
+
+apiVersion: v1
+kind: ReplicationController
+metadata:
+ name: postgres
+spec:
+ replicas: 1
+ selector:
+ name: postgres
+ template:
+ metadata:
+ name: postgres
+ labels:
+ name: postgres
+ spec:
+ containers:
+ - name: postgres
+ image: postgres
+ env:
+ - name: POSTGRES_PASSWORD
+ value: uuinkks
+ - name: PGDATA
+ value: /var/lib/postgresql/data/pgdata
+ ports:
+ - containerPort: 5432
http://git-wip-us.apache.org/repos/asf/beam/blob/4d1feebf/sdks/java/io/jdbc/src/test/resources/kubernetes/setup.sh
----------------------------------------------------------------------
diff --git a/sdks/java/io/jdbc/src/test/resources/kubernetes/setup.sh b/sdks/java/io/jdbc/src/test/resources/kubernetes/setup.sh
index 7d9b2e2..76e4e34 100644
--- a/sdks/java/io/jdbc/src/test/resources/kubernetes/setup.sh
+++ b/sdks/java/io/jdbc/src/test/resources/kubernetes/setup.sh
@@ -16,5 +16,4 @@
# limitations under the License.
#
-kubectl create -f postgres-pod-no-vol.yml
-kubectl create -f postgres-service-public.yml
+kubectl create -f ./sdks/java/io/jdbc/src/test/resources/kubernetes/postgres.yml
http://git-wip-us.apache.org/repos/asf/beam/blob/4d1feebf/sdks/java/io/jdbc/src/test/resources/kubernetes/teardown.sh
----------------------------------------------------------------------
diff --git a/sdks/java/io/jdbc/src/test/resources/kubernetes/teardown.sh b/sdks/java/io/jdbc/src/test/resources/kubernetes/teardown.sh
index 451562e..4fca74a 100644
--- a/sdks/java/io/jdbc/src/test/resources/kubernetes/teardown.sh
+++ b/sdks/java/io/jdbc/src/test/resources/kubernetes/teardown.sh
@@ -16,5 +16,4 @@
# limitations under the License.
#
-kubectl delete service postgres-no-pv
-kubectl delete pod postgres-no-pv
+kubectl delete -f ./sdks/java/io/jdbc/src/test/resources/kubernetes/postgres.yml
[25/50] [abbrv] beam git commit: This closes #2499
Posted by jb...@apache.org.
This closes #2499
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/ccf247c1
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/ccf247c1
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/ccf247c1
Branch: refs/heads/DSL_SQL
Commit: ccf247c1daa868cdaf441f918775c94ae6e3aa0b
Parents: a2d328a 041b636
Author: Thomas Groh <tg...@google.com>
Authored: Tue Apr 11 16:17:39 2017 -0700
Committer: Thomas Groh <tg...@google.com>
Committed: Tue Apr 11 16:17:39 2017 -0700
----------------------------------------------------------------------
.../sdk/io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java | 1 +
1 file changed, 1 insertion(+)
----------------------------------------------------------------------
[46/50] [abbrv] beam git commit: Revert "Revert "Revert "Add
ValueProvider class for FileBasedSource I/O Transforms"""
Posted by jb...@apache.org.
Revert "Revert "Revert "Add ValueProvider class for FileBasedSource I/O Transforms"""
This reverts commit 28a0ea86721871a97f68b573019e818217266356.
Manually resolved Conflicts:
sdks/python/apache_beam/io/filebasedsource.py
sdks/python/apache_beam/runners/direct/direct_runner.py
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/0f5c3639
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/0f5c3639
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/0f5c3639
Branch: refs/heads/DSL_SQL
Commit: 0f5c36398aa321d25b3a045cedc29203d75956f7
Parents: 4854291
Author: Ahmet Altay <al...@google.com>
Authored: Wed Apr 12 10:48:25 2017 -0700
Committer: Ahmet Altay <al...@google.com>
Committed: Wed Apr 12 11:50:39 2017 -0700
----------------------------------------------------------------------
sdks/python/apache_beam/examples/wordcount.py | 32 ++--
.../apache_beam/internal/gcp/json_value.py | 6 -
sdks/python/apache_beam/io/filebasedsource.py | 54 ++----
.../apache_beam/io/filebasedsource_test.py | 24 ---
sdks/python/apache_beam/io/fileio.py | 56 ++-----
sdks/python/apache_beam/io/fileio_test.py | 45 ++---
.../runners/dataflow/internal/apiclient.py | 1 -
.../apache_beam/runners/direct/direct_runner.py | 9 -
sdks/python/apache_beam/transforms/display.py | 1 -
.../apache_beam/transforms/display_test.py | 36 ----
.../apache_beam/utils/pipeline_options.py | 92 +----------
.../apache_beam/utils/pipeline_options_test.py | 52 +-----
sdks/python/apache_beam/utils/value_provider.py | 110 -------------
.../apache_beam/utils/value_provider_test.py | 165 -------------------
14 files changed, 65 insertions(+), 618 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/0f5c3639/sdks/python/apache_beam/examples/wordcount.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/examples/wordcount.py b/sdks/python/apache_beam/examples/wordcount.py
index a155148..e4792cb 100644
--- a/sdks/python/apache_beam/examples/wordcount.py
+++ b/sdks/python/apache_beam/examples/wordcount.py
@@ -19,6 +19,7 @@
from __future__ import absolute_import
+import argparse
import logging
import re
@@ -66,29 +67,24 @@ class WordExtractingDoFn(beam.DoFn):
def run(argv=None):
"""Main entry point; defines and runs the wordcount pipeline."""
- class WordcountOptions(PipelineOptions):
- @classmethod
- def _add_argparse_args(cls, parser):
- parser.add_value_provider_argument(
- '--input',
- dest='input',
- default='gs://dataflow-samples/shakespeare/kinglear.txt',
- help='Input file to process.')
- parser.add_value_provider_argument(
- '--output',
- dest='output',
- required=True,
- help='Output file to write results to.')
- pipeline_options = PipelineOptions(argv)
- wordcount_options = pipeline_options.view_as(WordcountOptions)
-
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--input',
+ dest='input',
+ default='gs://dataflow-samples/shakespeare/kinglear.txt',
+ help='Input file to process.')
+ parser.add_argument('--output',
+ dest='output',
+ required=True,
+ help='Output file to write results to.')
+ known_args, pipeline_args = parser.parse_known_args(argv)
# We use the save_main_session option because one or more DoFn's in this
# workflow rely on global context (e.g., a module imported at module level).
+ pipeline_options = PipelineOptions(pipeline_args)
pipeline_options.view_as(SetupOptions).save_main_session = True
p = beam.Pipeline(options=pipeline_options)
# Read the text file[pattern] into a PCollection.
- lines = p | 'read' >> ReadFromText(wordcount_options.input)
+ lines = p | 'read' >> ReadFromText(known_args.input)
# Count the occurrences of each word.
counts = (lines
@@ -103,7 +99,7 @@ def run(argv=None):
# Write the output using a "Write" transform that has side effects.
# pylint: disable=expression-not-assigned
- output | 'write' >> WriteToText(wordcount_options.output)
+ output | 'write' >> WriteToText(known_args.output)
# Actually run the pipeline (all operations above are deferred).
result = p.run()
http://git-wip-us.apache.org/repos/asf/beam/blob/0f5c3639/sdks/python/apache_beam/internal/gcp/json_value.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/internal/gcp/json_value.py b/sdks/python/apache_beam/internal/gcp/json_value.py
index 4099c1a..c8b5393 100644
--- a/sdks/python/apache_beam/internal/gcp/json_value.py
+++ b/sdks/python/apache_beam/internal/gcp/json_value.py
@@ -25,8 +25,6 @@ except ImportError:
extra_types = None
# pylint: enable=wrong-import-order, wrong-import-position
-from apache_beam.utils.value_provider import ValueProvider
-
_MAXINT64 = (1 << 63) - 1
_MININT64 = - (1 << 63)
@@ -106,10 +104,6 @@ def to_json_value(obj, with_type=False):
raise TypeError('Can not encode {} as a 64-bit integer'.format(obj))
elif isinstance(obj, float):
return extra_types.JsonValue(double_value=obj)
- elif isinstance(obj, ValueProvider):
- if obj.is_accessible():
- return to_json_value(obj.get())
- return extra_types.JsonValue(is_null=True)
else:
raise TypeError('Cannot convert %s to a JSON value.' % repr(obj))
http://git-wip-us.apache.org/repos/asf/beam/blob/0f5c3639/sdks/python/apache_beam/io/filebasedsource.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/io/filebasedsource.py b/sdks/python/apache_beam/io/filebasedsource.py
index 2e7043f..3a53886 100644
--- a/sdks/python/apache_beam/io/filebasedsource.py
+++ b/sdks/python/apache_beam/io/filebasedsource.py
@@ -32,9 +32,6 @@ from apache_beam.io import range_trackers
from apache_beam.io.filesystem import CompressionTypes
from apache_beam.io.filesystems_util import get_filesystem
from apache_beam.transforms.display import DisplayDataItem
-from apache_beam.utils.value_provider import ValueProvider
-from apache_beam.utils.value_provider import StaticValueProvider
-from apache_beam.utils.value_provider import check_accessible
MAX_NUM_THREADS_FOR_SIZE_ESTIMATION = 25
@@ -54,8 +51,7 @@ class FileBasedSource(iobase.BoundedSource):
"""Initializes ``FileBasedSource``.
Args:
- file_pattern: the file glob to read a string or a ValueProvider
- (placeholder to inject a runtime value).
+ file_pattern: the file glob to read.
min_bundle_size: minimum size of bundles that should be generated when
performing initial splitting on this source.
compression_type: compression type to use
@@ -73,25 +69,17 @@ class FileBasedSource(iobase.BoundedSource):
creation time.
Raises:
TypeError: when compression_type is not valid or if file_pattern is not a
- string or a ValueProvider.
+ string.
ValueError: when compression and splittable files are specified.
IOError: when the file pattern specified yields an empty result.
"""
+ if not isinstance(file_pattern, basestring):
+ raise TypeError(
+ '%s: file_pattern must be a string; got %r instead' %
+ (self.__class__.__name__, file_pattern))
- if (not (isinstance(file_pattern, basestring)
- or isinstance(file_pattern, ValueProvider))):
- raise TypeError('%s: file_pattern must be of type string'
- ' or ValueProvider; got %r instead'
- % (self.__class__.__name__, file_pattern))
-
- if isinstance(file_pattern, basestring):
- file_pattern = StaticValueProvider(str, file_pattern)
self._pattern = file_pattern
- if file_pattern.is_accessible():
- self._file_system = get_filesystem(file_pattern.get())
- else:
- self._file_system = None
-
+ self._file_system = get_filesystem(file_pattern)
self._concat_source = None
self._min_bundle_size = min_bundle_size
if not CompressionTypes.is_valid_compression_type(compression_type):
@@ -104,24 +92,19 @@ class FileBasedSource(iobase.BoundedSource):
else:
# We can't split compressed files efficiently so turn off splitting.
self._splittable = False
- if validate and file_pattern.is_accessible():
+ if validate:
self._validate()
def display_data(self):
- return {'file_pattern': DisplayDataItem(str(self._pattern),
+ return {'file_pattern': DisplayDataItem(self._pattern,
label="File Pattern"),
'compression': DisplayDataItem(str(self._compression_type),
label='Compression Type')}
- @check_accessible(['_pattern'])
def _get_concat_source(self):
if self._concat_source is None:
- pattern = self._pattern.get()
-
single_file_sources = []
- if self._file_system is None:
- self._file_system = get_filesystem(pattern)
- match_result = self._file_system.match([pattern])[0]
+ match_result = self._file_system.match([self._pattern])[0]
files_metadata = match_result.metadata_list
# We create a reference for FileBasedSource that will be serialized along
@@ -160,19 +143,14 @@ class FileBasedSource(iobase.BoundedSource):
file_name, 'application/octet-stream',
compression_type=self._compression_type)
- @check_accessible(['_pattern'])
def _validate(self):
"""Validate if there are actual files in the specified glob pattern
"""
- pattern = self._pattern.get()
- if self._file_system is None:
- self._file_system = get_filesystem(pattern)
-
# Limit the responses as we only want to check if something exists
- match_result = self._file_system.match([pattern], limits=[1])[0]
+ match_result = self._file_system.match([self._pattern], limits=[1])[0]
if len(match_result.metadata_list) <= 0:
raise IOError(
- 'No files found based on the file pattern %s' % pattern)
+ 'No files found based on the file pattern %s' % self._pattern)
def split(
self, desired_bundle_size=None, start_position=None, stop_position=None):
@@ -181,12 +159,8 @@ class FileBasedSource(iobase.BoundedSource):
start_position=start_position,
stop_position=stop_position)
- @check_accessible(['_pattern'])
def estimate_size(self):
- pattern = self._pattern.get()
- if self._file_system is None:
- self._file_system = get_filesystem(pattern)
- match_result = self._file_system.match([pattern])[0]
+ match_result = self._file_system.match([self._pattern])[0]
return sum([f.size_in_bytes for f in match_result.metadata_list])
def read(self, range_tracker):
@@ -211,7 +185,7 @@ class FileBasedSource(iobase.BoundedSource):
defined by a given ``RangeTracker``.
Returns:
- an iterator that gives the records read from the given file.
+ a iterator that gives the records read from the given file.
"""
raise NotImplementedError
http://git-wip-us.apache.org/repos/asf/beam/blob/0f5c3639/sdks/python/apache_beam/io/filebasedsource_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/io/filebasedsource_test.py b/sdks/python/apache_beam/io/filebasedsource_test.py
index c25ca5d..7b7ec8a 100644
--- a/sdks/python/apache_beam/io/filebasedsource_test.py
+++ b/sdks/python/apache_beam/io/filebasedsource_test.py
@@ -43,8 +43,6 @@ from apache_beam.transforms.display import DisplayData
from apache_beam.transforms.display_test import DisplayDataItemMatcher
from apache_beam.transforms.util import assert_that
from apache_beam.transforms.util import equal_to
-from apache_beam.utils.value_provider import StaticValueProvider
-from apache_beam.utils.value_provider import RuntimeValueProvider
class LineSource(FileBasedSource):
@@ -223,28 +221,6 @@ class TestFileBasedSource(unittest.TestCase):
# environments with limited amount of resources.
filebasedsource.MAX_NUM_THREADS_FOR_SIZE_ESTIMATION = 2
- def test_string_or_value_provider_only(self):
- str_file_pattern = tempfile.NamedTemporaryFile(delete=False).name
- self.assertEqual(str_file_pattern,
- FileBasedSource(str_file_pattern)._pattern.value)
-
- static_vp_file_pattern = StaticValueProvider(value_type=str,
- value=str_file_pattern)
- self.assertEqual(static_vp_file_pattern,
- FileBasedSource(static_vp_file_pattern)._pattern)
-
- runtime_vp_file_pattern = RuntimeValueProvider(
- option_name='arg',
- value_type=str,
- default_value=str_file_pattern,
- options_id=1)
- self.assertEqual(runtime_vp_file_pattern,
- FileBasedSource(runtime_vp_file_pattern)._pattern)
-
- invalid_file_pattern = 123
- with self.assertRaises(TypeError):
- FileBasedSource(invalid_file_pattern)
-
def test_validation_file_exists(self):
file_name, _ = write_data(10)
LineSource(file_name)
http://git-wip-us.apache.org/repos/asf/beam/blob/0f5c3639/sdks/python/apache_beam/io/fileio.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/io/fileio.py b/sdks/python/apache_beam/io/fileio.py
index 84949dc..f33942a 100644
--- a/sdks/python/apache_beam/io/fileio.py
+++ b/sdks/python/apache_beam/io/fileio.py
@@ -14,7 +14,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
-
"""File-based sources and sinks."""
from __future__ import absolute_import
@@ -31,9 +30,6 @@ from apache_beam.io.filesystem import CompressedFile as _CompressedFile
from apache_beam.io.filesystem import CompressionTypes
from apache_beam.io.filesystems_util import get_filesystem
from apache_beam.transforms.display import DisplayDataItem
-from apache_beam.utils.value_provider import ValueProvider
-from apache_beam.utils.value_provider import StaticValueProvider
-from apache_beam.utils.value_provider import check_accessible
MAX_BATCH_OPERATION_SIZE = 100
DEFAULT_SHARD_NAME_TEMPLATE = '-SSSSS-of-NNNNN'
@@ -153,30 +149,25 @@ class FileSink(iobase.Sink):
compression_type=CompressionTypes.AUTO):
"""
Raises:
- TypeError: if file path parameters are not a string or ValueProvider,
- or if compression_type is not member of CompressionTypes.
+ TypeError: if file path parameters are not a string or if compression_type
+ is not member of CompressionTypes.
ValueError: if shard_name_template is not of expected format.
"""
- if not (isinstance(file_path_prefix, basestring)
- or isinstance(file_path_prefix, ValueProvider)):
- raise TypeError('file_path_prefix must be a string or ValueProvider;'
- 'got %r instead' % file_path_prefix)
- if not (isinstance(file_name_suffix, basestring)
- or isinstance(file_name_suffix, ValueProvider)):
- raise TypeError('file_name_suffix must be a string or ValueProvider;'
- 'got %r instead' % file_name_suffix)
+ if not isinstance(file_path_prefix, basestring):
+ raise TypeError('file_path_prefix must be a string; got %r instead' %
+ file_path_prefix)
+ if not isinstance(file_name_suffix, basestring):
+ raise TypeError('file_name_suffix must be a string; got %r instead' %
+ file_name_suffix)
if not CompressionTypes.is_valid_compression_type(compression_type):
raise TypeError('compression_type must be CompressionType object but '
'was %s' % type(compression_type))
+
if shard_name_template is None:
shard_name_template = DEFAULT_SHARD_NAME_TEMPLATE
elif shard_name_template is '':
num_shards = 1
- if isinstance(file_path_prefix, basestring):
- file_path_prefix = StaticValueProvider(str, file_path_prefix)
- if isinstance(file_name_suffix, basestring):
- file_name_suffix = StaticValueProvider(str, file_name_suffix)
self.file_path_prefix = file_path_prefix
self.file_name_suffix = file_name_suffix
self.num_shards = num_shards
@@ -184,10 +175,7 @@ class FileSink(iobase.Sink):
self.shard_name_format = self._template_to_format(shard_name_template)
self.compression_type = compression_type
self.mime_type = mime_type
- if file_path_prefix.is_accessible():
- self._file_system = get_filesystem(file_path_prefix.get())
- else:
- self._file_system = None
+ self._file_system = get_filesystem(file_path_prefix)
def display_data(self):
return {'shards':
@@ -201,15 +189,12 @@ class FileSink(iobase.Sink):
self.file_name_suffix),
label='File Pattern')}
- @check_accessible(['file_path_prefix'])
def open(self, temp_path):
"""Opens ``temp_path``, returning an opaque file handle object.
The returned file handle is passed to ``write_[encoded_]record`` and
``close``.
"""
- if self._file_system is None:
- self._file_system = get_filesystem(self.file_path_prefix.get())
return self._file_system.create(temp_path, self.mime_type,
self.compression_type)
@@ -236,33 +221,22 @@ class FileSink(iobase.Sink):
if file_handle is not None:
file_handle.close()
- @check_accessible(['file_path_prefix', 'file_name_suffix'])
def initialize_write(self):
- file_path_prefix = self.file_path_prefix.get()
- file_name_suffix = self.file_name_suffix.get()
- tmp_dir = file_path_prefix + file_name_suffix + time.strftime(
+ tmp_dir = self.file_path_prefix + self.file_name_suffix + time.strftime(
'-temp-%Y-%m-%d_%H-%M-%S')
- if self._file_system is None:
- self._file_system = get_filesystem(file_path_prefix)
self._file_system.mkdirs(tmp_dir)
return tmp_dir
- @check_accessible(['file_path_prefix', 'file_name_suffix'])
def open_writer(self, init_result, uid):
# A proper suffix is needed for AUTO compression detection.
# We also ensure there will be no collisions with uid and a
# (possibly unsharded) file_path_prefix and a (possibly empty)
# file_name_suffix.
- file_path_prefix = self.file_path_prefix.get()
- file_name_suffix = self.file_name_suffix.get()
suffix = (
- '.' + os.path.basename(file_path_prefix) + file_name_suffix)
+ '.' + os.path.basename(self.file_path_prefix) + self.file_name_suffix)
return FileSinkWriter(self, os.path.join(init_result, uid) + suffix)
- @check_accessible(['file_path_prefix', 'file_name_suffix'])
def finalize_write(self, init_result, writer_results):
- file_path_prefix = self.file_path_prefix.get()
- file_name_suffix = self.file_name_suffix.get()
writer_results = sorted(writer_results)
num_shards = len(writer_results)
min_threads = min(num_shards, FileSink._MAX_RENAME_THREADS)
@@ -272,8 +246,8 @@ class FileSink(iobase.Sink):
destination_files = []
for shard_num, shard in enumerate(writer_results):
final_name = ''.join([
- file_path_prefix, self.shard_name_format % dict(
- shard_num=shard_num, num_shards=num_shards), file_name_suffix
+ self.file_path_prefix, self.shard_name_format % dict(
+ shard_num=shard_num, num_shards=num_shards), self.file_name_suffix
])
source_files.append(shard)
destination_files.append(final_name)
@@ -296,8 +270,6 @@ class FileSink(iobase.Sink):
"""_rename_batch executes batch rename operations."""
source_files, destination_files = batch
exceptions = []
- if self._file_system is None:
- self._file_system = get_filesystem(file_path_prefix)
try:
self._file_system.rename(source_files, destination_files)
return exceptions
http://git-wip-us.apache.org/repos/asf/beam/blob/0f5c3639/sdks/python/apache_beam/io/fileio_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/io/fileio_test.py b/sdks/python/apache_beam/io/fileio_test.py
index 2409873..3b4f61f 100644
--- a/sdks/python/apache_beam/io/fileio_test.py
+++ b/sdks/python/apache_beam/io/fileio_test.py
@@ -34,8 +34,6 @@ from apache_beam.test_pipeline import TestPipeline
from apache_beam.transforms.display import DisplayData
from apache_beam.transforms.display_test import DisplayDataItemMatcher
-from apache_beam.utils.value_provider import StaticValueProvider
-
# TODO: Refactor code so all io tests are using same library
# TestCaseWithTempDirCleanup class.
@@ -95,7 +93,7 @@ class TestFileSink(_TestCaseWithTempDirCleanUp):
def test_file_sink_writing(self):
temp_path = os.path.join(self._new_tempdir(), 'filesink')
sink = MyFileSink(
- temp_path, file_name_suffix='.output', coder=coders.ToStringCoder())
+ temp_path, file_name_suffix='.foo', coder=coders.ToStringCoder())
# Manually invoke the generic Sink API.
init_token = sink.initialize_write()
@@ -116,8 +114,8 @@ class TestFileSink(_TestCaseWithTempDirCleanUp):
res = list(sink.finalize_write(init_token, [res1, res2]))
# Check the results.
- shard1 = temp_path + '-00000-of-00002.output'
- shard2 = temp_path + '-00001-of-00002.output'
+ shard1 = temp_path + '-00000-of-00002.foo'
+ shard2 = temp_path + '-00001-of-00002.foo'
self.assertEqual(res, [shard1, shard2])
self.assertEqual(open(shard1).read(), '[start][a][b][end]')
self.assertEqual(open(shard2).read(), '[start][x][y][z][end]')
@@ -128,48 +126,33 @@ class TestFileSink(_TestCaseWithTempDirCleanUp):
def test_file_sink_display_data(self):
temp_path = os.path.join(self._new_tempdir(), 'display')
sink = MyFileSink(
- temp_path, file_name_suffix='.output', coder=coders.ToStringCoder())
+ temp_path, file_name_suffix='.foo', coder=coders.ToStringCoder())
dd = DisplayData.create_from(sink)
expected_items = [
DisplayDataItemMatcher(
'compression', 'auto'),
DisplayDataItemMatcher(
'file_pattern',
- '{}{}'.format(
- temp_path,
- '-%(shard_num)05d-of-%(num_shards)05d.output'))]
+ '{}{}'.format(temp_path,
+ '-%(shard_num)05d-of-%(num_shards)05d.foo'))]
+
hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
def test_empty_write(self):
temp_path = tempfile.NamedTemporaryFile().name
sink = MyFileSink(
- temp_path, file_name_suffix='.output', coder=coders.ToStringCoder()
- )
- p = TestPipeline()
- p | beam.Create([]) | beam.io.Write(sink) # pylint: disable=expression-not-assigned
- p.run()
- self.assertEqual(
- open(temp_path + '-00000-of-00001.output').read(), '[start][end]')
-
- def test_static_value_provider_empty_write(self):
- temp_path = StaticValueProvider(value_type=str,
- value=tempfile.NamedTemporaryFile().name)
- sink = MyFileSink(
- temp_path,
- file_name_suffix=StaticValueProvider(value_type=str, value='.output'),
- coder=coders.ToStringCoder()
- )
+ temp_path, file_name_suffix='.foo', coder=coders.ToStringCoder())
p = TestPipeline()
p | beam.Create([]) | beam.io.Write(sink) # pylint: disable=expression-not-assigned
p.run()
self.assertEqual(
- open(temp_path.get() + '-00000-of-00001.output').read(), '[start][end]')
+ open(temp_path + '-00000-of-00001.foo').read(), '[start][end]')
def test_fixed_shard_write(self):
temp_path = os.path.join(self._new_tempdir(), 'empty')
sink = MyFileSink(
temp_path,
- file_name_suffix='.output',
+ file_name_suffix='.foo',
num_shards=3,
shard_name_template='_NN_SSS_',
coder=coders.ToStringCoder())
@@ -179,7 +162,7 @@ class TestFileSink(_TestCaseWithTempDirCleanUp):
p.run()
concat = ''.join(
- open(temp_path + '_03_%03d_.output' % shard_num).read()
+ open(temp_path + '_03_%03d_.foo' % shard_num).read()
for shard_num in range(3))
self.assertTrue('][a][' in concat, concat)
self.assertTrue('][b][' in concat, concat)
@@ -187,7 +170,7 @@ class TestFileSink(_TestCaseWithTempDirCleanUp):
def test_file_sink_multi_shards(self):
temp_path = os.path.join(self._new_tempdir(), 'multishard')
sink = MyFileSink(
- temp_path, file_name_suffix='.output', coder=coders.ToStringCoder())
+ temp_path, file_name_suffix='.foo', coder=coders.ToStringCoder())
# Manually invoke the generic Sink API.
init_token = sink.initialize_write()
@@ -210,7 +193,7 @@ class TestFileSink(_TestCaseWithTempDirCleanUp):
res = sorted(res_second)
for i in range(num_shards):
- shard_name = '%s-%05d-of-%05d.output' % (temp_path, i, num_shards)
+ shard_name = '%s-%05d-of-%05d.foo' % (temp_path, i, num_shards)
uuid = 'uuid-%05d' % i
self.assertEqual(res[i], shard_name)
self.assertEqual(
@@ -222,7 +205,7 @@ class TestFileSink(_TestCaseWithTempDirCleanUp):
def test_file_sink_io_error(self):
temp_path = os.path.join(self._new_tempdir(), 'ioerror')
sink = MyFileSink(
- temp_path, file_name_suffix='.output', coder=coders.ToStringCoder())
+ temp_path, file_name_suffix='.foo', coder=coders.ToStringCoder())
# Manually invoke the generic Sink API.
init_token = sink.initialize_write()
http://git-wip-us.apache.org/repos/asf/beam/blob/0f5c3639/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 2b6f3fd..6a8aa93 100644
--- a/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py
+++ b/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py
@@ -233,7 +233,6 @@ class Environment(object):
options_dict = {k: v
for k, v in sdk_pipeline_options.iteritems()
if v is not None}
- options_dict['_options_id'] = options._options_id
self.proto.sdkPipelineOptions.additionalProperties.append(
dataflow.Environment.SdkPipelineOptionsValue.AdditionalProperty(
key='options', value=to_json_value(options_dict)))
http://git-wip-us.apache.org/repos/asf/beam/blob/0f5c3639/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 9b4e1ac..357e2d9 100644
--- a/sdks/python/apache_beam/runners/direct/direct_runner.py
+++ b/sdks/python/apache_beam/runners/direct/direct_runner.py
@@ -34,7 +34,6 @@ from apache_beam.runners.runner import PipelineState
from apache_beam.runners.runner import PValueCache
from apache_beam.runners.runner import group_by_key_input_visitor
from apache_beam.utils.pipeline_options import DirectOptions
-from apache_beam.utils.value_provider import RuntimeValueProvider
class DirectRunner(PipelineRunner):
@@ -89,9 +88,6 @@ class DirectRunner(PipelineRunner):
evaluation_context)
# Start the executor. This is a non-blocking call, it will start the
# execution in background threads and return.
-
- if pipeline.options:
- RuntimeValueProvider.set_runtime_options(pipeline.options._options_id, {})
executor.start(self.consumer_tracking_visitor.root_transforms)
result = DirectPipelineResult(executor, evaluation_context)
@@ -101,11 +97,6 @@ class DirectRunner(PipelineRunner):
result.wait_until_finish()
self._cache.finalize()
- # Unset runtime options after the pipeline finishes.
- # TODO: Move this to a post finish hook and clean for all cases.
- if pipeline.options:
- RuntimeValueProvider.unset_runtime_options(pipeline.options._options_id)
-
return result
@property
http://git-wip-us.apache.org/repos/asf/beam/blob/0f5c3639/sdks/python/apache_beam/transforms/display.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/transforms/display.py b/sdks/python/apache_beam/transforms/display.py
index f2ce0fc..2ced1af 100644
--- a/sdks/python/apache_beam/transforms/display.py
+++ b/sdks/python/apache_beam/transforms/display.py
@@ -40,7 +40,6 @@ from datetime import datetime, timedelta
import inspect
import json
-
__all__ = ['HasDisplayData', 'DisplayDataItem', 'DisplayData']
http://git-wip-us.apache.org/repos/asf/beam/blob/0f5c3639/sdks/python/apache_beam/transforms/display_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/transforms/display_test.py b/sdks/python/apache_beam/transforms/display_test.py
index 7d1130b..5e106e5 100644
--- a/sdks/python/apache_beam/transforms/display_test.py
+++ b/sdks/python/apache_beam/transforms/display_test.py
@@ -114,42 +114,6 @@ class DisplayDataTest(unittest.TestCase):
with self.assertRaises(ValueError):
DisplayData.create_from_options(MyDisplayComponent())
- def test_value_provider_display_data(self):
- class TestOptions(PipelineOptions):
- @classmethod
- def _add_argparse_args(cls, parser):
- parser.add_value_provider_argument(
- '--int_flag',
- type=int,
- help='int_flag description')
- parser.add_value_provider_argument(
- '--str_flag',
- type=str,
- default='hello',
- help='str_flag description')
- parser.add_value_provider_argument(
- '--float_flag',
- type=float,
- help='float_flag description')
- options = TestOptions(['--int_flag', '1'])
- items = DisplayData.create_from_options(options).items
- expected_items = [
- DisplayDataItemMatcher(
- 'int_flag',
- '1'),
- DisplayDataItemMatcher(
- 'str_flag',
- 'RuntimeValueProvider(option: str_flag,'
- ' type: str, default_value: \'hello\')'
- ),
- DisplayDataItemMatcher(
- 'float_flag',
- 'RuntimeValueProvider(option: float_flag,'
- ' type: float, default_value: None)'
- )
- ]
- hc.assert_that(items, hc.contains_inanyorder(*expected_items))
-
def test_create_list_display_data(self):
flags = ['--extra_package', 'package1', '--extra_package', 'package2']
pipeline_options = PipelineOptions(flags=flags)
http://git-wip-us.apache.org/repos/asf/beam/blob/0f5c3639/sdks/python/apache_beam/utils/pipeline_options.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/utils/pipeline_options.py b/sdks/python/apache_beam/utils/pipeline_options.py
index 769beb3..c2a44ad 100644
--- a/sdks/python/apache_beam/utils/pipeline_options.py
+++ b/sdks/python/apache_beam/utils/pipeline_options.py
@@ -18,81 +18,8 @@
"""Pipeline options obtained from command line parsing."""
import argparse
-import itertools
from apache_beam.transforms.display import HasDisplayData
-from apache_beam.utils.value_provider import StaticValueProvider
-from apache_beam.utils.value_provider import RuntimeValueProvider
-from apache_beam.utils.value_provider import ValueProvider
-
-
-def _static_value_provider_of(value_type):
- """"Helper function to plug a ValueProvider into argparse.
-
- Args:
- value_type: the type of the value. Since the type param of argparse's
- add_argument will always be ValueProvider, we need to
- preserve the type of the actual value.
- Returns:
- A partially constructed StaticValueProvider in the form of a function.
-
- """
- def _f(value):
- _f.func_name = value_type.__name__
- return StaticValueProvider(value_type, value)
- return _f
-
-
-class BeamArgumentParser(argparse.ArgumentParser):
- """An ArgumentParser that supports ValueProvider options.
-
- Example Usage::
-
- class TemplateUserOptions(PipelineOptions):
- @classmethod
-
- def _add_argparse_args(cls, parser):
- parser.add_value_provider_argument('--vp-arg1', default='start')
- parser.add_value_provider_argument('--vp-arg2')
- parser.add_argument('--non-vp-arg')
-
- """
- def __init__(self, options_id, *args, **kwargs):
- self._options_id = options_id
- super(BeamArgumentParser, self).__init__(*args, **kwargs)
-
- def add_value_provider_argument(self, *args, **kwargs):
- """ValueProvider arguments can be either of type keyword or positional.
- At runtime, even positional arguments will need to be supplied in the
- key/value form.
- """
- # Extract the option name from positional argument ['pos_arg']
- assert args != () and len(args[0]) >= 1
- if args[0][0] != '-':
- option_name = args[0]
- if kwargs.get('nargs') is None: # make them optionally templated
- kwargs['nargs'] = '?'
- else:
- # or keyword arguments like [--kw_arg, -k, -w] or [--kw-arg]
- option_name = [i.replace('--', '') for i in args if i[:2] == '--'][0]
-
- # reassign the type to make room for using
- # StaticValueProvider as the type for add_argument
- value_type = kwargs.get('type') or str
- kwargs['type'] = _static_value_provider_of(value_type)
-
- # reassign default to default_value to make room for using
- # RuntimeValueProvider as the default for add_argument
- default_value = kwargs.get('default')
- kwargs['default'] = RuntimeValueProvider(
- option_name=option_name,
- value_type=value_type,
- default_value=default_value,
- options_id=self._options_id
- )
-
- # have add_argument do most of the work
- self.add_argument(*args, **kwargs)
class PipelineOptions(HasDisplayData):
@@ -122,9 +49,8 @@ class PipelineOptions(HasDisplayData):
By default the options classes will use command line arguments to initialize
the options.
"""
- _options_id_generator = itertools.count(1)
- def __init__(self, flags=None, options_id=None, **kwargs):
+ def __init__(self, flags=None, **kwargs):
"""Initialize an options class.
The initializer will traverse all subclasses, add all their argparse
@@ -141,10 +67,7 @@ class PipelineOptions(HasDisplayData):
"""
self._flags = flags
self._all_options = kwargs
- self._options_id = (
- options_id or PipelineOptions._options_id_generator.next())
- parser = BeamArgumentParser(self._options_id)
-
+ parser = argparse.ArgumentParser()
for cls in type(self).mro():
if cls == PipelineOptions:
break
@@ -196,12 +119,13 @@ class PipelineOptions(HasDisplayData):
# TODO(BEAM-1319): PipelineOption sub-classes in the main session might be
# repeated. Pick last unique instance of each subclass to avoid conflicts.
+ parser = argparse.ArgumentParser()
subset = {}
- parser = BeamArgumentParser(self._options_id)
for cls in PipelineOptions.__subclasses__():
subset[str(cls)] = cls
for cls in subset.values():
cls._add_argparse_args(parser) # pylint: disable=protected-access
+
known_args, _ = parser.parse_known_args(self._flags)
result = vars(known_args)
@@ -209,9 +133,7 @@ class PipelineOptions(HasDisplayData):
for k in result.keys():
if k in self._all_options:
result[k] = self._all_options[k]
- if (drop_default and
- parser.get_default(k) == result[k] and
- not isinstance(parser.get_default(k), ValueProvider)):
+ if drop_default and parser.get_default(k) == result[k]:
del result[k]
return result
@@ -220,7 +142,7 @@ class PipelineOptions(HasDisplayData):
return self.get_all_options(True)
def view_as(self, cls):
- view = cls(self._flags, options_id=self._options_id)
+ view = cls(self._flags)
view._all_options = self._all_options
return view
@@ -244,7 +166,7 @@ class PipelineOptions(HasDisplayData):
(type(self).__name__, name))
def __setattr__(self, name, value):
- if name in ('_flags', '_all_options', '_visible_options', '_options_id'):
+ if name in ('_flags', '_all_options', '_visible_options'):
super(PipelineOptions, self).__setattr__(name, value)
elif name in self._visible_option_list():
self._all_options[name] = value
http://git-wip-us.apache.org/repos/asf/beam/blob/0f5c3639/sdks/python/apache_beam/utils/pipeline_options_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/utils/pipeline_options_test.py b/sdks/python/apache_beam/utils/pipeline_options_test.py
index 633d7da..507a827 100644
--- a/sdks/python/apache_beam/utils/pipeline_options_test.py
+++ b/sdks/python/apache_beam/utils/pipeline_options_test.py
@@ -24,13 +24,9 @@ import hamcrest as hc
from apache_beam.transforms.display import DisplayData
from apache_beam.transforms.display_test import DisplayDataItemMatcher
from apache_beam.utils.pipeline_options import PipelineOptions
-from apache_beam.utils.value_provider import StaticValueProvider
-from apache_beam.utils.value_provider import RuntimeValueProvider
class PipelineOptionsTest(unittest.TestCase):
- def setUp(self):
- RuntimeValueProvider.runtime_options_map = {}
TEST_CASES = [
{'flags': ['--num_workers', '5'],
@@ -135,7 +131,7 @@ class PipelineOptionsTest(unittest.TestCase):
options.view_as(PipelineOptionsTest.MockOptions).mock_flag = True
self.assertEqual(options.get_all_options()['num_workers'], 5)
- self.assertTrue(options.get_all_options()['mock_flag'])
+ self.assertEqual(options.get_all_options()['mock_flag'], True)
def test_experiments(self):
options = PipelineOptions(['--experiment', 'abc', '--experiment', 'def'])
@@ -189,51 +185,7 @@ class PipelineOptionsTest(unittest.TestCase):
parser.add_argument('--redefined_flag', action='store_true')
options = PipelineOptions(['--redefined_flag'])
- self.assertTrue(options.get_all_options()['redefined_flag'])
-
- def test_value_provider_options(self):
- class UserOptions(PipelineOptions):
- @classmethod
- def _add_argparse_args(cls, parser):
- parser.add_value_provider_argument(
- '--vp_arg',
- help='This flag is a value provider')
-
- parser.add_value_provider_argument(
- '--vp_arg2',
- default=1,
- type=int)
-
- parser.add_argument(
- '--non_vp_arg',
- default=1,
- type=int
- )
-
- # Provide values: if not provided, the option becomes of the type runtime vp
- 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(vp_arg=5,
- vp_arg2=StaticValueProvider(value_type=str,
- value='bye'),
- non_vp_arg=RuntimeValueProvider(
- option_name='foo',
- value_type=int,
- default_value=10,
- options_id=10))
- 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.non_vp_arg.get()
-
+ self.assertEqual(options.get_all_options()['redefined_flag'], True)
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
http://git-wip-us.apache.org/repos/asf/beam/blob/0f5c3639/sdks/python/apache_beam/utils/value_provider.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/utils/value_provider.py b/sdks/python/apache_beam/utils/value_provider.py
deleted file mode 100644
index a72fc4c..0000000
--- a/sdks/python/apache_beam/utils/value_provider.py
+++ /dev/null
@@ -1,110 +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 ValueProvider class to implement templates with both statically
-and dynamically provided values.
-"""
-
-from functools import wraps
-
-
-class ValueProvider(object):
- def is_accessible(self):
- raise NotImplementedError(
- 'ValueProvider.is_accessible implemented in derived classes'
- )
-
- def get(self):
- raise NotImplementedError(
- 'ValueProvider.get implemented in derived classes'
- )
-
-
-class StaticValueProvider(ValueProvider):
- def __init__(self, value_type, value):
- self.value_type = value_type
- self.value = value_type(value)
-
- def is_accessible(self):
- return True
-
- def get(self):
- return self.value
-
- def __str__(self):
- return str(self.value)
-
-
-class RuntimeValueProvider(ValueProvider):
- runtime_options_map = {}
-
- def __init__(self, option_name, value_type, default_value, options_id):
- assert options_id is not None
- self.option_name = option_name
- self.default_value = default_value
- self.value_type = value_type
- self.options_id = options_id
-
- def is_accessible(self):
- return RuntimeValueProvider.runtime_options_map.get(
- self.options_id) is not None
-
- def get(self):
- runtime_options = (
- RuntimeValueProvider.runtime_options_map.get(self.options_id))
- if runtime_options is None:
- raise RuntimeError('%s.get() not called from a runtime context' % self)
-
- candidate = runtime_options.get(self.option_name)
- if candidate:
- value = self.value_type(candidate)
- else:
- value = self.default_value
- return value
-
- @classmethod
- def set_runtime_options(cls, options_id, pipeline_options):
- assert options_id not in RuntimeValueProvider.runtime_options_map
- RuntimeValueProvider.runtime_options_map[options_id] = pipeline_options
-
- @classmethod
- def unset_runtime_options(cls, options_id):
- assert options_id in RuntimeValueProvider.runtime_options_map
- del RuntimeValueProvider.runtime_options_map[options_id]
-
- def __str__(self):
- return '%s(option: %s, type: %s, default_value: %s)' % (
- self.__class__.__name__,
- self.option_name,
- self.value_type.__name__,
- repr(self.default_value)
- )
-
-
-def check_accessible(value_provider_list):
- """Check accessibility of a list of ValueProvider objects."""
- assert isinstance(value_provider_list, list)
-
- def _check_accessible(fnc):
- @wraps(fnc)
- def _f(self, *args, **kwargs):
- for obj in [getattr(self, vp) for vp in value_provider_list]:
- if not obj.is_accessible():
- raise RuntimeError('%s not accessible' % obj)
- return fnc(self, *args, **kwargs)
- return _f
- return _check_accessible
http://git-wip-us.apache.org/repos/asf/beam/blob/0f5c3639/sdks/python/apache_beam/utils/value_provider_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/utils/value_provider_test.py b/sdks/python/apache_beam/utils/value_provider_test.py
deleted file mode 100644
index 83cb5e9..0000000
--- a/sdks/python/apache_beam/utils/value_provider_test.py
+++ /dev/null
@@ -1,165 +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.
-#
-
-"""Unit tests for the ValueProvider class."""
-
-import unittest
-
-from apache_beam.utils.pipeline_options import PipelineOptions
-from apache_beam.utils.value_provider import RuntimeValueProvider
-from apache_beam.utils.value_provider import StaticValueProvider
-
-
-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(
- '--vp_arg',
- help='This keyword argument is a value provider',
- default='some value')
- 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(
- '--vp_arg',
- help='This keyword argument is a value provider')
- options = UserDefinedOptions()
- self.assertTrue(isinstance(options.vp_arg, RuntimeValueProvider))
- self.assertFalse(options.vp_arg.is_accessible())
- with self.assertRaises(RuntimeError):
- 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(
- 'vp_pos_arg',
- help='This positional argument is a value provider',
- default='some value')
- options = UserDefinedOptions(['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(
- 'vp_pos_arg',
- help='This positional argument is a value provider')
- options = UserDefinedOptions([])
- self.assertTrue(isinstance(options.vp_pos_arg, RuntimeValueProvider))
- self.assertFalse(options.vp_pos_arg.is_accessible())
- with self.assertRaises(RuntimeError):
- 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(
- '--vp_arg',
- type=int,
- help='This flag is a value provider')
-
- 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
- class UserDefinedOptions1(PipelineOptions):
- @classmethod
- def _add_argparse_args(cls, parser):
- parser.add_value_provider_argument(
- '--vp_arg',
- help='This keyword argument is a value provider') # set at runtime
-
- parser.add_value_provider_argument( # not set, had default int
- '-v', '--vp_arg2', # with short form
- default=123,
- type=int)
-
- parser.add_value_provider_argument( # not set, had default str
- '--vp-arg3', # with dash in name
- default='123',
- type=str)
-
- parser.add_value_provider_argument( # not set and no default
- '--vp_arg4',
- type=float)
-
- parser.add_value_provider_argument( # positional argument set
- 'vp_pos_arg', # default & runtime ignored
- help='This positional argument is a value provider',
- type=float,
- default=5.4)
-
- # provide values at graph-construction time
- # (options not provided here become of the type RuntimeValueProvider)
- options = UserDefinedOptions1(['1.2'])
- 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(
- options._options_id, {'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)
-
- def test_options_id(self):
- class Opt1(PipelineOptions):
- @classmethod
- def _add_argparse_args(cls, parser):
- parser.add_value_provider_argument('--arg1')
-
- class Opt2(PipelineOptions):
- @classmethod
- def _add_argparse_args(cls, parser):
- parser.add_value_provider_argument('--arg2')
-
- opt1 = Opt1()
- opt2 = Opt2()
- self.assertFalse(opt1.arg1.is_accessible())
- self.assertFalse(opt2.arg2.is_accessible())
- RuntimeValueProvider.set_runtime_options(
- opt1.arg1.options_id, {'arg1': 'val1'})
- self.assertTrue(opt1.arg1.is_accessible())
- self.assertFalse(opt2.arg2.is_accessible())
[15/50] [abbrv] beam git commit: This closes #2493
Posted by jb...@apache.org.
This closes #2493
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/7587d29f
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/7587d29f
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/7587d29f
Branch: refs/heads/DSL_SQL
Commit: 7587d29f50654397f00a754596bfc5a225d4a64f
Parents: 7c169a6 db81205
Author: Thomas Groh <tg...@google.com>
Authored: Tue Apr 11 13:44:34 2017 -0700
Committer: Thomas Groh <tg...@google.com>
Committed: Tue Apr 11 13:44:34 2017 -0700
----------------------------------------------------------------------
.../direct/TransformExecutorServices.java | 37 +++++++++++----
.../direct/TransformExecutorServicesTest.java | 48 ++++++++++++++++++++
2 files changed, 77 insertions(+), 8 deletions(-)
----------------------------------------------------------------------
[31/50] [abbrv] beam git commit: [BEAM-1722] Move PubsubIO into the
google-cloud-platform module
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/src/test/java/org/apache/beam/sdk/util/PubsubGrpcClientTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/PubsubGrpcClientTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/PubsubGrpcClientTest.java
deleted file mode 100644
index 6d4cf4e..0000000
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/PubsubGrpcClientTest.java
+++ /dev/null
@@ -1,207 +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.util;
-
-import static org.junit.Assert.assertEquals;
-
-import com.google.auth.Credentials;
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Iterables;
-import com.google.protobuf.ByteString;
-import com.google.protobuf.Timestamp;
-import com.google.pubsub.v1.PublishRequest;
-import com.google.pubsub.v1.PublishResponse;
-import com.google.pubsub.v1.PublisherGrpc.PublisherImplBase;
-import com.google.pubsub.v1.PubsubMessage;
-import com.google.pubsub.v1.PullRequest;
-import com.google.pubsub.v1.PullResponse;
-import com.google.pubsub.v1.ReceivedMessage;
-import com.google.pubsub.v1.SubscriberGrpc.SubscriberImplBase;
-import io.grpc.ManagedChannel;
-import io.grpc.Server;
-import io.grpc.inprocess.InProcessChannelBuilder;
-import io.grpc.inprocess.InProcessServerBuilder;
-import io.grpc.stub.StreamObserver;
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-import java.util.concurrent.ThreadLocalRandom;
-import org.apache.beam.sdk.util.PubsubClient.IncomingMessage;
-import org.apache.beam.sdk.util.PubsubClient.OutgoingMessage;
-import org.apache.beam.sdk.util.PubsubClient.SubscriptionPath;
-import org.apache.beam.sdk.util.PubsubClient.TopicPath;
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-/**
- * Tests for PubsubGrpcClient.
- */
-@RunWith(JUnit4.class)
-public class PubsubGrpcClientTest {
- private ManagedChannel inProcessChannel;
- private Credentials testCredentials;
-
- private PubsubClient client;
- private String channelName;
-
- private static final TopicPath TOPIC = PubsubClient.topicPathFromName("testProject", "testTopic");
- private static final SubscriptionPath SUBSCRIPTION =
- PubsubClient.subscriptionPathFromName("testProject", "testSubscription");
- private static final long REQ_TIME = 1234L;
- private static final long PUB_TIME = 3456L;
- private static final long MESSAGE_TIME = 6789L;
- private static final String TIMESTAMP_LABEL = "timestamp";
- private static final String ID_LABEL = "id";
- private static final String MESSAGE_ID = "testMessageId";
- private static final String DATA = "testData";
- private static final String RECORD_ID = "testRecordId";
- private static final String ACK_ID = "testAckId";
- private static final Map<String, String> ATTRIBUTES =
- ImmutableMap.<String, String>builder().put("a", "b").put("c", "d").build();
-
- @Before
- public void setup() {
- channelName = String.format("%s-%s",
- PubsubGrpcClientTest.class.getName(), ThreadLocalRandom.current().nextInt());
- inProcessChannel = InProcessChannelBuilder.forName(channelName).directExecutor().build();
- testCredentials = new TestCredential();
- client = new PubsubGrpcClient(TIMESTAMP_LABEL, ID_LABEL, 10, inProcessChannel, testCredentials);
- }
-
- @After
- public void teardown() throws IOException {
- client.close();
- inProcessChannel.shutdownNow();
- }
-
- @Test
- public void pullOneMessage() throws IOException {
- String expectedSubscription = SUBSCRIPTION.getPath();
- final PullRequest expectedRequest =
- PullRequest.newBuilder()
- .setSubscription(expectedSubscription)
- .setReturnImmediately(true)
- .setMaxMessages(10)
- .build();
- Timestamp timestamp = Timestamp.newBuilder()
- .setSeconds(PUB_TIME / 1000)
- .setNanos((int) (PUB_TIME % 1000) * 1000)
- .build();
- PubsubMessage expectedPubsubMessage =
- PubsubMessage.newBuilder()
- .setMessageId(MESSAGE_ID)
- .setData(
- ByteString.copyFrom(DATA.getBytes()))
- .setPublishTime(timestamp)
- .putAllAttributes(ATTRIBUTES)
- .putAllAttributes(
- ImmutableMap.of(TIMESTAMP_LABEL,
- String.valueOf(MESSAGE_TIME),
- ID_LABEL, RECORD_ID))
- .build();
- ReceivedMessage expectedReceivedMessage =
- ReceivedMessage.newBuilder()
- .setMessage(expectedPubsubMessage)
- .setAckId(ACK_ID)
- .build();
- final PullResponse response =
- PullResponse.newBuilder()
- .addAllReceivedMessages(ImmutableList.of(expectedReceivedMessage))
- .build();
-
- final List<PullRequest> requestsReceived = new ArrayList<>();
- SubscriberImplBase subscriberImplBase = new SubscriberImplBase() {
- @Override
- public void pull(PullRequest request, StreamObserver<PullResponse> responseObserver) {
- requestsReceived.add(request);
- responseObserver.onNext(response);
- responseObserver.onCompleted();
- }
- };
- Server server = InProcessServerBuilder.forName(channelName)
- .addService(subscriberImplBase)
- .build()
- .start();
- try {
- List<IncomingMessage> acutalMessages = client.pull(REQ_TIME, SUBSCRIPTION, 10, true);
- assertEquals(1, acutalMessages.size());
- IncomingMessage actualMessage = acutalMessages.get(0);
- assertEquals(ACK_ID, actualMessage.ackId);
- assertEquals(DATA, new String(actualMessage.elementBytes));
- assertEquals(RECORD_ID, actualMessage.recordId);
- assertEquals(REQ_TIME, actualMessage.requestTimeMsSinceEpoch);
- assertEquals(MESSAGE_TIME, actualMessage.timestampMsSinceEpoch);
- assertEquals(expectedRequest, Iterables.getOnlyElement(requestsReceived));
- } finally {
- server.shutdownNow();
- }
- }
-
- @Test
- public void publishOneMessage() throws IOException {
- String expectedTopic = TOPIC.getPath();
- PubsubMessage expectedPubsubMessage =
- PubsubMessage.newBuilder()
- .setData(ByteString.copyFrom(DATA.getBytes()))
- .putAllAttributes(ATTRIBUTES)
- .putAllAttributes(
- ImmutableMap.of(TIMESTAMP_LABEL, String.valueOf(MESSAGE_TIME),
- ID_LABEL, RECORD_ID))
- .build();
- final PublishRequest expectedRequest =
- PublishRequest.newBuilder()
- .setTopic(expectedTopic)
- .addAllMessages(
- ImmutableList.of(expectedPubsubMessage))
- .build();
- final PublishResponse response =
- PublishResponse.newBuilder()
- .addAllMessageIds(ImmutableList.of(MESSAGE_ID))
- .build();
-
- final List<PublishRequest> requestsReceived = new ArrayList<>();
- PublisherImplBase publisherImplBase = new PublisherImplBase() {
- @Override
- public void publish(
- PublishRequest request, StreamObserver<PublishResponse> responseObserver) {
- requestsReceived.add(request);
- responseObserver.onNext(response);
- responseObserver.onCompleted();
- }
- };
- Server server = InProcessServerBuilder.forName(channelName)
- .addService(publisherImplBase)
- .build()
- .start();
- try {
- OutgoingMessage actualMessage = new OutgoingMessage(
- DATA.getBytes(), ATTRIBUTES, MESSAGE_TIME, RECORD_ID);
- int n = client.publish(TOPIC, ImmutableList.of(actualMessage));
- assertEquals(1, n);
- assertEquals(expectedRequest, Iterables.getOnlyElement(requestsReceived));
- } finally {
- server.shutdownNow();
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/src/test/java/org/apache/beam/sdk/util/PubsubJsonClientTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/PubsubJsonClientTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/PubsubJsonClientTest.java
deleted file mode 100644
index 019190b..0000000
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/PubsubJsonClientTest.java
+++ /dev/null
@@ -1,140 +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.util;
-
-import static org.junit.Assert.assertEquals;
-
-import com.google.api.services.pubsub.Pubsub;
-import com.google.api.services.pubsub.model.PublishRequest;
-import com.google.api.services.pubsub.model.PublishResponse;
-import com.google.api.services.pubsub.model.PubsubMessage;
-import com.google.api.services.pubsub.model.PullRequest;
-import com.google.api.services.pubsub.model.PullResponse;
-import com.google.api.services.pubsub.model.ReceivedMessage;
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
-import java.io.IOException;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-
-import org.apache.beam.sdk.util.PubsubClient.IncomingMessage;
-import org.apache.beam.sdk.util.PubsubClient.OutgoingMessage;
-import org.apache.beam.sdk.util.PubsubClient.SubscriptionPath;
-import org.apache.beam.sdk.util.PubsubClient.TopicPath;
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-import org.mockito.Mockito;
-
-/**
- * Tests for PubsubJsonClient.
- */
-@RunWith(JUnit4.class)
-public class PubsubJsonClientTest {
- private Pubsub mockPubsub;
- private PubsubClient client;
-
- private static final TopicPath TOPIC = PubsubClient.topicPathFromName("testProject", "testTopic");
- private static final SubscriptionPath SUBSCRIPTION =
- PubsubClient.subscriptionPathFromName("testProject", "testSubscription");
- private static final long REQ_TIME = 1234L;
- private static final long PUB_TIME = 3456L;
- private static final long MESSAGE_TIME = 6789L;
- private static final String TIMESTAMP_LABEL = "timestamp";
- private static final String ID_LABEL = "id";
- private static final String MESSAGE_ID = "testMessageId";
- private static final String DATA = "testData";
- private static final String RECORD_ID = "testRecordId";
- private static final String ACK_ID = "testAckId";
-
- @Before
- public void setup() throws IOException {
- mockPubsub = Mockito.mock(Pubsub.class, Mockito.RETURNS_DEEP_STUBS);
- client = new PubsubJsonClient(TIMESTAMP_LABEL, ID_LABEL, mockPubsub);
- }
-
- @After
- public void teardown() throws IOException {
- client.close();
- client = null;
- mockPubsub = null;
- }
-
- @Test
- public void pullOneMessage() throws IOException {
- String expectedSubscription = SUBSCRIPTION.getPath();
- PullRequest expectedRequest =
- new PullRequest().setReturnImmediately(true).setMaxMessages(10);
- PubsubMessage expectedPubsubMessage = new PubsubMessage()
- .setMessageId(MESSAGE_ID)
- .encodeData(DATA.getBytes())
- .setPublishTime(String.valueOf(PUB_TIME))
- .setAttributes(
- ImmutableMap.of(TIMESTAMP_LABEL, String.valueOf(MESSAGE_TIME),
- ID_LABEL, RECORD_ID));
- ReceivedMessage expectedReceivedMessage =
- new ReceivedMessage().setMessage(expectedPubsubMessage)
- .setAckId(ACK_ID);
- PullResponse expectedResponse =
- new PullResponse().setReceivedMessages(ImmutableList.of(expectedReceivedMessage));
- Mockito.when((Object) (mockPubsub.projects()
- .subscriptions()
- .pull(expectedSubscription, expectedRequest)
- .execute()))
- .thenReturn(expectedResponse);
- List<IncomingMessage> acutalMessages = client.pull(REQ_TIME, SUBSCRIPTION, 10, true);
- assertEquals(1, acutalMessages.size());
- IncomingMessage actualMessage = acutalMessages.get(0);
- assertEquals(ACK_ID, actualMessage.ackId);
- assertEquals(DATA, new String(actualMessage.elementBytes));
- assertEquals(RECORD_ID, actualMessage.recordId);
- assertEquals(REQ_TIME, actualMessage.requestTimeMsSinceEpoch);
- assertEquals(MESSAGE_TIME, actualMessage.timestampMsSinceEpoch);
- }
-
- @Test
- public void publishOneMessage() throws IOException {
- String expectedTopic = TOPIC.getPath();
- PubsubMessage expectedPubsubMessage = new PubsubMessage()
- .encodeData(DATA.getBytes())
- .setAttributes(
- ImmutableMap.<String, String> builder()
- .put(TIMESTAMP_LABEL, String.valueOf(MESSAGE_TIME))
- .put(ID_LABEL, RECORD_ID)
- .put("k", "v").build());
- PublishRequest expectedRequest = new PublishRequest()
- .setMessages(ImmutableList.of(expectedPubsubMessage));
- PublishResponse expectedResponse = new PublishResponse()
- .setMessageIds(ImmutableList.of(MESSAGE_ID));
- Mockito.when((Object) (mockPubsub.projects()
- .topics()
- .publish(expectedTopic, expectedRequest)
- .execute()))
- .thenReturn(expectedResponse);
- Map<String, String> attrs = new HashMap<>();
- attrs.put("k", "v");
- OutgoingMessage actualMessage = new OutgoingMessage(
- DATA.getBytes(), attrs, MESSAGE_TIME, RECORD_ID);
- int n = client.publish(TOPIC, ImmutableList.of(actualMessage));
- assertEquals(1, n);
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/src/test/java/org/apache/beam/sdk/util/PubsubTestClientTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/PubsubTestClientTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/PubsubTestClientTest.java
deleted file mode 100644
index a1b7daf..0000000
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/PubsubTestClientTest.java
+++ /dev/null
@@ -1,114 +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.util;
-
-import static org.junit.Assert.assertEquals;
-
-import com.google.api.client.util.Clock;
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.Lists;
-import com.google.common.collect.Sets;
-import java.io.IOException;
-import java.util.List;
-import java.util.concurrent.atomic.AtomicLong;
-import org.apache.beam.sdk.util.PubsubClient.IncomingMessage;
-import org.apache.beam.sdk.util.PubsubClient.OutgoingMessage;
-import org.apache.beam.sdk.util.PubsubClient.SubscriptionPath;
-import org.apache.beam.sdk.util.PubsubClient.TopicPath;
-import org.apache.beam.sdk.util.PubsubTestClient.PubsubTestClientFactory;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-/**
- * Tests for PubsubTestClient.
- */
-@RunWith(JUnit4.class)
-public class PubsubTestClientTest {
- private static final TopicPath TOPIC = PubsubClient.topicPathFromName("testProject", "testTopic");
- private static final SubscriptionPath SUBSCRIPTION =
- PubsubClient.subscriptionPathFromName("testProject", "testSubscription");
- private static final long REQ_TIME = 1234L;
- private static final long MESSAGE_TIME = 6789L;
- private static final String MESSAGE_ID = "testMessageId";
- private static final String DATA = "testData";
- private static final String ACK_ID = "testAckId";
- private static final int ACK_TIMEOUT_S = 60;
-
- @Test
- public void pullOneMessage() throws IOException {
- final AtomicLong now = new AtomicLong();
- Clock clock = new Clock() {
- @Override
- public long currentTimeMillis() {
- return now.get();
- }
- };
- IncomingMessage expectedIncomingMessage =
- new IncomingMessage(DATA.getBytes(), null, MESSAGE_TIME, REQ_TIME, ACK_ID, MESSAGE_ID);
- try (PubsubTestClientFactory factory =
- PubsubTestClient.createFactoryForPull(clock, SUBSCRIPTION, ACK_TIMEOUT_S,
- Lists.newArrayList(expectedIncomingMessage))) {
- try (PubsubTestClient client = (PubsubTestClient) factory.newClient(null, null, null)) {
- now.set(REQ_TIME);
- client.advance();
- List<IncomingMessage> incomingMessages = client.pull(now.get(), SUBSCRIPTION, 1, true);
- assertEquals(1, incomingMessages.size());
- assertEquals(expectedIncomingMessage, incomingMessages.get(0));
- // Timeout on ACK.
- now.addAndGet((ACK_TIMEOUT_S + 10) * 1000);
- client.advance();
- incomingMessages = client.pull(now.get(), SUBSCRIPTION, 1, true);
- assertEquals(1, incomingMessages.size());
- assertEquals(expectedIncomingMessage.withRequestTime(now.get()), incomingMessages.get(0));
- now.addAndGet(10 * 1000);
- client.advance();
- // Extend ack
- client.modifyAckDeadline(SUBSCRIPTION, ImmutableList.of(ACK_ID), 20);
- // Timeout on extended ACK
- now.addAndGet(30 * 1000);
- client.advance();
- incomingMessages = client.pull(now.get(), SUBSCRIPTION, 1, true);
- assertEquals(1, incomingMessages.size());
- assertEquals(expectedIncomingMessage.withRequestTime(now.get()), incomingMessages.get(0));
- // Extend ack
- client.modifyAckDeadline(SUBSCRIPTION, ImmutableList.of(ACK_ID), 20);
- // Ack
- now.addAndGet(15 * 1000);
- client.advance();
- client.acknowledge(SUBSCRIPTION, ImmutableList.of(ACK_ID));
- }
- }
- }
-
- @Test
- public void publishOneMessage() throws IOException {
- OutgoingMessage expectedOutgoingMessage =
- new OutgoingMessage(DATA.getBytes(), null, MESSAGE_TIME, MESSAGE_ID);
- try (PubsubTestClientFactory factory =
- PubsubTestClient.createFactoryForPublish(
- TOPIC,
- Sets.newHashSet(expectedOutgoingMessage),
- ImmutableList.<OutgoingMessage>of())) {
- try (PubsubTestClient client = (PubsubTestClient) factory.newClient(null, null, null)) {
- client.publish(TOPIC, ImmutableList.of(expectedOutgoingMessage));
- }
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/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 1a2e100..d22c6c5 100644
--- a/sdks/java/io/google-cloud-platform/pom.xml
+++ b/sdks/java/io/google-cloud-platform/pom.xml
@@ -84,6 +84,16 @@
</dependency>
<dependency>
+ <groupId>com.google.apis</groupId>
+ <artifactId>google-api-services-pubsub</artifactId>
+ </dependency>
+
+ <dependency>
+ <groupId>com.google.api.grpc</groupId>
+ <artifactId>grpc-google-pubsub-v1</artifactId>
+ </dependency>
+
+ <dependency>
<groupId>com.google.auto.service</groupId>
<artifactId>auto-service</artifactId>
<optional>true</optional>
@@ -106,10 +116,44 @@
<dependency>
<groupId>io.grpc</groupId>
+ <artifactId>grpc-auth</artifactId>
+ </dependency>
+
+ <dependency>
+ <groupId>io.grpc</groupId>
<artifactId>grpc-core</artifactId>
</dependency>
<dependency>
+ <groupId>io.grpc</groupId>
+ <artifactId>grpc-netty</artifactId>
+ </dependency>
+
+ <dependency>
+ <groupId>io.netty</groupId>
+ <artifactId>netty-handler</artifactId>
+ </dependency>
+
+ <dependency>
+ <groupId>io.grpc</groupId>
+ <artifactId>grpc-stub</artifactId>
+ </dependency>
+
+ <!-- grpc-all does not obey IWYU, so we need to exclude from compile
+ scope and depend on it at runtime. -->
+ <dependency>
+ <groupId>io.grpc</groupId>
+ <artifactId>grpc-all</artifactId>
+ <scope>runtime</scope>
+ </dependency>
+
+ <dependency>
+ <groupId>io.grpc</groupId>
+ <artifactId>grpc-protobuf</artifactId>
+ <scope>runtime</scope>
+ </dependency>
+
+ <dependency>
<groupId>joda-time</groupId>
<artifactId>joda-time</artifactId>
</dependency>
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubClient.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubClient.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubClient.java
new file mode 100644
index 0000000..750178c
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubClient.java
@@ -0,0 +1,544 @@
+/*
+ * 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.pubsub;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkState;
+
+import com.google.api.client.util.DateTime;
+import com.google.common.base.Objects;
+import com.google.common.base.Strings;
+import java.io.Closeable;
+import java.io.IOException;
+import java.io.Serializable;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ThreadLocalRandom;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.options.PubsubOptions;
+
+/**
+ * An (abstract) helper class for talking to Pubsub via an underlying transport.
+ */
+abstract class PubsubClient implements Closeable {
+ /**
+ * Factory for creating clients.
+ */
+ public interface PubsubClientFactory extends Serializable {
+ /**
+ * Construct a new Pubsub client. It should be closed via {@link #close} in order
+ * to ensure tidy cleanup of underlying netty resources (or use the try-with-resources
+ * construct). Uses {@code options} to derive pubsub endpoints and application credentials.
+ * If non-{@literal null}, use {@code timestampLabel} and {@code idLabel} to store custom
+ * timestamps/ids within message metadata.
+ */
+ PubsubClient newClient(
+ @Nullable String timestampLabel,
+ @Nullable String idLabel,
+ PubsubOptions options) throws IOException;
+
+ /**
+ * Return the display name for this factory. Eg "Json", "gRPC".
+ */
+ String getKind();
+ }
+
+ /**
+ * Return timestamp as ms-since-unix-epoch corresponding to {@code timestamp}.
+ * Return {@literal null} if no timestamp could be found. Throw {@link IllegalArgumentException}
+ * if timestamp cannot be recognized.
+ */
+ @Nullable
+ private static Long asMsSinceEpoch(@Nullable String timestamp) {
+ if (Strings.isNullOrEmpty(timestamp)) {
+ return null;
+ }
+ try {
+ // Try parsing as milliseconds since epoch. Note there is no way to parse a
+ // string in RFC 3339 format here.
+ // Expected IllegalArgumentException if parsing fails; we use that to fall back
+ // to RFC 3339.
+ return Long.parseLong(timestamp);
+ } catch (IllegalArgumentException e1) {
+ // Try parsing as RFC3339 string. DateTime.parseRfc3339 will throw an
+ // IllegalArgumentException if parsing fails, and the caller should handle.
+ return DateTime.parseRfc3339(timestamp).getValue();
+ }
+ }
+
+ /**
+ * Return the timestamp (in ms since unix epoch) to use for a Pubsub message with {@code
+ * attributes} and {@code pubsubTimestamp}.
+ *
+ * <p>If {@code timestampLabel} is non-{@literal null} then the message attributes must contain
+ * that label, and the value of that label will be taken as the timestamp.
+ * Otherwise the timestamp will be taken from the Pubsub publish timestamp {@code
+ * pubsubTimestamp}.
+ *
+ * @throws IllegalArgumentException if the timestamp cannot be recognized as a ms-since-unix-epoch
+ * or RFC3339 time.
+ */
+ protected static long extractTimestamp(
+ @Nullable String timestampLabel,
+ @Nullable String pubsubTimestamp,
+ @Nullable Map<String, String> attributes) {
+ Long timestampMsSinceEpoch;
+ if (Strings.isNullOrEmpty(timestampLabel)) {
+ timestampMsSinceEpoch = asMsSinceEpoch(pubsubTimestamp);
+ checkArgument(timestampMsSinceEpoch != null,
+ "Cannot interpret PubSub publish timestamp: %s",
+ pubsubTimestamp);
+ } else {
+ String value = attributes == null ? null : attributes.get(timestampLabel);
+ checkArgument(value != null,
+ "PubSub message is missing a value for timestamp label %s",
+ timestampLabel);
+ timestampMsSinceEpoch = asMsSinceEpoch(value);
+ checkArgument(timestampMsSinceEpoch != null,
+ "Cannot interpret value of label %s as timestamp: %s",
+ timestampLabel, value);
+ }
+ return timestampMsSinceEpoch;
+ }
+
+ /**
+ * Path representing a cloud project id.
+ */
+ static class ProjectPath implements Serializable {
+ private final String projectId;
+
+ /**
+ * Creates a {@link ProjectPath} from a {@link String} representation, which
+ * must be of the form {@code "projects/" + projectId}.
+ */
+ ProjectPath(String path) {
+ String[] splits = path.split("/");
+ checkArgument(
+ splits.length == 2 && splits[0].equals("projects"),
+ "Malformed project path \"%s\": must be of the form \"projects/\" + <project id>",
+ path);
+ this.projectId = splits[1];
+ }
+
+ public String getPath() {
+ return String.format("projects/%s", projectId);
+ }
+
+ public String getId() {
+ return projectId;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+
+ ProjectPath that = (ProjectPath) o;
+
+ return projectId.equals(that.projectId);
+ }
+
+ @Override
+ public int hashCode() {
+ return projectId.hashCode();
+ }
+
+ @Override
+ public String toString() {
+ return getPath();
+ }
+ }
+
+ public static ProjectPath projectPathFromPath(String path) {
+ return new ProjectPath(path);
+ }
+
+ public static ProjectPath projectPathFromId(String projectId) {
+ return new ProjectPath(String.format("projects/%s", projectId));
+ }
+
+ /**
+ * Path representing a Pubsub subscription.
+ */
+ public static class SubscriptionPath implements Serializable {
+ private final String projectId;
+ private final String subscriptionName;
+
+ SubscriptionPath(String path) {
+ String[] splits = path.split("/");
+ checkState(
+ splits.length == 4 && splits[0].equals("projects") && splits[2].equals("subscriptions"),
+ "Malformed subscription path %s: "
+ + "must be of the form \"projects/\" + <project id> + \"subscriptions\"", path);
+ this.projectId = splits[1];
+ this.subscriptionName = splits[3];
+ }
+
+ public String getPath() {
+ return String.format("projects/%s/subscriptions/%s", projectId, subscriptionName);
+ }
+
+ public String getName() {
+ return subscriptionName;
+ }
+
+ public String getV1Beta1Path() {
+ return String.format("/subscriptions/%s/%s", projectId, subscriptionName);
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ SubscriptionPath that = (SubscriptionPath) o;
+ return this.subscriptionName.equals(that.subscriptionName)
+ && this.projectId.equals(that.projectId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hashCode(projectId, subscriptionName);
+ }
+
+ @Override
+ public String toString() {
+ return getPath();
+ }
+ }
+
+ public static SubscriptionPath subscriptionPathFromPath(String path) {
+ return new SubscriptionPath(path);
+ }
+
+ public static SubscriptionPath subscriptionPathFromName(
+ String projectId, String subscriptionName) {
+ return new SubscriptionPath(String.format("projects/%s/subscriptions/%s",
+ projectId, subscriptionName));
+ }
+
+ /**
+ * Path representing a Pubsub topic.
+ */
+ public static class TopicPath implements Serializable {
+ private final String path;
+
+ TopicPath(String path) {
+ this.path = path;
+ }
+
+ public String getPath() {
+ return path;
+ }
+
+ public String getName() {
+ String[] splits = path.split("/");
+ checkState(splits.length == 4, "Malformed topic path %s", path);
+ return splits[3];
+ }
+
+ public String getV1Beta1Path() {
+ String[] splits = path.split("/");
+ checkState(splits.length == 4, "Malformed topic path %s", path);
+ return String.format("/topics/%s/%s", splits[1], splits[3]);
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ TopicPath topicPath = (TopicPath) o;
+ return path.equals(topicPath.path);
+ }
+
+ @Override
+ public int hashCode() {
+ return path.hashCode();
+ }
+
+ @Override
+ public String toString() {
+ return path;
+ }
+ }
+
+ public static TopicPath topicPathFromPath(String path) {
+ return new TopicPath(path);
+ }
+
+ public static TopicPath topicPathFromName(String projectId, String topicName) {
+ return new TopicPath(String.format("projects/%s/topics/%s", projectId, topicName));
+ }
+
+ /**
+ * A message to be sent to Pubsub.
+ *
+ * <p>NOTE: This class is {@link Serializable} only to support the {@link PubsubTestClient}.
+ * Java serialization is never used for non-test clients.
+ */
+ static class OutgoingMessage implements Serializable {
+ /**
+ * Underlying (encoded) element.
+ */
+ public final byte[] elementBytes;
+
+ public final Map<String, String> attributes;
+
+ /**
+ * Timestamp for element (ms since epoch).
+ */
+ public final long timestampMsSinceEpoch;
+
+ /**
+ * If using an id label, the record id to associate with this record's metadata so the receiver
+ * can reject duplicates. Otherwise {@literal null}.
+ */
+ @Nullable
+ public final String recordId;
+
+ public OutgoingMessage(byte[] elementBytes, Map<String, String> attributes,
+ long timestampMsSinceEpoch, @Nullable String recordId) {
+ this.elementBytes = elementBytes;
+ this.attributes = attributes;
+ this.timestampMsSinceEpoch = timestampMsSinceEpoch;
+ this.recordId = recordId;
+ }
+
+ @Override
+ public String toString() {
+ return String.format("OutgoingMessage(%db, %dms)",
+ elementBytes.length, timestampMsSinceEpoch);
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+
+ OutgoingMessage that = (OutgoingMessage) o;
+
+ return timestampMsSinceEpoch == that.timestampMsSinceEpoch
+ && Arrays.equals(elementBytes, that.elementBytes)
+ && Objects.equal(attributes, that.attributes)
+ && Objects.equal(recordId, that.recordId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hashCode(Arrays.hashCode(elementBytes), attributes, timestampMsSinceEpoch,
+ recordId);
+ }
+ }
+
+ /**
+ * A message received from Pubsub.
+ *
+ * <p>NOTE: This class is {@link Serializable} only to support the {@link PubsubTestClient}.
+ * Java serialization is never used for non-test clients.
+ */
+ static class IncomingMessage implements Serializable {
+ /**
+ * Underlying (encoded) element.
+ */
+ public final byte[] elementBytes;
+
+ public Map<String, String> attributes;
+
+ /**
+ * Timestamp for element (ms since epoch). Either Pubsub's processing time,
+ * or the custom timestamp associated with the message.
+ */
+ public final long timestampMsSinceEpoch;
+
+ /**
+ * Timestamp (in system time) at which we requested the message (ms since epoch).
+ */
+ public final long requestTimeMsSinceEpoch;
+
+ /**
+ * Id to pass back to Pubsub to acknowledge receipt of this message.
+ */
+ public final String ackId;
+
+ /**
+ * Id to pass to the runner to distinguish this message from all others.
+ */
+ public final String recordId;
+
+ public IncomingMessage(
+ byte[] elementBytes,
+ Map<String, String> attributes,
+ long timestampMsSinceEpoch,
+ long requestTimeMsSinceEpoch,
+ String ackId,
+ String recordId) {
+ this.elementBytes = elementBytes;
+ this.attributes = attributes;
+ this.timestampMsSinceEpoch = timestampMsSinceEpoch;
+ this.requestTimeMsSinceEpoch = requestTimeMsSinceEpoch;
+ this.ackId = ackId;
+ this.recordId = recordId;
+ }
+
+ public IncomingMessage withRequestTime(long requestTimeMsSinceEpoch) {
+ return new IncomingMessage(elementBytes, attributes, timestampMsSinceEpoch,
+ requestTimeMsSinceEpoch, ackId, recordId);
+ }
+
+ @Override
+ public String toString() {
+ return String.format("IncomingMessage(%db, %dms)",
+ elementBytes.length, timestampMsSinceEpoch);
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+
+ IncomingMessage that = (IncomingMessage) o;
+
+ return timestampMsSinceEpoch == that.timestampMsSinceEpoch
+ && requestTimeMsSinceEpoch == that.requestTimeMsSinceEpoch
+ && ackId.equals(that.ackId)
+ && recordId.equals(that.recordId)
+ && Arrays.equals(elementBytes, that.elementBytes)
+ && Objects.equal(attributes, that.attributes);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hashCode(Arrays.hashCode(elementBytes), attributes, timestampMsSinceEpoch,
+ requestTimeMsSinceEpoch,
+ ackId, recordId);
+ }
+ }
+
+ /**
+ * Publish {@code outgoingMessages} to Pubsub {@code topic}. Return number of messages
+ * published.
+ */
+ public abstract int publish(TopicPath topic, List<OutgoingMessage> outgoingMessages)
+ throws IOException;
+
+ /**
+ * Request the next batch of up to {@code batchSize} messages from {@code subscription}.
+ * Return the received messages, or empty collection if none were available. Does not
+ * wait for messages to arrive if {@code returnImmediately} is {@literal true}.
+ * Returned messages will record their request time as {@code requestTimeMsSinceEpoch}.
+ */
+ public abstract List<IncomingMessage> pull(
+ long requestTimeMsSinceEpoch,
+ SubscriptionPath subscription,
+ int batchSize,
+ boolean returnImmediately)
+ throws IOException;
+
+ /**
+ * Acknowldege messages from {@code subscription} with {@code ackIds}.
+ */
+ public abstract void acknowledge(SubscriptionPath subscription, List<String> ackIds)
+ throws IOException;
+
+ /**
+ * Modify the ack deadline for messages from {@code subscription} with {@code ackIds} to
+ * be {@code deadlineSeconds} from now.
+ */
+ public abstract void modifyAckDeadline(
+ SubscriptionPath subscription, List<String> ackIds,
+ int deadlineSeconds) throws IOException;
+
+ /**
+ * Create {@code topic}.
+ */
+ public abstract void createTopic(TopicPath topic) throws IOException;
+
+ /*
+ * Delete {@code topic}.
+ */
+ public abstract void deleteTopic(TopicPath topic) throws IOException;
+
+ /**
+ * Return a list of topics for {@code project}.
+ */
+ public abstract List<TopicPath> listTopics(ProjectPath project) throws IOException;
+
+ /**
+ * Create {@code subscription} to {@code topic}.
+ */
+ public abstract void createSubscription(
+ TopicPath topic, SubscriptionPath subscription, int ackDeadlineSeconds) throws IOException;
+
+ /**
+ * Create a random subscription for {@code topic}. Return the {@link SubscriptionPath}. It
+ * is the responsibility of the caller to later delete the subscription.
+ */
+ public SubscriptionPath createRandomSubscription(
+ ProjectPath project, TopicPath topic, int ackDeadlineSeconds) throws IOException {
+ // Create a randomized subscription derived from the topic name.
+ String subscriptionName = topic.getName() + "_beam_" + ThreadLocalRandom.current().nextLong();
+ SubscriptionPath subscription =
+ PubsubClient.subscriptionPathFromName(project.getId(), subscriptionName);
+ createSubscription(topic, subscription, ackDeadlineSeconds);
+ return subscription;
+ }
+
+ /**
+ * Delete {@code subscription}.
+ */
+ public abstract void deleteSubscription(SubscriptionPath subscription) throws IOException;
+
+ /**
+ * Return a list of subscriptions for {@code topic} in {@code project}.
+ */
+ public abstract List<SubscriptionPath> listSubscriptions(ProjectPath project, TopicPath topic)
+ throws IOException;
+
+ /**
+ * Return the ack deadline, in seconds, for {@code subscription}.
+ */
+ public abstract int ackDeadlineSeconds(SubscriptionPath subscription) throws IOException;
+
+ /**
+ * Return {@literal true} if {@link #pull} will always return empty list. Actual clients
+ * will return {@literal false}. Test clients may return {@literal true} to signal that all
+ * expected messages have been pulled and the test may complete.
+ */
+ public abstract boolean isEOF();
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubGrpcClient.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubGrpcClient.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubGrpcClient.java
new file mode 100644
index 0000000..912d59c
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubGrpcClient.java
@@ -0,0 +1,424 @@
+/*
+ * 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.pubsub;
+
+import static com.google.common.base.Preconditions.checkState;
+
+import com.google.auth.Credentials;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Strings;
+import com.google.common.collect.ImmutableList;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Timestamp;
+import com.google.pubsub.v1.AcknowledgeRequest;
+import com.google.pubsub.v1.DeleteSubscriptionRequest;
+import com.google.pubsub.v1.DeleteTopicRequest;
+import com.google.pubsub.v1.GetSubscriptionRequest;
+import com.google.pubsub.v1.ListSubscriptionsRequest;
+import com.google.pubsub.v1.ListSubscriptionsResponse;
+import com.google.pubsub.v1.ListTopicsRequest;
+import com.google.pubsub.v1.ListTopicsResponse;
+import com.google.pubsub.v1.ModifyAckDeadlineRequest;
+import com.google.pubsub.v1.PublishRequest;
+import com.google.pubsub.v1.PublishResponse;
+import com.google.pubsub.v1.PublisherGrpc;
+import com.google.pubsub.v1.PublisherGrpc.PublisherBlockingStub;
+import com.google.pubsub.v1.PubsubMessage;
+import com.google.pubsub.v1.PullRequest;
+import com.google.pubsub.v1.PullResponse;
+import com.google.pubsub.v1.ReceivedMessage;
+import com.google.pubsub.v1.SubscriberGrpc;
+import com.google.pubsub.v1.SubscriberGrpc.SubscriberBlockingStub;
+import com.google.pubsub.v1.Subscription;
+import com.google.pubsub.v1.Topic;
+import io.grpc.Channel;
+import io.grpc.ClientInterceptors;
+import io.grpc.ManagedChannel;
+import io.grpc.auth.ClientAuthInterceptor;
+import io.grpc.netty.GrpcSslContexts;
+import io.grpc.netty.NegotiationType;
+import io.grpc.netty.NettyChannelBuilder;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.options.GcpOptions;
+import org.apache.beam.sdk.options.PubsubOptions;
+
+/**
+ * A helper class for talking to Pubsub via grpc.
+ *
+ * <p>CAUTION: Currently uses the application default credentials and does not respect any
+ * credentials-related arguments in {@link GcpOptions}.
+ */
+class PubsubGrpcClient extends PubsubClient {
+ private static final String PUBSUB_ADDRESS = "pubsub.googleapis.com";
+ private static final int PUBSUB_PORT = 443;
+ private static final int LIST_BATCH_SIZE = 1000;
+
+ private static final int DEFAULT_TIMEOUT_S = 15;
+
+ private static class PubsubGrpcClientFactory implements PubsubClientFactory {
+ @Override
+ public PubsubClient newClient(
+ @Nullable String timestampLabel, @Nullable String idLabel, PubsubOptions options)
+ throws IOException {
+ ManagedChannel channel = NettyChannelBuilder
+ .forAddress(PUBSUB_ADDRESS, PUBSUB_PORT)
+ .negotiationType(NegotiationType.TLS)
+ .sslContext(GrpcSslContexts.forClient().ciphers(null).build())
+ .build();
+
+ return new PubsubGrpcClient(timestampLabel,
+ idLabel,
+ DEFAULT_TIMEOUT_S,
+ channel,
+ options.getGcpCredential());
+ }
+
+ @Override
+ public String getKind() {
+ return "Grpc";
+ }
+ }
+
+ /**
+ * Factory for creating Pubsub clients using gRCP transport.
+ */
+ public static final PubsubClientFactory FACTORY = new PubsubGrpcClientFactory();
+
+ /**
+ * Timeout for grpc calls (in s).
+ */
+ private final int timeoutSec;
+
+ /**
+ * Underlying netty channel, or {@literal null} if closed.
+ */
+ @Nullable
+ private ManagedChannel publisherChannel;
+
+ /**
+ * Credentials determined from options and environment.
+ */
+ private final Credentials credentials;
+
+ /**
+ * Label to use for custom timestamps, or {@literal null} if should use Pubsub publish time
+ * instead.
+ */
+ @Nullable
+ private final String timestampLabel;
+
+ /**
+ * Label to use for custom ids, or {@literal null} if should use Pubsub provided ids.
+ */
+ @Nullable
+ private final String idLabel;
+
+
+ /**
+ * Cached stubs, or null if not cached.
+ */
+ @Nullable
+ private PublisherGrpc.PublisherBlockingStub cachedPublisherStub;
+ private SubscriberGrpc.SubscriberBlockingStub cachedSubscriberStub;
+
+ @VisibleForTesting
+ PubsubGrpcClient(
+ @Nullable String timestampLabel,
+ @Nullable String idLabel,
+ int timeoutSec,
+ ManagedChannel publisherChannel,
+ Credentials credentials) {
+ this.timestampLabel = timestampLabel;
+ this.idLabel = idLabel;
+ this.timeoutSec = timeoutSec;
+ this.publisherChannel = publisherChannel;
+ this.credentials = credentials;
+ }
+
+ /**
+ * Gracefully close the underlying netty channel.
+ */
+ @Override
+ public void close() {
+ if (publisherChannel == null) {
+ // Already closed.
+ return;
+ }
+ // Can gc the underlying stubs.
+ cachedPublisherStub = null;
+ cachedSubscriberStub = null;
+ // Mark the client as having been closed before going further
+ // in case we have an exception from the channel.
+ ManagedChannel publisherChannel = this.publisherChannel;
+ this.publisherChannel = null;
+ // Gracefully shutdown the channel.
+ publisherChannel.shutdown();
+ try {
+ publisherChannel.awaitTermination(timeoutSec, TimeUnit.SECONDS);
+ } catch (InterruptedException e) {
+ // Ignore.
+ Thread.currentThread().interrupt();
+ }
+ }
+
+ /**
+ * Return channel with interceptor for returning credentials.
+ */
+ private Channel newChannel() throws IOException {
+ checkState(publisherChannel != null, "PubsubGrpcClient has been closed");
+ ClientAuthInterceptor interceptor =
+ new ClientAuthInterceptor(credentials, Executors.newSingleThreadExecutor());
+ return ClientInterceptors.intercept(publisherChannel, interceptor);
+ }
+
+ /**
+ * Return a stub for making a publish request with a timeout.
+ */
+ private PublisherBlockingStub publisherStub() throws IOException {
+ if (cachedPublisherStub == null) {
+ cachedPublisherStub = PublisherGrpc.newBlockingStub(newChannel());
+ }
+ return cachedPublisherStub.withDeadlineAfter(timeoutSec, TimeUnit.SECONDS);
+ }
+
+ /**
+ * Return a stub for making a subscribe request with a timeout.
+ */
+ private SubscriberBlockingStub subscriberStub() throws IOException {
+ if (cachedSubscriberStub == null) {
+ cachedSubscriberStub = SubscriberGrpc.newBlockingStub(newChannel());
+ }
+ return cachedSubscriberStub.withDeadlineAfter(timeoutSec, TimeUnit.SECONDS);
+ }
+
+ @Override
+ public int publish(TopicPath topic, List<OutgoingMessage> outgoingMessages)
+ throws IOException {
+ PublishRequest.Builder request = PublishRequest.newBuilder()
+ .setTopic(topic.getPath());
+ for (OutgoingMessage outgoingMessage : outgoingMessages) {
+ PubsubMessage.Builder message =
+ PubsubMessage.newBuilder()
+ .setData(ByteString.copyFrom(outgoingMessage.elementBytes));
+
+ if (outgoingMessage.attributes != null) {
+ message.putAllAttributes(outgoingMessage.attributes);
+ }
+
+ if (timestampLabel != null) {
+ message.getMutableAttributes()
+ .put(timestampLabel, String.valueOf(outgoingMessage.timestampMsSinceEpoch));
+ }
+
+ if (idLabel != null && !Strings.isNullOrEmpty(outgoingMessage.recordId)) {
+ message.getMutableAttributes().put(idLabel, outgoingMessage.recordId);
+ }
+
+ request.addMessages(message);
+ }
+
+ PublishResponse response = publisherStub().publish(request.build());
+ return response.getMessageIdsCount();
+ }
+
+ @Override
+ public List<IncomingMessage> pull(
+ long requestTimeMsSinceEpoch,
+ SubscriptionPath subscription,
+ int batchSize,
+ boolean returnImmediately) throws IOException {
+ PullRequest request = PullRequest.newBuilder()
+ .setSubscription(subscription.getPath())
+ .setReturnImmediately(returnImmediately)
+ .setMaxMessages(batchSize)
+ .build();
+ PullResponse response = subscriberStub().pull(request);
+ if (response.getReceivedMessagesCount() == 0) {
+ return ImmutableList.of();
+ }
+ List<IncomingMessage> incomingMessages = new ArrayList<>(response.getReceivedMessagesCount());
+ for (ReceivedMessage message : response.getReceivedMessagesList()) {
+ PubsubMessage pubsubMessage = message.getMessage();
+ @Nullable Map<String, String> attributes = pubsubMessage.getAttributes();
+
+ // Payload.
+ byte[] elementBytes = pubsubMessage.getData().toByteArray();
+
+ // Timestamp.
+ String pubsubTimestampString = null;
+ Timestamp timestampProto = pubsubMessage.getPublishTime();
+ if (timestampProto != null) {
+ pubsubTimestampString = String.valueOf(timestampProto.getSeconds()
+ + timestampProto.getNanos() / 1000L);
+ }
+ long timestampMsSinceEpoch =
+ extractTimestamp(timestampLabel, pubsubTimestampString, attributes);
+
+ // Ack id.
+ String ackId = message.getAckId();
+ checkState(!Strings.isNullOrEmpty(ackId));
+
+ // Record id, if any.
+ @Nullable String recordId = null;
+ if (idLabel != null && attributes != null) {
+ recordId = attributes.get(idLabel);
+ }
+ if (Strings.isNullOrEmpty(recordId)) {
+ // Fall back to the Pubsub provided message id.
+ recordId = pubsubMessage.getMessageId();
+ }
+
+ incomingMessages.add(new IncomingMessage(elementBytes, attributes, timestampMsSinceEpoch,
+ requestTimeMsSinceEpoch, ackId, recordId));
+ }
+ return incomingMessages;
+ }
+
+ @Override
+ public void acknowledge(SubscriptionPath subscription, List<String> ackIds)
+ throws IOException {
+ AcknowledgeRequest request = AcknowledgeRequest.newBuilder()
+ .setSubscription(subscription.getPath())
+ .addAllAckIds(ackIds)
+ .build();
+ subscriberStub().acknowledge(request); // ignore Empty result.
+ }
+
+ @Override
+ public void modifyAckDeadline(
+ SubscriptionPath subscription, List<String> ackIds, int deadlineSeconds)
+ throws IOException {
+ ModifyAckDeadlineRequest request =
+ ModifyAckDeadlineRequest.newBuilder()
+ .setSubscription(subscription.getPath())
+ .addAllAckIds(ackIds)
+ .setAckDeadlineSeconds(deadlineSeconds)
+ .build();
+ subscriberStub().modifyAckDeadline(request); // ignore Empty result.
+ }
+
+ @Override
+ public void createTopic(TopicPath topic) throws IOException {
+ Topic request = Topic.newBuilder()
+ .setName(topic.getPath())
+ .build();
+ publisherStub().createTopic(request); // ignore Topic result.
+ }
+
+ @Override
+ public void deleteTopic(TopicPath topic) throws IOException {
+ DeleteTopicRequest request = DeleteTopicRequest.newBuilder()
+ .setTopic(topic.getPath())
+ .build();
+ publisherStub().deleteTopic(request); // ignore Empty result.
+ }
+
+ @Override
+ public List<TopicPath> listTopics(ProjectPath project) throws IOException {
+ ListTopicsRequest.Builder request =
+ ListTopicsRequest.newBuilder()
+ .setProject(project.getPath())
+ .setPageSize(LIST_BATCH_SIZE);
+ ListTopicsResponse response = publisherStub().listTopics(request.build());
+ if (response.getTopicsCount() == 0) {
+ return ImmutableList.of();
+ }
+ List<TopicPath> topics = new ArrayList<>(response.getTopicsCount());
+ while (true) {
+ for (Topic topic : response.getTopicsList()) {
+ topics.add(topicPathFromPath(topic.getName()));
+ }
+ if (response.getNextPageToken().isEmpty()) {
+ break;
+ }
+ request.setPageToken(response.getNextPageToken());
+ response = publisherStub().listTopics(request.build());
+ }
+ return topics;
+ }
+
+ @Override
+ public void createSubscription(
+ TopicPath topic, SubscriptionPath subscription,
+ int ackDeadlineSeconds) throws IOException {
+ Subscription request = Subscription.newBuilder()
+ .setTopic(topic.getPath())
+ .setName(subscription.getPath())
+ .setAckDeadlineSeconds(ackDeadlineSeconds)
+ .build();
+ subscriberStub().createSubscription(request); // ignore Subscription result.
+ }
+
+ @Override
+ public void deleteSubscription(SubscriptionPath subscription) throws IOException {
+ DeleteSubscriptionRequest request =
+ DeleteSubscriptionRequest.newBuilder()
+ .setSubscription(subscription.getPath())
+ .build();
+ subscriberStub().deleteSubscription(request); // ignore Empty result.
+ }
+
+ @Override
+ public List<SubscriptionPath> listSubscriptions(ProjectPath project, TopicPath topic)
+ throws IOException {
+ ListSubscriptionsRequest.Builder request =
+ ListSubscriptionsRequest.newBuilder()
+ .setProject(project.getPath())
+ .setPageSize(LIST_BATCH_SIZE);
+ ListSubscriptionsResponse response = subscriberStub().listSubscriptions(request.build());
+ if (response.getSubscriptionsCount() == 0) {
+ return ImmutableList.of();
+ }
+ List<SubscriptionPath> subscriptions = new ArrayList<>(response.getSubscriptionsCount());
+ while (true) {
+ for (Subscription subscription : response.getSubscriptionsList()) {
+ if (subscription.getTopic().equals(topic.getPath())) {
+ subscriptions.add(subscriptionPathFromPath(subscription.getName()));
+ }
+ }
+ if (response.getNextPageToken().isEmpty()) {
+ break;
+ }
+ request.setPageToken(response.getNextPageToken());
+ response = subscriberStub().listSubscriptions(request.build());
+ }
+ return subscriptions;
+ }
+
+ @Override
+ public int ackDeadlineSeconds(SubscriptionPath subscription) throws IOException {
+ GetSubscriptionRequest request =
+ GetSubscriptionRequest.newBuilder()
+ .setSubscription(subscription.getPath())
+ .build();
+ Subscription response = subscriberStub().getSubscription(request);
+ return response.getAckDeadlineSeconds();
+ }
+
+ @Override
+ public boolean isEOF() {
+ return false;
+ }
+}
[14/50] [abbrv] beam git commit: Improve Work Rejection handling
Posted by jb...@apache.org.
Improve Work Rejection handling
The timing between checking for a shutdown state and submitting work is
racy. Re-check to see if a work rejection is acceptable because the
underlying executor is shut down before throwing a
RejectedExecutionException.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/db81205a
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/db81205a
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/db81205a
Branch: refs/heads/DSL_SQL
Commit: db81205af9b38245ff85ce2801af433cc31bce42
Parents: 7c169a6
Author: Thomas Groh <tg...@google.com>
Authored: Tue Apr 11 11:33:53 2017 -0700
Committer: Thomas Groh <tg...@google.com>
Committed: Tue Apr 11 13:44:34 2017 -0700
----------------------------------------------------------------------
.../direct/TransformExecutorServices.java | 37 +++++++++++----
.../direct/TransformExecutorServicesTest.java | 48 ++++++++++++++++++++
2 files changed, 77 insertions(+), 8 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/db81205a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformExecutorServices.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformExecutorServices.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformExecutorServices.java
index 6733758..53087bf 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformExecutorServices.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformExecutorServices.java
@@ -21,8 +21,11 @@ import com.google.common.base.MoreObjects;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ExecutorService;
+import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
* Static factory methods for constructing instances of {@link TransformExecutorService}.
@@ -37,7 +40,7 @@ final class TransformExecutorServices {
* parallel.
*/
public static TransformExecutorService parallel(ExecutorService executor) {
- return new ParallelEvaluationState(executor);
+ return new ParallelTransformExecutor(executor);
}
/**
@@ -45,7 +48,7 @@ final class TransformExecutorServices {
* serial.
*/
public static TransformExecutorService serial(ExecutorService executor) {
- return new SerialEvaluationState(executor);
+ return new SerialTransformExecutor(executor);
}
/**
@@ -55,18 +58,36 @@ final class TransformExecutorServices {
* <p>A principal use of this is for the evaluation of an unkeyed Step. Unkeyed computations are
* processed in parallel.
*/
- private static class ParallelEvaluationState implements TransformExecutorService {
+ private static class ParallelTransformExecutor implements TransformExecutorService {
+ private static final Logger LOG = LoggerFactory.getLogger(ParallelTransformExecutor.class);
+
private final ExecutorService executor;
private final AtomicBoolean active = new AtomicBoolean(true);
- private ParallelEvaluationState(ExecutorService executor) {
+ private ParallelTransformExecutor(ExecutorService executor) {
this.executor = executor;
}
@Override
public void schedule(TransformExecutor<?> work) {
if (active.get()) {
- executor.submit(work);
+ try {
+ executor.submit(work);
+ } catch (RejectedExecutionException rejected) {
+ boolean stillActive = active.get();
+ if (stillActive) {
+ throw new IllegalStateException(
+ String.format(
+ "Execution of Work %s was rejected, but the %s is still active",
+ work, ParallelTransformExecutor.class.getSimpleName()));
+ } else {
+ LOG.debug(
+ "Rejected execution of Work {} on executor {}. "
+ + "Suppressed exception because evaluator is not active",
+ work,
+ this);
+ }
+ }
}
}
@@ -88,14 +109,14 @@ final class TransformExecutorServices {
* <p>A principal use of this is for the serial evaluation of a (Step, Key) pair.
* Keyed computations are processed serially per step.
*/
- private static class SerialEvaluationState implements TransformExecutorService {
+ private static class SerialTransformExecutor implements TransformExecutorService {
private final ExecutorService executor;
private AtomicReference<TransformExecutor<?>> currentlyEvaluating;
private final Queue<TransformExecutor<?>> workQueue;
private boolean active = true;
- private SerialEvaluationState(ExecutorService executor) {
+ private SerialTransformExecutor(ExecutorService executor) {
this.executor = executor;
this.currentlyEvaluating = new AtomicReference<>();
this.workQueue = new ConcurrentLinkedQueue<>();
@@ -149,7 +170,7 @@ final class TransformExecutorServices {
@Override
public String toString() {
- return MoreObjects.toStringHelper(SerialEvaluationState.class)
+ return MoreObjects.toStringHelper(SerialTransformExecutor.class)
.add("currentlyEvaluating", currentlyEvaluating)
.add("workQueue", workQueue)
.toString();
http://git-wip-us.apache.org/repos/asf/beam/blob/db81205a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/TransformExecutorServicesTest.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/TransformExecutorServicesTest.java b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/TransformExecutorServicesTest.java
index b085723..77652b2 100644
--- a/runners/direct-java/src/test/java/org/apache/beam/runners/direct/TransformExecutorServicesTest.java
+++ b/runners/direct-java/src/test/java/org/apache/beam/runners/direct/TransformExecutorServicesTest.java
@@ -64,6 +64,31 @@ public class TransformExecutorServicesTest {
}
@Test
+ public void parallelRejectedStillActiveThrows() {
+ @SuppressWarnings("unchecked")
+ TransformExecutor<Object> first = mock(TransformExecutor.class);
+
+ TransformExecutorService parallel =
+ TransformExecutorServices.parallel(executorService);
+ executorService.shutdown();
+ thrown.expect(IllegalStateException.class);
+ thrown.expectMessage("still active");
+ parallel.schedule(first);
+ }
+
+ @Test
+ public void parallelRejectedShutdownSucceeds() {
+ @SuppressWarnings("unchecked")
+ TransformExecutor<Object> first = mock(TransformExecutor.class);
+
+ TransformExecutorService parallel =
+ TransformExecutorServices.parallel(executorService);
+ executorService.shutdown();
+ parallel.shutdown();
+ parallel.schedule(first);
+ }
+
+ @Test
public void serialScheduleTwoWaitsForFirstToComplete() {
@SuppressWarnings("unchecked")
TransformExecutor<Object> first = mock(TransformExecutor.class);
@@ -97,4 +122,27 @@ public class TransformExecutorServicesTest {
serial.complete(second);
}
+
+ /**
+ * Tests that a Serial {@link TransformExecutorService} does not schedule follow up work if the
+ * executor is shut down when the initial work completes.
+ */
+ @Test
+ public void serialShutdownCompleteActive() {
+ @SuppressWarnings("unchecked")
+ TransformExecutor<Object> first = mock(TransformExecutor.class);
+ @SuppressWarnings("unchecked")
+ TransformExecutor<Object> second = mock(TransformExecutor.class);
+
+ TransformExecutorService serial = TransformExecutorServices.serial(executorService);
+ serial.schedule(first);
+ verify(first).run();
+
+ serial.schedule(second);
+ verify(second, never()).run();
+
+ serial.shutdown();
+ serial.complete(first);
+ verify(second, never()).run();
+ }
}
[10/50] [abbrv] beam git commit: This closes #2430
Posted by jb...@apache.org.
This closes #2430
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/7c169a61
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/7c169a61
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/7c169a61
Branch: refs/heads/DSL_SQL
Commit: 7c169a614f339ed829c6d8b6bbc5c4868ff33ef5
Parents: a5a5bf9 0e5737f
Author: Thomas Groh <tg...@google.com>
Authored: Tue Apr 11 12:40:46 2017 -0700
Committer: Thomas Groh <tg...@google.com>
Committed: Tue Apr 11 12:40:46 2017 -0700
----------------------------------------------------------------------
.../FlattenPCollectionTranslator.java | 13 +-
.../apex/translation/ParDoTranslator.java | 13 +-
.../apex/translation/TranslationContext.java | 10 +-
.../DeduplicatedFlattenFactory.java | 18 ++-
.../EmptyFlattenAsCreateFactory.java | 7 +-
.../core/construction/PTransformMatchers.java | 5 +-
.../core/construction/PrimitiveCreate.java | 7 +-
.../core/construction/ReplacementOutputs.java | 63 ++++-----
.../SingleInputOutputOverrideFactory.java | 11 +-
.../UnsupportedOverrideFactory.java | 8 +-
.../DeduplicatedFlattenFactoryTest.java | 6 +-
.../EmptyFlattenAsCreateFactoryTest.java | 8 +-
.../construction/PTransformMatchersTest.java | 132 +++++++++----------
.../construction/ReplacementOutputsTest.java | 109 ++-------------
.../SingleInputOutputOverrideFactoryTest.java | 6 +-
.../UnsupportedOverrideFactoryTest.java | 7 +-
.../direct/BoundedReadEvaluatorFactory.java | 2 +-
.../beam/runners/direct/DirectGraphVisitor.java | 5 +-
.../direct/ExecutorServiceParallelExecutor.java | 4 +-
.../runners/direct/FlattenEvaluatorFactory.java | 2 +-
.../GroupAlsoByWindowEvaluatorFactory.java | 2 +-
.../direct/GroupByKeyOnlyEvaluatorFactory.java | 4 +-
.../direct/KeyedPValueTrackingVisitor.java | 14 +-
.../runners/direct/ParDoEvaluatorFactory.java | 10 +-
.../direct/ParDoMultiOverrideFactory.java | 9 +-
.../direct/StatefulParDoEvaluatorFactory.java | 8 +-
.../direct/TestStreamEvaluatorFactory.java | 8 +-
.../direct/UnboundedReadEvaluatorFactory.java | 4 +-
.../runners/direct/ViewEvaluatorFactory.java | 4 +-
.../runners/direct/ViewOverrideFactory.java | 9 +-
.../beam/runners/direct/WatermarkManager.java | 19 +--
.../runners/direct/WindowEvaluatorFactory.java | 2 +-
.../direct/WriteWithShardingFactory.java | 10 +-
.../runners/direct/DirectGraphVisitorTest.java | 7 +-
.../beam/runners/direct/ParDoEvaluatorTest.java | 2 +-
.../StatefulParDoEvaluatorFactoryTest.java | 2 +-
.../direct/TestStreamEvaluatorFactoryTest.java | 5 +-
.../runners/direct/ViewOverrideFactoryTest.java | 2 +-
.../flink/FlinkBatchTransformTranslators.java | 36 ++---
.../flink/FlinkBatchTranslationContext.java | 11 +-
.../flink/FlinkStreamingPipelineTranslator.java | 9 +-
.../FlinkStreamingTransformTranslators.java | 32 ++---
.../flink/FlinkStreamingTranslationContext.java | 12 +-
.../dataflow/BatchStatefulParDoOverrides.java | 15 +--
.../dataflow/DataflowPipelineTranslator.java | 20 +--
.../beam/runners/dataflow/DataflowRunner.java | 27 ++--
.../runners/dataflow/TransformTranslator.java | 6 +-
.../dataflow/DataflowPipelineJobTest.java | 7 +-
.../apache/beam/runners/spark/SparkRunner.java | 20 +--
.../beam/runners/spark/TestSparkRunner.java | 7 +-
.../spark/translation/EvaluationContext.java | 11 +-
.../spark/translation/TransformTranslator.java | 25 ++--
.../streaming/StreamingTransformTranslator.java | 20 +--
.../sdk/runners/PTransformOverrideFactory.java | 10 +-
.../beam/sdk/runners/TransformHierarchy.java | 80 +++++------
.../beam/sdk/transforms/AppliedPTransform.java | 17 ++-
.../transforms/join/KeyedPCollectionTuple.java | 12 +-
.../java/org/apache/beam/sdk/values/PBegin.java | 6 +-
.../apache/beam/sdk/values/PCollectionList.java | 27 ++--
.../beam/sdk/values/PCollectionTuple.java | 10 +-
.../java/org/apache/beam/sdk/values/PDone.java | 6 +-
.../java/org/apache/beam/sdk/values/PInput.java | 4 +-
.../org/apache/beam/sdk/values/POutput.java | 4 +-
.../java/org/apache/beam/sdk/values/PValue.java | 4 +-
.../org/apache/beam/sdk/values/PValueBase.java | 6 +-
.../apache/beam/sdk/values/TaggedPValue.java | 5 +
.../java/org/apache/beam/sdk/PipelineTest.java | 20 ++-
.../sdk/runners/TransformHierarchyTest.java | 72 +++++-----
.../beam/sdk/values/PCollectionListTest.java | 70 ++++------
.../beam/sdk/values/PCollectionTupleTest.java | 5 +-
.../beam/sdk/io/gcp/bigquery/WriteResult.java | 11 +-
71 files changed, 526 insertions(+), 658 deletions(-)
----------------------------------------------------------------------
[48/50] [abbrv] beam git commit: [BEAM-911] Mark IO APIs as
@Experimental
Posted by jb...@apache.org.
[BEAM-911] Mark IO APIs as @Experimental
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/024b6e05
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/024b6e05
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/024b6e05
Branch: refs/heads/DSL_SQL
Commit: 024b6e05a99b12142a2b10c0612288c6177eff96
Parents: a405a2c
Author: Isma�l Mej�a <ie...@apache.org>
Authored: Wed Apr 12 17:12:05 2017 +0200
Committer: Jean-Baptiste Onofr� <jb...@apache.org>
Committed: Wed Apr 12 21:35:13 2017 +0200
----------------------------------------------------------------------
.../org/apache/beam/sdk/io/elasticsearch/ElasticsearchIO.java | 2 ++
.../beam/sdk/io/hadoop/inputformat/HadoopInputFormatIO.java | 3 ++-
.../src/main/java/org/apache/beam/sdk/io/hdfs/HDFSFileSink.java | 2 ++
.../src/main/java/org/apache/beam/sdk/io/hdfs/HDFSFileSource.java | 2 ++
.../io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java | 2 ++
.../io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java | 2 ++
.../kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java | 2 ++
.../src/main/java/org/apache/beam/sdk/io/kinesis/KinesisIO.java | 2 ++
.../main/java/org/apache/beam/sdk/io/mongodb/MongoDbGridFSIO.java | 2 ++
.../src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java | 2 ++
.../io/mqtt/src/main/java/org/apache/beam/sdk/io/mqtt/MqttIO.java | 2 ++
11 files changed, 22 insertions(+), 1 deletion(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/024b6e05/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 54692b9..baf0cc2 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
@@ -41,6 +41,7 @@ import java.util.Map;
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.StringUtf8Coder;
import org.apache.beam.sdk.io.BoundedSource;
@@ -114,6 +115,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
public class ElasticsearchIO {
public static Read read() {
http://git-wip-us.apache.org/repos/asf/beam/blob/024b6e05/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 61dc1bf..d776ea0 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
@@ -42,6 +42,7 @@ import java.util.concurrent.atomic.AtomicLong;
import javax.annotation.Nullable;
+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.CoderException;
@@ -170,7 +171,7 @@ import org.slf4j.LoggerFactory;
* }
* </pre>
*/
-
+@Experimental
public class HadoopInputFormatIO {
private static final Logger LOG = LoggerFactory.getLogger(HadoopInputFormatIO.class);
http://git-wip-us.apache.org/repos/asf/beam/blob/024b6e05/sdks/java/io/hdfs/src/main/java/org/apache/beam/sdk/io/hdfs/HDFSFileSink.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hdfs/src/main/java/org/apache/beam/sdk/io/hdfs/HDFSFileSink.java b/sdks/java/io/hdfs/src/main/java/org/apache/beam/sdk/io/hdfs/HDFSFileSink.java
index 10ff788..15d61cb 100644
--- a/sdks/java/io/hdfs/src/main/java/org/apache/beam/sdk/io/hdfs/HDFSFileSink.java
+++ b/sdks/java/io/hdfs/src/main/java/org/apache/beam/sdk/io/hdfs/HDFSFileSink.java
@@ -33,6 +33,7 @@ import org.apache.avro.Schema;
import org.apache.avro.generic.GenericRecord;
import org.apache.avro.mapred.AvroKey;
import org.apache.avro.mapreduce.AvroKeyOutputFormat;
+import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.coders.AvroCoder;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
@@ -92,6 +93,7 @@ import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl;
* @param <V> the type of values to be written to the sink via {@link FileOutputFormat}.
*/
@AutoValue
+@Experimental
public abstract class HDFSFileSink<T, K, V> extends Sink<T> {
private static final JobID jobId = new JobID(
http://git-wip-us.apache.org/repos/asf/beam/blob/024b6e05/sdks/java/io/hdfs/src/main/java/org/apache/beam/sdk/io/hdfs/HDFSFileSource.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hdfs/src/main/java/org/apache/beam/sdk/io/hdfs/HDFSFileSource.java b/sdks/java/io/hdfs/src/main/java/org/apache/beam/sdk/io/hdfs/HDFSFileSource.java
index e317c6e..b55944b 100644
--- a/sdks/java/io/hdfs/src/main/java/org/apache/beam/sdk/io/hdfs/HDFSFileSource.java
+++ b/sdks/java/io/hdfs/src/main/java/org/apache/beam/sdk/io/hdfs/HDFSFileSource.java
@@ -39,6 +39,7 @@ import org.apache.avro.Schema;
import org.apache.avro.generic.GenericRecord;
import org.apache.avro.mapred.AvroKey;
import org.apache.avro.mapreduce.AvroKeyInputFormat;
+import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.coders.AvroCoder;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderException;
@@ -104,6 +105,7 @@ import org.slf4j.LoggerFactory;
* @param <V> the type of values to be read from the source via {@link FileInputFormat}.
*/
@AutoValue
+@Experimental
public abstract class HDFSFileSource<T, K, V> extends BoundedSource<T> {
private static final long serialVersionUID = 0L;
http://git-wip-us.apache.org/repos/asf/beam/blob/024b6e05/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 dc31130..8fdbeb0 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
@@ -31,6 +31,7 @@ import java.util.Random;
import javax.annotation.Nullable;
import javax.sql.DataSource;
+import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
@@ -135,6 +136,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
public class JdbcIO {
/**
* Read data from a JDBC datasource.
http://git-wip-us.apache.org/repos/asf/beam/blob/024b6e05/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 a935b56..89016ac 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
@@ -37,6 +37,7 @@ import javax.jms.MessageConsumer;
import javax.jms.MessageProducer;
import javax.jms.Session;
import javax.jms.TextMessage;
+import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.coders.AvroCoder;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.SerializableCoder;
@@ -97,6 +98,7 @@ import org.slf4j.LoggerFactory;
*
* }</pre>
*/
+@Experimental
public class JmsIO {
private static final Logger LOG = LoggerFactory.getLogger(JmsIO.class);
http://git-wip-us.apache.org/repos/asf/beam/blob/024b6e05/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 80b40be..69d82bc 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
@@ -52,6 +52,7 @@ import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import javax.annotation.Nullable;
+import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.coders.AtomicCoder;
import org.apache.beam.sdk.coders.AvroCoder;
import org.apache.beam.sdk.coders.ByteArrayCoder;
@@ -211,6 +212,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
public class KafkaIO {
/**
* Creates an uninitialized {@link Read} {@link PTransform}. Before use, basic Kafka
http://git-wip-us.apache.org/repos/asf/beam/blob/024b6e05/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 91939d4..45a7b2d 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
@@ -27,6 +27,7 @@ import com.amazonaws.regions.Regions;
import com.amazonaws.services.kinesis.AmazonKinesis;
import com.amazonaws.services.kinesis.AmazonKinesisClient;
import com.amazonaws.services.kinesis.clientlibrary.lib.worker.InitialPositionInStream;
+import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.transforms.PTransform;
import org.joda.time.Instant;
@@ -99,6 +100,7 @@ import org.joda.time.Instant;
* }</pre>
*
*/
+@Experimental
public final class KinesisIO {
/**
* A {@link PTransform} that reads from a Kinesis stream.
http://git-wip-us.apache.org/repos/asf/beam/blob/024b6e05/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 d924c14..919fda3 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
@@ -43,6 +43,7 @@ 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.coders.StringUtf8Coder;
@@ -119,6 +120,7 @@ import org.joda.time.Instant;
* to the file separated with line feeds.
* </p>
*/
+@Experimental
public class MongoDbGridFSIO {
/**
http://git-wip-us.apache.org/repos/asf/beam/blob/024b6e05/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 bd262e7..09b8505 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
@@ -30,6 +30,7 @@ import com.mongodb.client.MongoDatabase;
import java.util.ArrayList;
import java.util.List;
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.BoundedSource;
@@ -91,6 +92,7 @@ import org.slf4j.LoggerFactory;
*
* }</pre>
*/
+@Experimental
public class MongoDbIO {
private static final Logger LOG = LoggerFactory.getLogger(MongoDbIO.class);
http://git-wip-us.apache.org/repos/asf/beam/blob/024b6e05/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 46f2dcc..820b265 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
@@ -31,6 +31,7 @@ import java.util.UUID;
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.SerializableCoder;
@@ -99,6 +100,7 @@ import org.slf4j.LoggerFactory;
*
* }</pre>
*/
+@Experimental
public class MqttIO {
private static final Logger LOG = LoggerFactory.getLogger(MqttIO.class);
[45/50] [abbrv] beam git commit: Revert "Throw specialized exception
in value providers"
Posted by jb...@apache.org.
Revert "Throw specialized exception in value providers"
This reverts commit ee92b9642bb6b6e42bb701ab638c55539163bb69.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/4854291f
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/4854291f
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/4854291f
Branch: refs/heads/DSL_SQL
Commit: 4854291fe38df7b624446311956d0ae0677f6b8c
Parents: b7ddab2
Author: Ahmet Altay <al...@google.com>
Authored: Wed Apr 12 10:43:24 2017 -0700
Committer: Ahmet Altay <al...@google.com>
Committed: Wed Apr 12 11:50:39 2017 -0700
----------------------------------------------------------------------
.../apache_beam/runners/dataflow/dataflow_runner.py | 6 ------
sdks/python/apache_beam/utils/value_provider.py | 14 ++------------
2 files changed, 2 insertions(+), 18 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/4854291f/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 4e81788..1a935c1 100644
--- a/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
+++ b/sdks/python/apache_beam/runners/dataflow/dataflow_runner.py
@@ -43,7 +43,6 @@ from apache_beam.runners.runner import PipelineState
from apache_beam.transforms.display import DisplayData
from apache_beam.typehints import typehints
from apache_beam.utils.pipeline_options import StandardOptions
-from apache_beam.utils.value_provider import RuntimeValueProviderError
class DataflowRunner(PipelineRunner):
@@ -477,11 +476,6 @@ class DataflowRunner(PipelineRunner):
'estimated_size_bytes': json_value.get_typed_value_descriptor(
transform.source.estimate_size())
}
- except RuntimeValueProviderError:
- # Size estimation is best effort, and this error is by value provider.
- logging.info(
- 'Could not estimate size of source %r due to ' + \
- 'RuntimeValueProviderError', transform.source)
except Exception: # pylint: disable=broad-except
# Size estimation is best effort. So we log the error and continue.
logging.info(
http://git-wip-us.apache.org/repos/asf/beam/blob/4854291f/sdks/python/apache_beam/utils/value_provider.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/utils/value_provider.py b/sdks/python/apache_beam/utils/value_provider.py
index 271202d..a72fc4c 100644
--- a/sdks/python/apache_beam/utils/value_provider.py
+++ b/sdks/python/apache_beam/utils/value_provider.py
@@ -22,15 +22,6 @@ and dynamically provided values.
from functools import wraps
-class RuntimeValueProviderError(RuntimeError):
- def __init__(self, msg):
- """Class representing the errors thrown during runtime by the valueprovider
- Args:
- msg: Message string for the exception thrown
- """
- super(RuntimeValueProviderError, self).__init__(msg)
-
-
class ValueProvider(object):
def is_accessible(self):
raise NotImplementedError(
@@ -76,8 +67,7 @@ class RuntimeValueProvider(ValueProvider):
runtime_options = (
RuntimeValueProvider.runtime_options_map.get(self.options_id))
if runtime_options is None:
- raise RuntimeValueProviderError(
- '%s.get() not called from a runtime context' % self)
+ raise RuntimeError('%s.get() not called from a runtime context' % self)
candidate = runtime_options.get(self.option_name)
if candidate:
@@ -114,7 +104,7 @@ def check_accessible(value_provider_list):
def _f(self, *args, **kwargs):
for obj in [getattr(self, vp) for vp in value_provider_list]:
if not obj.is_accessible():
- raise RuntimeValueProviderError('%s not accessible' % obj)
+ raise RuntimeError('%s not accessible' % obj)
return fnc(self, *args, **kwargs)
return _f
return _check_accessible
[09/50] [abbrv] beam git commit: Closes #2494
Posted by jb...@apache.org.
Closes #2494
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/a5a5bf94
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/a5a5bf94
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/a5a5bf94
Branch: refs/heads/DSL_SQL
Commit: a5a5bf946f69d75897279f66aeeb0dcea3babfb2
Parents: 77712c9 a168357
Author: Robert Bradshaw <ro...@gmail.com>
Authored: Tue Apr 11 12:12:52 2017 -0700
Committer: Robert Bradshaw <ro...@gmail.com>
Committed: Tue Apr 11 12:12:52 2017 -0700
----------------------------------------------------------------------
.gitignore | 3 +++
1 file changed, 3 insertions(+)
----------------------------------------------------------------------
[22/50] [abbrv] beam git commit: DataflowRunner: remove dead code
Posted by jb...@apache.org.
DataflowRunner: remove dead code
It is not possible to try to use an unbounded source in batch mode any more, as the
runner will automatically enable streaming mode.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/f4736729
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/f4736729
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/f4736729
Branch: refs/heads/DSL_SQL
Commit: f4736729a5a04dc9d1cb30a44b2dd4692be7e935
Parents: 132d4c5
Author: Dan Halperin <dh...@google.com>
Authored: Tue Apr 11 14:12:43 2017 -0700
Committer: Dan Halperin <dh...@google.com>
Committed: Tue Apr 11 15:39:40 2017 -0700
----------------------------------------------------------------------
.../beam/runners/dataflow/DataflowRunner.java | 16 ----------------
1 file changed, 16 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/f4736729/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 1708d8a..8726635 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
@@ -39,7 +39,6 @@ import com.google.common.base.Strings;
import com.google.common.base.Utf8;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import java.io.File;
import java.io.IOException;
@@ -65,7 +64,6 @@ import org.apache.beam.runners.core.construction.PTransformMatchers;
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.UnsupportedOverrideFactory;
import org.apache.beam.runners.dataflow.BatchViewOverrides.BatchCombineGloballyAsSingletonViewFactory;
import org.apache.beam.runners.dataflow.DataflowPipelineTranslator.JobSpecification;
import org.apache.beam.runners.dataflow.StreamingViewOverrides.StreamingCreatePCollectionViewFactory;
@@ -330,14 +328,6 @@ public class DataflowRunner extends PipelineRunner<DataflowPipelineJob> {
PTransformMatchers.classEqualTo(View.CreatePCollectionView.class),
new StreamingCreatePCollectionViewFactory()));
} else {
- // In batch mode must use the custom Pubsub bounded source/sink.
- for (Class<? extends PTransform> unsupported :
- ImmutableSet.of(PubsubUnboundedSink.class, PubsubUnboundedSource.class)) {
- overridesBuilder.add(
- PTransformOverride.of(
- PTransformMatchers.classEqualTo(unsupported),
- UnsupportedOverrideFactory.withMessage(getUnsupportedMessage(unsupported, false))));
- }
overridesBuilder
// State and timer pardos are implemented by expansion to GBK-then-ParDo
.add(
@@ -399,12 +389,6 @@ public class DataflowRunner extends PipelineRunner<DataflowPipelineJob> {
return overridesBuilder.build();
}
- private String getUnsupportedMessage(Class<?> unsupported, boolean streaming) {
- return String.format(
- "%s is not supported in %s",
- NameUtils.approximateSimpleName(unsupported), streaming ? "streaming" : "batch");
- }
-
private static class ReflectiveOneToOneOverrideFactory<
InputT extends PValue,
OutputT extends PValue,
[23/50] [abbrv] beam git commit: This closes #2498
Posted by jb...@apache.org.
This closes #2498
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/a2d328a2
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/a2d328a2
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/a2d328a2
Branch: refs/heads/DSL_SQL
Commit: a2d328a2d5d22def7901eb1d814118a36d58b5b7
Parents: 132d4c5 f473672
Author: Dan Halperin <dh...@google.com>
Authored: Tue Apr 11 15:39:45 2017 -0700
Committer: Dan Halperin <dh...@google.com>
Committed: Tue Apr 11 15:39:45 2017 -0700
----------------------------------------------------------------------
.../beam/runners/dataflow/DataflowRunner.java | 16 ----------------
1 file changed, 16 deletions(-)
----------------------------------------------------------------------
[35/50] [abbrv] beam git commit: This closes #2503
Posted by jb...@apache.org.
This closes #2503
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/7e603d5c
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/7e603d5c
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/7e603d5c
Branch: refs/heads/DSL_SQL
Commit: 7e603d5c763ee4b95dffb94e129603de55e5c33d
Parents: 82f2f2c 5bcb8c5
Author: Dan Halperin <dh...@google.com>
Authored: Wed Apr 12 08:16:47 2017 -0700
Committer: Dan Halperin <dh...@google.com>
Committed: Wed Apr 12 08:16:47 2017 -0700
----------------------------------------------------------------------
.../beam/examples/complete/game/GameStats.java | 2 +-
.../examples/complete/game/LeaderBoard.java | 2 +-
.../triggers/AfterWatermarkStateMachine.java | 14 +-
runners/google-cloud-dataflow-java/pom.xml | 5 +-
.../beam/runners/dataflow/DataflowRunner.java | 4 +-
sdks/java/core/pom.xml | 44 -
.../java/org/apache/beam/sdk/io/PubsubIO.java | 1016 ------------
.../apache/beam/sdk/io/PubsubUnboundedSink.java | 494 ------
.../beam/sdk/io/PubsubUnboundedSource.java | 1463 ------------------
.../apache/beam/sdk/transforms/GroupByKey.java | 4 +-
.../transforms/windowing/AfterWatermark.java | 14 +-
.../org/apache/beam/sdk/util/PubsubClient.java | 544 -------
.../apache/beam/sdk/util/PubsubGrpcClient.java | 424 -----
.../apache/beam/sdk/util/PubsubJsonClient.java | 317 ----
.../apache/beam/sdk/util/PubsubTestClient.java | 436 ------
.../org/apache/beam/sdk/util/Transport.java | 3 +-
.../org/apache/beam/sdk/io/PubsubIOTest.java | 189 ---
.../beam/sdk/io/PubsubUnboundedSinkTest.java | 190 ---
.../beam/sdk/io/PubsubUnboundedSourceTest.java | 411 -----
.../apache/beam/sdk/util/PubsubClientTest.java | 189 ---
.../beam/sdk/util/PubsubGrpcClientTest.java | 207 ---
.../beam/sdk/util/PubsubJsonClientTest.java | 140 --
.../beam/sdk/util/PubsubTestClientTest.java | 114 --
sdks/java/io/google-cloud-platform/pom.xml | 44 +
.../beam/sdk/io/gcp/pubsub/PubsubClient.java | 544 +++++++
.../sdk/io/gcp/pubsub/PubsubGrpcClient.java | 424 +++++
.../apache/beam/sdk/io/gcp/pubsub/PubsubIO.java | 1014 ++++++++++++
.../sdk/io/gcp/pubsub/PubsubJsonClient.java | 319 ++++
.../sdk/io/gcp/pubsub/PubsubTestClient.java | 436 ++++++
.../sdk/io/gcp/pubsub/PubsubUnboundedSink.java | 490 ++++++
.../io/gcp/pubsub/PubsubUnboundedSource.java | 1463 ++++++++++++++++++
.../beam/sdk/io/gcp/pubsub/package-info.java | 24 +
.../beam/sdk/io/gcp/GcpApiSurfaceTest.java | 5 +-
.../sdk/io/gcp/pubsub/PubsubClientTest.java | 189 +++
.../sdk/io/gcp/pubsub/PubsubGrpcClientTest.java | 208 +++
.../beam/sdk/io/gcp/pubsub/PubsubIOTest.java | 189 +++
.../sdk/io/gcp/pubsub/PubsubJsonClientTest.java | 139 ++
.../sdk/io/gcp/pubsub/PubsubTestClientTest.java | 114 ++
.../io/gcp/pubsub/PubsubUnboundedSinkTest.java | 188 +++
.../gcp/pubsub/PubsubUnboundedSourceTest.java | 409 +++++
40 files changed, 6218 insertions(+), 6207 deletions(-)
----------------------------------------------------------------------
[07/50] [abbrv] beam git commit: Fix build breaks caused by overlaps
between b615013 and c08b7b1
Posted by jb...@apache.org.
Fix build breaks caused by overlaps between b615013 and c08b7b1
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/732610d8
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/732610d8
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/732610d8
Branch: refs/heads/DSL_SQL
Commit: 732610d8eb5c682934d6a1bec0c5c8d8f25c635d
Parents: 38cf23e
Author: Stephen Sisk <si...@google.com>
Authored: Tue Apr 11 11:07:48 2017 -0700
Committer: Thomas Groh <tg...@google.com>
Committed: Tue Apr 11 11:12:13 2017 -0700
----------------------------------------------------------------------
.../io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/732610d8/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java b/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java
index 97addcf..402c19c 100644
--- a/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java
+++ b/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java
@@ -26,7 +26,7 @@ import com.datastax.driver.mapping.annotations.Table;
import java.io.Serializable;
-import org.apache.beam.sdk.io.hadoop.inputformat.hashing.HashingFn;
+import org.apache.beam.sdk.io.common.HashingFn;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Combine;
@@ -78,7 +78,7 @@ public class HIFIOWithEmbeddedCassandraTest implements Serializable {
@Test
public void testHIFReadForCassandra() throws Exception {
// Expected hashcode is evaluated during insertion time one time and hardcoded here.
- String expectedHashCode = "4651110ba1ef2cd3a7315091ca27877b18fceb0e";
+ String expectedHashCode = "1b9780833cce000138b9afa25ba63486";
Configuration conf = getConfiguration();
PCollection<KV<Long, String>> cassandraData =
p.apply(HadoopInputFormatIO.<Long, String>read().withConfiguration(conf)
@@ -109,7 +109,7 @@ public class HIFIOWithEmbeddedCassandraTest implements Serializable {
@Test
public void testHIFReadForCassandraQuery() throws Exception {
Long expectedCount = 1L;
- String expectedChecksum = "6a62f24ccce0713004889aec1cf226949482d188";
+ String expectedChecksum = "f11caabc7a9fc170e22b41218749166c";
Configuration conf = getConfiguration();
conf.set("cassandra.input.cql", "select * from " + CASSANDRA_KEYSPACE + "." + CASSANDRA_TABLE
+ " where token(id) > ? and token(id) <= ? and scientist='Faraday1' allow filtering");
[26/50] [abbrv] beam git commit: Triggers: handle missing case
Posted by jb...@apache.org.
Triggers: handle missing case
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/e7f43372
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/e7f43372
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/e7f43372
Branch: refs/heads/DSL_SQL
Commit: e7f433720e229dcc296d279c459ed65c08fbb531
Parents: ccf247c
Author: Dan Halperin <dh...@google.com>
Authored: Tue Apr 11 16:26:41 2017 -0700
Committer: Dan Halperin <dh...@google.com>
Committed: Tue Apr 11 16:26:41 2017 -0700
----------------------------------------------------------------------
.../java/org/apache/beam/sdk/transforms/windowing/Triggers.java | 2 ++
1 file changed, 2 insertions(+)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/e7f43372/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/Triggers.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/Triggers.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/Triggers.java
index ec82114..591af37 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/Triggers.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/Triggers.java
@@ -287,6 +287,8 @@ public class Triggers implements Serializable {
return trigger;
case AFTER_SYNCHRONIZED_PROCESSING_TIME:
return AfterSynchronizedProcessingTime.ofFirstElement();
+ case ALWAYS:
+ return new ReshuffleTrigger();
case ELEMENT_COUNT:
return AfterPane.elementCountAtLeast(triggerProto.getElementCount().getElementCount());
case NEVER:
[37/50] [abbrv] beam git commit: This closes #2490
Posted by jb...@apache.org.
This closes #2490
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/571631a5
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/571631a5
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/571631a5
Branch: refs/heads/DSL_SQL
Commit: 571631a5ef83ef31ea9f652c5698cea3cc95ad43
Parents: 7e603d5 e0fd0a2
Author: Dan Halperin <dh...@google.com>
Authored: Wed Apr 12 09:20:46 2017 -0700
Committer: Dan Halperin <dh...@google.com>
Committed: Wed Apr 12 09:20:46 2017 -0700
----------------------------------------------------------------------
.../beam/runners/core/construction/Coders.java | 6 +-
.../core/construction/SdkComponents.java | 5 +-
.../core/construction/WindowingStrategies.java | 94 ++++++++------------
.../core/construction/SdkComponentsTest.java | 6 +-
.../construction/WindowingStrategiesTest.java | 18 ++++
5 files changed, 66 insertions(+), 63 deletions(-)
----------------------------------------------------------------------
[02/50] [abbrv] beam git commit: Improve HadoopInputFormatIO
DisplayData and Cassandra tests
Posted by jb...@apache.org.
Improve HadoopInputFormatIO DisplayData and Cassandra tests
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/c08b7b17
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/c08b7b17
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/c08b7b17
Branch: refs/heads/DSL_SQL
Commit: c08b7b1771481b77c94ec78a96db5b34fec29841
Parents: 8beea73
Author: Dipti Kulkarni <di...@persistent.co.in>
Authored: Mon Apr 10 15:43:12 2017 +0530
Committer: Dan Halperin <dh...@google.com>
Committed: Tue Apr 11 10:20:15 2017 -0700
----------------------------------------------------------------------
.../hadoop/inputformat/HadoopInputFormatIO.java | 33 +--
.../inputformat/HadoopInputFormatIOTest.java | 99 ++++++---
sdks/java/io/hadoop/jdk1.8-tests/.toDelete | 0
.../HIFIOWithEmbeddedCassandraTest.java | 214 +++++++++++++++++++
4 files changed, 306 insertions(+), 40 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/c08b7b17/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 675f4bf..61dc1bf 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
@@ -36,9 +36,7 @@ import java.math.BigInteger;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
-import java.util.Iterator;
import java.util.List;
-import java.util.Map.Entry;
import java.util.Set;
import java.util.concurrent.atomic.AtomicLong;
@@ -360,20 +358,6 @@ public class HadoopInputFormatIO {
+ e.getMessage(), e);
}
}
-
- @Override
- public void populateDisplayData(DisplayData.Builder builder) {
- super.populateDisplayData(builder);
- if (getConfiguration().getHadoopConfiguration() != null) {
- Iterator<Entry<String, String>> configProperties = getConfiguration()
- .getHadoopConfiguration().iterator();
- while (configProperties.hasNext()) {
- Entry<String, String> property = configProperties.next();
- builder.addIfNotNull(DisplayData.item(property.getKey(), property.getValue())
- .withLabel(property.getKey()));
- }
- }
- }
}
/**
@@ -447,6 +431,23 @@ public class HadoopInputFormatIO {
}
@Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ super.populateDisplayData(builder);
+ Configuration hadoopConfig = getConfiguration().getHadoopConfiguration();
+ if (hadoopConfig != null) {
+ builder.addIfNotNull(DisplayData.item("mapreduce.job.inputformat.class",
+ hadoopConfig.get("mapreduce.job.inputformat.class"))
+ .withLabel("InputFormat Class"));
+ builder.addIfNotNull(DisplayData.item("key.class",
+ hadoopConfig.get("key.class"))
+ .withLabel("Key Class"));
+ builder.addIfNotNull(DisplayData.item("value.class",
+ hadoopConfig.get("value.class"))
+ .withLabel("Value Class"));
+ }
+ }
+
+ @Override
public List<BoundedSource<KV<K, V>>> splitIntoBundles(long desiredBundleSizeBytes,
PipelineOptions options) throws Exception {
// desiredBundleSizeBytes is not being considered as splitting based on this
http://git-wip-us.apache.org/repos/asf/beam/blob/c08b7b17/sdks/java/io/hadoop/input-format/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HadoopInputFormatIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hadoop/input-format/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HadoopInputFormatIOTest.java b/sdks/java/io/hadoop/input-format/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HadoopInputFormatIOTest.java
index 2f2857b..3a4a99d 100644
--- a/sdks/java/io/hadoop/input-format/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HadoopInputFormatIOTest.java
+++ b/sdks/java/io/hadoop/input-format/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HadoopInputFormatIOTest.java
@@ -21,9 +21,7 @@ import static org.junit.Assert.assertThat;
import java.io.IOException;
import java.util.ArrayList;
-import java.util.Iterator;
import java.util.List;
-import java.util.Map.Entry;
import org.apache.beam.sdk.coders.AvroCoder;
import org.apache.beam.sdk.coders.Coder;
@@ -157,24 +155,6 @@ public class HadoopInputFormatIOTest {
}
/**
- * This test validates functionality of {@link HadoopInputFormatIO.Read#populateDisplayData()
- * populateDisplayData()}.
- */
- @Test
- public void testReadDisplayData() {
- HadoopInputFormatIO.Read<String, String> read = HadoopInputFormatIO.<String, String>read()
- .withConfiguration(serConf.getHadoopConfiguration())
- .withKeyTranslation(myKeyTranslate)
- .withValueTranslation(myValueTranslate);
- DisplayData displayData = DisplayData.from(read);
- Iterator<Entry<String, String>> propertyElement = serConf.getHadoopConfiguration().iterator();
- while (propertyElement.hasNext()) {
- Entry<String, String> element = propertyElement.next();
- assertThat(displayData, hasDisplayItem(element.getKey(), element.getValue()));
- }
- }
-
- /**
* This test validates {@link HadoopInputFormatIO.Read Read} transform object creation fails with
* null configuration. {@link HadoopInputFormatIO.Read#withConfiguration() withConfiguration()}
* method checks configuration is null and throws exception if it is null.
@@ -415,6 +395,32 @@ public class HadoopInputFormatIOTest {
}
/**
+ * This test validates functionality of
+ * {@link HadoopInputFormatIO.HadoopInputFormatBoundedSource#populateDisplayData()
+ * populateDisplayData()}.
+ */
+ @Test
+ public void testReadDisplayData() {
+ HadoopInputFormatBoundedSource<Text, Employee> boundedSource =
+ new HadoopInputFormatBoundedSource<Text, Employee>(
+ serConf,
+ WritableCoder.of(Text.class),
+ AvroCoder.of(Employee.class),
+ null, // No key translation required.
+ null, // No value translation required.
+ new SerializableSplit());
+ DisplayData displayData = DisplayData.from(boundedSource);
+ assertThat(
+ displayData,
+ hasDisplayItem("mapreduce.job.inputformat.class",
+ serConf.getHadoopConfiguration().get("mapreduce.job.inputformat.class")));
+ assertThat(displayData,
+ hasDisplayItem("key.class", serConf.getHadoopConfiguration().get("key.class")));
+ assertThat(displayData,
+ hasDisplayItem("value.class", serConf.getHadoopConfiguration().get("value.class")));
+ }
+
+ /**
* This test validates behavior of {@link HadoopInputFormatBoundedSource} if RecordReader object
* creation fails.
*/
@@ -472,7 +478,8 @@ public class HadoopInputFormatIOTest {
*/
@Test
public void testReadersStartWhenZeroRecords() throws Exception {
- InputFormat<Text, Employee> mockInputFormat = Mockito.mock(EmployeeInputFormat.class);
+
+ InputFormat mockInputFormat = Mockito.mock(EmployeeInputFormat.class);
EmployeeRecordReader mockReader = Mockito.mock(EmployeeRecordReader.class);
Mockito.when(
mockInputFormat.createRecordReader(Mockito.any(InputSplit.class),
@@ -487,9 +494,11 @@ public class HadoopInputFormatIOTest {
null, // No key translation required.
null, // No value translation required.
new SerializableSplit(mockInputSplit));
- BoundedReader<KV<Text, Employee>> boundedReader = boundedSource.createReader(p.getOptions());
- assertEquals(false, boundedReader.start());
- assertEquals(Double.valueOf(1), boundedReader.getFractionConsumed());
+ boundedSource.setInputFormatObj(mockInputFormat);
+ BoundedReader<KV<Text, Employee>> reader = boundedSource.createReader(p.getOptions());
+ assertEquals(false, reader.start());
+ assertEquals(Double.valueOf(1), reader.getFractionConsumed());
+ reader.close();
}
/**
@@ -546,6 +555,48 @@ public class HadoopInputFormatIOTest {
assertThat(bundleRecords, containsInAnyOrder(referenceRecords.toArray()));
}
+/**
+ * This test validates the method getFractionConsumed()- when a bad progress value is returned by
+ * the inputformat.
+ */
+ @Test
+ public void testGetFractionConsumedForBadProgressValue() throws Exception {
+ InputFormat<Text, Employee> mockInputFormat = Mockito.mock(EmployeeInputFormat.class);
+ EmployeeRecordReader mockReader = Mockito.mock(EmployeeRecordReader.class);
+ Mockito.when(
+ mockInputFormat.createRecordReader(Mockito.any(InputSplit.class),
+ Mockito.any(TaskAttemptContext.class))).thenReturn(mockReader);
+ Mockito.when(mockReader.nextKeyValue()).thenReturn(true);
+ // Set to a bad value , not in range of 0 to 1
+ Mockito.when(mockReader.getProgress()).thenReturn(2.0F);
+ InputSplit mockInputSplit = Mockito.mock(NewObjectsEmployeeInputSplit.class);
+ HadoopInputFormatBoundedSource<Text, Employee> boundedSource =
+ new HadoopInputFormatBoundedSource<Text, Employee>(
+ serConf,
+ WritableCoder.of(Text.class),
+ AvroCoder.of(Employee.class),
+ null, // No key translation required.
+ null, // No value translation required.
+ new SerializableSplit(mockInputSplit));
+ boundedSource.setInputFormatObj(mockInputFormat);
+ BoundedReader<KV<Text, Employee>> reader = boundedSource.createReader(p.getOptions());
+ assertEquals(Double.valueOf(0), reader.getFractionConsumed());
+ boolean start = reader.start();
+ assertEquals(true, start);
+ if (start) {
+ boolean advance = reader.advance();
+ assertEquals(null, reader.getFractionConsumed());
+ assertEquals(true, advance);
+ if (advance) {
+ advance = reader.advance();
+ assertEquals(null, reader.getFractionConsumed());
+ }
+ }
+ // Validate if getFractionConsumed() returns null after few number of reads as getProgress
+ // returns invalid value '2' which is not in the range of 0 to 1.
+ assertEquals(null, reader.getFractionConsumed());
+ reader.close();
+ }
/**
* This test validates that reader and its parent source reads the same records.
*/
http://git-wip-us.apache.org/repos/asf/beam/blob/c08b7b17/sdks/java/io/hadoop/jdk1.8-tests/.toDelete
----------------------------------------------------------------------
diff --git a/sdks/java/io/hadoop/jdk1.8-tests/.toDelete b/sdks/java/io/hadoop/jdk1.8-tests/.toDelete
new file mode 100644
index 0000000..e69de29
http://git-wip-us.apache.org/repos/asf/beam/blob/c08b7b17/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java b/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java
new file mode 100644
index 0000000..97addcf
--- /dev/null
+++ b/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java
@@ -0,0 +1,214 @@
+/*
+ * 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.hadoop.inputformat;
+
+import com.datastax.driver.core.Cluster;
+import com.datastax.driver.core.Row;
+import com.datastax.driver.core.Session;
+import com.datastax.driver.core.SocketOptions;
+import com.datastax.driver.mapping.annotations.Column;
+import com.datastax.driver.mapping.annotations.Table;
+
+import java.io.Serializable;
+
+import org.apache.beam.sdk.io.hadoop.inputformat.hashing.HashingFn;
+import org.apache.beam.sdk.testing.PAssert;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Combine;
+import org.apache.beam.sdk.transforms.Count;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.transforms.Values;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.cassandra.service.EmbeddedCassandraService;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.mapreduce.InputFormat;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Tests to validate HadoopInputFormatIO for embedded Cassandra instance.
+ */
+@RunWith(JUnit4.class)
+public class HIFIOWithEmbeddedCassandraTest implements Serializable {
+ private static final long serialVersionUID = 1L;
+ private static final String CASSANDRA_KEYSPACE = "beamdb";
+ private static final String CASSANDRA_HOST = "127.0.0.1";
+ private static final String CASSANDRA_TABLE = "scientists";
+ private static final String CASSANDRA_THRIFT_PORT_PROPERTY = "cassandra.input.thrift.port";
+ private static final String CASSANDRA_THRIFT_ADDRESS_PROPERTY = "cassandra.input.thrift.address";
+ private static final String CASSANDRA_PARTITIONER_CLASS_PROPERTY =
+ "cassandra.input.partitioner.class";
+ private static final String CASSANDRA_PARTITIONER_CLASS_VALUE = "Murmur3Partitioner";
+ private static final String CASSANDRA_KEYSPACE_PROPERTY = "cassandra.input.keyspace";
+ private static final String CASSANDRA_COLUMNFAMILY_PROPERTY = "cassandra.input.columnfamily";
+ private static final String CASSANDRA_PORT = "9061";
+ private static transient Cluster cluster;
+ private static transient Session session;
+ private static final long TEST_DATA_ROW_COUNT = 10L;
+ private static EmbeddedCassandraService cassandra = new EmbeddedCassandraService();
+
+ @Rule
+ public final transient TestPipeline p = TestPipeline.create();
+
+ /**
+ * Test to read data from embedded Cassandra instance and verify whether data is read
+ * successfully.
+ * @throws Exception
+ */
+ @Test
+ public void testHIFReadForCassandra() throws Exception {
+ // Expected hashcode is evaluated during insertion time one time and hardcoded here.
+ String expectedHashCode = "4651110ba1ef2cd3a7315091ca27877b18fceb0e";
+ Configuration conf = getConfiguration();
+ PCollection<KV<Long, String>> cassandraData =
+ p.apply(HadoopInputFormatIO.<Long, String>read().withConfiguration(conf)
+ .withValueTranslation(myValueTranslate));
+ // Verify the count of data retrieved from Cassandra matches expected count.
+ PAssert.thatSingleton(cassandraData.apply("Count", Count.<KV<Long, String>>globally()))
+ .isEqualTo(TEST_DATA_ROW_COUNT);
+ PCollection<String> textValues = cassandraData.apply(Values.<String>create());
+ // Verify the output values using checksum comparison.
+ PCollection<String> consolidatedHashcode =
+ textValues.apply(Combine.globally(new HashingFn()).withoutDefaults());
+ PAssert.that(consolidatedHashcode).containsInAnyOrder(expectedHashCode);
+ p.run().waitUntilFinish();
+ }
+
+ SimpleFunction<Row, String> myValueTranslate = new SimpleFunction<Row, String>() {
+ @Override
+ public String apply(Row input) {
+ String scientistRecord = input.getInt("id") + "|" + input.getString("scientist");
+ return scientistRecord;
+ }
+ };
+
+ /**
+ * Test to read data from embedded Cassandra instance based on query and verify whether data is
+ * read successfully.
+ */
+ @Test
+ public void testHIFReadForCassandraQuery() throws Exception {
+ Long expectedCount = 1L;
+ String expectedChecksum = "6a62f24ccce0713004889aec1cf226949482d188";
+ Configuration conf = getConfiguration();
+ conf.set("cassandra.input.cql", "select * from " + CASSANDRA_KEYSPACE + "." + CASSANDRA_TABLE
+ + " where token(id) > ? and token(id) <= ? and scientist='Faraday1' allow filtering");
+ PCollection<KV<Long, String>> cassandraData =
+ p.apply(HadoopInputFormatIO.<Long, String>read().withConfiguration(conf)
+ .withValueTranslation(myValueTranslate));
+ // Verify the count of data retrieved from Cassandra matches expected count.
+ PAssert.thatSingleton(cassandraData.apply("Count", Count.<KV<Long, String>>globally()))
+ .isEqualTo(expectedCount);
+ PCollection<String> textValues = cassandraData.apply(Values.<String>create());
+ // Verify the output values using checksum comparison.
+ PCollection<String> consolidatedHashcode =
+ textValues.apply(Combine.globally(new HashingFn()).withoutDefaults());
+ PAssert.that(consolidatedHashcode).containsInAnyOrder(expectedChecksum);
+ p.run().waitUntilFinish();
+ }
+
+ /**
+ * Returns configuration of CqlInutFormat. Mandatory parameters required apart from inputformat
+ * class name, key class, value class are thrift port, thrift address, partitioner class, keyspace
+ * and columnfamily name
+ */
+ public Configuration getConfiguration() {
+ Configuration conf = new Configuration();
+ conf.set(CASSANDRA_THRIFT_PORT_PROPERTY, CASSANDRA_PORT);
+ conf.set(CASSANDRA_THRIFT_ADDRESS_PROPERTY, CASSANDRA_HOST);
+ conf.set(CASSANDRA_PARTITIONER_CLASS_PROPERTY, CASSANDRA_PARTITIONER_CLASS_VALUE);
+ conf.set(CASSANDRA_KEYSPACE_PROPERTY, CASSANDRA_KEYSPACE);
+ conf.set(CASSANDRA_COLUMNFAMILY_PROPERTY, CASSANDRA_TABLE);
+ conf.setClass("mapreduce.job.inputformat.class",
+ org.apache.cassandra.hadoop.cql3.CqlInputFormat.class, InputFormat.class);
+ conf.setClass("key.class", java.lang.Long.class, Object.class);
+ conf.setClass("value.class", com.datastax.driver.core.Row.class, Object.class);
+ return conf;
+ }
+
+ public static void createCassandraData() throws Exception {
+ session.execute("CREATE KEYSPACE " + CASSANDRA_KEYSPACE
+ + " WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':1};");
+ session.execute("USE " + CASSANDRA_KEYSPACE);
+ session.execute("CREATE TABLE " + CASSANDRA_TABLE
+ + "(id int, scientist text, PRIMARY KEY(id));");
+ for (int i = 0; i < TEST_DATA_ROW_COUNT; i++) {
+ session.execute("INSERT INTO " + CASSANDRA_TABLE + "(id, scientist) values(" + i
+ + ", 'Faraday" + i + "');");
+ }
+ }
+
+ @BeforeClass
+ public static void startCassandra() throws Exception {
+ //Start the Embedded Cassandra Service
+ cassandra.start();
+ final SocketOptions socketOptions = new SocketOptions();
+ // Setting this to 0 disables read timeouts.
+ socketOptions.setReadTimeoutMillis(0);
+ // This defaults to 5 s. Increase to a minute.
+ socketOptions.setConnectTimeoutMillis(60 * 1000);
+ cluster =
+ Cluster.builder().addContactPoint(CASSANDRA_HOST).withClusterName("beam")
+ .withSocketOptions(socketOptions).build();
+ session = cluster.connect();
+ createCassandraData();
+ }
+
+ @AfterClass
+ public static void stopEmbeddedCassandra() throws Exception {
+ session.close();
+ cluster.close();
+ }
+
+ /**
+ * POJO class for scientist data.
+ */
+ @Table(name = CASSANDRA_TABLE, keyspace = CASSANDRA_KEYSPACE)
+ public static class Scientist implements Serializable {
+ private static final long serialVersionUID = 1L;
+ @Column(name = "scientist")
+ private String name;
+ @Column(name = "id")
+ private int id;
+
+ public String getName() {
+ return name;
+ }
+
+ public void setName(String name) {
+ this.name = name;
+ }
+
+ public int getId() {
+ return id;
+ }
+
+ public void setId(int id) {
+ this.id = id;
+ }
+
+ public String toString() {
+ return id + ":" + name;
+ }
+ }
+}
[05/50] [abbrv] beam git commit: This closes #2491
Posted by jb...@apache.org.
This closes #2491
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/38cf23eb
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/38cf23eb
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/38cf23eb
Branch: refs/heads/DSL_SQL
Commit: 38cf23ebe34724b24db390e1bca774c7c77fdd21
Parents: d255fa2 ee7965c
Author: Thomas Groh <tg...@google.com>
Authored: Tue Apr 11 10:40:01 2017 -0700
Committer: Thomas Groh <tg...@google.com>
Committed: Tue Apr 11 10:40:01 2017 -0700
----------------------------------------------------------------------
sdks/java/io/hadoop/jdk1.8-tests/.toDelete | 0
1 file changed, 0 insertions(+), 0 deletions(-)
----------------------------------------------------------------------
[33/50] [abbrv] beam git commit: [BEAM-1722] Move PubsubIO into the
google-cloud-platform module
Posted by jb...@apache.org.
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/src/main/java/org/apache/beam/sdk/io/PubsubUnboundedSource.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/PubsubUnboundedSource.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/PubsubUnboundedSource.java
deleted file mode 100644
index 0e6bec8..0000000
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/PubsubUnboundedSource.java
+++ /dev/null
@@ -1,1463 +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 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.api.client.util.Clock;
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Charsets;
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Lists;
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
-import java.security.GeneralSecurityException;
-import java.util.ArrayDeque;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.HashSet;
-import java.util.LinkedHashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.NoSuchElementException;
-import java.util.Queue;
-import java.util.Set;
-import java.util.concurrent.ConcurrentLinkedQueue;
-import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.atomic.AtomicReference;
-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.ListCoder;
-import org.apache.beam.sdk.coders.NullableCoder;
-import org.apache.beam.sdk.coders.StringUtf8Coder;
-import org.apache.beam.sdk.io.PubsubIO.PubsubMessage;
-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.PubsubOptions;
-import org.apache.beam.sdk.options.ValueProvider;
-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.SimpleFunction;
-import org.apache.beam.sdk.transforms.Sum;
-import org.apache.beam.sdk.transforms.display.DisplayData;
-import org.apache.beam.sdk.transforms.display.DisplayData.Builder;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.util.BucketingFunction;
-import org.apache.beam.sdk.util.CoderUtils;
-import org.apache.beam.sdk.util.MovingFunction;
-import org.apache.beam.sdk.util.PubsubClient;
-import org.apache.beam.sdk.util.PubsubClient.ProjectPath;
-import org.apache.beam.sdk.util.PubsubClient.PubsubClientFactory;
-import org.apache.beam.sdk.util.PubsubClient.SubscriptionPath;
-import org.apache.beam.sdk.util.PubsubClient.TopicPath;
-import org.apache.beam.sdk.values.PBegin;
-import org.apache.beam.sdk.values.PCollection;
-import org.joda.time.Duration;
-import org.joda.time.Instant;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/**
- * A PTransform which streams messages from Pubsub.
- * <ul>
- * <li>The underlying implementation in an {@link UnboundedSource} which receives messages
- * in batches and hands them out one at a time.
- * <li>The watermark (either in Pubsub processing time or custom timestamp time) is estimated
- * by keeping track of the minimum of the last minutes worth of messages. This assumes Pubsub
- * delivers the oldest (in Pubsub processing time) available message at least once a minute,
- * and that custom timestamps are 'mostly' monotonic with Pubsub processing time. Unfortunately
- * both of those assumptions are fragile. Thus the estimated watermark may get ahead of
- * the 'true' watermark and cause some messages to be late.
- * <li>Checkpoints are used both to ACK received messages back to Pubsub (so that they may
- * be retired on the Pubsub end), and to NACK already consumed messages should a checkpoint
- * need to be restored (so that Pubsub will resend those messages promptly).
- * <li>The backlog is determined by each reader using the messages which have been pulled from
- * Pubsub but not yet consumed downstream. The backlog does not take account of any messages queued
- * by Pubsub for the subscription. Unfortunately there is currently no API to determine
- * the size of the Pubsub queue's backlog.
- * <li>The subscription must already exist.
- * <li>The subscription timeout is read whenever a reader is started. However it is not
- * checked thereafter despite the timeout being user-changeable on-the-fly.
- * <li>We log vital stats every 30 seconds.
- * <li>Though some background threads may be used by the underlying transport all Pubsub calls
- * are blocking. We rely on the underlying runner to allow multiple
- * {@link UnboundedSource.UnboundedReader} instances to execute concurrently and thus hide latency.
- * </ul>
- */
-public class PubsubUnboundedSource<T> extends PTransform<PBegin, PCollection<T>> {
- private static final Logger LOG = LoggerFactory.getLogger(PubsubUnboundedSource.class);
-
- /**
- * Default ACK timeout for created subscriptions.
- */
- private static final int DEAULT_ACK_TIMEOUT_SEC = 60;
-
- /**
- * Coder for checkpoints.
- */
- private static final PubsubCheckpointCoder<?> CHECKPOINT_CODER = new PubsubCheckpointCoder<>();
-
- /**
- * Maximum number of messages per pull.
- */
- private static final int PULL_BATCH_SIZE = 1000;
-
- /**
- * Maximum number of ACK ids per ACK or ACK extension call.
- */
- private static final int ACK_BATCH_SIZE = 2000;
-
- /**
- * Maximum number of messages in flight.
- */
- private static final int MAX_IN_FLIGHT = 20000;
-
- /**
- * Timeout for round trip from receiving a message to finally ACKing it back to Pubsub.
- */
- private static final Duration PROCESSING_TIMEOUT = Duration.standardSeconds(120);
-
- /**
- * Percentage of ack timeout by which to extend acks when they are near timeout.
- */
- private static final int ACK_EXTENSION_PCT = 50;
-
- /**
- * Percentage of ack timeout we should use as a safety margin. We'll try to extend acks
- * by this margin before the ack actually expires.
- */
- private static final int ACK_SAFETY_PCT = 20;
-
- /**
- * For stats only: How close we can get to an ack deadline before we risk it being already
- * considered passed by Pubsub.
- */
- private static final Duration ACK_TOO_LATE = Duration.standardSeconds(2);
-
- /**
- * Period of samples to determine watermark and other stats.
- */
- private static final Duration SAMPLE_PERIOD = Duration.standardMinutes(1);
-
- /**
- * Period of updates to determine watermark and other stats.
- */
- private static final Duration SAMPLE_UPDATE = Duration.standardSeconds(5);
-
- /**
- * Period for logging stats.
- */
- private static final Duration LOG_PERIOD = Duration.standardSeconds(30);
-
- /**
- * Minimum number of unread messages required before considering updating watermark.
- */
- private static final int MIN_WATERMARK_MESSAGES = 10;
-
- /**
- * Minimum number of SAMPLE_UPDATE periods over which unread messages shoud be spread
- * before considering updating watermark.
- */
- private static final int MIN_WATERMARK_SPREAD = 2;
-
- /**
- * Additional sharding so that we can hide read message latency.
- */
- private static final int SCALE_OUT = 4;
-
- // TODO: Would prefer to use MinLongFn but it is a BinaryCombineFn<Long> rather
- // than a BinaryCombineLongFn. [BEAM-285]
- private static final Combine.BinaryCombineLongFn MIN =
- new Combine.BinaryCombineLongFn() {
- @Override
- public long apply(long left, long right) {
- return Math.min(left, right);
- }
-
- @Override
- public long identity() {
- return Long.MAX_VALUE;
- }
- };
-
- private static final Combine.BinaryCombineLongFn MAX =
- new Combine.BinaryCombineLongFn() {
- @Override
- public long apply(long left, long right) {
- return Math.max(left, right);
- }
-
- @Override
- public long identity() {
- return Long.MIN_VALUE;
- }
- };
-
- private static final Combine.BinaryCombineLongFn SUM = Sum.ofLongs();
-
- // ================================================================================
- // Checkpoint
- // ================================================================================
-
- /**
- * Which messages have been durably committed and thus can now be ACKed.
- * Which messages have been read but not yet committed, in which case they should be NACKed if
- * we need to restore.
- */
- @VisibleForTesting
- static class PubsubCheckpoint<T> implements UnboundedSource.CheckpointMark {
- /**
- * The {@link SubscriptionPath} to the subscription the reader is reading from. May be
- * {@code null} if the {@link PubsubUnboundedSource} contains the subscription.
- */
- @VisibleForTesting
- @Nullable String subscriptionPath;
-
- /**
- * If the checkpoint is for persisting: the reader who's snapshotted state we are persisting.
- * If the checkpoint is for restoring: {@literal null}.
- * Not persisted in durable checkpoint.
- * CAUTION: Between a checkpoint being taken and {@link #finalizeCheckpoint()} being called
- * the 'true' active reader may have changed.
- */
- @Nullable
- private PubsubReader<T> reader;
-
- /**
- * If the checkpoint is for persisting: The ACK ids of messages which have been passed
- * downstream since the last checkpoint.
- * If the checkpoint is for restoring: {@literal null}.
- * Not persisted in durable checkpoint.
- */
- @Nullable
- private List<String> safeToAckIds;
-
- /**
- * If the checkpoint is for persisting: The ACK ids of messages which have been received
- * from Pubsub but not yet passed downstream at the time of the snapshot.
- * If the checkpoint is for restoring: Same, but recovered from durable storage.
- */
- @VisibleForTesting
- final List<String> notYetReadIds;
-
- public PubsubCheckpoint(
- @Nullable String subscriptionPath,
- @Nullable PubsubReader<T> reader,
- @Nullable List<String> safeToAckIds,
- List<String> notYetReadIds) {
- this.subscriptionPath = subscriptionPath;
- this.reader = reader;
- this.safeToAckIds = safeToAckIds;
- this.notYetReadIds = notYetReadIds;
- }
-
- @Nullable
- private SubscriptionPath getSubscription() {
- return subscriptionPath == null
- ? null
- : PubsubClient.subscriptionPathFromPath(subscriptionPath);
- }
-
- /**
- * BLOCKING
- * All messages which have been passed downstream have now been durably committed.
- * We can ACK them upstream.
- * CAUTION: This may never be called.
- */
- @Override
- public void finalizeCheckpoint() throws IOException {
- checkState(reader != null && safeToAckIds != null, "Cannot finalize a restored checkpoint");
- // Even if the 'true' active reader has changed since the checkpoint was taken we are
- // fine:
- // - The underlying Pubsub topic will not have changed, so the following ACKs will still
- // go to the right place.
- // - We'll delete the ACK ids from the readers in-flight state, but that only effects
- // flow control and stats, neither of which are relevant anymore.
- try {
- int n = safeToAckIds.size();
- List<String> batchSafeToAckIds = new ArrayList<>(Math.min(n, ACK_BATCH_SIZE));
- for (String ackId : safeToAckIds) {
- batchSafeToAckIds.add(ackId);
- if (batchSafeToAckIds.size() >= ACK_BATCH_SIZE) {
- reader.ackBatch(batchSafeToAckIds);
- n -= batchSafeToAckIds.size();
- // CAUTION: Don't reuse the same list since ackBatch holds on to it.
- batchSafeToAckIds = new ArrayList<>(Math.min(n, ACK_BATCH_SIZE));
- }
- }
- if (!batchSafeToAckIds.isEmpty()) {
- reader.ackBatch(batchSafeToAckIds);
- }
- } finally {
- int remainingInFlight = reader.numInFlightCheckpoints.decrementAndGet();
- checkState(remainingInFlight >= 0,
- "Miscounted in-flight checkpoints");
- reader.maybeCloseClient();
- reader = null;
- safeToAckIds = null;
- }
- }
-
- /**
- * Return current time according to {@code reader}.
- */
- private static long now(PubsubReader<?> reader) {
- if (reader.outer.outer.clock == null) {
- return System.currentTimeMillis();
- } else {
- return reader.outer.outer.clock.currentTimeMillis();
- }
- }
-
- /**
- * BLOCKING
- * NACK all messages which have been read from Pubsub but not passed downstream.
- * This way Pubsub will send them again promptly.
- */
- public void nackAll(PubsubReader<T> reader) throws IOException {
- checkState(this.reader == null, "Cannot nackAll on persisting checkpoint");
- List<String> batchYetToAckIds =
- new ArrayList<>(Math.min(notYetReadIds.size(), ACK_BATCH_SIZE));
- for (String ackId : notYetReadIds) {
- batchYetToAckIds.add(ackId);
- if (batchYetToAckIds.size() >= ACK_BATCH_SIZE) {
- long nowMsSinceEpoch = now(reader);
- reader.nackBatch(nowMsSinceEpoch, batchYetToAckIds);
- batchYetToAckIds.clear();
- }
- }
- if (!batchYetToAckIds.isEmpty()) {
- long nowMsSinceEpoch = now(reader);
- reader.nackBatch(nowMsSinceEpoch, batchYetToAckIds);
- }
- }
- }
-
- /** The coder for our checkpoints. */
- private static class PubsubCheckpointCoder<T> extends AtomicCoder<PubsubCheckpoint<T>> {
- private static final Coder<String> SUBSCRIPTION_PATH_CODER =
- NullableCoder.of(StringUtf8Coder.of());
- private static final Coder<List<String>> LIST_CODER = ListCoder.of(StringUtf8Coder.of());
-
- @Override
- public void encode(PubsubCheckpoint<T> value, OutputStream outStream, Context context)
- throws IOException {
- SUBSCRIPTION_PATH_CODER.encode(
- value.subscriptionPath,
- outStream,
- context.nested());
- LIST_CODER.encode(value.notYetReadIds, outStream, context);
- }
-
- @Override
- public PubsubCheckpoint<T> decode(InputStream inStream, Context context) throws IOException {
- String path = SUBSCRIPTION_PATH_CODER.decode(inStream, context.nested());
- List<String> notYetReadIds = LIST_CODER.decode(inStream, context);
- return new PubsubCheckpoint<>(path, null, null, notYetReadIds);
- }
- }
-
- // ================================================================================
- // Reader
- // ================================================================================
-
- /**
- * A reader which keeps track of which messages have been received from Pubsub
- * but not yet consumed downstream and/or ACKed back to Pubsub.
- */
- @VisibleForTesting
- static class PubsubReader<T> extends UnboundedSource.UnboundedReader<T> {
- /**
- * For access to topic and checkpointCoder.
- */
- private final PubsubSource<T> outer;
- @VisibleForTesting
- final SubscriptionPath subscription;
-
- private final SimpleFunction<PubsubIO.PubsubMessage, T> parseFn;
-
- /**
- * Client on which to talk to Pubsub. Contains a null value if the client has been closed.
- */
- private AtomicReference<PubsubClient> pubsubClient;
-
- /**
- * The closed state of this {@link PubsubReader}. If true, the reader has not yet been closed,
- * and it will have a non-null value within {@link #pubsubClient}.
- */
- private AtomicBoolean active = new AtomicBoolean(true);
-
- /**
- * Ack timeout, in ms, as set on subscription when we first start reading. Not
- * updated thereafter. -1 if not yet determined.
- */
- private int ackTimeoutMs;
-
- /**
- * ACK ids of messages we have delivered downstream but not yet ACKed.
- */
- private Set<String> safeToAckIds;
-
- /**
- * Messages we have received from Pubsub and not yet delivered downstream.
- * We preserve their order.
- */
- private final Queue<PubsubClient.IncomingMessage> notYetRead;
-
- private static class InFlightState {
- /**
- * When request which yielded message was issues.
- */
- long requestTimeMsSinceEpoch;
-
- /**
- * When Pubsub will consider this message's ACK to timeout and thus it needs to be
- * extended.
- */
- long ackDeadlineMsSinceEpoch;
-
- public InFlightState(long requestTimeMsSinceEpoch, long ackDeadlineMsSinceEpoch) {
- this.requestTimeMsSinceEpoch = requestTimeMsSinceEpoch;
- this.ackDeadlineMsSinceEpoch = ackDeadlineMsSinceEpoch;
- }
- }
-
- /**
- * Map from ACK ids of messages we have received from Pubsub but not yet ACKed to their
- * in flight state. Ordered from earliest to latest ACK deadline.
- */
- private final LinkedHashMap<String, InFlightState> inFlight;
-
- /**
- * Batches of successfully ACKed ids which need to be pruned from the above.
- * CAUTION: Accessed by both reader and checkpointing threads.
- */
- private final Queue<List<String>> ackedIds;
-
- /**
- * Byte size of undecoded elements in {@link #notYetRead}.
- */
- private long notYetReadBytes;
-
- /**
- * Bucketed map from received time (as system time, ms since epoch) to message
- * timestamps (mssince epoch) of all received but not-yet read messages.
- * Used to estimate watermark.
- */
- private BucketingFunction minUnreadTimestampMsSinceEpoch;
-
- /**
- * Minimum of timestamps (ms since epoch) of all recently read messages.
- * Used to estimate watermark.
- */
- private MovingFunction minReadTimestampMsSinceEpoch;
-
- /**
- * System time (ms since epoch) we last received a message from Pubsub, or -1 if
- * not yet received any messages.
- */
- private long lastReceivedMsSinceEpoch;
-
- /**
- * The last reported watermark (ms since epoch), or beginning of time if none yet reported.
- */
- private long lastWatermarkMsSinceEpoch;
-
- /**
- * The current message, or {@literal null} if none.
- */
- @Nullable
- private PubsubClient.IncomingMessage current;
-
- /**
- * Stats only: System time (ms since epoch) we last logs stats, or -1 if never.
- */
- private long lastLogTimestampMsSinceEpoch;
-
- /**
- * Stats only: Total number of messages received.
- */
- private long numReceived;
-
- /**
- * Stats only: Number of messages which have recently been received.
- */
- private MovingFunction numReceivedRecently;
-
- /**
- * Stats only: Number of messages which have recently had their deadline extended.
- */
- private MovingFunction numExtendedDeadlines;
-
- /**
- * Stats only: Number of messages which have recenttly had their deadline extended even
- * though it may be too late to do so.
- */
- private MovingFunction numLateDeadlines;
-
-
- /**
- * Stats only: Number of messages which have recently been ACKed.
- */
- private MovingFunction numAcked;
-
- /**
- * Stats only: Number of messages which have recently expired (ACKs were extended for too
- * long).
- */
- private MovingFunction numExpired;
-
- /**
- * Stats only: Number of messages which have recently been NACKed.
- */
- private MovingFunction numNacked;
-
- /**
- * Stats only: Number of message bytes which have recently been read by downstream consumer.
- */
- private MovingFunction numReadBytes;
-
- /**
- * Stats only: Minimum of timestamp (ms since epoch) of all recently received messages.
- * Used to estimate timestamp skew. Does not contribute to watermark estimator.
- */
- private MovingFunction minReceivedTimestampMsSinceEpoch;
-
- /**
- * Stats only: Maximum of timestamp (ms since epoch) of all recently received messages.
- * Used to estimate timestamp skew.
- */
- private MovingFunction maxReceivedTimestampMsSinceEpoch;
-
- /**
- * Stats only: Minimum of recent estimated watermarks (ms since epoch).
- */
- private MovingFunction minWatermarkMsSinceEpoch;
-
- /**
- * Stats ony: Maximum of recent estimated watermarks (ms since epoch).
- */
- private MovingFunction maxWatermarkMsSinceEpoch;
-
- /**
- * Stats only: Number of messages with timestamps strictly behind the estimated watermark
- * at the time they are received. These may be considered 'late' by downstream computations.
- */
- private MovingFunction numLateMessages;
-
- /**
- * Stats only: Current number of checkpoints in flight.
- * CAUTION: Accessed by both checkpointing and reader threads.
- */
- private AtomicInteger numInFlightCheckpoints;
-
- /**
- * Stats only: Maximum number of checkpoints in flight at any time.
- */
- private int maxInFlightCheckpoints;
-
- private static MovingFunction newFun(Combine.BinaryCombineLongFn function) {
- return new MovingFunction(SAMPLE_PERIOD.getMillis(),
- SAMPLE_UPDATE.getMillis(),
- MIN_WATERMARK_SPREAD,
- MIN_WATERMARK_MESSAGES,
- function);
- }
-
- /**
- * Construct a reader.
- */
- public PubsubReader(PubsubOptions options, PubsubSource<T> outer, SubscriptionPath subscription,
- SimpleFunction<PubsubIO.PubsubMessage, T> parseFn)
- throws IOException, GeneralSecurityException {
- this.outer = outer;
- this.subscription = subscription;
- this.parseFn = parseFn;
- pubsubClient =
- new AtomicReference<>(
- outer.outer.pubsubFactory.newClient(
- outer.outer.timestampLabel, outer.outer.idLabel, options));
- ackTimeoutMs = -1;
- safeToAckIds = new HashSet<>();
- notYetRead = new ArrayDeque<>();
- inFlight = new LinkedHashMap<>();
- ackedIds = new ConcurrentLinkedQueue<>();
- notYetReadBytes = 0;
- minUnreadTimestampMsSinceEpoch = new BucketingFunction(SAMPLE_UPDATE.getMillis(),
- MIN_WATERMARK_SPREAD,
- MIN_WATERMARK_MESSAGES,
- MIN);
- minReadTimestampMsSinceEpoch = newFun(MIN);
- lastReceivedMsSinceEpoch = -1;
- lastWatermarkMsSinceEpoch = BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis();
- current = null;
- lastLogTimestampMsSinceEpoch = -1;
- numReceived = 0L;
- numReceivedRecently = newFun(SUM);
- numExtendedDeadlines = newFun(SUM);
- numLateDeadlines = newFun(SUM);
- numAcked = newFun(SUM);
- numExpired = newFun(SUM);
- numNacked = newFun(SUM);
- numReadBytes = newFun(SUM);
- minReceivedTimestampMsSinceEpoch = newFun(MIN);
- maxReceivedTimestampMsSinceEpoch = newFun(MAX);
- minWatermarkMsSinceEpoch = newFun(MIN);
- maxWatermarkMsSinceEpoch = newFun(MAX);
- numLateMessages = newFun(SUM);
- numInFlightCheckpoints = new AtomicInteger();
- maxInFlightCheckpoints = 0;
- }
-
- @VisibleForTesting
- PubsubClient getPubsubClient() {
- return pubsubClient.get();
- }
-
- /**
- * Acks the provided {@code ackIds} back to Pubsub, blocking until all of the messages are
- * ACKed.
- *
- * <p>CAUTION: May be invoked from a separate thread.
- *
- * <p>CAUTION: Retains {@code ackIds}.
- */
- void ackBatch(List<String> ackIds) throws IOException {
- pubsubClient.get().acknowledge(subscription, ackIds);
- ackedIds.add(ackIds);
- }
-
- /**
- * BLOCKING
- * NACK (ie request deadline extension of 0) receipt of messages from Pubsub
- * with the given {@code ockIds}. Does not retain {@code ackIds}.
- */
- public void nackBatch(long nowMsSinceEpoch, List<String> ackIds) throws IOException {
- pubsubClient.get().modifyAckDeadline(subscription, ackIds, 0);
- numNacked.add(nowMsSinceEpoch, ackIds.size());
- }
-
- /**
- * BLOCKING
- * Extend the processing deadline for messages from Pubsub with the given {@code ackIds}.
- * Does not retain {@code ackIds}.
- */
- private void extendBatch(long nowMsSinceEpoch, List<String> ackIds) throws IOException {
- int extensionSec = (ackTimeoutMs * ACK_EXTENSION_PCT) / (100 * 1000);
- pubsubClient.get().modifyAckDeadline(subscription, ackIds, extensionSec);
- numExtendedDeadlines.add(nowMsSinceEpoch, ackIds.size());
- }
-
- /**
- * Return the current time, in ms since epoch.
- */
- private long now() {
- if (outer.outer.clock == null) {
- return System.currentTimeMillis();
- } else {
- return outer.outer.clock.currentTimeMillis();
- }
- }
-
- /**
- * Messages which have been ACKed (via the checkpoint finalize) are no longer in flight.
- * This is only used for flow control and stats.
- */
- private void retire() throws IOException {
- long nowMsSinceEpoch = now();
- while (true) {
- List<String> ackIds = ackedIds.poll();
- if (ackIds == null) {
- return;
- }
- numAcked.add(nowMsSinceEpoch, ackIds.size());
- for (String ackId : ackIds) {
- inFlight.remove(ackId);
- safeToAckIds.remove(ackId);
- }
- }
- }
-
- /**
- * BLOCKING
- * Extend deadline for all messages which need it.
- * CAUTION: If extensions can't keep up with wallclock then we'll never return.
- */
- private void extend() throws IOException {
- while (true) {
- long nowMsSinceEpoch = now();
- List<String> assumeExpired = new ArrayList<>();
- List<String> toBeExtended = new ArrayList<>();
- List<String> toBeExpired = new ArrayList<>();
- // Messages will be in increasing deadline order.
- for (Map.Entry<String, InFlightState> entry : inFlight.entrySet()) {
- if (entry.getValue().ackDeadlineMsSinceEpoch - (ackTimeoutMs * ACK_SAFETY_PCT) / 100
- > nowMsSinceEpoch) {
- // All remaining messages don't need their ACKs to be extended.
- break;
- }
-
- if (entry.getValue().ackDeadlineMsSinceEpoch - ACK_TOO_LATE.getMillis()
- < nowMsSinceEpoch) {
- // Pubsub may have already considered this message to have expired.
- // If so it will (eventually) be made available on a future pull request.
- // If this message ends up being committed then it will be considered a duplicate
- // when re-pulled.
- assumeExpired.add(entry.getKey());
- continue;
- }
-
- if (entry.getValue().requestTimeMsSinceEpoch + PROCESSING_TIMEOUT.getMillis()
- < nowMsSinceEpoch) {
- // This message has been in-flight for too long.
- // Give up on it, otherwise we risk extending its ACK indefinitely.
- toBeExpired.add(entry.getKey());
- continue;
- }
-
- // Extend the ACK for this message.
- toBeExtended.add(entry.getKey());
- if (toBeExtended.size() >= ACK_BATCH_SIZE) {
- // Enough for one batch.
- break;
- }
- }
-
- if (assumeExpired.isEmpty() && toBeExtended.isEmpty() && toBeExpired.isEmpty()) {
- // Nothing to be done.
- return;
- }
-
- if (!assumeExpired.isEmpty()) {
- // If we didn't make the ACK deadline assume expired and no longer in flight.
- numLateDeadlines.add(nowMsSinceEpoch, assumeExpired.size());
- for (String ackId : assumeExpired) {
- inFlight.remove(ackId);
- }
- }
-
- if (!toBeExpired.isEmpty()) {
- // Expired messages are no longer considered in flight.
- numExpired.add(nowMsSinceEpoch, toBeExpired.size());
- for (String ackId : toBeExpired) {
- inFlight.remove(ackId);
- }
- }
-
- if (!toBeExtended.isEmpty()) {
- // Pubsub extends acks from it's notion of current time.
- // We'll try to track that on our side, but note the deadlines won't necessarily agree.
- long newDeadlineMsSinceEpoch = nowMsSinceEpoch + (ackTimeoutMs * ACK_EXTENSION_PCT) / 100;
- for (String ackId : toBeExtended) {
- // Maintain increasing ack deadline order.
- InFlightState state = inFlight.remove(ackId);
- inFlight.put(ackId,
- new InFlightState(state.requestTimeMsSinceEpoch, newDeadlineMsSinceEpoch));
- }
- // BLOCKs until extended.
- extendBatch(nowMsSinceEpoch, toBeExtended);
- }
- }
- }
-
- /**
- * BLOCKING
- * Fetch another batch of messages from Pubsub.
- */
- private void pull() throws IOException {
- if (inFlight.size() >= MAX_IN_FLIGHT) {
- // Wait for checkpoint to be finalized before pulling anymore.
- // There may be lag while checkpoints are persisted and the finalizeCheckpoint method
- // is invoked. By limiting the in-flight messages we can ensure we don't end up consuming
- // messages faster than we can checkpoint them.
- return;
- }
-
- long requestTimeMsSinceEpoch = now();
- long deadlineMsSinceEpoch = requestTimeMsSinceEpoch + ackTimeoutMs;
-
- // Pull the next batch.
- // BLOCKs until received.
- Collection<PubsubClient.IncomingMessage> receivedMessages =
- pubsubClient.get().pull(requestTimeMsSinceEpoch, subscription, PULL_BATCH_SIZE, true);
- if (receivedMessages.isEmpty()) {
- // Nothing available yet. Try again later.
- return;
- }
-
- lastReceivedMsSinceEpoch = requestTimeMsSinceEpoch;
-
- // Capture the received messages.
- for (PubsubClient.IncomingMessage incomingMessage : receivedMessages) {
- notYetRead.add(incomingMessage);
- notYetReadBytes += incomingMessage.elementBytes.length;
- inFlight.put(incomingMessage.ackId,
- new InFlightState(requestTimeMsSinceEpoch, deadlineMsSinceEpoch));
- numReceived++;
- numReceivedRecently.add(requestTimeMsSinceEpoch, 1L);
- minReceivedTimestampMsSinceEpoch.add(requestTimeMsSinceEpoch,
- incomingMessage.timestampMsSinceEpoch);
- maxReceivedTimestampMsSinceEpoch.add(requestTimeMsSinceEpoch,
- incomingMessage.timestampMsSinceEpoch);
- minUnreadTimestampMsSinceEpoch.add(requestTimeMsSinceEpoch,
- incomingMessage.timestampMsSinceEpoch);
- }
- }
-
- /**
- * Log stats if time to do so.
- */
- private void stats() {
- long nowMsSinceEpoch = now();
- if (lastLogTimestampMsSinceEpoch < 0) {
- lastLogTimestampMsSinceEpoch = nowMsSinceEpoch;
- return;
- }
- long deltaMs = nowMsSinceEpoch - lastLogTimestampMsSinceEpoch;
- if (deltaMs < LOG_PERIOD.getMillis()) {
- return;
- }
-
- String messageSkew = "unknown";
- long minTimestamp = minReceivedTimestampMsSinceEpoch.get(nowMsSinceEpoch);
- long maxTimestamp = maxReceivedTimestampMsSinceEpoch.get(nowMsSinceEpoch);
- if (minTimestamp < Long.MAX_VALUE && maxTimestamp > Long.MIN_VALUE) {
- messageSkew = (maxTimestamp - minTimestamp) + "ms";
- }
-
- String watermarkSkew = "unknown";
- long minWatermark = minWatermarkMsSinceEpoch.get(nowMsSinceEpoch);
- long maxWatermark = maxWatermarkMsSinceEpoch.get(nowMsSinceEpoch);
- if (minWatermark < Long.MAX_VALUE && maxWatermark > Long.MIN_VALUE) {
- watermarkSkew = (maxWatermark - minWatermark) + "ms";
- }
-
- String oldestInFlight = "no";
- String oldestAckId = Iterables.getFirst(inFlight.keySet(), null);
- if (oldestAckId != null) {
- oldestInFlight =
- (nowMsSinceEpoch - inFlight.get(oldestAckId).requestTimeMsSinceEpoch) + "ms";
- }
-
- LOG.info("Pubsub {} has "
- + "{} received messages, "
- + "{} current unread messages, "
- + "{} current unread bytes, "
- + "{} current in-flight msgs, "
- + "{} oldest in-flight, "
- + "{} current in-flight checkpoints, "
- + "{} max in-flight checkpoints, "
- + "{}B/s recent read, "
- + "{} recent received, "
- + "{} recent extended, "
- + "{} recent late extended, "
- + "{} recent ACKed, "
- + "{} recent NACKed, "
- + "{} recent expired, "
- + "{} recent message timestamp skew, "
- + "{} recent watermark skew, "
- + "{} recent late messages, "
- + "{} last reported watermark",
- subscription,
- numReceived,
- notYetRead.size(),
- notYetReadBytes,
- inFlight.size(),
- oldestInFlight,
- numInFlightCheckpoints.get(),
- maxInFlightCheckpoints,
- numReadBytes.get(nowMsSinceEpoch) / (SAMPLE_PERIOD.getMillis() / 1000L),
- numReceivedRecently.get(nowMsSinceEpoch),
- numExtendedDeadlines.get(nowMsSinceEpoch),
- numLateDeadlines.get(nowMsSinceEpoch),
- numAcked.get(nowMsSinceEpoch),
- numNacked.get(nowMsSinceEpoch),
- numExpired.get(nowMsSinceEpoch),
- messageSkew,
- watermarkSkew,
- numLateMessages.get(nowMsSinceEpoch),
- new Instant(lastWatermarkMsSinceEpoch));
-
- lastLogTimestampMsSinceEpoch = nowMsSinceEpoch;
- }
-
- @Override
- public boolean start() throws IOException {
- // Determine the ack timeout.
- ackTimeoutMs = pubsubClient.get().ackDeadlineSeconds(subscription) * 1000;
- return advance();
- }
-
- /**
- * BLOCKING
- * Return {@literal true} if a Pubsub messaage is available, {@literal false} if
- * none is available at this time or we are over-subscribed. May BLOCK while extending
- * ACKs or fetching available messages. Will not block waiting for messages.
- */
- @Override
- public boolean advance() throws IOException {
- // Emit stats.
- stats();
-
- if (current != null) {
- // Current is consumed. It can no longer contribute to holding back the watermark.
- minUnreadTimestampMsSinceEpoch.remove(current.requestTimeMsSinceEpoch);
- current = null;
- }
-
- // Retire state associated with ACKed messages.
- retire();
-
- // Extend all pressing deadlines.
- // Will BLOCK until done.
- // If the system is pulling messages only to let them sit in a downsteam queue then
- // this will have the effect of slowing down the pull rate.
- // However, if the system is genuinely taking longer to process each message then
- // the work to extend ACKs would be better done in the background.
- extend();
-
- if (notYetRead.isEmpty()) {
- // Pull another batch.
- // Will BLOCK until fetch returns, but will not block until a message is available.
- pull();
- }
-
- // Take one message from queue.
- current = notYetRead.poll();
- if (current == null) {
- // Try again later.
- return false;
- }
- notYetReadBytes -= current.elementBytes.length;
- checkState(notYetReadBytes >= 0);
- long nowMsSinceEpoch = now();
- numReadBytes.add(nowMsSinceEpoch, current.elementBytes.length);
- minReadTimestampMsSinceEpoch.add(nowMsSinceEpoch, current.timestampMsSinceEpoch);
- if (current.timestampMsSinceEpoch < lastWatermarkMsSinceEpoch) {
- numLateMessages.add(nowMsSinceEpoch, 1L);
- }
-
- // Current message can be considered 'read' and will be persisted by the next
- // checkpoint. So it is now safe to ACK back to Pubsub.
- safeToAckIds.add(current.ackId);
- return true;
- }
-
- @Override
- public T getCurrent() throws NoSuchElementException {
- if (current == null) {
- throw new NoSuchElementException();
- }
- try {
- if (parseFn != null) {
- return parseFn.apply(new PubsubIO.PubsubMessage(
- current.elementBytes, current.attributes));
- } else {
- return CoderUtils.decodeFromByteArray(outer.outer.elementCoder, current.elementBytes);
- }
- } catch (CoderException e) {
- throw new RuntimeException("Unable to decode element from Pubsub message: ", e);
- }
- }
-
- @Override
- public Instant getCurrentTimestamp() throws NoSuchElementException {
- if (current == null) {
- throw new NoSuchElementException();
- }
- return new Instant(current.timestampMsSinceEpoch);
- }
-
- @Override
- public byte[] getCurrentRecordId() throws NoSuchElementException {
- if (current == null) {
- throw new NoSuchElementException();
- }
- return current.recordId.getBytes(Charsets.UTF_8);
- }
-
- /**
- * {@inheritDoc}.
- *
- * <p>Marks this {@link PubsubReader} as no longer active. The {@link PubsubClient}
- * continue to exist and be active beyond the life of this call if there are any in-flight
- * checkpoints. When no in-flight checkpoints remain, the reader will be closed.
- */
- @Override
- public void close() throws IOException {
- active.set(false);
- maybeCloseClient();
- }
-
- /**
- * Close this reader's underlying {@link PubsubClient} if the reader has been closed and there
- * are no outstanding checkpoints.
- */
- private void maybeCloseClient() throws IOException {
- if (!active.get() && numInFlightCheckpoints.get() == 0) {
- // The reader has been closed and it has no more outstanding checkpoints. The client
- // must be closed so it doesn't leak
- PubsubClient client = pubsubClient.getAndSet(null);
- if (client != null) {
- client.close();
- }
- }
- }
-
- @Override
- public PubsubSource<T> getCurrentSource() {
- return outer;
- }
-
- @Override
- public Instant getWatermark() {
- if (pubsubClient.get().isEOF() && notYetRead.isEmpty()) {
- // For testing only: Advance the watermark to the end of time to signal
- // the test is complete.
- return BoundedWindow.TIMESTAMP_MAX_VALUE;
- }
-
- // NOTE: We'll allow the watermark to go backwards. The underlying runner is responsible
- // for aggregating all reported watermarks and ensuring the aggregate is latched.
- // If we attempt to latch locally then it is possible a temporary starvation of one reader
- // could cause its estimated watermark to fast forward to current system time. Then when
- // the reader resumes its watermark would be unable to resume tracking.
- // By letting the underlying runner latch we avoid any problems due to localized starvation.
- long nowMsSinceEpoch = now();
- long readMin = minReadTimestampMsSinceEpoch.get(nowMsSinceEpoch);
- long unreadMin = minUnreadTimestampMsSinceEpoch.get();
- if (readMin == Long.MAX_VALUE
- && unreadMin == Long.MAX_VALUE
- && lastReceivedMsSinceEpoch >= 0
- && nowMsSinceEpoch > lastReceivedMsSinceEpoch + SAMPLE_PERIOD.getMillis()) {
- // We don't currently have any unread messages pending, we have not had any messages
- // read for a while, and we have not received any new messages from Pubsub for a while.
- // Advance watermark to current time.
- // TODO: Estimate a timestamp lag.
- lastWatermarkMsSinceEpoch = nowMsSinceEpoch;
- } else if (minReadTimestampMsSinceEpoch.isSignificant()
- || minUnreadTimestampMsSinceEpoch.isSignificant()) {
- // Take minimum of the timestamps in all unread messages and recently read messages.
- lastWatermarkMsSinceEpoch = Math.min(readMin, unreadMin);
- }
- // else: We're not confident enough to estimate a new watermark. Stick with the old one.
- minWatermarkMsSinceEpoch.add(nowMsSinceEpoch, lastWatermarkMsSinceEpoch);
- maxWatermarkMsSinceEpoch.add(nowMsSinceEpoch, lastWatermarkMsSinceEpoch);
- return new Instant(lastWatermarkMsSinceEpoch);
- }
-
- @Override
- public PubsubCheckpoint<T> getCheckpointMark() {
- int cur = numInFlightCheckpoints.incrementAndGet();
- maxInFlightCheckpoints = Math.max(maxInFlightCheckpoints, cur);
- // It's possible for a checkpoint to be taken but never finalized.
- // So we simply copy whatever safeToAckIds we currently have.
- // It's possible a later checkpoint will be taken before an earlier one is finalized,
- // in which case we'll double ACK messages to Pubsub. However Pubsub is fine with that.
- List<String> snapshotSafeToAckIds = Lists.newArrayList(safeToAckIds);
- List<String> snapshotNotYetReadIds = new ArrayList<>(notYetRead.size());
- for (PubsubClient.IncomingMessage incomingMessage : notYetRead) {
- snapshotNotYetReadIds.add(incomingMessage.ackId);
- }
- if (outer.subscriptionPath == null) {
- // need to include the subscription in case we resume, as it's not stored in the source.
- return new PubsubCheckpoint<>(
- subscription.getPath(), this, snapshotSafeToAckIds, snapshotNotYetReadIds);
- }
- return new PubsubCheckpoint<>(null, this, snapshotSafeToAckIds, snapshotNotYetReadIds);
- }
-
- @Override
- public long getSplitBacklogBytes() {
- return notYetReadBytes;
- }
- }
-
- // ================================================================================
- // Source
- // ================================================================================
-
- @VisibleForTesting
- static class PubsubSource<T> extends UnboundedSource<T, PubsubCheckpoint<T>> {
- public final PubsubUnboundedSource<T> outer;
- // The subscription to read from.
- @VisibleForTesting
- final SubscriptionPath subscriptionPath;
-
- public PubsubSource(PubsubUnboundedSource<T> outer) {
- this(outer, outer.getSubscription());
- }
-
- private PubsubSource(PubsubUnboundedSource<T> outer, SubscriptionPath subscriptionPath) {
- this.outer = outer;
- this.subscriptionPath = subscriptionPath;
- }
-
- @Override
- public List<PubsubSource<T>> generateInitialSplits(
- int desiredNumSplits, PipelineOptions options) throws Exception {
- List<PubsubSource<T>> result = new ArrayList<>(desiredNumSplits);
- PubsubSource<T> splitSource = this;
- if (subscriptionPath == null) {
- splitSource = new PubsubSource<>(outer, outer.createRandomSubscription(options));
- }
- for (int i = 0; i < desiredNumSplits * SCALE_OUT; i++) {
- // Since the source is immutable and Pubsub automatically shards we simply
- // replicate ourselves the requested number of times
- result.add(splitSource);
- }
- return result;
- }
-
- @Override
- public PubsubReader<T> createReader(
- PipelineOptions options,
- @Nullable PubsubCheckpoint<T> checkpoint) {
- PubsubReader<T> reader;
- SubscriptionPath subscription = subscriptionPath;
- if (subscription == null) {
- if (checkpoint == null) {
- // This reader has never been started and there was no call to #splitIntoBundles; create
- // a single random subscription, which will be kept in the checkpoint.
- subscription = outer.createRandomSubscription(options);
- } else {
- subscription = checkpoint.getSubscription();
- }
- }
- try {
- reader = new PubsubReader<>(options.as(PubsubOptions.class), this, subscription,
- outer.parseFn);
- } catch (GeneralSecurityException | IOException e) {
- throw new RuntimeException("Unable to subscribe to " + subscriptionPath + ": ", e);
- }
- if (checkpoint != null) {
- // NACK all messages we may have lost.
- try {
- // Will BLOCK until NACKed.
- checkpoint.nackAll(reader);
- } catch (IOException e) {
- LOG.error("Pubsub {} cannot have {} lost messages NACKed, ignoring: {}",
- subscriptionPath, checkpoint.notYetReadIds.size(), e);
- }
- }
- return reader;
- }
-
- @Nullable
- @Override
- public Coder<PubsubCheckpoint<T>> getCheckpointMarkCoder() {
- @SuppressWarnings("unchecked") PubsubCheckpointCoder<T> typedCoder =
- (PubsubCheckpointCoder<T>) CHECKPOINT_CODER;
- return typedCoder;
- }
-
- @Override
- public Coder<T> getDefaultOutputCoder() {
- return outer.elementCoder;
- }
-
- @Override
- public void validate() {
- // Nothing to validate.
- }
-
- @Override
- public boolean requiresDeduping() {
- // We cannot prevent re-offering already read messages after a restore from checkpoint.
- return true;
- }
- }
-
- // ================================================================================
- // StatsFn
- // ================================================================================
-
- private static class StatsFn<T> extends DoFn<T, T> {
- private final Counter elementCounter = Metrics.counter(StatsFn.class, "elements");
-
- private final PubsubClientFactory pubsubFactory;
- @Nullable
- private final ValueProvider<SubscriptionPath> subscription;
- @Nullable
- private final ValueProvider<TopicPath> topic;
- @Nullable
- private final String timestampLabel;
- @Nullable
- private final String idLabel;
-
- public StatsFn(
- PubsubClientFactory pubsubFactory,
- @Nullable ValueProvider<SubscriptionPath> subscription,
- @Nullable ValueProvider<TopicPath> topic,
- @Nullable String timestampLabel,
- @Nullable String idLabel) {
- checkArgument(pubsubFactory != null, "pubsubFactory should not be null");
- this.pubsubFactory = pubsubFactory;
- this.subscription = subscription;
- this.topic = topic;
- this.timestampLabel = timestampLabel;
- this.idLabel = idLabel;
- }
-
- @ProcessElement
- public void processElement(ProcessContext c) throws Exception {
- elementCounter.inc();
- c.output(c.element());
- }
-
- @Override
- public void populateDisplayData(Builder builder) {
- super.populateDisplayData(builder);
- if (subscription != null) {
- String subscriptionString = subscription.isAccessible()
- ? subscription.get().getPath()
- : subscription.toString();
- builder.add(DisplayData.item("subscription", subscriptionString));
- }
- if (topic != null) {
- String topicString = topic.isAccessible()
- ? topic.get().getPath()
- : topic.toString();
- builder.add(DisplayData.item("topic", topicString));
- }
- builder.add(DisplayData.item("transport", pubsubFactory.getKind()));
- builder.addIfNotNull(DisplayData.item("timestampLabel", timestampLabel));
- builder.addIfNotNull(DisplayData.item("idLabel", idLabel));
- }
- }
-
- // ================================================================================
- // PubsubUnboundedSource
- // ================================================================================
-
- /**
- * For testing only: Clock to use for all timekeeping. If {@literal null} use system clock.
- */
- @Nullable
- private Clock clock;
-
- /**
- * Factory for creating underlying Pubsub transport.
- */
- private final PubsubClientFactory pubsubFactory;
-
- /**
- * Project under which to create a subscription if only the {@link #topic} was given.
- */
- @Nullable
- private final ValueProvider<ProjectPath> project;
-
- /**
- * Topic to read from. If {@literal null}, then {@link #subscription} must be given.
- * Otherwise {@link #subscription} must be null.
- */
- @Nullable
- private final ValueProvider<TopicPath> topic;
-
- /**
- * Subscription to read from. If {@literal null} then {@link #topic} must be given.
- * Otherwise {@link #topic} must be null.
- *
- * <p>If no subscription is given a random one will be created when the transorm is
- * applied. This field will be update with that subscription's path. The created
- * subscription is never deleted.
- */
- @Nullable
- private ValueProvider<SubscriptionPath> subscription;
-
- /**
- * Coder for elements. Elements are effectively double-encoded: first to a byte array
- * using this checkpointCoder, then to a base-64 string to conform to Pubsub's payload
- * conventions.
- */
- private final Coder<T> elementCoder;
-
- /**
- * Pubsub metadata field holding timestamp of each element, or {@literal null} if should use
- * Pubsub message publish timestamp instead.
- */
- @Nullable
- private final String timestampLabel;
-
- /**
- * Pubsub metadata field holding id for each element, or {@literal null} if need to generate
- * a unique id ourselves.
- */
- @Nullable
- private final String idLabel;
-
- /**
- * If not {@literal null}, the user is asking for PubSub attributes. This parse function will be
- * used to parse {@link PubsubIO.PubsubMessage}s containing a payload and attributes.
- */
- @Nullable
- SimpleFunction<PubsubMessage, T> parseFn;
-
- @VisibleForTesting
- PubsubUnboundedSource(
- Clock clock,
- PubsubClientFactory pubsubFactory,
- @Nullable ValueProvider<ProjectPath> project,
- @Nullable ValueProvider<TopicPath> topic,
- @Nullable ValueProvider<SubscriptionPath> subscription,
- Coder<T> elementCoder,
- @Nullable String timestampLabel,
- @Nullable String idLabel,
- @Nullable SimpleFunction<PubsubIO.PubsubMessage, T> parseFn) {
- checkArgument((topic == null) != (subscription == null),
- "Exactly one of topic and subscription must be given");
- checkArgument((topic == null) == (project == null),
- "Project must be given if topic is given");
- this.clock = clock;
- this.pubsubFactory = checkNotNull(pubsubFactory);
- this.project = project;
- this.topic = topic;
- this.subscription = subscription;
- this.elementCoder = checkNotNull(elementCoder);
- this.timestampLabel = timestampLabel;
- this.idLabel = idLabel;
- this.parseFn = parseFn;
- }
-
- /**
- * Construct an unbounded source to consume from the Pubsub {@code subscription}.
- */
- public PubsubUnboundedSource(
- PubsubClientFactory pubsubFactory,
- @Nullable ValueProvider<ProjectPath> project,
- @Nullable ValueProvider<TopicPath> topic,
- @Nullable ValueProvider<SubscriptionPath> subscription,
- Coder<T> elementCoder,
- @Nullable String timestampLabel,
- @Nullable String idLabel,
- @Nullable SimpleFunction<PubsubIO.PubsubMessage, T> parseFn) {
- this(null, pubsubFactory, project, topic, subscription, elementCoder, timestampLabel, idLabel,
- parseFn);
- }
-
- /**
- * Get the coder used for elements.
- */
- public Coder<T> getElementCoder() {
- return elementCoder;
- }
-
- /**
- * Get the project path.
- */
- @Nullable
- public ProjectPath getProject() {
- return project == null ? null : project.get();
- }
-
- /**
- * Get the topic being read from.
- */
- @Nullable
- public TopicPath getTopic() {
- return topic == null ? null : topic.get();
- }
-
- /**
- * Get the {@link ValueProvider} for the topic being read from.
- */
- @Nullable
- public ValueProvider<TopicPath> getTopicProvider() {
- return topic;
- }
-
- /**
- * Get the subscription being read from.
- */
- @Nullable
- public SubscriptionPath getSubscription() {
- return subscription == null ? null : subscription.get();
- }
-
- /**
- * Get the {@link ValueProvider} for the subscription being read from.
- */
- @Nullable
- public ValueProvider<SubscriptionPath> getSubscriptionProvider() {
- return subscription;
- }
-
- /**
- * Get the timestamp label.
- */
- @Nullable
- public String getTimestampLabel() {
- return timestampLabel;
- }
-
- /**
- * Get the id label.
- */
- @Nullable
- public String getIdLabel() {
- return idLabel;
- }
-
- /**
- * Get the parsing function for PubSub attributes.
- */
- @Nullable
- public SimpleFunction<PubsubIO.PubsubMessage, T> getWithAttributesParseFn() {
- return parseFn;
- }
-
- @Override
- public PCollection<T> expand(PBegin input) {
- return input.getPipeline().begin()
- .apply(Read.from(new PubsubSource<T>(this)))
- .apply("PubsubUnboundedSource.Stats",
- ParDo.of(new StatsFn<T>(
- pubsubFactory, subscription, topic, timestampLabel, idLabel)));
- }
-
- private SubscriptionPath createRandomSubscription(PipelineOptions options) {
- try {
- try (PubsubClient pubsubClient =
- pubsubFactory.newClient(timestampLabel, idLabel, options.as(PubsubOptions.class))) {
- checkState(project.isAccessible(), "createRandomSubscription must be called at runtime.");
- checkState(topic.isAccessible(), "createRandomSubscription must be called at runtime.");
- SubscriptionPath subscriptionPath =
- pubsubClient.createRandomSubscription(
- project.get(), topic.get(), DEAULT_ACK_TIMEOUT_SEC);
- LOG.warn(
- "Created subscription {} to topic {}."
- + " Note this subscription WILL NOT be deleted when the pipeline terminates",
- subscriptionPath,
- topic);
- return subscriptionPath;
- }
- } catch (Exception e) {
- throw new RuntimeException("Failed to create subscription: ", e);
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupByKey.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupByKey.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupByKey.java
index d228dbb..cc92102 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupByKey.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupByKey.java
@@ -101,9 +101,7 @@ import org.apache.beam.sdk.values.PCollection.IsBounded;
* The output {@code PCollection} will have the same {@link WindowFn}
* as the input.
*
- * <p>If the input {@code PCollection} contains late data (see
- * {@link org.apache.beam.sdk.io.PubsubIO.Read#timestampLabel}
- * for an example of how this can occur) or the
+ * <p>If the input {@code PCollection} contains late data or the
* {@link Window#triggering requested TriggerFn} can fire before
* the watermark, then there may be multiple elements
* output by a {@code GroupByKey} that correspond to the same key and window.
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/AfterWatermark.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/AfterWatermark.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/AfterWatermark.java
index 7cd2601..fb60538 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/AfterWatermark.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/AfterWatermark.java
@@ -33,18 +33,16 @@ import org.joda.time.Instant;
* lower-bound, sometimes heuristically established, on event times that have been fully processed
* by the pipeline.
*
- * <p>For sources that provide non-heuristic watermarks (e.g.
- * {@link org.apache.beam.sdk.io.PubsubIO} when using arrival times as event times), the
- * watermark is a strict guarantee that no data with an event time earlier than
+ * <p>For sources that provide non-heuristic watermarks (e.g. PubsubIO when using arrival times as
+ * event times), the watermark is a strict guarantee that no data with an event time earlier than
* that watermark will ever be observed in the pipeline. In this case, it's safe to assume that any
* pane triggered by an {@code AfterWatermark} trigger with a reference point at or beyond the end
* of the window will be the last pane ever for that window.
*
- * <p>For sources that provide heuristic watermarks (e.g.
- * {@link org.apache.beam.sdk.io.PubsubIO} when using user-supplied event times), the
- * watermark itself becomes an <i>estimate</i> that no data with an event time earlier than that
- * watermark (i.e. "late data") will ever be observed in the pipeline. These heuristics can
- * often be quite accurate, but the chance of seeing late data for any given window is non-zero.
+ * <p>For sources that provide heuristic watermarks (e.g. PubsubIO when using user-supplied event
+ * times), the watermark itself becomes an <i>estimate</i> that no data with an event time earlier
+ * than that watermark (i.e. "late data") will ever be observed in the pipeline. These heuristics
+ * can often be quite accurate, but the chance of seeing late data for any given window is non-zero.
* Thus, if absolute correctness over time is important to your use case, you may want to consider
* using a trigger that accounts for late data. The default trigger,
* {@code Repeatedly.forever(AfterWatermark.pastEndOfWindow())}, which fires
http://git-wip-us.apache.org/repos/asf/beam/blob/5bcb8c57/sdks/java/core/src/main/java/org/apache/beam/sdk/util/PubsubClient.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/PubsubClient.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/PubsubClient.java
deleted file mode 100644
index fc84057..0000000
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/PubsubClient.java
+++ /dev/null
@@ -1,544 +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.util;
-
-import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Preconditions.checkState;
-
-import com.google.api.client.util.DateTime;
-import com.google.common.base.Objects;
-import com.google.common.base.Strings;
-import java.io.Closeable;
-import java.io.IOException;
-import java.io.Serializable;
-import java.util.Arrays;
-import java.util.List;
-import java.util.Map;
-import java.util.concurrent.ThreadLocalRandom;
-import javax.annotation.Nullable;
-import org.apache.beam.sdk.options.PubsubOptions;
-
-/**
- * An (abstract) helper class for talking to Pubsub via an underlying transport.
- */
-public abstract class PubsubClient implements Closeable {
- /**
- * Factory for creating clients.
- */
- public interface PubsubClientFactory extends Serializable {
- /**
- * Construct a new Pubsub client. It should be closed via {@link #close} in order
- * to ensure tidy cleanup of underlying netty resources (or use the try-with-resources
- * construct). Uses {@code options} to derive pubsub endpoints and application credentials.
- * If non-{@literal null}, use {@code timestampLabel} and {@code idLabel} to store custom
- * timestamps/ids within message metadata.
- */
- PubsubClient newClient(
- @Nullable String timestampLabel,
- @Nullable String idLabel,
- PubsubOptions options) throws IOException;
-
- /**
- * Return the display name for this factory. Eg "Json", "gRPC".
- */
- String getKind();
- }
-
- /**
- * Return timestamp as ms-since-unix-epoch corresponding to {@code timestamp}.
- * Return {@literal null} if no timestamp could be found. Throw {@link IllegalArgumentException}
- * if timestamp cannot be recognized.
- */
- @Nullable
- private static Long asMsSinceEpoch(@Nullable String timestamp) {
- if (Strings.isNullOrEmpty(timestamp)) {
- return null;
- }
- try {
- // Try parsing as milliseconds since epoch. Note there is no way to parse a
- // string in RFC 3339 format here.
- // Expected IllegalArgumentException if parsing fails; we use that to fall back
- // to RFC 3339.
- return Long.parseLong(timestamp);
- } catch (IllegalArgumentException e1) {
- // Try parsing as RFC3339 string. DateTime.parseRfc3339 will throw an
- // IllegalArgumentException if parsing fails, and the caller should handle.
- return DateTime.parseRfc3339(timestamp).getValue();
- }
- }
-
- /**
- * Return the timestamp (in ms since unix epoch) to use for a Pubsub message with {@code
- * attributes} and {@code pubsubTimestamp}.
- *
- * <p>If {@code timestampLabel} is non-{@literal null} then the message attributes must contain
- * that label, and the value of that label will be taken as the timestamp.
- * Otherwise the timestamp will be taken from the Pubsub publish timestamp {@code
- * pubsubTimestamp}.
- *
- * @throws IllegalArgumentException if the timestamp cannot be recognized as a ms-since-unix-epoch
- * or RFC3339 time.
- */
- protected static long extractTimestamp(
- @Nullable String timestampLabel,
- @Nullable String pubsubTimestamp,
- @Nullable Map<String, String> attributes) {
- Long timestampMsSinceEpoch;
- if (Strings.isNullOrEmpty(timestampLabel)) {
- timestampMsSinceEpoch = asMsSinceEpoch(pubsubTimestamp);
- checkArgument(timestampMsSinceEpoch != null,
- "Cannot interpret PubSub publish timestamp: %s",
- pubsubTimestamp);
- } else {
- String value = attributes == null ? null : attributes.get(timestampLabel);
- checkArgument(value != null,
- "PubSub message is missing a value for timestamp label %s",
- timestampLabel);
- timestampMsSinceEpoch = asMsSinceEpoch(value);
- checkArgument(timestampMsSinceEpoch != null,
- "Cannot interpret value of label %s as timestamp: %s",
- timestampLabel, value);
- }
- return timestampMsSinceEpoch;
- }
-
- /**
- * Path representing a cloud project id.
- */
- public static class ProjectPath implements Serializable {
- private final String projectId;
-
- /**
- * Creates a {@link ProjectPath} from a {@link String} representation, which
- * must be of the form {@code "projects/" + projectId}.
- */
- ProjectPath(String path) {
- String[] splits = path.split("/");
- checkArgument(
- splits.length == 2 && splits[0].equals("projects"),
- "Malformed project path \"%s\": must be of the form \"projects/\" + <project id>",
- path);
- this.projectId = splits[1];
- }
-
- public String getPath() {
- return String.format("projects/%s", projectId);
- }
-
- public String getId() {
- return projectId;
- }
-
- @Override
- public boolean equals(Object o) {
- if (this == o) {
- return true;
- }
- if (o == null || getClass() != o.getClass()) {
- return false;
- }
-
- ProjectPath that = (ProjectPath) o;
-
- return projectId.equals(that.projectId);
- }
-
- @Override
- public int hashCode() {
- return projectId.hashCode();
- }
-
- @Override
- public String toString() {
- return getPath();
- }
- }
-
- public static ProjectPath projectPathFromPath(String path) {
- return new ProjectPath(path);
- }
-
- public static ProjectPath projectPathFromId(String projectId) {
- return new ProjectPath(String.format("projects/%s", projectId));
- }
-
- /**
- * Path representing a Pubsub subscription.
- */
- public static class SubscriptionPath implements Serializable {
- private final String projectId;
- private final String subscriptionName;
-
- SubscriptionPath(String path) {
- String[] splits = path.split("/");
- checkState(
- splits.length == 4 && splits[0].equals("projects") && splits[2].equals("subscriptions"),
- "Malformed subscription path %s: "
- + "must be of the form \"projects/\" + <project id> + \"subscriptions\"", path);
- this.projectId = splits[1];
- this.subscriptionName = splits[3];
- }
-
- public String getPath() {
- return String.format("projects/%s/subscriptions/%s", projectId, subscriptionName);
- }
-
- public String getName() {
- return subscriptionName;
- }
-
- public String getV1Beta1Path() {
- return String.format("/subscriptions/%s/%s", projectId, subscriptionName);
- }
-
- @Override
- public boolean equals(Object o) {
- if (this == o) {
- return true;
- }
- if (o == null || getClass() != o.getClass()) {
- return false;
- }
- SubscriptionPath that = (SubscriptionPath) o;
- return this.subscriptionName.equals(that.subscriptionName)
- && this.projectId.equals(that.projectId);
- }
-
- @Override
- public int hashCode() {
- return Objects.hashCode(projectId, subscriptionName);
- }
-
- @Override
- public String toString() {
- return getPath();
- }
- }
-
- public static SubscriptionPath subscriptionPathFromPath(String path) {
- return new SubscriptionPath(path);
- }
-
- public static SubscriptionPath subscriptionPathFromName(
- String projectId, String subscriptionName) {
- return new SubscriptionPath(String.format("projects/%s/subscriptions/%s",
- projectId, subscriptionName));
- }
-
- /**
- * Path representing a Pubsub topic.
- */
- public static class TopicPath implements Serializable {
- private final String path;
-
- TopicPath(String path) {
- this.path = path;
- }
-
- public String getPath() {
- return path;
- }
-
- public String getName() {
- String[] splits = path.split("/");
- checkState(splits.length == 4, "Malformed topic path %s", path);
- return splits[3];
- }
-
- public String getV1Beta1Path() {
- String[] splits = path.split("/");
- checkState(splits.length == 4, "Malformed topic path %s", path);
- return String.format("/topics/%s/%s", splits[1], splits[3]);
- }
-
- @Override
- public boolean equals(Object o) {
- if (this == o) {
- return true;
- }
- if (o == null || getClass() != o.getClass()) {
- return false;
- }
- TopicPath topicPath = (TopicPath) o;
- return path.equals(topicPath.path);
- }
-
- @Override
- public int hashCode() {
- return path.hashCode();
- }
-
- @Override
- public String toString() {
- return path;
- }
- }
-
- public static TopicPath topicPathFromPath(String path) {
- return new TopicPath(path);
- }
-
- public static TopicPath topicPathFromName(String projectId, String topicName) {
- return new TopicPath(String.format("projects/%s/topics/%s", projectId, topicName));
- }
-
- /**
- * A message to be sent to Pubsub.
- *
- * <p>NOTE: This class is {@link Serializable} only to support the {@link PubsubTestClient}.
- * Java serialization is never used for non-test clients.
- */
- public static class OutgoingMessage implements Serializable {
- /**
- * Underlying (encoded) element.
- */
- public final byte[] elementBytes;
-
- public final Map<String, String> attributes;
-
- /**
- * Timestamp for element (ms since epoch).
- */
- public final long timestampMsSinceEpoch;
-
- /**
- * If using an id label, the record id to associate with this record's metadata so the receiver
- * can reject duplicates. Otherwise {@literal null}.
- */
- @Nullable
- public final String recordId;
-
- public OutgoingMessage(byte[] elementBytes, Map<String, String> attributes,
- long timestampMsSinceEpoch, @Nullable String recordId) {
- this.elementBytes = elementBytes;
- this.attributes = attributes;
- this.timestampMsSinceEpoch = timestampMsSinceEpoch;
- this.recordId = recordId;
- }
-
- @Override
- public String toString() {
- return String.format("OutgoingMessage(%db, %dms)",
- elementBytes.length, timestampMsSinceEpoch);
- }
-
- @Override
- public boolean equals(Object o) {
- if (this == o) {
- return true;
- }
- if (o == null || getClass() != o.getClass()) {
- return false;
- }
-
- OutgoingMessage that = (OutgoingMessage) o;
-
- return timestampMsSinceEpoch == that.timestampMsSinceEpoch
- && Arrays.equals(elementBytes, that.elementBytes)
- && Objects.equal(attributes, that.attributes)
- && Objects.equal(recordId, that.recordId);
- }
-
- @Override
- public int hashCode() {
- return Objects.hashCode(Arrays.hashCode(elementBytes), attributes, timestampMsSinceEpoch,
- recordId);
- }
- }
-
- /**
- * A message received from Pubsub.
- *
- * <p>NOTE: This class is {@link Serializable} only to support the {@link PubsubTestClient}.
- * Java serialization is never used for non-test clients.
- */
- public static class IncomingMessage implements Serializable {
- /**
- * Underlying (encoded) element.
- */
- public final byte[] elementBytes;
-
- public Map<String, String> attributes;
-
- /**
- * Timestamp for element (ms since epoch). Either Pubsub's processing time,
- * or the custom timestamp associated with the message.
- */
- public final long timestampMsSinceEpoch;
-
- /**
- * Timestamp (in system time) at which we requested the message (ms since epoch).
- */
- public final long requestTimeMsSinceEpoch;
-
- /**
- * Id to pass back to Pubsub to acknowledge receipt of this message.
- */
- public final String ackId;
-
- /**
- * Id to pass to the runner to distinguish this message from all others.
- */
- public final String recordId;
-
- public IncomingMessage(
- byte[] elementBytes,
- Map<String, String> attributes,
- long timestampMsSinceEpoch,
- long requestTimeMsSinceEpoch,
- String ackId,
- String recordId) {
- this.elementBytes = elementBytes;
- this.attributes = attributes;
- this.timestampMsSinceEpoch = timestampMsSinceEpoch;
- this.requestTimeMsSinceEpoch = requestTimeMsSinceEpoch;
- this.ackId = ackId;
- this.recordId = recordId;
- }
-
- public IncomingMessage withRequestTime(long requestTimeMsSinceEpoch) {
- return new IncomingMessage(elementBytes, attributes, timestampMsSinceEpoch,
- requestTimeMsSinceEpoch, ackId, recordId);
- }
-
- @Override
- public String toString() {
- return String.format("IncomingMessage(%db, %dms)",
- elementBytes.length, timestampMsSinceEpoch);
- }
-
- @Override
- public boolean equals(Object o) {
- if (this == o) {
- return true;
- }
- if (o == null || getClass() != o.getClass()) {
- return false;
- }
-
- IncomingMessage that = (IncomingMessage) o;
-
- return timestampMsSinceEpoch == that.timestampMsSinceEpoch
- && requestTimeMsSinceEpoch == that.requestTimeMsSinceEpoch
- && ackId.equals(that.ackId)
- && recordId.equals(that.recordId)
- && Arrays.equals(elementBytes, that.elementBytes)
- && Objects.equal(attributes, that.attributes);
- }
-
- @Override
- public int hashCode() {
- return Objects.hashCode(Arrays.hashCode(elementBytes), attributes, timestampMsSinceEpoch,
- requestTimeMsSinceEpoch,
- ackId, recordId);
- }
- }
-
- /**
- * Publish {@code outgoingMessages} to Pubsub {@code topic}. Return number of messages
- * published.
- */
- public abstract int publish(TopicPath topic, List<OutgoingMessage> outgoingMessages)
- throws IOException;
-
- /**
- * Request the next batch of up to {@code batchSize} messages from {@code subscription}.
- * Return the received messages, or empty collection if none were available. Does not
- * wait for messages to arrive if {@code returnImmediately} is {@literal true}.
- * Returned messages will record their request time as {@code requestTimeMsSinceEpoch}.
- */
- public abstract List<IncomingMessage> pull(
- long requestTimeMsSinceEpoch,
- SubscriptionPath subscription,
- int batchSize,
- boolean returnImmediately)
- throws IOException;
-
- /**
- * Acknowldege messages from {@code subscription} with {@code ackIds}.
- */
- public abstract void acknowledge(SubscriptionPath subscription, List<String> ackIds)
- throws IOException;
-
- /**
- * Modify the ack deadline for messages from {@code subscription} with {@code ackIds} to
- * be {@code deadlineSeconds} from now.
- */
- public abstract void modifyAckDeadline(
- SubscriptionPath subscription, List<String> ackIds,
- int deadlineSeconds) throws IOException;
-
- /**
- * Create {@code topic}.
- */
- public abstract void createTopic(TopicPath topic) throws IOException;
-
- /*
- * Delete {@code topic}.
- */
- public abstract void deleteTopic(TopicPath topic) throws IOException;
-
- /**
- * Return a list of topics for {@code project}.
- */
- public abstract List<TopicPath> listTopics(ProjectPath project) throws IOException;
-
- /**
- * Create {@code subscription} to {@code topic}.
- */
- public abstract void createSubscription(
- TopicPath topic, SubscriptionPath subscription, int ackDeadlineSeconds) throws IOException;
-
- /**
- * Create a random subscription for {@code topic}. Return the {@link SubscriptionPath}. It
- * is the responsibility of the caller to later delete the subscription.
- */
- public SubscriptionPath createRandomSubscription(
- ProjectPath project, TopicPath topic, int ackDeadlineSeconds) throws IOException {
- // Create a randomized subscription derived from the topic name.
- String subscriptionName = topic.getName() + "_beam_" + ThreadLocalRandom.current().nextLong();
- SubscriptionPath subscription =
- PubsubClient.subscriptionPathFromName(project.getId(), subscriptionName);
- createSubscription(topic, subscription, ackDeadlineSeconds);
- return subscription;
- }
-
- /**
- * Delete {@code subscription}.
- */
- public abstract void deleteSubscription(SubscriptionPath subscription) throws IOException;
-
- /**
- * Return a list of subscriptions for {@code topic} in {@code project}.
- */
- public abstract List<SubscriptionPath> listSubscriptions(ProjectPath project, TopicPath topic)
- throws IOException;
-
- /**
- * Return the ack deadline, in seconds, for {@code subscription}.
- */
- public abstract int ackDeadlineSeconds(SubscriptionPath subscription) throws IOException;
-
- /**
- * Return {@literal true} if {@link #pull} will always return empty list. Actual clients
- * will return {@literal false}. Test clients may return {@literal true} to signal that all
- * expected messages have been pulled and the test may complete.
- */
- public abstract boolean isEOF();
-}
[16/50] [abbrv] beam git commit: DataflowRunner misc cleanups
Posted by jb...@apache.org.
DataflowRunner misc cleanups
Bad use of format string, very old debugging assert
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/3fb3a0c5
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/3fb3a0c5
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/3fb3a0c5
Branch: refs/heads/DSL_SQL
Commit: 3fb3a0c5f200bd04387a3fbdc89f0855e737b8c8
Parents: 7587d29
Author: Dan Halperin <dh...@google.com>
Authored: Tue Apr 11 13:47:34 2017 -0700
Committer: Dan Halperin <dh...@google.com>
Committed: Tue Apr 11 15:01:34 2017 -0700
----------------------------------------------------------------------
.../beam/runners/dataflow/DataflowPipelineTranslator.java | 3 ---
.../java/org/apache/beam/runners/dataflow/DataflowRunner.java | 7 ++++---
2 files changed, 4 insertions(+), 6 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/3fb3a0c5/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 1a2e663..34da996 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
@@ -23,7 +23,6 @@ import static com.google.common.base.Preconditions.checkState;
import static com.google.common.base.Strings.isNullOrEmpty;
import static org.apache.beam.sdk.util.SerializableUtils.serializeToByteArray;
import static org.apache.beam.sdk.util.StringUtils.byteArrayToJsonString;
-import static org.apache.beam.sdk.util.StringUtils.jsonStringToByteArray;
import static org.apache.beam.sdk.util.Structs.addBoolean;
import static org.apache.beam.sdk.util.Structs.addDictionary;
import static org.apache.beam.sdk.util.Structs.addList;
@@ -49,7 +48,6 @@ import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import java.io.IOException;
import java.util.ArrayList;
-import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedList;
@@ -895,7 +893,6 @@ public class DataflowPipelineTranslator {
WindowingStrategy<?, ?> strategy = context.getOutput(transform).getWindowingStrategy();
byte[] serializedBytes = serializeToByteArray(strategy);
String serializedJson = byteArrayToJsonString(serializedBytes);
- assert Arrays.equals(serializedBytes, jsonStringToByteArray(serializedJson));
stepContext.addInput(PropertyNames.SERIALIZED_FN, serializedJson);
}
});
http://git-wip-us.apache.org/repos/asf/beam/blob/3fb3a0c5/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 9b993f4..1708d8a 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
@@ -580,9 +580,10 @@ public class DataflowRunner extends PipelineRunner<DataflowPipelineJob> {
}
String fileLocation = firstNonNull(
options.getTemplateLocation(), options.getDataflowJobFile());
- checkArgument(fileLocation.startsWith("/") || fileLocation.startsWith("gs://"),
- String.format(
- "Location must be local or on Cloud Storage, got {}.", fileLocation));
+ checkArgument(
+ fileLocation.startsWith("/") || fileLocation.startsWith("gs://"),
+ "Location must be local or on Cloud Storage, got %s.",
+ fileLocation);
String workSpecJson = DataflowPipelineTranslator.jobToString(newJob);
try (PrintWriter printWriter = new PrintWriter(
Channels.newOutputStream(IOChannelUtils.create(fileLocation, MimeTypes.TEXT)))) {
[41/50] [abbrv] beam git commit: [BEAM-1882] This closes #2431
Posted by jb...@apache.org.
[BEAM-1882] This closes #2431
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/f3143540
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/f3143540
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/f3143540
Branch: refs/heads/DSL_SQL
Commit: f3143540378fbd84f216597f168ae0a158ddc819
Parents: 435c3b4 4d1feeb
Author: Jean-Baptiste Onofr� <jb...@apache.org>
Authored: Wed Apr 12 20:07:30 2017 +0200
Committer: Jean-Baptiste Onofr� <jb...@apache.org>
Committed: Wed Apr 12 20:07:30 2017 +0200
----------------------------------------------------------------------
.../kubernetes/postgres-pod-no-vol.yml | 32 -----------
.../postgres-service-for-local-dev.yml | 28 ++++++++++
.../kubernetes/postgres-service-public.yml | 28 ----------
.../src/test/resources/kubernetes/postgres.yml | 56 ++++++++++++++++++++
.../jdbc/src/test/resources/kubernetes/setup.sh | 3 +-
.../src/test/resources/kubernetes/teardown.sh | 3 +-
6 files changed, 86 insertions(+), 64 deletions(-)
----------------------------------------------------------------------
[08/50] [abbrv] beam git commit: Ignore more python build artifacts.
Posted by jb...@apache.org.
Ignore more python build artifacts.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/a1683578
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/a1683578
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/a1683578
Branch: refs/heads/DSL_SQL
Commit: a168357885b17bde688a548ed1fa12231a49658d
Parents: 77712c9
Author: Robert Bradshaw <ro...@gmail.com>
Authored: Tue Apr 11 11:49:16 2017 -0700
Committer: Robert Bradshaw <ro...@gmail.com>
Committed: Tue Apr 11 11:49:16 2017 -0700
----------------------------------------------------------------------
.gitignore | 3 +++
1 file changed, 3 insertions(+)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/a1683578/.gitignore
----------------------------------------------------------------------
diff --git a/.gitignore b/.gitignore
index bc9f675..69946a9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -19,6 +19,9 @@ build/
dist/
distribute-*
env/
+sdks/python/**/*.c
+sdks/python/**/*.so
+sdks/python/**/*.egg
# Ignore IntelliJ files.
.idea/
[24/50] [abbrv] beam git commit: Clean HFIOWithEmbeddedCassandraTest
before Execution
Posted by jb...@apache.org.
Clean HFIOWithEmbeddedCassandraTest before Execution
HIFIO Cassandra tests were failing if run twice in a row without a clean
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/041b636b
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/041b636b
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/041b636b
Branch: refs/heads/DSL_SQL
Commit: 041b636b0d912242ae1292c6fb2e5d0fa12fdaab
Parents: a2d328a
Author: Stephen Sisk <si...@google.com>
Authored: Tue Apr 11 14:36:05 2017 -0700
Committer: Thomas Groh <tg...@google.com>
Committed: Tue Apr 11 16:17:23 2017 -0700
----------------------------------------------------------------------
.../sdk/io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java | 1 +
1 file changed, 1 insertion(+)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/041b636b/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java b/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java
index 402c19c..50b2c4f 100644
--- a/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java
+++ b/sdks/java/io/hadoop/jdk1.8-tests/src/test/java/org/apache/beam/sdk/io/hadoop/inputformat/HIFIOWithEmbeddedCassandraTest.java
@@ -147,6 +147,7 @@ public class HIFIOWithEmbeddedCassandraTest implements Serializable {
}
public static void createCassandraData() throws Exception {
+ session.execute("DROP KEYSPACE IF EXISTS " + CASSANDRA_KEYSPACE);
session.execute("CREATE KEYSPACE " + CASSANDRA_KEYSPACE
+ " WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':1};");
session.execute("USE " + CASSANDRA_KEYSPACE);