You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by jb...@apache.org on 2017/07/20 17:09:37 UTC
[10/28] beam git commit: Revert "[BEAM-2610] This closes #3553"
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/ShardedKeyCoder.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/ShardedKeyCoder.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/ShardedKeyCoder.java
new file mode 100644
index 0000000..c2b62b7
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/ShardedKeyCoder.java
@@ -0,0 +1,74 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.beam.sdk.io.gcp.bigquery;
+
+import com.google.common.annotations.VisibleForTesting;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.Arrays;
+import java.util.List;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.StructuredCoder;
+import org.apache.beam.sdk.coders.VarIntCoder;
+
+
+/**
+ * A {@link Coder} for {@link ShardedKey}, using a wrapped key {@link Coder}.
+ */
+@VisibleForTesting
+class ShardedKeyCoder<KeyT>
+ extends StructuredCoder<ShardedKey<KeyT>> {
+ public static <KeyT> ShardedKeyCoder<KeyT> of(Coder<KeyT> keyCoder) {
+ return new ShardedKeyCoder<>(keyCoder);
+ }
+
+ private final Coder<KeyT> keyCoder;
+ private final VarIntCoder shardNumberCoder;
+
+ protected ShardedKeyCoder(Coder<KeyT> keyCoder) {
+ this.keyCoder = keyCoder;
+ this.shardNumberCoder = VarIntCoder.of();
+ }
+
+ @Override
+ public List<? extends Coder<?>> getCoderArguments() {
+ return Arrays.asList(keyCoder);
+ }
+
+ @Override
+ public void encode(ShardedKey<KeyT> key, OutputStream outStream)
+ throws IOException {
+ keyCoder.encode(key.getKey(), outStream);
+ shardNumberCoder.encode(key.getShardNumber(), outStream);
+ }
+
+ @Override
+ public ShardedKey<KeyT> decode(InputStream inStream)
+ throws IOException {
+ return new ShardedKey<>(
+ keyCoder.decode(inStream),
+ shardNumberCoder.decode(inStream));
+ }
+
+ @Override
+ public void verifyDeterministic() throws NonDeterministicException {
+ keyCoder.verifyDeterministic();
+ }
+}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteFn.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteFn.java
index a210858..63e5bc1 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteFn.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteFn.java
@@ -33,7 +33,6 @@ import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.SystemDoFnInternal;
import org.apache.beam.sdk.values.KV;
-import org.apache.beam.sdk.values.ShardedKey;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.ValueInSingleWindow;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java
index fa5b3ce..18b2033 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java
@@ -19,7 +19,6 @@ package org.apache.beam.sdk.io.gcp.bigquery;
import com.google.api.services.bigquery.model.TableRow;
import org.apache.beam.sdk.coders.KvCoder;
-import org.apache.beam.sdk.coders.ShardedKeyCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
@@ -30,7 +29,6 @@ import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionTuple;
-import org.apache.beam.sdk.values.ShardedKey;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TagWithUniqueIds.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TagWithUniqueIds.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TagWithUniqueIds.java
index 51b9375..cd88222 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TagWithUniqueIds.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TagWithUniqueIds.java
@@ -26,7 +26,6 @@ import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.values.KV;
-import org.apache.beam.sdk.values.ShardedKey;
/**
* Fn that tags each table row with a unique id and destination table. To avoid calling
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteBundlesToFiles.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteBundlesToFiles.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteBundlesToFiles.java
index e1ed746..d68779a 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteBundlesToFiles.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteBundlesToFiles.java
@@ -19,7 +19,6 @@
package org.apache.beam.sdk.io.gcp.bigquery;
import static com.google.common.base.Preconditions.checkNotNull;
-
import com.google.api.services.bigquery.model.TableRow;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
@@ -41,7 +40,6 @@ import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.ShardedKey;
import org.apache.beam.sdk.values.TupleTag;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteGroupedRecordsToFiles.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteGroupedRecordsToFiles.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteGroupedRecordsToFiles.java
index 887cb93..45dc2a8 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteGroupedRecordsToFiles.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteGroupedRecordsToFiles.java
@@ -22,7 +22,6 @@ import com.google.api.services.bigquery.model.TableRow;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.ShardedKey;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WritePartition.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WritePartition.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WritePartition.java
index 451d1bd..acd1132 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WritePartition.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WritePartition.java
@@ -26,7 +26,6 @@ import org.apache.beam.sdk.io.gcp.bigquery.WriteBundlesToFiles.Result;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.ShardedKey;
import org.apache.beam.sdk.values.TupleTag;
/**
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java
index 9ed2916..c5494d8 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java
@@ -42,7 +42,6 @@ import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.JobService;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
-import org.apache.beam.sdk.values.ShardedKey;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java
index 0a90dde..1692cda 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java
@@ -175,7 +175,7 @@ import org.slf4j.LoggerFactory;
* pipeline. Please refer to the documentation of corresponding
* {@link PipelineRunner PipelineRunners} for more details.
*/
-@Experimental(Experimental.Kind.SOURCE_SINK)
+@Experimental
public class BigtableIO {
private static final Logger LOG = LoggerFactory.getLogger(BigtableIO.class);
@@ -211,7 +211,7 @@ public class BigtableIO {
*
* @see BigtableIO
*/
- @Experimental(Experimental.Kind.SOURCE_SINK)
+ @Experimental
@AutoValue
public abstract static class Read extends PTransform<PBegin, PCollection<Row>> {
@@ -415,7 +415,7 @@ public class BigtableIO {
*
* @see BigtableIO
*/
- @Experimental(Experimental.Kind.SOURCE_SINK)
+ @Experimental
@AutoValue
public abstract static class Write
extends PTransform<PCollection<KV<ByteString, Iterable<Mutation>>>, PDone> {
@@ -1027,7 +1027,7 @@ public class BigtableIO {
"{}: Failed to interpolate key for fraction {}.", rangeTracker.getRange(), fraction, e);
return null;
}
- LOG.info(
+ LOG.debug(
"Proposing to split {} at fraction {} (key {})", rangeTracker, fraction, splitKey);
BigtableSource primary;
BigtableSource residual;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImpl.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImpl.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImpl.java
index 07476e2..d1a17fe 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImpl.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableServiceImpl.java
@@ -168,13 +168,14 @@ class BigtableServiceImpl implements BigtableService {
private BigtableSession session;
private AsyncExecutor executor;
private BulkMutation bulkMutation;
- private final String tableName;
+ private final MutateRowRequest.Builder partialBuilder;
public BigtableWriterImpl(BigtableSession session, BigtableTableName tableName) {
this.session = session;
executor = session.createAsyncExecutor();
bulkMutation = session.createBulkMutation(tableName, executor);
- this.tableName = tableName.toString();
+
+ partialBuilder = MutateRowRequest.newBuilder().setTableName(tableName.toString());
}
@Override
@@ -207,8 +208,8 @@ class BigtableServiceImpl implements BigtableService {
KV<ByteString, Iterable<Mutation>> record)
throws IOException {
MutateRowRequest r =
- MutateRowRequest.newBuilder()
- .setTableName(tableName)
+ partialBuilder
+ .clone()
.setRowKey(record.getKey())
.addAllMutations(record.getValue())
.build();
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottler.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottler.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottler.java
deleted file mode 100644
index ce6ebe6..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/AdaptiveThrottler.java
+++ /dev/null
@@ -1,103 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.beam.sdk.io.gcp.datastore;
-
-import com.google.common.annotations.VisibleForTesting;
-import java.util.Random;
-import org.apache.beam.sdk.transforms.Sum;
-import org.apache.beam.sdk.util.MovingFunction;
-
-
-/**
- * An implementation of client-side adaptive throttling. See
- * https://landing.google.com/sre/book/chapters/handling-overload.html#client-side-throttling-a7sYUg
- * for a full discussion of the use case and algorithm applied.
- */
-class AdaptiveThrottler {
- private final MovingFunction successfulRequests;
- private final MovingFunction allRequests;
-
- /** The target ratio between requests sent and successful requests. This is "K" in the formula in
- * https://landing.google.com/sre/book/chapters/handling-overload.html */
- private final double overloadRatio;
-
- /** The target minimum number of requests per samplePeriodMs, even if no requests succeed. Must be
- * greater than 0, else we could throttle to zero. Because every decision is probabilistic, there
- * is no guarantee that the request rate in any given interval will not be zero. (This is the +1
- * from the formula in https://landing.google.com/sre/book/chapters/handling-overload.html */
- private static final double MIN_REQUESTS = 1;
- private final Random random;
-
- /**
- * @param samplePeriodMs the time window to keep of request history to inform throttling
- * decisions.
- * @param sampleUpdateMs the length of buckets within this time window.
- * @param overloadRatio the target ratio between requests sent and successful requests. You should
- * always set this to more than 1, otherwise the client would never try to send more requests than
- * succeeded in the past - so it could never recover from temporary setbacks.
- */
- public AdaptiveThrottler(long samplePeriodMs, long sampleUpdateMs,
- double overloadRatio) {
- this(samplePeriodMs, sampleUpdateMs, overloadRatio, new Random());
- }
-
- @VisibleForTesting
- AdaptiveThrottler(long samplePeriodMs, long sampleUpdateMs,
- double overloadRatio, Random random) {
- allRequests =
- new MovingFunction(samplePeriodMs, sampleUpdateMs,
- 1 /* numSignificantBuckets */, 1 /* numSignificantSamples */, Sum.ofLongs());
- successfulRequests =
- new MovingFunction(samplePeriodMs, sampleUpdateMs,
- 1 /* numSignificantBuckets */, 1 /* numSignificantSamples */, Sum.ofLongs());
- this.overloadRatio = overloadRatio;
- this.random = random;
- }
-
- @VisibleForTesting
- double throttlingProbability(long nowMsSinceEpoch) {
- if (!allRequests.isSignificant()) {
- return 0;
- }
- long allRequestsNow = allRequests.get(nowMsSinceEpoch);
- long successfulRequestsNow = successfulRequests.get(nowMsSinceEpoch);
- return Math.max(0,
- (allRequestsNow - overloadRatio * successfulRequestsNow) / (allRequestsNow + MIN_REQUESTS));
- }
-
- /**
- * Call this before sending a request to the remote service; if this returns true, drop the
- * request (treating it as a failure or trying it again at a later time).
- */
- public boolean throttleRequest(long nowMsSinceEpoch) {
- double delayProbability = throttlingProbability(nowMsSinceEpoch);
- // Note that we increment the count of all requests here, even if we return true - so even if we
- // tell the client not to send a request at all, it still counts as a failed request.
- allRequests.add(nowMsSinceEpoch, 1);
-
- return (random.nextDouble() < delayProbability);
- }
-
- /**
- * Call this after {@link throttleRequest} if your request was successful.
- */
- public void successfulRequest(long nowMsSinceEpoch) {
- successfulRequests.add(nowMsSinceEpoch, 1);
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
index 5f65428..b198a6f 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
@@ -71,8 +71,6 @@ import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.annotations.Experimental.Kind;
import org.apache.beam.sdk.coders.SerializableCoder;
import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
-import org.apache.beam.sdk.metrics.Counter;
-import org.apache.beam.sdk.metrics.Metrics;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
@@ -203,31 +201,11 @@ public class DatastoreV1 {
DatastoreV1() {}
/**
- * The number of entity updates written per RPC, initially. We buffer updates in the connector and
- * write a batch to Datastore once we have collected a certain number. This is the initial batch
- * size; it is adjusted at runtime based on the performance of previous writes (see {@link
- * DatastoreV1.WriteBatcher}).
- *
- * <p>Testing has found that a batch of 200 entities will generally finish within the timeout even
- * in adverse conditions.
- */
- @VisibleForTesting
- static final int DATASTORE_BATCH_UPDATE_ENTITIES_START = 200;
-
- /**
- * When choosing the number of updates in a single RPC, never exceed the maximum allowed by the
- * API.
+ * Cloud Datastore has a limit of 500 mutations per batch operation, so we flush
+ * changes to Datastore every 500 entities.
*/
@VisibleForTesting
- static final int DATASTORE_BATCH_UPDATE_ENTITIES_LIMIT = 500;
-
- /**
- * When choosing the number of updates in a single RPC, do not go below this value. The actual
- * number of entities per request may be lower when we flush for the end of a bundle or if we hit
- * {@link DatastoreV1.DATASTORE_BATCH_UPDATE_BYTES_LIMIT}.
- */
- @VisibleForTesting
- static final int DATASTORE_BATCH_UPDATE_ENTITIES_MIN = 10;
+ static final int DATASTORE_BATCH_UPDATE_LIMIT = 500;
/**
* Cloud Datastore has a limit of 10MB per RPC, so we also flush if the total size of mutations
@@ -235,7 +213,7 @@ public class DatastoreV1 {
* the mutations themselves and not the CommitRequest wrapper around them.
*/
@VisibleForTesting
- static final int DATASTORE_BATCH_UPDATE_BYTES_LIMIT = 9_000_000;
+ static final int DATASTORE_BATCH_UPDATE_BYTES_LIMIT = 5_000_000;
/**
* Returns an empty {@link DatastoreV1.Read} builder. Configure the source {@code projectId},
@@ -1129,74 +1107,18 @@ public class DatastoreV1 {
}
}
- /** Determines batch sizes for commit RPCs. */
- @VisibleForTesting
- interface WriteBatcher {
- /** Call before using this WriteBatcher. */
- void start();
-
- /**
- * Reports the latency of a previous commit RPC, and the number of mutations that it contained.
- */
- void addRequestLatency(long timeSinceEpochMillis, long latencyMillis, int numMutations);
-
- /** Returns the number of entities to include in the next CommitRequest. */
- int nextBatchSize(long timeSinceEpochMillis);
- }
-
- /**
- * Determines batch sizes for commit RPCs based on past performance.
- *
- * <p>It aims for a target response time per RPC: it uses the response times for previous RPCs
- * and the number of entities contained in them, calculates a rolling average time-per-entity, and
- * chooses the number of entities for future writes to hit the target time.
- *
- * <p>This enables us to send large batches without sending over-large requests in the case of
- * expensive entity writes that may timeout before the server can apply them all.
- */
- @VisibleForTesting
- static class WriteBatcherImpl implements WriteBatcher, Serializable {
- /** Target time per RPC for writes. */
- static final int DATASTORE_BATCH_TARGET_LATENCY_MS = 5000;
-
- @Override
- public void start() {
- meanLatencyPerEntityMs = new MovingAverage(
- 120000 /* sample period 2 minutes */, 10000 /* sample interval 10s */,
- 1 /* numSignificantBuckets */, 1 /* numSignificantSamples */);
- }
-
- @Override
- public void addRequestLatency(long timeSinceEpochMillis, long latencyMillis, int numMutations) {
- meanLatencyPerEntityMs.add(timeSinceEpochMillis, latencyMillis / numMutations);
- }
-
- @Override
- public int nextBatchSize(long timeSinceEpochMillis) {
- if (!meanLatencyPerEntityMs.hasValue(timeSinceEpochMillis)) {
- return DATASTORE_BATCH_UPDATE_ENTITIES_START;
- }
- long recentMeanLatency = Math.max(meanLatencyPerEntityMs.get(timeSinceEpochMillis), 1);
- return (int) Math.max(DATASTORE_BATCH_UPDATE_ENTITIES_MIN,
- Math.min(DATASTORE_BATCH_UPDATE_ENTITIES_LIMIT,
- DATASTORE_BATCH_TARGET_LATENCY_MS / recentMeanLatency));
- }
-
- private transient MovingAverage meanLatencyPerEntityMs;
- }
-
/**
* {@link DoFn} that writes {@link Mutation}s to Cloud Datastore. Mutations are written in
- * batches; see {@link DatastoreV1.WriteBatcherImpl}.
+ * batches, where the maximum batch size is {@link DatastoreV1#DATASTORE_BATCH_UPDATE_LIMIT}.
*
* <p>See <a
* href="https://cloud.google.com/datastore/docs/concepts/entities">
* Datastore: Entities, Properties, and Keys</a> for information about entity keys and mutations.
*
* <p>Commits are non-transactional. If a commit fails because of a conflict over an entity
- * group, the commit will be retried (up to {@link DatastoreV1.DatastoreWriterFn#MAX_RETRIES}
+ * group, the commit will be retried (up to {@link DatastoreV1#DATASTORE_BATCH_UPDATE_LIMIT}
* times). This means that the mutation operation should be idempotent. Thus, the writer should
- * only be used for {@code upsert} and {@code delete} mutation operations, as these are the only
+ * only be used for {code upsert} and {@code delete} mutation operations, as these are the only
* two Cloud Datastore mutations that are idempotent.
*/
@VisibleForTesting
@@ -1210,14 +1132,6 @@ public class DatastoreV1 {
// Current batch of mutations to be written.
private final List<Mutation> mutations = new ArrayList<>();
private int mutationsSize = 0; // Accumulated size of protos in mutations.
- private WriteBatcher writeBatcher;
- private transient AdaptiveThrottler throttler;
- private final Counter throttledSeconds =
- Metrics.counter(DatastoreWriterFn.class, "cumulativeThrottlingSeconds");
- private final Counter rpcErrors =
- Metrics.counter(DatastoreWriterFn.class, "datastoreRpcErrors");
- private final Counter rpcSuccesses =
- Metrics.counter(DatastoreWriterFn.class, "datastoreRpcSuccesses");
private static final int MAX_RETRIES = 5;
private static final FluentBackoff BUNDLE_WRITE_BACKOFF =
@@ -1225,31 +1139,24 @@ public class DatastoreV1 {
.withMaxRetries(MAX_RETRIES).withInitialBackoff(Duration.standardSeconds(5));
DatastoreWriterFn(String projectId, @Nullable String localhost) {
- this(StaticValueProvider.of(projectId), localhost, new V1DatastoreFactory(),
- new WriteBatcherImpl());
+ this(StaticValueProvider.of(projectId), localhost, new V1DatastoreFactory());
}
DatastoreWriterFn(ValueProvider<String> projectId, @Nullable String localhost) {
- this(projectId, localhost, new V1DatastoreFactory(), new WriteBatcherImpl());
+ this(projectId, localhost, new V1DatastoreFactory());
}
@VisibleForTesting
DatastoreWriterFn(ValueProvider<String> projectId, @Nullable String localhost,
- V1DatastoreFactory datastoreFactory, WriteBatcher writeBatcher) {
+ V1DatastoreFactory datastoreFactory) {
this.projectId = checkNotNull(projectId, "projectId");
this.localhost = localhost;
this.datastoreFactory = datastoreFactory;
- this.writeBatcher = writeBatcher;
}
@StartBundle
public void startBundle(StartBundleContext c) {
datastore = datastoreFactory.getDatastore(c.getPipelineOptions(), projectId.get(), localhost);
- writeBatcher.start();
- if (throttler == null) {
- // Initialize throttler at first use, because it is not serializable.
- throttler = new AdaptiveThrottler(120000, 10000, 1.25);
- }
}
@ProcessElement
@@ -1262,7 +1169,7 @@ public class DatastoreV1 {
}
mutations.add(c.element());
mutationsSize += size;
- if (mutations.size() >= writeBatcher.nextBatchSize(System.currentTimeMillis())) {
+ if (mutations.size() >= DatastoreV1.DATASTORE_BATCH_UPDATE_LIMIT) {
flushBatch();
}
}
@@ -1292,42 +1199,18 @@ public class DatastoreV1 {
while (true) {
// Batch upsert entities.
- CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
- commitRequest.addAllMutations(mutations);
- commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
- long startTime = System.currentTimeMillis(), endTime;
-
- if (throttler.throttleRequest(startTime)) {
- LOG.info("Delaying request due to previous failures");
- throttledSeconds.inc(WriteBatcherImpl.DATASTORE_BATCH_TARGET_LATENCY_MS / 1000);
- sleeper.sleep(WriteBatcherImpl.DATASTORE_BATCH_TARGET_LATENCY_MS);
- continue;
- }
-
try {
+ CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
+ commitRequest.addAllMutations(mutations);
+ commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
datastore.commit(commitRequest.build());
- endTime = System.currentTimeMillis();
-
- writeBatcher.addRequestLatency(endTime, endTime - startTime, mutations.size());
- throttler.successfulRequest(startTime);
- rpcSuccesses.inc();
-
// Break if the commit threw no exception.
break;
} catch (DatastoreException exception) {
- if (exception.getCode() == Code.DEADLINE_EXCEEDED) {
- /* Most errors are not related to request size, and should not change our expectation of
- * the latency of successful requests. DEADLINE_EXCEEDED can be taken into
- * consideration, though. */
- endTime = System.currentTimeMillis();
- writeBatcher.addRequestLatency(endTime, endTime - startTime, mutations.size());
- }
// Only log the code and message for potentially-transient errors. The entire exception
// will be propagated upon the last retry.
- LOG.error("Error writing batch of {} mutations to Datastore ({}): {}", mutations.size(),
- exception.getCode(), exception.getMessage());
- rpcErrors.inc();
-
+ LOG.error("Error writing to the Datastore ({}): {}", exception.getCode(),
+ exception.getMessage());
if (!BackOffUtils.next(sleeper, backoff)) {
LOG.error("Aborting after {} retries.", MAX_RETRIES);
throw exception;
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/MovingAverage.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/MovingAverage.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/MovingAverage.java
deleted file mode 100644
index 0890e79..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/MovingAverage.java
+++ /dev/null
@@ -1,50 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.beam.sdk.io.gcp.datastore;
-
-import org.apache.beam.sdk.transforms.Sum;
-import org.apache.beam.sdk.util.MovingFunction;
-
-
-class MovingAverage {
- private final MovingFunction sum;
- private final MovingFunction count;
-
- public MovingAverage(long samplePeriodMs, long sampleUpdateMs,
- int numSignificantBuckets, int numSignificantSamples) {
- sum = new MovingFunction(samplePeriodMs, sampleUpdateMs,
- numSignificantBuckets, numSignificantSamples, Sum.ofLongs());
- count = new MovingFunction(samplePeriodMs, sampleUpdateMs,
- numSignificantBuckets, numSignificantSamples, Sum.ofLongs());
- }
-
- public void add(long nowMsSinceEpoch, long value) {
- sum.add(nowMsSinceEpoch, value);
- count.add(nowMsSinceEpoch, 1);
- }
-
- public long get(long nowMsSinceEpoch) {
- return sum.get(nowMsSinceEpoch) / count.get(nowMsSinceEpoch);
- }
-
- public boolean hasValue(long nowMsSinceEpoch) {
- return sum.isSignificant() && count.isSignificant()
- && count.get(nowMsSinceEpoch) > 0;
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/AbstractSpannerFn.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/AbstractSpannerFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/AbstractSpannerFn.java
deleted file mode 100644
index 00008f1..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/AbstractSpannerFn.java
+++ /dev/null
@@ -1,58 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner;
-
-import com.google.cloud.spanner.DatabaseClient;
-import com.google.cloud.spanner.DatabaseId;
-import com.google.cloud.spanner.Spanner;
-import com.google.cloud.spanner.SpannerOptions;
-import org.apache.beam.sdk.transforms.DoFn;
-
-/**
- * Abstract {@link DoFn} that manages {@link Spanner} lifecycle. Use {@link
- * AbstractSpannerFn#databaseClient} to access the Cloud Spanner database client.
- */
-abstract class AbstractSpannerFn<InputT, OutputT> extends DoFn<InputT, OutputT> {
- private transient Spanner spanner;
- private transient DatabaseClient databaseClient;
-
- abstract SpannerConfig getSpannerConfig();
-
- @Setup
- public void setup() throws Exception {
- SpannerConfig spannerConfig = getSpannerConfig();
- SpannerOptions options = spannerConfig.buildSpannerOptions();
- spanner = options.getService();
- databaseClient = spanner.getDatabaseClient(DatabaseId
- .of(options.getProjectId(), spannerConfig.getInstanceId().get(),
- spannerConfig.getDatabaseId().get()));
- }
-
- @Teardown
- public void teardown() throws Exception {
- if (spanner == null) {
- return;
- }
- spanner.close();
- spanner = null;
- }
-
- protected DatabaseClient databaseClient() {
- return databaseClient;
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/CreateTransactionFn.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/CreateTransactionFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/CreateTransactionFn.java
deleted file mode 100644
index da8e8b1..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/CreateTransactionFn.java
+++ /dev/null
@@ -1,51 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner;
-
-import com.google.cloud.spanner.ReadOnlyTransaction;
-import com.google.cloud.spanner.ResultSet;
-import com.google.cloud.spanner.Statement;
-
-/** Creates a batch transaction. */
-class CreateTransactionFn extends AbstractSpannerFn<Object, Transaction> {
-
- private final SpannerIO.CreateTransaction config;
-
- CreateTransactionFn(SpannerIO.CreateTransaction config) {
- this.config = config;
- }
-
- @ProcessElement
- public void processElement(ProcessContext c) throws Exception {
- try (ReadOnlyTransaction readOnlyTransaction =
- databaseClient().readOnlyTransaction(config.getTimestampBound())) {
- // Run a dummy sql statement to force the RPC and obtain the timestamp from the server.
- ResultSet resultSet = readOnlyTransaction.executeQuery(Statement.of("SELECT 1"));
- while (resultSet.next()) {
- // do nothing
- }
- Transaction tx = Transaction.create(readOnlyTransaction.getReadTimestamp());
- c.output(tx);
- }
- }
-
- @Override
- SpannerConfig getSpannerConfig() {
- return config.getSpannerConfig();
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationGroup.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationGroup.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationGroup.java
deleted file mode 100644
index 5b08da2..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationGroup.java
+++ /dev/null
@@ -1,67 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner;
-
-import com.google.cloud.spanner.Mutation;
-import com.google.common.collect.ImmutableList;
-
-import java.io.Serializable;
-import java.util.Arrays;
-import java.util.Iterator;
-import java.util.List;
-
-/**
- * A bundle of mutations that must be submitted atomically.
- *
- * <p>One of the mutations is chosen to be "primary", and can be used to determine partitions.
- */
-public final class MutationGroup implements Serializable, Iterable<Mutation> {
- private final ImmutableList<Mutation> mutations;
-
- /**
- * Creates a new group.
- *
- * @param primary a primary mutation.
- * @param other other mutations, usually interleaved in parent.
- * @return new mutation group.
- */
- public static MutationGroup create(Mutation primary, Mutation... other) {
- return create(primary, Arrays.asList(other));
- }
-
- public static MutationGroup create(Mutation primary, Iterable<Mutation> other) {
- return new MutationGroup(ImmutableList.<Mutation>builder().add(primary).addAll(other).build());
- }
-
- @Override
- public Iterator<Mutation> iterator() {
- return mutations.iterator();
- }
-
- private MutationGroup(ImmutableList<Mutation> mutations) {
- this.mutations = mutations;
- }
-
- public Mutation primary() {
- return mutations.get(0);
- }
-
- public List<Mutation> attached() {
- return mutations.subList(1, mutations.size());
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimator.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimator.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimator.java
index 2418816..61652e7 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimator.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/MutationSizeEstimator.java
@@ -44,15 +44,6 @@ class MutationSizeEstimator {
return result;
}
- /** Estimates a size of the mutation group in bytes. */
- public static long sizeOf(MutationGroup group) {
- long result = 0;
- for (Mutation m : group) {
- result += sizeOf(m);
- }
- return result;
- }
-
private static long estimatePrimitiveValue(Value v) {
switch (v.getType().getCode()) {
case BOOL:
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/NaiveSpannerReadFn.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/NaiveSpannerReadFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/NaiveSpannerReadFn.java
deleted file mode 100644
index d193b95..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/NaiveSpannerReadFn.java
+++ /dev/null
@@ -1,65 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner;
-
-import com.google.cloud.spanner.ReadOnlyTransaction;
-import com.google.cloud.spanner.ResultSet;
-import com.google.cloud.spanner.Struct;
-import com.google.cloud.spanner.TimestampBound;
-import com.google.common.annotations.VisibleForTesting;
-
-/** A simplest read function implementation. Parallelism support is coming. */
-@VisibleForTesting
-class NaiveSpannerReadFn extends AbstractSpannerFn<Object, Struct> {
- private final SpannerIO.Read config;
-
- NaiveSpannerReadFn(SpannerIO.Read config) {
- this.config = config;
- }
-
- SpannerConfig getSpannerConfig() {
- return config.getSpannerConfig();
- }
-
- @ProcessElement
- public void processElement(ProcessContext c) throws Exception {
- TimestampBound timestampBound = TimestampBound.strong();
- if (config.getTransaction() != null) {
- Transaction transaction = c.sideInput(config.getTransaction());
- timestampBound = TimestampBound.ofReadTimestamp(transaction.timestamp());
- }
- try (ReadOnlyTransaction readOnlyTransaction =
- databaseClient().readOnlyTransaction(timestampBound)) {
- ResultSet resultSet = execute(readOnlyTransaction);
- while (resultSet.next()) {
- c.output(resultSet.getCurrentRowAsStruct());
- }
- }
- }
-
- private ResultSet execute(ReadOnlyTransaction readOnlyTransaction) {
- if (config.getQuery() != null) {
- return readOnlyTransaction.executeQuery(config.getQuery());
- }
- if (config.getIndex() != null) {
- return readOnlyTransaction.readUsingIndex(
- config.getTable(), config.getIndex(), config.getKeySet(), config.getColumns());
- }
- return readOnlyTransaction.read(config.getTable(), config.getKeySet(), config.getColumns());
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java
deleted file mode 100644
index 02716fb..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerConfig.java
+++ /dev/null
@@ -1,137 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner;
-
-import static com.google.common.base.Preconditions.checkNotNull;
-
-import com.google.auto.value.AutoValue;
-import com.google.cloud.ServiceFactory;
-import com.google.cloud.spanner.Spanner;
-import com.google.cloud.spanner.SpannerOptions;
-import com.google.common.annotations.VisibleForTesting;
-import java.io.Serializable;
-import javax.annotation.Nullable;
-import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.options.ValueProvider;
-import org.apache.beam.sdk.transforms.display.DisplayData;
-
-/** Configuration for a Cloud Spanner client. */
-@AutoValue
-public abstract class SpannerConfig implements Serializable {
-
- @Nullable
- abstract ValueProvider<String> getProjectId();
-
- @Nullable
- abstract ValueProvider<String> getInstanceId();
-
- @Nullable
- abstract ValueProvider<String> getDatabaseId();
-
- @Nullable
- @VisibleForTesting
- abstract ServiceFactory<Spanner, SpannerOptions> getServiceFactory();
-
- abstract Builder toBuilder();
-
- SpannerOptions buildSpannerOptions() {
- SpannerOptions.Builder builder = SpannerOptions.newBuilder();
- if (getProjectId() != null) {
- builder.setProjectId(getProjectId().get());
- }
- if (getServiceFactory() != null) {
- builder.setServiceFactory(getServiceFactory());
- }
- return builder.build();
- }
-
- public static SpannerConfig create() {
- return builder().build();
- }
-
- static Builder builder() {
- return new AutoValue_SpannerConfig.Builder();
- }
-
- public void validate(PipelineOptions options) {
- checkNotNull(
- getInstanceId(),
- "SpannerIO.read() requires instance id to be set with withInstanceId method");
- checkNotNull(
- getDatabaseId(),
- "SpannerIO.read() requires database id to be set with withDatabaseId method");
- }
-
- public void populateDisplayData(DisplayData.Builder builder) {
- builder
- .addIfNotNull(DisplayData.item("projectId", getProjectId()).withLabel("Output Project"))
- .addIfNotNull(DisplayData.item("instanceId", getInstanceId()).withLabel("Output Instance"))
- .addIfNotNull(DisplayData.item("databaseId", getDatabaseId()).withLabel("Output Database"));
-
- if (getServiceFactory() != null) {
- builder.addIfNotNull(
- DisplayData.item("serviceFactory", getServiceFactory().getClass().getName())
- .withLabel("Service Factory"));
- }
- }
-
- /** Builder for {@link SpannerConfig}. */
- @AutoValue.Builder
- public abstract static class Builder {
-
- abstract Builder setProjectId(ValueProvider<String> projectId);
-
- abstract Builder setInstanceId(ValueProvider<String> instanceId);
-
- abstract Builder setDatabaseId(ValueProvider<String> databaseId);
-
- abstract Builder setServiceFactory(ServiceFactory<Spanner, SpannerOptions> serviceFactory);
-
- public abstract SpannerConfig build();
- }
-
- public SpannerConfig withProjectId(ValueProvider<String> projectId) {
- return toBuilder().setProjectId(projectId).build();
- }
-
- public SpannerConfig withProjectId(String projectId) {
- return withProjectId(ValueProvider.StaticValueProvider.of(projectId));
- }
-
- public SpannerConfig withInstanceId(ValueProvider<String> instanceId) {
- return toBuilder().setInstanceId(instanceId).build();
- }
-
- public SpannerConfig withInstanceId(String instanceId) {
- return withInstanceId(ValueProvider.StaticValueProvider.of(instanceId));
- }
-
- public SpannerConfig withDatabaseId(ValueProvider<String> databaseId) {
- return toBuilder().setDatabaseId(databaseId).build();
- }
-
- public SpannerConfig withDatabaseId(String databaseId) {
- return withDatabaseId(ValueProvider.StaticValueProvider.of(databaseId));
- }
-
- @VisibleForTesting
- SpannerConfig withServiceFactory(ServiceFactory<Spanner, SpannerOptions> serviceFactory) {
- return toBuilder().setServiceFactory(serviceFactory).build();
- }
-
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
index a247d4c..5058d13 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
@@ -17,39 +17,38 @@
*/
package org.apache.beam.sdk.io.gcp.spanner;
-import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import com.google.auto.value.AutoValue;
import com.google.cloud.ServiceFactory;
-import com.google.cloud.Timestamp;
-import com.google.cloud.spanner.KeySet;
+import com.google.cloud.ServiceOptions;
+import com.google.cloud.spanner.AbortedException;
+import com.google.cloud.spanner.DatabaseClient;
+import com.google.cloud.spanner.DatabaseId;
import com.google.cloud.spanner.Mutation;
import com.google.cloud.spanner.Spanner;
import com.google.cloud.spanner.SpannerOptions;
-import com.google.cloud.spanner.Statement;
-import com.google.cloud.spanner.Struct;
-import com.google.cloud.spanner.TimestampBound;
import com.google.common.annotations.VisibleForTesting;
-
-import java.util.Arrays;
-import java.util.Collections;
+import java.io.IOException;
+import java.util.ArrayList;
import java.util.List;
import javax.annotation.Nullable;
-
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.options.PipelineOptions;
-import org.apache.beam.sdk.options.ValueProvider;
-import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
-import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.display.DisplayData;
-import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.transforms.display.DisplayData.Builder;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.BackOffUtils;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.sdk.util.Sleeper;
import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PDone;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
* Experimental {@link PTransform Transforms} for reading from and writing to <a
@@ -57,69 +56,7 @@ import org.apache.beam.sdk.values.PDone;
*
* <h3>Reading from Cloud Spanner</h3>
*
- * <p>To read from Cloud Spanner, apply {@link SpannerIO.Read} transformation. It will return a
- * {@link PCollection} of {@link Struct Structs}, where each element represents
- * an individual row returned from the read operation. Both Query and Read APIs are supported.
- * See more information about <a href="https://cloud.google.com/spanner/docs/reads">reading from
- * Cloud Spanner</a>
- *
- * <p>To execute a <strong>query</strong>, specify a {@link SpannerIO.Read#withQuery(Statement)} or
- * {@link SpannerIO.Read#withQuery(String)} during the construction of the transform.
- *
- * <pre>{@code
- * PCollection<Struct> rows = p.apply(
- * SpannerIO.read()
- * .withInstanceId(instanceId)
- * .withDatabaseId(dbId)
- * .withQuery("SELECT id, name, email FROM users"));
- * }</pre>
- *
- * <p>To use the Read API, specify a {@link SpannerIO.Read#withTable(String) table name} and
- * a {@link SpannerIO.Read#withColumns(List) list of columns}.
- *
- * <pre>{@code
- * PCollection<Struct> rows = p.apply(
- * SpannerIO.read()
- * .withInstanceId(instanceId)
- * .withDatabaseId(dbId)
- * .withTable("users")
- * .withColumns("id", "name", "email"));
- * }</pre>
- *
- * <p>To optimally read using index, specify the index name using {@link SpannerIO.Read#withIndex}.
- *
- * <p>The transform is guaranteed to be executed on a consistent snapshot of data, utilizing the
- * power of read only transactions. Staleness of data can be controlled using
- * {@link SpannerIO.Read#withTimestampBound} or {@link SpannerIO.Read#withTimestamp(Timestamp)}
- * methods. <a href="https://cloud.google.com/spanner/docs/transactions">Read more</a> about
- * transactions in Cloud Spanner.
- *
- * <p>It is possible to read several {@link PCollection PCollections} within a single transaction.
- * Apply {@link SpannerIO#createTransaction()} transform, that lazily creates a transaction. The
- * result of this transformation can be passed to read operation using
- * {@link SpannerIO.Read#withTransaction(PCollectionView)}.
- *
- * <pre>{@code
- * SpannerConfig spannerConfig = ...
- *
- * PCollectionView<Transaction> tx =
- * p.apply(
- * SpannerIO.createTransaction()
- * .withSpannerConfig(spannerConfig)
- * .withTimestampBound(TimestampBound.strong()));
- *
- * PCollection<Struct> users = p.apply(
- * SpannerIO.read()
- * .withSpannerConfig(spannerConfig)
- * .withQuery("SELECT name, email FROM users")
- * .withTransaction(tx));
- *
- * PCollection<Struct> tweets = p.apply(
- * SpannerIO.read()
- * .withSpannerConfig(spannerConfig)
- * .withQuery("SELECT user, tweet, date FROM tweets")
- * .withTransaction(tx));
- * }</pre>
+ * <p>This functionality is not yet implemented.
*
* <h3>Writing to Cloud Spanner</h3>
*
@@ -151,11 +88,6 @@ import org.apache.beam.sdk.values.PDone;
* <li>If the pipeline was unexpectedly stopped, mutations that were already applied will not get
* rolled back.
* </ul>
- *
- * <p>Use {@link MutationGroup} to ensure that a small set mutations is bundled together. It is
- * guaranteed that mutations in a group are submitted in the same transaction. Build
- * {@link SpannerIO.Write} transform, and call {@link Write#grouped()} method. It will return a
- * transformation that can be applied to a PCollection of MutationGroup.
*/
@Experimental(Experimental.Kind.SOURCE_SINK)
public class SpannerIO {
@@ -163,33 +95,6 @@ public class SpannerIO {
private static final long DEFAULT_BATCH_SIZE_BYTES = 1024 * 1024; // 1 MB
/**
- * Creates an uninitialized instance of {@link Read}. Before use, the {@link Read} must be
- * configured with a {@link Read#withInstanceId} and {@link Read#withDatabaseId} that identify the
- * Cloud Spanner database.
- */
- @Experimental(Experimental.Kind.SOURCE_SINK)
- public static Read read() {
- return new AutoValue_SpannerIO_Read.Builder()
- .setSpannerConfig(SpannerConfig.create())
- .setTimestampBound(TimestampBound.strong())
- .setKeySet(KeySet.all())
- .build();
- }
-
- /**
- * Returns a transform that creates a batch transaction. By default,
- * {@link TimestampBound#strong()} transaction is created, to override this use
- * {@link CreateTransaction#withTimestampBound(TimestampBound)}.
- */
- @Experimental
- public static CreateTransaction createTransaction() {
- return new AutoValue_SpannerIO_CreateTransaction.Builder()
- .setSpannerConfig(SpannerConfig.create())
- .setTimestampBound(TimestampBound.strong())
- .build();
- }
-
- /**
* Creates an uninitialized instance of {@link Write}. Before use, the {@link Write} must be
* configured with a {@link Write#withInstanceId} and {@link Write#withDatabaseId} that identify
* the Cloud Spanner database being written.
@@ -197,408 +102,247 @@ public class SpannerIO {
@Experimental
public static Write write() {
return new AutoValue_SpannerIO_Write.Builder()
- .setSpannerConfig(SpannerConfig.create())
.setBatchSizeBytes(DEFAULT_BATCH_SIZE_BYTES)
.build();
}
/**
- * A {@link PTransform} that reads data from Google Cloud Spanner.
+ * A {@link PTransform} that writes {@link Mutation} objects to Google Cloud Spanner.
*
* @see SpannerIO
*/
@Experimental(Experimental.Kind.SOURCE_SINK)
@AutoValue
- public abstract static class Read extends PTransform<PBegin, PCollection<Struct>> {
-
- abstract SpannerConfig getSpannerConfig();
-
- @Nullable
- abstract TimestampBound getTimestampBound();
-
- @Nullable
- abstract Statement getQuery();
+ public abstract static class Write extends PTransform<PCollection<Mutation>, PDone> {
@Nullable
- abstract String getTable();
+ abstract String getProjectId();
@Nullable
- abstract String getIndex();
+ abstract String getInstanceId();
@Nullable
- abstract List<String> getColumns();
+ abstract String getDatabaseId();
- @Nullable
- abstract KeySet getKeySet();
+ abstract long getBatchSizeBytes();
@Nullable
- abstract PCollectionView<Transaction> getTransaction();
+ @VisibleForTesting
+ abstract ServiceFactory<Spanner, SpannerOptions> getServiceFactory();
abstract Builder toBuilder();
@AutoValue.Builder
abstract static class Builder {
- abstract Builder setSpannerConfig(SpannerConfig spannerConfig);
-
- abstract Builder setTimestampBound(TimestampBound timestampBound);
-
- abstract Builder setQuery(Statement statement);
-
- abstract Builder setTable(String table);
-
- abstract Builder setIndex(String index);
+ abstract Builder setProjectId(String projectId);
- abstract Builder setColumns(List<String> columns);
+ abstract Builder setInstanceId(String instanceId);
- abstract Builder setKeySet(KeySet keySet);
+ abstract Builder setDatabaseId(String databaseId);
- abstract Builder setTransaction(PCollectionView<Transaction> transaction);
-
- abstract Read build();
- }
-
- /** Specifies the Cloud Spanner configuration. */
- public Read withSpannerConfig(SpannerConfig spannerConfig) {
- return toBuilder().setSpannerConfig(spannerConfig).build();
- }
+ abstract Builder setBatchSizeBytes(long batchSizeBytes);
- /** Specifies the Cloud Spanner project. */
- public Read withProjectId(String projectId) {
- return withProjectId(ValueProvider.StaticValueProvider.of(projectId));
- }
+ @VisibleForTesting
+ abstract Builder setServiceFactory(ServiceFactory<Spanner, SpannerOptions> serviceFactory);
- /** Specifies the Cloud Spanner project. */
- public Read withProjectId(ValueProvider<String> projectId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withProjectId(projectId));
+ abstract Write build();
}
- /** Specifies the Cloud Spanner instance. */
- public Read withInstanceId(String instanceId) {
- return withInstanceId(ValueProvider.StaticValueProvider.of(instanceId));
+ /**
+ * Returns a new {@link SpannerIO.Write} that will write to the specified Cloud Spanner project.
+ *
+ * <p>Does not modify this object.
+ */
+ public Write withProjectId(String projectId) {
+ return toBuilder().setProjectId(projectId).build();
}
- /** Specifies the Cloud Spanner instance. */
- public Read withInstanceId(ValueProvider<String> instanceId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withInstanceId(instanceId));
+ /**
+ * Returns a new {@link SpannerIO.Write} that will write to the specified Cloud Spanner
+ * instance.
+ *
+ * <p>Does not modify this object.
+ */
+ public Write withInstanceId(String instanceId) {
+ return toBuilder().setInstanceId(instanceId).build();
}
- /** Specifies the Cloud Spanner database. */
- public Read withDatabaseId(String databaseId) {
- return withDatabaseId(ValueProvider.StaticValueProvider.of(databaseId));
+ /**
+ * Returns a new {@link SpannerIO.Write} with a new batch size limit.
+ *
+ * <p>Does not modify this object.
+ */
+ public Write withBatchSizeBytes(long batchSizeBytes) {
+ return toBuilder().setBatchSizeBytes(batchSizeBytes).build();
}
- /** Specifies the Cloud Spanner database. */
- public Read withDatabaseId(ValueProvider<String> databaseId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withDatabaseId(databaseId));
+ /**
+ * Returns a new {@link SpannerIO.Write} that will write to the specified Cloud Spanner
+ * database.
+ *
+ * <p>Does not modify this object.
+ */
+ public Write withDatabaseId(String databaseId) {
+ return toBuilder().setDatabaseId(databaseId).build();
}
@VisibleForTesting
- Read withServiceFactory(ServiceFactory<Spanner, SpannerOptions> serviceFactory) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withServiceFactory(serviceFactory));
- }
-
- public Read withTransaction(PCollectionView<Transaction> transaction) {
- return toBuilder().setTransaction(transaction).build();
- }
-
- public Read withTimestamp(Timestamp timestamp) {
- return withTimestampBound(TimestampBound.ofReadTimestamp(timestamp));
- }
-
- public Read withTimestampBound(TimestampBound timestampBound) {
- return toBuilder().setTimestampBound(timestampBound).build();
- }
-
- public Read withTable(String table) {
- return toBuilder().setTable(table).build();
- }
-
- public Read withColumns(String... columns) {
- return withColumns(Arrays.asList(columns));
- }
-
- public Read withColumns(List<String> columns) {
- return toBuilder().setColumns(columns).build();
- }
-
- public Read withQuery(Statement statement) {
- return toBuilder().setQuery(statement).build();
- }
-
- public Read withQuery(String sql) {
- return withQuery(Statement.of(sql));
- }
-
- public Read withKeySet(KeySet keySet) {
- return toBuilder().setKeySet(keySet).build();
- }
-
- public Read withIndex(String index) {
- return toBuilder().setIndex(index).build();
+ Write withServiceFactory(ServiceFactory<Spanner, SpannerOptions> serviceFactory) {
+ return toBuilder().setServiceFactory(serviceFactory).build();
}
-
@Override
public void validate(PipelineOptions options) {
- getSpannerConfig().validate(options);
checkNotNull(
- getTimestampBound(),
- "SpannerIO.read() runs in a read only transaction and requires timestamp to be set "
- + "with withTimestampBound or withTimestamp method");
-
- if (getQuery() != null) {
- // TODO: validate query?
- } else if (getTable() != null) {
- // Assume read
- checkNotNull(
- getColumns(),
- "For a read operation SpannerIO.read() requires a list of "
- + "columns to set with withColumns method");
- checkArgument(
- !getColumns().isEmpty(),
- "For a read operation SpannerIO.read() requires a"
- + " list of columns to set with withColumns method");
- } else {
- throw new IllegalArgumentException(
- "SpannerIO.read() requires configuring query or read operation.");
- }
+ getInstanceId(),
+ "SpannerIO.write() requires instance id to be set with withInstanceId method");
+ checkNotNull(
+ getDatabaseId(),
+ "SpannerIO.write() requires database id to be set with withDatabaseId method");
}
@Override
- public PCollection<Struct> expand(PBegin input) {
- Read config = this;
- List<PCollectionView<Transaction>> sideInputs = Collections.emptyList();
- if (getTimestampBound() != null) {
- PCollectionView<Transaction> transaction =
- input.apply(createTransaction().withSpannerConfig(getSpannerConfig()));
- config = config.withTransaction(transaction);
- sideInputs = Collections.singletonList(transaction);
- }
- return input
- .apply(Create.of(1))
- .apply(
- "Execute query", ParDo.of(new NaiveSpannerReadFn(config)).withSideInputs(sideInputs));
- }
- }
-
- /**
- * A {@link PTransform} that create a transaction.
- *
- * @see SpannerIO
- */
- @Experimental(Experimental.Kind.SOURCE_SINK)
- @AutoValue
- public abstract static class CreateTransaction
- extends PTransform<PBegin, PCollectionView<Transaction>> {
-
- abstract SpannerConfig getSpannerConfig();
-
- @Nullable
- abstract TimestampBound getTimestampBound();
-
- abstract Builder toBuilder();
-
- @Override
- public PCollectionView<Transaction> expand(PBegin input) {
- return input.apply(Create.of(1))
- .apply("Create transaction", ParDo.of(new CreateTransactionFn(this)))
- .apply("As PCollectionView", View.<Transaction>asSingleton());
- }
-
- /** Specifies the Cloud Spanner configuration. */
- public CreateTransaction withSpannerConfig(SpannerConfig spannerConfig) {
- return toBuilder().setSpannerConfig(spannerConfig).build();
- }
-
- /** Specifies the Cloud Spanner project. */
- public CreateTransaction withProjectId(String projectId) {
- return withProjectId(ValueProvider.StaticValueProvider.of(projectId));
- }
-
- /** Specifies the Cloud Spanner project. */
- public CreateTransaction withProjectId(ValueProvider<String> projectId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withProjectId(projectId));
- }
-
- /** Specifies the Cloud Spanner instance. */
- public CreateTransaction withInstanceId(String instanceId) {
- return withInstanceId(ValueProvider.StaticValueProvider.of(instanceId));
- }
-
- /** Specifies the Cloud Spanner instance. */
- public CreateTransaction withInstanceId(ValueProvider<String> instanceId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withInstanceId(instanceId));
- }
-
- /** Specifies the Cloud Spanner database. */
- public CreateTransaction withDatabaseId(String databaseId) {
- return withDatabaseId(ValueProvider.StaticValueProvider.of(databaseId));
- }
-
- /** Specifies the Cloud Spanner database. */
- public CreateTransaction withDatabaseId(ValueProvider<String> databaseId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withDatabaseId(databaseId));
- }
-
- @VisibleForTesting
- CreateTransaction withServiceFactory(
- ServiceFactory<Spanner, SpannerOptions> serviceFactory) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withServiceFactory(serviceFactory));
- }
-
- public CreateTransaction withTimestampBound(TimestampBound timestampBound) {
- return toBuilder().setTimestampBound(timestampBound).build();
+ public PDone expand(PCollection<Mutation> input) {
+ input.apply("Write mutations to Cloud Spanner", ParDo.of(new SpannerWriteFn(this)));
+ return PDone.in(input.getPipeline());
}
@Override
- public void validate(PipelineOptions options) {
- getSpannerConfig().validate(options);
- }
-
- /** A builder for {@link CreateTransaction}. */
- @AutoValue.Builder public abstract static class Builder {
-
- public abstract Builder setSpannerConfig(SpannerConfig spannerConfig);
-
- public abstract Builder setTimestampBound(TimestampBound newTimestampBound);
-
- public abstract CreateTransaction build();
+ public void populateDisplayData(DisplayData.Builder builder) {
+ super.populateDisplayData(builder);
+ builder
+ .addIfNotNull(DisplayData.item("projectId", getProjectId()).withLabel("Output Project"))
+ .addIfNotNull(
+ DisplayData.item("instanceId", getInstanceId()).withLabel("Output Instance"))
+ .addIfNotNull(
+ DisplayData.item("databaseId", getDatabaseId()).withLabel("Output Database"))
+ .add(DisplayData.item("batchSizeBytes", getBatchSizeBytes())
+ .withLabel("Batch Size in Bytes"));
+ if (getServiceFactory() != null) {
+ builder.addIfNotNull(
+ DisplayData.item("serviceFactory", getServiceFactory().getClass().getName())
+ .withLabel("Service Factory"));
+ }
}
}
+ /** Batches together and writes mutations to Google Cloud Spanner. */
+ @VisibleForTesting
+ static class SpannerWriteFn extends DoFn<Mutation, Void> {
+ private static final Logger LOG = LoggerFactory.getLogger(SpannerWriteFn.class);
+ private final Write spec;
+ private transient Spanner spanner;
+ private transient DatabaseClient dbClient;
+ // Current batch of mutations to be written.
+ private List<Mutation> mutations;
+ private long batchSizeBytes = 0;
+
+ private static final int MAX_RETRIES = 5;
+ private static final FluentBackoff BUNDLE_WRITE_BACKOFF =
+ FluentBackoff.DEFAULT
+ .withMaxRetries(MAX_RETRIES)
+ .withInitialBackoff(Duration.standardSeconds(5));
- /**
- * A {@link PTransform} that writes {@link Mutation} objects to Google Cloud Spanner.
- *
- * @see SpannerIO
- */
- @Experimental(Experimental.Kind.SOURCE_SINK)
- @AutoValue
- public abstract static class Write extends PTransform<PCollection<Mutation>, PDone> {
-
- abstract SpannerConfig getSpannerConfig();
-
- abstract long getBatchSizeBytes();
-
- abstract Builder toBuilder();
-
- @AutoValue.Builder
- abstract static class Builder {
-
- abstract Builder setSpannerConfig(SpannerConfig spannerConfig);
-
- abstract Builder setBatchSizeBytes(long batchSizeBytes);
-
- abstract Write build();
- }
-
- /** Specifies the Cloud Spanner configuration. */
- public Write withSpannerConfig(SpannerConfig spannerConfig) {
- return toBuilder().setSpannerConfig(spannerConfig).build();
- }
-
- /** Specifies the Cloud Spanner project. */
- public Write withProjectId(String projectId) {
- return withProjectId(ValueProvider.StaticValueProvider.of(projectId));
+ @VisibleForTesting
+ SpannerWriteFn(Write spec) {
+ this.spec = spec;
}
- /** Specifies the Cloud Spanner project. */
- public Write withProjectId(ValueProvider<String> projectId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withProjectId(projectId));
+ @Setup
+ public void setup() throws Exception {
+ SpannerOptions spannerOptions = getSpannerOptions();
+ spanner = spannerOptions.getService();
+ dbClient = spanner.getDatabaseClient(
+ DatabaseId.of(projectId(), spec.getInstanceId(), spec.getDatabaseId()));
+ mutations = new ArrayList<>();
+ batchSizeBytes = 0;
}
- /** Specifies the Cloud Spanner instance. */
- public Write withInstanceId(String instanceId) {
- return withInstanceId(ValueProvider.StaticValueProvider.of(instanceId));
+ @ProcessElement
+ public void processElement(ProcessContext c) throws Exception {
+ Mutation m = c.element();
+ mutations.add(m);
+ batchSizeBytes += MutationSizeEstimator.sizeOf(m);
+ if (batchSizeBytes >= spec.getBatchSizeBytes()) {
+ flushBatch();
+ }
}
- /** Specifies the Cloud Spanner instance. */
- public Write withInstanceId(ValueProvider<String> instanceId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withInstanceId(instanceId));
+ private String projectId() {
+ return spec.getProjectId() == null
+ ? ServiceOptions.getDefaultProjectId()
+ : spec.getProjectId();
}
- /** Specifies the Cloud Spanner database. */
- public Write withDatabaseId(String databaseId) {
- return withDatabaseId(ValueProvider.StaticValueProvider.of(databaseId));
+ @FinishBundle
+ public void finishBundle() throws Exception {
+ if (!mutations.isEmpty()) {
+ flushBatch();
+ }
}
- /** Specifies the Cloud Spanner database. */
- public Write withDatabaseId(ValueProvider<String> databaseId) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withDatabaseId(databaseId));
+ @Teardown
+ public void teardown() throws Exception {
+ if (spanner == null) {
+ return;
+ }
+ spanner.closeAsync().get();
+ spanner = null;
}
- @VisibleForTesting
- Write withServiceFactory(ServiceFactory<Spanner, SpannerOptions> serviceFactory) {
- SpannerConfig config = getSpannerConfig();
- return withSpannerConfig(config.withServiceFactory(serviceFactory));
+ private SpannerOptions getSpannerOptions() {
+ SpannerOptions.Builder spannerOptionsBuider = SpannerOptions.newBuilder();
+ if (spec.getServiceFactory() != null) {
+ spannerOptionsBuider.setServiceFactory(spec.getServiceFactory());
+ }
+ if (spec.getProjectId() != null) {
+ spannerOptionsBuider.setProjectId(spec.getProjectId());
+ }
+ return spannerOptionsBuider.build();
}
/**
- * Same transform but can be applied to {@link PCollection} of {@link MutationGroup}.
+ * Writes a batch of mutations to Cloud Spanner.
+ *
+ * <p>If a commit fails, it will be retried up to {@link #MAX_RETRIES} times. If the retry limit
+ * is exceeded, the last exception from Cloud Spanner will be thrown.
+ *
+ * @throws AbortedException if the commit fails or IOException or InterruptedException if
+ * backing off between retries fails.
*/
- public WriteGrouped grouped() {
- return new WriteGrouped(this);
- }
-
- /** Specifies the batch size limit. */
- public Write withBatchSizeBytes(long batchSizeBytes) {
- return toBuilder().setBatchSizeBytes(batchSizeBytes).build();
- }
-
- @Override
- public void validate(PipelineOptions options) {
- getSpannerConfig().validate(options);
- }
-
- @Override
- public PDone expand(PCollection<Mutation> input) {
- input
- .apply("To mutation group", ParDo.of(new ToMutationGroupFn()))
- .apply("Write mutations to Cloud Spanner", ParDo.of(new SpannerWriteGroupFn(this)));
- return PDone.in(input.getPipeline());
+ private void flushBatch() throws AbortedException, IOException, InterruptedException {
+ LOG.debug("Writing batch of {} mutations", mutations.size());
+ Sleeper sleeper = Sleeper.DEFAULT;
+ BackOff backoff = BUNDLE_WRITE_BACKOFF.backoff();
+
+ while (true) {
+ // Batch upsert rows.
+ try {
+ dbClient.writeAtLeastOnce(mutations);
+
+ // Break if the commit threw no exception.
+ break;
+ } catch (AbortedException exception) {
+ // Only log the code and message for potentially-transient errors. The entire exception
+ // will be propagated upon the last retry.
+ LOG.error(
+ "Error writing to Spanner ({}): {}", exception.getCode(), exception.getMessage());
+ if (!BackOffUtils.next(sleeper, backoff)) {
+ LOG.error("Aborting after {} retries.", MAX_RETRIES);
+ throw exception;
+ }
+ }
+ }
+ LOG.debug("Successfully wrote {} mutations", mutations.size());
+ mutations = new ArrayList<>();
+ batchSizeBytes = 0;
}
-
@Override
- public void populateDisplayData(DisplayData.Builder builder) {
+ public void populateDisplayData(Builder builder) {
super.populateDisplayData(builder);
- getSpannerConfig().populateDisplayData(builder);
- builder.add(
- DisplayData.item("batchSizeBytes", getBatchSizeBytes()).withLabel("Batch Size in Bytes"));
- }
- }
-
- /** Same as {@link Write} but supports grouped mutations. */
- public static class WriteGrouped extends PTransform<PCollection<MutationGroup>, PDone> {
- private final Write spec;
-
- public WriteGrouped(Write spec) {
- this.spec = spec;
- }
-
- @Override public PDone expand(PCollection<MutationGroup> input) {
- input.apply("Write mutations to Cloud Spanner", ParDo.of(new SpannerWriteGroupFn(spec)));
- return PDone.in(input.getPipeline());
- }
- }
-
- private static class ToMutationGroupFn extends DoFn<Mutation, MutationGroup> {
- @ProcessElement
- public void processElement(ProcessContext c) throws Exception {
- Mutation value = c.element();
- c.output(MutationGroup.create(value));
+ spec.populateDisplayData(builder);
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteGroupFn.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteGroupFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteGroupFn.java
deleted file mode 100644
index 34a11da..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerWriteGroupFn.java
+++ /dev/null
@@ -1,125 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner;
-
-import com.google.cloud.spanner.AbortedException;
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.Iterables;
-
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
-
-import org.apache.beam.sdk.transforms.display.DisplayData;
-import org.apache.beam.sdk.util.BackOff;
-import org.apache.beam.sdk.util.BackOffUtils;
-import org.apache.beam.sdk.util.FluentBackoff;
-import org.apache.beam.sdk.util.Sleeper;
-import org.joda.time.Duration;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-/** Batches together and writes mutations to Google Cloud Spanner. */
-@VisibleForTesting class SpannerWriteGroupFn extends AbstractSpannerFn<MutationGroup, Void> {
- private static final Logger LOG = LoggerFactory.getLogger(SpannerWriteGroupFn.class);
- private final SpannerIO.Write spec;
- // Current batch of mutations to be written.
- private List<MutationGroup> mutations;
- private long batchSizeBytes = 0;
-
- private static final int MAX_RETRIES = 5;
- private static final FluentBackoff BUNDLE_WRITE_BACKOFF =
- FluentBackoff.DEFAULT
- .withMaxRetries(MAX_RETRIES)
- .withInitialBackoff(Duration.standardSeconds(5));
-
- @VisibleForTesting SpannerWriteGroupFn(SpannerIO.Write spec) {
- this.spec = spec;
- }
-
- @Override SpannerConfig getSpannerConfig() {
- return spec.getSpannerConfig();
- }
-
- @Setup
- public void setup() throws Exception {
- super.setup();
- mutations = new ArrayList<>();
- batchSizeBytes = 0;
- }
-
- @ProcessElement
- public void processElement(ProcessContext c) throws Exception {
- MutationGroup m = c.element();
- mutations.add(m);
- batchSizeBytes += MutationSizeEstimator.sizeOf(m);
- if (batchSizeBytes >= spec.getBatchSizeBytes()) {
- flushBatch();
- }
- }
-
- @FinishBundle
- public void finishBundle() throws Exception {
- if (!mutations.isEmpty()) {
- flushBatch();
- }
- }
-
- /**
- * Writes a batch of mutations to Cloud Spanner.
- *
- * <p>If a commit fails, it will be retried up to {@link #MAX_RETRIES} times. If the retry limit
- * is exceeded, the last exception from Cloud Spanner will be thrown.
- *
- * @throws AbortedException if the commit fails or IOException or InterruptedException if
- * backing off between retries fails.
- */
- private void flushBatch() throws AbortedException, IOException, InterruptedException {
- LOG.debug("Writing batch of {} mutations", mutations.size());
- Sleeper sleeper = Sleeper.DEFAULT;
- BackOff backoff = BUNDLE_WRITE_BACKOFF.backoff();
-
- while (true) {
- // Batch upsert rows.
- try {
- databaseClient().writeAtLeastOnce(Iterables.concat(mutations));
-
- // Break if the commit threw no exception.
- break;
- } catch (AbortedException exception) {
- // Only log the code and message for potentially-transient errors. The entire exception
- // will be propagated upon the last retry.
- LOG.error(
- "Error writing to Spanner ({}): {}", exception.getCode(), exception.getMessage());
- if (!BackOffUtils.next(sleeper, backoff)) {
- LOG.error("Aborting after {} retries.", MAX_RETRIES);
- throw exception;
- }
- }
- }
- LOG.debug("Successfully wrote {} mutations", mutations.size());
- mutations = new ArrayList<>();
- batchSizeBytes = 0;
- }
-
- @Override
- public void populateDisplayData(DisplayData.Builder builder) {
- super.populateDisplayData(builder);
- spec.populateDisplayData(builder);
- }
-}
http://git-wip-us.apache.org/repos/asf/beam/blob/c1b2b96a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/Transaction.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/Transaction.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/Transaction.java
deleted file mode 100644
index 22af3b8..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/Transaction.java
+++ /dev/null
@@ -1,33 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.spanner;
-
-import com.google.auto.value.AutoValue;
-import com.google.cloud.Timestamp;
-import java.io.Serializable;
-
-/** A transaction object. */
-@AutoValue
-public abstract class Transaction implements Serializable {
-
- abstract Timestamp timestamp();
-
- public static Transaction create(Timestamp timestamp) {
- return new AutoValue_Transaction(timestamp);
- }
-}