Repository: GoogleCloudPlatform/cloud-bigtable-client
Pull request #2073 (2019-02-06 19:26:40)

Reducing the responsiblity of BatchExecutor
Using the more of the BufferedMutator infrastructure for Table batch operations like `Table.put(List<Put> puts)`

This should help with the transition to google-cloud-bigtable


Patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/core/IBigtableDataClient.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/core/IBigtableDataClient.java
index 33b48f833..6d35cea88 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/core/IBigtableDataClient.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/core/IBigtableDataClient.java
@@ -88,8 +88,6 @@ public interface IBigtableDataClient {
    *
    * @param conditionalRowMutation a {@link ConditionalRowMutation} model object.
    * @return returns true if predicate returns any result.
-   * @throws ExecutionException if any.
-   * @throws InterruptedException if any.
    */
   Boolean checkAndMutateRow(ConditionalRowMutation conditionalRowMutation);
 
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/BigtableTableName.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/BigtableTableName.java
index a2bcbc522..5f018dde4 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/BigtableTableName.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/BigtableTableName.java
@@ -15,7 +15,6 @@
  */
 package com.google.cloud.bigtable.grpc;
 
-import com.google.cloud.bigtable.data.v2.models.InstanceName;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/AsyncExecutor.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/AsyncExecutor.java
index f427804eb..9908d51f2 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/AsyncExecutor.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/AsyncExecutor.java
@@ -16,25 +16,14 @@
 package com.google.cloud.bigtable.grpc.async;
 
 import java.io.IOException;
-import java.util.List;
 
-import com.google.bigtable.v2.CheckAndMutateRowRequest;
-import com.google.bigtable.v2.CheckAndMutateRowResponse;
-import com.google.bigtable.v2.MutateRowRequest;
-import com.google.bigtable.v2.MutateRowResponse;
-import com.google.bigtable.v2.MutateRowsRequest;
-import com.google.bigtable.v2.MutateRowsResponse;
 import com.google.bigtable.v2.ReadModifyWriteRowRequest;
 import com.google.bigtable.v2.ReadModifyWriteRowResponse;
-import com.google.bigtable.v2.ReadRowsRequest;
-import com.google.bigtable.v2.Row;
 import com.google.cloud.bigtable.config.Logger;
 import com.google.cloud.bigtable.grpc.BigtableDataClient;
-import com.google.cloud.bigtable.grpc.scanner.FlatRow;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
-import com.google.protobuf.MessageLite;
 
 /**
  * This class provides management of asynchronous Bigtable RPCs. It ensures that there aren't too
@@ -49,80 +38,6 @@ public class AsyncExecutor {
   /** Constant <code>LOG</code> */
   protected static final Logger LOG = new Logger(AsyncExecutor.class);
 
-  protected interface AsyncCall<RequestT, ResponseT> {
-    ListenableFuture<ResponseT> call(BigtableDataClient client, RequestT request);
-  }
-
-  /**
-   * Calls {@link BigtableDataClient#mutateRowAsync(MutateRowRequest)}.
-   */
-  protected static AsyncCall<MutateRowRequest, MutateRowResponse> MUTATE_ROW_ASYNC =
-      new AsyncCall<MutateRowRequest, MutateRowResponse>() {
-        @Override
-        public ListenableFuture<MutateRowResponse> call(BigtableDataClient client, MutateRowRequest request) {
-          return client.mutateRowAsync(request);
-        }
-      };
-
-  /**
-   * Calls {@link BigtableDataClient#mutateRowsAsync(MutateRowsRequest)}.
-   */
-  protected static AsyncCall<MutateRowsRequest, List<MutateRowsResponse>> MUTATE_ROWS_ASYNC =
-      new AsyncCall<MutateRowsRequest, List<MutateRowsResponse>>() {
-        @Override
-        public ListenableFuture<List<MutateRowsResponse>> call(BigtableDataClient client,
-            MutateRowsRequest request) {
-          return client.mutateRowsAsync(request);
-        }
-      };
-
-  /**
-   * Calls {@link BigtableDataClient#readModifyWriteRowAsync(ReadModifyWriteRowRequest)}.
-   */
-  protected static AsyncCall<ReadModifyWriteRowRequest, ReadModifyWriteRowResponse> READ_MODIFY_WRITE_ASYNC =
-      new AsyncCall<ReadModifyWriteRowRequest, ReadModifyWriteRowResponse>() {
-        @Override
-        public ListenableFuture<ReadModifyWriteRowResponse> call(BigtableDataClient client,
-            ReadModifyWriteRowRequest request) {
-          return client.readModifyWriteRowAsync(request);
-        }
-      };
-
-  /**
-   * Calls {@link BigtableDataClient#checkAndMutateRowAsync(CheckAndMutateRowRequest)}.
-   */
-  protected static AsyncCall<CheckAndMutateRowRequest, CheckAndMutateRowResponse> CHECK_AND_MUTATE_ASYNC =
-      new AsyncCall<CheckAndMutateRowRequest, CheckAndMutateRowResponse>() {
-        @Override
-        public ListenableFuture<CheckAndMutateRowResponse> call(BigtableDataClient client,
-            CheckAndMutateRowRequest request) {
-          return client.checkAndMutateRowAsync(request);
-        }
-      };
-
-  /**
-   * Calls {@link BigtableDataClient#readRowsAsync(ReadRowsRequest)}.
-   */
-  protected static AsyncCall<ReadRowsRequest, List<Row>> READ_ROWS_ASYNC =
-      new AsyncCall<ReadRowsRequest, List<Row>>() {
-        @Override
-        public ListenableFuture<List<Row>> call(BigtableDataClient client, ReadRowsRequest request) {
-          return client.readRowsAsync(request);
-        }
-      };
-
-
-  /**
-   * Calls {@link BigtableDataClient#readRowsAsync(ReadRowsRequest)}.
-   */
-  protected static AsyncCall<ReadRowsRequest, List<FlatRow>> READ_FLAT_ROWS_ASYNC =
-      new AsyncCall<ReadRowsRequest, List<FlatRow>>() {
-        @Override
-        public ListenableFuture<List<FlatRow>> call(BigtableDataClient client, ReadRowsRequest request) {
-          return client.readFlatRowsAsync(request);
-        }
-      };
-
   private final BigtableDataClient client;
   private final OperationAccountant operationsAccountant;
 
@@ -152,48 +67,6 @@ public class AsyncExecutor {
     this.operationsAccountant = operationAccountant;
   }
 
-  /**
-   * Performs a {@link com.google.cloud.bigtable.grpc.BigtableDataClient#mutateRowAsync(MutateRowRequest)} on the
-   * {@link com.google.bigtable.v2.MutateRowRequest}. This method may block if
-   * {@link OperationAccountant#registerOperation(ListenableFuture)} blocks.
-   *
-   * @param request The {@link com.google.bigtable.v2.MutateRowRequest} to send.
-   * @return a {@link com.google.common.util.concurrent.ListenableFuture} which can be listened to for completion events.
-   * @throws java.lang.InterruptedException if any.
-   */
-  public ListenableFuture<MutateRowResponse> mutateRowAsync(MutateRowRequest request)
-      throws InterruptedException {
-    return call(MUTATE_ROW_ASYNC, request);
-  }
-
-  /**
-   * Performs a {@link com.google.cloud.bigtable.grpc.BigtableDataClient#mutateRowsAsync(MutateRowsRequest)} on the
-   * {@link com.google.bigtable.v2.MutateRowsRequest}. This method may block if
-   * {@link OperationAccountant#registerOperation(ListenableFuture)} blocks.
-   *
-   * @param request The {@link com.google.bigtable.v2.MutateRowRequest} to send.
-   * @return a {@link com.google.common.util.concurrent.ListenableFuture} which can be listened to for completion events.
-   * @throws java.lang.InterruptedException if any.
-   */
-  public ListenableFuture<List<MutateRowsResponse>> mutateRowsAsync(MutateRowsRequest request)
-      throws InterruptedException {
-    return call(MUTATE_ROWS_ASYNC, request);
-  }
-
-  /**
-   * Performs a {@link com.google.cloud.bigtable.grpc.BigtableDataClient#checkAndMutateRowAsync(CheckAndMutateRowRequest)} on the
-   * {@link com.google.bigtable.v2.CheckAndMutateRowRequest}. This method may block if
-   * {@link OperationAccountant#registerOperation(ListenableFuture)} blocks.
-   *
-   * @param request The {@link com.google.bigtable.v2.CheckAndMutateRowRequest} to send.
-   * @return a {@link com.google.common.util.concurrent.ListenableFuture} which can be listened to for completion events.
-   * @throws java.lang.InterruptedException if any.
-   */
-  public ListenableFuture<CheckAndMutateRowResponse> checkAndMutateRowAsync(
-      CheckAndMutateRowRequest request) throws InterruptedException {
-    return call(CHECK_AND_MUTATE_ASYNC, request);
-  }
-
   /**
    * Performs a {@link com.google.cloud.bigtable.grpc.BigtableDataClient#readModifyWriteRow(ReadModifyWriteRowRequest)} on the
    * {@link com.google.bigtable.v2.ReadModifyWriteRowRequest}. This method may block if
@@ -201,48 +74,14 @@ public class AsyncExecutor {
    *
    * @param request The {@link com.google.bigtable.v2.ReadModifyWriteRowRequest} to send.
    * @return a {@link com.google.common.util.concurrent.ListenableFuture} which can be listened to for completion events.
-   * @throws java.lang.InterruptedException if any.
    */
   public ListenableFuture<ReadModifyWriteRowResponse>
-      readModifyWriteRowAsync(ReadModifyWriteRowRequest request) throws InterruptedException {
-    return call(READ_MODIFY_WRITE_ASYNC, request);
-  }
-
-  /**
-   * Performs a {@link com.google.cloud.bigtable.grpc.BigtableDataClient#readRowsAsync(ReadRowsRequest)} on the
-   * {@link com.google.bigtable.v2.ReadRowsRequest}. This method may block if
-   * {@link OperationAccountant#registerOperation(ListenableFuture)} blocks.
-   *
-   * @param request The {@link com.google.bigtable.v2.ReadRowsRequest} to send.
-   * @return a {@link com.google.common.util.concurrent.ListenableFuture} which can be listened to for completion events.
-   * @throws java.lang.InterruptedException if any.
-   */
-  public ListenableFuture<List<Row>> readRowsAsync(ReadRowsRequest request)
-      throws InterruptedException {
-    return call(READ_ROWS_ASYNC, request);
-  }
-
-  /**
-   * Performs a {@link com.google.cloud.bigtable.grpc.BigtableDataClient#readRowsAsync(ReadRowsRequest)} on the
-   * {@link com.google.bigtable.v2.ReadRowsRequest}. This method may block if
-   * {@link OperationAccountant#registerOperation(ListenableFuture)} blocks.
-   *
-   * @param request The {@link com.google.bigtable.v2.ReadRowsRequest} to send.
-   * @return a {@link com.google.common.util.concurrent.ListenableFuture} which can be listened to for completion events.
-   * @throws java.lang.InterruptedException if any.
-   */
-  public ListenableFuture<List<FlatRow>> readFlatRowsAsync(ReadRowsRequest request)
-      throws InterruptedException {
-    return call(READ_FLAT_ROWS_ASYNC, request);
-  }
-
-  private <RequestT extends MessageLite, ResponseT> ListenableFuture<ResponseT> call(
-      AsyncCall<RequestT, ResponseT> rpc, RequestT request) {
+      readModifyWriteRowAsync(ReadModifyWriteRowRequest request) {
     // Wait until both the memory and rpc count maximum requirements are achieved before getting a
     // unique id used to track this request.
-    ListenableFuture<ResponseT> future;
+    ListenableFuture<ReadModifyWriteRowResponse> future;
     try {
-      future = rpc.call(client, request);
+      future = client.readModifyWriteRowAsync(request);
     } catch (Throwable e) {
       future = Futures.immediateFailedFuture(e);
     }
@@ -278,23 +117,4 @@ public class AsyncExecutor {
     return operationsAccountant.hasInflightOperations();
   }
 
-  /**
-   * <p>
-   * Getter for the field <code>client</code>.
-   * </p>
-   * @return a {@link com.google.cloud.bigtable.grpc.BigtableDataClient} object.
-   */
-  public BigtableDataClient getClient() {
-    return client;
-  }
-
-  /**
-   * <p>
-   * Getter for the field <code>operationsAccountant</code>.
-   * </p>
-   * @return a {@link com.google.cloud.bigtable.grpc.async.OperationAccountant} object.
-   */
-  public OperationAccountant getOperationAccountant() {
-    return operationsAccountant;
-  }
 }
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/CallController.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/CallController.java
index dcc89d66e..0205dd2f3 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/CallController.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/CallController.java
@@ -24,7 +24,7 @@ import io.grpc.stub.ClientCallStreamObserver;
 import javax.annotation.Nullable;
 
 /**
- * Wraps a {@Link ClientCall}, and implements {@ClientCallStreamObserver} to allow access to the call's
+ * Wraps a {@link ClientCall}, and implements {@link ClientCallStreamObserver} to allow access to the call's
  * underlying functionality.
  * <p>
  *   This class is intended to be used by the user to control flow and the life of the call.
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BatchExecutor.java b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BatchExecutor.java
index 83f5ba0f9..08a1c35af 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BatchExecutor.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BatchExecutor.java
@@ -23,11 +23,8 @@ import java.util.concurrent.ExecutionException;
 
 import javax.annotation.Nullable;
 
-import org.apache.hadoop.hbase.client.Append;
-import org.apache.hadoop.hbase.client.Delete;
 import org.apache.hadoop.hbase.client.Get;
-import org.apache.hadoop.hbase.client.Increment;
-import org.apache.hadoop.hbase.client.Put;
+import org.apache.hadoop.hbase.client.Mutation;
 import org.apache.hadoop.hbase.client.Result;
 import org.apache.hadoop.hbase.client.RetriesExhaustedWithDetailsException;
 import org.apache.hadoop.hbase.client.Row;
@@ -39,9 +36,6 @@ import com.google.bigtable.v2.ReadModifyWriteRowResponse;
 import com.google.cloud.bigtable.config.BigtableOptions;
 import com.google.cloud.bigtable.config.Logger;
 import com.google.cloud.bigtable.grpc.BigtableSession;
-import com.google.cloud.bigtable.grpc.BigtableTableName;
-import com.google.cloud.bigtable.grpc.async.AsyncExecutor;
-import com.google.cloud.bigtable.grpc.async.BulkMutation;
 import com.google.cloud.bigtable.grpc.async.BulkRead;
 import com.google.cloud.bigtable.grpc.scanner.FlatRow;
 import com.google.cloud.bigtable.hbase.adapters.Adapters;
@@ -57,7 +51,11 @@ import com.google.common.util.concurrent.SettableFuture;
 import org.apache.hadoop.hbase.util.Bytes;
 
 /**
- * Class to help BigtableTable with batch operations on an BigtableClient.
+ * Class to help BigtableTable with batch operations on an BigtableClient, such as
+ * {@link org.apache.hadoop.hbase.client.Table#batch(List, Object[])}.
+ * {@link org.apache.hadoop.hbase.client.Table#put(List)} and
+ * {@link org.apache.hadoop.hbase.client.Table#get(List)}.  This class relies on implementations found
+ * in {@link BulkRead} and in {@link BigtableBufferedMutatorHelper}.
  *
  * @author sduskis
  * @version $Id: $Id
@@ -135,31 +133,13 @@ public class BatchExecutor {
     }
   }
 
-  protected static class BulkOperation {
-    private BulkMutation bulkMutation;
-    private BulkRead bulkRead;
-
-    protected BulkOperation(
-        BigtableSession session,
-        BigtableTableName tableName) {
-      this.bulkRead = session.createBulkRead(tableName);
-      this.bulkMutation = session.createBulkMutation(tableName);
-    }
-
-    protected void flush() {
-      // If there is a bulk mutation in progress, then send it.
-      bulkMutation.sendUnsent();
-      bulkRead.flush();
-    }
-  }
-
-  protected final BigtableSession session;
-  protected final AsyncExecutor asyncExecutor;
   protected final BigtableOptions options;
   protected final HBaseRequestAdapter requestAdapter;
   protected final Timer batchTimer = BigtableClientMetrics.timer(MetricLevel.Info, "batch.latency");
   // Once the IBigtableDataClient interface is implemented this will be removed
   private final RequestContext requestContext;
+  private final BigtableBufferedMutatorHelper bufferedMutatorHelper;
+  private final BulkRead bulkRead;
 
   /**
    * Constructor for BatchExecutor.
@@ -169,19 +149,20 @@ public class BatchExecutor {
    *     object.
    */
   public BatchExecutor(BigtableSession session, HBaseRequestAdapter requestAdapter) {
-    this.session = session;
-    this.asyncExecutor = session.createAsyncExecutor();
-    this.options = session.getOptions();
     this.requestAdapter = requestAdapter;
+    this.options = session.getOptions();
     this.requestContext = RequestContext
         .create(options.getProjectId(), options.getInstanceId(), options.getAppProfileId());
+    this.bulkRead = session.createBulkRead(requestAdapter.getBigtableTableName());
+    this.bufferedMutatorHelper = new BigtableBufferedMutatorHelper(
+        requestAdapter,
+        null, // configuration isn't passed in, but also isn't used in BigtableBufferedMutatorHelper
+        session);
   }
 
   /**
    * Issue a single RPC recording the result into {@code results[index]} and if not-null, invoking
    * the supplied callback.
-   * @param bulkOperation An object that encapsulates a set of Mutations and Reads that happen via a
-   *          bulk API.
    * @param row The action to perform
    * @param callback The callback to invoke when the RPC completes and we have results
    * @param results An array of results, into which we should store the result of the operation
@@ -191,38 +172,27 @@ public class BatchExecutor {
    * @return A ListenableFuture that will have the result when the RPC completes.
    */
   private <R extends Row, T> ListenableFuture<Result> issueAsyncRowRequest(
-      BulkOperation bulkOperation, Row row, Batch.Callback<T> callback, Object[] results,
+      Row row, Batch.Callback<T> callback, Object[] results,
       int index) {
-    LOG.trace("issueRowRequest(BulkOperation, Row, Batch.Callback, Object[], index");
+    LOG.trace("issueRowRequest(Row, Batch.Callback, Object[], index");
     SettableFuture<Result> resultFuture = SettableFuture.create();
     RpcResultFutureCallback<T> futureCallback =
         new RpcResultFutureCallback<T>(row, callback, index, results, resultFuture);
     results[index] = null;
-    Futures.addCallback(issueAsyncRequest(bulkOperation, row),
+    Futures.addCallback(issueAsyncRequest(row),
         futureCallback, MoreExecutors.directExecutor());
     return resultFuture;
   }
 
-  private ListenableFuture<?> issueAsyncRequest(BulkOperation bulkOperation, Row row) {
+  private ListenableFuture<?> issueAsyncRequest(Row row) {
     try {
       if (row instanceof Get) {
-        return bulkOperation.bulkRead.add(requestAdapter.adapt((Get) row).toProto(requestContext));
-      } else if (row instanceof Put) {
-        return bulkOperation.bulkMutation.add(requestAdapter.adaptEntry((Put) row));
-      } else if (row instanceof Delete) {
-        return bulkOperation.bulkMutation.add(requestAdapter.adaptEntry((Delete) row));
-      } else if (row instanceof Append) {
-        return asyncExecutor.readModifyWriteRowAsync(
-            requestAdapter.adapt((Append) row).toProto(requestContext));
-      } else if (row instanceof Increment) {
-        return asyncExecutor.readModifyWriteRowAsync(
-            requestAdapter.adapt((Increment) row).toProto(requestContext));
+        return bulkRead.add(requestAdapter.adapt((Get) row).toProto(requestContext));
+      } else if (row instanceof Mutation) {
+        return bufferedMutatorHelper.mutate((Mutation) row);
       } else if (row instanceof RowMutations) {
-        return bulkOperation.bulkMutation.add(requestAdapter.adaptEntry((RowMutations) row));
+        return bufferedMutatorHelper.mutate((RowMutations) row);
       }
-    } catch (InterruptedException e) {
-      Thread.currentThread().interrupt();
-      return Futures.immediateFailedFuture(new IOException("Could not process the batch due to interrupt", e));
     } catch (Throwable e) {
       return Futures.immediateFailedFuture(new IOException("Could not process the batch", e));
     }
@@ -249,16 +219,16 @@ public class BatchExecutor {
 
   public <R> List<ListenableFuture<?>> issueAsyncRowRequests(List<? extends Row> actions,
       Object[] results, Batch.Callback<R> callback) {
-    BulkOperation bulkOperation = new BulkOperation(session, requestAdapter.getBigtableTableName());
     try {
       List<ListenableFuture<?>> resultFutures = new ArrayList<>(actions.size());
       for (int i = 0; i < actions.size(); i++) {
-        resultFutures
-            .add(issueAsyncRowRequest(bulkOperation, actions.get(i), callback, results, i));
+        resultFutures.add(issueAsyncRowRequest(actions.get(i), callback, results, i));
       }
       return resultFutures;
     } finally {
-      bulkOperation.flush();
+      // If there is a bulk mutation in progress, then send it.
+      bufferedMutatorHelper.sendUnsent();
+      bulkRead.flush();
     }
   }
 
@@ -315,7 +285,7 @@ public class BatchExecutor {
         } catch (ExecutionException e) {
           problemActions.add(actions.get(i));
           problems.add(e.getCause());
-          hosts.add(options.getDataHost().toString());
+          hosts.add(options.getDataHost());
         }
       }
       if (problems.size() > 0) {
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BigtableBufferedMutator.java b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BigtableBufferedMutator.java
index e3b56a5c4..dff3ed7aa 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BigtableBufferedMutator.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BigtableBufferedMutator.java
@@ -78,7 +78,7 @@ public class BigtableBufferedMutator implements BufferedMutator {
       BufferedMutator.ExceptionListener listener) {
     helper = new BigtableBufferedMutatorHelper(adapter, configuration, session);
     this.listener = listener;
-    this.host = session.getOptions().getDataHost().toString();
+    this.host = session.getOptions().getDataHost();
   }
 
   /** {@inheritDoc} */
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BigtableBufferedMutatorHelper.java b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BigtableBufferedMutatorHelper.java
index 0ff14a2c9..129ebe9f6 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BigtableBufferedMutatorHelper.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BigtableBufferedMutatorHelper.java
@@ -40,6 +40,7 @@ import com.google.cloud.bigtable.grpc.async.BulkMutation;
 import com.google.cloud.bigtable.hbase.adapters.HBaseRequestAdapter;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
+import org.apache.hadoop.hbase.client.RowMutations;
 
 /**
  * A helper for Bigtable's {@link org.apache.hadoop.hbase.client.BufferedMutator} implementations.
@@ -66,9 +67,9 @@ public class BigtableBufferedMutatorHelper {
   private final HBaseRequestAdapter adapter;
   private final AsyncExecutor asyncExecutor;
 
-  private BulkMutation bulkMutation = null;
+  private final BulkMutation bulkMutation;
 
-  private BigtableOptions options;
+  private final BigtableOptions options;
   // Once the IBigtableDataClient interface is implemented this will be removed
   protected final RequestContext requestContext;
 
@@ -143,6 +144,9 @@ public class BigtableBufferedMutatorHelper {
   public List<ListenableFuture<?>> mutate(List<? extends Mutation> mutations) {
     closedReadLock.lock();
     try {
+      if (closed) {
+        throw new IllegalStateException("Cannot mutate when the BufferedMutator is closed.");
+      }
       List<ListenableFuture<?>> futures = new ArrayList<>(mutations.size());
       for (Mutation mutation : mutations) {
         futures.add(offer(mutation));
@@ -171,16 +175,34 @@ public class BigtableBufferedMutatorHelper {
     }
   }
 
+  /**
+   *
+   * @param mutation
+   * @return
+   */
+  public ListenableFuture<?> mutate(final RowMutations mutation) {
+    closedReadLock.lock();
+    try {
+      if (closed) {
+        throw new IllegalStateException("Cannot mutate when the BufferedMutator is closed.");
+      }
+      if (mutation == null) {
+        return Futures.immediateFailedFuture(
+            new IllegalArgumentException("Cannot perform a mutation on a null object."));
+      } else {
+        return bulkMutation.add(adapter.adaptEntry(mutation));
+      }
+    } finally {
+      closedReadLock.unlock();
+    }
+  }
+
   /**
    * Send the operations to the async executor asynchronously.  The conversion from hbase
    * object to cloud bigtable proto and the async call both take time (microseconds worth) that
    * could be parallelized, or at least removed from the user's thread.
    */
   private ListenableFuture<?> offer(Mutation mutation) {
-    if (closed) {
-      Futures.immediateFailedFuture(
-        new IllegalStateException("Cannot mutate when the BufferedMutator is closed."));
-    }
     ListenableFuture<?> future = null;
     try {
       if (mutation == null) {
diff --git a/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/batch/common/CloudBigtableService.java b/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/batch/common/CloudBigtableService.java
index 1f2490901..20da82636 100755
--- a/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/batch/common/CloudBigtableService.java
+++ b/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/batch/common/CloudBigtableService.java
@@ -15,7 +15,7 @@
  */
 package com.google.cloud.bigtable.batch.common;
 
-import com.google.bigtable.repackaged.com.google.bigtable.v2.SampleRowKeysResponse;
+import com.google.bigtable.repackaged.com.google.cloud.bigtable.data.v2.models.KeyOffset;
 import com.google.cloud.bigtable.beam.CloudBigtableIO;
 import com.google.cloud.bigtable.beam.CloudBigtableTableConfiguration;
 
@@ -27,6 +27,5 @@ import java.util.List;
  * Cloud Bigtable
  */
 public interface CloudBigtableService {
-  List<SampleRowKeysResponse> getSampleRowKeys(CloudBigtableTableConfiguration config)
-      throws IOException;
+  List<KeyOffset> getSampleRowKeys(CloudBigtableTableConfiguration config) throws IOException;
 }
diff --git a/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/batch/common/CloudBigtableServiceImpl.java b/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/batch/common/CloudBigtableServiceImpl.java
index f07457eed..712876315 100755
--- a/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/batch/common/CloudBigtableServiceImpl.java
+++ b/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/batch/common/CloudBigtableServiceImpl.java
@@ -15,11 +15,8 @@
  */
 package com.google.cloud.bigtable.batch.common;
 
-import com.google.bigtable.repackaged.com.google.bigtable.v2.SampleRowKeysRequest;
-import com.google.bigtable.repackaged.com.google.bigtable.v2.SampleRowKeysResponse;
-import com.google.bigtable.repackaged.com.google.cloud.bigtable.config.BigtableOptions;
+import com.google.bigtable.repackaged.com.google.cloud.bigtable.data.v2.models.KeyOffset;
 import com.google.bigtable.repackaged.com.google.cloud.bigtable.grpc.BigtableSession;
-import com.google.bigtable.repackaged.com.google.cloud.bigtable.grpc.BigtableTableName;
 import com.google.cloud.bigtable.beam.CloudBigtableTableConfiguration;
 
 import java.io.IOException;
@@ -28,15 +25,10 @@ import java.util.List;
 public class CloudBigtableServiceImpl implements CloudBigtableService {
 
   @Override
-  public List<SampleRowKeysResponse> getSampleRowKeys(CloudBigtableTableConfiguration config)
+  public List<KeyOffset> getSampleRowKeys(CloudBigtableTableConfiguration config)
       throws IOException {
-    BigtableOptions bigtableOptions = config.toBigtableOptions();
-    try (BigtableSession session = new BigtableSession(bigtableOptions)) {
-      BigtableTableName tableName =
-          bigtableOptions.getInstanceName().toTableName(config.getTableId());
-      SampleRowKeysRequest request =
-          SampleRowKeysRequest.newBuilder().setTableName(tableName.toString()).build();
-      return session.getDataClient().sampleRowKeys(request);
+    try (BigtableSession session = new BigtableSession(config.toBigtableOptions())) {
+      return session.getClientWrapper().sampleRowKeys(config.getTableId());
     }
   }
 }
diff --git a/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/beam/CloudBigtableIO.java b/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/beam/CloudBigtableIO.java
index 651c0c54f..74c314957 100755
--- a/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/beam/CloudBigtableIO.java
+++ b/bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/beam/CloudBigtableIO.java
@@ -26,6 +26,8 @@ import java.util.NoSuchElementException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 
+import com.google.bigtable.repackaged.com.google.cloud.bigtable.data.v2.models.KeyOffset;
+import com.google.bigtable.repackaged.com.google.cloud.bigtable.data.v2.models.Query;
 import org.apache.beam.sdk.annotations.Experimental;
 import org.apache.beam.sdk.coders.CannotProvideCoderException;
 import org.apache.beam.sdk.coders.Coder;
@@ -66,7 +68,6 @@ import org.slf4j.LoggerFactory;
 
 import com.google.bigtable.repackaged.com.google.bigtable.v2.SampleRowKeysResponse;
 import com.google.bigtable.repackaged.com.google.cloud.bigtable.config.BulkOptions;
-import com.google.bigtable.repackaged.com.google.cloud.bigtable.grpc.BigtableDataClient;
 import com.google.bigtable.repackaged.com.google.cloud.bigtable.grpc.BigtableInstanceName;
 import com.google.bigtable.repackaged.com.google.cloud.bigtable.grpc.BigtableSession;
 import com.google.bigtable.repackaged.com.google.cloud.bigtable.grpc.BigtableSessionSharedThreadPools;
@@ -163,7 +164,7 @@ public class CloudBigtableIO {
      */
     private final CloudBigtableScanConfiguration configuration;
 
-    private transient List<SampleRowKeysResponse> sampleRowKeys;
+    private transient List<KeyOffset> sampleRowKeys;
 
     AbstractSource(CloudBigtableScanConfiguration configuration) {
       this.configuration = configuration;
@@ -185,8 +186,8 @@ public class CloudBigtableIO {
       List<SourceWithKeys> splits = new ArrayList<>();
       byte[] startKey = HConstants.EMPTY_START_ROW;
       long lastOffset = 0;
-      for (SampleRowKeysResponse response : getSampleRowKeys()) {
-        byte[] endKey = response.getRowKey().toByteArray();
+      for (KeyOffset response : getSampleRowKeys()) {
+        byte[] endKey = response.getKey().toByteArray();
         // Avoid empty regions.
         if (Bytes.equals(startKey, endKey) && startKey.length > 0) {
           continue;
@@ -274,11 +275,11 @@ public class CloudBigtableIO {
 
     /**
      * Performs a call to get sample row keys from
-     * {@link BigtableDataClient#sampleRowKeys(com.google.bigtable.repackaged.com.google.bigtable.v2.SampleRowKeysRequest)}
+     * {@link CloudBigtableServiceImpl#getSampleRowKeys(CloudBigtableTableConfiguration)}
      * if they are not yet cached. The sample row keys give information about tablet key boundaries
      * and estimated sizes.
      */
-    public synchronized List<SampleRowKeysResponse> getSampleRowKeys() throws IOException {
+    public synchronized List<KeyOffset> getSampleRowKeys() throws IOException {
       if (sampleRowKeys == null) {
         sampleRowKeys = new CloudBigtableServiceImpl().getSampleRowKeys(getConfiguration());
       }
@@ -286,7 +287,7 @@ public class CloudBigtableIO {
     }
 
     @VisibleForTesting
-    void setSampleRowKeys(List<SampleRowKeysResponse> sampleRowKeys) {
+    void setSampleRowKeys(List<KeyOffset> sampleRowKeys) {
       this.sampleRowKeys = sampleRowKeys;
     }
 
@@ -317,8 +318,8 @@ public class CloudBigtableIO {
 
       byte[] startKey = HConstants.EMPTY_START_ROW;
       long lastOffset = 0;
-      for (SampleRowKeysResponse response : getSampleRowKeys()) {
-        byte[] currentEndKey = response.getRowKey().toByteArray();
+      for (KeyOffset response : getSampleRowKeys()) {
+        byte[] currentEndKey = response.getKey().toByteArray();
         // Avoid empty regions.
         if (Bytes.equals(startKey, currentEndKey) && startKey.length != 0) {
           continue;
@@ -418,7 +419,7 @@ public class CloudBigtableIO {
     /**
      * Splits the table based on keys that belong to tablets, known as "regions" in the HBase API.
      * The current implementation uses the HBase {@link RegionLocator} interface, which calls
-     * {@link BigtableDataClient#sampleRowKeys(com.google.bigtable.repackaged.com.google.bigtable.v2.SampleRowKeysRequest)}
+     * {@link  CloudBigtableServiceImpl#getSampleRowKeys(CloudBigtableTableConfiguration)}
      * under the covers. A {@link SourceWithKeys} may correspond to a single region or a portion of
      * a region.
      * <p>
@@ -451,7 +452,7 @@ public class CloudBigtableIO {
 
     /**
      * Gets an estimated size based on data returned from
-     * {@link BigtableDataClient#sampleRowKeys(com.google.bigtable.repackaged.com.google.bigtable.v2.SampleRowKeysRequest)}.
+     * {@link CloudBigtableServiceImpl#getSampleRowKeys(CloudBigtableTableConfiguration)}.
      * The estimate will be high if a {@link Scan} is set on the
      * {@link CloudBigtableScanConfiguration}; in such cases, the estimate will not take the Scan
      * into account, and will return a larger estimate than what the {@link CloudBigtableIO.Reader}
@@ -468,8 +469,8 @@ public class CloudBigtableIO {
 
       byte[] startKey = HConstants.EMPTY_START_ROW;
       long lastOffset = 0;
-      for (SampleRowKeysResponse response : getSampleRowKeys()) {
-        byte[] currentEndKey = response.getRowKey().toByteArray();
+      for (KeyOffset response : getSampleRowKeys()) {
+        byte[] currentEndKey = response.getKey().toByteArray();
         // Avoid empty regions.
         if (Bytes.equals(startKey, currentEndKey) && startKey.length != 0) {
           continue;
@@ -620,7 +621,7 @@ public class CloudBigtableIO {
 
       // This will use cached data channels under the covers.
       session = new BigtableSession(BigtableOptionsFactory.fromConfiguration(config));
-      scanner = session.getDataClient().readFlatRows(source.getConfiguration().getRequest());
+      scanner = session.getClientWrapper().readFlatRows(Query.fromProto(source.getConfiguration().getRequest()));
     }
 
     /**
diff --git a/pom.xml b/pom.xml
index d8a602bf9..3faba1386 100644
--- a/pom.xml
+++ b/pom.xml
@@ -181,7 +181,7 @@ limitations under the License.
 
                         <offlineLinks>
                             <offlineLink>
-                                <url>https://google.github.io/guava/releases/20.0/api/docs/</url>
+                                <url>https://google.github.io/guava/releases/27.0-jre/api/docs/</url>
                                 <location>${basedir}/javadoc/guava-docs</location>
                             </offlineLink>
                             <offlineLink>


Test patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestAsyncExecutor.java b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestAsyncExecutor.java
index e2e2c5683..8cec21fde 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestAsyncExecutor.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestAsyncExecutor.java
@@ -62,26 +62,9 @@ public class TestAsyncExecutor {
     Assert.assertFalse(underTest.hasInflightRequests());
   }
 
-  @Test
-  public void testMutation() throws InterruptedException {
-    when(client.mutateRowAsync(any(MutateRowRequest.class))).thenReturn(future);
-    underTest.mutateRowAsync(MutateRowRequest.getDefaultInstance());
-    Assert.assertTrue(underTest.hasInflightRequests());
-    future.set("");
-    Assert.assertFalse(underTest.hasInflightRequests());
-  }
-
-  @Test
-  public void testCheckAndMutate() throws InterruptedException {
-    when(client.checkAndMutateRowAsync(any(CheckAndMutateRowRequest.class))).thenReturn(future);
-    underTest.checkAndMutateRowAsync(CheckAndMutateRowRequest.getDefaultInstance());
-    Assert.assertTrue(underTest.hasInflightRequests());
-    future.set("");
-    Assert.assertFalse(underTest.hasInflightRequests());
-  }
 
   @Test
-  public void testReadWriteModify() throws InterruptedException {
+  public void testReadWriteModify()  {
     when(client.readModifyWriteRowAsync(any(ReadModifyWriteRowRequest.class))).thenReturn(future);
     underTest.readModifyWriteRowAsync(ReadModifyWriteRowRequest.getDefaultInstance());
     Assert.assertTrue(underTest.hasInflightRequests());
@@ -90,30 +73,9 @@ public class TestAsyncExecutor {
   }
 
   @Test
-  public void testReadRowsAsync() throws InterruptedException {
-    when(client.readRowsAsync(any(ReadRowsRequest.class))).thenReturn(future);
-    underTest.readRowsAsync(ReadRowsRequest.getDefaultInstance());
-    Assert.assertTrue(underTest.hasInflightRequests());
-    future.set("");
-    Assert.assertFalse(underTest.hasInflightRequests());
-  }
-
-  @Test
-  public void testReadFlatRowsAsync() throws InterruptedException {
-    when(client.readFlatRowsAsync(any(ReadRowsRequest.class))).thenReturn(future);
-    underTest.readFlatRowsAsync(ReadRowsRequest.getDefaultInstance());
-    Assert.assertTrue(underTest.hasInflightRequests());
-    future.set("");
-    Assert.assertFalse(underTest.hasInflightRequests());
-  }
-
-  @Test
-  public void testInvalidMutation() throws Exception {
-    try {
-      when(client.mutateRowAsync(any(MutateRowRequest.class))).thenThrow(new RuntimeException());
-      underTest.mutateRowAsync(MutateRowRequest.getDefaultInstance());
-    } catch(Exception ignored) {
-    }
+  public void testInvalidMutation() {
+    when(client.readModifyWriteRowAsync(any(ReadModifyWriteRowRequest.class))).thenThrow(new RuntimeException());
+    underTest.readModifyWriteRowAsync(ReadModifyWriteRowRequest.getDefaultInstance());
     future.set(ReadRowsResponse.getDefaultInstance());
     Assert.assertFalse(underTest.hasInflightRequests());
   }
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBatchExecutor.java b/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBatchExecutor.java
index 3e7b71b6b..9164a0d2d 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBatchExecutor.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBatchExecutor.java
@@ -24,7 +24,6 @@ import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import com.google.bigtable.v2.MutateRowRequest;
 import com.google.bigtable.v2.MutateRowsRequest;
 import com.google.bigtable.v2.ReadModifyWriteRowRequest;
 import com.google.bigtable.v2.ReadRowsRequest;
@@ -70,7 +69,6 @@ import org.mockito.MockitoAnnotations;
 import org.mockito.invocation.InvocationOnMock;
 import org.mockito.stubbing.Answer;
 
-import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -134,8 +132,9 @@ public class TestBatchExecutor {
   private HBaseRequestAdapter requestAdapter;
 
   private BigtableOptions options;
+
   @Before
-  public void setup() throws InterruptedException {
+  public void setup() {
     options = BigtableOptions.builder()
         .setProjectId("projectId")
         .setInstanceId("instanceId")
@@ -196,28 +195,25 @@ public class TestBatchExecutor {
 
   @Test
   public void testShutdownService() throws Exception {
-    when(mockAsyncExecutor.mutateRowAsync(any(MutateRowRequest.class)))
-        .thenThrow(new IllegalStateException("closed"));
+    when(mockFuture.get()).thenThrow(new IllegalStateException("closed"));
     try {
       batch(Arrays.asList(randomPut()));
     } catch (RetriesExhaustedWithDetailsException e) {
       Assert.assertEquals(1, e.getCauses().size());
-      Assert.assertEquals(IOException.class, e.getCause(0).getClass());
+      Assert.assertEquals(IllegalStateException.class, e.getCause(0).getClass());
     }
   }
 
   @Test
   public void testAsyncException() throws Exception {
     String message = "Something bad happened";
-    when(mockAsyncExecutor.mutateRowAsync(any(MutateRowRequest.class)))
-        .thenThrow(new RuntimeException(message));
+    when(mockFuture.get()).thenThrow(new RuntimeException(message));
     try {
       batch(Arrays.asList(randomPut()));
     } catch (RetriesExhaustedWithDetailsException e) {
       Assert.assertEquals(1, e.getCauses().size());
-      Assert.assertEquals(IOException.class, e.getCause(0).getClass());
-      Assert.assertEquals(RuntimeException.class, e.getCause(0).getCause().getClass());
-      Assert.assertEquals(message, e.getCause(0).getCause().getMessage());
+      Assert.assertEquals(RuntimeException.class, e.getCause(0).getClass());
+      Assert.assertEquals(message, e.getCause(0).getMessage());
     }
   }
 
diff --git a/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBigtableBufferedMutator.java b/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBigtableBufferedMutator.java
index 7ebc3fa3c..a9f6eb444 100644
--- a/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBigtableBufferedMutator.java
+++ b/bigtable-client-core-parent/bigtable-hbase/src/test/java/com/google/cloud/bigtable/hbase/TestBigtableBufferedMutator.java
@@ -42,7 +42,6 @@ import org.junit.runners.JUnit4;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import com.google.bigtable.v2.MutateRowRequest;
 import com.google.bigtable.v2.MutateRowResponse;
 import com.google.bigtable.v2.MutateRowsRequest;
 import com.google.bigtable.v2.ReadModifyWriteRowRequest;
@@ -112,25 +111,23 @@ public class TestBigtableBufferedMutator {
   }
 
   @Test
-  public void testPut() throws IOException, InterruptedException {
+  public void testPut() throws IOException{
     when(mockBulkMutation.add(any(MutateRowsRequest.Entry.class))).thenReturn(future);
     BigtableBufferedMutator underTest = createMutator(new Configuration(false));
     underTest.mutate(SIMPLE_PUT);
-    verify(mockAsyncExecutor, times(0)).mutateRowAsync(any(MutateRowRequest.class));
     verify(mockBulkMutation, times(1)).add(any(MutateRowsRequest.Entry.class));
   }
 
   @Test
-  public void testDelete() throws IOException, InterruptedException {
+  public void testDelete() throws IOException {
     when(mockBulkMutation.add(any(MutateRowsRequest.Entry.class))).thenReturn(future);
     BigtableBufferedMutator underTest = createMutator(new Configuration(false));
     underTest.mutate(new Delete(EMPTY_BYTES));
-    verify(mockAsyncExecutor, times(0)).mutateRowAsync(any(MutateRowRequest.class));
     verify(mockBulkMutation, times(1)).add(any(MutateRowsRequest.Entry.class));
   }
 
   @Test
-  public void testIncrement() throws IOException, InterruptedException {
+  public void testIncrement() throws IOException {
     when(mockAsyncExecutor.readModifyWriteRowAsync(any(ReadModifyWriteRowRequest.class)))
         .thenReturn(future);
     BigtableBufferedMutator underTest = createMutator(new Configuration(false));
@@ -140,7 +137,7 @@ public class TestBigtableBufferedMutator {
   }
 
   @Test
-  public void testAppend() throws IOException, InterruptedException {
+  public void testAppend() throws IOException {
     when(mockAsyncExecutor.readModifyWriteRowAsync(any(ReadModifyWriteRowRequest.class)))
         .thenReturn(future);
     BigtableBufferedMutator underTest = createMutator(new Configuration(false));
@@ -168,7 +165,6 @@ public class TestBigtableBufferedMutator {
     when(mockBulkMutation.add(any(MutateRowsRequest.Entry.class))).thenReturn(future);
     final BigtableBufferedMutator underTest = createMutator(config);
     underTest.mutate(SIMPLE_PUT);
-    verify(mockAsyncExecutor, times(0)).mutateRowAsync(any(MutateRowRequest.class));
     verify(mockBulkMutation, times(1)).add(any(MutateRowsRequest.Entry.class));
     underTest.flush();
     verify(mockBulkMutation, times(1)).flush();
diff --git a/bigtable-dataflow-parent/bigtable-hbase-beam/src/test/java/com/google/cloud/bigtable/beam/CloudBigtableIOIntegrationTest.java b/bigtable-dataflow-parent/bigtable-hbase-beam/src/test/java/com/google/cloud/bigtable/beam/CloudBigtableIOIntegrationTest.java
index 315218788..621e9604a 100755
--- a/bigtable-dataflow-parent/bigtable-hbase-beam/src/test/java/com/google/cloud/bigtable/beam/CloudBigtableIOIntegrationTest.java
+++ b/bigtable-dataflow-parent/bigtable-hbase-beam/src/test/java/com/google/cloud/bigtable/beam/CloudBigtableIOIntegrationTest.java
@@ -17,6 +17,7 @@
 
 import com.google.bigtable.repackaged.com.google.bigtable.v2.SampleRowKeysResponse;
 import com.google.bigtable.repackaged.com.google.cloud.bigtable.config.Logger;
+import com.google.bigtable.repackaged.com.google.cloud.bigtable.data.v2.models.KeyOffset;
 import com.google.cloud.bigtable.hbase.BigtableConfiguration;
 import org.apache.beam.sdk.io.BoundedSource;
 import org.apache.beam.sdk.io.BoundedSource.BoundedReader;
@@ -252,10 +253,10 @@ public class CloudBigtableIOIntegrationTest {
       try {
         CloudBigtableIO.Source source =
             (CloudBigtableIO.Source) CloudBigtableIO.read(createScanConfig(tableName));
-        List<SampleRowKeysResponse> sampleRowKeys = source.getSampleRowKeys();
+        List<KeyOffset> sampleRowKeys = source.getSampleRowKeys();
         LOG.info("Creating BoundedSource in testEstimatedAndSplitForSmallTable()");
         long estimatedSizeBytes = source.getEstimatedSizeBytes(null);
-        SampleRowKeysResponse lastSample = sampleRowKeys.get(sampleRowKeys.size() - 1);
+        KeyOffset lastSample = sampleRowKeys.get(sampleRowKeys.size() - 1);
         Assert.assertEquals(lastSample.getOffsetBytes(), estimatedSizeBytes);
 
         LOG.info("Creating Bundles in testEstimatedAndSplitForSmallTable()");
@@ -294,10 +295,10 @@ public class CloudBigtableIOIntegrationTest {
         LOG.info("Getting Source in testEstimatedAndSplitForLargeTable()");
         CloudBigtableIO.Source source =
             (CloudBigtableIO.Source) CloudBigtableIO.read(createScanConfig(tableName));
-        List<SampleRowKeysResponse> sampleRowKeys = source.getSampleRowKeys();
+        List<KeyOffset> sampleRowKeys = source.getSampleRowKeys();
         LOG.info("Getting estimated size in testEstimatedAndSplitForLargeTable()");
         long estimatedSizeBytes = source.getEstimatedSizeBytes(null);
-        SampleRowKeysResponse lastSample = sampleRowKeys.get(sampleRowKeys.size() - 1);
+        KeyOffset lastSample = sampleRowKeys.get(sampleRowKeys.size() - 1);
         Assert.assertEquals(lastSample.getOffsetBytes(), estimatedSizeBytes);
 
         LOG.info("Getting Bundles in testEstimatedAndSplitForLargeTable()");
diff --git a/bigtable-dataflow-parent/bigtable-hbase-beam/src/test/java/com/google/cloud/bigtable/beam/CloudBigtableIOTest.java b/bigtable-dataflow-parent/bigtable-hbase-beam/src/test/java/com/google/cloud/bigtable/beam/CloudBigtableIOTest.java
index 27589c057..0ec463651 100755
--- a/bigtable-dataflow-parent/bigtable-hbase-beam/src/test/java/com/google/cloud/bigtable/beam/CloudBigtableIOTest.java
+++ b/bigtable-dataflow-parent/bigtable-hbase-beam/src/test/java/com/google/cloud/bigtable/beam/CloudBigtableIOTest.java
@@ -25,6 +25,7 @@ import java.util.Comparator;
 import java.util.Iterator;
 import java.util.List;
 
+import com.google.bigtable.repackaged.com.google.cloud.bigtable.data.v2.models.KeyOffset;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.coders.CannotProvideCoderException;
 import org.apache.beam.sdk.coders.Coder;
@@ -109,7 +110,7 @@ public class CloudBigtableIOTest {
 
   @Test
   public void testSampleRowKeys() throws Exception {
-    List<SampleRowKeysResponse> sampleRowKeys = new ArrayList<>();
+    List<KeyOffset> sampleRowKeys = new ArrayList<>();
     int count = (int) (AbstractSource.COUNT_MAX_SPLIT_COUNT * 3 - 5);
     byte[][] keys = Bytes.split("A".getBytes(), "Z".getBytes(), count-2);
     long tabletSize = 2L * 1024L * 1024L * 1024L;
@@ -117,10 +118,7 @@ public class CloudBigtableIOTest {
     for (byte[] currentKey : keys) {
       boundary += tabletSize;
       try {
-        sampleRowKeys.add(SampleRowKeysResponse.newBuilder()
-          .setRowKey(ByteString.copyFrom(currentKey))
-          .setOffsetBytes(boundary)
-          .build());
+        sampleRowKeys.add(KeyOffset.create(ByteString.copyFrom(currentKey), boundary));
       } catch (NoClassDefFoundError e) {
         // This could cause some problems for javadoc or cobertura because of the shading magic we
         // do.