Skip to content

Commit c2dd5fe

Browse files
rahulKQLrahulKQL
committed
Updated as per gax BatchingDescriptor changes
This is still in progess. Adding these changes to showcase the implementation, test cases for BulkRead API.
1 parent fdcf724 commit c2dd5fe

File tree

8 files changed

+271
-42
lines changed

8 files changed

+271
-42
lines changed

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

Lines changed: 39 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -945,16 +945,52 @@ public Batcher<RowMutationEntry, Void> newBulkMutationBatcher(@Nonnull String ta
945945
*
946946
* <pre>{@code
947947
* try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
948-
* Filters.Filter filter = Filters.FILTERS.chain().filter()..;
949948
* List<ApiFuture<Row>> rows = new ArrayList<>();
950949
*
951-
* try (Batcher<ByteString, Row> batcher = bigtableDataClient.newBulkReadRowsBatcher("[TABLE]", filter )) {
950+
* try (Batcher<ByteString, Row> batcher = bigtableDataClient.newBulkReadRowsBatcher("[TABLE]")) {
952951
* for (String someValue : someCollection) {
953952
* ApiFuture<Row> rowFuture =
954953
* batcher.add(ByteString.copyFromUtf8("[ROW KEY]"));
955954
* rows.add(rowFuture);
956955
* }
957956
*
957+
* // Sends collected elements for batching asynchronously.
958+
* batcher.sendOutstanding();
959+
*
960+
* // Blocks until mutations are applied on all submitted row entries.
961+
* batcher.flush();
962+
* }
963+
* // Before `batcher` is closed, all remaining(If any) mutations are applied.
964+
*
965+
* List<Row> actualRows = ApiFutures.allAsList(rows).get();
966+
* }
967+
* }</pre>
968+
*/
969+
public Batcher<ByteString, Row> newBulkReadRowsBatcher(@Nonnull String tableId) {
970+
return stub.newBulkReadRowsBatcher(Query.create(tableId));
971+
}
972+
973+
/**
974+
* Reads rows in batches for given tableId and filter criteria. This method should be called in
975+
* single thread.
976+
*
977+
* <p>Sample Code:
978+
*
979+
* <pre>{@code
980+
* try (BigtableDataClient bigtableDataClient = BigtableDataClient.create("[PROJECT]", "[INSTANCE]")) {
981+
* Filters.Filter rowKeyPrefix = Filters.FILTERS.key().regex("prefix");
982+
* List<ApiFuture<Row>> rows = new ArrayList<>();
983+
*
984+
* try (Batcher<ByteString, Row> batcher = bigtableDataClient.newBulkReadRowsBatcher("[TABLE]", rowKeyPrefix)) {
985+
* for (String someValue : someCollection) {
986+
* ApiFuture<Row> rowFuture =
987+
* batcher.add(ByteString.copyFromUtf8("[ROW KEY]"));
988+
* rows.add(rowFuture);
989+
* }
990+
*
991+
* // Sends collected elements for batching asynchronously.
992+
* batcher.sendOutstanding();
993+
*
958994
* // Blocks until mutations are applied on all submitted row entries.
959995
* batcher.flush();
960996
* }
@@ -966,7 +1002,7 @@ public Batcher<RowMutationEntry, Void> newBulkMutationBatcher(@Nonnull String ta
9661002
*/
9671003
public Batcher<ByteString, Row> newBulkReadRowsBatcher(
9681004
@Nonnull String tableId, Filters.Filter filter) {
969-
return stub.newBulkReadRowsBatcher(tableId, filter);
1005+
return stub.newBulkReadRowsBatcher(Query.create(tableId).filter(filter));
9701006
}
9711007

9721008
/**

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

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,6 @@
3838
import com.google.cloud.bigtable.data.v2.models.BulkMutation;
3939
import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation;
4040
import com.google.cloud.bigtable.data.v2.models.DefaultRowAdapter;
41-
import com.google.cloud.bigtable.data.v2.models.Filters;
4241
import com.google.cloud.bigtable.data.v2.models.KeyOffset;
4342
import com.google.cloud.bigtable.data.v2.models.Query;
4443
import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
@@ -412,13 +411,11 @@ public Batcher<RowMutationEntry, Void> newMutateRowsBatcher(@Nonnull String tabl
412411
}
413412

414413
// TODO: javadoc
415-
public Batcher<ByteString, Row> newBulkReadRowsBatcher(
416-
@Nonnull String tableId, @Nonnull Filters.Filter filter) {
417-
414+
public Batcher<ByteString, Row> newBulkReadRowsBatcher(@Nonnull Query request) {
418415
return new BatcherImpl<>(
419416
settings.bulkReadRowsSettings().getBatchingDescriptor(),
420417
readRowsCallable.all(),
421-
Query.create(tableId).filter(filter),
418+
request,
422419
settings.bulkReadRowsSettings().getBatchingSettings(),
423420
clientContext.getExecutor());
424421
}

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

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -513,19 +513,24 @@ private Builder() {
513513

514514
bulkReadRowsSettings =
515515
BigtableBulkReadRowsSettings.newBuilder(new ReadRowsBatchingDescriptor())
516-
.setRetryableCodes(IDEMPOTENT_RETRY_CODES)
517-
.setRetrySettings(MUTATE_ROWS_RETRY_SETTINGS)
516+
.setRetryableCodes(readRowSettings.getRetryableCodes())
517+
.setRetrySettings(
518+
readRowSettings
519+
.getRetrySettings()
520+
.toBuilder()
521+
// Is this too high number?
522+
.setMaxRetryDelay(Duration.ofSeconds(60))
523+
.setTotalTimeout(Duration.ofMinutes(30))
524+
.build())
518525
.setBatchingSettings(
519-
// TODO: Need to finalize initial thresholds for Read
520526
BatchingSettings.newBuilder()
521527
.setIsEnabled(true)
522528
.setElementCountThreshold(100L)
523-
.setRequestByteThreshold(20L * 1024 * 1024)
529+
.setRequestByteThreshold(10_000L)
524530
.setDelayThreshold(Duration.ofSeconds(1))
525531
.setFlowControlSettings(
526532
FlowControlSettings.newBuilder()
527533
.setLimitExceededBehavior(LimitExceededBehavior.Block)
528-
.setMaxOutstandingRequestBytes(100L * 1024 * 1024)
529534
.setMaxOutstandingElementCount(1_000L)
530535
.build())
531536
.build());

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsBatchingDescriptor.java

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@
2323
import com.google.cloud.bigtable.data.v2.models.MutateRowsException;
2424
import com.google.cloud.bigtable.data.v2.models.MutateRowsException.FailedMutation;
2525
import com.google.cloud.bigtable.data.v2.models.RowMutationEntry;
26+
import com.google.common.base.Preconditions;
27+
import com.google.common.collect.ImmutableList;
2628
import com.google.common.collect.Maps;
2729
import java.util.List;
2830
import java.util.Map;
@@ -45,7 +47,9 @@ public BatchingRequestBuilder<RowMutationEntry, BulkMutation> newRequestBuilder(
4547
}
4648

4749
@Override
48-
public void splitResponse(Void response, List<SettableApiFuture<Void>> batch) {
50+
public void splitResponse(
51+
List<RowMutationEntry> elements, Void response, List<SettableApiFuture<Void>> batch) {
52+
Preconditions.checkState(elements.size() == batch.size());
4953
for (SettableApiFuture<Void> batchResponse : batch) {
5054
batchResponse.set(null);
5155
}
@@ -95,14 +99,22 @@ public long countBytes(RowMutationEntry entry) {
9599
*/
96100
static class RequestBuilder implements BatchingRequestBuilder<RowMutationEntry, BulkMutation> {
97101
private BulkMutation bulkMutation;
102+
private ImmutableList.Builder<RowMutationEntry> elements;
98103

99104
RequestBuilder(BulkMutation prototype) {
100105
this.bulkMutation = prototype.clone();
106+
this.elements = ImmutableList.builder();
101107
}
102108

103109
@Override
104110
public void add(RowMutationEntry entry) {
105111
bulkMutation.add(entry);
112+
elements.add(entry);
113+
}
114+
115+
@Override
116+
public List<RowMutationEntry> getElements() {
117+
return elements.build();
106118
}
107119

108120
@Override

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsBatchingDescriptor.java

Lines changed: 36 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,19 @@
2121
import com.google.cloud.bigtable.data.v2.models.Query;
2222
import com.google.cloud.bigtable.data.v2.models.Row;
2323
import com.google.common.base.Preconditions;
24+
import com.google.common.collect.ImmutableList;
2425
import com.google.protobuf.ByteString;
26+
import java.util.HashMap;
2527
import java.util.List;
28+
import java.util.Map;
2629

27-
// TODO: JavaDoc
30+
/**
31+
* Implementation for {@link BatchingDescriptor} to split batch response or exception into
32+
* individual row request.
33+
*
34+
* <p>This class is considered an internal implementation detail and not meant to be used by
35+
* applications directly.
36+
*/
2837
public class ReadRowsBatchingDescriptor
2938
implements BatchingDescriptor<ByteString, Row, Query, List<Row>> {
3039

@@ -34,18 +43,30 @@ public BatchingRequestBuilder<ByteString, Query> newRequestBuilder(Query query)
3443
}
3544

3645
@Override
37-
public void splitResponse(List<Row> rows, List<SettableApiFuture<Row>> futureResults) {
38-
// TODO: verify if this check
39-
Preconditions.checkState(rows.size() == futureResults.size());
46+
public void splitResponse(
47+
List<ByteString> requestElements,
48+
List<Row> rowsResponse,
49+
List<SettableApiFuture<Row>> futureResults) {
50+
Preconditions.checkState(requestElements.size() == futureResults.size());
51+
52+
Map<ByteString, Row> rowKeysMap = new HashMap<>();
53+
for (Row row : rowsResponse) {
54+
rowKeysMap.put(row.getKey(), row);
55+
}
56+
57+
for (int index = 0; index < futureResults.size(); index++) {
58+
Row row = rowKeysMap.get(requestElements.get(index));
4059

41-
for (int rowCount = 0; rowCount < rows.size(); rowCount++) {
42-
futureResults.get(rowCount).set(rows.get(rowCount));
60+
if (row == null) {
61+
// RowKey were not existent
62+
futureResults.get(index).set(null);
63+
}
64+
futureResults.get(index).set(row);
4365
}
4466
}
4567

4668
@Override
4769
public void splitException(Throwable throwable, List<SettableApiFuture<Row>> list) {
48-
// TODO: do we have to consider specific exception?
4970
for (SettableApiFuture<Row> row : list) {
5071
row.setException(throwable);
5172
}
@@ -57,16 +78,23 @@ public long countBytes(ByteString bytes) {
5778
}
5879

5980
static class BulkReadRowsRequestBuilder implements BatchingRequestBuilder<ByteString, Query> {
60-
6181
private final Query query;
82+
private final ImmutableList.Builder<ByteString> rowKeys;
6283

6384
BulkReadRowsRequestBuilder(Query query) {
6485
this.query = query.clone();
86+
this.rowKeys = ImmutableList.builder();
6587
}
6688

6789
@Override
6890
public void add(ByteString rowKey) {
6991
query.rowKey(rowKey);
92+
rowKeys.add(rowKey);
93+
}
94+
95+
@Override
96+
public List<ByteString> getElements() {
97+
return rowKeys.build();
7098
}
7199

72100
@Override

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

Lines changed: 21 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -15,20 +15,19 @@
1515
*/
1616
package com.google.cloud.bigtable.data.v2.it;
1717

18+
import static com.google.common.truth.Truth.assertThat;
19+
1820
import com.google.api.core.ApiFuture;
1921
import com.google.api.core.ApiFutures;
2022
import com.google.api.gax.batching.Batcher;
2123
import com.google.cloud.bigtable.data.v2.BigtableDataClient;
2224
import com.google.cloud.bigtable.data.v2.models.BulkMutation;
23-
import com.google.cloud.bigtable.data.v2.models.Filters;
2425
import com.google.cloud.bigtable.data.v2.models.Row;
2526
import com.google.cloud.bigtable.data.v2.models.RowCell;
2627
import com.google.cloud.bigtable.data.v2.models.RowMutationEntry;
2728
import com.google.cloud.bigtable.test_helpers.env.TestEnvRule;
2829
import com.google.common.collect.ImmutableList;
29-
import com.google.common.truth.Truth;
3030
import com.google.protobuf.ByteString;
31-
import java.io.IOException;
3231
import java.util.ArrayList;
3332
import java.util.List;
3433
import java.util.UUID;
@@ -46,18 +45,16 @@ public class BulkReadIT {
4645
@Test
4746
public void testBulkRead() throws InterruptedException, ExecutionException {
4847
BigtableDataClient client = testEnvRule.env().getDataClient();
49-
String tableId = testEnvRule.env().getTableId();
5048
String family = testEnvRule.env().getFamilyId();
5149
String rowPrefix = UUID.randomUUID().toString();
5250
int numRows = 10;
5351

54-
BulkMutation bulkMutation = BulkMutation.create(tableId);
52+
BulkMutation bulkMutation = BulkMutation.create(testEnvRule.env().getTableId());
5553
List<Row> expectedRows = new ArrayList<>();
5654
for (int i = 0; i < numRows; i++) {
5755
bulkMutation.add(
5856
RowMutationEntry.create(rowPrefix + "-" + i)
5957
.setCell(family, "qualifier", 10_000L, "value-" + i));
60-
6158
expectedRows.add(
6259
Row.create(
6360
ByteString.copyFromUtf8(rowPrefix + "-" + i),
@@ -69,20 +66,32 @@ public void testBulkRead() throws InterruptedException, ExecutionException {
6966
ImmutableList.<String>of(),
7067
ByteString.copyFromUtf8("value-" + i)))));
7168
}
72-
7369
client.bulkMutateRows(bulkMutation);
7470

7571
List<ApiFuture<Row>> rowFuture = new ArrayList<>();
76-
Filters.Filter myFilters = Filters.FILTERS.family().exactMatch(family);
77-
try (Batcher<ByteString, Row> batcher = client.newBulkReadRowsBatcher(tableId, myFilters)) {
72+
try (Batcher<ByteString, Row> batcher =
73+
client.newBulkReadRowsBatcher(testEnvRule.env().getTableId())) {
7874

7975
for (int rowCount = 0; rowCount < numRows; rowCount++) {
8076
rowFuture.add(batcher.add(ByteString.copyFromUtf8(rowPrefix + "-" + rowCount)));
8177
}
82-
}
78+
batcher.sendOutstanding();
79+
List<Row> actualRows = ApiFutures.allAsList(rowFuture).get();
80+
assertThat(actualRows).isEqualTo(expectedRows);
81+
82+
rowFuture = new ArrayList<>();
83+
// duplicate row key
84+
rowFuture.add(batcher.add(ByteString.copyFromUtf8(rowPrefix + "-" + 0)));
85+
rowFuture.add(batcher.add(ByteString.copyFromUtf8(rowPrefix + "-" + 0)));
8386

84-
List<Row> actualRows = ApiFutures.allAsList(rowFuture).get();
87+
// non-existent row key
88+
rowFuture.add(batcher.add(ByteString.copyFromUtf8(UUID.randomUUID().toString())));
8589

86-
Truth.assertThat(actualRows).isEqualTo(expectedRows);
90+
batcher.sendOutstanding();
91+
actualRows = ApiFutures.allAsList(rowFuture).get();
92+
assertThat(actualRows.get(0)).isEqualTo(expectedRows.get(0));
93+
assertThat(actualRows.get(1)).isEqualTo(expectedRows.get(0));
94+
assertThat(actualRows.get(2)).isEqualTo(null);
95+
}
8796
}
8897
}

‎google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/mutaterows/MutateRowsBatchingDescriptorTest.java

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,14 @@ public class MutateRowsBatchingDescriptorTest {
4141
private static final String FAMILY = "fake-family";
4242
private static final String QUALIFIER = "fake-qualifier";
4343
private static final String VALUE = "fake-value";
44+
private static final long TIMESTAMP = 10_000L;
4445

4546
private static final RequestContext requestContext =
4647
RequestContext.create("fake-project", "fake-instance", "fake-profile");
48+
private static final List<RowMutationEntry> ENTRIES =
49+
ImmutableList.of(
50+
RowMutationEntry.create(ROW_KEY).setCell(FAMILY, QUALIFIER, TIMESTAMP, VALUE),
51+
RowMutationEntry.create("rowKey-2").setCell("family-2", "q", TIMESTAMP, VALUE));
4752

4853
@Test
4954
public void countBytesTest() {
@@ -57,22 +62,20 @@ public void countBytesTest() {
5762
@Test
5863
public void requestBuilderTest() {
5964
MutateRowsBatchingDescriptor underTest = new MutateRowsBatchingDescriptor();
60-
long timestamp = 10_000L;
6165
BulkMutation bulkMutation = BulkMutation.create("fake-table");
6266
BatchingRequestBuilder<RowMutationEntry, BulkMutation> requestBuilder =
6367
underTest.newRequestBuilder(bulkMutation);
64-
requestBuilder.add(
65-
RowMutationEntry.create(ROW_KEY).setCell(FAMILY, QUALIFIER, timestamp, VALUE));
66-
requestBuilder.add(
67-
RowMutationEntry.create("rowKey-2").setCell("family-2", "q", 20_000L, "some-value"));
68+
requestBuilder.add(ENTRIES.get(0));
69+
requestBuilder.add(ENTRIES.get(1));
6870

6971
BulkMutation actualBulkMutation = requestBuilder.build();
7072
assertThat(actualBulkMutation.toProto(requestContext))
7173
.isEqualTo(
7274
BulkMutation.create("fake-table")
73-
.add(ROW_KEY, Mutation.create().setCell(FAMILY, QUALIFIER, timestamp, VALUE))
74-
.add("rowKey-2", Mutation.create().setCell("family-2", "q", 20_000L, "some-value"))
75+
.add(ROW_KEY, Mutation.create().setCell(FAMILY, QUALIFIER, TIMESTAMP, VALUE))
76+
.add("rowKey-2", Mutation.create().setCell("family-2", "q", TIMESTAMP, VALUE))
7577
.toProto(requestContext));
78+
assertThat(requestBuilder.getElements()).isEqualTo(ENTRIES);
7679
}
7780

7881
@Test
@@ -83,7 +86,7 @@ public void splitResponseTest() {
8386
assertThat(batchResponse.get(1).isDone()).isFalse();
8487

8588
MutateRowsBatchingDescriptor underTest = new MutateRowsBatchingDescriptor();
86-
underTest.splitResponse(null, batchResponse);
89+
underTest.splitResponse(ENTRIES, null, batchResponse);
8790
assertThat(batchResponse.get(0).isDone()).isTrue();
8891
assertThat(batchResponse.get(1).isDone()).isTrue();
8992
}

0 commit comments

Comments
 (0)