feat: Implement row_range parameter support in SampleRowKeys API

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java

@@ -1519,6 +1519,81 @@ public ApiFuture<List<KeyOffset>> sampleRowKeysAsync(TargetId targetId) {
     return sampleRowKeysCallableWithRequest().futureCall(SampleRowKeysRequest.create(targetId));
   }
 
+  /**
+   * Convenience method to synchronously return a sample of row keys on the specified {@link
+   * TargetId} within the specified {@link ByteStringRange}.
+   *
+   * <p>The returned row keys will delimit contiguous sections of the table of approximately equal
+   * size, which can be used to break up the data for distributed tasks like mapreduces.
+   *
+   * <p>The returned samples are constrained by the provided {@link ByteStringRange}, and the last
+   * sample returned will always match the end key of the range.
+   *
+   * <p>Sample code:
+   *
+   * <pre>{@code
+   * try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
+   *   String tableId = "[TABLE_ID]";
+   *   ByteStringRange range = ByteStringRange.create("[START_KEY]", "[END_KEY]");
+   *
+   *   List<KeyOffset> keyOffsets = bigtableDataClient.sampleRowKeys(TableId.of(tableId), range);
+   *   for(KeyOffset keyOffset : keyOffsets) {
+   *     // Do something with keyOffset
+   *   }
+   * } catch(ApiException e) {
+   *   e.printStackTrace();
+   * }
+   * }</pre>
+   *
+   * @throws com.google.api.gax.rpc.ApiException when a serverside error occurs
+   */
+  public List<KeyOffset> sampleRowKeys(TargetId targetId, ByteStringRange range) {
+    return ApiExceptions.callAndTranslateApiException(sampleRowKeysAsync(targetId, range));
+  }
+
+  /**
+   * Convenience method to asynchronously return a sample of row keys on the specified {@link
+   * TargetId} within the specified {@link ByteStringRange}.
+   *
+   * <p>The returned row keys will delimit contiguous sections of the table of approximately equal
+   * size, which can be used to break up the data for distributed tasks like mapreduces.
+   *
+   * <p>The returned samples are constrained by the provided {@link ByteStringRange}, and the last
+   * sample returned will always match the end key of the range.
+   *
+   * <p>Sample code:
+   *
+   * <pre>{@code
+   * try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
+   *   String tableId = "[TABLE_ID]";
+   *   ByteStringRange range = ByteStringRange.create("[START_KEY]", "[END_KEY]");
+   *   ApiFuture<List<KeyOffset>> keyOffsetsFuture = bigtableDataClient.sampleRowKeysAsync(TableId.of(tableId), range);
+   *
+   *   ApiFutures.addCallback(keyOffsetsFuture, new ApiFutureCallback<List<KeyOffset>>() {
+   *     public void onFailure(Throwable t) {
+   *       if (t instanceof NotFoundException) {
+   *         System.out.println("Tried to sample keys of a non-existent table");
+   *       } else {
+   *         t.printStackTrace();
+   *       }
+   *     }
+   *     public void onSuccess(List<KeyOffset> keyOffsets) {
+   *       System.out.println("Got key offsets: " + keyOffsets);
+   *     }
+   *   }, MoreExecutors.directExecutor());
+   * }
+   * }</pre>
+   *
+   * @see com.google.cloud.bigtable.data.v2.models.AuthorizedViewId
+   * @see TableId
+   */
+  public ApiFuture<List<KeyOffset>> sampleRowKeysAsync(TargetId targetId, ByteStringRange range) {
+    com.google.common.base.Preconditions.checkNotNull(range, "range can't be null.");
+    return sampleRowKeysCallableWithRequest()
+        .futureCall(
+            SampleRowKeysRequest.newBuilder().setTargetId(targetId).setRowRange(range).build());
+  }
+
   /**
    * Returns a sample of row keys in the table. The returned row keys will delimit contiguous
    * sections of the table of approximately equal size, which can be used to break up the data for

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Range.java

@@ -414,6 +414,66 @@ public static ByteStringRange toByteStringRange(ByteString byteString)
       return ByteStringRange.create(rowRange.getStartKeyClosed(), rowRange.getEndKeyOpen());
     }
 
+    @InternalApi
+    public RowRange toProto() {
+      RowRange.Builder rangeBuilder = RowRange.newBuilder();
+      switch (getStartBound()) {
+        case OPEN:
+          rangeBuilder.setStartKeyOpen(getStart());
+          break;
+        case CLOSED:
+          rangeBuilder.setStartKeyClosed(getStart());
+          break;
+        case UNBOUNDED:
+          rangeBuilder.clearStartKey();
+          break;
+        default:
+          throw new IllegalStateException("Unknown start bound: " + getStartBound());
+      }
+      switch (getEndBound()) {
+        case OPEN:
+          rangeBuilder.setEndKeyOpen(getEnd());
+          break;
+        case CLOSED:
+          rangeBuilder.setEndKeyClosed(getEnd());
+          break;
+        case UNBOUNDED:
+          rangeBuilder.clearEndKey();
+          break;
+        default:
+          throw new IllegalStateException("Unknown end bound: " + getEndBound());
+      }
+      return rangeBuilder.build();
+    }
+
+    @InternalApi
+    public static ByteStringRange fromProto(RowRange rowRange) {
+      ByteStringRange range = ByteStringRange.unbounded();
+      switch (rowRange.getStartKeyCase()) {
+        case START_KEY_CLOSED:
+          range.startClosed(rowRange.getStartKeyClosed());
+          break;
+        case START_KEY_OPEN:
+          range.startOpen(rowRange.getStartKeyOpen());
+          break;
+        case STARTKEY_NOT_SET:
+          range.startUnbounded();
+          break;
+      }
+      switch (rowRange.getEndKeyCase()) {
+        case END_KEY_CLOSED:
+          range.endClosed(rowRange.getEndKeyClosed());
+          break;
+        case END_KEY_OPEN:
+          range.endOpen(rowRange.getEndKeyOpen());
+          break;
+        case ENDKEY_NOT_SET:
+          range.endUnbounded();
+          break;
+      }
+      return range;
+    }
+
     @Override
     public boolean equals(Object o) {
       if (this == o) {

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/SampleRowKeysRequest.java

@@ -19,6 +19,7 @@
 import com.google.api.core.InternalApi;
 import com.google.cloud.bigtable.data.v2.internal.NameUtil;
 import com.google.cloud.bigtable.data.v2.internal.RequestContext;
+import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
 import com.google.common.base.Objects;
 import com.google.common.base.Preconditions;
 import java.io.Serializable;
@@ -27,15 +28,32 @@
 /** Wraps a {@link com.google.bigtable.v2.SampleRowKeysRequest}. */
 public final class SampleRowKeysRequest implements Serializable {
   private final TargetId targetId;
+  private final ByteStringRange rowRange;
 
-  private SampleRowKeysRequest(TargetId targetId) {
-    Preconditions.checkNotNull(targetId, "target id can't be null.");
-    this.targetId = targetId;
+  private SampleRowKeysRequest(Builder builder) {
+    this.targetId = Preconditions.checkNotNull(builder.targetId, "target id can't be null.");
+    this.rowRange = builder.rowRange;
   }
 
   /** Creates a new instance of the sample row keys builder for the given target with targetId */
   public static SampleRowKeysRequest create(TargetId targetId) {
-    return new SampleRowKeysRequest(targetId);
+    return newBuilder().setTargetId(targetId).build();
+  }
+
+  public static Builder newBuilder() {
+    return new Builder();
+  }
+
+  public Builder toBuilder() {
+    return new Builder(this);
+  }
+
+  public TargetId getTargetId() {
+    return targetId;
+  }
+
+  public ByteStringRange getRowRange() {
+    return rowRange;
   }
 
   @InternalApi
@@ -51,6 +69,9 @@ public com.google.bigtable.v2.SampleRowKeysRequest toProto(RequestContext reques
     } else {
       builder.setTableName(resourceName);
     }
+    if (rowRange != null && !rowRange.equals(ByteStringRange.unbounded())) {
+      builder.setRowRange(rowRange.toProto());
+    }
     return builder.setAppProfileId(requestContext.getAppProfileId()).build();
   }
 
@@ -67,11 +88,14 @@ public static SampleRowKeysRequest fromProto(
     String authorizedViewName = request.getAuthorizedViewName();
     String materializedViewName = request.getMaterializedViewName();
 
-    SampleRowKeysRequest sampleRowKeysRequest =
-        SampleRowKeysRequest.create(
-            NameUtil.extractTargetId(tableName, authorizedViewName, materializedViewName));
-
-    return sampleRowKeysRequest;
+    Builder builder =
+        newBuilder()
+            .setTargetId(
+                NameUtil.extractTargetId(tableName, authorizedViewName, materializedViewName));
+    if (request.hasRowRange()) {
+      builder.setRowRange(ByteStringRange.fromProto(request.getRowRange()));
+    }
+    return builder.build();
   }
 
   @Override
@@ -83,11 +107,38 @@ public boolean equals(Object o) {
       return false;
     }
     SampleRowKeysRequest sampleRowKeysRequest = (SampleRowKeysRequest) o;
-    return Objects.equal(targetId, sampleRowKeysRequest.targetId);
+    return Objects.equal(targetId, sampleRowKeysRequest.targetId)
+        && Objects.equal(rowRange, sampleRowKeysRequest.rowRange);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hashCode(targetId);
+    return Objects.hashCode(targetId, rowRange);
+  }
+
+  public static final class Builder {
+    private TargetId targetId;
+    private ByteStringRange rowRange = ByteStringRange.unbounded();
+
+    private Builder() {}
+
+    private Builder(SampleRowKeysRequest request) {
+      this.targetId = request.targetId;
+      this.rowRange = request.rowRange;
+    }
+
+    public Builder setTargetId(TargetId targetId) {
+      this.targetId = targetId;
+      return this;
+    }
+
+    public Builder setRowRange(ByteStringRange rowRange) {
+      this.rowRange = Preconditions.checkNotNull(rowRange, "rowRange can't be null.");
+      return this;
+    }
+
+    public SampleRowKeysRequest build() {
+      return new SampleRowKeysRequest(this);
+    }
   }
 }

java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java

@@ -730,6 +730,42 @@ public void sampleRowKeysOnAuthorizedViewTest() {
             SampleRowKeysRequest.create(AuthorizedViewId.of("fake-table", "fake-authorized-view")));
   }
 
+  @Test
+  public void proxySampleRowKeysWithRangeTest() {
+    Mockito.when(mockStub.sampleRowKeysCallableWithRequest())
+        .thenReturn(mockSampleRowKeysCallableWithRequest);
+
+    ByteStringRange range = ByteStringRange.create("a", "b");
+    @SuppressWarnings("VariableUnused")
+    ApiFuture<?> ignored = bigtableDataClient.sampleRowKeysAsync(TableId.of("fake-table"), range);
+
+    Mockito.verify(mockSampleRowKeysCallableWithRequest)
+        .futureCall(
+            SampleRowKeysRequest.newBuilder()
+                .setTargetId(TableId.of("fake-table"))
+                .setRowRange(range)
+                .build());
+  }
+
+  @Test
+  public void sampleRowKeysWithRangeTest() {
+    Mockito.when(mockStub.sampleRowKeysCallableWithRequest())
+        .thenReturn(mockSampleRowKeysCallableWithRequest);
+
+    Mockito.when(
+            mockSampleRowKeysCallableWithRequest.futureCall(
+                ArgumentMatchers.any(SampleRowKeysRequest.class)))
+        .thenReturn(ApiFutures.immediateFuture(Collections.<KeyOffset>emptyList()));
+    ByteStringRange range = ByteStringRange.create("a", "b");
+    bigtableDataClient.sampleRowKeys(TableId.of("fake-table"), range);
+    Mockito.verify(mockSampleRowKeysCallableWithRequest)
+        .futureCall(
+            SampleRowKeysRequest.newBuilder()
+                .setTargetId(TableId.of("fake-table"))
+                .setRowRange(range)
+                .build());
+  }
+
   @Test
   public void proxyMutateRowCallableTest() {
     Mockito.when(mockStub.mutateRowCallable()).thenReturn(mockMutateRowCallable);

java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/SampleRowsIT.java

@@ -27,7 +27,9 @@
 import com.google.cloud.bigtable.data.v2.BigtableDataClient;
 import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId;
 import com.google.cloud.bigtable.data.v2.models.KeyOffset;
+import com.google.cloud.bigtable.data.v2.models.Range;
 import com.google.cloud.bigtable.data.v2.models.RowMutation;
+import com.google.cloud.bigtable.data.v2.models.TableId;
 import com.google.cloud.bigtable.test_helpers.env.EmulatorEnv;
 import com.google.cloud.bigtable.test_helpers.env.PrefixGenerator;
 import com.google.cloud.bigtable.test_helpers.env.TestEnvRule;
@@ -127,4 +129,41 @@ private static AuthorizedView createPreSplitTableAndAuthorizedView() {
             .setDeletionProtection(false);
     return testEnvRule.env().getTableAdminClient().createAuthorizedView(request);
   }
+
+  @Test
+  public void testWithRowRange() throws InterruptedException, ExecutionException, TimeoutException {
+    String tableId =
+        createPreSplitTable(
+            "SampleRowsIT#RowRange", "apple", "banana", "cherry", "date", "eggplant");
+    BigtableDataClient client = testEnvRule.env().getDataClient();
+
+    try {
+      Range.ByteStringRange range = Range.ByteStringRange.create("banana", "date");
+
+      ApiFuture<List<KeyOffset>> future = client.sampleRowKeysAsync(TableId.of(tableId), range);
+
+      List<KeyOffset> results = future.get(1, TimeUnit.MINUTES);
+
+      List<ByteString> resultKeys = new ArrayList<>();
+      for (KeyOffset keyOffset : results) {
+        resultKeys.add(keyOffset.getKey());
+      }
+
+      assertThat(resultKeys)
+          .containsExactly(ByteString.copyFromUtf8("cherry"), ByteString.copyFromUtf8("date"));
+
+    } finally {
+      testEnvRule.env().getTableAdminClient().deleteTable(tableId);
+    }
+  }
+
+  private static String createPreSplitTable(String prefix, String... splitKeys) {
+    String tableId = PrefixGenerator.newPrefix(prefix);
+    CreateTableRequest request = CreateTableRequest.of(tableId);
+    for (String splitKey : splitKeys) {
+      request.addSplit(ByteString.copyFromUtf8(splitKey));
+    }
+    testEnvRule.env().getTableAdminClient().createTable(request);
+    return tableId;
+  }
 }