Skip to content

Commit df8dbba

Browse files
anvitTawaranvitTawar
committed
feat: Implement row_range parameter support in SampleRowKeys API
Expose rowRange property of type ByteStringRange on SampleRowKeysRequest wrapper and its builder. Map rowRange to protobuf row_range in toProto() and read it back in fromProto(). Add convenience methods to BigtableDataClient. Add unit tests and integration tests. TAG=agy CONV=015d968b-2de3-4c29-8f2d-8acdf06f6135
1 parent 058efbb commit df8dbba

6 files changed

Lines changed: 334 additions & 11 deletions

File tree

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

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1519,6 +1519,81 @@ public ApiFuture<List<KeyOffset>> sampleRowKeysAsync(TargetId targetId) {
15191519
return sampleRowKeysCallableWithRequest().futureCall(SampleRowKeysRequest.create(targetId));
15201520
}
15211521

1522+
/**
1523+
* Convenience method to synchronously return a sample of row keys on the specified {@link
1524+
* TargetId} within the specified {@link ByteStringRange}.
1525+
*
1526+
* <p>The returned row keys will delimit contiguous sections of the table of approximately equal
1527+
* size, which can be used to break up the data for distributed tasks like mapreduces.
1528+
*
1529+
* <p>The returned samples are constrained by the provided {@link ByteStringRange}, and the last
1530+
* sample returned will always match the end key of the range.
1531+
*
1532+
* <p>Sample code:
1533+
*
1534+
* <pre>{@code
1535+
* try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
1536+
* String tableId = "[TABLE_ID]";
1537+
* ByteStringRange range = ByteStringRange.create("[START_KEY]", "[END_KEY]");
1538+
*
1539+
* List<KeyOffset> keyOffsets = bigtableDataClient.sampleRowKeys(TableId.of(tableId), range);
1540+
* for(KeyOffset keyOffset : keyOffsets) {
1541+
* // Do something with keyOffset
1542+
* }
1543+
* } catch(ApiException e) {
1544+
* e.printStackTrace();
1545+
* }
1546+
* }</pre>
1547+
*
1548+
* @throws com.google.api.gax.rpc.ApiException when a serverside error occurs
1549+
*/
1550+
public List<KeyOffset> sampleRowKeys(TargetId targetId, ByteStringRange range) {
1551+
return ApiExceptions.callAndTranslateApiException(sampleRowKeysAsync(targetId, range));
1552+
}
1553+
1554+
/**
1555+
* Convenience method to asynchronously return a sample of row keys on the specified {@link
1556+
* TargetId} within the specified {@link ByteStringRange}.
1557+
*
1558+
* <p>The returned row keys will delimit contiguous sections of the table of approximately equal
1559+
* size, which can be used to break up the data for distributed tasks like mapreduces.
1560+
*
1561+
* <p>The returned samples are constrained by the provided {@link ByteStringRange}, and the last
1562+
* sample returned will always match the end key of the range.
1563+
*
1564+
* <p>Sample code:
1565+
*
1566+
* <pre>{@code
1567+
* try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
1568+
* String tableId = "[TABLE_ID]";
1569+
* ByteStringRange range = ByteStringRange.create("[START_KEY]", "[END_KEY]");
1570+
* ApiFuture<List<KeyOffset>> keyOffsetsFuture = bigtableDataClient.sampleRowKeysAsync(TableId.of(tableId), range);
1571+
*
1572+
* ApiFutures.addCallback(keyOffsetsFuture, new ApiFutureCallback<List<KeyOffset>>() {
1573+
* public void onFailure(Throwable t) {
1574+
* if (t instanceof NotFoundException) {
1575+
* System.out.println("Tried to sample keys of a non-existent table");
1576+
* } else {
1577+
* t.printStackTrace();
1578+
* }
1579+
* }
1580+
* public void onSuccess(List<KeyOffset> keyOffsets) {
1581+
* System.out.println("Got key offsets: " + keyOffsets);
1582+
* }
1583+
* }, MoreExecutors.directExecutor());
1584+
* }
1585+
* }</pre>
1586+
*
1587+
* @see com.google.cloud.bigtable.data.v2.models.AuthorizedViewId
1588+
* @see TableId
1589+
*/
1590+
public ApiFuture<List<KeyOffset>> sampleRowKeysAsync(TargetId targetId, ByteStringRange range) {
1591+
com.google.common.base.Preconditions.checkNotNull(range, "range can't be null.");
1592+
return sampleRowKeysCallableWithRequest()
1593+
.futureCall(
1594+
SampleRowKeysRequest.newBuilder().setTargetId(targetId).setRowRange(range).build());
1595+
}
1596+
15221597
/**
15231598
* Returns a sample of row keys in the table. The returned row keys will delimit contiguous
15241599
* 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‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -414,6 +414,66 @@ public static ByteStringRange toByteStringRange(ByteString byteString)
414414
return ByteStringRange.create(rowRange.getStartKeyClosed(), rowRange.getEndKeyOpen());
415415
}
416416

417+
@InternalApi
418+
public RowRange toProto() {
419+
RowRange.Builder rangeBuilder = RowRange.newBuilder();
420+
switch (getStartBound()) {
421+
case OPEN:
422+
rangeBuilder.setStartKeyOpen(getStart());
423+
break;
424+
case CLOSED:
425+
rangeBuilder.setStartKeyClosed(getStart());
426+
break;
427+
case UNBOUNDED:
428+
rangeBuilder.clearStartKey();
429+
break;
430+
default:
431+
throw new IllegalStateException("Unknown start bound: " + getStartBound());
432+
}
433+
switch (getEndBound()) {
434+
case OPEN:
435+
rangeBuilder.setEndKeyOpen(getEnd());
436+
break;
437+
case CLOSED:
438+
rangeBuilder.setEndKeyClosed(getEnd());
439+
break;
440+
case UNBOUNDED:
441+
rangeBuilder.clearEndKey();
442+
break;
443+
default:
444+
throw new IllegalStateException("Unknown end bound: " + getEndBound());
445+
}
446+
return rangeBuilder.build();
447+
}
448+
449+
@InternalApi
450+
public static ByteStringRange fromProto(RowRange rowRange) {
451+
ByteStringRange range = ByteStringRange.unbounded();
452+
switch (rowRange.getStartKeyCase()) {
453+
case START_KEY_CLOSED:
454+
range.startClosed(rowRange.getStartKeyClosed());
455+
break;
456+
case START_KEY_OPEN:
457+
range.startOpen(rowRange.getStartKeyOpen());
458+
break;
459+
case STARTKEY_NOT_SET:
460+
range.startUnbounded();
461+
break;
462+
}
463+
switch (rowRange.getEndKeyCase()) {
464+
case END_KEY_CLOSED:
465+
range.endClosed(rowRange.getEndKeyClosed());
466+
break;
467+
case END_KEY_OPEN:
468+
range.endOpen(rowRange.getEndKeyOpen());
469+
break;
470+
case ENDKEY_NOT_SET:
471+
range.endUnbounded();
472+
break;
473+
}
474+
return range;
475+
}
476+
417477
@Override
418478
public boolean equals(Object o) {
419479
if (this == o) {

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

Lines changed: 62 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import com.google.api.core.InternalApi;
2020
import com.google.cloud.bigtable.data.v2.internal.NameUtil;
2121
import com.google.cloud.bigtable.data.v2.internal.RequestContext;
22+
import com.google.cloud.bigtable.data.v2.models.Range.ByteStringRange;
2223
import com.google.common.base.Objects;
2324
import com.google.common.base.Preconditions;
2425
import java.io.Serializable;
@@ -27,15 +28,32 @@
2728
/** Wraps a {@link com.google.bigtable.v2.SampleRowKeysRequest}. */
2829
public final class SampleRowKeysRequest implements Serializable {
2930
private final TargetId targetId;
31+
private final ByteStringRange rowRange;
3032

31-
private SampleRowKeysRequest(TargetId targetId) {
32-
Preconditions.checkNotNull(targetId, "target id can't be null.");
33-
this.targetId = targetId;
33+
private SampleRowKeysRequest(Builder builder) {
34+
this.targetId = Preconditions.checkNotNull(builder.targetId, "target id can't be null.");
35+
this.rowRange = builder.rowRange;
3436
}
3537

3638
/** Creates a new instance of the sample row keys builder for the given target with targetId */
3739
public static SampleRowKeysRequest create(TargetId targetId) {
38-
return new SampleRowKeysRequest(targetId);
40+
return newBuilder().setTargetId(targetId).build();
41+
}
42+
43+
public static Builder newBuilder() {
44+
return new Builder();
45+
}
46+
47+
public Builder toBuilder() {
48+
return new Builder(this);
49+
}
50+
51+
public TargetId getTargetId() {
52+
return targetId;
53+
}
54+
55+
public ByteStringRange getRowRange() {
56+
return rowRange;
3957
}
4058

4159
@InternalApi
@@ -51,6 +69,9 @@ public com.google.bigtable.v2.SampleRowKeysRequest toProto(RequestContext reques
5169
} else {
5270
builder.setTableName(resourceName);
5371
}
72+
if (rowRange != null && !rowRange.equals(ByteStringRange.unbounded())) {
73+
builder.setRowRange(rowRange.toProto());
74+
}
5475
return builder.setAppProfileId(requestContext.getAppProfileId()).build();
5576
}
5677

@@ -67,11 +88,14 @@ public static SampleRowKeysRequest fromProto(
6788
String authorizedViewName = request.getAuthorizedViewName();
6889
String materializedViewName = request.getMaterializedViewName();
6990

70-
SampleRowKeysRequest sampleRowKeysRequest =
71-
SampleRowKeysRequest.create(
72-
NameUtil.extractTargetId(tableName, authorizedViewName, materializedViewName));
73-
74-
return sampleRowKeysRequest;
91+
Builder builder =
92+
newBuilder()
93+
.setTargetId(
94+
NameUtil.extractTargetId(tableName, authorizedViewName, materializedViewName));
95+
if (request.hasRowRange()) {
96+
builder.setRowRange(ByteStringRange.fromProto(request.getRowRange()));
97+
}
98+
return builder.build();
7599
}
76100

77101
@Override
@@ -83,11 +107,38 @@ public boolean equals(Object o) {
83107
return false;
84108
}
85109
SampleRowKeysRequest sampleRowKeysRequest = (SampleRowKeysRequest) o;
86-
return Objects.equal(targetId, sampleRowKeysRequest.targetId);
110+
return Objects.equal(targetId, sampleRowKeysRequest.targetId)
111+
&& Objects.equal(rowRange, sampleRowKeysRequest.rowRange);
87112
}
88113

89114
@Override
90115
public int hashCode() {
91-
return Objects.hashCode(targetId);
116+
return Objects.hashCode(targetId, rowRange);
117+
}
118+
119+
public static final class Builder {
120+
private TargetId targetId;
121+
private ByteStringRange rowRange = ByteStringRange.unbounded();
122+
123+
private Builder() {}
124+
125+
private Builder(SampleRowKeysRequest request) {
126+
this.targetId = request.targetId;
127+
this.rowRange = request.rowRange;
128+
}
129+
130+
public Builder setTargetId(TargetId targetId) {
131+
this.targetId = targetId;
132+
return this;
133+
}
134+
135+
public Builder setRowRange(ByteStringRange rowRange) {
136+
this.rowRange = Preconditions.checkNotNull(rowRange, "rowRange can't be null.");
137+
return this;
138+
}
139+
140+
public SampleRowKeysRequest build() {
141+
return new SampleRowKeysRequest(this);
142+
}
92143
}
93144
}

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -730,6 +730,42 @@ public void sampleRowKeysOnAuthorizedViewTest() {
730730
SampleRowKeysRequest.create(AuthorizedViewId.of("fake-table", "fake-authorized-view")));
731731
}
732732

733+
@Test
734+
public void proxySampleRowKeysWithRangeTest() {
735+
Mockito.when(mockStub.sampleRowKeysCallableWithRequest())
736+
.thenReturn(mockSampleRowKeysCallableWithRequest);
737+
738+
ByteStringRange range = ByteStringRange.create("a", "b");
739+
@SuppressWarnings("VariableUnused")
740+
ApiFuture<?> ignored = bigtableDataClient.sampleRowKeysAsync(TableId.of("fake-table"), range);
741+
742+
Mockito.verify(mockSampleRowKeysCallableWithRequest)
743+
.futureCall(
744+
SampleRowKeysRequest.newBuilder()
745+
.setTargetId(TableId.of("fake-table"))
746+
.setRowRange(range)
747+
.build());
748+
}
749+
750+
@Test
751+
public void sampleRowKeysWithRangeTest() {
752+
Mockito.when(mockStub.sampleRowKeysCallableWithRequest())
753+
.thenReturn(mockSampleRowKeysCallableWithRequest);
754+
755+
Mockito.when(
756+
mockSampleRowKeysCallableWithRequest.futureCall(
757+
ArgumentMatchers.any(SampleRowKeysRequest.class)))
758+
.thenReturn(ApiFutures.immediateFuture(Collections.<KeyOffset>emptyList()));
759+
ByteStringRange range = ByteStringRange.create("a", "b");
760+
bigtableDataClient.sampleRowKeys(TableId.of("fake-table"), range);
761+
Mockito.verify(mockSampleRowKeysCallableWithRequest)
762+
.futureCall(
763+
SampleRowKeysRequest.newBuilder()
764+
.setTargetId(TableId.of("fake-table"))
765+
.setRowRange(range)
766+
.build());
767+
}
768+
733769
@Test
734770
public void proxyMutateRowCallableTest() {
735771
Mockito.when(mockStub.mutateRowCallable()).thenReturn(mockMutateRowCallable);

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,9 @@
2727
import com.google.cloud.bigtable.data.v2.BigtableDataClient;
2828
import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId;
2929
import com.google.cloud.bigtable.data.v2.models.KeyOffset;
30+
import com.google.cloud.bigtable.data.v2.models.Range;
3031
import com.google.cloud.bigtable.data.v2.models.RowMutation;
32+
import com.google.cloud.bigtable.data.v2.models.TableId;
3133
import com.google.cloud.bigtable.test_helpers.env.EmulatorEnv;
3234
import com.google.cloud.bigtable.test_helpers.env.PrefixGenerator;
3335
import com.google.cloud.bigtable.test_helpers.env.TestEnvRule;
@@ -127,4 +129,41 @@ private static AuthorizedView createPreSplitTableAndAuthorizedView() {
127129
.setDeletionProtection(false);
128130
return testEnvRule.env().getTableAdminClient().createAuthorizedView(request);
129131
}
132+
133+
@Test
134+
public void testWithRowRange() throws InterruptedException, ExecutionException, TimeoutException {
135+
String tableId =
136+
createPreSplitTable(
137+
"SampleRowsIT#RowRange", "apple", "banana", "cherry", "date", "eggplant");
138+
BigtableDataClient client = testEnvRule.env().getDataClient();
139+
140+
try {
141+
Range.ByteStringRange range = Range.ByteStringRange.create("banana", "date");
142+
143+
ApiFuture<List<KeyOffset>> future = client.sampleRowKeysAsync(TableId.of(tableId), range);
144+
145+
List<KeyOffset> results = future.get(1, TimeUnit.MINUTES);
146+
147+
List<ByteString> resultKeys = new ArrayList<>();
148+
for (KeyOffset keyOffset : results) {
149+
resultKeys.add(keyOffset.getKey());
150+
}
151+
152+
assertThat(resultKeys)
153+
.containsExactly(ByteString.copyFromUtf8("cherry"), ByteString.copyFromUtf8("date"));
154+
155+
} finally {
156+
testEnvRule.env().getTableAdminClient().deleteTable(tableId);
157+
}
158+
}
159+
160+
private static String createPreSplitTable(String prefix, String... splitKeys) {
161+
String tableId = PrefixGenerator.newPrefix(prefix);
162+
CreateTableRequest request = CreateTableRequest.of(tableId);
163+
for (String splitKey : splitKeys) {
164+
request.addSplit(ByteString.copyFromUtf8(splitKey));
165+
}
166+
testEnvRule.env().getTableAdminClient().createTable(request);
167+
return tableId;
168+
}
130169
}

0 commit comments

Comments
 (0)