diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java index 85529de8d904..af946ac5c267 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java @@ -1519,6 +1519,81 @@ public ApiFuture> 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}. + * + *

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. + * + *

The returned samples are constrained by the provided {@link ByteStringRange}, and the last + * sample returned will always match the end key of the range. + * + *

Sample code: + * + *

{@code
+   * try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
+   *   String tableId = "[TABLE_ID]";
+   *   ByteStringRange range = ByteStringRange.create("[START_KEY]", "[END_KEY]");
+   *
+   *   List keyOffsets = bigtableDataClient.sampleRowKeys(TableId.of(tableId), range);
+   *   for(KeyOffset keyOffset : keyOffsets) {
+   *     // Do something with keyOffset
+   *   }
+   * } catch(ApiException e) {
+   *   e.printStackTrace();
+   * }
+   * }
+ * + * @throws com.google.api.gax.rpc.ApiException when a serverside error occurs + */ + public List 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}. + * + *

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. + * + *

The returned samples are constrained by the provided {@link ByteStringRange}, and the last + * sample returned will always match the end key of the range. + * + *

Sample code: + * + *

{@code
+   * try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
+   *   String tableId = "[TABLE_ID]";
+   *   ByteStringRange range = ByteStringRange.create("[START_KEY]", "[END_KEY]");
+   *   ApiFuture> keyOffsetsFuture = bigtableDataClient.sampleRowKeysAsync(TableId.of(tableId), range);
+   *
+   *   ApiFutures.addCallback(keyOffsetsFuture, new ApiFutureCallback>() {
+   *     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 keyOffsets) {
+   *       System.out.println("Got key offsets: " + keyOffsets);
+   *     }
+   *   }, MoreExecutors.directExecutor());
+   * }
+   * }
+ * + * @see com.google.cloud.bigtable.data.v2.models.AuthorizedViewId + * @see TableId + */ + public ApiFuture> 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 diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Range.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Range.java index a3cdff5912f6..2727661baf18 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Range.java +++ b/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) { diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/SampleRowKeysRequest.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/SampleRowKeysRequest.java index 78a444019cb6..1b3bd7216e1c 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/SampleRowKeysRequest.java +++ b/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); + } } } diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java index ec82f1c12c8a..d317bc5a9d94 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientTests.java +++ b/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.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); diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/SampleRowsIT.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/SampleRowsIT.java index 063d0d1f5095..cb7cdcc30d35 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/it/SampleRowsIT.java +++ b/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> future = client.sampleRowKeysAsync(TableId.of(tableId), range); + + List results = future.get(1, TimeUnit.MINUTES); + + List 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; + } } diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/SampleRowKeysRequestTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/SampleRowKeysRequestTest.java index 5b9c0cee4b80..6add519705ef 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/SampleRowKeysRequestTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/SampleRowKeysRequestTest.java @@ -187,5 +187,67 @@ public void testEquality() { assertThat(SampleRowKeysRequest.create(AuthorizedViewId.of(TABLE_ID, AUTHORIZED_VIEW_ID))) .isNotEqualTo(SampleRowKeysRequest.create(TableId.of(TABLE_ID))); + + // Test with row range + Range.ByteStringRange range1 = Range.ByteStringRange.create("a", "b"); + Range.ByteStringRange range2 = Range.ByteStringRange.create("a", "c"); + + assertThat( + SampleRowKeysRequest.newBuilder() + .setTargetId(TableId.of(TABLE_ID)) + .setRowRange(range1) + .build()) + .isEqualTo( + SampleRowKeysRequest.newBuilder() + .setTargetId(TableId.of(TABLE_ID)) + .setRowRange(range1) + .build()); + + assertThat( + SampleRowKeysRequest.newBuilder() + .setTargetId(TableId.of(TABLE_ID)) + .setRowRange(range1) + .build()) + .isNotEqualTo( + SampleRowKeysRequest.newBuilder() + .setTargetId(TableId.of(TABLE_ID)) + .setRowRange(range2) + .build()); + } + + @Test + public void toProtoWithRowRangeTest() { + Range.ByteStringRange range = Range.ByteStringRange.create("start", "end"); + SampleRowKeysRequest sampleRowKeysRequest = + SampleRowKeysRequest.newBuilder() + .setTargetId(TableId.of(TABLE_ID)) + .setRowRange(range) + .build(); + + com.google.bigtable.v2.SampleRowKeysRequest actualRequest = + sampleRowKeysRequest.toProto(REQUEST_CONTEXT); + + assertThat(actualRequest.getTableName()) + .isEqualTo(NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID)); + assertThat(actualRequest.hasRowRange()).isTrue(); + assertThat(actualRequest.getRowRange().getStartKeyClosed()).isEqualTo(range.getStart()); + assertThat(actualRequest.getRowRange().getEndKeyOpen()).isEqualTo(range.getEnd()); + } + + @Test + public void fromProtoWithRowRangeTest() { + Range.ByteStringRange range = Range.ByteStringRange.create("start", "end"); + SampleRowKeysRequest sampleRowKeysRequest = + SampleRowKeysRequest.newBuilder() + .setTargetId(TableId.of(TABLE_ID)) + .setRowRange(range) + .build(); + + com.google.bigtable.v2.SampleRowKeysRequest protoRequest = + sampleRowKeysRequest.toProto(REQUEST_CONTEXT); + SampleRowKeysRequest actualRequest = SampleRowKeysRequest.fromProto(protoRequest); + + assertThat(actualRequest.toProto(REQUEST_CONTEXT)).isEqualTo(protoRequest); + assertThat(actualRequest.getRowRange()).isEqualTo(range); } }